TOWARDS LAYER-WISE PERSONALIZED FEDERATED LEARNING: ADAPTIVE LAYER DISENTANGLEMENT VIA CONFLICTING GRADIENTS

## Abstract

Abstract In personalized Federated Learning (pFL), high data heterogeneity can cause significant gradient divergence across devices, adversely affecting the learning process. This divergence, especially when gradients from different users form an obtuse angle during aggregation, can negate progress, leading to severe weight and gradient update degradation.
To address this issue, we introduce a new approach to pFL design, namely Federated Learning with Layer-wise Aggregation via Gradient Analysis (FedLAG), utilizing the concept of gradient conflict at the layer level.
Specifically, when layer-wise gradients of different clients form acute angles, those gradients align in the same direction, enabling updates across different clients toward identifying client-invariant features.
Conversely, when layer-wise gradient pairs make create obtuse angles, the layers tend to focus on client-specific tasks.
In hindsights, FedLAG assigns layers for personalization based on the extent of layer-wise gradient conflicts. Specifically, layers with gradient conflicts are excluded from the global aggregation process.
The theoretical evaluation demonstrates that when integrated into other pFL baselines, FedLAG enhances pFL performance by a certain margin. Therefore, our proposed method achieves superior convergence behavior compared with other baselines. Extensive experiments show that our FedLAG outperforms several state-of-the-art methods and can be easily incorporated with many existing methods to further enhance performance.

## 1 Introduction

The challenge of non-independent and non-identically distributed (non-IID) data significantly impacts personalized Federated Learning (pFL). Addressing the aforementioned problem, numerous researchers have delved deeply into various directions, e.g., 1) consisting adding regularization (T. Dinh et al., ; Li et al., ), 2) pseudo representations generation (Zhu et al., ; Zhang et al., ), 3) on-client data pre-processing (Huang et al., ; ), 4) gradient update modification (Reddi et al., ; Wang et al., ; Sun et al., ; Huang et al., ), 5) cross-client feature alignments (Zhang et al., ; Dinh et al., ), and 6) adaptive global update by leveraging on-server gradients (Jhunjhunwala et al., ; Panchal et al., ). These approaches are *orthogonal* in nature, enabling their *combined integration to achieve further performance enhancement*.

Lately, model layer disentanglement has emerged as a promising approach to further enhancing the performance of pFL (Oh et al., ; Collins et al., ; Chen & Chao, ; Xu et al., ). This approach involves the disentanglement of local model into two distinct components: global aggregation layers (GAL) and personalized layers (PL). GAL manages common tasks among all users, while PL handles the specific tasks according to each user. However, existing methods require extensive fine-tuning to determine which layers should be used for global aggregation and personalization. Consequently, the optimal selection for layer disentanglement remains an open question.

Addressing this challenge, we consider the gradient conflict (Yu et al., ) in multi-task learning (MTL). The principle of gradient conflict posits that if the angle between the gradients from two users is less than $\pi/2$, the gradients are aligned in the *same direction*. Consequently, the gradient progress of each user does not negatively impact the progress of the others (see Fig. [1](#S1.F1.6)). By utilizing gradient conflict, it becomes possible to monitor user interactions directly from the server without local data accessibility, i.e., the extent to which one user’s progress may adversely affect others. Moreover, when analyzing gradient conflict at the layer level, we discovered that layer-wise gradient conflicts do not align in a structured way. Specifically, the initial layers tend to exhibit low gradient conflicts, indicating their role in learning generic features, whereas the deeper layers show high gradient conflicts, reflecting their role in learning more personalized tasks (see Section [3](#S3)). This observation is also aligned with the phenomenon in MTL (Shi et al., ).

Figure: Figure 1: The issue of gradient conflicts, we denote ideal gradient as the aggregated gradient when component gradients do not make any conflicts, i.e., when the conflicted gradient with angle more than $\pi/2$ is projected into $\pi/2$ as mentioned in (Yu et al., ).
Refer to caption: /html/2410.02845/assets/x1.png

- •
Ability to evaluate the performance of deep learning models at the layer level.
- •
Ability to analyze statistics of local user data via users’ gradients.
- •
Efficiently reduce communication overhead by utilizing the gradients.

Our contributions can be summarized as follows: 1) We develop an algorithm, namely Federated Learning with Layer-wise Aggregation via Gradient analysis (FedLAG), which can automatically differentiate between the GAL and PL layers without requiring extensive tuning. 2) We give a theoretical evaluation and present evidence showcasing the convergence of the proposed algorithm. We illustrate its sustained superior convergence compared to conventional FL algorithms without the incorporation of FedLAG. 3) We conduct extensive experimental evaluations to show the superiority of the proposed algorithm over conventional FL algorithms.

## 2 Problem Formulation & Preliminaries

### 2.1 Notations

We use $\mathrm{span}(g^{(r)}_{1},\ldots,g^{(r)}_{J})$ as the subspaces that are contained by all vectors $\{g^{(r)}_{1},\ldots,g^{(r)}_{J}\}$.
We consider an FL system comprising a set of users denoted by $\mathcal{U}=\{u|u=1,2,\ldots,U\}$. Each user gains access to its local data, which remains inaccessible to others. Specifically, the data collected by the $u$-th user can be represented by $\mathcal{D}_{u}\in\mathbb{R}^{N_{u}}$, where $N_{u}$ denotes the number of data instances for user $u$. The entire dataset available across all users can be denoted by $\mathcal{D}=\bigcup_{u\in\mathcal{U}}\mathcal{D}_{u}$, and we have $N=\sum^{U}_{u=1}N_{u}$. We use the term $E$ to represent the number of local epochs. We abuse the notations $w_{u}^{(r)}$ and $w_{g}^{(r)}$ to refer to the local model of user $u$ and the global model at round $r$, respectively.

### 2.2 Problem Setup

During FL training, users iteratively conduct local training and communicate with the server for model updating. To be specific, our FL concept works as follows:

##### Local Updates.

In each round $r$, users’ local models are updated by the global model, i.e., $w^{(r)}_{u}\leftarrow w^{(r)}_{g},~{}\forall u$. Then, users conduct local training in parallel. We assume that the local models parallelly solve the empirical loss over data distribution of the user $u$. For instance, $\mathcal{L}(w^{(r,e)}_{u})=\mathbb{E}_{B_{k}\in\mathcal{D}_{u}}\Big{[}\ell(B_{k},w^{(r,e)}_{u})\Big{]},$ where $B_{k}$ is the $k^{\textrm{th}}$ mini-batch sampled from $\mathcal{D}_{u}$, $k$ represents the data batch index, and $e$ as the local epoch. In each training epoch, the $\mathcal{L}(w^{(r,e)}_{u})$ is computed by averaging over the batch-wise loss computation $\ell(B_{k},w^{(r,e)}_{u})$ over the all batches $B_{k},~{}\forall k\in\{1,\ldots,K\}$. Subsequently, at every FL communication round, users update with a local learning rate $\eta$ as $w_{u}^{(r,E)}\leftarrow w_{u}^{(r)}-\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}\left(w_{u}^{(r,e)}\right),$ where $\mathcal{L}$, $\nabla\mathcal{L}$ are the empirical loss function and its gradient, respectively.

Global Aggregation. After local training, the server aggregates the local models $w^{(r+1)}_{g}\leftarrow\frac{1}{U}\sum_{u=1}^{U}w^{(r,E)}_{u}$ where $N_{u}$ is the number of samples from user $u$ and $N=\sum^{U}_{u=1}N_{u}$. The server disseminates the global parameters $w^{(r+1)}_{g}$ to the local users chosen for the following round $r+1$.

Layer Disentanglement. Assume user models $w^{(r)}_{u}$ consisting of $L$ layers, i.e., $w^{(r)}_{u}=\{\theta^{(r)}_{l,u}\}_{l=1}^{L}$, where $\theta^{(r)}_{l,u}$ represents the parameter weights at $l^{th}$ layer. Our objective is to disentangle the $L$ layers into GAL and PL. The GAL is defined as $\theta^{(r)}_{s,u}=\{\theta^{(r)}_{l,u}|~{}\forall l\in\mathbb{L}_{g}\}$. On the other hand, PL is defined as $\theta^{(r)}_{p,u}=\{\theta^{(r)}_{l,u}|~{}\forall l\in\mathbb{L}_{p}\}$. Here, $\mathbb{L}_{g}$ and $\mathbb{L}_{p}$ represent the layers assigned to GAL and PL, respectively. We have $w^{(r)}_{u}=\{\theta^{(r)}_{l,u}\}_{l=1}^{L}=\{\theta^{(r)}_{s,u},\theta^{(r)}_{p,u}\}$. Similarly, we denote $w^{(r)}_{g}=\{\theta^{(r)}_{l,g}\}_{l=1}^{L}=\{\theta^{(r)}_{s,g},\theta^{(r)}_{p,g}\}$.

### 2.3 Negative Transfer and Gradient Conflicts in Multi-task Learning

A major challenge for MTL is negative transfer, which refers to the performance drop on a task caused by the learning from other tasks, resulting in degradation in overall performance. A rationale of this phenomenon is the conflicting gradients (Yu et al., ). Specifically, gradients from different tasks may point in different directions so that directly optimizing the average loss is detrimental to a specific task’s performance. Denote $g_{i}=\nabla\mathcal{L}\left(w_{i}\right)$ as the gradient of task $i$, and $\phi_{ij}$ as the angle between two task gradients $g_{i}$ and $g_{j}$, we have

###### Definition 2.1 (Conflicting gradients (Yu et al., 2020a ) ) .

Given two gradients $g_{i}$, $g_{j}$ $(i\neq j)$, and $\cos{\phi_{ij}}=\frac{g_{i}\cdot g_{j}}{\|g_{i}\|\|g_{j}\|}$ is the cosine between two vectors. $g_{i}$ and $g_{j}$ $(i\neq j)$ are said to be conflicting with each other if $\cos{\phi_{ij}}<0$.

Figure: (a) Illustration of gradient conflict with toy dataset.
Refer to caption: /html/2410.02845/assets/x3.png

In (Shi et al., ), the concept of conflicting gradients is extended to the layer level.

###### Definition 2.2 (Layer-wise Conflicting Gradients Shi et al. ( 2023 ) ) .

The gradients $g_{l,i}$, $g_{l,j}$ $(i\neq j)$ of layer $l$ are said to be conflicting with each other if $\cos{\phi^{l}_{ij}}<0$.

The definition suggests that gradient conflicts within a model can be analyzed at the layer level, offering a foundation for exploring and determining optimal strategies for layer disentanglement.

###### Definition 2.1 (Conflicting gradients (Yu et al., 2020a ) ) .

###### Definition 2.2 (Layer-wise Conflicting Gradients Shi et al. ( 2023 ) ) .

## 3 Validation of Layer-wise gradient conflicts in FL

We argue that the aggregation of layers with layer-wise gradient conflicts is detrimental in FL. To validate our arguments, we conduct two experiments to answer two questions:

###### Question 3.1 .

Is gradient conflicts appear among users in federated learning?

