We proposed a layer-wise model aggregation scheme that adaptively adjusts the model aggregation interval at run-time.
Breaking the convention of aggregating the whole model parameters at once, this novel model aggregation scheme introduces a flexible communication strategy for scalable Federated Learning.
Furthermore, we provide a solid convergence guarantee of FedLAMA under the assumptions on the non-convex objective functions and the non-IID data distribution.
Our empirical study also demonstrates the efficacy of FedLAMA for scalable and accurate Federated Learning.
Harmonizing FedLAMA with other advanced optimizers, gradient compression, and low-rank approximation methods is a promising future work.

Appendix
Convergence Analysis
Herein, we provide the proofs of the lemmas and theorem shown in Section .
Preliminaries
FedLAMA periodically chooses a few layers that will be less frequently synchronized.
We call these layers Least Critical Layers (LCL) for short.
Notations –
All vectors in this paper are column vectors.
\(\mathbf {x} \in \mathbb {R}^d\)  denotes the parameters of one local model and \(m\)  is the number of workers.
The stochastic gradient computed from a single training data point \(\xi \)  is denoted by \(g(\mathbf {x},\xi )\) .
For convenience, we use \(g(\mathbf {x})\)  instead.
The full batch gradient is denoted by \(\nabla F(\mathbf {x})\) .
We use \(\Vert \cdot \Vert \)  and \(\Vert  \cdot \Vert _{op}\)  to denote \(l2\)  norm and matrix operator norm, respectively.
Objective Function –
In this paper, we consider federated optimization problems as follows.
\(\underset{\mathbf {x} \in \mathbb {R}^d}{\min }\left[F(\mathbf {x}) := \sum _{i=1}^{m} p_i F_i(\mathbf {x}) \right] ,\) 

where \(p_i = n_i / n\)  is the ratio of local data to the total dataset, and \(F_i(\mathbf {x}) = \frac{1}{n_i} \sum _{\xi \in \mathcal {D}} f_i(\mathbf {x}, \xi )\) 
is the local objective function of client \(i\) .
\(n\)  is the global dataset size and \(n_i\)  is the local dataset size.
Note that, by definition, \(\sum _{i=1}^{m} p_i = 1\) .
Averaging Matrix –
We define a time-varying averaging matrix \(\mathbf {W_k} \in \mathbb {R}^{dm \times dm}\) .
First, \(\mathbf {W}\)  is defined as follows.
\(\mathbf {W_k} ={\left\lbrace \begin{array}{ll}\mathbf {P}, \hfill \hspace{14.22636pt} \textrm { if } k \textrm { mod } \tau _{min} \textrm { is } 0 \\\mathbf {I}, \hfill \textrm { otherwise }\end{array}\right.}\) 
\(\mathbf {P}\)  is also a time-varying averaging matrix and \(\mathbf {I}\)  is an identity matrix.
\(\mathbf {P}\)  consists of \(m \times m\)  blocks each of which is a diagonal matrix of size \(d \times d\) .
\(\mathbf {P} ={\left\lbrace \begin{array}{ll}\mathbf {P}^1, \textrm { for } m \textrm { diagonal blocks} \\\mathbf {P}^0, \textrm { for all the other blocks}\end{array}\right.}\) 
where \(\mathbf {P}^1\)  is a \(d \times d\)  diagonal matrix that has 1 for the diagonal elements that correspond to the LCL parameters and \(1 / m\)  for all the other diagonal elements.
Likewise, \(\mathbf {P}^0\)  is a \(d \times d\)  diagonal matrix that has 0 for the diagonal elements that correspond to the LCL parameters and \(1 / m\)  for all the other diagonal elements.
Note that, regardless of which layers are chosen as LCL, \(\mathbf {P}\)  is a symmetric matrix of size \(dm \times dm\) .
The averaging matrix \(\mathbf {P}\)  has the following properties:

\(\mathbf {P}\mathbf {1}_{dm} = \mathbf {1}_{dm}\) .

The product of any different two \(\mathbf {P}\)  matrices consists only of diagonal block matrices because all the blocks in any \(\mathbf {P}\)  are diagonal.

The product of any different two \(\mathbf {P}\)  matrices is symmetric since all possible \(\mathbf {P}\)  matrices are symmetric and the diagonal blocks within each \(\mathbf {P}\)  are all the same.

The product of any \(\mathbf {P}_i\)  and \(\mathbf {P}_j\)  contains all \(\frac{1}{m}\)  elements in both \(\mathbf {P}_i\)  and \(\mathbf {P}_j\) .

