# FedSoL: Bridging Global Alignment and Local Generality in Federated Learning

Gihun Lee1, Minchan Jeong1, Sangmook Kim1, Jaehoon Oh2, Se-Young Yun1

###### Abstract

Federated Learning (FL) aggregates locally trained models from individual
clients to construct a global model. While FL enables learning a model with
data privacy, it often suffers from significant performance degradation when
client data distributions are heterogeneous. Many previous FL algorithms have
addressed this issue by introducing various proximal restrictions. These
restrictions aim to encourage global alignment by constraining the deviation
of local learning from the global objective. However, they inherently limit
local learning by interfering with the original local objectives. Recently, an
alternative approach has emerged to improve local learning generality. By
obtaining local models within a smooth loss landscape, this approach mitigates
conflicts among different local objectives of the clients. Yet, it does not
ensure stable global alignment, as local learning does not take the global
objective into account. In this study, we propose Federated Stability on
Learning (FedSoL), which combines both the concepts of global alignment and
local generality. In FedSoL, the local learning seeks a parameter region
robust against proximal perturbations. This strategy introduces an implicit
proximal restriction effect in local learning while maintaining the original
local objective for parameter update. Our experiments show that FedSoL
consistently achieves state-of-the-art performance on various setups.

##  1 Introduction

Federated Learning (FL) is an emerging distributed learning framework that
preserves data privacy while leveraging client data for training (Konečnỳ et
al. 2016b, a). In this approach, individual clients train their local models
using their private data, while the server aggregates these models into a
global model. By precluding the need for direct access to private data, FL
enables the utilization of extensive data collected from edge devices such as
mobile phones, vehicles, and facilities (Aledhari et al. 2020; Yang et al.
2019).

However, FL encounters a notorious challenge known as data heterogeneity
(Kairouz et al. 2019). Due to the diverse underlying distributions of the
clients, the local datasets are non-independent and identically distributed
(Non-IID). Its inevitable occurrence in many real-world scenarios leads to an
inconsistency between global and local objectives, often significantly
degrading performance (Li et al. 2019, 2020a).

To tackle the data heterogeneity problem, most prior studies have introduced
various proximal restriction into the local objective (Karimireddy et al.
2020; Li, He, and Song 2021; Lee et al. 2021; Li et al. 2020b). These
restrictions are designed to maintain alignment between global and local
objectives by preventing the deviation of local learning from the global
objective. Nevertheless, this approach inherently limits local learning by
interfering with the original local objectives (Mendieta et al. 2022; Zhang et
al. 2022). Furthermore, it often falls into a sharp global landscape under
high heterogeneity, which results in unreliable minima and poor stability (Qu
et al. 2022; Sun et al. 2023a) (Figure1(a)).

In contrast, an alternative approach has emerged that focuses on the local
generality in FL (Qu et al. 2022; Mendieta et al. 2022). Building on recent
findings that highlight the benefits of smooth loss landscape for better
generalization (Izmailov et al. 2018; Jiang et al. 2019; Chaudhari et al.
2019), this approach aims to seek flatness during local learning by employing
the recently proposed Sharpness-Aware Minimization (SAM) (Foret et al. 2020)
as the local optimizer (Qu et al. 2022). By improving local generality, this
approach mitigates the conflicts between individual local objectives,
contributing to the overall smoothness of the aggregated global model
(Caldarola, Caputo, and Ciccone 2022; Sun et al. 2023a). Although these
approaches have demonstrated competitive performance without the proximal
restrictions, their ability to generalize well within their respective local
distributions does not necessarily ensure alignment with the global objective
(Figure1(b)).

Figure 1: An overview of FL scenarios. Shaded regions stand for the degree of
local generality of the trained local models. (a) shows that the local models
are aligned with the global objective via proximal restriction, but the local
generality is limited. On the other hand, (b) presents enhanced local
generality but the local learning deviates from the global objective. In (c),
FedSoL integrates the strengths of both scenarios, ensuring global alignment
while promoting local generality.

Our key motivation is to tackle data heterogeneity by harnessing both the
strengths of global alignment and local generality. To this end, we propose a
novel algorithm Federated Stability on Learning (FedSoL) (Figure1(c)). FedSoL
seamlessly incorporates the proximal restriction effect into the SAM
optimization, without interfering with the original local objective. More
specifically, FedSoL updates the local model using the gradient of the
original local objective, which is determined at the weights perturbed by the
gradient of the proximal restriction objective. By identifying a parameter
region that is minimally influenced by the proximal perturbation, FedSoL
diminishes the negative impact of local updates on global alignment, while
maintaining the original local objective to preserve local generality.
Furthermore, we control the perturbation strength for each parameter by
reflecting the discrepancy between the global and local model parameters. We
comprehensively demonstrate the efficacy of FedSoL on global alignment and
local generality. Experimental results show FedSoL achieves state-of-the-art
performance in various setups. To summarize, our main contributions are as
follows:

  * •

We propose FedSoL, a novel and effective FL algorithm that leverages both the
strengths of global alignment and local generality. FedSoL conducts proximal
perturbations with the SAM strategy, yielding an implicit regularization
effect during local learning. (Section 3)

  * •

We validate the efficacy of FedSoL on various setups and show that it
consistently achieves state-of-the-art performances. We highlight FedSoL
performs exceptionally well under high levels of heterogeneity. (Section 4)

  * •

We provide a comprehensive analysis of the benefits that FedSoL brings to FL.
Not only does it enhance the smoothness of the global model, but it also
preserves global knowledge during local learning. (Section 5)

##  2 Background

###  2.1 Proximal Restriction

Consider an FL system that consists of $K$ clients and a central server. Each
client $k$ has a local dataset $\mathcal{D}^{k}$, where the entire dataset is
a union of the local datasets as
$\mathcal{D}={\bigcup}_{k\in[K]}\mathcal{D}^{k}$. FL aims to train a global
server model with weights $\boldsymbol{w}$ that minimize the loss across all
clients:

$\vspace{-1pt}\mathcal{L}_{\mathrm{global}}(\boldsymbol{w})=\underset{{k\in[K]}}{{\sum}}\,\frac{|\mathcal{D}_{k}|}{|\mathcal{D}|}\,\mathcal{L}_{\mathrm{local}}^{k}(\boldsymbol{w})\,,$
(1)

where $|\mathcal{D}^{k}|$ and $|\mathcal{D}|$ are the number of instances in
each datasets. When using a proximal restriction objective, the loss function
for each client $k$ is a linear combination of its original local loss,
$\mathcal{L}_{\mathrm{local}}^{k}(\boldsymbol{w}_{k})$, and a proximal loss,
$\mathcal{L}_{p}^{k}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})$, controlled by a
hyperparameter $\beta$:

$\mathcal{L}^{k}(\boldsymbol{w}_{k})=\mathcal{L}_{\mathrm{local}}^{k}(\boldsymbol{w}_{k})+\beta\cdot\mathcal{L}_{p}^{k}(\boldsymbol{w}_{k};\boldsymbol{w}_{g}).$
(2)

Here, $\mathcal{L}_{\mathrm{local}}^{k}(\boldsymbol{w}_{k})$ is the loss on
the client’s local distribution (e.g., cross-entropy loss), and
$\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})$ quantifies the
discrepancy between the global model $\boldsymbol{w}_{g}$ and the local model
$\boldsymbol{w}_{k}$. This discrepancy can be measured in various ways, such
as the Euclidean distance between the two parameters (Li et al. 2020b) or the
KL-divergence between probability vectors computed using the client’s data
(Lee et al. 2021). Introducing proximal restriction within the local objective
constrains the deviation of local learning from the global objective
(Karimireddy et al. 2020). However, such modified local objective inherently
limits the acquisition of new knowledge on the local distribution (Qu et al.
2022).

###  2.2 Overview of SAM

The SAM (Foret et al. 2020) optimizer pursues flatter minima for the given
loss $\mathcal{L}$ by solving the minimax problem:

$\underset{\boldsymbol{w}}{\text{min}}\;\underset{{\|\boldsymbol{\epsilon}\|}_{2}<\rho}{\text{max}}\mathcal{L}(\boldsymbol{w}+\boldsymbol{\epsilon}).$
(3)

In the above equation, the inner maximization finds a parameter perturbation
$\boldsymbol{\epsilon}$ that induces maximal loss change within the
$\rho$-ball neighborhood. In practice, it is approximated by a single re-
scaled gradient step
${\boldsymbol{\epsilon}}^{*}$$=$$\rho{\nabla}_{\boldsymbol{w}}\mathcal{L}(\boldsymbol{w})/\|{\nabla}_{\boldsymbol{w}}\mathcal{L}(\boldsymbol{w})\|_{2}$.
Then, the outer minimization is conducted by a base optimizer such as SGD
(Nesterov 1983), by taking the gradient
${\nabla}_{\boldsymbol{w}}\mathcal{L}(\boldsymbol{w}+{\boldsymbol{\epsilon}}^{*})$
at the perturbed weights. SAM demonstrates an exceptional ability to perform
well across different model structures (Chen, Hsieh, and Gong 2021; Zhuang et
al. 2022) and tasks (Wang et al. 2023; Abbas et al. 2022) with high
generalization performance. In FL, using SAM enhances the generalization of
each client’s local model (Qu et al. 2022; Caldarola, Caputo, and Ciccone
2022). However, the deviation of local learning hinders the trained local
models to align with each other, preventing the aggregated global model from
fully leveraging the advantages of their smoother landscapes (Sun et al.
2023b, a).

##  3 FedSoL: Federated Stability on Learning

In this section, we introduce Federated Stability on Learning (FedSoL). Our
primary motivation is to integrate the proximal restriction effect into the
SAM optimization without interfering with the original local objective. The
detailed procedure is outlined in Algorithm 1.

###  3.1 Proximal Perturbation

In local learning, the local model $\boldsymbol{w}_{k}$ begins with the same
parameters as the distributed global model $\boldsymbol{w}_{g}$, thereby
initially having a minimal proximal loss. The main challenge is guiding the
local learning to reduce the original local loss
$\mathcal{L}^{k}_{\mathrm{local}}$ without inducing an increase in the
proximal loss $\mathcal{L}^{k}_{p}$. We address this problem in the context of
SAM optimization by seeking a gradient that not only minimizes the original
local loss $\mathcal{L}^{k}_{\mathrm{local}}$, but also is robust against the
increases in proximal loss $\mathcal{L}^{k}_{p}$. To this end, FedSoL aims to
minimize the original local loss, which is minimally affected by the weight
perturbation that maximizes proximal loss. By decoupling the roles of these
two types of losses, FedSoL conducts the following steps for each local update
of client $k$:

#### Step1: Weight Perturbation

FedSoL finds a weight perturbation $\boldsymbol{\epsilon}^{*}_{p}$ that causes
the most significant change for any given proximal loss $\mathcal{L}^{k}_{p}$:

$\boldsymbol{\epsilon}^{*}_{p}=\rho\scalebox{0.9}{$\displaystyle\frac{{\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})}{\|{\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})\|_{2}}$}\\!\approx\underset{\|\boldsymbol{\epsilon}\|_{2}\leq\rho}{\text{argmax}}\;\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k}\\!+\\!\boldsymbol{\epsilon};\boldsymbol{w}_{g})\,.$
(4)

#### Step2: Parameter Update

After perturbation, FedSoL updates parameters by computing the gradient of the
original local loss $\mathcal{L}^{k}_{\mathrm{local}}$ at this perturbed
weights:

$\boldsymbol{w}_{k}\leftarrow\boldsymbol{w}_{k}-\gamma\cdot{\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}^{*}_{p})\,,$
(5)

where $\gamma$ is a learning rate. In the above procedures, the update
gradient is computed on the original local loss
$\mathcal{L}^{k}_{\mathrm{local}}$, whereas the proximal loss
$\mathcal{L}_{{p}}^{k}$ only plays in an implicit role. Note that
$\boldsymbol{\epsilon}^{*}_{p}$ is used solely for weight perturbation,
thereby we do not need to compute its gradient. To clarify how FedSoL
influences local learning, we analyze the update gradient $\boldsymbol{g}_{u}$
of FedSoL by employing the first-order Taylor approximation of
$\mathcal{L}^{k}_{\mathrm{local}}$ at $\boldsymbol{w}_{k}$, with perturbation
$\boldsymbol{\epsilon}^{*}_{p}$:

$\displaystyle\boldsymbol{g}_{u}(\boldsymbol{w}_{k})$
$\displaystyle={\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}^{*}_{p})$
$\displaystyle\approx\boldsymbol{g}_{l}(\boldsymbol{w}_{k})+\rho\,\nabla^{2}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\hat{\boldsymbol{g}}_{p}(\boldsymbol{w}_{k})\,.$
(6)

We denote
$\boldsymbol{g}_{l}$$=$$\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})$,
$\boldsymbol{g}_{p}$$=$$\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};{\boldsymbol{w}_{g}})$,
and
$\hat{\boldsymbol{g}}_{p}$$=$$\boldsymbol{g}_{p}/\|\boldsymbol{g}_{p}\|_{2}$,
omitting $\boldsymbol{w}_{k}$ and ${\boldsymbol{w}_{g}}$ if there is no
conflict. Based on Equation (6), we examine the change of each loss induced by
a single local update with a learning rate $\gamma$, as defined in Equation
(7):

$\displaystyle\Delta^{\mathrm{algo}}\mathcal{L}^{k}(\boldsymbol{w}_{k})$
$\displaystyle=\mathcal{L}^{k}\big{(}\boldsymbol{w}_{k}-\gamma\boldsymbol{g}_{u}(\boldsymbol{w}_{k})\big{)}-\mathcal{L}^{k}(\boldsymbol{w}_{k})$
$\displaystyle\approx-\gamma\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}(\boldsymbol{w}_{k}),\boldsymbol{g}_{u}(\boldsymbol{w}_{k})\rangle\,.$
(7)

In the above equation, we apply a first-order Taylor approximation to
$\mathcal{L}^{k}$ at $\boldsymbol{w}_{k}$, taking into account a local update
by $-\gamma\boldsymbol{g}_{u}(\boldsymbol{w}_{k})$. Here, $\mathcal{L}^{k}$
can be either local loss or proximal loss. By combining the approximation for
$\boldsymbol{g_{u}}$ from Equation (6) into the loss difference in Equation
(7), we derive the following two key propositions.

###### Proposition 1

Given a convex local loss $\mathcal{L}^{k}_{\mathrm{local}}$, the change of
proximal loss $\mathcal{L}^{k}_{p}$ by FedSoL update with a learning rate
$\gamma$ reduces the conflicts between $\boldsymbol{g}_{l}$ and
$\boldsymbol{g}_{p}$ as $\rho$ grows:

$\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{p}\approx-\gamma\Big{(}\langle\boldsymbol{g}_{l}\,,\boldsymbol{g}_{p}\rangle+\rho\cdot\mathop{\vtop{\halign{#\cr$\hfil\displaystyle{\hat{\boldsymbol{g}}_{p}{\\!\\!}^{\top}\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}\,{\boldsymbol{g}_{p}}}\hfil$\crcr\kern
3.0pt\nointerlineskip\cr$\bracelu\leaders\vrule
height=0.39998pt,depth=0.0pt\hfill\,\lower 1.00006pt\hbox{$\geq
0$}\,\leaders\vrule height=0.39998pt,depth=0.0pt\hfill\braceru$\crcr\kern
3.0pt\cr}}}\limits\Big{)}\,,$ (8)

where $\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}$ is the Hessian at
$\boldsymbol{w}_{k}$. In Proposition 1, we study the change of proximal loss
after the FedSoL update, denoted as
$\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{p}$. The proposition suggests that
FedSoL implicitly regularizes local learning for global alignment, reducing
the negative impact of local updates on proximal loss. This regularization
effect grows as the curvature of local loss
$\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}$ becomes steeper. Note that when
$\rho$ is set to 0, Equation (8) becomes the change of proximal loss after
FedAvg update, $\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{p}$.

Algorithm 1 Federated Stability on Learning (FedSoL)

Input: local loss $\mathcal{L}^{k}_{\mathrm{local}}$ and proximal loss
$\mathcal{L}^{k}_{p}$ for each client $k\in[K]$, learning rate $\gamma$, and
base perturbation radius $\rho$

Initialize global server weight $\boldsymbol{w}_{g}$

for each communication round $t$ do

Server samples clients $K^{(t)}\subset[K]$

Server broadcasts ${\boldsymbol{w}_{g}}$ for all $k\in K^{(t)}$

Client replaces ${\boldsymbol{w}_{k}}\leftarrow{\boldsymbol{w}_{g}}$

for each client $k\in K^{(t)}$ in parallel do

for each local step do

 # Set Adaptive Perturbation Radius (Sec 3.2)

$\boldsymbol{\rho}_{\mathrm{adaptive}}=\rho\cdot\boldsymbol{\Lambda}$
(element-wise rescale)

# Perturb using Proximal Gradient (Sec 3.1)

${\boldsymbol{\epsilon}}^{*}_{p}=\boldsymbol{\rho}_{\mathrm{adaptive}}\odot\frac{\nabla_{\boldsymbol{w}_{k}}{\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})}}{\|\nabla_{\boldsymbol{w}_{k}}{\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})}\|}$

# Update Local Model Parameters (Sec 3.1)

${\boldsymbol{w}_{k}}\leftarrow{\boldsymbol{w}_{k}}-\gamma\cdot\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}\big{(}\boldsymbol{w}_{k}+{\boldsymbol{\epsilon}}^{*}_{p}\big{)}$

end for

end for

Upload ${\boldsymbol{w}_{k}}$ to server

Server Aggregation
:${\boldsymbol{w}_{g}}\leftarrow\frac{1}{|K^{(t)}|}\sum_{k\in
K^{(t)}}{\boldsymbol{w}_{k}}$

end for

Server output : ${\boldsymbol{w}_{g}}$

###### Proposition 2

The change of original local loss $\mathcal{L}^{k}_{\mathrm{local}}$ by FedSoL
update with a learning rate $\gamma$ is equivalent to the FedAvg update at
$\nabla\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k}+\frac{\rho}{2}\boldsymbol{\epsilon}^{*}_{p})$
as:

$\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\approx\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{\mathrm{local}}\Big{(}\boldsymbol{w}_{k}+\frac{\rho}{2}\boldsymbol{\epsilon}^{*}_{p}\Big{)}\,.$
(9)

Meanwhile, Proposition 2 analyzes how the original local loss changes after
the FedSoL update, $\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{\mathrm{local}}$,
and compare it to its counterpart in FedAvg,
$\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{\mathrm{local}}$. The proposition
suggests that even though FedSoL calculates the gradient of the original local
loss at the weight perturbed by proximal gradient, its behavior on the
original local loss is almost identical to that of FedAvg with standard
gradient descent using $\mathcal{L}^{k}_{\mathrm{local}}$. This implies that
FedSoL does not noticeably interfere with or slow down the learning process on
the local data distribution. The detailed proofs for the propositions are
provided in Appendix.

###  3.2 Adaptive Perturbation Radius

While SAM defines a fixed radius $\rho$, it is often insufficient in capturing
the loss landscape dynamics (Kwon et al. 2021; Kim et al. 2022). In FedSoL, we
introduce an adaptive radius reflecting the global and local parameter
discrepancies. For each layer $m$, we construct a scaling vector
$\boldsymbol{\lambda}^{(m)}$, where the $i$-th entry corresponds to each
parameter in that layer:

$\vspace{-1pt}\boldsymbol{\lambda}^{(m)}[i]=\frac{\big{|}{\boldsymbol{w}}^{(m)}_{k}[i]-{\boldsymbol{w}}^{(m)}_{g}[i]\big{|}}{\big{\|}{\boldsymbol{w}}^{(m)}_{k}-{\boldsymbol{w}}^{(m)}_{g}\big{\|}_{2}}.$
(10)

Here, $\boldsymbol{w}_{g}^{(m)}$ and $\boldsymbol{w}_{k}^{(m)}$ denote layer
$m$ in the global and local model respectively. The denominator represents the
normalization of the discrepancy within the layer, accounting for the layer-
wise scale variance. The adaptive radius allows more perturbation for the
parameter with large difference, and vice versa. It fits with the typical
behavior of proximal loss, which increases as
$\|\boldsymbol{w}_{k}-\boldsymbol{w}_{g}\|_{2}$ grows. By concatenating these
layer-specific vectors,
$\boldsymbol{\Lambda}=(\boldsymbol{\lambda}^{(1)},\ldots,\boldsymbol{\lambda}^{(m)},\ldots,\boldsymbol{\lambda}^{(\mathrm{last})})$,
and incorporating it into the Equation (4)), the proximal perturbation
$\boldsymbol{\epsilon}^{*}_{p}$ becomes:

$\displaystyle\vspace{-1pt}{\boldsymbol{\epsilon}}^{*}_{p}$
$\displaystyle=\>\>\rho\cdot\boldsymbol{\Lambda}\odot\frac{{\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})}{\|{\nabla}_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})\|_{2}}$
$\displaystyle\approx\underset{\scalebox{0.7}{$\left\lVert\boldsymbol{\Lambda}^{-1}\\!\odot\boldsymbol{\epsilon}\right\rVert_{2}\leq\rho$}}{\text{argmax}}\;\mathcal{L}^{k}_{p}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon};\boldsymbol{w}_{g})\,,$
(11)

where $\odot$ denotes the element-wise product. Intuitively, the adaptive
radius allows local learning to deviate certain parameters from the global
model, only if they are crucial enough to withstand larger weight
perturbations.

##  4 Experiment

Table 1:  Test accuracy@1(%) comparison among baselines and FedSoL on
different datasets. The values in the parenthesis are the standard deviation.
The arrow ($\downarrow$, $\uparrow$) shows the comparison to the FedAvg. We
set $s\in\\{2,3,5,10\\}$ and $\alpha\in\\{0.05,0.1,0.3,0.5\\}$ for CIFAR-10
datasets, whereas $s=2$ and $\alpha=0.1$ for the others. Non-IID Partition
Strategy : Sharding
---
Method | MNIST | CIFAR-10 | SVHN | CINIC-10 | PathMNIST | TissueMNIST
$s$ = 2 |  $s$ = 3 |  $s$ = 5 |  $s$ = 10
FedAvg | 96.16${}_{(0.19)}\;\;\;$ | 51.48${}_{(3.41)}\;\;\;$ | 62.94${}_{(0.00)}\;\;\;$ | 70.96${}_{(0.91)}\;\;\;$ | 74.60${}_{(0.88)}\;\;\;$ | 73.63${}_{(3.16)}\;\;\;$ | 42.40${}_{(2.70)}\;\;\;$ | 57.40${}_{(1.48)}\;\;\;$ | 49.36${}_{(1.64)}\;\;\;$
FedProx | 95.86${}_{(0.12)}\;$$\downarrow$ | 52.80${}_{(2.66)}\;$$\uparrow$ | 58.19${}_{(0.55)}\;$$\downarrow$ | 64.71${}_{(0.74)}\;$$\downarrow$ | 69.37${}_{(1.21)}\;$$\downarrow$ | 71.09${}_{(3.13)}\;$$\downarrow$ | 40.00${}_{(3.01)}\;$$\downarrow$ | 60.77${}_{(3.64)}\;$$\uparrow$ | 48.20${}_{(1.95)}\;$$\downarrow$
FedNova | 94.13${}_{(0.36)}\;$$\downarrow$ | 46.89${}_{(2.57)}\;$$\downarrow$ | 61.12${}_{(0.88)}\;$$\downarrow$ | 67.11${}_{(0.25)}\;$$\downarrow$ | 70.59${}_{(0.52)}\;$$\downarrow$ | 67.35${}_{(2.84)}\;$$\downarrow$ | 40.94${}_{(2.29)}\;$$\downarrow$ | 58.85${}_{(4.10)}\;$$\uparrow$ | 36.44${}_{(0.95)}\;$$\downarrow$
Scaffold | 95.91${}_{(0.18)}\;$$\downarrow$ | 62.60${}_{(0.70)}\;$$\uparrow$ | 68.53${}_{(0.99)}\;$$\uparrow$ | 74.28${}_{(0.39)}\;$$\uparrow$ | 76.71${}_{(0.16)}\;$$\uparrow$ | 77.84${}_{(2.28)}\;$$\uparrow$ | 47.76${}_{(0.45)}\;$$\uparrow$ | 71.12${}_{(1.04)}\;$$\uparrow$ | 30.99${}_{(6.09)}\;$$\downarrow$
FedNTD | 96.62${}_{(0.06)}\;$$\uparrow$ |  67.25${}_{(1.08)}\;$$\uparrow$ |  70.47${}_{(0.33)}\;$$\uparrow$ | 75.21${}_{(0.39)}\;$$\uparrow$ | 76.46${}_{(0.07)}\;$$\uparrow$ |  85.30${}_{(0.78)}\;$$\uparrow$ | 52.72${}_{(1.12)}\;$$\uparrow$ | 65.00${}_{(1.26)}\;$$\uparrow$ | 52.63${}_{(0.59)}\;$$\uparrow$
FedSAM | 96.12${}_{(0.19)}\;$$\downarrow$ | 51.85${}_{(3.14)}\;$$\uparrow$ | 60.90${}_{(0.93)}\;$$\downarrow$ | 69.29${}_{(0.39)}\;$$\downarrow$ | 72.98${}_{(0.34)}\;$$\downarrow$ | 65.85${}_{(3.77)}\;$$\downarrow$ | 45.91${}_{(2.02)}\;$$\uparrow$ | 67.32${}_{(3.15)}\;$$\uparrow$ | 49.62${}_{(1.61)}\;$$\uparrow$
FedASAM |  97.08${}_{(0.15)}\;$$\uparrow$ | 52.08${}_{(2.19)}\;$$\uparrow$ | 63.24${}_{(1.16)}\;$$\uparrow$ | 70.95${}_{(0.76)}\;$$\downarrow$ | 74.74${}_{(0.88)}\;$$\uparrow$ | 79.48${}_{(2.17)}\;$$\downarrow$ | 43.15${}_{(2.73)}\;$$\uparrow$ | 59.47${}_{(2.91)}\;$$\uparrow$ | 49.46${}_{(1.91)}\;$$\uparrow$
FedSoL (Ours) |  97.15${}_{(0.08)}\;$$\uparrow$ |  66.72${}_{(0.61)}\;$$\uparrow$ |  69.88${}_{(0.15)}\;$$\uparrow$ |  75.82${}_{(0.34)}\;$$\uparrow$ |  77.79${}_{(0.19)}\;$$\uparrow$ |  85.18${}_{(0.37)}\;$$\uparrow$ |  55.17${}_{(0.32)}\;$$\uparrow$ |  73.85${}_{(1.55)}\;$$\uparrow$ |  53.42${}_{(0.46)}\;$$\uparrow$
Non-IID Partition Strategy : LDA
Method | MNIST | CIFAR-10 | SVHN | CINIC-10 | PathMNIST | TissueMNIST
$\alpha$ = 0.05 |  $\alpha$ = 0.1 |  $\alpha$ = 0.3 |  $\alpha$ = 0.5
FedAvg | 96.11${}_{(0.19)}\;\;\;$ | 42.27${}_{(1.34)}\;\;\;$ | 56.13${}_{(0.78)}\;\;\;$ | 67.32${}_{(0.94)}\;\;\;$ | 73.90${}_{(0.66)}\;\;\;$ | 55.36${}_{(4.85)}\;\;\;$ | 36.49${}_{(4.37)}\;\;\;$ | 65.98${}_{(4.76)}\;\;\;$ | 42.78${}_{(2.03)}\;\;\;$
FedProx | 96.05${}_{(0.13)}\;$$\downarrow$ | 50.58${}_{(0.57)}\;$$\uparrow$ | 59.80${}_{(1.12)}\;$$\uparrow$ | 68.39${}_{(0.81)}\;$$\uparrow$ | 72.87${}_{(0.55)}\;$$\uparrow$ | 72.40${}_{(3.15)}\;$$\uparrow$ | 40.09${}_{(3.97)}\;$$\uparrow$ | 70.44${}_{(1.92)}\;$$\uparrow$ | 52.25${}_{(1.40)}\;$$\uparrow$
FedNova | 88.24${}_{(1.37)}\;$$\downarrow$ | 10.00${}_{({\textit{Failed}})}\;$$\downarrow$ | 10.00${}_{({\textit{Failed}})}\;$$\downarrow$ | 64.67${}_{(0.77)}\;$$\downarrow$ | 70.04${}_{(0.45)}\;$$\downarrow$ | 53.07${}_{(3.30)}\;$$\downarrow$ | 21.89${}_{(1.71)}\;$$\downarrow$ | 38.94${}_{(2.34)}\;$$\downarrow$ | 15.03${}_{(3.74)}\;$$\downarrow$
Scaffold | 94.18${}_{(0.32)}\;$$\downarrow$ | 10.00${}_{({\textit{Failed}})}\;$$\downarrow$ | 10.00${}_{({\textit{Failed}})}\;$$\downarrow$ |  71.92${}_{(0.17)}\;$$\uparrow$ |  75.49${}_{(0.21)}\;$$\uparrow$ | 21.46${}_{(1.75)}\;$$\downarrow$ | 16.89${}_{(2.25)}\;$$\downarrow$ | 18.07${}_{(0.04)}\;$$\downarrow$ | 32.04${}_{(0.07)}\;$$\downarrow$
FedNTD | 96.97${}_{(0.27)}\;$$\uparrow$ | 58.08${}_{(0.48)}\;$$\uparrow$ |  63.16${}_{(1.02)}\;$$\uparrow$ | 71.56${}_{(0.26)}\;$$\uparrow$ | 74.91${}_{(0.33)}\;$$\uparrow$ | 79.25${}_{(0.61)}\;$$\uparrow$ | 50.22${}_{(3.71)}\;$$\uparrow$ | 74.26${}_{(1.25)}\;$$\uparrow$ | 44.55${}_{(1.95)}\;$$\uparrow$
FedSAM | 95.72${}_{(0.43)}\;$$\downarrow$ | 36.14${}_{(1.21)}\;$$\downarrow$ | 52.14${}_{(0.94)}\;$$\downarrow$ | 64.83${}_{(0.56)}\;$$\downarrow$ | 70.74${}_{(0.40)}\;$$\downarrow$ | 13.27${}_{(2.78)}\;$$\downarrow$ | 36.70${}_{(4.28)}\;$$\uparrow$ | 66.64${}_{(3.76)}\;$$\uparrow$ | 44.07${}_{(3.02)}\;$$\uparrow$
FedASAM | 96.60${}_{(0.10)}\;$$\uparrow$ | 43.12${}_{(1.25)}\;$$\uparrow$ | 57.00${}_{(0.30)}\;$$\uparrow$ | 67.45${}_{(0.92)}\;$$\uparrow$ | 73.91${}_{(0.51)}\;$$\uparrow$ | 60.25${}_{(4.56)}\;$$\uparrow$ | 36.93${}_{(4.60)}\;$$\uparrow$ | 69.45${}_{(3.19)}\;$$\uparrow$ | 42.73${}_{(2.35)}\;$$\uparrow$
FedSoL (Ours) |  97.44${}_{(0.11)}\;$$\uparrow$ |  60.01${}_{(0.30)}\;$$\uparrow$ |  64.13${}_{(0.46)}\;$$\uparrow$ |  71.94${}_{(0.57)}\;$$\uparrow$ |  75.60${}_{(0.32)}\;$$\uparrow$ |  83.92${}_{(0.29)}\;$$\uparrow$ |  55.07${}_{(1.48)}\;$$\uparrow$ |  78.88${}_{(0.46)}\;$$\uparrow$ |  53.40${}_{(0.85)}\;$$\uparrow$