To answer the question [3.1](#S3.Thmtheorem1), we conduct the experiments on 2-D toy dataset (the details of the toy dataset is reported in Appendix [G](#A7)). The result is visualized Fig. [2(a)](#S2.F2.sf1). In the non-IID setting, the divergences between two pairs of gradients (e.g., $g^{(r)}_{1}$ vs. $g^{(r)}_{3}$ and $g^{(r)}_{2}$ vs. $g^{(r)}_{4}$) result in the divergence of the FL process (Fig. [2(a)](#S2.F2.sf1), top). Due to gradient conflicts (angles exceeding $\pi/2$), the model diverges along two hyper-spaces, $\mathrm{span}(g^{(r)}_{1},g^{(r)}_{3})$ and $\mathrm{span}(g^{(r)}_{2},g^{(r)}_{4})$ (Fig. [2(a)](#S2.F2.sf1), bottom). In contrast, under IID settings, the gradients form smaller angles and align toward consistent directions. Consequently, the model progresses more directly toward the optimal solution. As noted in Appendix [G](#A7), the optimal solution is located at $y=0$, implying that $w_{1}=0$, $w_{2}\neq 0$, and $b=0$.

###### Question 3.2 .

Is the gradient conflict distribution aligned in an ascending way (i.e., from the input layers to the output layers) during the training of federated learning?

To answer the question [3.2](#S3.Thmtheorem2), we conduct the experiments on Cifar-10 dataset and visualize in Fig. [2(b)](#S2.F2.sf2). As observed in the figure, the distribution of gradient conflicts does not correspond with the assumptions regarding layer disentanglement found in current SOTA, e.g., Oh et al. (), Collins et al. (). Specifically, the density of conflicting gradients among clients does not progressively increase from the input layer to the output layer, as these approaches apply to design the layer disentanglement (i.e., personalized layers are assigned at the very last layers).

Figure: Figure 3: The FedLAG architecture. First, calculate the previous layer-wise gradient $h^{(r-1)}_{u}$ using received and stored models. Second, measure angles between pairs of gradient vectors, considering angles above $90$ degrees as conflicts. Thirdly, the $GC_{\epsilon}(l)$ (in Definition [4.1](#S4.Thmtheorem1)) score for layer $l$ increases by $1$ for each conflicted pair. The $k$ layers with highest $GC_{\epsilon}(l)$ are assigned to PL. The GAL and PL are used to assign the personalized and global aggregated layers.
Refer to caption: /html/2410.02845/assets/x7.png

###### Question 3.1 .

###### Question 3.2 .

## 4 Methodology

Building upon the validations presented in Section [3](#S3), we introduce FedLAG, a method that exploits layer-wise gradient conflicts to identify the optimal selection for disentangling personalized and generic layers. The process of FedLAG is briefly illustrated as in Figure [3](#S3.F3).
The core contribution of FedLAG is its ability to utilize users’ gradients as a proxy for their behavior on the server, enabling an analysis of user relationships *without direct access to local datasets*. By taking user gradients as input, FedLAG applies the gradient divergence analysis (GDA) to examine gradients at the layer level (see Section [4.1](#S4.SS1)). Upon performing GDA, the $GC_{\xi}(l)$ score is computed, which quantifies the degree to which a given layer should be personalized. Consequently, we can leverage $GC_{\xi}(l)$ to adaptively disentangle generic and personalized layers (see Section [4.2](#S4.SS2)).

### 4.1 Gradient Divergence Analysis

Layer-wise gradient. To implement the GDA, we calculate the gradient of each user $u$ on round $r$:

$$ $\displaystyle h^{(r)}_{u}$ $\displaystyle=\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}(w^{(r,e)}_{u})=w^{(r,E)}_{u}-w^{(r)}_{u},$ (1) $$

where the layer-wise gradient $h^{(r)}_{l,u}$ of $h^{(r)}_{u}$ can be represented as

$$ $\displaystyle h^{(r)}_{l,u}$ $\displaystyle=\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}(\theta^{(r,e)}_{l,u})=\theta^{(r,E)}_{l,u}-\theta^{(r)}_{l,u}.$ (2) $$

Here, $h^{(r)}_{u}$ and $h^{(r)}_{l,u}$ denote the gradient of $w^{(r)}_{u}$ and the $l^{th}$ layer of $w^{(r)}_{u}$, respectively. We further define $h^{(r)}_{u}=[h^{(r)}_{1,u},h^{(r)}_{2,u},\ldots,h^{(r)}_{L,u}]$ where $h^{(r)}_{l,u}$ is a $l$-th layer-wise gradient of user $u$. The gradient is calculated by utilizing the models from the previous round in conjunction with the recently received model. As a result, FedLAG is both *communication-efficient* and capable of *information aggregation*.

Layer-wise Gradient Divergence Analysis.
To evaluate the layer-wise gradient divergence between two users, we introduce $\phi^{(r)}_{l}(u,v)$ to signify the angle between $h^{(r)}_{l,u}$ and $h^{(r)}_{l,v}$. Based on Definition [2.1](#S2.Thmtheorem1), we define gradient conflict of layer $l$ in the round $r$ with hyper-parameter $\xi$ (w.r.t $-1<\xi\leq 0$) as the *layer-wise gradient conflict score*, denoted by $GC_{\xi}(l)$.

###### Definition 4.1 ( $GC_{\xi}(l)$ score) .

Given the threshold $\xi$, the $GC_{\xi}(l)$ of the $l$-th layer is calculated as the number of distinct user pairs $(u,v)$ (where $u\neq v$) that satisfy $\cos\phi_{l}^{(r)}(u,v)<\xi$. For instance,

$$ $\displaystyle GC_{\xi}(l)=\frac{1}{2}\sum^{U}_{u=1}\sum^{U}_{v=1,v\neq u}{\mathbb{I}_{u,v}},\quad\text{s.t.}\quad\mathbb{I}_{u,v}=\begin{cases}1,&\text{if}\cos\phi_{l}^{(r)}(u,v)<\xi,\\ 0,&\text{otherwise},\end{cases}$ (3) $$

where $\cos\phi^{(r)}_{l}(u,v)=\cos(h^{(r)}_{l,u},h^{(r)}_{l,v}),~{}\forall v,u\in U$; $\xi$ denotes the extent of conflict severity. To elaborate, by setting a smaller value for $\xi$, the angle between the two vectors becomes more obtuse.

Consequently, the Definition [4.1](#S4.Thmtheorem1) allows us to focus primarily on the count of more prominent conflicts. Subsequently, $GC_{\xi}(l)$ acts as an indicator for conflicting gradients across various severity levels within the layers. If $GC_{\xi}(l)$ takes on the value $\binom{U}{2}$, it implies that for any two users, there is a conflict in their gradients w.r.t the $l$-th layer. By computing these layer-wise conflict scores, we can pinpoint the layers where conflicts occur most frequently.

We describe our method to find personalized layers $\mathbb{L}_{p}$ in Algorithm [2](#alg2). First, we calculate $h^{(r)}_{l,u}$ for each layer by Eq. equation [2](#S4.E2). Because we utilize the users’ model parameters to measure users’ gradient trajectories $h^{(r)}_{u}$. *No additional communication cost is incurred for the execution of Algorithm [2](#alg2)*. Afterwards, we calculate cosine $\cos\phi^{(r)}_{l}(u,v)$ by Definition [2.1](#S2.Thmtheorem1) and $GC_{\xi}(l)$ score via the Definition [4.1](#S4.Thmtheorem1). Finally, we find the $k$ layers with the highest scores and assign their index to $\mathbb{L}_{p}$.

###### Definition 4.1 ( G ​ C ξ ​ ( l ) 𝐺 subscript 𝐶 𝜉 𝑙 GC_{\xi}(l) score) .

### 4.2 Layer-wise Personalized Model Aggregation

To achieve adaptive layer disentanglement, we base our approach on two key principles: (1) maintain the global aggregation of FL on non-conflict layers, and (2) motivate the personalized learning on conflict layers. Specifically, when a layer is significantly affected by gradient conflict, we transform it into a personalized layer to prevent negative transfer resulting from the aggregation process at the global server. To this end, rather than broadcasting the entire model to the local users, we restrict the local model update to the global layer (i.e., layers assigned to GAL):

$$ $\displaystyle\theta^{(r+1)}_{l,u}=\frac{1}{U}\sum^{U}_{u=1}\theta^{(r,E)}_{l,u},~{}~{}~{}$ $\displaystyle\text{if }l\in\mathbb{L}_{g}.$ (4) $$

To motivate the personalized learning, the users do not update the conflict layers:

$$ $\displaystyle\theta^{(r+1)}_{l,u}=\theta^{(r,E)}_{l,u},~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}$ $\displaystyle\text{if }l\in\mathbb{L}_{p}.$ (5) $$

Here, $\mathbb{L}_{p}$ represents the layers that suffer from the gradient conflict and need to be converted into personalized layers. The detailed description of our method is demonstrated in Algorithm [1](#alg1). Compared to the conventional pFL, FedLAG only alters the aggregation process on the global server, leaving the local training unaffected. Consequently, our technique is applicable to various existing FL or pFL approaches.

## 5 Theoretical Analysis

In this section, we discuss the convergence of our algorithm. Our theoretical analysis aims to show the improvement in terms of upper bound reduction for the convergence upper boundary.

### 5.1 Layer-wise Loss Improvement

To show the robustness and prove the convergence of FedLAG, we first want to analyze the performance improvement when using our LAG algorithm in FL.

###### Lemma 5.1 (Personalization Improvement) .

Each user $u$ achieve an improvement in loss when using FedLAG over the vanilla FL approach as follows:

$$ $\displaystyle\mathcal{L}$ $\displaystyle\Big{(}w^{(r)}_{\textrm{LAG},u}\Big{)}-\mathcal{L}\Big{(}w^{(r)}_{\textrm{VFL},u}\Big{)}=-\eta\frac{1}{U}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r-1)}_{l,u}\|\times(\|h^{(r-1)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r-1)}_{l,v}\|)<0,$ $$

###### Lemma 5.2 (Generalization Improvement) .

For any sufficiently small learning rate $\eta$, the following holds:

$$ $\displaystyle\mathcal{L}$ $\displaystyle\Big{(}w^{(r)}_{\textrm{LAG},g}\Big{)}-\mathcal{L}\Big{(}w^{(r)}_{\textrm{VFL},g}\Big{)}=-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r-1)}_{l,u}\|\times(\|h^{(r-1)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r-1)}_{l,v}\|)<0,$ $$

where $w^{(r)}_{\textrm{LAG},g}=\{\theta_{s,g},\theta_{p,g}\}$ and $w^{(r)}_{\textrm{VFL},g}=\{\theta_{s,u},\theta_{p,u}\}$ stand for the FL with and without the integration of FedLAG, respectively.

According to Lemma [5.2](#S5.Thmtheorem2), at each round $r$, the loss function of the layer-wise pFL model consistently surpasses that of the vanilla pFL model by a specific margin. This enhancement is directly proportional to both the layer-wise gradient norm from the previous round $r-1$ and the layer-wise angle between pairs of users in the FL system. Consequently, the proof establishes that the application of FedLAG is viable by relying on the previous gradients for estimating gradient conflicts, *eliminating the need for using the upcoming local gradients*, which is unavailable at the global server. This approach avoids communication overheads associated with exchanging information among local users. Furthermore, the improvement remains constant in each round upon integrating the LAG, in contrast to the FL algorithm lacking LAG integration. The proof of Lemma [5.2](#S5.Thmtheorem2) is provided in Appendix [H.4](#A8.SS4).

###### Lemma 5.1 (Personalization Improvement) .

###### Lemma 5.2 (Generalization Improvement) .

### 5.2 Convergence Analysis

We use Assumptions [H.1](#A8.Thmtheorem1), [H.2](#A8.Thmtheorem2), [H.4](#A8.Thmtheorem4), [H.3](#A8.Thmtheorem3) and have the following theorem:

###### Theorem 5.3 .

Assuming users compute full-batch gradient with full participation and $\frac{1}{2\sqrt{6}E^{2}L}<\eta<\frac{1}{6EL}$, the series $\{w^{(r)}\}$ generated by FedLAG satisfy:

$$ $\displaystyle\mathcal{L}(w^{(R)})-\mathcal{L}(w^{*})$ $\displaystyle\leq\mathcal{O}\Big{(}\frac{\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ (6) $\displaystyle-\mathcal{O}\Big{(}A\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|(\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|)\Big{)},$ $$

where $w^{*}\triangleq\operatorname*{\arg\min}_{w}\mathcal{L}(w)$ is the optimal global model, and $A=\Big{(}24E^{2}(E-1)\eta^{4}L^{2}-\eta^{2}E\Big{)}>0,~{}\forall\eta,E,L$.

The convergence rate in Theorem [5.3](#S5.Thmtheorem3) consists of four terms:

- •
The first term $\mathcal{O}\Big{(}\frac{\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}$ is the initialization error term that depends on the total communication rounds $R$ and local learning epoch $E$. This term is fixed among all FL algorithms and independent of the FedLAG hyperparameters.
- •
The second term $\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}$ is the noise at optimum. This term reveals that the prediction at the optimum always make a certain variance.
- •
The third term $\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ refers to the user drift error, which is the error induced by the divergence when the user drift toward their specific domain characteristics. It affects by the data characteristics and FL hyperparameters such as the number of epoch $E$ and learning rate $\eta$.
- •
The last term $\mathcal{O}\Big{(}A\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)\Big{)}$ is the personalized loss improvement term. This loss shows the improvement in loss thanks to the disentanglement in FL models into GAP and PL. We can see that this term consistently reduces the bound and thus creates an absolute improvement to the FL system regardless of whether the FL algorithm is integrated into it.

From the Theorem [5.3](#S5.Thmtheorem3), we can have the following remarks:

###### Remark 5.4 .

When the conflict does not occur (e.g., the data is IID), no layer being assigned to the PL subset, which makes $\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)$ becomes 0 0. Therefore, the convergence of the FedLAG reduces to $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$.

###### Remark 5.5 .

When we set the top $k$ layers to 0 0, the sum of $\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)$ becomes 0 0, therefore, the convergence of FedLAG reduces to $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$.

###### Remark 5.6 .

From the equation, as the number of layers becomes large (i.e., over-parameterization), the right hand side may close to 0 0. However, in Appendix [H.6](#A8.SS6), we prove that the value of the last term is agnostic to the number of parameters in the network model. The reason is because as the number of layers increase, the layer-wise gradient norm will decrease respectively.

###### Remark 5.7 .

Increasing $E$ and $\eta$ amplifies the value of $A$, boosting the improvement term. However, this enhancement affects local learning in distributed users. For instance, a high value of $E$ introduces bias towards local characteristics, causing users to forget global knowledge learned through aggregation, which loses the generality of the pFL concept.

###### Remark 5.8 .

The FedLAG’s main contribution is on the server. Therefore, FedLAG can be integrated with any other FL algorithms to improve other algorithms’ performance. The theoretical results of integration between FedLAG and other FL algorithms are demonstrated in Appendix [D](#A4).

**Table 1: Test accuracy comparison among baselines and our proposed method on $4$ datasets (i.e., MNIST, CIFAR10/100, EMNIST). The data set splitting method is selected from the Dirichlet sampling with replacement. The experimental setups are $100$ users at $20\%$ user participation (i.e. $\alpha\in\{0.1,0.5\}$). Each result is averaged after 5 times.**
| Setting | non-IID ($\alpha=0.1$) | non-IID ($\alpha=0.5$) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Problem<br>Method | MNIST | CIFAR10 | CIFAR100 | EMNIST | MNIST | CIFAR10 | CIFAR100 | EMNIST |
| FedLAG<br>(Ours) | 97.18 $\pm$ 0.13 | 85.21 $\pm$ 0.29 | 51.35 $\pm$ 0.02 | 96.34 $\pm$ 0.25 | 96.68 $\pm$ 0.01 | 69.26 $\pm$ 0.09 | 33.75 $\pm$ 0.14 | 91.42 $\pm$ 0.18 |
| PerAvg | 89.11 $\pm$ 0.17 | 81.25 $\pm$ 0.20 | 43.21 $\pm$ 0.28 | 84.84 $\pm$ 0.20 | 94.89 $\pm$ 0.06 | 61.25 $\pm$ 0.14 | 23.98 $\pm$ 0.25 | 89.70 $\pm$ 0.14 |
| FedROD | 97.02 $\pm$ 0.29 | 81.72 $\pm$ 0.16 | 46.17 $\pm$ 0.06 | 96.02 $\pm$ 0.08 | 94.87 $\pm$ 0.12 | 68.47 $\pm$ 0.19 | 26.76 $\pm$ 0.08 | 90.70 $\pm$ 0.28 |
| FedPAC | 89.12 $\pm$ 0.04 | 83.13 $\pm$ 0.01 | 44.77 $\pm$ 0.06 | 89.63 $\pm$ 0.04 | 94.60 $\pm$ 0.18 | 63.01 $\pm$ 0.04 | 25.42 $\pm$ 0.25 | 88.44 $\pm$ 0.61 |
| FedBABU | 95.92 $\pm$ 0.26 | 80.75 $\pm$ 0.08 | 42.59 $\pm$ 0.03 | 84.63 $\pm$ 0.19 | 91.42 $\pm$ 0.03 | 65.12 $\pm$ 0.18 | 21.54 $\pm$ 0.30 | 87.35 $\pm$ 0.10 |
| FedAvg | 90.61 $\pm$ 0.11 | 65.47 $\pm$ 0.17 | 41.28 $\pm$ 0.09 | 85.02 $\pm$ 0.30 | 88.43 $\pm$ 0.07 | 59.01 $\pm$ 0.26 | 15.99 $\pm$ 0.10 | 86.35 $\pm$ 0.16 |
| FedCAC | 96.77 $\pm$ 0.03 | 84.62 $\pm$ 0.19 | 47.22 $\pm$ 1.52 | 96.01 $\pm$ 0.15 | 96.33 $\pm$ 0.11 | 68.93 $\pm$ 0.19 | 26.12 $\pm$ 0.04 | 91.10 $\pm$ 0.18 |
| FedDBE | 95.73 $\pm$ 0.11 | 83.76 $\pm$ 0.09 | 50.12 $\pm$ 0.12 | 95.04 $\pm$ 0.05 | 95.52 $\pm$ 0.11 | 67.75 $\pm$ 0.04 | 25.43 $\pm$ 0.04 | 90.14 $\pm$ 0.08 |
| GPFL | 94.15 $\pm$ 0.10 | 82.11 $\pm$ 0.25 | 49.85 $\pm$ 0.01 | 93.25 $\pm$ 0.21 | 93.56 $\pm$ 0.08 | 66.38 $\pm$ 0.07 | 24.01 $\pm$ 0.11 | 88.45 $\pm$ 0.15 |
| FedAS | 95.36 $\pm$ 0.12 | 84.11 $\pm$ 0.02 | 50.26 $\pm$ 0.03 | 93.25 $\pm$ 0.19 | 94.28 $\pm$ 0.18 | 67.38 $\pm$ 0.14 | 32.01 $\pm$ 0.04 | 89.45 $\pm$ 0.32 |
| FedAF | 95.24 $\pm$ 0.15 | 84.35 $\pm$ 0.03 | 50.11 $\pm$ 0.04 | 93.10 $\pm$ 0.20 | 94.40 $\pm$ 0.17 | 67.55 $\pm$ 0.13 | 32.07 $\pm$ 0.05 | 89.50 $\pm$ 0.30 |

###### Theorem 5.3 .

###### Remark 5.4 .

###### Remark 5.5 .

###### Remark 5.6 .

###### Remark 5.7 .

###### Remark 5.8 .

## 6 Experiment Setup

### 6.1 Datasets

To conduct a fair comparison among methods, we consider four different context data sets including MNIST ($28\times 28$, 10 modalities) (LeCun et al., ), CIFAR10 ($32\times 32\times 3$, 10 modalities) (Krizhevsky, ), EMNIST ($28\times 28$, 62 modalities) (Cohen et al., ), and CIFAR100 ($32\times 32\times 3$, 100 modalities) (Krizhevsky, ). Otherwise, each data set is divided into several parts corresponding to the number of users by randomly distinguished distribution. Each divided part is identical to the others to ensure that non-IID term of FL problem. Each user owns each part containing a train and test set, to which no data augmentation method is applied.

### 6.2 Baselines

To assess the robustness of our proposed FedLAG, we run evaluations on $5$ other baselines (i.e., FedAvg (McMahan et al., ), PerAvg (Fallah et al., ), FedBABU (Oh et al., ), FedPAC (Xu et al., ), and FedRoD (Chen & Chao, ), FedCAC (Wu et al., ), FedDBE (Zhang et al., ), GBFL (Zhang et al., ), which are all trained from scratch using ResNet18 (He et al., ). We apply the same settings on all baselines to achieve the fairest experimental evaluations. In detail, each method is conducted in with unbalanced data distribution and non-IID scenarios, along with different sampling rates $\alpha=0.1\%$, and $0.5\%$. Otherwise, various numbers of users are applied to perform a fair comparison among methods (i.e. $20,40,60,80,100$), corresponding with the different number of global rounds: $100,200,400,600$, and $800$ rounds, respectively.

## 7 Experimental Evaluations

### 7.1 Overall Performance

To evaluate the overall performance of the FL system, we compute average results across users. Table [1](#S5.T1) provides a detailed overview of comprehensive performance metrics, where $\alpha$ represents the Dirichlet coefficient. The table explores the FL performance of two settings: 1) varying participation ratio and 2) different levels of heterogeneity.

#### 7.1.1 Different heterogeneity levels

We assess at two heterogeneity levels, i.e., $\alpha=0.1,1$. The table reveals that FedLAG outperforms other baselines, with improvements ranging from an average of $5-7\%$ to $15\%$ when $\alpha=0.5$. The resilience of FedLAG becomes more evident in a more challenging setting, namely $\alpha=0.1$, where it demonstrates a substantial improvement over other baselines, averaging $10\%$ to $40\%$. Among the competing baselines, FedRoD poses a notable challenge. This is because, in FedRoD, the authors employ disentanglement in personalized and generic models across different data settings (same as our work). In contrast, our research incorporates adaptive control over the two sub-models, resulting in superior performance compared to FedRoD.

#### 7.1.2 Different user participation ratio

We evaluated under five different participation ratio, i.e., $\{20\%,40\%,60\%,80\%,100\%\}$, and illustrated as in Table [7](#A5.T7). As it can easily be seen from the table, our proposed FedLAG can achieve significantly higher performance, as opposed to other baselines (i.e., an average of $20\%$ up to $40\%$ in performance improvement). The improvement is more significant when we employ the algorithm in more challenging data sets, which showcases a more divergence in data characteristics among users. The detailed comparisons between FedLAG and other baseline models in terms of comprehensive training are presented in Appendix [E](#A5).

### 7.2 Ablation Test

#### 7.2.1 Integratability

In this section, we evaluavate FedLAG ’s compatibility into other FL baselines on the Cifar10 dataset, and illustrate the overall performance as in Figure [5](#S7.F5). The integration of FedLAG showcase a significantly better accuracy over baselines.

Figure: Figure 4: The comparison of FedLAG with different fixed layer disentanglement. First-$K$ means we fix first $K$ layers for the disentanglement, Last-$K$ means we fix last $K$ layers, Middle-$K$ means we fix $K$ layers at the middle of the network.
Refer to caption: /html/2410.02845/assets/x8.png

#### 7.2.2 Is FedLAG more efficient than fixed layer disentanglement?

To prove the robustness of FedLAG over fixed layer disentanglement techniques, we conduct experiments on FedBABU, and fix $K$ layers for personalized layers. Tab. [5](#S7.F5) demonstrates that the fixing the last $K$ layers generally yields the most efficient performance.
Other settings show significant drop in performance, compared to that of the FedLAG. This observation aligns with the assumption that gradient conflicts tend to be concentrated in the final layers, though not all layers with high conflict scores are necessarily the last layers.
In total, we can obviously see that the adaptive layer disentanglement of FedLAG shows a significant robustness over fixed layer disentanglement.

#### 7.2.3 Top k 𝑘 k score layer

We conducted experiments with varying values of $k$, and the results are presented in Table [3](#S7.T3). As shown in the table, selecting the top $k$ layers leads to significantly improved performance compared to the case where $k=0$. This enhancement is notable because when $k=0$, the algorithm reduces to the FedAvg algorithm, lacking the robustness inherent in the GDA algorithm. Our analysis reveals that gradient conflicts predominantly affect only a small number of layers in the FL model. Consequently, setting $k=5,10,15$ does not result in a significant performance improvement, as the most conflicted gradients are already captured in the initial layers with the highest conflicts. Detailed results are presented in Appendix [F.2](#A6.SS2).

#### 7.2.4 Conflict gradient score

Table [3](#S7.T3) demonstrates the performance of FedLAG under different conflict score $\xi$. As observed in the table, there is negligible performance disparity despite different gradient scores being applied. This is because the score primarily impacts the frequency of gradient conflicts across various model layers, yet it does not alter the distribution of conflicted gradients throughout the model layers. A high value of $\xi$ (e.g., 0 0) results in performance drops due to an excessive inclusion of gradient pairs with small conflicts. Some of these pairs are beneficial for learning common structures and should not be excluded. Conversely, a too-small value for $S$ (e.g., $-0.3$) also leads to performance degradation by ignoring many gradient pairs with significant conflicts, which can be detrimental to the learning process. Detailed results are presented in Appendix [F.3](#A6.SS3).

Figure: Figure 6: The figures illustrate the performance without pretrained model of FL-LAG vs. various baselines on CIFAR10. The evaluation is implemented on different numbers of users (from left to right, the number of users is $\{100\%,80\%,60\%,40\%,20\%\}$, respectively). $\alpha=0.1$.
Refer to caption: /html/2410.02845/assets/x9.png

Figure: Figure 7: The figures illustrate the performance without pretrained model of FL-LAG vs. various baselines on CIFAR100. The evaluation is implemented on different numbers of users (from left to right, the number of users is $\{100\%,80\%,60\%,40\%,20\%\}$, respectively). $\alpha=0.1$.
Refer to caption: /html/2410.02845/assets/x14.png

#### 7.2.5 Efficiency when training without pretrained models

In this section, we evaluate FL without pretrained models on the CIFAR-10 and CIFAR-100 datasets to assess FL performance when training from scratch (see Figs. [6](#S7.F6), [7](#S7.F7)).
We only apply the FedLAG after first $30$ rounds to sample the very first gradient trajectories for the initial process of GDA.
Most baseline models encounter significant challenges when trained without pre-trained parameters. In contrast, our FedLAG consistently demonstrates superior performance, achieving the highest results alongside FedRod and FedPAC. The significant superiority of FedLAG is shown when training on the challenging CIFAR-100 dataset, which contains a large number of labels. While other federated learning baselines struggle to exceed random performance, FedLAG rapidly surpasses this threshold and achieves effective convergence in fewer than $100$ communication rounds.

## 8 Conclusion

Current researches on layer disentanglement in pFL typically require extensive fine-tuning to achieve optimal separation between generic and personalized layers. In our study, we introduce an adaptive approach to disentangle these layers by leveraging a well-established principle in multi-task learning, namely, conflicting gradients. To address conflicting gradients among users without incurring communication overhead from user-to-user interactions, we propose a novel data-free gradient divergence analysis method performed on the server. This technique enhances federated learning performance by enabling the selective assignment of network layers to personalization when layer-wise gradients are in conflict, and to generic layers otherwise. Our proposed method, FedLAG, demonstrates significant improvements over current baselines in both accuracy and convergence time.

## Appendix A Notations

**Table 4: Abbreviations**
| Symbol | Description |
| --- | --- |
| $U$ | Number of users |
| $\eta$ | Learning rate |
| $w$ | Model parameter |
| $E$ | Number of training epochs for each user per round |
| $\xi$ | Threshold for layer-wise gradient conflict score |
| $L$ | Number of model layers |
| $R$ | Number of rounds |
| $w*$ | Optimal global model |
| $\sigma^{2}_{*}$ | Upper bound of variance of user gradients at optimum |
| $w_{g}^{(r)}$ | Global parameter at round $t$ |
| $w_{u}^{(r,e)}$ | Local parameter at round $t$, training at epoch $e$ |
| $\cos\Phi^{(l)}_{u,v}$ | Cosine of gradient of user $u$ and $v$ w.r.t the $l^{th}$ layer |
| $\theta^{(r)}_{s,u},\theta^{(r)}_{s,g}$ | Model layers that is assigned for global aggregation |
| $\theta^{(r)}_{p,u},\theta^{(r)}_{p,g}$ | Model layers that is assigned for personalized purpose |
| $\theta_{l,u}^{(r)}$ | $l^{th}$ Layer in model parameter |
| $h_{u}^{(r)}$ | gradient of user $u$ on round $r$ |
| $h_{l,u}^{(r)}$ | gradient w.r.t the $l^{th}$ layer |
| $P$ | Number of parameters of each pFL model |
| $P_{l}$ | Number of parameters on layer $l^{th}$ |
| $\mathbb{L}_{p},\mathbb{L}_{g}$ | The set of coefficients in the personalized and global model |
| $k$ | Number of selected layers for personalization |
| $w^{(r)}_{\textrm{LAG},u}$ | Layer-wise Personalized Model of user $u$ |
| $GC_{\xi}(l)$ | Layer-wise gradient conflict score |

## Appendix B Related Works

Gradient-based Multi-task Learning. In MTL, the AI model consists of two distinguished groups: 1) shared encoder that is common to all tasks, and 2) task-specific decoders which are designed and learned independently for each task. Addressing the contemporary challenge of task conflicts in MTL, numerous approaches have been devised to mitigate the aforementioned issue, which can be categorized in two main directions, i.e., task loss balancing and gradient manipulation. PCGrad (Yu et al., ), projects each gradient onto the normal plane of another gradient and employs the average of these projected gradients for updates. GradDrop (Chen et al., ), randomly drops some elements of gradients based on element-wise conflicts. CAGrad (Liu et al., ), ensures convergence to a minimum of the average loss across tasks through gradient manipulation. RotoGrad (Javaloy & Valera, ), re-weights task gradients and rotates the shared feature space to mitigate conflicts. RECON (Shi et al., ), leverages gradient information to modify network structure and address task conflicts at their core.

Mitigating non-IID in Federated Learning.
Prior work has investigated how to improve the FL robustness against non-IID data. FedPAC (Xu et al., ) training personalized models by exploiting a better feature extractor and user-specific classifier collaboration. FedProx (Li et al., ) adds a proximal term to the local training objective to keep updated parameters close to the original downloaded model. SCAFFOLD (Karimireddy et al., ) introduces control variates to correct the drift in local updates. MOON (Li et al., ) adopts contrastive loss to improve representation learning.

Layer disentanglement in Federated Learning. Oh et al. () decompose the entire local network into the body (extractor), which is related to universality, and the head (classifier), which is related to personalization.
Collins et al. () proposes a method where the entire network is trained sequentially during local updates, but only the body is aggregated. During the local update phase, each client first trains the head using the aggregated representation. Then, within the same epoch, the client trains the body using its own head.
FedRoD (Chen & Chao, ) proposes to use the balanced softmax for learning generic models and vanilla softmax for personalized heads.
Xu et al. () design an objective function to constraint the body with the task of learning invariant features.
However, most layer disentanglement approaches demand significant effort to identify the optimal layer selection, often leading to arbitrary choices for the body and head layers. Currently, there is no clear method for determining which layers should be personalized and which should remain generic.

Mitigating negative transfer in Federated Learning.
Recently, negative transfer has been discovered to be one of the most crucial issue in FL.
FedCollab (Bao et al., ) minimizes the pair-wise distribution distances between users to alleviate the negative transfer among clients.
DisentAFL (Chen & Zhang, ) integrates mixtures of experts to selectively aggregate clients’ representations. To this end, the FL system can aggregate the fine-grained inter-client relationships to achieve sufficient positive transfer while avoiding negative transfer. However, we believe that current methods lack a strong theoretical foundation that directly addresses negative transfer among clients.

## Appendix C Detail of models used

**Table 5: Model employed for MNIST, EMNIST, CIFAR-10, CIFAR-100 Datasets.**
| Dataset | Model | 1 Initial<br>Conv | Residual Blocks | 3 FC Layers | Activation<br>Function | Total<br>Model<br>Size (MB) |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Output<br>Channel | Blocks/<br>Group | Output Channels/<br>Group | Output<br>Dimension |  |  |  |  |
| MNIST | Resnet-9 | 32 | (1,2,2,1) | (32,64,128,256) | (256,128,10) | ReLU | 6.81 |
| EMNIST | Resnet-9 | 32 | (1,2,2,1) | (32,64,128,256) | (256,128,10) | ReLU | 6.81 |
| CIFAR-10 | Resnet-20 | 64 | (3,3,3) | (64,128,256) | (256,128,10) | ReLU | 17.54 |
| CIFAR-100 | Resnet-20 | 64 | (2,3,4) | (64,128,256) | (256,128,100) | ReLU | 17.54 |

## Appendix D Extensive Theoretical Results

As current FL concepts are mostly gradient-based optimization approaches, the convergence proofs of all FL approaches are proved in a similar approach. Therefore, for simplicity, we do not prove all the theoretical results comprehensively. As being proved in Appendix [H.4](#A8.SS4), our LAG always give the improvement of $-\frac{1}{U}\sum^{U}_{u=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)$ on any global loss function. Therefore, by apply this theorem to the proof in other approaches (e.g., pFedMe (T. Dinh et al., ), FedEXP (Jhunjhunwala et al., ), SCAFFOLD (Karimireddy et al., )), and apply the model disentanglement to apply the loss improvement lemma, we can have the results as following table:

**Table 6: Table of convergence of baselines vs. baselines + LAG.**
| Method | Convergence Rate |  |  |  |
| --- | --- | --- | --- | --- |
| FedAvg (McMahan et al., ) | $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ |  |  |  |
| FedAvg + LAG (Ours) | $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ |  |  |  |
|  | $-\mathcal{O}\Big{(}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)\Big{)}$ |  |  |  |
| FedEXP (Jhunjhunwala et al., ) | $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{\eta_{l}E\sum^{R-1}_{r=0}\eta_{g}^{(r)}}\Big{)}+\mathcal{O}\Big{(}\eta_{l}\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta_{l}^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ |  |  |  |
| FedEXP + LAG (Ours) | $\mathcal{O}\Big{(}\frac{3\|w^{(0)}_{g}-w^{*}\|^{2}}{\eta_{l}E\sum^{R-1}_{r=0}\eta_{g}^{(r)}}\Big{)}+\mathcal{O}\Big{(}\eta_{l}\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta_{l}^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ |  |  |  |
|  | $-\mathcal{O}\Big{(}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)\Big{)}$ |  |  |  |
| pFedMe (T. Dinh et al., ) | $\mathcal{O}\Big{(}3\|w^{(0)}_{g}-w^{*}\|^{2}\mu_{F}e^{-\hat{\eta_{1}}\mu_{F}T/2}\Big{)}+\mathcal{O}\Big{(}\frac{(N/S-1)\sigma^{2}_{F,1}}{\mu_{F}TN}\Big{)}$ |  |  |  |
|  | $+\mathcal{O}\Big{(}\frac{(R\sigma^{2}_{F,1}+\delta^{2}\lambda^{2})\kappa_{F}}{R(T\beta\mu_{F})^{2}}\Big{)}+\mathcal{O}(\frac{\lambda^{2}\delta^{2}}{\mu_{F}})$ |  |  |  |
| pFedMe + LAG (Ours) | $\mathcal{O}\Big{(}3\|w^{(0)}_{g}-w^{*}\|^{2}\mu_{F}e^{-\hat{\eta_{1}}\mu_{F}T/2}\Big{)}+\mathcal{O}\Big{(}\frac{(N/S-1)\sigma^{2}_{F,1}}{\mu_{F}TN}\Big{)}$ |  |  |  |
|  | $+\mathcal{O}\Big{(}\frac{(R\sigma^{2}_{F,1}+\delta^{2}\lambda^{2})\kappa_{F}}{R(T\beta\mu_{F})^{2}}\Big{)}+\mathcal{O}(\frac{\lambda^{2}\delta^{2}}{\mu_{F}})$ |  |  |  |
|  | $-\mathcal{O}\Big{(}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)\Big{)}$ |  |  |  |
| SCAFFOLD (Karimireddy et al., ) | $\mathcal{O}(\frac{\sigma}{\mu SK_{\epsilon}})+\mathcal{O}(\frac{1}{\mu})+\mathcal{O}(\frac{N}{S})$ |  |  |  |
| SCAFFOLD + LAG (Ours) | $\mathcal{O}(\frac{\sigma}{\mu SK_{\epsilon}})+\mathcal{O}(\frac{1}{\mu})+\mathcal{O}(\frac{N}{S})$ |  |  |  |
|  | $-\mathcal{O}\Big{(}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|)\Big{)}$ |  |  |  |

## Appendix E Detailed Results

### E.1 Full table of overall performance

**Table 7: Test accuracy comparison among baselines and our proposed method on $4$ datasets (i.e., MNIST, CIFAR10/100, EMNIST). The data set splitting method is selected from the Dirichlet sampling with replacement (i.e. $\alpha\in\{0.1,0.5\}$). The experimental setups are at $20\%$, $40\%$, $60\%$, $80\%$, and $100\%$ user participation. Each result is averaged after 10 times.**
| Problem | non-IID ($\alpha=0.1$) | non-IID ($\alpha=0.5$) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Problem<br>Method | MNIST | CIFAR10 | CIFAR100 | EMNIST | MNIST | CIFAR10 | CIFAR100 | EMNIST |
| Participation ratio = $100\%$ |  |  |  |  |  |  |  |  |
| FedLAG<br>(Ours) | 97.97 $\pm$ 0.17 | 88.54 $\pm$ 0.1 | 57.08 $\pm$ 0.28 | 96.73 $\pm$ 0.15 | 99.57 $\pm$ 0.10 | 87.18 $\pm$ 0.03 | 41.81 $\pm$ 0.18 | 97.57 $\pm$ 0.27 |
| PerAvg | 87.05 $\pm$ 0.13 | 83.12 $\pm$ 0.26 | 51.27 $\pm$ 0.05 | 83.34 $\pm$ 0.24 | 93.36 $\pm$ 0.05 | 83.50 $\pm$ 0.08 | 34.13 $\pm$ 0.27 | 95.74 $\pm$ 0.12 |
| FedROD | 97.74 $\pm$ 0.11 | 82.09 $\pm$ 0.04 | 49.33 $\pm$ 0.20 | 95.25 $\pm$ 0.20 | 94.88 $\pm$ 0.29 | 81.94 $\pm$ 0.10 | 38.34 $\pm$ 0.30 | 96.46 $\pm$ 0.20 |
| FedPAC | 96.64 $\pm$ 0.21 | 81.75 $\pm$ 0.17 | 53.83 $\pm$ 0.11 | 92.22 $\pm$ 0.09 | 94.36 $\pm$ 0.19 | 80.30 $\pm$ 0.13 | 38.51 $\pm$ 0.04 | 93.44 $\pm$ 0.05 |
| FedBABU | 85.46 $\pm$ 0.15 | 79.55 $\pm$ 0.27 | 51.78 $\pm$ 0.12 | 83.77 $\pm$ 0.27 | 90.62 $\pm$ 0.03 | 86.60 $\pm$ 0.15 | 33.18 $\pm$ 0.16 | 93.48 $\pm$ 0.16 |
| FedAvg | 85.77 $\pm$ 0.19 | 60.45 $\pm$ 0.02 | 47.25 $\pm$ 0.21 | 83.24 $\pm$ 0.10 | 91.90 $\pm$ 0.11 | 78.15 $\pm$ 0.05 | 32.64 $\pm$ 0.02 | 93.83 $\pm$ 0.26 |
| FedCAC | 95.90 $\pm$ 0.15 | 86.50 $\pm$ 0.08 | 53.46 $\pm$ 0.25 | 94.65 $\pm$ 0.13 | 93.12 $\pm$ 0.08 | 85.73 $\pm$ 0.02 | 39.14 $\pm$ 0.12 | 95.50 $\pm$ 0.25 |
| FedDBE | 94.90 $\pm$ 0.12 | 85.50 $\pm$ 0.06 | 55.13 $\pm$ 0.22 | 93.65 $\pm$ 0.10 | 92.26 $\pm$ 0.06 | 84.26 $\pm$ 0.01 | 38.75 $\pm$ 0.26 | 93.50 $\pm$ 0.22 |
| GPFL | 93.90 $\pm$ 0.10 | 84.50 $\pm$ 0.04 | 54.11 $\pm$ 0.20 | 92.65 $\pm$ 0.08 | 91.37 $\pm$ 0.04 | 83.12 $\pm$ 0.01 | 38.23 $\pm$ 0.10 | 95.50 $\pm$ 0.17 |
| FedAS | 97.22 $\pm$ 0.02 | 87.73 $\pm$ 0.23 | 54.92 $\pm$ 0.06 | 97.77 $\pm$ 0.01 | 98.14 $\pm$ 0.18 | 85.92 $\pm$ 0.26 | 38.19 $\pm$ 0.05 | 96.89 $\pm$ 0.22 |
| FedAF | 97.25 $\pm$ 0.04 | 86.25 $\pm$ 0.01 | 51.80 $\pm$ 0.02 | 95.25 $\pm$ 0.11 | 97.00 $\pm$ 0.02 | 85.10 $\pm$ 0.03 | 38.20 $\pm$ 0.16 | 95.50 $\pm$ 0.23 |
| Participation ratio = $80\%$ |  |  |  |  |  |  |  |  |
| FedLAG<br>(Ours) | 97.95 $\pm$ 0.20 | 88.02 $\pm$ 0.24 | 55.98 $\pm$ 0.10 | 95.72 $\pm$ 0.16 | 99.28 $\pm$ 0.21 | 85.25 $\pm$ 0.12 | 39.91 $\pm$ 0.09 | 96.14 $\pm$ 0.29 |
| PerAvg | 89.44 $\pm$ 0.24 | 83.61 $\pm$ 0.06 | 48.93 $\pm$ 0.04 | 83.88 $\pm$ 0.09 | 93.18 $\pm$ 0.05 | 81.59 $\pm$ 0.16 | 35.37 $\pm$ 0.23 | 94.53 $\pm$ 0.22 |
| FedROD | 97.61 $\pm$ 0.24 | 87.42 $\pm$ 0.11 | 48.18 $\pm$ 0.27 | 95.51 $\pm$ 0.09 | 94.85 $\pm$ 0.27 | 82.03 $\pm$ 0.18 | 35.26 $\pm$ 0.24 | 97.92 $\pm$ 0.28 |
| FedPAC | 95.05 $\pm$ 0.23 | 83.81 $\pm$ 0.11 | 50.37 $\pm$ 0.21 | 93.92 $\pm$ 0.07 | 94.24 $\pm$ 0.15 | 80.87 $\pm$ 0.19 | 36.12 $\pm$ 0.24 | 93.55 $\pm$ 0.10 |
| FedBABU | 96.51 $\pm$ 0.08 | 86.7 $\pm$ 0.28 | 50.69 $\pm$ 0.29 | 83.58 $\pm$ 0.18 | 90.03 $\pm$ 0.25 | 83.52 $\pm$ 0.29 | 37.24 $\pm$ 0.03 | 92.53 $\pm$ 0.12 |
| FedAvg | 89.88 $\pm$ 0.06 | 63.36 $\pm$ 0.18 | 46.46 $\pm$ 0.05 | 83.71 $\pm$ 0.15 | 92.18 $\pm$ 0.01 | 75.75 $\pm$ 0.21 | 30.13 $\pm$ 0.26 | 92.94 $\pm$ 0.19 |
| FedCAC | 95.82 $\pm$ 0.14 | 86.47 $\pm$ 0.09 | 52.92 $\pm$ 0.27 | 94.60 $\pm$ 0.12 | 92.92 $\pm$ 0.07 | 85.04 $\pm$ 0.03 | 37.84 $\pm$ 0.14 | 94.46 $\pm$ 0.23 |
| FedDBE | 94.87 $\pm$ 0.13 | 85.44 $\pm$ 0.07 | 54.94 $\pm$ 0.23 | 93.58 $\pm$ 0.09 | 91.97 $\pm$ 0.05 | 84.04 $\pm$ 0.02 | 37.62 $\pm$ 0.13 | 93.44 $\pm$ 0.21 |
| GPFL | 93.84 $\pm$ 0.09 | 84.44 $\pm$ 0.05 | 53.93 $\pm$ 0.19 | 92.58 $\pm$ 0.07 | 90.95 $\pm$ 0.03 | 83.07 $\pm$ 0.14 | 38.71 $\pm$ 0.11 | 92.46 $\pm$ 0.18 |
| FedAS | 96.62 $\pm$ 0.07 | 86.23 $\pm$ 0.25 | 53.73 $\pm$ 0.01 | 95.77 $\pm$ 0.25 | 97.28 $\pm$ 0.08 | 84.92 $\pm$ 0.07 | 37.89 $\pm$ 0.11 | 93.45 $\pm$ 0.15 |
| FedAF | 97.25 $\pm$ 0.04 | 86.25 $\pm$ 0.01 | 51.80 $\pm$ 0.02 | 95.25 $\pm$ 0.11 | 96.00 $\pm$ 0.09 | 79.10 $\pm$ 0.07 | 37.20 $\pm$ 0.02 | 91.50 $\pm$ 0.20 |
| Participation ratio = $60\%$ |  |  |  |  |  |  |  |  |
| FedLAG<br>(Ours) | 97.39 $\pm$ 0.20 | 87.25 $\pm$ 0.15 | 55.03 $\pm$ 0.22 | 96.88 $\pm$ 0.14 | 98.87 $\pm$ 0.26 | 83.24 $\pm$ 0.26 | 37.98 $\pm$ 0.24 | 94.82 $\pm$ 0.26 |
| PerAvg | 89.04 $\pm$ 0.08 | 83.85 $\pm$ 0.28 | 46.52 $\pm$ 0.09 | 84.18 $\pm$ 0.04 | 94.52 $\pm$ 0.22 | 79.76 $\pm$ 0.10 | 33.09 $\pm$ 0.07 | 93.78 $\pm$ 0.09 |
| FedROD | 97.15 $\pm$ 0.19 | 83.15 $\pm$ 0.21 | 48.39 $\pm$ 0.09 | 96.5 $\pm$ 0.16 | 94.82 $\pm$ 0.30 | 79.52 $\pm$ 0.07 | 35.19 $\pm$ 0.18 | 93.46 $\pm$ 0.19 |
| FedPAC | 84.9 $\pm$ 0.04 | 83.35 $\pm$ 0.19 | 49.75 $\pm$ 0.06 | 94.74 $\pm$ 0.04 | 94.62 $\pm$ 0.17 | 76.17 $\pm$ 0.19 | 34.95 $\pm$ 0.20 | 93.10 $\pm$ 0.12 |
| FedBABU | 95.85 $\pm$ 0.07 | 80.17 $\pm$ 0.15 | 49.61 $\pm$ 0.15 | 83.76 $\pm$ 0.02 | 90.78 $\pm$ 0.07 | 74.91 $\pm$ 0.27 | 31.96 $\pm$ 0.06 | 92.11 $\pm$ 0.12 |
| FedAvg | 90.04 $\pm$ 0.14 | 64.9 $\pm$ 0.05 | 46.51 $\pm$ 0.22 | 84.23 $\pm$ 0.05 | 93.47 $\pm$ 0.20 | 68.12 $\pm$ 0.07 | 22.24 $\pm$ 0.03 | 92.45 $\pm$ 0.18 |
| FedCAC | 95.75 $\pm$ 0.16 | 86.43 $\pm$ 0.10 | 51.98 $\pm$ 0.45 | 94.55 $\pm$ 0.14 | 92.85 $\pm$ 0.09 | 75.14 $\pm$ 0.03 | 35.20 $\pm$ 0.16 | 93.40 $\pm$ 0.27 |
| FedDBE | 94.80 $\pm$ 0.14 | 85.40 $\pm$ 0.07 | 54.23 $\pm$ 0.24 | 93.50 $\pm$ 0.11 | 91.90 $\pm$ 0.07 | 76.42 $\pm$ 0.02 | 34.81 $\pm$ 0.14 | 92.40 $\pm$ 0.23 |
| GPFL | 93.80 $\pm$ 0.11 | 84.40 $\pm$ 0.06 | 53.85 $\pm$ 0.21 | 92.50 $\pm$ 0.09 | 90.90 $\pm$ 0.05 | 76.95 $\pm$ 0.03 | 34.65 $\pm$ 0.12 | 91.40 $\pm$ 0.19 |
| FedAS | 96.24 $\pm$ 0.11 | 85.79 $\pm$ 0.12 | 51.24 $\pm$ 0.23 | 95.34 $\pm$ 0.16 | 96.28 $\pm$ 0.22 | 76.92 $\pm$ 0.01 | 32.89 $\pm$ 0.09 | 91.26 $\pm$ 0.05 |
| FedAF | 96.25 $\pm$ 0.07 | 85.25 $\pm$ 0.01 | 51.00 $\pm$ 0.02 | 94.25 $\pm$ 0.15 | 95.20 $\pm$ 0.12 | 78.30 $\pm$ 0.09 | 32.50 $\pm$ 0.02 | 90.50 $\pm$ 0.24 |
| Participation ratio = $40\%$ |  |  |  |  |  |  |  |  |
| FedLAG<br>(Ours) | 97.55 $\pm$ 0.15 | 86.11 $\pm$ 0.62 | 53.19 $\pm$ 0.09 | 96.14 $\pm$ 0.07 | 96.17 $\pm$ 0.22 | 76.34 $\pm$ 0.08 | 36.39 $\pm$ 0.16 | 93.58 $\pm$ 0.13 |
| PerAvg | 89.43 $\pm$ 0.02 | 81.44 $\pm$ 0.01 | 44.38 $\pm$ 0.23 | 84.17 $\pm$ 0.02 | 91.71 $\pm$ 0.12 | 73.57 $\pm$ 0.07 | 32.10 $\pm$ 0.27 | 90.09 $\pm$ 0.29 |
| FedROD | 97.15 $\pm$ 0.07 | 82.53 $\pm$ 0.19 | 47.76 $\pm$ 0.42 | 95.73 $\pm$ 0.18 | 94.72 $\pm$ 0.03 | 78.03 $\pm$ 0.07 | 35.21 $\pm$ 0.16 | 93.30 $\pm$ 0.16 |
| FedPAC | 84.24 $\pm$ 0.29 | 82.63 $\pm$ 0.21 | 48.38 $\pm$ 0.13 | 86.42 $\pm$ 0.19 | 93.57 $\pm$ 0.07 | 75.91 $\pm$ 0.27 | 32.26 $\pm$ 0.28 | 93.27 $\pm$ 0.02 |
| FedBABU | 95.34 $\pm$ 0.09 | 76.19 $\pm$ 0.06 | 45.01 $\pm$ 0.04 | 83.97 $\pm$ 0.07 | 91.17 $\pm$ 0.25 | 72.21 $\pm$ 0.06 | 31.46 $\pm$ 0.13 | 91.44 $\pm$ 0.14 |
| FedAvg | 90.71 $\pm$ 0.14 | 69.66 $\pm$ 0.17 | 42.21 $\pm$ 0.18 | 84.45 $\pm$ 0.19 | 93.88 $\pm$ 0.28 | 65.63 $\pm$ 0.25 | 16.05 $\pm$ 0.25 | 91.05 $\pm$ 0.03 |
| FedCAC | 96.05 $\pm$ 0.18 | 77.19 $\pm$ 0.09 | 48.15 $\pm$ 0.12 | 94.80 $\pm$ 0.14 | 93.20 $\pm$ 0.09 | 75.20 $\pm$ 0.03 | 33.45 $\pm$ 0.16 | 92.80 $\pm$ 0.27 |
| FedDBE | 95.48 $\pm$ 0.13 | 75.80 $\pm$ 0.08 | 51.27 $\pm$ 0.26 | 94.25 $\pm$ 0.12 | 92.10 $\pm$ 0.07 | 74.20 $\pm$ 0.02 | 32.74 $\pm$ 0.13 | 91.70 $\pm$ 0.22 |
| GPFL | 94.36 $\pm$ 0.12 | 74.60 $\pm$ 0.05 | 50.86 $\pm$ 0.22 | 93.17 $\pm$ 0.10 | 91.20 $\pm$ 0.06 | 73.20 $\pm$ 0.02 | 32.26 $\pm$ 0.11 | 91.52 $\pm$ 0.20 |
| FedAS | 95.76 $\pm$ 0.15 | 84.79 $\pm$ 0.22 | 50.89 $\pm$ 0.01 | 94.56 $\pm$ 0.19 | 95.28 $\pm$ 0.03 | 74.22 $\pm$ 0.17 | 32.89 $\pm$ 0.16 | 90.45 $\pm$ 0.02 |
| FedAF | 95.50 $\pm$ 0.10 | 84.50 $\pm$ 0.02 | 50.40 $\pm$ 0.02 | 93.50 $\pm$ 0.18 | 94.60 $\pm$ 0.15 | 67.70 $\pm$ 0.12 | 32.15 $\pm$ 0.04 | 89.80 $\pm$ 0.28 |
| Participation ratio = $20\%$ |  |  |  |  |  |  |  |  |
| FedLAG<br>(Ours) | 97.18 $\pm$ 0.13 | 85.21 $\pm$ 0.29 | 51.35 $\pm$ 0.02 | 96.34 $\pm$ 0.25 | 96.68 $\pm$ 0.01 | 69.26 $\pm$ 0.09 | 33.75 $\pm$ 0.14 | 91.42 $\pm$ 0.18 |
| PerAvg | 89.11 $\pm$ 0.17 | 81.25 $\pm$ 0.20 | 43.21 $\pm$ 0.28 | 84.84 $\pm$ 0.20 | 94.89 $\pm$ 0.06 | 61.25 $\pm$ 0.14 | 23.98 $\pm$ 0.25 | 89.70 $\pm$ 0.14 |
| FedROD | 97.02 $\pm$ 0.29 | 81.72 $\pm$ 0.16 | 46.17 $\pm$ 0.06 | 96.02 $\pm$ 0.08 | 94.87 $\pm$ 0.12 | 68.47 $\pm$ 0.19 | 26.76 $\pm$ 0.08 | 90.70 $\pm$ 0.28 |
| FedPAC | 89.12 $\pm$ 0.04 | 83.13 $\pm$ 0.01 | 44.77 $\pm$ 0.06 | 89.63 $\pm$ 0.04 | 94.60 $\pm$ 0.18 | 63.01 $\pm$ 0.04 | 25.42 $\pm$ 0.25 | 88.44 $\pm$ 0.61 |
| FedBABU | 95.92 $\pm$ 0.26 | 80.75 $\pm$ 0.08 | 42.59 $\pm$ 0.03 | 84.63 $\pm$ 0.19 | 91.42 $\pm$ 0.03 | 65.12 $\pm$ 0.18 | 21.54 $\pm$ 0.30 | 87.35 $\pm$ 0.10 |
| FedAvg | 90.61 $\pm$ 0.11 | 65.47 $\pm$ 0.17 | 41.28 $\pm$ 0.09 | 85.02 $\pm$ 0.30 | 88.43 $\pm$ 0.07 | 59.01 $\pm$ 0.26 | 15.99 $\pm$ 0.10 | 86.35 $\pm$ 0.16 |
| FedCAC | 96.77 $\pm$ 0.03 | 84.62 $\pm$ 0.19 | 47.22 $\pm$ 1.52 | 96.01 $\pm$ 0.15 | 96.33 $\pm$ 0.11 | 68.93 $\pm$ 0.19 | 26.12 $\pm$ 0.04 | 91.10 $\pm$ 0.18 |
| FedDBE | 95.73 $\pm$ 0.11 | 83.76 $\pm$ 0.09 | 50.12 $\pm$ 0.12 | 95.04 $\pm$ 0.05 | 95.52 $\pm$ 0.11 | 67.75 $\pm$ 0.04 | 25.43 $\pm$ 0.04 | 90.14 $\pm$ 0.08 |
| GPFL | 94.15 $\pm$ 0.10 | 82.11 $\pm$ 0.25 | 49.85 $\pm$ 0.01 | 93.25 $\pm$ 0.21 | 93.56 $\pm$ 0.08 | 66.38 $\pm$ 0.07 | 24.01 $\pm$ 0.11 | 88.45 $\pm$ 0.15 |
| FedAS | 95.36 $\pm$ 0.12 | 84.11 $\pm$ 0.02 | 50.26 $\pm$ 0.03 | 93.25 $\pm$ 0.19 | 94.28 $\pm$ 0.18 | 67.38 $\pm$ 0.14 | 32.01 $\pm$ 0.04 | 89.45 $\pm$ 0.32 |
| FedAF | 95.24 $\pm$ 0.15 | 84.35 $\pm$ 0.03 | 50.11 $\pm$ 0.04 | 93.10 $\pm$ 0.20 | 94.40 $\pm$ 0.17 | 67.55 $\pm$ 0.13 | 32.07 $\pm$ 0.05 | 89.50 $\pm$ 0.30 |

## Appendix F Detailed Ablation Test

### F.1 Training Time

We have conducted a thorough evaluation of FedLAG’s robustness in comparison to other baselines, focusing on convergence rate and accuracy. Additionally, we evaluated the robustness of FedLAG by computing the average computation time per round for various FL algorithms, and the results are presented in Table [8](#A6.T8).

As depicted in the table, FedLAG demonstrates a lower computation cost compared to other baselines. This efficiency stems from our approach, which involves measuring the gradient and applying a straightforward analysis on the server. It is noteworthy that servers typically possess significantly more computational resources than individual devices. As a result, our proposed FedLAG effectively leverages the server’s computation capacity, leading to substantial time savings.

The efficiency gains achieved by FedLAG make it particularly well-suited for Internet of Things (IoT) systems with constrained computational resources. In contrast, FedPAC tends to be less suitable for such environments, as it demands substantial computation time on distributed devices. Our findings emphasize the practical advantages of FedLAG in scenarios where computational efficiency is crucial, such as IoT systems with limited resources.

**Table 8: Training time (second/round) of FedLAG vs. baselines.**
| Method<br>Dataset | FedBABU | PerAvg | FedROD | FedPAC | FedLAG<br>(Ours) |
| --- | --- | --- | --- | --- | --- |
| Mnist | 11 $\pm$ 0.11 | 9 $\pm$ 0.13 | 12 $\pm$ 0.32 | 342 $\pm$ 0.06 | 8 $\pm$ 0.37 |
| Cifar10 | 23 $\pm$ 0.25 | 19 $\pm$ 0.62 | 22 $\pm$ 0.18 | 474 $\pm$ 0.26 | 15 $\pm$ 0.77 |
| Cifar100 | 22 $\pm$ 0.13 | 20 $\pm$ 0.11 | 21 $\pm$ 0.55 | 698 $\pm$ 0.77 | 17 $\pm$ 0.34 |
| EMNIST | 117 $\pm$ 0.25 | 108 $\pm$ 0.20 | 114 $\pm$ 0.70 | 503 $\pm$ 0.43 | 112 $\pm$ 0.43 |

### F.2 Top K Layers

Figure: Figure 8: The impact of the top $k$ conflict layer hyperparameter on accuracy on cifar 10.
Refer to caption: /html/2410.02845/assets/x19.png

### F.3 Gradient Score

Figure: Figure 10: The impact of the conflict hyperparameter $\xi$ on accuracy on cifar 10.
Refer to caption: /html/2410.02845/assets/x21.png

### F.4 Conflicted Score during the Training of FedAvg

#### F.4.1 Mnist Dataset

Figure: Figure 12: The GC score on each layer (Dataset MNIST - IID).
Refer to caption: /html/2410.02845/assets/x23.png

Figure: Figure 13: The GC score on each layer (Dataset MNIST - NIID $\alpha=0.1$).
Refer to caption: /html/2410.02845/assets/x24.png

In the MNIST dataset under the IID setting shown in Fig. [12](#A6.F12), it’s clear that the gradient conflict score is close to 0 0 in the early rounds. This is due to training on IID settings and the simplicity of the MNIST dataset. Consequently, there’s minimal negative transfer of gradient conflict among users. When compared with that of the conflicted scores in MNIST dataset with non-IID settings, we can see that the problem of conflicted scores become significant.

#### F.4.2 Cifar10 Dataset

Figure: Figure 14: The GC score on each layer (Dataset Cifar10 - NonIID with $\alpha=0.1$).
Refer to caption: /html/2410.02845/assets/x25.png

Figure: Figure 15: The GC score on each layer (Dataset Cifar10 - IID).
Refer to caption: /html/2410.02845/assets/x26.png

Fig. [14](#A6.F14) shows the fluctuating scores observed during training on the Cifar10 dataset under non-IID conditions ($\alpha=0.1$). We excluded layers without trainable parameters, such as Batch Normalization layers, from visualization. Initially, from round 0 0 to $150$, the conflicting gradients are prominent in the later layers (i.e., layers $12$ to $16$). However, as training progresses, the intensity of conflicts in these layers diminishes, resulting in a more balanced distribution of conflicting gradients across layers $8$ to $16$. This indicates that the nature of conflicting scores varies throughout the training stages of FL. In the IID settings (Fig. [15](#A6.F15)), the conflicted scores is much reduced as the divergence among domains are small. However, we can easily see that the layers with high gradient conflicts are not densely distributed at the very first and the very last layers.

## Appendix G Toy dataset description

We visualize the training data in [16](#A7.F16), and [17](#A7.F17). In the FDG setting, the users are from different domains. To this end, we design the data where the point are distributed into rectangular with different size and shape. The rationale of designing the data distribution is as follows:

- •
The global dataset consists of two classes from two rectangular, which has the classification boundary is equal to $y=0$.
- •
Each domain-wise dataset has different classification boundary (e.g., $x=-6$ for domain $1$). We add the noisy data on every domains so that the user assign to each domain will tend to learn the local boundary instead of the global boundary. Thus, we can observe the gradient divergence more clearly, as the global boundary is not the optimal solution when learn on local dataset.
- •
All of the local classification boundary is orthogonal from the global classification boundary, thus, we can make the learning more challenging despite the simplicity of the toy dataset.

Figure: Figure 16: Illustration of users with different domains.
Refer to caption: /html/2410.02845/assets/image-lib/illustrative-toy/domain1.png

Figure: Figure 17: Illustration of users with same domains.
Refer to caption: /html/2410.02845/assets/image-lib/illustrative-toy/domain1-iid.png

## Appendix H Proof

### H.1 Assumptions and Definitions

To come up with our theoretical analysis, we first adopt the following assumptions:

###### Assumption H.1 ( $L$ -smooth) .

Local objective $F_{u}(w)$ is differentiable and $L$-smooth for all $u\in\{1,2,\ldots,U\}$, i.e., $\|\nabla\mathcal{L}_{u}(w)-\nabla\mathcal{L}_{u}(w^{\prime})\|\leq L\|w-w^{\prime}\|,~{}\forall w,w^{\prime}\in\mathbb{R}^{d}$.

###### Assumption H.2 ( $\mu$ -strongly convex) .

Local objective $F_{u}(w)$ is differentiable and $L$-smooth for all $u\in\{1,2,\ldots,U\}$, i.e., $\|\nabla\mathcal{L}_{u}(w)-\nabla\mathcal{L}_{u}(w^{\prime})\|\geq\mu\|w-w^{\prime}\|,~{}\forall w,w^{\prime}\in\mathbb{R}^{d}$.

###### Assumption H.3 (Bounded data heterogeneity at optimum) .

The norm of the user gradients at the global optima $w^{*}$ is bounded as follows: $\frac{1}{U}\sum^{U}_{u=1}\|\nabla\mathcal{L}_{u}(w^{*})\|^{2}\leq\sigma^{2}_{*}$.

###### Assumption H.4 (Bounded global gradient variance) .

There exists a constant $\sigma^{2}_{g}>0$ such that the global gradient variance is bounded as follows: $\frac{1}{U}\sum^{U}_{u=1}\|\nabla\mathcal{L}_{u}(w)-\nabla\mathcal{L}(w)\|^{2}\leq\sigma^{2}_{g}$.

###### Assumption H.1 ( L 𝐿 L -smooth) .

###### Assumption H.2 ( μ 𝜇 \mu -strongly convex) .

###### Assumption H.3 (Bounded data heterogeneity at optimum) .

###### Assumption H.4 (Bounded global gradient variance) .

### H.2 Preliminaries

###### Lemma H.5 (Jensen’s Inequality) .

For any $a_{u}\in\mathbb{R}^{d},u\in\{1,2,\ldots,U\}$:

$$ $\displaystyle\|\frac{1}{U}\sum^{U}_{u=1}\mathbf{a}_{u}\|^{2}$ $\displaystyle\leq\frac{1}{U}\sum^{U}_{u=1}\|\mathbf{a}_{u}\|^{2}$ (7) $\displaystyle\|\sum^{U}_{u=1}\mathbf{a}_{u}\|^{2}$ $\displaystyle\leq U\sum^{U}_{u=1}\|\mathbf{a}_{u}\|^{2}$ (8) $$

###### Lemma H.6 ( (Khaled et al., 2020 ) ) .

If $\mathcal{L}$ is smooth and convex, then

$$ $\displaystyle\|\nabla\mathcal{L}(w)-\nabla\mathcal{L}(w^{\prime})\|^{2}\leq 2L(\mathcal{L}(w)-\mathcal{L}(w^{\prime})-\langle\nabla\mathcal{L}(w^{\prime}),w-w^{\prime}\rangle)$ (9) $$

###### Lemma H.7 (Co-coercivity of convex $L$ -smooth function) .

If $F$ is $L$-smooth and convex then

$$ $\displaystyle\langle\nabla\mathcal{L}(w)-\nabla\mathcal{L}(w^{\prime}),w-w^{\prime}\rangle\geq\frac{1}{L}\|\nabla\mathcal{L}(w)-\nabla\mathcal{L}(w^{\prime})\|^{2}.$ (10) $$

A direct consequence of this lemma is:

$$ $\displaystyle\langle\nabla\mathcal{L}(w),w-w^{*}\rangle\geq\frac{1}{L}\|\nabla\mathcal{L}(w)\|^{2},$ (11) $$

where $w^{*}$ is a minimizer of $\mathcal{L}(w)$

###### Lemma H.8 (Co-coercivity of $\mu$ -strongly convex function) .

If $F$ is $L$-smooth and convex then

$$ $\displaystyle\langle\nabla\mathcal{L}(w)-\nabla\mathcal{L}(w^{\prime}),w-w^{\prime}\rangle\geq\frac{1}{L}\|\nabla\mathcal{L}(w)-\nabla\mathcal{L}(w^{\prime})\|^{2}.$ (12) $$

A direct consequence of this lemma is:

$$ $\displaystyle\langle\nabla\mathcal{L}(w),w-w^{*}\rangle\geq\frac{1}{L}\|\nabla\mathcal{L}(w)\|^{2},$ (13) $$

where $w^{*}$ is a minimizer of $F(w)$

###### Lemma H.5 (Jensen’s Inequality) .

###### Lemma H.6 ( (Khaled et al., 2020 ) ) .

###### Lemma H.7 (Co-coercivity of convex L 𝐿 L -smooth function) .

###### Lemma H.8 (Co-coercivity of μ 𝜇 \mu -strongly convex function) .

### H.3 Reformulation for Vanilla pFL and FedLAG

In the appendix, we aim to achieve two main objectives. Firstly, we intend to illustrate the distinction between FedLAG and other Federated Learning (FL) algorithms that do not incorporate FedLAG. Secondly, we aim to demonstrate the convergence under novel FedLAG settings. To accomplish these goals, we introduce additional notations in the appendix. It is essential to note that these extensively defined notations are exclusive to the appendix and do not compromise the generality of the main paper. Towards the conclusion of the appendix, our objective is to consolidate the formulations and revert the convergence formulation back to its original form.

We define $w^{(r)}_{u},w^{(r)}_{g}$ as local and global model parameters, respectively. In the appendix, to distinguish between two concepts Vanilla pFL and FedLAG (with our proposed layer-wise personalization technique), we also define the following auxiliary variables that will used in the proof:

- •
Layer-wise Personalized Model: To represent the layer disentanglement architecture, where the model are layer-wise personalized, we define the user’s model $w^{(r)}_{\textrm{LAG},u}=\{\theta^{(r)}_{s,u},\theta^{(r)}_{p,u}\}$, where $\{\theta^{(r)}_{s,u},\theta^{(r)}_{p,u}\}$ represents the model layers that is assigned for global aggregation, and the one that is assigned for personalized purpose, respectively.
- •
Vanilla FL Model: We define the user’s model $w^{(r)}_{\textrm{VFL},u}=\{\theta^{(r)}_{s,u},\theta^{(r)}_{q,u}\}$, where $\{\theta^{(r)}_{s,u},\theta^{(r)}_{q,u}\}$ represents the model layers that is assigned for global aggregation, and the one that is assigned for personalized purpose. However, in Vanilla settings, we do the shared global aggregation on the personalized model. Thus, the model update at global aggregation is similar to the shared layers.

As the model layer $\theta^{(r)}_{s,u}$ is assigned for the global aggregation. Without the loss of generality, the two model layers have same attributes and update rules. Consequently, we use the same notations $\theta^{(r)}_{s,u}$ for the two model layers according to Layer-wise Personalized Model and Vanilla FL Model.

- •
Accumulated Local Gradient: We have the local gradient update on the whole user’s network:
$\displaystyle h^{(r)}_{u}=\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}_{u}(w^{(r,e)}_{u})=w^{(r,E)}_{u}-w^{(r)}_{u}.$
(14)
By decomposing into shared aggregation layer set $\theta^{(r)}_{s,u}$, and personalized layer set $\theta^{(r)}_{p,u}$, we have the following:
$\displaystyle h^{(r)}_{s,u}=\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}_{u}(\theta^{(r,e)}_{s,u})=\theta^{(r,E)}_{s,u}-\theta^{(r)}_{s,u},$
(15)
$\displaystyle h^{(r)}_{p,u}=\sum^{E-1}_{e=0}\eta\nabla\mathcal{L}_{u}(\theta^{(r,e)}_{p,u})=\theta^{(r,E)}_{p,u}-\theta^{(r)}_{p,u}.$
(16)
As the local update on both vanilla and FedLAG remain the same, thus, we have the gradient trajectories on both setting similar. For instance,
$\displaystyle h^{(r)}_{p,u}=h^{(r)}_{q,u}.$
(17)
- •
Global Gradient Aggregation: For FedLAG, we have the two different update rules for two different type of layers:
$\displaystyle\theta^{(r+1,0)}_{s,u}=\frac{1}{U}\sum^{U}_{u=1}\theta^{(r)}_{s,u}=\theta^{(r,0)}_{s,u}-\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{s,u},$
$\displaystyle\theta^{(r+1,0)}_{p,u}=\frac{1}{U}\sum^{U}_{u=1}\theta^{(r)}_{p,u}=\theta^{(r,0)}_{p,u}-h^{(r)}_{p,u}.$
(18)
On Vanilla pFL, we have the same update rule for two different type of layers:
$\displaystyle\theta^{(r+1,0)}_{su}=\frac{1}{U}\sum^{U}_{u=1}\theta^{(r)}_{s,u}=\theta^{(r,0)}_{s,u}-\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{s,u},$
(19)
$\displaystyle\theta^{(r+1,0)}_{q,u}=\frac{1}{U}\sum^{U}_{u=1}\theta^{(r)}_{p,u}=\theta^{(r,0)}_{p,u}-\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{p,u}.$
(20)
- •
Loss Function: Due to the definition of two models $\theta_{s,u},\theta_{p,u}$, we can define the $\mathcal{L}_{u}(\theta^{(r)}_{s,g}),\mathcal{L}_{u}(\theta^{(r)}_{p,g})$ are the two empirical loss that aim to design a global optimal models $w^{*}=\{\theta^{*}_{s,u},\theta^{*}_{p,u}\}$, respectively. Therefore, we have $\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})=\mathcal{L}_{u}(\theta^{(r)}_{s,g})+\mathcal{L}_{u}(\theta^{(r)}_{p,g})$.

### H.4 Proof on Lemma 5.1

By leveraging Taylor approximation, we have the following loss estimation for joint model parameters at each communication round. For the layer-wise personalized aggregation, we have:

$$ $\displaystyle\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{p,u})$ $\displaystyle\overset{(a)}{=}\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r)}_{p,u})+(\theta^{(r+1)}_{p,u}-\theta^{(r)}_{p,u})\nabla_{\theta^{(r)}_{s,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+\mathcal{O}(\eta)$ $\displaystyle\overset{(b)}{=}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})\nabla_{\theta^{(r)}_{s,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})$ $\displaystyle~{}~{}~{}~{}+(\theta^{(r+1)}_{p,u}-\theta^{(r)}_{p,u})\nabla_{\theta^{(r)}_{p,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+\mathcal{O}(\eta)$ $\displaystyle=\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})\sum^{E-1}_{e=0}\eta\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(\theta^{(r,e)}_{s,u},\theta^{(r,e)}_{p,u})$ $\displaystyle~{}~{}~{}~{}+(\theta^{(r+1)}_{p,u}-\theta^{(r)}_{p,u})\sum^{E-1}_{e=0}\eta\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(\theta^{(r,e)}_{s,u},\theta^{(r,e)}_{p,u})$ $\displaystyle=\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})h^{(r)}_{s,u}+(\theta^{(r+1)}_{p,u}-\theta^{(r)}_{p,u})h^{(r)}_{p,u}+\mathcal{O}(\eta),$ (21) $$

where (a) and (b) hold due to the approximations of $\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{p,u})$ according to $\theta^{(r+1)}_{s,u}$ and $\theta^{(r+1)}_{p,u}$, respectively. For normal update, we have:

$$ $\displaystyle\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{q,u})$ $\displaystyle\overset{(a)}{=}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r+1)}_{q,u})+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})\nabla_{\theta^{(r)}_{s,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+\mathcal{O}(\eta)$ $\displaystyle\overset{(b)}{=}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})\nabla_{\theta^{(r)}_{s,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})$ $\displaystyle~{}~{}~{}~{}+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})\nabla_{\theta^{(r)}_{q,u}}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+\mathcal{O}(\eta)$ $\displaystyle=\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})\sum^{E-1}_{e=0}\eta\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(\theta^{(r,e)}_{s,u},\theta^{(r,e)}_{q,u})$ $\displaystyle~{}~{}~{}~{}+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})\sum^{E-1}_{e=0}\eta\nabla_{\theta^{(r,e)}_{q,u}}\mathcal{L}_{u}(\theta^{(r,e)}_{s,u},\theta^{(r,e)}_{q,u})$ $\displaystyle=\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})h^{(r)}_{s,u}+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})h^{(r)}_{q,u}+\mathcal{O}(\eta)$ $\displaystyle\overset{(c)}{=}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})h^{(r)}_{s,u}+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})h^{(r)}_{p,u}+\mathcal{O}(\eta),$ (22) $$