We also define a full-averaging matrix \(\mathbf {J}\)  as follows.
\(\begin{split}& \mathbf {J} = \frac{1}{m}\mathbf {1}_{m}\mathbf {1}_{m}^{\top } \otimes \mathbf {I}_{d \times d},\\\end{split}\) 
where \(\otimes \)  indicates Kronecker product, \(\mathbf {1}_{m}\in \mathbb {R}^{m}\)  is a vector of 1s, and \(\mathbf {I}_{d \times d}\in \mathbb {R}^{d\times d}\)  is an identity matrix.
Note that, due to the first property of \(\mathbf {P}\) , the product of \(\mathbf {J}\)  and any \(\mathbf {P}\)  is always \(\mathbf {J}\) .
Vectorization –
We define a vectorized form of \(m\)  local model parameters \(\mathbf {x}_k \in \mathbb {R}^{dm}\)  and its gradients \(\mathbf {g}_k \in \mathbb {R}^{dm}\)  as follows
\(\begin{split}& \mathbf {x}_k = vec\left\lbrace  \mathbf {x}_k^1, \mathbf {x}_k^2, \cdots , \mathbf {x}_k^m \right\rbrace  \\& \mathbf {g}_k = vec\left\lbrace  g_1(\mathbf {x}_k^1), g_2(\mathbf {x}_k^2), \cdots , g_m(\mathbf {x}_k^m) \right\rbrace  \\& \mathbf {f}_k = vec\left\lbrace  \nabla F_1(\mathbf {x}_k^1), \nabla F_2(\mathbf {x}_k^2), \cdots , \nabla F_m(\mathbf {x}_k^m) \right\rbrace .\end{split}\) 
Using the above notations, the weighted average of the distance between the averaged model \(\mathbf {u}_k\)  and the local models \(\mathbf {x}_{k}^{i}\)  can be written as follows.
\(\sum _{i=1}^{m} p_i \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 = \left\Vert  \mathbf {J} \mathbf {x}_{k} - \mathbf {x}_{k} \right\Vert ^2 = \left\Vert  (\mathbf {J} - \mathbf {I})\mathbf {x}_{k} \right\Vert ^2\) 
We also define the following additional vectorized forms of the weighted model parameters and gradients for convenience.
\(\begin{split}\mathbf {\hat{x}}_k &= \text{vec} \left\lbrace \frac{\mathbf {x}_{k}^{1}}{\sqrt{p_1}}, \frac{\mathbf {x}_{k}^{2}}{\sqrt{p_2}}, \cdots , \frac{\mathbf {x}_{k}^{m}}{\sqrt{p_m}} \right\rbrace  \\\mathbf {\hat{g}}_k &= \text{vec} \left\lbrace  \frac{g_{1}(\mathbf {x}_{k}^{1})}{\sqrt{p_1}}, \frac{g_{2}(\mathbf {x}_{k}^{2})}{\sqrt{p_2}}, \cdots , \frac{g_{m}(\mathbf {x}_{k}^{m})}{\sqrt{p_m}} \right\rbrace  \\\mathbf {\hat{f}}_k &= \text{vec} \left\lbrace  \frac{\nabla F_{1}(\mathbf {x}_{k}^{1})}{\sqrt{p_1}}, \frac{\nabla F_{2}(\mathbf {x}_{k}^{2})}{\sqrt{p_2}}, \cdots , \frac{\nabla F_{m}(\mathbf {x}_{k}^{m})}{\sqrt{p_m}} \right\rbrace \\\end{split}\) 
Assumptions –
We analyze the convergence rate of FedLAMA under the following assumptions.

(Smoothness). Each local objective function is L-smooth, that is, \(\Vert  \nabla F_i (\mathbf {x}) - \nabla F_i (\mathbf {y}) \Vert  \le L \Vert  \mathbf {x} - \mathbf {y} \Vert , \forall i \in \lbrace  1, \cdots , m \rbrace \) .

(Unbiased Gradient). The stochastic gradient at each client is an unbiased estimator of the local full-batch gradient: \(\mathop {\mathbb {E}}_{\xi }[ g_i(\mathbf {x}, \xi )] = \nabla F_i(\mathbf {x})\) .

(Bounded Variance). The stochastic gradient at each client has bounded variance: \(\mathop {\mathbb {E}}_\xi [\Vert  g_i(\mathbf {x}, \xi ) - \nabla F_i (\mathbf {x}) \Vert ^2 \le \sigma ^2], \forall i \in \lbrace  1, \cdots , m \rbrace , \sigma ^2 \ge 0\) .

(Bounded Dissimilarity). For any sets of weights \(\lbrace p_i \ge 0\rbrace _{i=1}^{m}, \sum _{i=1}^{m} p_i = 1\) , there exist constants \(\beta ^2 \ge 1\)  and \(\kappa ^2 \ge 0\)  such that \(\sum _{i=1}^{m} p_i \Vert  \nabla F_i(\mathbf {x}) \Vert ^2 \le \beta ^2 \Vert  \sum _{i=1}^{m} p_i \nabla F_i (\mathbf {x}) \Vert ^2 + \kappa ^2\) . If local objective functions are identical to each other, \(\beta ^2 = 1\)  and \(\kappa ^2 = 0\) .