###  4.1 Experimental Setups

#### Data Setups

We employ 6 datasets: MNIST (Deng 2012), CIFAR-10 (Krizhevsky, Nair, and
Hinton 2009), SVHN (Netzer et al. 2011), CINIC-10 (Darlow et al. 2018),
PathMNIST (Yang et al. 2023), and TissueMNIST (Yang et al. 2023). We
distribute data to clients via two strategies: Sharding (McMahan et al. 2017)
and Latent Dirichlet Allocation (LDA) (Wang et al. 2020a). Sharding sorts data
by label and assigns equal-size shards to clients. The heterogeneity increases
as the shard per user, $s$, becomes smaller. On the other hand, LDA assigns
class $c$ data samples to each client $k$ with probability
$p_{c}(\approx\text{Dir}(\alpha))$, where the heterogeneity increases as
$\alpha$, becomes smaller. Figure 2 illustrates the difference between these
partition strategies.

Figure 2: CIFAR-10 partition examples across 10 clients.

#### Learning Setups

We distribute MNIST, CIFAR-10, and SVHN datasets across 100 clients with a
sampling ratio of 0.1, while CINIC-10, PathMNIST, and TissueMNIST across 200
clients with a ratio of 0.05. We use a model architecture as described in
(McMahan et al. 2017), which consists of two convolutional layers, max-pooling
layers, and two fully connected layers. Each client optimizes its local
datasets for 5 local epochs using momentum SGD with a learning rate of 0.01,
momentum 0.9, and weight decay 1e-5. The learning rate is decayed by a factor
of 0.99 at every communication round. We conduct 300 communication rounds in
general, and 200 for MNIST, PathMNIST, and TissueMNIST, which is enough for
the server model’s performance to reach saturation. We use KL-divergence loss
as the proximal loss.

###  4.2 Effect on Local Learning

#### Proximal Restriction

We examine the proximal restriction effect in FedSoL by analyzing how its
update gradient $g_{u}$ interacts with the proximal loss $\mathcal{L}_{p}$, as
shown in Figure 3. As $\rho$ increases, $g_{u}$ becomes more orthogonal to the
proximal gradient $g_{p}$ (Figure 3(b)). This orthogonality helps in
maintaining the proximal loss low during the local learning process, which
implies better global alignment (Figure 3(a)).

Figure 3: Effect of FedSoL on local learning in CIFAR-10 ($\alpha$=0.1) by
varying $\rho$ values. (a) Average proximal loss of local models. (b) Cosine
similarity between FedSoL gradient ($g_{u}$) and proximal gradient ($g_{p}$)
during local learning.

#### Adaptive Radius

The advantage of the adaptive approach are depicted in Figure 4. As shown in
Figure 4(a), using the adaptive radius not only improve performance but also
reduces sensitivity to the selection of $\rho$. Meanwhile, Figure 4(b)
displays the averaged $\lambda$ values for each local model layer,
highlighting the increased deviation in the later layers, as a consequence of
the data heterogeneity (Luo et al. 2021). Note that using a fixed value for
$\rho$ corresponds to setting $\boldsymbol{\Lambda}$ in Equation (11) as a
vector with all entries to one.

Figure 4: Effect of adaptive radius in CIFAR-10 ($\alpha$=0.1). (a) Server
test accuracy after 300 rounds. (b) Layer-wisely averaged $\lambda$ values of
FedSoL ($\rho=1.0$) at round 200.

###  4.3 Performance on Data Heterogeneity

#### Heterogeneity Level

Table 1 presents a comparison between our approach, FedSoL, and other
baselines such as FedProx (Li et al. 2020b), FedNova (Wang et al. 2020b),
Scaffold (Karimireddy et al. 2020), FedNTD (Lee et al. 2021), FedSAM (Qu et
al. 2022), and FedASAM (Caldarola, Caputo, and Ciccone 2022) as baseline
methods. Notably, many recently proposed FL methods tend to underperform when
compared to the standard FedAvg baseline, where a similar observation is
reported in (Zhao et al. 2018; Lee et al. 2021). In contrast, FedSoL
consistently exceeds the performance of FedAvg across all evaluated scenarios.
FedSoL achieves state-of-the-art results in most cases, particularly showing
consistent improvement on high heterogeneity levels ($s$=2 and $\alpha$=0.05).

#### Learning Factors

In Figure 5, we examine the learning factors that influence FedSoL’s
performance: Partial Participation (Figure 5(a)), Number of Local Epochs
(Figure 5(b)), and Learning Rate (Figure 5(c)). Throughout experiments, FedSoL
consistently surpasses the FedAvg across varying factors. Most of all, FedSoL
enlarges its gain as the smaller portion of clients participate in each round.
For instance, FedAvg significantly declines in performance at a sampling ratio
of 0.02, reaching to a near-random accuracy. However, FedSoL maintains robust
performance under such condition.

Figure 5: Performance of FedAvg and FedSoL on CIFAR-10 ($\alpha$=0.1) with
various setups: (a) sampling ratio, (b) the number of local epochs, (c)
initial learning rate, and (d) perturbation strength. The error bars stand for
the standard deviations.

#### Perturbation Strength

In FedSoL, a hyperparameter $\rho$ controls the overall perturbation strength.
Figure 5(d) plots FedSoL’s performance against varying $\rho$ values. Although
the model often diverges when using the SAM strategy with high perturbation
strength (as shown in Table 1), our FedSoL remains relatively robust and
achieve its best performance within the $\rho$ range between 0.5 and 2.0.

#### Model Architecture

We conduct further experiments on different model architectures: VggNet-11
(Simonyan and Zisserman 2014), ResNet-18 (He et al. 2016), and SL-ViT (Lee,
Lee, and Song 2021), which is a specialized structure of ViT (Dosovitskiy et
al. 2020) for small-sized datasets. The results provided in Table 2 validates
the efficacy of FedSoL across varying model architectures.

Table 2: Comparison of method on different model architectures. The heterogeneity is set as LDA ($\alpha$ = 0.1). Model | Method | CIFAR-10 | SVNH | PathMNIST
---|---|---|---|---
Vgg11 | FedAvg | 41.30±1.07 | 50.02±4.25 | 61.79±9.88
FedProx | 40.45±1.41 | 31.07±6.72 | 63.47±2.68
FedNTD |  60.55±2.14 | 56.62±2.64 | 69.82±2.27
FedSoL | 56.39±1.40 |  74.74±0.04 |  78.38±1.12
Res18 | FedAvg | 49.92±0.62 | 76.98±2.90 | 57.91±1.27
FedProx | 59.00±2.58 | 82.09±2.35 | 75.84±1.58
FedNTD | 57.79±3.42 | 78.50±0.18 | 76.87±0.57
FedSoL |  66.32±0.48 |  85.97±0.04 |  80.59±0.11
SL-ViT | FedAvg | 35.48±2.09 | 53.94±5.17 | 72.44±1.91
FedProx | 38.73±1.23 | 58.25±4.23 | 74.10±1.23
FedNTD | 47.59±2.84 | 61.46±1.76 | 71.65±1.71
FedSoL |  47.95±1.51 |  67.19±0.33 |  77.96±0.47

###  4.4 Proximal Losses

In our primary experiments, we utilize KL-divergence loss as the proximal
loss. However, FedSoL can be combine with various other proximal objectives.
Table 3 shows the impact of incorporating FedSoL with other proximal
restrictions: FedProx (Li et al. 2020b), FedNova (Wang et al. 2020b), Scaffold
(Karimireddy et al. 2020), FedDyn (Acar et al. 2021), and Moon (Li, He, and
Song 2021). These methods are compared in two distinct scenarios: as an
auxiliary objective alongside the original local objective (Base) and as
proximal perturbation within FedSoL (Combined). The results show the enhanced
performances with FedSoL.

Table 3: Comparison of proximal methods combined with FedSoL ($\rho$=2.0). The heterogeneity is set as LDA ($\alpha$ = 0.1). Method | CIFAR-10 | SVHN | CINIC-10
---|---|---|---
Base | Combined | Base | Combined | Base | Combined
FedProx | 59.80 |  63.93 $\uparrow$ | 72.40 |  84.32 $\uparrow$ | 40.09 |  55.25 $\uparrow$
FedNova | 10.00 |  31.77 $\uparrow$ | 53.07 |  79.95 $\uparrow$ | 21.89 |  42.37 $\uparrow$
Scaffold | 10.00 |  62.70 $\uparrow$ | 21.46 |  77.52 $\uparrow$ | 16.89 |  49.96 $\uparrow$
FedDyn | 60.80 |  62.85 $\uparrow$ | 78.15 |  79.43 $\uparrow$ | 48.25 |  52.17 $\uparrow$
MOON | 55.72 |  60.91 $\uparrow$ | 29.67 |  76.82 $\uparrow$ | 38.15 |  49.14 $\uparrow$

###  4.5 Partial Perturbation

As the data heterogeneity does not affect all layers equally (Luo et al.
2021), we investigate the use of partial perturbation in FedSoL, by
selectively perturbing specific layers instead of the entire model. The
results in Table 4 reveal that perturbing only the last classifier layer (Head
in Table 4) is sufficient for FedSoL. The performance is nearly as high as the
full-model perturbation, but the computational requirement is significantly
lower by avoiding multiple forward and backward computations across all layers
when using the standard SAM strategy. Interestingly, perturbing all layers
except the classifier head (Body in Table 4) consumes almost the same amount
of computation but rather drops in performance, showing the importance of the
later layers.

Table 4: Effect of partial weight perturbation in CIFAR10 ($\alpha$=0.1). The FLOPs shows relative computation w.r.t. FedAvg. $\delta$ stands for the computation for the proximal loss. Target Position |  Perturbation ($\boldsymbol{\rho}$) | FLOPs
---|---|---
0.0 | 0.5 | 1.0 | 1.5 | 2.0
All (full) | 56.13 | 61.17 | 64.16 | 64.38 | 63.94 | 2$\times$ $+\delta$
Body (partial) | 60.98 | 62.95 | 63.94 | 63.80 | 1.96$\times$ $+\delta$
Head (partial) | 62.65 | 63.62 | 64.13 | 63.25 | 1.33$\times$ $+\delta$