where $(a)$ and $(b)$ hold due to the approximations of $\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{q,u})$ according to $\theta^{(r+1)}_{s,u}$ and $\theta^{(r+1)}_{q,u}$, respectively. $(c)$ holds according to the Equation [17](#A8.E17). Specifically, in the one-step layer-wise personalized aggregation, we want to estimate the divergence that loss function from LAG move from the vanilla pFL. The difference between the two loss function after the update is measured by subtracting $\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{p,u})$ in Equation equation [21](#A8.E21) from $\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{q,u})$ in Equation equation [22](#A8.E22) as follows:

$$ $\displaystyle~{}~{}~{}~{}~{}\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{p,u})-\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{q,u})$ $\displaystyle=\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{p,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})h^{(r)}_{s,u}+(\theta^{(r+1)}_{p,u}-\theta^{(r)}_{p,u})h^{(r)}_{p,u}\Big{)}$ $\displaystyle-\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{s,u},\theta^{(r)}_{q,u})+(\theta^{(r+1)}_{s,u}-\theta^{(r)}_{s,u})h^{(r)}_{s,u}+(\theta^{(r+1)}_{q,u}-\theta^{(r)}_{q,u})h^{(r)}_{p,u}\Big{)}+\mathcal{O}(\eta)$ $\displaystyle=(\theta^{(r+1)}_{p,u}-\theta^{(r+1)}_{q,u})h^{(r)}_{p,u}+(\theta^{(r)}_{p,u}-\theta^{(r)}_{q,u})h^{(r)}_{p,u}+\mathcal{O}(\eta),$ $\displaystyle\overset{(a)}{=}(\theta^{(r+1)}_{p,u}-\theta^{(r+1)}_{q,u})h^{(r)}_{p,u}+\mathcal{O}(\eta),$ $\displaystyle=-\eta(h^{(r)}_{p,u}-\frac{1}{U}\sum^{U}_{v=1}h^{(r)}_{p,v})^{\top}h^{(r)}_{p,u}+\mathcal{O}(\eta),$ $\displaystyle=-\eta\frac{1}{U}\sum^{U}_{v=1}(h^{(r)}_{p,u}-h^{(r)}_{p,v})^{\top}h^{(r)}_{p,u}+\mathcal{O}(\eta),$ $\displaystyle=-\eta\frac{1}{U}\sum^{U}_{v=1}(\|h^{(r)}_{p,u}\|^{2}-h^{(r)\top}_{p,v}h^{(r)}_{p,u})+\mathcal{O}(\eta),$ (23) $$