Proofs
Theorem 4.1. Suppose all \(m\)  local models are initialized to the same point \(\mathbf {u}_1\) . Under Assumption \(1 \sim 4\) , if FedLAMA runs for \(K\)  iterations and the learning rate satisfies \(\eta \le \min \left\lbrace \frac{1}{2L}, \frac{1}{L \sqrt{2\tau _{max} (\tau _{max} - 1) (2\beta ^2 + 1)}} \right\rbrace \) , the average-squared gradient norm of \(\mathbf {u}_k\)  is bounded as follows
\(\mathop {\mathbb {E}}\left[\frac{1}{K}\sum _{i=1}^{K} \Vert  \nabla F(\mathbf {u}_k) \Vert ^2\right] & \le \frac{4}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 4\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + 3 \eta ^2 (\tau _{max} - 1) L^2 \sigma ^2 + 6 \eta ^2 \tau _{max} (\tau _{max} - 1) L^2 \kappa ^2,\) 

where \(\mathbf {u}_{*}\)  indicates a local minimum.
Based on Lemma REF  and REF , we have
\(\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] & \le \frac{2}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 2\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\&\quad + L^2 \left( \frac{\eta ^2 (\tau _{max} - 1) \sigma _j^2}{1 - A} + \frac{A \beta ^2}{K L^2 (1 - A)} \sum _{k=1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 \right] + \frac{A \kappa ^2}{L^2 (1 - A)} \right). \nonumber \) 
After re-writing the left-hand side and a minor rearrangement, we have
\(\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] & \le \frac{2}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 2\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + \frac{1}{K} \sum _{k = 1}^{K} \frac{A \beta ^2}{1 - A} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 \right] \nonumber \\& \quad \quad + L^2 \left( \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A \kappa ^2}{L^2 (1 - A)} \right). \nonumber \) 
By moving the third term on the right-hand side to the left-hand side, we have
\(\frac{1}{K} \sum _{k=1}^{K} \left(1 - \frac{A \beta ^2}{1 - A} \right) \mathop {\mathbb {E}}\left[ \left\Vert  \nabla _j F(\mathbf {u}_k) \right\Vert ^2 \right] & \le \frac{2}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 2\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + L^2 \left( \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A \kappa ^2}{L^2 (1 - A)} \right). \) 
If \(A \le \frac{1}{2 \beta ^2 + 1}\) , then \(\frac{A \beta ^2}{1 - A} \le \frac{1}{2}\) .
Therefore, (REF ) can be simplified as follows.
\(\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] & \le \frac{4}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 4\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + 2 L^2 \left( \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} \right) + 2 \frac{A \kappa ^2}{1 - A}. \nonumber \) 
The learning rate condition \(A \le \frac{1}{2 \beta ^2 + 1}\)  also ensures that \(\frac{1}{1 - A} \le 1 + \frac{1}{2 \beta ^2}\) .
Based on Assumption 4, \(\frac{1}{2 \beta ^2} \le \frac{2}{3}\) , and thus \(\frac{1}{1 - A} \le \frac{2}{3}\) .
Therefore, we have
\(\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] & \le \frac{4}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 4\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + 3 L^2 \eta ^2 (\tau _{max} - 1) \sigma ^2 + 6 \eta ^2 \tau _{max} (\tau _{max} - 1) L^2 \kappa ^2 \nonumber \\& = \frac{4}{\eta K}\left( \mathop {\mathbb {E}}\left[ F(\mathbf {u}_{1}) - F(\mathbf {u}_{*}) \right] \right) + 4\eta \sum _{i=1}^{m} p_i^2 L \sigma ^2 \nonumber \\& \quad \quad + 3 \eta ^2 (\tau _{max} - 1) L^2 \sigma ^2 + 6 \eta ^2 \tau _{max} (\tau _{max} - 1) L^2 \kappa ^2. \nonumber \) 