##  5 Analysis

In this section, we explore how FedSoL benefits in FL. We demonstrate that
FedSoL not only smooths the loss landscape of the learned model but also
ensures that local learning stays closely aligned with the global model. This
alignment retains the knowledge of the global distribution.

###  5.1 Weight Divergence

To assess the deviation of local learning from the global model, we measure
the L2 distance between models: $\lVert{w}_{g}-{w}_{k}\rVert$ where $w_{g}$ is
the global model and $w_{k}$ is the client $k$’s trained local model. The
results, averaged across sampled clients, are shown in Figure 6. In Figure
6(a), FedSoL effectively reduces the divergence, ensuring that local models
remain closely aligned with the global model, verifying the proximal
restriction effect in FedSoL. Figure 6(b) illustrates that this alignment also
fosters increased consistency among local models, reducing their mutual
divergence.

Figure 6: Comparative analysis of weight divergence in FedAvg and FedSoL
($\rho$=2.0) on CIFAR-10 LDA ($\alpha$=0.1). (a) shows global-local model
divergence, while (b) presents the divergence across local models.

###  5.2 Knowledge Preservation

To understand how FedSoL stabilizes the local learning at the prediction
level, we examine how well a local model maintains its performance on the
global distribution after local learning. As illustrated in Figure 7(a),
FedAvg’s local models undergo a significant drop in performance on the global
distribution after local learning. Conversely, FedSoL maintains high
performance, indicating better alignment of local learning with the global
objective, and thus stabilizing the learning process. We further analyze the
class-wise accuracy of FedAvg and FedSoL server models. As Figure 7(b)
demonstrates, while FedAvg exhibits significant fluctuations and inconsistent
class-wise performance, FedSoL preserves its class-wise accuracy as the
communication proceeds.

Figure 7: Comparison of FedAvg (blue lines) and FedSoL ($\rho$=2.0) (red
lines) on CIFAR-10 (s=2). (a) shows learning curves for global and local
models, with shaded areas reflecting standard deviation across clients. (b)
exhibits the class-wise accuracy of the global model.

###  5.3 Smoothness of Loss Landscape

We visualize the loss landscapes (Li et al. 2018) of global models obtained
from FedAvg, FedASAM, and FedSoL in Figure 8. In these plots, each axis
corresponds to one of the two dominant eigenvectors (top-1 and top-2) of the
Hessian matrix, representing the directions of the most significant shifts in
the loss landscape. Along with each landscape, we provide the value of the
dominant eigenvalue (${\lambda}_{1}$) and its ratio to the fifth largest
eigenvalue (${\lambda}_{1}/{\lambda}_{5}$), following the criteria used in
(Foret et al. 2020; Mi et al. 2022). Here, FedSoL’s smaller ratio indicates
that the variations in loss are more evenly distributed across various
directions. Both the landscape visualization and Hessian eigenvalues
underscore the efficacy of FedSoL in smoothing the loss landscape.

Figure 8: Loss landscape visualization of global model on CIFAR-10 LDA
($\alpha$=0.1). The ${\lambda}_{1}$ and ${\lambda}_{5}$ in each figure stand
for the top-1 and top-5 eigenvalues of the Hessian matrix.

##  6 Related Work

###  6.1 Federated Learning (FL)

Federated learning is a distributed learning paradigm to train models without
directly accessing private client data (Konečnỳ et al. 2016b, a). The standard
algorithm, FedAvg (McMahan et al. 2017), aggregates locally trained models by
averaging their parameters. While a variety of FL algorithms have been
introduced, they commonly conduct parameter averaging in a certain manner (Li
et al. 2020b; Karimireddy et al. 2020; Lee et al. 2021; Zhang et al. 2022).
Although FedAvg ideally performs well when all client devices are active and
IID distributed (Stich 2018; Woodworth et al. 2020), its performance
significantly degrades when clients have heterogeneous data distributions
(Zhao et al. 2018; Li et al. 2019; Kairouz et al. 2019). Our work focuses on
mitigating this data heterogeneity issue by modifying the local learning
strategy.

###  6.2 Proximal Restriction in FL

A prevalent strategy to address data heterogeneity in FL is the introduction
of a proximal term into local learning objectives (Lee et al. 2021; Li et al.
2020b; Karimireddy et al. 2020; Li, He, and Song 2021). This approach aims to
restrict the local learning deviation induced by the biased local
distributions. For example, FedProx (Li et al. 2020b) employs $L_{2}$ distance
between models, while MOON (Li, He, and Song 2021) uses the contrastive loss
(Oord, Li, and Vinyals 2018), regarding the previously trained local model’s
representations as negative pairs. Meanwhile, Scaffold (Karimireddy et al.
2020) use estimated global direction as a control variate to adjust local
gradients. However, such explicit alteration of local objectives may hinder
the acquisition of new knowledge during local learning (Qu et al. 2022;
Mendieta et al. 2022). In our study, we aim to leverage the benefits of
proximal restriction effect during local learning, but without changing the
original local objectives.

###  6.3 SAM Optimization in FL

Recent studies have begun to suggest that enhancing local learning generality
can significantly boost FL performance (Mendieta et al. 2022; Qu et al. 2022;
Caldarola, Caputo, and Ciccone 2022), aiding the global model in generalizing
more effectively. Inspired by the latest findings that connect loss geometry
to the generalization gap (Izmailov et al. 2018; Keskar et al. 2016; Jiang et
al. 2019; Chaudhari et al. 2019), those works seek for flat minima, utilizing
the recently proposed Sharpness-Aware Minimization (SAM) (Foret et al. 2020)
as the local optimizer. For instance, FedSAM (Qu et al. 2022) and FedASAM
(Caldarola, Caputo, and Ciccone 2022) demonstrate the benefits of using SAM
and its variants as local optimizer. Meanwhile, FedSMOO (Sun et al. 2023a)
incorporates a global-level SAM optimizer, and FedSpeed,(Sun et al. 2023b)
employs multiple gradient calculations to encourage global consistency. In our
work, we introduce the proximal restriction effect into SAM in an implicit
manner, by adjusting the perturbation direction and magnitude during local
learning.

##  7 Conclusion

We believe that FL serves as a crucial learning paradigm that opens up
extensive data utilization under the growing concerns on data privacy. In this
study, we emphasize the importance of both global alignment and local
generality in tackling data heterogeneity within FL. To unify these essential
components, we propose Federated Stability on Learning (FedSoL), a novel and
versatile method that seeks a robust parameter region against proximal weight
perturbations. This allows for an implicit proximal restriction effect on
local learning, without interfering with the original local objective. We
present a comprehensive analysis of FedSoL and demonstrate its benefits in FL.
The insights provided in our work may inspire new research directions,
furthering our understanding of FL and its potential applications.