where $\mathcal{O}(\eta)$ represents the remainder which depends on the variable $\eta$. $(a)$ holds due to the model are considered to be trained at the same starting position at each round $r$, thus $\theta^{(r)}_{p,u}=\theta^{(r)}_{q,u}$ . Without loss of generality, we assume that $\|h^{(r)}_{p,u}\|\neq 0$, and $h^{(r)}_{p,u}=\{h^{(r)}_{l,u}|\forall l\in\mathbb{L}_{p}\}$ then:

$$ $\displaystyle\|h^{(r)}_{p,u}\|^{2}-h^{(r)\top}_{p,v}h^{(r)}_{p,u}$ $\displaystyle=\sum_{l\in\mathbb{L}_{p}}(\|h^{(r)}_{l,u}\|^{2}-h^{(r)\top}_{l,v}h^{(r)}_{l,u}),$ $\displaystyle=\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|(\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|)\stackrel{{\scriptstyle(a)}}{{>}}0,$ (24) $$

where $\mathbb{L}_{p}$ represents the layers that have gradient conflict and turned become personalized layers. Due to the gradient conflict definition which has proposed in Definition [2.1](#S2.Thmtheorem1), we have: $\cos\Phi^{(l)}_{u,v}<0$ as $\Phi^{(l)}_{u,v}>\pi/2$, which makes the inequality (a) hold. Hence, the above difference is negative, if $\eta$ is sufficiently small. As such, the difference between the vanilla FL and LAG loss functions is also negative, if $\eta$ is sufficiently small.

### H.5 Proof on Lemma 5.2

From Lemma [5.1](#A11.EGx6), we have

$$ $\displaystyle\mathcal{L}(\theta^{(r+1)}_{s,g},\theta^{(r+1)}_{p,g})-\mathcal{L}(\theta^{(r+1)}_{s,g},\theta^{(r+1)}_{q,g})$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{p,u})-\frac{1}{U}\sum^{U}_{u=1}\mathcal{L}_{u}(\theta^{(r+1)}_{s,u},\theta^{(r+1)}_{q,u})$ $\displaystyle=-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\Big{(}\|h^{(r)}_{p,u}\|^{2}-h^{(r)\top}_{p,v}h^{(r)}_{p,u}\Big{)}$ $\displaystyle\overset{(a)}{=}-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}}\Big{(}\|h^{(r)}_{l,u}\|^{2}-h^{(r)}_{l,v}\cdot h^{(r)}_{l,u}\Big{)}$ $\displaystyle\overset{(b)}{=}-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\Big{(}\|h^{(r)}_{l,u}\|^{2}-\frac{(h^{(r)}_{l,u})\cdot h^{(r)}_{l,v}}{\|h^{(r)}_{l,u}\|\|h^{(r)}_{l,v}\|}\|h^{(r)}_{l,u}\|\|h^{(r)}_{l,v}\|\Big{)}$ $\displaystyle=-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\frac{h^{(r)}_{l,u}\cdot h^{(r)}_{l,v}}{\|h^{(r)}_{l,u}\|\|h^{(r)}_{l,v}\|}\|h^{(r)}_{l,v}\|\Big{)}$ $\displaystyle\overset{(c)}{=}-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}<0,$ (25) $$