We complete the proof.
Here, we provide proofs of the two key Lemmas as follows.
Lemma 4.1. (Framework) Under Assumption \(1 \sim 3\) , if the learning rate satisfies \(\eta \le \frac{1}{2L}\) , FedLAMA ensures
\(\begin{split}\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] &\le \frac{2}{\eta K} \mathop {\mathbb {E}}\left[F(\mathbf {u}_1) - F(\mathbf {u}_{*}) \right] + 2 \eta L \sigma ^2 \sum _{i=1}^{m} (p_i)^2 \\& \quad + \frac{L^2}{K} \sum _{k=1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 \right].\\\end{split}\) 
Based on Assumption 1, we have
\(\mathop {\mathbb {E}}\left[ F(\mathbf {u}_{k+1}) - F(\mathbf {u}_{k}) \right] &\le - \eta \underset{T_1}{\underbrace{ \mathop {\mathbb {E}} \left[ \langle \nabla F(\mathbf {u}_{k}), \sum _{i=1}^{m} p_i g_i(\mathbf {x}_{k}^{i}) \rangle \right] }} + \frac{\eta ^2 L}{2} \underset{T_2}{\underbrace{ \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i g_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] }} \) 
First, \(T_1\)  can be rewritten as follows.
\(T_1 & = \mathop {\mathbb {E}} \left[ \langle \nabla F(\mathbf {u}_{k}), \sum _{i=1}^{m} p_i \left( g_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {x}_{k}^{i}) \right) \rangle \right] + \mathop {\mathbb {E}} \left[ \langle \nabla F(\mathbf {u}_{k}), \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \rangle \right] \nonumber \\& = \mathop {\mathbb {E}} \left[ \langle \nabla F(\mathbf {u}_{k}), \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \rangle \right] \nonumber \\& = \frac{1}{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 + \frac{1}{2} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] - \frac{1}{2} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) - \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right], \) 

where the last equality holds based on a basic equality: \(2\mathbf {a}^{\top }\mathbf {b} = \Vert  \mathbf {a} \Vert ^2 + \Vert  \mathbf {b} \Vert ^2 - \Vert  \mathbf {a} - \mathbf {b} \Vert ^2\)  .
Then, \(T_2\)  can be bounded as follows.
\(T_2 & = \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \left( g_i(\mathbf {x}_{k}^{i}) - \mathop {\mathbb {E}} \left[ g_i(\mathbf {x}_{k}^{i}) \right] \right) + \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ g_i(\mathbf {x}_{k}^{i}) \right] \right\Vert ^2 \right] \nonumber \\& = \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \left( g_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {x}_{k}^{i}) \right) + \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \le 2 \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \left( g_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {x}_{k}^{i}) \right) \right\Vert ^2 \right] + 2 \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& = 2 \sum _{i=1}^{m} p_i^2 \mathop {\mathbb {E}}\left[ \left\Vert  g_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] + 2 \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \le 2 \sigma ^2 \sum _{i=1}^{m} p_i^2 + 2 \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right], \) 

where the last equality holds because \(g_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {x}_{k}^{i})\)  has 0 mean and is independent across \(i\) , and the last inequality follows Assumption 3.
By plugging in (REF ) and (REF ) into (REF ), we have the following.
\(\mathop {\mathbb {E}}\left[ F(\mathbf {u}_{k+1}) - F(\mathbf {u}_{k}) \right] &\le - \frac{\eta }{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 - \frac{\eta }{2} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \quad + \frac{\eta }{2} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) - \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] + \eta ^2 L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \\& \quad + \eta ^2 L \mathop {\mathbb {E}}\left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\&= - \frac{\eta }{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 - \frac{\eta }{2}(1 - 2 \eta L) \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \quad + \frac{\eta }{2} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) - \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] + \eta ^2 L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \) 
If \(\eta \le \frac{1}{2L}\) , it follows
\(\frac{\mathop {\mathbb {E}}\left[ F(\mathbf {u}_{k+1}) - F(\mathbf {u}_{k}) \right]}{\eta } & \le - \frac{1}{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 + \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \\& \quad + \frac{1}{2} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) - \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \le - \frac{1}{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 + \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2  \\& \quad + \frac{1}{2} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) - \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \nonumber \\& \le - \frac{1}{2} \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 + \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2 + \frac{L^2}{2} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber ,\) 