## References

  * Abbas et al. (2022) Abbas, M.; Xiao, Q.; Chen, L.; Chen, P.-Y.; and Chen, T. 2022.  Sharp-maml: Sharpness-aware model-agnostic meta learning.  In _International conference on machine learning_ , 10–32. PMLR.
  * Acar et al. (2021) Acar, D. A. E.; Zhao, Y.; Navarro, R. M.; Mattina, M.; Whatmough, P. N.; and Saligrama, V. 2021.  Federated learning based on dynamic regularization.  _arXiv preprint arXiv:2111.04263_.
  * Aledhari et al. (2020) Aledhari, M.; Razzak, R.; Parizi, R. M.; and Saeed, F. 2020.  Federated learning: A survey on enabling technologies, protocols, and applications.  _IEEE Access_ , 8: 140699–140725.
  * Caldarola, Caputo, and Ciccone (2022) Caldarola, D.; Caputo, B.; and Ciccone, M. 2022.  Improving Generalization in Federated Learning by Seeking Flat Minima.  _arXiv preprint arXiv:2203.11834_.
  * Chaudhari et al. (2019) Chaudhari, P.; Choromanska, A.; Soatto, S.; LeCun, Y.; Baldassi, C.; Borgs, C.; Chayes, J.; Sagun, L.; and Zecchina, R. 2019.  Entropy-sgd: Biasing gradient descent into wide valleys.  _Journal of Statistical Mechanics: Theory and Experiment_ , 2019(12): 124018.
  * Chen, Hsieh, and Gong (2021) Chen, X.; Hsieh, C.-J.; and Gong, B. 2021.  When vision transformers outperform resnets without pre-training or strong data augmentations.  _arXiv preprint arXiv:2106.01548_.
  * Darlow et al. (2018) Darlow, L. N.; Crowley, E. J.; Antoniou, A.; and Storkey, A. J. 2018.  Cinic-10 is not imagenet or cifar-10.  _arXiv preprint arXiv:1810.03505_.
  * Deng et al. (2009) Deng, J.; Dong, W.; Socher, R.; Li, L.-J.; Li, K.; and Fei-Fei, L. 2009.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , 248–255. Ieee.
  * Deng (2012) Deng, L. 2012.  The mnist database of handwritten digit images for machine learning research.  _IEEE Signal Processing Magazine_ , 29(6): 141–142.
  * DeVries and Taylor (2017) DeVries, T.; and Taylor, G. W. 2017.  Improved regularization of convolutional neural networks with cutout.  _arXiv preprint arXiv:1708.04552_.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; et al. 2020\.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Fallah, Mokhtari, and Ozdaglar (2020) Fallah, A.; Mokhtari, A.; and Ozdaglar, A. 2020.  Personalized federated learning with theoretical guarantees: A model-agnostic meta-learning approach.  _Advances in Neural Information Processing Systems_ , 33: 3557–3568.
  * Foret et al. (2020) Foret, P.; Kleiner, A.; Mobahi, H.; and Neyshabur, B. 2020.  Sharpness-aware minimization for efficiently improving generalization.  _arXiv preprint arXiv:2010.01412_.
  * He et al. (2020) He, C.; Li, S.; So, J.; Zeng, X.; Zhang, M.; Wang, H.; Wang, X.; Vepakomma, P.; Singh, A.; Qiu, H.; et al. 2020.  Fedml: A research library and benchmark for federated machine learning.  _arXiv preprint arXiv:2007.13518_.
  * He et al. (2016) He, K.; Zhang, X.; Ren, S.; and Sun, J. 2016.  Identity mappings in deep residual networks.  In _European conference on computer vision_ , 630–645. Springer.
  * Hinton, Vinyals, and Dean (2015) Hinton, G.; Vinyals, O.; and Dean, J. 2015.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_.
  * Izmailov et al. (2018) Izmailov, P.; Podoprikhin, D.; Garipov, T.; Vetrov, D.; and Wilson, A. G. 2018.  Averaging weights leads to wider optima and better generalization.  _arXiv preprint arXiv:1803.05407_.
  * Jiang et al. (2019) Jiang, Y.; Neyshabur, B.; Mobahi, H.; Krishnan, D.; and Bengio, S. 2019.  Fantastic generalization measures and where to find them.  _arXiv preprint arXiv:1912.02178_.
  * Kairouz et al. (2019) Kairouz, P.; McMahan, H. B.; Avent, B.; Bellet, A.; Bennis, M.; Bhagoji, A. N.; Bonawitz, K.; Charles, Z.; Cormode, G.; Cummings, R.; et al. 2019.  Advances and open problems in federated learning.  _arXiv preprint arXiv:1912.04977_.
  * Karimireddy et al. (2020) Karimireddy, S. P.; Kale, S.; Mohri, M.; Reddi, S.; Stich, S.; and Suresh, A. T. 2020.  SCAFFOLD: Stochastic controlled averaging for federated learning.  In _International Conference on Machine Learning_ , 5132–5143. PMLR.
  * Keskar et al. (2016) Keskar, N. S.; Mudigere, D.; Nocedal, J.; Smelyanskiy, M.; and Tang, P. T. P. 2016\.  On large-batch training for deep learning: Generalization gap and sharp minima.  _arXiv preprint arXiv:1609.04836_.
  * Kim et al. (2022) Kim, M.; Li, D.; Hu, S. X.; and Hospedales, T. 2022.  Fisher SAM: Information Geometry and Sharpness Aware Minimisation.  In _International Conference on Machine Learning_ , 11148–11161. PMLR.
  * Konečnỳ et al. (2016a) Konečnỳ, J.; McMahan, H. B.; Ramage, D.; and Richtárik, P. 2016a.  Federated optimization: Distributed machine learning for on-device intelligence.  _arXiv preprint arXiv:1610.02527_.
  * Konečnỳ et al. (2016b) Konečnỳ, J.; McMahan, H. B.; Yu, F. X.; Richtárik, P.; Suresh, A. T.; and Bacon, D. 2016b.  Federated learning: Strategies for improving communication efficiency.  _arXiv preprint arXiv:1610.05492_.
  * Krizhevsky, Nair, and Hinton (2009) Krizhevsky, A.; Nair, V.; and Hinton, G. 2009.  Cifar-10 and cifar-100 datasets.  _URl: https://www. cs. toronto. edu/kriz/cifar. html_ , 6.
  * Kwon et al. (2021) Kwon, J.; Kim, J.; Park, H.; and Choi, I. K. 2021.  Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.  In _International Conference on Machine Learning_ , 5905–5914. PMLR.
  * Lee et al. (2021) Lee, G.; Shin, Y.; Jeong, M.; and Yun, S.-Y. 2021.  Preservation of the global knowledge by not-true self knowledge distillation in federated learning.  _arXiv preprint arXiv:2106.03097_.
  * Lee, Lee, and Song (2021) Lee, S. H.; Lee, S.; and Song, B. C. 2021.  Vision transformer for small-size datasets.  _arXiv preprint arXiv:2112.13492_.
  * Li et al. (2018) Li, H.; Xu, Z.; Taylor, G.; Studer, C.; and Goldstein, T. 2018.  Visualizing the loss landscape of neural nets.  _Advances in neural information processing systems_ , 31.
  * Li, He, and Song (2021) Li, Q.; He, B.; and Song, D. 2021.  Model-Contrastive Federated Learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 10713–10722.
  * Li et al. (2020a) Li, T.; Sahu, A. K.; Talwalkar, A.; and Smith, V. 2020a.  Federated learning: Challenges, methods, and future directions.  _IEEE Signal Processing Magazine_ , 37(3): 50–60.
  * Li et al. (2020b) Li, T.; Sahu, A. K.; Zaheer, M.; Sanjabi, M.; Talwalkar, A.; and Smith, V. 2020b.  Federated optimization in heterogeneous networks.  _Proceedings of Machine Learning and Systems_ , 2: 429–450.
  * Li et al. (2019) Li, X.; Huang, K.; Yang, W.; Wang, S.; and Zhang, Z. 2019.  On the convergence of fedavg on non-iid data.  _arXiv preprint arXiv:1907.02189_.
  * Luo et al. (2021) Luo, M.; Chen, F.; Hu, D.; Zhang, Y.; Liang, J.; and Feng, J. 2021.  No Fear of Heterogeneity: Classifier Calibration for Federated Learning with Non-IID Data.  _arXiv preprint arXiv:2106.05001_.
  * Marfoq et al. (2022) Marfoq, O.; Neglia, G.; Vidal, R.; and Kameni, L. 2022.  Personalized federated learning through local memorization.  In _International Conference on Machine Learning_ , 15070–15092. PMLR.
  * McMahan et al. (2017) McMahan, B.; Moore, E.; Ramage, D.; Hampson, S.; and y Arcas, B. A. 2017.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial Intelligence and Statistics_ , 1273–1282. PMLR.
  * Mendieta et al. (2022) Mendieta, M.; Yang, T.; Wang, P.; Lee, M.; Ding, Z.; and Chen, C. 2022.  Local Learning Matters: Rethinking Data Heterogeneity in Federated Learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 8397–8406.
  * Mi et al. (2022) Mi, P.; Shen, L.; Ren, T.; Zhou, Y.; Sun, X.; Ji, R.; and Tao, D. 2022.  Make sharpness-aware minimization stronger: A sparsified perturbation approach.  _arXiv preprint arXiv:2210.05177_.
  * Nesterov (1983) Nesterov, Y. E. 1983.  A method for solving the convex programming problem with convergence rate O (1/k^ 2).  In _Dokl. akad. nauk Sssr_ , volume 269, 543–547.
  * Netzer et al. (2011) Netzer, Y.; Wang, T.; Coates, A.; Bissacco, A.; Wu, B.; and Ng, A. Y. 2011.  Reading digits in natural images with unsupervised feature learning.  _NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011._
  * Oh, Kim, and Yun (2021) Oh, J.; Kim, S.; and Yun, S.-Y. 2021.  Fedbabu: Towards enhanced representation for federated image classification.  _arXiv preprint arXiv:2106.06042_.
  * Oord, Li, and Vinyals (2018) Oord, A. v. d.; Li, Y.; and Vinyals, O. 2018.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_.
  * Paszke et al. (2019) Paszke, A.; Gross, S.; Massa, F.; Lerer, A.; Bradbury, J.; Chanan, G.; Killeen, T.; Lin, Z.; Gimelshein, N.; Antiga, L.; Desmaison, A.; Kopf, A.; Yang, E.; DeVito, Z.; Raison, M.; Tejani, A.; Chilamkurthy, S.; Steiner, B.; Fang, L.; Bai, J.; and Chintala, S. 2019.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  In Wallach, H.; Larochelle, H.; Beygelzimer, A.; d'Alché-Buc, F.; Fox, E.; and Garnett, R., eds., _Advances in Neural Information Processing Systems 32_ , 8024–8035. Curran Associates, Inc.
  * Qu et al. (2022) Qu, Z.; Li, X.; Duan, R.; Liu, Y.; Tang, B.; and Lu, Z. 2022.  Generalized Federated Learning via Sharpness Aware Minimization.  _arXiv preprint arXiv:2206.02618_.
  * Simonyan and Zisserman (2014) Simonyan, K.; and Zisserman, A. 2014.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_.
  * Stich (2018) Stich, S. U. 2018.  Local SGD converges fast and communicates little.  _arXiv preprint arXiv:1805.09767_.
  * Sun et al. (2023a) Sun, Y.; Shen, L.; Chen, S.; Ding, L.; and Tao, D. 2023a.  Dynamic Regularized Sharpness Aware Minimization in Federated Learning: Approaching Global Consistency and Smooth Landscape.  _arXiv preprint arXiv:2305.11584_.
  * Sun et al. (2023b) Sun, Y.; Shen, L.; Huang, T.; Ding, L.; and Tao, D. 2023b.  Fedspeed: Larger local interval, less communication round, and higher generalization accuracy.  _arXiv preprint arXiv:2302.10429_.
  * Torralba, Fergus, and Freeman (2008) Torralba, A.; Fergus, R.; and Freeman, W. T. 2008.  80 million tiny images: A large data set for nonparametric object and scene recognition.  _IEEE transactions on pattern analysis and machine intelligence_ , 30(11): 1958–1970.
  * Wang et al. (2020a) Wang, H.; Yurochkin, M.; Sun, Y.; Papailiopoulos, D.; and Khazaeni, Y. 2020a.  Federated learning with matched averaging.  _arXiv preprint arXiv:2002.06440_.
  * Wang et al. (2020b) Wang, J.; Liu, Q.; Liang, H.; Joshi, G.; and Poor, H. V. 2020b.  Tackling the objective inconsistency problem in heterogeneous federated optimization.  _arXiv preprint arXiv:2007.07481_.
  * Wang et al. (2023) Wang, P.; Zhang, Z.; Lei, Z.; and Zhang, L. 2023.  Sharpness-aware gradient matching for domain generalization.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 3769–3778.
  * Woodworth et al. (2020) Woodworth, B.; Patel, K. K.; Stich, S.; Dai, Z.; Bullins, B.; Mcmahan, B.; Shamir, O.; and Srebro, N. 2020.  Is local SGD better than minibatch SGD?  In _International Conference on Machine Learning_ , 10334–10343. PMLR.
  * Yang et al. (2023) Yang, J.; Shi, R.; Wei, D.; Liu, Z.; Zhao, L.; Ke, B.; Pfister, H.; and Ni, B. 2023\.  MedMNIST v2-A large-scale lightweight benchmark for 2D and 3D biomedical image classification.  _Scientific Data_ , 10(1): 41.
  * Yang et al. (2019) Yang, Q.; Liu, Y.; Chen, T.; and Tong, Y. 2019.  Federated machine learning: Concept and applications.  _ACM Transactions on Intelligent Systems and Technology (TIST)_ , 10(2): 1–19.
  * Zhang et al. (2022) Zhang, L.; Shen, L.; Ding, L.; Tao, D.; and Duan, L.-Y. 2022.  Fine-tuning global model via data-free knowledge distillation for non-iid federated learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 10174–10183.
  * Zhao et al. (2018) Zhao, Y.; Li, M.; Lai, L.; Suda, N.; Civin, D.; and Chandra, V. 2018.  Federated learning with non-iid data.  _arXiv preprint arXiv:1806.00582_.
  * Zhuang et al. (2022) Zhuang, J.; Gong, B.; Yuan, L.; Cui, Y.; Adam, H.; Dvornek, N.; Tatikonda, S.; Duncan, J.; and Liu, T. 2022.  Surrogate gap minimization improves sharpness-aware training.  _arXiv preprint arXiv:2203.08065_.