where $(a)$ holds according to the Lemma [H.9](#A8.Thmtheorem9) and $(b)$ holds as the difference only be made on the personalized layers $\theta^{(r)}_{p,u}$. $(c)$ holds according to Definition [2.1](#S2.Thmtheorem1)

###### Lemma H.9 .

Given two vector $h^{(r)}_{p,v},h^{(r)}_{p,u}\in\mathbb{R}^{P}$, we have following relationship:

$$ $\displaystyle h^{(r)\top}_{p,v}h^{(r)}_{p,u}=h^{(r)}_{p,v}\cdot h^{(r)}_{p,u}=\sum^{\mathbb{L}}_{l=1}h^{(r)}_{l,u}\cdot h^{(r)}_{l,v},$ (26) $$

where $P=\sum^{L}_{l=1}P_{l}$ is the number of parameters of the FL model, $P_{l}$ is the number of parameters on each layer $l$.

###### Lemma H.9 .

### H.6 Upper boundary on Layer-wise Loss Improvement Approximation

Consider the inequality in Eq. equation [25](#A8.E25). To find an upper boundary on the loss improvement, we consider the maximal gradient norm $\widetilde{h}^{(r)}_{p,u}=\arg\max_{h^{(r)}_{l,u}}\|h^{(r)}_{l,u}\|$. Thus, we have:

$$ $\displaystyle\mathcal{L}(\theta^{(r+1)}_{s,g},\theta^{(r+1)}_{p,g})-\mathcal{L}(\theta^{(r+1)}_{s,g},\theta^{(r+1)}_{q,g})$ $\displaystyle=-\eta\frac{1}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}$ $\displaystyle=-\eta\mathbb{E}_{u,v}\Bigg{[}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Bigg{]}$ $\displaystyle\leq-\eta\mathbb{E}_{u,v}\Bigg{[}\sum_{l\in\mathbb{L}_{p}}\|\widetilde{h}^{(r)}_{p,u}\|\Big{(}\|\widetilde{h}^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|\widetilde{h}^{(r)}_{p,v}\|\Big{)}\Bigg{]}$ $\displaystyle=-\eta\mathbb{E}_{u,v}\Bigg{[}L_{p}\|\widetilde{h}^{(r)}_{p,u}\|\Big{(}\|\widetilde{h}^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|\widetilde{h}^{(r)}_{p,v}\|\Big{)}\Bigg{]}$ $\displaystyle=-\eta\mathbb{E}_{u,v}\Bigg{[}\frac{L_{p}}{L^{2}}\|\widetilde{h}^{(r)}_{u}\|\Big{(}\|\widetilde{h}^{(r)}_{u}\|-\cos\Phi^{(l)}_{u,v}\|\widetilde{h}^{(r)}_{v}\|\Big{)}\Bigg{]}$ $\displaystyle\leq-\eta\mathbb{E}_{u,v}\Bigg{[}\|\widetilde{h}^{(r)}_{u}\|\Big{(}\|\widetilde{h}^{(r)}_{u}\|-\cos\Phi^{(l)}_{u,v}\|\widetilde{h}^{(r)}_{v}\|\Big{)}\Bigg{]}.$ (27) $$

Inequality in Eq. equation [27](#A8.E27) proves that the improvement of layer-wise loss is upper-bounded by the model gradient norm. Furthermore, the improvement on AI model depends on the percentage of layers that required to be personalized rather than that of the number of layers. As a consequence, the formula is agnostic to the over-parameterized of the AI model.

### H.7 Bounding user aggregate gradients

###### Lemma H.10 ( (Jhunjhunwala et al., 2023 ) , Bounding user aggregate gradients on Vanilla FL) .

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{s,u})\|^{2}$ $\displaystyle\leq\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{s,g}-w^{(r)}_{\textrm{VFL},g}\|^{2}$ (28) $\displaystyle+6EL(\nabla\mathcal{L}(w^{(r)}_{\textrm{VFL},g})-\nabla\mathcal{L}(w^{*}_{g}))+3E\sigma^{2}_{*}.$ (29) $$

The lemma shows that the local gradient variance $\mathcal{L}_{u}(w^{(r,e)}_{s,u})$ on user $u$ after $E$ local epochs is always bounded by a certain threshold.

###### Lemma H.11 (Bounding user aggregate gradients on Layer-wise Personalized FL) .

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\|^{2}$ $\displaystyle\leq\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}+6EL\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle-\frac{6EL\eta}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+3E\sigma^{2}_{*}.$ (30) $$