where (REF ) holds based on the convexity of \(l2\)  norm and Jensen's inequality.
By taking expectation and averaging across \(K\)  iterations, we have.
\(\frac{1}{K} \sum _{k=1}^{K} \frac{\mathop {\mathbb {E}}\left[ F(\mathbf {u}_{k+1}) - F(\mathbf {u}_{k}) \right]}{\eta } & \le -\frac{1}{2K} \sum _{k=1}^{K} \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 + \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \\& \quad \quad + \frac{L^2}{2K} \sum _{k-1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right]. \nonumber \) 
After a minor rearrangement, we have a telescoping sum as follows.
\(\begin{split}\frac{1}{K} \sum _{k=1}^{K} \mathop {\mathbb {E}}\left[ \left\Vert  \nabla F(\mathbf {u}_k) \right\Vert ^2 \right] &\le \frac{2}{\eta K} \mathop {\mathbb {E}}\left[F(\mathbf {u}_1) - F(\mathbf {u}_{k+1}) \right] + 2 \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \\& \quad + \frac{L^2}{K} \sum _{k=1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\&\le \frac{2}{\eta K} \mathop {\mathbb {E}}\left[F(\mathbf {u}_1) - F(\mathbf {u}_{*}) \right] + 2 \eta L \sigma ^2 \sum _{i=1}^{m} p_i^2 \nonumber \\& \quad + \frac{L^2}{K} \sum _{k=1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 \right], \nonumber \end{split}\) 
where \(\mathbf {u}_{*}\)  indicates the local minimum.
Here, we complete the proof.
Lemma 4.2. (Model Discrepancy) Under Assumption \(1 \sim 4\) , FedLAMA ensures
\(\begin{split}\frac{1}{K} \sum _{k=1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] & \le \frac{2 \eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A \kappa ^2}{L^2 (1 - A)} \\& \quad + \frac{A \beta ^2}{K L^2 (1 - A)} \sum _{k=1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F (\mathbf {u}_k) \right\Vert ^2 \right], \\\end{split}\) 
where \(A = 4 \eta ^2 (\tau _{max} - 1)^2 L^2\)  and \(\tau _{max}\)  is the largest averaging interval across all the layers.
We begin with rewriting the weighted average of the squared distance using the vectorized form of the local models as follows.
\(\sum _{i=1}^{m} p_i \left\Vert  \mathbf {u}_k - \mathbf {x}_{k}^{i} \right\Vert ^2 = \left\Vert  \mathbf {J} \mathbf {\hat{x}}_{k} - \mathbf {\hat{x}}_{k} \right\Vert ^2 = \left\Vert  (\mathbf {J} - \mathbf {I})\mathbf {\hat{x}}_{k} \right\Vert ^2 \\\) 
Then, according to the parameter update rule, we have
\((\mathbf {J} - \mathbf {I})\mathbf {\hat{x}}_{k} & = (\mathbf {J} - \mathbf {I})\mathbf {W}_{k-1}(\mathbf {\hat{x}}_{k-1} - \eta \mathbf {\hat{g}}_{k-1})\\& = (\mathbf {J} - \mathbf {I}) \mathbf {W}_{k-1} \mathbf {\hat{x}}_{k-1} - (\mathbf {J} - \mathbf {W}_{k-1}) \eta \mathbf {\hat{g}}_{k-1},\) 

where the second equality holds because \(\mathbf {J} \mathbf {W} = \mathbf {J}\)  and \(\mathbf {I} \mathbf {W} = \mathbf {W}\) .
Then, expanding the expression of \(\mathbf {x}_{k-1}\) , we have
\((\mathbf {J} - \mathbf {I}) \mathbf {\hat{x}}_{k} & = (\mathbf {J} - \mathbf {I}) \mathbf {W}_{k-1} (\mathbf {W}_{k-2} (\mathbf {\hat{x}}_{k-2} - \eta \mathbf {\hat{g}}_{k-2})) - (\mathbf {J} - \mathbf {W}_{k-1}) \eta \mathbf {\hat{g}}_{k-1} \nonumber \\& = (\mathbf {J} - \mathbf {I}) \mathbf {W}_{k-1} \mathbf {W}_{k-2} \mathbf {\hat{x}}_{k-2} - (\mathbf {J} - \mathbf {W}_{k-1} \mathbf {W}_{k-2}) \eta \mathbf {\hat{g}}_{k-2} - (\mathbf {J} - \mathbf {W}_{k-1}) \eta \mathbf {\hat{g}}_{k-1}. \nonumber \) 
Repeating the same procedure for \(\mathbf {\hat{x}}_{k-2}\) , \(\mathbf {\hat{x}}_{k-3}\) , \(\cdots \) , \(\mathbf {\hat{x}}_{2}\) , we have
\((\mathbf {J} - \mathbf {I}) \mathbf {\hat{x}}_{k} & = (\mathbf {J} - \mathbf {I}) \prod _{s=1}^{k-1} \mathbf {W}_{s} \mathbf {\hat{x}}_{1} - \eta \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1}\mathbf {W}_{(j,l)}) \mathbf {\hat{g}}_{s} \nonumber \\& = - \eta \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) \mathbf {\hat{g}}_{s}, \) 