##  Appendix A Table of Notations

Table 5: Table of Notations throughout the paper. Indices: |
---|---
$k$ |  Index for clients ($k\in[K]$)
$g$ |  Index for global server
Environment: |
$\mathcal{D}$ |  Whole dataset
$\mathcal{D}^{k}$ |  Local dataset of the $k$-th client
$\alpha$ |  Parameter for the Dirichlet Distribution
$s$ |  The number of shards per user
FL algorithms: |
$\beta,\mu$ |  Multiplicative coefficient for the proximal loss
$\gamma$ |  Learning rate
$\tau$ |  Temperature of probability distribution
$\rho$ |  Perturbation Radius for SAM-related algorithms
$\boldsymbol{\Lambda}$ |  Vector consists with scaling parameters for perturbation vector in SAM-related algorithms
Weights: |
$\boldsymbol{w}_{g}$ |  Weight of the global server model on the round $t$
$\boldsymbol{w}_{k}$ |  Weight of the $k$-th client model on the round $t$
$\|\boldsymbol{w}_{g}-\boldsymbol{w}_{k}\|$ |  Collection of $L^{2}$-norm between server and client models, among all rounds.
Objective Functions: |
$\mathcal{L}^{k}_{\mathrm{local}}$ |  Local objective for the $k$-th client
$\mathcal{L}^{k}_{p}$ |  Proximal Loss for the $k$-th client

##  Appendix B Experimental Setups

The code is implemented by PyTorch (Paszke et al. 2019). The overall code
structure is based on FedML (He et al. 2020) library with some modifications
for simplicity. We use 2 A6000 GPU cards, but without Multi-GPU training.

###  B.1 Model Architecture

For the primary experiments, we use the model architecture used in FedAvg
(McMahan et al. 2017), which consists of two convolutional layers with
subsequent max-pooling layers, and two fully-connected layers. The same model
is also used in (Lee et al. 2021; Li, He, and Song 2021; Luo et al. 2021). We
also conduct experiments on ResNet-18 (He et al. 2016), Vgg-11 (Simonyan and
Zisserman 2014), and SL-ViT (Lee, Lee, and Song 2021). For SL-ViT, we resize
$28\times 28$-sized images into $32\times 32$ to fit in the required minimum
patch size.

###  B.2 Datasets

To validate our algorithm, we employ 6 distinct datasets, as listed below. The
values in the parentheses denote the number of samples used to train and test,
respectively.

  * •

MNIST (Deng 2012) (60,000 / 10,000): contains hand-written digits images,
ranging from 0 to 9. The data is augmented using Random Cropping, Random
Horizontal Flipping, and Normalization. The data is converted to 3-channel RGB
images.

  * •

CIFAR-10 (Krizhevsky, Nair, and Hinton 2009) (50,000 / 10,000): contains a
labeled subset of 80 Million Tiny Images (Torralba, Fergus, and Freeman 2008)
for 10 different classes. The data is augmented using Random Cropping,
Horizontal Flipping, Normalization, and Cutout (DeVries and Taylor 2017).

  * •

SVHN (Netzer et al. 2011) (73,257 / 26,032): contains digits of house numbers
obtained from Google Street View. The data is augmented using Random Cropping,
Random Horizontal Flipping, and Normalization.

  * •

CINIC-10 (Darlow et al. 2018) (90,000 / 90,000): contains a combination of
CIFAR and downsized ImageNet (Deng et al. 2009), which is compiled to serve as
a bridge between the two datasets. The data is augmented using Random
Cropping, Random Horizontal Flipping, and Normalization.

  * •

PathMNIST (Yang et al. 2023) (110,000 / 7,180): contains non-overlapping
patches from Hematoxylin & Eosin stained colorectal cancer histology slide
images. The data is augmented using Random Horizontal Flipping, and
Normalization.

  * •

TissueMNIST (Yang et al. 2023) (189,106 / 47,280): contains microscope images
of human kidney cortex cells, which are segmented from 3 reference tissue
specimens. The data is augmented using Random Horizontal Flipping, and
Normalization. The data is converted to 3-channel RGB images.

Figure 9: Example images from PathMNIST datasets and TissueMNIST datasets.

Note that we evaluate our algorithm is on medical imaging datasets - a crucial
practical application of federated learning (Aledhari et al. 2020; Yang et al.
2019). Illustrative examples of the images are in Figure 9.

###  B.3 Non-IID Partition Strategy

To comprehensively address the data heterogeneity issue in federated learning,
we distribute the local datasets using the following two distinct data
partition strategies: (i) Sharding and (ii) Latent Dirichlet Allocation (LDA).

  * •

(i) Sharding (McMahan et al. 2017; Oh, Kim, and Yun 2021; Lee et al. 2021):
sorts the data by label and divide the data into shards of the same size, and
distribute them to the clients. In this strategy, the heterogeneity level
increases as the shard per user, $s$, becomes smaller, and vice versa. As the
number of shards is the same across all the clients, the dataset size is
identical for each client.

  * •

(ii) Latent Dirichlet Allocation (LDA) (Luo et al. 2021; Wang et al. 2020a;
Li, He, and Song 2021): allocates the data samples of class $c$ to each client
$k$ with the probability $p_{c}$, where $p_{c}\approx\text{Dir}(\alpha)$. In
this strategy, both the distribution and dataset size differ for each client.
The heterogeneity level increases as the concentration parameter, $\alpha$,
becomes smaller, and vice versa.

Note that although only the statistical distributions varies across the
clients in Sharding strategy, both the distribution and dataset size differ in
LDA strategy.

###  B.4 Learning Setups

We use a momentum SGD optimizer with an initial learning rate of 0.01, a
momentum value of 0.9, and weight decay 1e-5. The momentum is employed only
for local learning and is not uploaded to the server. Note that SAM
optimization also requires its base optimizer, which performs the parameter
update using the obtained gradient at the perturbed weights. The learning rate
is decays with a factor of 0.99. As we are assuming a synchronized FL
scenario, we simulate the parallel distributed learning by sequentially
conducting local learning for the sampled clients and then aggregate them into
a global model. The standard deviation is measured over 3 runs. The detailed
learning setups for each datasets is provided in Table 6.

Figure 10: Learning curves of FL methods on LDA ($\alpha$=0.1). The curves are smoothed for clear visualization. Table 6: Learning scenarios for each datasets. Datasets | Clients | Comm. Rounds | Sampling Ratio
---|---|---|---
MNIST | 100 | 200 | 0.1
CIFAR-10 | 100 | 300 | 0.1
SVHN | 100 | 200 | 0.1
CINIC-10 | 200 | 300 | 0.05
PathMNIST | 200 | 200 | 0.05
TissueMNIST | 200 | 200 | 0.05

###  B.5 Algorithm Implementation Details

We search for hyperparameters and select the best among the candidates. The
hyperparameters for each method is provided in Table 7. In the primary
experiments, we use KL-divergence loss (Hinton, Vinyals, and Dean 2015) with
softened logits with temperature $\tau$=3 for the proximal loss for the weight
perturbation in FedSoL.

Table 7: Algorithm-specific hyperparameters. Methods | Selected | Searched Candidates
---|---|---
FedAvg | None | None
FedProx |  $\mu$=1.0 | $\mu\in\\{0.1,0.5,1.0,2.0\\}$
Scaffold | None | None
FedNova | None | None
FedNTD |  $\beta$=1.0, $\tau$=1.0 |  $\beta\in\\{0.5,1.0\\}$, $\tau\in\\{1.0,3.0\\}$
FedSAM |  $\rho$=0.1 | $\rho\in\\{0.1,0.5,1.0,2.0\\}$
FedASAM |  $\rho$=1.0 | $\rho\in\\{0.1,0.5,1.0,2.0\\}$
FedDyn | None | None
MOON |  $\mu$=0.1, $\tau$=0.5 |  $\mu\in\\{0.1,0.5\\}$, $\tau\in\\{0.5,1.0\\}$
FedSoL | $\rho=2.0$ | $\rho\in\\{0.1,0.5,1.0,2.0\\}$

##  Appendix C Learning Curves

To provide further insights into the learning process, we illustrate the
learning curves of different FL methods in Figure 10. Although we utilize
different communication rounds for each dataset, the performance of the model
becomes sufficiently saturated at the end of communication rounds. For all
datasets, FedSoL not only achieves a superior final model at the end of the
communication round but also demonstrates much faster convergence. Moreover,
although some algorithms that perform well on a dataset fail on another (ex.
FedNTD (Lee et al. 2021) underperforms compared to FedProx (Li et al. 2020b)
on the TissueMNIST datasets), FedSoL consistently exhibits significant
improvements when compared to the other baselines.

##  Appendix D Personalized Performance

In Table 8, we compare FedSoL with several methods specifically designed for
personalized federated learning (pFL): PerFedAvg (Fallah, Mokhtari, and
Ozdaglar 2020), FedBabu (Oh, Kim, and Yun 2021), and kNN-Per (Marfoq et al.
2022). Each method is assessed by fine-tuning them for $e$ local epochs from
the global model after the final communication round. As global alignment is
unnecessary for the personalized model, we fine-tune FedSoL using original the
local objective without perturbation and denote it as FedSoL-FT. The standard
deviation is measured across the clients. The results reveal that our FedSoL-
FT consistently outperforms other pFL methods under various scenarios.
Furthermore, the gap is enlarged when local ($e$=1), implying that the global
model obtained by FedSoL adapts more quickly to local distributions. We
suggest that by integrating FedSoL with other methods specialized for pFL, we
can attain superior performance for both the global server model and client
local models.