The lemma shows that the local gradient variance $\mathcal{L}_{u}(w^{(r,e)}_{u})$ on user $u$ after $E$ local epochs is always bounded by a certain threshold.

Proof:

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\|^{2}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})+\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{*})+\nabla\mathcal{L}_{u}(w^{*})\|^{2}$ $\displaystyle\leq\underbrace{\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})\|^{2}}_{Q_{1}}+\underbrace{\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{*})\|^{2}}_{Q_{2}}$ (31) $\displaystyle+\underbrace{\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{*})\|^{2}}_{Q_{3}}.$ (32) $$

We derive $Q_{1},Q_{2},Q_{3}$ as follows:

$$ $\displaystyle Q_{1}$ $\displaystyle=\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})\|^{2}\leq\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}$ (33) $$

$$ $\displaystyle Q_{2}$ $\displaystyle=\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{*})\|^{2}$ $\displaystyle\leq\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{[}2L\Big{(}\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{*})\Big{)}-\langle\nabla\mathcal{L}_{u}(w^{*}),w^{(r)}_{\textrm{LAG},u}-w^{*}\rangle\Big{]}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{=}}\frac{6L}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{(}\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\nabla\mathcal{L}_{u}(w^{*})\Big{)}$ $\displaystyle=6EL\Big{(}\frac{1}{U}\sum^{U}_{u=1}\nabla\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\frac{1}{U}\sum^{U}_{u=1}\nabla\mathcal{L}_{u}(w^{*})\Big{)}=6EL\Big{(}\nabla\mathcal{L}(w^{(r)}_{\textrm{LAG},g})-\nabla\mathcal{L}(w^{*})\Big{)}$ (34) $$

$$ $\displaystyle Q_{3}$ $\displaystyle=\frac{3}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{*})\|^{2}\leq 3E\sigma^{2}_{*}$ (35) $$

Therefore, we have:

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{u})\|^{2}$ $\displaystyle\leq\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}+6EL\Big{(}\mathcal{L}(w^{(r)}_{\textrm{LAG},g})-\mathcal{L}(w^{*})\Big{)}+3E\sigma^{2}_{*}.$ (36) $$

Base on [H.4](#A8.SS4), we have:

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\|^{2}$ $\displaystyle\leq\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}+6EL\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle-\frac{6EL\eta}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+3E\sigma^{2}_{*}.$ (37) $$

###### Lemma H.10 ( (Jhunjhunwala et al., 2023 ) , Bounding user aggregate gradients on Vanilla FL) .

###### Lemma H.11 (Bounding user aggregate gradients on Layer-wise Personalized FL) .

### H.8 Bounding user drift

###### Lemma H.12 ( (Jhunjhunwala et al., 2023 ) , Bounding user drift on Vanilla FL) .

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r)}_{u}-w^{(r,e)}_{u}\|^{2}\leq 12\eta^{2}E^{2}(E-1)L(F(w^{(r)})-F(w^{*}))+6\eta^{2}E^{2}(E-1)\sigma^{2}_{*}.$ (38) $$

###### Lemma H.13 (Bounding user drift on Layer-wise Personalized FL) .

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r)}_{\textrm{LAG},u}-w^{(r,e)}_{\textrm{LAG},u}\|^{2}\leq 12\eta^{2}E^{2}(E-1)L(F(w^{(r)})-F(w^{*}))+6\eta^{2}E^{2}(E-1)\sigma^{2}_{*}.$ (39) $$

Proof:

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r)}_{\textrm{LAG},u}-w^{(r,e)}_{\textrm{LAG},u}\|^{2}$ $\displaystyle=\eta^{2}\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\sum^{e}_{i=0}\nabla\mathcal{L}_{u}(w^{r,i}_{\textrm{LAG},u})\|^{2}$ $\displaystyle\leq\eta^{2}\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}e\sum^{e}_{i=0}\|\nabla\mathcal{L}_{u}(w^{r,i}_{\textrm{LAG},u})\|^{2}$ $\displaystyle\leq\eta^{2}E(E-1)\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla\mathcal{L}_{u}(w^{r,e}_{\textrm{LAG},u})\|^{2}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\eta^{2}E(E-1)\Big{[}\frac{3L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}+6EL\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle~{}~{}~{}-\frac{6EL\eta}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+3E\sigma^{2}_{*}\Big{]}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{1}{2U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},g}-w^{(r)}_{\textrm{LAG},g}\|^{2}+6E^{2}(E-1)\eta^{2}L\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle~{}~{}~{}-6E^{2}(E-1)\eta^{3}\frac{L}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+3\eta^{2}E^{2}(E-1)\sigma^{2}_{*},$ (40) $$

where (a) holds due to the Lemma [H.11](#A8.Thmtheorem11) and we have (b) due to the assumption with constraints $\eta^{2}\leq\frac{1}{6EL}$. Therefore, we have:

$$ $\displaystyle\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r)}_{\textrm{LAG},u}-w^{(r,e)}_{\textrm{LAG},u}\|^{2}\leq 12E^{2}(E-1)\eta^{2}L\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}-12E^{2}(E-1)\eta^{3}\frac{L}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+6\eta^{2}E^{2}(E-1)\sigma^{2}_{*}$ (41) $$

###### Lemma H.12 ( (Jhunjhunwala et al., 2023 ) , Bounding user drift on Vanilla FL) .

###### Lemma H.13 (Bounding user drift on Layer-wise Personalized FL) .

### H.9 Proof on Theorem 5.3

We define $\bar{h}^{(r)}=\{\bar{h}^{(r)}_{s},\bar{h}^{(r)}_{p}\}$, where $\bar{h}^{(r)}_{s},\bar{h}^{(r)}_{p}$ represents the gradient update rules for shared layers and personalized layers, respectively.
Recall that the update of the global model can be written as $w^{(r+1)}_{\textrm{LAG},g}=w^{(r)}_{\textrm{LAG},g}-\eta\frac{1}{U}\sum^{U}_{u=1}\bar{h}^{(r)}_{u}$. When disentangle into two distinguished models, we have the update for shared layers as $\theta^{(r+1)}_{s,g}=\theta^{(r)}_{s,g}-\eta\frac{1}{U}\sum^{U}_{u=1}\bar{h}^{(r)}_{s,u}$ and personalized layers as $\theta^{(r+1)}_{p,g}=\theta^{(r)}_{p,u}-\eta\bar{h}^{(r)}_{p,u}$. Therefore, we have:

$$ $\displaystyle\mathbb{E}\Big{[}\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}$ $\displaystyle=\mathbb{E}_{u\in U}\Big{[}\|\theta^{(r+1)}_{s,g}-\theta^{*}_{s,g}\|^{2}\Big{]}+\mathbb{E}_{u\in U}\Big{[}\|\theta^{(r+1)}_{p,u}-\theta^{*}_{p,g}\|^{2}\Big{]}$ $\displaystyle=\|\theta^{(r)}_{s,g}-\frac{\eta}{U}\sum^{U}_{u=1}h^{(r)}_{s,u}-\theta^{*}_{s,g}\|^{2}+\mathbb{E}_{u\in U}\Big{[}\|\theta^{(r)}_{p,u}-\eta h^{(r)}_{p,u}-\theta^{*}_{p,g}\|^{2}\Big{]}$ $\displaystyle=\|\theta^{(r)}_{s,g}-\theta^{*}_{s,g}\|^{2}\ +\mathbb{E}_{u\in U}\Big{[}\|\theta^{(r)}_{p,u}-\theta^{*}_{p,g}\|^{2}\Big{]}-2\eta\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{s,u}\rangle$ $\displaystyle~{}-2\eta\mathbb{E}_{u\in U}\Big{[}\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},h^{(r)}_{p,u}\rangle\Big{]}+\|\frac{\eta}{U}\sum^{U}_{u=1}h^{(r)}_{s,u}\|^{2}+\mathbb{E}_{u\in U}\Big{[}\|\eta h^{(r)}_{p,u}\|^{2}\Big{]}$ $\displaystyle\leq\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}-2\eta\underbrace{\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{s,u}\rangle}_{Q_{1}}+\eta^{2}\underbrace{\frac{1}{U}\sum^{U}_{u=1}\|h^{(r)}_{s,u}\|^{2}}_{Q_{3}}$ $\displaystyle~{}-2\eta\underbrace{\frac{1}{U}\sum^{U}_{u=1}\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},h^{(r)}_{p,u}\rangle}_{Q_{2}}+\eta^{2}\underbrace{\frac{1}{U}\sum^{U}_{u=1}\|h^{(r)}_{p,u}\|^{2}}_{Q_{4}}$ (42) $$

Bounding $Q_{1}$: We have:

$$ $\displaystyle Q_{1}$ $\displaystyle=\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\frac{1}{U}\sum^{U}_{u=1}h^{(r)}_{s,u}\rangle=\frac{1}{U}\sum^{U}_{u=1}\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},h^{(r)}_{s,u}\rangle$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle$ (43) $$

We have:

$$ $\displaystyle\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle=\underbrace{\langle\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle}_{Q^{1}_{1}}+\underbrace{\langle\theta^{(r,e)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle}_{Q^{2}_{1}}$ (44) $$

Next, we need to consider two terms $Q_{1}$ and $Q_{2}$. Take $Q^{1}_{1}$ into consideration, due to the $L$-smooth in Assumption [H.1](#A8.Thmtheorem1), we have:

$$ $\displaystyle Q^{1}_{1}$ $\displaystyle=\langle\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle\geq\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{(r,e)}_{s,g})-\frac{L}{2}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}.$ (45) $$

Take $Q^{2}_{1}$ into consideration, from Assumption [H.2](#A8.Thmtheorem2), we have:

$$ $\displaystyle Q^{2}_{1}$ $\displaystyle=\langle\theta^{(r,e)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle\geq\mathcal{L}_{u}(\theta^{(r,e)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})+\frac{\mu}{2}\|\theta^{(r,e)}_{s,g}-\theta^{*}_{s,g}\|^{2}$ $\displaystyle\geq\mathcal{L}_{u}(\theta^{(r,e)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g}).$ (46) $$

Therefore, adding the above inequalities equation [45](#A8.E45) and equation [46](#A8.E46) together, we have:

$$ $\displaystyle\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle\geq\mathcal{L}_{u}(\theta^{(r)}_{s,g})-(\theta^{*}_{s,g})-\frac{L}{2}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}$ (47) $$

Substituting equation [47](#A8.E47) into equation [43](#A8.E43), we have:

$$ $\displaystyle Q_{1}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{[}\langle\theta^{(r)}_{s,g}-\theta^{*}_{s,g},\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle\Big{]}$ $\displaystyle\geq\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})-\frac{L}{2}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}\Big{]}$ $\displaystyle=E\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})\Big{]}-\frac{L}{2U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}$ (48) $$

Bounding $Q_{2}$: We have:

$$ $\displaystyle Q_{2}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},h^{(r)}_{p,u}\rangle=\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle$ (49) $$

We have:

$$ $\displaystyle\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle=\underbrace{\langle\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle}_{Q^{1}_{2}}+\underbrace{\langle\theta^{(r,e)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{\textrm{LAG},u})\rangle}_{Q^{2}_{2}}$ (50) $$

Take $Q^{1}_{2}$ into consideration, due to the $L$-smooth in Assumption [H.1](#A8.Thmtheorem1), we have:

$$ $\displaystyle Q^{1}_{2}=\langle\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\rangle\geq\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{(r,e)}_{p,u})-\frac{L}{2}\|\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u}\|^{2}$ (51) $$

Take $Q^{2}_{2}$ into consideration, from Assumption [H.2](#A8.Thmtheorem2), we have:

$$ $\displaystyle Q^{2}_{2}$ $\displaystyle=\langle\theta^{(r,e)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\rangle\geq\mathcal{L}_{u}(\theta^{(r,e)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})+\frac{\mu}{2}\|\theta^{(r,e)}_{p,u}-\theta^{*}_{p,g}\|^{2}$ $\displaystyle\geq\mathcal{L}_{u}(\theta^{(r,e)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})$ (52) $$

Therefore, adding the above inequalities equation [51](#A8.E51) and equation [52](#A8.E52) together, we have:

$$ $\displaystyle\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\rangle\geq\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})-\frac{L}{2}\|\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u}\|^{2}$ (53) $$

Substituting equation [53](#A8.E53) into equation [49](#A8.E49), we have:

$$ $\displaystyle Q_{2}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\langle\theta^{(r)}_{p,u}-\theta^{*}_{p,g},\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\rangle$ $\displaystyle\geq\frac{1}{U}\sum^{U}_{u=1}\Big{[}\sum^{E-1}_{e=0}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})-\frac{L}{2}\|\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u}\|^{2}\Big{]}$ $\displaystyle=\frac{E}{U}\sum^{U}_{u=1}\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})\Big{]}-\frac{L}{2U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r)}_{p,u}-\theta^{(r,e)}_{p,u}\|^{2}$ (54) $$

Bounding $Q_{3}$: We have:

$$ $\displaystyle Q_{3}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\|h^{(r)}_{s,u}\|^{2}=\frac{1}{U}\sum^{U}_{u=1}\Big{\|}\sum^{E-1}_{e=0}\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\Big{\|}^{2}\stackrel{{\scriptstyle(a)}}{{\leq}}\frac{E}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\|}\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\Big{\|}^{2}$ $\displaystyle\stackrel{{\scriptstyle(b)}}{{\leq}}\frac{3EL^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r,e)}_{s,u}-\theta^{(r)}_{s,u}\|^{2}+6E^{2}L(\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g}))+3E^{2}\sigma^{2}_{s,*}.$ (55) $$

where (a) holds due to lemma [H.5](#A11.EGx9), and (b) holds due to lemma [H.10](#A8.Thmtheorem10).

Bounding $Q_{4}$: We have:

$$ $\displaystyle Q_{4}$ $\displaystyle=\frac{1}{U}\sum^{U}_{u=1}\|h^{(r)}_{p,u}\|^{2}=\frac{1}{U}\sum^{U}_{u=1}\Big{\|}\sum^{E-1}_{e=0}\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\Big{\|}^{2}\leq\frac{E}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\|}\nabla_{\theta^{(r,e)}_{s,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})\Big{\|}^{2}$ $\displaystyle\stackrel{{\scriptstyle}}{{\leq}}\frac{3E}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r,e)}_{u})-\nabla_{\theta^{(r,e)}_{p,u}}\mathcal{L}_{u}(w^{(r)}_{p,u})\|^{2}$ $\displaystyle~{}~{}+\frac{3E}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla_{\theta^{(r)}_{p,u}}\mathcal{L}_{u}(w^{(r)}_{p,u})-\nabla_{\theta^{*}_{p,g}}\mathcal{L}_{u}(w^{*})\|^{2}+\frac{3E}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\nabla_{\theta^{*}_{p,g}}\mathcal{L}_{u}(w^{*})\|^{2}$ $\displaystyle\leq\frac{3EL^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r,e)}_{p,u}-\theta^{(r)}_{p,u}\|^{2}+\frac{6E^{2}L}{U}\sum^{U}_{u=1}\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*})\Big{)}+3E^{2}\sigma^{2}_{p,*}$ (56) $$

Combining equation [48](#A8.E48), equation [54](#A8.E54), equation [55](#A8.E55), equation [56](#A8.E56) together, we have:

$$ $\displaystyle\mathbb{E}\Big{[}\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}-Q_{1}+Q_{3}-Q_{2}+Q_{4}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}-2\eta\underbrace{\Big{\{}E\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})\Big{]}-\frac{L}{2U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}\Big{\}}}_{Q_{1}}$ $\displaystyle+\eta^{2}\underbrace{\Big{\{}\frac{3EL^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r,e)}_{s,u}-\theta^{(r)}_{s,u}\|^{2}+6E^{2}L(\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g}))+3E^{2}\sigma^{2}_{s,*}\Big{\}}}_{Q_{3}}$ $\displaystyle~{}-2\eta\underbrace{\Big{\{}\frac{E}{U}\sum^{U}_{u=1}\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})\Big{]}-\frac{L}{2U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r)}_{p,g}-\theta^{(r,e)}_{p,g}\|^{2}\Big{\}}}_{Q_{2}}$ $\displaystyle+\eta^{2}\underbrace{\Big{\{}\frac{3EL^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|\theta^{(r,e)}_{p,u}-\theta^{(r)}_{p,u}\|^{2}+\frac{6E^{2}L}{U}\sum^{U}_{u=1}\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*})\Big{)}+3E^{2}\sigma^{2}_{p,*}\Big{\}}}_{Q_{4}}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+\frac{\eta^{2}L}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\{}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}+\|\theta^{(r)}_{p,g}-\theta^{(r,e)}_{p,g}\|^{2}\Big{\}}$ $\displaystyle+\frac{3E\eta^{2}L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\{}\|\theta^{(r,e)}_{s,u}-\theta^{(r)}_{s,u}\|^{2}+\|\theta^{(r,e)}_{p,u}-\theta^{(r)}_{p,u}\|^{2}\Big{\}}$ $\displaystyle~{}-2\frac{\eta E}{U}\sum^{U}_{u=1}\Big{\{}\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*}_{p,g})\Big{]}+\Big{[}\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})\Big{]}\Big{\}}$ $\displaystyle+\frac{6E^{2}\eta^{2}L}{U}\sum^{U}_{u=1}\Big{\{}\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{s,g})-\mathcal{L}_{u}(\theta^{*}_{s,g})\Big{)}+\Big{(}\mathcal{L}_{u}(\theta^{(r)}_{p,u})-\mathcal{L}_{u}(\theta^{*})\Big{)}\Big{\}}+3E^{2}\eta^{2}\sigma^{2}_{*}.$ (57) $$

Combine together, we have:

$$ $\displaystyle\mathbb{E}\Big{[}\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+\frac{\eta L}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\{}\|\theta^{(r)}_{s,g}-\theta^{(r,e)}_{s,g}\|^{2}+\|\theta^{(r)}_{p,g}-\theta^{(r,e)}_{p,g}\|^{2}\Big{\}}$ $\displaystyle+\frac{3E\eta^{2}L^{2}}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\Big{\{}\|\theta^{(r,e)}_{s,u}-\theta^{(r)}_{s,u}\|^{2}+\|\theta^{(r,e)}_{p,u}-\theta^{(r)}_{p,u}\|^{2}\Big{\}}$ $\displaystyle~{}-2\frac{\eta E}{U}\sum^{U}_{u=1}\Big{[}\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\mathcal{L}_{u}(w^{*})\Big{]}+\frac{6E^{2}\eta^{2}L}{U}\sum^{U}_{u=1}\Big{(}\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\mathcal{L}_{u}(w^{*})\Big{)}+3E^{2}\eta^{2}\sigma^{2}_{*}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+(3E\eta^{2}L^{2}+\eta L)\frac{1}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},u}-w^{(r)}_{\textrm{LAG},u}\|^{2}$ $\displaystyle~{}-2\frac{\eta E}{U}(1-3E\eta L)\sum^{U}_{u=1}\Big{[}\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\mathcal{L}_{u}(w^{*})\Big{]}+3E^{2}\eta^{2}\sigma^{2}_{*}$ $\displaystyle\stackrel{{\scriptstyle(a)}}{{\leq}}\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+\frac{2\eta L}{U}\sum^{U}_{u=1}\sum^{E-1}_{e=0}\|w^{(r,e)}_{\textrm{LAG},u}-w^{(r)}_{\textrm{LAG},u}\|^{2}$ $\displaystyle~{}-\frac{\eta E}{U}\sum^{U}_{u=1}\Big{[}\mathcal{L}_{u}(w^{(r)}_{\textrm{LAG},u})-\mathcal{L}_{u}(w^{*})\Big{]}+3E^{2}\eta^{2}\sigma^{2}_{*}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+2\eta L\Big{\{}12E^{2}(E-1)\eta^{2}L\Big{(}\mathcal{L}(w^{(r)}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle-12E^{2}(E-1)\eta^{3}\frac{L}{U^{2}}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+6\eta^{2}E^{2}(E-1)\sigma^{2}_{*}\Big{\}}$ $\displaystyle~{}-\frac{\eta E}{U}\Big{\{}\sum^{U}_{u=1}\Big{(}\mathcal{L}(w^{(r)}_{g}-\mathcal{L}(w^{*})\Big{)}-\frac{\eta}{U}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}+3E^{2}\eta^{2}\sigma^{2}_{*}.$ (58) $$

Rearrange terms, we have:

$$ $\displaystyle\mathbb{E}\Big{[}\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}$ $\displaystyle=\mathbb{E}\Big{[}\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}\Big{]}+\Big{(}24E^{2}(E-1)\eta^{3}L^{2}-\eta E\Big{)}\frac{1}{U}\sum^{U}_{u=1}\Big{(}\mathcal{L}(w^{(r)}_{\textrm{LAG},u}-\mathcal{L}(w^{*})\Big{)}$ $\displaystyle-\frac{1}{U^{2}}\Big{(}24E^{2}(E-1)\eta^{4}L^{2}-\eta^{2}E\Big{)}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}$ $\displaystyle+12\eta^{3}E^{2}(E-1)L\sigma^{2}_{*}+3E^{2}\eta^{2}\sigma^{2}_{*}.$ (59) $$

Here, (a) due to the assumption that $\eta<\frac{1}{6EL}$. For local epoch $E>1$, we have $E\geq\sqrt{3}$, which means that $\frac{1}{2\sqrt{6}E^{2}L}<\eta<\frac{1}{6EL}$. Thus, we have $\frac{1}{U^{2}}\Big{(}24E^{2}(E-1)\eta^{4}L^{2}-\eta^{2}E\Big{)}>0,~{}\forall\eta>\frac{1}{2\sqrt{6}E^{2}L}$ (the proof is provided in Appendix [I](#A9)). Therefore, we can rewrite the term as:

$$ $\displaystyle\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}$ $\displaystyle\leq\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}-\frac{\eta E}{3}\Big{(}\mathcal{L}(w^{(r)}_{g}-\mathcal{L}(w^{*})\Big{)}+\Big{[}3E\eta^{2}+12\eta^{3}E^{2}(E-1)L\Big{]}\sigma^{2}_{*}$ $\displaystyle~{}~{}~{}~{}~{}-\frac{24}{U^{2}}A\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{p,u}\|(\|h^{(r)}_{p,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{p,v}\|),$ (60) $$

where $A=\Big{(}24E^{2}(E-1)\eta^{4}L^{2}-\eta^{2}E\Big{)}>0$, thus, the algorithm always give a consistent improvement gap to the global learning performance.

Due to Lemma [H.13](#A11.EGx36) and Lemma [H.11](#A8.Thmtheorem11). Combine together, we have:

$$ $\displaystyle\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}$ $\displaystyle\leq\|w^{(r)}_{\textrm{LAG},g}-w^{*}\|^{2}-\frac{\eta E}{3}\Big{(}\mathcal{L}(w^{(r)}_{g}-\mathcal{L}(w^{*})\Big{)}+\Big{[}3E\eta^{2}+12\eta^{3}E^{2}(E-1)L\Big{]}\sigma^{2}_{*}$ $\displaystyle~{}~{}~{}~{}~{}-\mathcal{O}\Big{(}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Big{)}.$ (61) $$

Moreover, we have $\eta>\frac{1}{2\sqrt{6}E^{2}L}>\frac{1}{2\sqrt{6}EL}$ (as $E>1$). Therefore, we have
Rearranging terms and averaging over all rounds, we have:

$$ $\displaystyle\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}$ $\displaystyle\leq\|w^{(0)}_{\textrm{LAG},g}-w^{*}\|^{2}-\sum^{R-1}_{r=0}\frac{\eta E}{3}\Big{(}\mathcal{L}(w^{(r)}_{g}-\mathcal{L}(w^{*})\Big{)}+\sum^{R-1}_{r=0}\Big{[}3E\eta^{2}+12\eta^{3}E^{2}(E-1)L\Big{]}\sigma^{2}_{*}$ $\displaystyle~{}~{}~{}~{}~{}-\mathcal{O}\Big{(}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Big{)}.$ (62) $$

$$ $\displaystyle\mathcal{L}(w^{(r)}_{g})-\mathcal{L}(w^{*})$ $\displaystyle\leq\frac{3}{R\eta E}\|w^{(0)}_{\textrm{LAG},g}-w^{*}\|^{2}-\frac{3}{R\eta E}\|w^{(r+1)}_{\textrm{LAG},g}-w^{*}\|^{2}+\frac{3}{\eta E}\Big{[}3E\eta^{2}+12\eta^{3}E^{2}(E-1)L\Big{]}\sigma^{2}_{*}$ $\displaystyle~{}~{}~{}~{}~{}-\mathcal{O}\Big{(}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Big{)}$ $\displaystyle\leq\frac{3\|w^{(0)}_{\textrm{LAG},g}-w^{*}\|^{2}}{R\eta E}+9\eta\sigma^{2}_{*}+36\eta^{2}E(E-1)L\sigma^{2}_{*}$ $\displaystyle~{}~{}~{}~{}~{}-\mathcal{O}\Big{(}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Big{)}.$ (63) $$

As we have aggregated the layer-wise personalized model and global aggregate model back into $w^{(r)}_{\textrm{LAG},g}$, we can have:

$$ $\displaystyle\mathcal{L}(w^{(R)}_{g})-\mathcal{L}(w^{*})$ $\displaystyle\leq\mathcal{O}\Big{(}\frac{\|w^{(0)}_{g}-w^{*}\|^{2}}{R\eta E}\Big{)}+\mathcal{O}\Big{(}\eta\sigma^{2}_{*}\Big{)}+\mathcal{O}\Big{(}\eta^{2}E(E-1)L\sigma^{2}_{*}\Big{)}$ $\displaystyle~{}~{}~{}~{}~{}-\mathcal{O}\Big{(}\sum^{U}_{u=1}\sum^{U}_{v=1}\sum_{l\in\mathbb{L}_{p}}\|h^{(r)}_{l,u}\|\Big{(}\|h^{(r)}_{l,u}\|-\cos\Phi^{(l)}_{u,v}\|h^{(r)}_{l,v}\|\Big{)}\Big{)}.$ (64) $$

This completes the proof.

## Appendix I Proof on consistent improvement gap

Consider a second-degree polynomial function given by $Y=AX^{2}-BX$, where $A>0$. Consequently, we ensure $Y>0$ under the condition $X\in(-\infty,0]\cup[B/A;+\infty)$.
Next, we consider the function $Y=E^{2}(E-1)\eta^{4}L^{2}-\eta^{2}E$, and take $\eta^{2}=X$. This yields:

$$ $\displaystyle Y=24E^{2}(E-1)L^{2}X^{2}-EX.$ (65) $$

Given that $E^{2}(E-1)L^{2}\geq 0$ for all $E,L$. To ensure $Y>0$, we have to satisfy the following condition:

$$ $\displaystyle X\geq\frac{E}{24E^{2}(E-1)L^{2}}=\frac{1}{24E(E-1)L^{2}}.$ (66) $$

Substituting $X$ back to $\eta^{2}$, then we have

$$ $\displaystyle\eta^{2}\geq\frac{1}{24E(E-1)L^{2}}.$ (67) $$

Accordingly, we

$$ $\displaystyle\eta\geq\frac{1}{2\sqrt{6}EL}.$ (68) $$

Given that $\frac{1}{2\sqrt{6}E^{2}L}<\eta<\frac{1}{6EL}$, then, $\eta\geq\frac{1}{2\sqrt{6}E^{2}L}>\frac{1}{2\sqrt{6}EL}$, which holds that $\eta$ always satisfy the condition to have $Y>0$.

## Appendix J Analysis on Computational Complexity of On-server Gradient Analysis

We consider the time complexity of the feed-forward process of an AI model. Consider a model with $L$ hidden layers, denoted as $\{w_{1},w_{2},\ldots,w_{d_{l}}\}$ representing the weights of each layer. For simple approximation, we assume that all layers have same number of parameters, i.e., $d_{1}=d_{2}=\ldots=d_{L}=D$.

We consider the feed-forwarding task of an AI model with a batch size of $B$.

To go through $L$ layers, each layers have $D$ parameters, we have to do $y=W_{LD}\times x$. Therefore, we have the feed-forward process for a mini-batch with batch size $B$ has time complexity $\mathcal{O}(L\times D\times B\times S)$, where $S$ is the data size.

Second, according to the personalized layers selection, the time complexity is around $\mathcal{O}(U\times(U-1)/2\times D\times L)$. If we consider a large number of users $U=1000$, and due to the large DNN, we use ImageNet as evaluations, where the average image dimensionality is $3\times 469\times 387=544509$. Therefore, we have the two complexity two process as follows:

- •
One iteration of feed-forward in a DNN with only one sample: $\mathcal{O}(L\times D\times 1\times 544509)$.
- •
Personalized layers selection: $\mathcal{O}(500\times 999\times D\times L)$.

It is obvious that the personalized layers selection is approximately same with the time consumption of $1$ iteration of feed-forward in DNN.

Finally, we survey and found that the inference time in Large DNN Model is trivial ($\ll 1$s/sample). The detailed results can be found in (Liu et al., , Table 1).

## Appendix K Detailed Algorithms

Figure: Algorithm 1 Layer-wise Gradient Analysis supported pFL

Figure: Algorithm 2 Layer-wise Gradient Divergence Analysis

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 1
questions: Q1: What is the meaning of vertical axis in Figure 2(a)?

Q2: Does the conclusion of Figure 2(b) contradict the textual interpretation in main text?

Q3: Is it necessary to re-compute the conflict layer in each computation round?

Q4: What is the purpose of the experiments in Section 7.2.5?
soundness: 2
strengths: S1: This paper focuses on an important issue in the federated learning, i.e., build layer-wise personalized federated optimization framework to solve the data heterogeneity challenge.

S2: The idea is meaningful that the system can disentangle the personalized parameters by monitoring the gradient conflict.

S3: Experimental and theoretical analysis demonstrate the effectiveness of the proposed method.
summary: This paper presents a layer disentanglement mechanism to achieve personalized federated learning. Specifically, the paper distinguishes between personal layers and global layers by examining the conflicts in local model gradients. Consequently, the model can adaptively separate the model parameters into personal and global components. Experimental evaluations demonstrate performance improvements compared to baseline models.
weaknesses: W1: Compared with common layer-wise personalized federated learning framework, the proposed method has a higher communication overhead. Common layer-wise methods upload partial layers to the server for global aggregation and keep other parameters as personal component locally. However, the proposed method demands clients to upload full model parameters to the server for distinguishing gradient conflict.

W2: The proposed method incurs high computational costs due to the need to calculate the parameter correlations between any two users.

W3: The presentation is poor and there are many confusing clarifications, please refer to section "Questions" for details.

-----------------------------------------------------------------------After Rebuttal------------------------------------------------------------------------

I have acknowledged the authors' response, but the main concerns have not been adequately addressed or explained. Therefore, I will maintain my original rating and recommend that the authors carefully revise the paper in accordance with all reviewers' suggestions.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see the weakness.
soundness: 2
strengths: **S1**. The authors provide a thorough theoretical analysis to show the strength of FedLAG in convergence and model performance.

**S2**.  Experimental results show that FedLAG outperforms baselines in terms of accuracy and convergence time, especially under high data heterogeneity and low user participation.

**S3**. The proposed method is simple bug effective and can be easily integrated into existing methods.
summary: This paper focuses on enhancing model performance in personalized federated learning (FL).  Recent pFL's methods usually face the problem of selecting which layers of the neural network to be personalized. To this end, the authors propose FedLAG to adaptively personalize layers whose gradients exhibit large conflicts across different clients.  Comprehensive theoretical and empirical results confirm the advantage of the proposed method in terms of accuracy.
weaknesses: **W1** . It is quite weird that the modern neural network ResNet9 can only achieve at most 97% testing accuracy on MNIST across all the methods and some results are even lower than 90% in Table 2 (e.g., PerAvg and FedPAC). A simple LeNet-5 model can already achieve 90% testing accuracy on MNIST by FedAvg even in the case of extremely non-i.i.d. partition (e.g., one label per client) [a]. In addition, the performance of FedAvg on MNIST-Dir(0.5) is lower than its performance on MNIST-Dir(0.1), which is also strange since a higher degree of data heterogeneity leads to a better performance for a non-personalized method. These factors make the experimental results less convincing. Please explain the potential reasons for these anomalous results.

**W2**. There may be some typo errors. For example, in the first line of Sec. 7.1.1, the α in the sentence "We assess at two heterogeneity levels, i.e., α = 0.1, 1."  is inconsistent with α 's values in Table 2.  In the Sec. 6.2, the method name GBFL (Zhang et al., 2023c) should be GPFL. Please correct these issues and check over the manuscript.

**W3** Some works related to leveraging gradient conflicts in FL [b][c][d] were missing. More discuss on how the approach relates to or differs from the gradient conflict handling methods is needed. One highly related baseline [e] similarly tackling the layer personalization problem in FL was also ignored. Please compare with this method in the experiments.

[a] Zhao Y, Li M, Lai L, et al. Federated learning with non-iid data[J]. arXiv preprint arXiv:1806.00582, 2018.

[b] Pan Z, Li C, Yu F, et al. FedLF: Layer-Wise Fair Federated Learning[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2024, 38(13): 14527-14535.

[c] Wang Z, Fan X, Qi J, et al. Federated learning with fair averaging[J]. arXiv preprint arXiv:2104.14937, 2021.

[d] Hu Z, Shaloudegi K, Zhang G, et al. Federated learning meets multi-objective optimization[J]. IEEE Transactions on Network Science and Engineering, 2022, 9(4): 2039-2051.

[e] Ma X, Zhang J, Guo S, et al. Layer-wised model aggregation for personalized federated learning[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 10092-10101.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Section 3:

1. In Question 3.2, the authors state that “the density of conflicting gradients among clients does not progressively increase from the input layer to the output layer.” However, from Fig. 2(b), it appears that deeper layers have higher conflict scores. Could the authors clarify this discrepancy? I suggest that the authors provide a quantitative analysis of the trend in conflict scores across layers, rather than relying on a visual inspection of the figure.
2. In Section 4.1, why is FedLAG considered communication-efficient?

Section 7:

1. In 7.1.1, the authors mention “with improvements ranging from an average of 5–7% to 15%.” How was this calculated? From Table 1, it seems that FedLAG does not exceed the performance of the best method by more than 1% in most cases.
2. In 7.2.2, what would happen if  K  layers were randomly selected instead?
3. Which layers does FedLAG typically select for personalization?
soundness: 3
strengths: 1. The paper is well-written, and the main claim is easy to understand.
2. Using gradient divergence as a criterion for selecting personalized layers seems reasonable.
3. Extensive experiments demonstrate the effectiveness of FedLAG.
summary: This paper proposes a novel layer-wise personalized federated learning (PFL) method called FedLAG, which uses the divergence of client gradients as a criterion for selecting personalized layers. Experimental results demonstrate the effectiveness of this approach.
weaknesses: 1. Although gradient conflict is easy to understand, how gradient alignment specifically influences collaboration effectiveness in federated learning remains unclear. For instance, how large of an angle actually affects performance? Is an angle greater than 90 degrees necessarily detrimental? The paper primarily tunes the hyperparameter  \xi  to identify a threshold, but exploring the relationship between angle and performance further, along with providing a guideline for adjusting  \xi , would greatly enhance the quality of the work.
2. From the results in Tables 1 and 7, the performance improvement of FedLAG over the SOTA methods appears limited, with gains of less than 1%. Moreover, existing methods can already decouple the learning of personalized and generic knowledge in each parameter [1], I suggest that the authors compare FedLAG to this method.
3. As shown in Table 3, FedLAG appears highly sensitive to the hyperparameter  \xi . Combined with the results in Table 1, FedLAG requires very careful tuning of  \xi  to achieve performance superior to SOTA, which limits its usability.

[1] Decoupling General and Personalized Knowledge in Federated Learning via Additive and Low-Rank Decomposition, ACM MM 2024

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: See the weakness.
soundness: 3
strengths: 1)Innovation: This paper addresses a key challenge in personalized federated learning, specifically, how to implement layer-wise model partitioning under data heterogeneity, and introduces an innovative method based on gradient conflict analysis.
2)Theoretical Rigor: The paper provides a theoretical analysis demonstrating the convergence advantages of FedLAG and includes mathematical proofs to validate the effectiveness of its gradient conflict-based layer separation.
3)Comprehensive Experimental Design: The experiments encompass various datasets of data heterogeneity, with results affirming FedLAG's strengths in both accuracy and convergence speed.
summary: This paper presents a novel layer-wise aggregation method for personalized federated learning, named FedLAG. By analyzing gradient conflicts, FedLAG dynamically determines which layers should undergo global aggregation (GAL) and which should be personalized (PL) to mitigate the negative impact of data heterogeneity on federated learning. FedLAG adapts to the personalized needs of the data, reducing the reliance on prior knowledge that traditional methods require, thereby enhancing model convergence and adaptability. Experimental results indicate that FedLAG outperforms existing baseline methods across multiple datasets.
weaknesses: 1)Gradient Conflict Concept Insufficiently Explained: Although the paper defines gradient conflict, it lacks examples or illustrations to show how gradient conflicts are detected and quantified in practice, especially in the context of layer separation. Additionally, the choice of the gradient conflict threshold , as well as its impact on model performance, is not clearly explained. The paper merely describes conflict-based layer separation without providing an intuitive example. Supplementary examples or visual aids are recommended to enhance reader comprehension.
2)Inadequate Ablation Studies: While the paper compares the performance of fixed layers, it does not explain why certain layers are more suitable for personalization than others. This section lacks theoretical support or empirical evidence, rendering the rationale for layer selection ambiguous. Further experiments or analyses could substantiate the impact of layer choice on model performance, strengthening the logic behind model design.
3)Inaccurate Terminology: The paper states that “existing methods require extensive fine-tuning to determine which layers should be used for global aggregation and personalization” characterizing a limitation of existing methods. However, the term “fine-tuning” is inappropriate in this context. References [1] and [2] emphasize the role of prior knowledge in layer separation, and clarifying this terminology would help prevent potential misunderstandings.
4)Ambiguous Notation: Several notations in Section 2 are unclear, such as the definition of global communication rounds. This paper uses “each round” to describe this symbol, which may confuse readers regarding whether it refers to global communication rounds or local iteration rounds. Furthermore, in the “Layer Disentanglement” section, the complex and nested notation complicates the reader's understanding of the hierarchical structure. A clearer description of the relationships among symbols is recommended, along with a well-organized explanation of global and personalized layers. Visual aids, such as diagrams, could also serve as useful tools to help readers intuitively grasp the layered model structure.
5)Insufficient Discussion on Limitations and Future Directions: The paper does not discuss the limitations of the proposed approach or provide insights into future research directions in the conclusion. It is recommended to analyze the limitations, such as the uncertain performance on more complex types of heterogeneous data (e.g., feature distribution skewed datasets) or potential applications in NLP and time-series analysis. Future research directions could also include exploring the method on multi-type heterogeneous datasets to enhance generalizability.
6)Lack of In-depth Analysis of Experimental Results: Table 1 and Figures 6 and 7 demonstrate FedLAG's performance across various datasets, but the analysis of performance differences is not sufficiently in-depth. For instance, FedLAG's substantial improvement over other methods on CIFAR-10 and CIFAR-100 is not explained in terms of potential causes or contributing factors. A detailed analysis of these phenomena would increase the paper's credibility and persuasive power.
7)Insufficient Validation of Gradient Conflict Distribution: Although a 2D toy dataset experiment was conducted in Section 3, there is insufficient evidence to demonstrate the consistency of gradient conflict distribution across different datasets.
[1] Shen Y, Zhou Y, Yu L. Cd2-pfed: Cyclic distillation-guided channel decoupling for model personalization in federated learning[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 10041-10050.
[2] Wu X, Liu X, Niu J, et al. Bold but cautious: Unlocking the potential of personalized federated learning through cautiously aggressive collaboration[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 19375-19384.