where (REF ) holds because \(\mathbf {x}_{1}^i\)  is the same across all the workers and thus \((\mathbf {J} - \mathbf {I})\mathbf {\hat{x}}_{1} = 0\) .
Based on (REF ), we have
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad = \frac{1}{K} \sum _{k = 1}^{K} \left( \mathop {\mathbb {E}} \left[ \left\Vert  (\mathbf {J} - \mathbf {I}) \mathbf {\hat{x}}_{k} \right\Vert ^2 \right] \right) \nonumber \\& \quad \quad = \frac{1}{K} \sum _{k = 1}^{K} \left( \eta ^2 \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) \mathbf {\hat{g}}_{s} \right\Vert ^2 \right] \right) \nonumber \\& \quad \quad = \frac{1}{K} \sum _{k = 1}^{K} \left( \eta ^2 \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) ( \mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s} ) + \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) \mathbf {\hat{f}}_{s} \right\Vert ^2 \right] \right) \nonumber \\& \quad \quad \le \frac{2 \eta ^2}{K} \left( \underset{T_3}{\underbrace{ \sum _{k = 1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s} ) \right\Vert ^2 \right] }} + \underset{T_4}{\underbrace{ \sum _{k = 1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{s=1}^{k-1} ( \mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l} ) \mathbf {\hat{f}}_{s} \right\Vert ^2 \right] }} \right) \) 

where (REF ) holds based on the convexity of \(\ell 2\)  norm and Jensen's inequality.
Now, we focus on bounding \(T_3\)  and \(T_4\) , separately.
Bounding \(T_3\)
\(& \sum _{k = 1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l})(\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \right] \nonumber \\& \quad \quad = \sum _{k = 1}^{K} \sum _{s = 1}^{k - 1} \mathop {\mathbb {E}} \left[ \left\Vert  (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l})(\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \right]  \\& \quad \quad \le \sum _{k = 1}^{K} \sum _{s = 1}^{k - 1} \mathop {\mathbb {E}} \left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) \right\Vert _{op}^2 \right], \) 

where (REF ) holds because \(\hat{\mathbf {g}}_s - \hat{\mathbf {f}}_s\)  has 0 mean and independent across \(s\) , and () holds based on Lemma REF .
Without loss of generality, we replace \(k\)  with \(a\tau _{max} + b\) , where \(a\)  is the communication round index and \(b\)  is the iteration index within each communication round.
Then, we have
\(& \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}} \left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{k-1} \mathbf {W}_{l}) \right\Vert _{op}^2 \right] \nonumber \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = 1}^{a\tau } \mathop {\mathbb {E}}\left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1} \mathbf {W}_{l}) \right\Vert _{op}^2\right] \nonumber \\& \quad \quad \quad \quad + \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}}\left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1} \mathbf {W}_{l}) \right\Vert _{op}^2\right] \nonumber \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}}\left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1} \mathbf {W}_{l}) \right\Vert _{op}^2\right]  \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = a\tau _{max} + 1}^{a\tau + b - 1} \mathop {\mathbb {E}}\left[ \left\Vert  (\mathbf {\hat{g}}_{s} - \mathbf {\hat{f}}_{s}) \right\Vert ^2 \right]  \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}}\left[ \left\Vert  (g_{i}(\mathbf {x}_{s}^{i}) - \nabla F_i (\mathbf {x}_{s}^{i})) \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \sum _{i = 1}^{m} p_i \sigma ^2  \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} (b - 1) \sigma ^2 \nonumber \\& \quad \quad = \sum _{a = 0}^{K / \tau _{max} - 1} \frac{\tau _{max} (\tau _{max} - 1)}{2} \sigma ^2 \nonumber \\& \quad \quad \le K\frac{(\tau _{max} - 1)}{2} \sigma ^2, \) 

where (REF ) holds because \(\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1} \mathbf {W}_{l}\)  becomes 0 when \(s \le a\tau _{max}\) .
Remember FedLAMA guarantees that the whole parameters are synchronized at least once within every \(\tau _{max}\)  iterations.
() holds based on Lemma REF .
() holds based on Assumption 3.
Bounding \(T_4\)
\(& \sum _{k = 1}^{K - \tau _{max}} \mathop {\mathbb {E}}\left[\left\Vert  \sum _{s=1}^{k-1} (\mathbf {J} - \prod _{l=s}^{k-1}\mathbf {W}_{l}) \mathbf {\hat{f}}_{s} \right\Vert ^2\right] \nonumber \\& = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \mathop {\mathbb {E}}\left[\left\Vert  \sum _{s=1}^{a \tau + b - 1} (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1}\mathbf {W}_{l}) \mathbf {\hat{f}}_{s} \right\Vert ^2\right] \nonumber \\& = \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \mathop {\mathbb {E}}\left[\left\Vert  \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1}\mathbf {P}_{l}) \mathbf {\hat{f}}_{s} \right\Vert ^2\right] \\& \le \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \left( (b - 1) \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}}\left[\left\Vert  (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1} \mathbf {P}_{l}) \mathbf {\hat{f}}_{s} \right\Vert ^2\right] \right) \\& \le \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \left( (b - 1) \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}}\left[\left\Vert  \mathbf {\hat{f}}_{s} \right\Vert ^2 \left\Vert  (\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1}\mathbf {P}_{l}) \right\Vert _{op}^2 \right] \right) \\& \le \sum _{a = 0}^{K / \tau _{max} - 1} \sum _{b = 1}^{\tau _{max}} \left( (b - 1) \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + b - 1} \mathop {\mathbb {E}}\left[\left\Vert  \mathbf {\hat{f}}_{s} \right\Vert ^2 \right] \right)  \\& \le \frac{\tau _{max} (\tau _{max} - 1)}{2} \sum _{a = 0}^{K / \tau _{max} - 1} \left( \sum _{s = a\tau _{max} + 1}^{a\tau _{max} + \tau _{max} - 1} \mathop {\mathbb {E}}\left[\left\Vert  \mathbf {\hat{f}}_{s} \right\Vert ^2 \right] \right) \nonumber \\& \le \frac{\tau _{max} (\tau _{max} - 1)}{2} \sum _{k = 1}^{K} \mathop {\mathbb {E}}\left[\left\Vert  \mathbf {\hat{f}}_{k} \right\Vert ^2 \right] \nonumber \\& = \frac{\tau _{max}(\tau _{max} - 1)}{2} \sum _{k=1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_k^i) \right\Vert ^2 \right] , \) 