Table 8: Personalized FL performance after $\tau$ epochs of fine-tuning. The heterogeneity level is set as LDA ($\alpha=0.1$). Method | $\boldsymbol{e}$ | CIFAR-10 | SVHN | TissueMNIST
---|---|---|---|---
Local-only | - | 84.7${}_{\;\pm 12.8}$ | 87.4${}_{\;\pm 13.0}$ | 82.4${}_{\;\pm 15.5}$
FedAvg | 1 | 84.1${}_{\;\pm 13.4}$ | 86.6${}_{\;\pm 15.5}$ | 82.2${}_{\;\pm 17.5}$
5 | 88.9${}_{\;\pm 8.9}\;\,$ | 92.1${}_{\;\pm 5.7}$ | 89.2${}_{\;\pm 10.1}$
PerFedAvg | 1 | 80.5${}_{\;\pm 16.2}$ | 64.1${}_{\;\pm 30.3}$ | 82.3${}_{\;\pm 18.9}$
5 | 86.3${}_{\;\pm 10.4}$ | 72.4${}_{\;\pm 21.2}$ | 88.8${}_{\;\pm 10.2}$
FedBabu | 1 | 84.6${}_{\;\pm 12.7}$ | 88.7${}_{\;\pm 9.6}$ | 85.7${}_{\;\pm 14.3}$
5 | 89.2${}_{\;\pm 8.4}$ | 92.7${}_{\;\pm 6.2}$ | 90.5${}_{\;\pm 8.8}$
kNN-Per | 1 | 85.7${}_{\;\pm 12.3}$ | 86.4${}_{\;\pm 15.0}$ | 86.5${}_{\;\pm 14.2}$
5 | 89.7${}_{\;\pm 8.1}$ | 92.8${}_{\;\pm 6.2}$ | 91.4${}_{\;\pm 7.5}$
FedSoL-FT (Ours) | 1 |  87.5${}_{\;\pm 9.7}$ |  92.5${}_{\;\pm 7.4}$ |  88.1${}_{\;\pm 12.2}$
5 |  90.5${}_{\;\pm 7.8}$ |  95.0${}_{\;\pm 3.9}$ |  91.6${}_{\;\pm 6.9}$

##  Appendix E Proximal Perturbation with SAM

In our work, we consider how to combine the proximal restriction into SAM
optimization to improve overall FL performance. A straightforward approach
might involve using the linearly combined local objective between the original
local loss $\mathcal{L}^{k}_{\mathrm{local}}$ and the proximal loss
$\mathcal{L}^{k}_{\mathrm{p}}$ in Equation (2) for the SAM optimization in
Equation (3) as follows:

$\underset{\boldsymbol{w}_{k}}{\text{min}}\;\underset{{\|\boldsymbol{\epsilon}\|}_{2}<\rho}{\text{max}}\;\big{[}\mathcal{L}_{\mathrm{local}}^{k}(\boldsymbol{w}_{k})+\beta\cdot\mathcal{L}_{p}^{k}(\boldsymbol{w}_{k};\boldsymbol{w}_{g})\big{]}\,.$
(12)

In the above equation, the gradients for weight perturbation and parameter
update are obtained from the same objective.

However, this approach encounters the same drawbacks as when using each method
on its own. The combined loss also varies considerably across clients due to
heterogeneous local distributions, causing the smoothness to largely rely on
individual local distributions. Furthermore, the negative correlation between
the gradients of the two objectives within the combined loss still limits
local learning. Consequently, this approach neither encourages global
alignment nor preserves the local generality as desired.

Instead in FedSoL, we overcome this issue by decoupling this directly combined
loss into the proximal loss $\mathcal{L}^{k}_{\mathrm{p}}$ for weight
perturbation and the original local loss $\mathcal{L}^{k}_{\mathrm{local}}$
for weight updates. To further analyze the relationship between loss functions
and weight perturbation in SAM optimization, we conduct an ablation study on
the following strategies.

  * •

$A_{0}$: Use original local loss without any weight perturbation (FedAvg).

  * •

$A_{1}$: Use original local loss, but get the original local loss gradient at
weights perturbed by the proximal gradient (FedSoL).

  * •

$A_{2}$: Use combined loss, but get the proximal loss gradient at weights
perturbed by the proximal gradient.

  * •

$A_{3}$: Use combined loss, but get the proximal loss gradient at weights
perturbed by the proximal gradient.

  * •

$A_{4}$: Use combined loss, but get the combined loss gradient at weights
perturbed by the combined gradient.

  * •

$A_{5}$: Use combined loss without any weight perturbation (Proximal
Restriction).

  * •

$A_{6}$: Use combined loss, but get the original local loss gradient at
proximally perturbed weight loss.

We exclude the strategies that obtaining proximal loss at the perturbed
weights using the original local loss gradient i.e.,
$\mathcal{L}_{p}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{c}^{*}$), where
$\boldsymbol{\epsilon}_{c}^{*}$ $=$
$\rho\,\frac{\boldsymbol{g}_{p}+\boldsymbol{g}_{l}}{\lVert\boldsymbol{g}_{p}+\boldsymbol{g}_{l}\rVert}$,
as it leads the learning to diverge. The detailed formulation for each method
is provided in Table 9 with its corresponding performance. The results in
Table 9 demonstrates that utilizing the original local loss gradient at
weights perturbed by the proximal loss gradient ($A_{1}$ in Table 9) yields
outperforms the other approaches. We suggest that our FedSoL is an effective
way to integrate proximal restriction effect into SAM optimization in FL.

Table 9: Detailed formulation for each method and their performance on CIFAR-10 datasets (LDA $\alpha$=0.1). Name | Method Formulation | Performance
---|---|---
$A_{0}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k})$ | 56.13
$A_{1}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{p}^{*})$ | 64.13
$A_{2}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k})+\beta\cdot\mathcal{L}_{p}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{p}^{*})$ | 53.85
$A_{3}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{p}^{*})+\beta\cdot\mathcal{L}_{p}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{p}^{*})$ | 60.28
$A_{4}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{\text{c}}^{*})+\beta\cdot\mathcal{L}_{p}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{\text{c}}^{*})$ | 45.72
$A_{5}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k})+\mathcal{L}_{p}(\boldsymbol{w}_{k})$ | 61.76
$A_{6}$ | $\mathcal{L}_{\text{local}}(\boldsymbol{w}_{k}+\boldsymbol{\epsilon}_{p}^{*})+\beta\cdot\mathcal{L}_{p}(\boldsymbol{w}_{k})$ | 44.12

##  Appendix F Proof of Proposition

We begin by organizing Equation (7), substituting Equation (6) into FedSoL:

$\displaystyle\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{\\{\mathrm{local},p\\}}(\boldsymbol{w}_{k})\approx-\gamma\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\\{\mathrm{local},p\\}}(\boldsymbol{w}_{k}),\boldsymbol{g}_{u}(\boldsymbol{w}_{k})\rangle$
$\displaystyle\approx-\gamma\left(\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\\{\mathrm{local},p\\}}(\boldsymbol{w}_{k}),\boldsymbol{g}_{l}\rangle+\rho\,\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\\{\mathrm{local},p\\}}(\boldsymbol{w}_{k}),\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}\,\hat{\boldsymbol{g}}_{p}\rangle\right)\,,$
(13)

and for the update of $\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})$
in FedAvg:

$\displaystyle\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\approx-\gamma\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}(\boldsymbol{w}_{k}),\boldsymbol{g}_{l}\rangle=-\gamma\|\nabla\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}\,.$
(14)

###  F.1 Proof of Proposition 1

Regarding Proposition 1, from Equation (13), we obtain:

$\displaystyle\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{p}\approx-\gamma\left(\langle\boldsymbol{g}_{p},\boldsymbol{g}_{l}\rangle+\rho\,\langle\boldsymbol{g}_{p},\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}\,\hat{\boldsymbol{g}}_{p}\rangle\right)\approx-\gamma\Big{(}\langle\boldsymbol{g}_{l}\,,\boldsymbol{g}_{p}\rangle+\rho\cdot\hat{\boldsymbol{g}}_{p}{\\!\\!}^{\top}\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}\,{\boldsymbol{g}_{p}}\Big{)}\,.$

Furthermore, if the local objective is convex, then the Hessian is always
positive semi-definite. Consequently, $\boldsymbol{g}_{p}{\\!\\!}^{\top}$
$\nabla^{2}\mathcal{L}^{k}_{\mathrm{local}}$ $\,{\boldsymbol{g}_{p}}\geq 0$,
and we can guarantee that the second term is nonnegative as well.

###  F.2 Proof of Proposition 2

Similarly, from Equation (13), we derive:

$\displaystyle\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})$
$\displaystyle\approx-\gamma\left(\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k}),\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\rangle+\rho\,\langle\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k}),\nabla_{\boldsymbol{w}_{k}}^{2}\mathcal{L}^{k}_{\mathrm{local}}\,\hat{\boldsymbol{g}}_{p}\rangle\right)$
$\displaystyle=-\gamma\left(\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}+\frac{\rho}{2}\,\langle\hat{\boldsymbol{g}}_{p},2\nabla_{\boldsymbol{w}_{k}}^{2}\mathcal{L}^{k}_{\mathrm{local}}\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\rangle\right)\,.$
$\displaystyle=-\gamma\left(\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}+\nabla_{\boldsymbol{w}_{k}}\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}\cdot\frac{\rho}{2}\,\hat{\boldsymbol{g}}_{p}\right)\,.$

On the other hand, from Equation (14), applying the first-order Taylor
approximation to
$\boldsymbol{w}_{k}\mapsto\|\nabla\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}$,
we have:

$\displaystyle\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{\mathrm{local}}\left(\boldsymbol{w}_{k}+\frac{\rho}{2}\hat{\boldsymbol{g}}_{p}\right)\approx-\gamma\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}\left(\boldsymbol{w}_{k}+\frac{\rho}{2}\hat{\boldsymbol{g}}_{p}\right)\|^{2}$
$\displaystyle\approx\gamma\left(\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\|^{2}+\nabla_{\boldsymbol{w}_{k}}\left\|\nabla_{\boldsymbol{w}_{k}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})\right\|^{2}\cdot\frac{\rho}{2}\,\hat{\boldsymbol{g}}_{p}\right)\,.$

Therefore, $\Delta^{\mathrm{FedAvg}}\mathcal{L}^{k}_{\mathrm{local}}$
$(\boldsymbol{w}_{k}+\frac{\rho}{2}\hat{\boldsymbol{g}}_{p})$ is equal to
$\Delta^{\mathrm{FedSoL}}\mathcal{L}^{k}_{\mathrm{local}}(\boldsymbol{w}_{k})$,
up to the first-order Taylor approximation.