where (REF ) holds because \(\mathbf {J} - \prod _{l=s}^{a\tau _{max} + b - 1}\mathbf {P}_{l}\)  becomes 0 when \(s \le a\tau \) .
() holds based on the convexity of \(\ell \) 2 norm and Jensen's inequality.
() holds based on Lemma REF .
() holds based on Lemma REF .
Final Result
By plugging in () and () into (REF ), we have
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \frac{2 \eta ^2}{K} \left( K \frac{(\tau _{max} - 1)}{2} \sigma ^2 + \frac{\tau _{max} (\tau _{max} - 1)}{2} \left( \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \right) \right) \nonumber \\& \quad \quad = \eta ^2 (\tau _{max} - 1) \sigma ^2 + \frac{\eta ^2 \tau _{max} (\tau _{max} - 1)}{K} \left( \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] \right) \) 
The local gradient term on the right-hand side in (REF ) can be rewritten using the following inequality.
\(\mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_{k}^{i}) \right\Vert ^2 \right] & = \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {u}_{k}) + \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right] \nonumber \\& \le 2 \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {x}_{k}^{i}) - \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right] + 2 \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right]  \\& \le 2 L^2 \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] + 2 \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right], \) 

where (REF ) holds based on the convexity of \(\ell 2\)  norm and Jensen's inequality.
Plugging in () into (REF ), we have
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \eta ^2 (\tau _{max} - 1) \sigma ^2 + \frac{2 \eta ^2 \tau _{max} (\tau _{max} - 1) L^2}{K} \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \quad \quad + \frac{2 \eta ^2 \tau _{max} (\tau _{max} - 1)}{K} \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right]\) 
After a minor rearranging, we have
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - 2\eta ^2 \tau _{max} (\tau _{max} - 1) L^2} + \frac{2 \eta ^2 \tau _{max} (\tau _{max} - 1)}{K (1 - 2\eta ^2 \tau _{max} (\tau _{max} - 1) L^2)} \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right]  \\\) 
Let us define \(A = 2\eta ^2 \tau _{max} (\tau _{max} - 1) L^2\) . Then (REF ) is simplified as follows.
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A}{K L^2 (1 - A)} \sum _{k=1}^{K} \sum _{i = 1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right] \nonumber \) 
Based on Assumption 4, we have
\(& \frac{1}{K} \sum _{k = 1}^{K} \sum _{i=1}^{m} p_i \mathop {\mathbb {E}} \left[ \left\Vert  \mathbf {u}_{k} - \mathbf {x}_{k}^{i} \right\Vert ^2 \right] \nonumber \\& \quad \quad \le \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A \beta ^2}{K L^2 (1 - A)} \sum _{k=1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \sum _{i=1}^{m} p_i \nabla F_i(\mathbf {u}_{k}) \right\Vert ^2 \right] + \frac{A \kappa ^2}{L^2 (1 - A)} \nonumber \\& \quad \quad = \frac{\eta ^2 (\tau _{max} - 1) \sigma ^2}{1 - A} + \frac{A \beta ^2}{K L^2 (1 - A)} \sum _{k=1}^{K} \mathop {\mathbb {E}} \left[ \left\Vert  \nabla F(\mathbf {u}_{k}) \right\Vert ^2 \right] + \frac{A \kappa ^2}{L^2 (1 - A)}, \) 

where (REF ) holds based on the definition of the objective function (REF ).
Here, we complete the proof.

Proof of Other Lemmas
Lemma A.1 
Consider a real matrix \(\mathbf {A} \in \mathbb {R}^{md_j \times md_j}\)  and a real vector \(\mathbf {b} \in \mathbb {R}^{md_j}\) . If \(A\)  is symmetric and \(\mathbf {b} \ne \mathbf {0}_{md_j}\) , we have
\(\Vert  \mathbf {Ab} \Vert  \le \Vert  \mathbf {A} \Vert _{op} \Vert  \mathbf {b} \Vert \) 

\(\Vert  \mathbf {Ab} \Vert ^2 & = \frac{\Vert  \mathbf {Ab} \Vert ^2}{\Vert \mathbf {b} \Vert ^2} \Vert \mathbf {b}\Vert ^2 \nonumber \\& \le \Vert \mathbf {A}\Vert _{op}^{2}\Vert \mathbf {b}\Vert ^2 \) 

where (REF ) holds based on the definition of operator norm.
Lemma A.2 
Suppose there is a \(md \times md\)  averaging matrix \(\mathbf {P}\)  and \(t < \tau _{max}\) , then
\(\Vert  \mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_{i} \Vert _{op}^2 = 1.\) 

By the property 2 and 3 of the averaging matrix \(P\) , \(\prod _{i=1}^{t}\mathbf {P}_i\)  is symmetric for all \(t < \tau \) .
Since \(\mathbf {J}\)  is symmetric, \(\prod _{i=1}^{t}\mathbf {P}_i - \mathbf {J}\)  is also symmetric and it can be decomposed to \(Q\Lambda _t Q^\top \) , where \(Q\)  is an orthogonal matrix and \(\Lambda _t\)  is an eigenvalue vector.
By the definition of \(\mathbf {P}_i\)  and the property 4, the columns of \(\prod _{i=1}^{t}\mathbf {P}_i\)  can be grouped to two subsets:
The first group is the columns that are the same as the corresponding columns of \(\mathbf {J}\) .
Each of these columns contain \(m\)  elements of \(\frac{1}{m}\) .
The second group is the columns that do not contain \(\frac{1}{m}\) .
These columns strictly have only one element of 1 and all the other elements are 0.
Thus, in \(\mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_i\) , the first group columns are all zeroed out while the second group columns remain.
Because every block in both \(\mathbf {J}\)  and \(\prod _{i=1}^{t}\mathbf {P}_i\)  is diagonal, all the second group columns in \(\mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_i\)  are orthogonal and their eigenvalues are either 1 or \(-1\) .
Since \(t < \tau \) , \(\mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_i\)  has at least \(\frac{md}{\tau }\)  second group columns.
By the definition of the matrix operator norm, \(\Vert  \mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_i \Vert _{op}^2 = max\lbrace | \lambda (\mathbf {J} - \prod _{i=1}^{t}\mathbf {P}_i) |\rbrace  = 1\) .

Additional Experimental Results
In this section, we provide extra experimental results with extensive hyper-parameter settings.
We commonly use 128 clients and a lcoal batch size of 32 in all the experiments.
The gradual learning rate warmup ([1]}) is also applied to the first 10 epochs in all the experiments.
CIFAR-10 –
Table REF  and REF  show the CIFAR-10 classification performance of FedLAMA across different \(\phi \)  settings.
As expected, the accuracy is reduced as \(\phi \)  increases.
The IID and non-IID data settings show the common trend.
Depending on the system network bandwidth, \(\phi \)  can be tuned to be an appropriate value.
When \(\phi =2\) , the accuracy is almost the same as or even slightly higher than FedAvg accuracy.
If the network bandwidth is limited, one can increase \(\phi \)  and slightly increase the epoch budget to achieve a good accuracy.
Table REF  shows the CIFAR-10 accuracy across different \(\tau ^{\prime }\)  settings.
We see that the accuracy is significantly dropped as \(\tau ^{\prime }\)  increases.
CIFAR-100 –
Table REF  and REF  show the CIFAR-100 classification performance of FedLAMA across different \(\phi \)  settings.
FedLAMA achieves a comparable accuracy to FedAvg with a short aggregation interval, even when the degree of data heterogeneity is extreamly high (\(25\%\)  device sampling and Direchlet's coefficient of \(0.1\) ).
Table REF  shows the FedAvg accuracy with different \(\tau ^{\prime }\)  settings.
Under the strongly heterogeneous data distributions, FedAvg with a large aggregation interval (\(\tau \ge 12\) ) do not achieve a reasonable accuracy.
FEMNIST –
Table REF  shows the FEMNIST classification performance of FedLAMA across different \(\phi \)  settings.
FedLAMA achieves a similar accuracy to the baseline (FedAvg with \(\tau ^{\prime }=10\) ) even when using a large interval increase factor \(\phi \ge 4\) .
These results demonstrate the effectiveness of the proposed layer-wise adaptive model aggregation method on the problems with heterogeneous data distributions.
<TABLE><TABLE><TABLE><TABLE><TABLE><TABLE><TABLE>