SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The University of Sydney ysun9899@uni.sydney.edu.au
Li Shen*
JD Explore Academy mathshenli@gmail.com
The University of Sydney dacheng.tao@gmail.com
Federated learning (FL) is a distributed paradigm that coordinates massive local clients to collaboratively train a global model via stage-wise local training processes on the heterogeneous dataset. Previous works have implicitly studied that FL suffers from the "client-drift" problem, which is caused by the inconsistent optimum across local clients. However, till now it still lacks solid theoretical analysis to explain the impact of this local inconsistency. To alleviate the negative impact of the "client drift" and explore its substance in FL, in this paper, we first design an efficient FL algorithm FedInit, which allows employing the personalized relaxed initialization state at the beginning of each local training stage. Specifically, FedInit initializes the local state by moving away from the current global state towards the reverse direction of the latest local state. This relaxed initialization helps to revise the local divergence and enhance the local consistency level. Moreover, to further understand how inconsistency disrupts performance in FL, we introduce the excess risk analysis and study the divergence term to investigate the test error of the proposed FedInit method. Our studies show that optimization error is not sensitive to this local inconsistency, while it mainly affects the generalization error bound in FedInit. Extensive experiments are conducted to validate this conclusion. Our proposed FedInit could achieve state-of-the-art (SOTA) results compared to several advanced benchmarks without any additional costs. Meanwhile, stage-wise relaxed initialization could also be incorporated into the current advanced algorithms to achieve higher performance in the FL paradigm.
Since McMahan et al. [26] developed federated learning, it becomes a promising paradigm to effectively make full use of the computational ability of massive edge devices. Kairouz et al. [17] further classify the modes based on the specific tasks and different environmental setups. Different from centralized training, FL utilizes a central server to coordinate the clients to perform several local training stages and aggregate local models as one global model. However, due to the heterogeneous dataset, it still suffers from significant performance degradation in practical scenarios.
Several previous studies explore the essence of performance limitations in FL and summarize it as the "client-drift" problem [1, 19, 22, 38, 44, 46, 48]. From the perspective of the global target, Karimireddy et al. [19] claim that the aggregated local optimum is far away from the global optimum due to the heterogeneity of the local dataset, which introduces the "client-drift" in FL. However, under limited local training steps, local clients can not genuinely approach the local optimum. To describe this negative impact more accurately, Acar et al. [1] and Wang et al. [44] point out that each locally optimized objective should be regularized to be aligned with the global objective. Moreover, beyond the guarantees of local consistent objective, Xu et al. [46] indicate that the performance
*Li Shen is the corresponding author.
degradation could be further eliminated in FL if it guarantees the local consistent updates at each communication round, which is more similar to the centralized scenarios. These arguments intuitively provide forward-looking guidance for improving the performance in FL. However, in the existing analysis, there is still no solid theoretical support to understand the impact of the consistency term, which also severely hinders the further development of the FL paradigm.
To alleviate the negative impact of the "client-drift" problem and strengthen consistency in the FL paradigm, in this paper, we take into account adopting the personalized relaxed initialization at the beginning of each communication round, dubbed FedInit method. Specifically, FedInit initializes the selected local state by moving away from the current global state towards the reverse direction of the current latest local state. Personalized relaxed initialization helps each local model to revise its divergence and gather together with each other during the local training process. This flexible approach is surprisingly effective in FL and only adopts a constant coefficient to control the divergence level of the initialization. It could also be easily incorporated as a plug-in into other advanced benchmarks to further improve their performance.
Moreover, to explicitly understand how local inconsistency disrupts performance, we introduce the excess risk analysis to investigate the test error of FedInit under the smooth non-convex objective, which includes an optimization error bound and a generalization error bound. Our theoretical studies indicate that the optimization error is insensitive to local inconsistency, while it mainly affects the generalization performance. Under  $PL$ -condition, consistency performs as the dominant term in the excess risk. Extensive empirical studies are conducted to validate the efficiency of the FedInit method. On the CIFAR-10/100 dataset, it could achieve SOTA results compared to several advanced benchmarks without additional costs. It also helps to enhance the consistency level in FL.
In summary, the main contributions of this work are stated as follows:
- We propose an efficient and novel FL method, dubbed FedInit, which adopts the personalized relaxed initialization state on the selected local clients at each communication round. Relaxed initialization is dedicated to enhancing local consistency during training, and it is also a practical plug-in that could easily to incorporated into other methods.
- One important contribution is that we introduce the excess risk analysis in the proposed FedInit method to understand the intrinsic impact of local consistency. Our theoretical studies prove that the optimization error is insensitive to consistency, while it mainly affects the test error and generalization error bound.
- Extensive numerical studies are conducted on the real-world dataset to validate the efficiency of the FedInit method, which outperforms several SOTA benchmarks without additional training costs. Meanwhile, as an efficient plug-in, relaxed initialization (FedInit) could also help the other benchmarks in our paper to achieve higher performance with effortlessness.
Consistency in FL. FL employs an enormous number of edge devices to jointly train a single model among the isolated heterogeneous dataset [17, 26]. As a standard benchmark, FedAvg [2, 26, 48] allows the local stochastic gradient descent (local SGD) [10, 23, 45] based updates and uniformly selected partial clients' participation to alleviate the communication bottleneck. The stage-wise local training processes lead to significant divergence for each client [5, 25, 43, 44]. To improve the efficiency of the FL paradigm, a series of methods are proposed. Karimireddy et al. [19] indicate that inconsistent local optimums cause the severe "client drift" problem and propose the SCAFFOLD method which adopts the variance reduction [6, 16] technique to mitigate it. Li et al. [22] penalize the prox-term on the local objective to force the local update towards both the local optimum and the last global state. Zhang et al. [49] utilize the primal-dual method to improve consistency via solving local objectives under the equality constraint. Specifically, a series of works further adopt the alternating direction method of multipliers (ADMM) to optimize the global objective [1, 9, 41, 52], which could also enhance the consistency term. Beyond these, a series of momentum-based methods are proposed to strengthen local consistency. Wang et al. [42] study a global momentum update method to stabilize the global model. Further, Gao et al. [8] use a local drift correction via a momentum-based term to revise the local gradient, efficiently reducing inconsistency. Ozfatura et al. [28], Xu et al. [46], Sun et al. [39] propose a similar client-level momentum to force the local update towards the last global direction. A variant of client-level momentum that adopts the inertial momentum to further improve
the local consistency level [24, 40]. At present, improving the consistency in FL remains a very important and promising research direction. Though these studies involve the heuristic discussion on consistency, in this paper we focus on the personalized relaxed initialization.
Generalization in FL. A lot of works have studied the properties of generalization in FL. Based on the margin loss [3, 7, 27], Reiszadeh et al. [31] develop a robust FL paradigm to alleviate the distribution shifts across the heterogeneous clients. Shi et al. [32] study the efficient and stable model technique of model ensembling. Yagli et al. [47] prove the information-theoretic bounds on the generalization error and privacy leakage in the general FL paradigm. Qu et al. [29] propose to adopt the sharpness aware minimization (SAM) optimizer on the local client to improve the flatness of the loss landscape. Caldarola et al. [4], Sun et al. [37, 38], Shi et al. [33, 34] propose several variants based on SAM that could achieve higher performance. However, these works only focus on the generalization efficiency in FL, while in this paper we prove that its generalization error bound is dominated by consistency.
Under the cross-device FL setups, there are a very large number of local clients to collaboratively train a global model. Due to privacy protection and unreliable network bandwidth, only a fraction of devices are open-accessed at any one time [17, 29]. Therefore, we define each client stores a private dataset  $S_{i} = \{z_{j}\}$  where  $z_{j}$  is drawn from an unknown unique distribution  $\mathcal{D}_i$ . The whole local clients constitute a set  $\mathcal{C} = \{i\}$  where  $i$  is the index of each local client and  $|\mathcal{C}| = C$ . Actually, in the training process, we expect to approach the optimum of the population risk  $F$ :
$$
w _ {\mathcal {D}} ^ {\star} \in \arg \min  _ {w} \left\{F (w) \triangleq \frac {1}{C} \sum_ {i \in \mathcal {C}} F _ {i} (w) \right\}, \tag {1}
$$
where  $F_{i}(w) = \mathbb{E}_{z_{j}\sim \mathcal{D}_{i}}F_{i}(w,z_{j})$  is the local population risk. While in practice, we usually consider the empirical risk minimization of the non-convex finite-sum problem in FL as:
$$
w ^ {\star} \in \arg \min  _ {w} \left\{f (w) \triangleq \frac {1}{C} \sum_ {i \in \mathcal {C}} f _ {i} (w) \right\}, \tag {2}
$$
where  $f_{i}(w) = \frac{1}{S_{i}}\sum_{z_{j}\in S_{i}}f_{i}(w;z_{j})$  is the local empirical risk. In Section 4.1, we will analyze the difference between these two results. Furthermore, we introduce the excess risk analysis to upper bound the test error and further understand how consistency works in the FL paradigm.
3.2 Personalized Relaxed Initialization
In this part, we introduce the relaxed initialization in FedInit method. FedAvg proposes the localSGD-based implementation in the FL paradigm with a partial participation selection. It allows uniformly selecting a subset of clients  $\mathcal{N}$  to participate in the current training. In each round, it initializes the local model as the last global model. Therefore, after each round, the local models are always far away from each other. The local offset  $w_{i,K}^{t - 1} - w^t$  is the main culprit leading to inconsistency. Moreover, for different clients, their impacts vary with local heterogeneity. To alleviate this divergence, we propose the FedInit method which adopts the personalized relaxed initialization at the beginning of each round. Concretely, on the selected active clients, it begins the local training from a new personalized state, which moves away from the last global model towards the reverse direction from the latest local state (Line.6 in
Algorithm 1: FedInit Algorithm
Input: model  $w$ , local model  $w_{i}, T, K, \beta$ .
Output: model  $w^T$ .
Initialize states: initialize  $w^{-1} = w_{i,0}^{-1} = w^{0}$ .
for  $t = 0,1,\dots,T - 1$  do
randomly select active clients set  $\mathcal{N}$  from  $\mathcal{C}$  for  $i\in \mathcal{N}$  in parallel do
send the  $w^t$  to the active clients
set the  $w^{t} + \beta (w^{t} - w_{i,K}^{t - 1})$  as  $w_{i,0}^{t}$
for  $k = 0,1,\dots,K - 1$  do
compute gradient  $g_{i,k}^{t}$  at  $w_{i,k}^{t}$
$$
w _ {i, k + 1} ^ {t} = w _ {i, k} ^ {t} - \eta g _ {i, k} ^ {t}
$$
send the  $w_{i}^{t} = w_{i,K}^{t}$  to the server
$$
w ^ {t + 1} = \frac {1}{N} \sum_ {i \in \mathcal {N}} w _ {i} ^ {t}
$$
Algorithm 1). A coefficient  $\beta$  is adopted to control the level of personality. This offset  $\beta (w^{t} - w_{i,K}^{t - 1})$  in the relaxed initialization (RI) provides a correction that could help local models gather together after the local training process. Furthermore, this relaxed initialization is irrelevant to the local optimizer, which means, it could be easily incorporated into other methods. Additionally, FedInit does not require extra auxiliary information to communicate. It is a practical technique in FL.
In this section, we first introduce the excess risk in FL which could provide a comprehensive analysis on the joint performance of both optimization and generalization. In the second part, we introduce the main assumptions adopted in our proofs and discuss them in different situations. Then we state the main theorems on the analysis of the excess risk of our proposed FedInit method.
4.1 Excess Risk Error in FL
Since Karimireddy et al. [19] pointed out that client-drift problem may seriously damage the performance in the FL paradigm, many previous works [15, 18, 19, 30, 36, 44, 46, 48] have learned its inefficiency in the FL paradigm. However, most of the analyses focus on the studies from the onefold perspective of optimization convergence but ignore investigating its impact on generality. To further provide a comprehensive understanding of how client-drift affects the performance in FL, we adopt the well-known excess risk in the analysis of our proposed FedInit method.
We denote  $w^T$  as the model generated by FedInit method after  $T$  communication rounds. Compared with  $f(w^T)$ , we mainly focus on the efficiency of  $F(w^T)$  which corresponds to its generalization performance. Therefore, we analyze the  $\mathbb{E}[F(w^T)]$  from the excess risk  $\mathcal{E}_E$  as:
$$
\mathcal {E} _ {E} = \mathbb {E} [ F (w ^ {T}) ] - \mathbb {E} [ f (w ^ {*}) ] = \underbrace {\mathbb {E} [ F (w ^ {T}) - f (w ^ {T}) ]} _ {\mathcal {E} _ {G}: \text {g e n e r a l i z a t i o n e r r o r}} + \underbrace {\mathbb {E} [ f (w ^ {T}) - f (w ^ {*}) ]} _ {\mathcal {E} _ {O}: \text {o p t i m i z a t i o n e r r o r}}. \tag {3}
$$
Generally, the  $\mathbb{E}[f(w^{*})]$  is expected to be very small and even to zero if the model could well-fit the dataset. Thus  $\mathcal{E}_E$  could be considered as the joint efficiency of the generated model  $w^{T}$ . Thereinto,  $\mathcal{E}_G$  means the different performance of  $w^{T}$  between the training dataset and the test dataset, and  $\mathcal{E}_O$  means the similarity between  $w^{T}$  and optimization optimum  $w^{\star}$  on the training dataset.
In this part, we introduce some assumptions adopted in our analysis. We will discuss their properties and distinguish the proofs they are used in.
Assumption 1 For  $\forall w_{1},w_{2}\in \mathbb{R}^{d}$ , the non-convex local function  $f_{i}$  satisfies  $L$ -smooth if:
$$
\left\| \nabla f _ {i} \left(w _ {1}\right) - \nabla f _ {i} \left(w _ {2}\right) \right\| \leq L \| w _ {1} - w _ {2} \|. \tag {4}
$$
Assumption 2 For  $\forall w\in \mathbb{R}^d$  , the stochastic gradient is bounded by its expectation and variance as:
$$
\mathbb {E} \left[ g _ {i, k} ^ {t} \right] = \nabla f _ {i} \left(w _ {i, k} ^ {t}\right), \quad \mathbb {E} \| g _ {i, k} ^ {t} - \nabla f _ {i} \left(w _ {i, k} ^ {t}\right) \| ^ {2} \leq \sigma_ {l} ^ {2}. \tag {5}
$$
Assumption 3 For  $\forall w\in \mathbb{R}^d$  , the heterogeneous similarity is bounded on the gradient norm as:
$$
\mathbb {E} \| \nabla f _ {i} (w) \| ^ {2} \leq G ^ {2} + B ^ {2} \mathbb {E} \| \nabla f (w) \| ^ {2}. \tag {6}
$$
Assumption 4 For  $\forall w_{1},w_{2}\in \mathbb{R}^{d}$ , the global function  $f$  satisfies  $L_{G}$ -Lipschitz if:
$$
\left\| f \left(w _ {1}\right) - f \left(w _ {2}\right) \right\| \leq L _ {G} \| w _ {1} - w _ {2} \|. \tag {7}
$$
Assumption 5 For  $\forall w\in \mathbb{R}^d$  , let  $w^{\star}\in \arg \min_{w}f(w)$  , the function  $f$  satisfies PL-condition if:
$$
2 \mu (f (w) - f (w ^ {\star})) \leq \| \nabla f (w) \| ^ {2}. \tag {8}
$$
**Discussions.** Assumptions  $1 \sim 3$  are three general assumptions to analyze the non-convex objective in FL, which is widely used in the previous works [15, 18, 19, 30, 36, 44, 46, 48]. Assumption 4 is used to bound the uniform stability for the non-convex objective, which is used in [11, 51]. Different from the analysis in the margin-based generalization bound [27, 29, 31, 38] that focus on understanding how the designed objective affects the final generalization performance, our work focuses on understanding
how the generalization performance changes in the training process. We consider the entire training process and adopt uniform stability to measure the global generality in FL. For the general non-convex objective, one often uses the gradient norm  $\mathbb{E}\| \nabla f(w)\|^2$  instead of bounding the loss difference  $\mathbb{E}\left[f(w^T) - f(w^\star)\right]$  to measure the optimization convergence. To construct and analyze the excess risk, and further understand how the consistency affects the FL paradigm, we follow [51] to use Assumption 5 to bound the loss distance. Through this, we can establish a theoretical framework to jointly analyze the trade-off on the optimization and generalization in the FL paradigm.
4.3 Main Theorems
4.3.1 Optimization Error  $\varepsilon_{O}$
Theorem 1 Under Assumptions  $1 \sim 3$ , let participation ratio is  $N / C$  where  $1 < N < C$ , let the learning rate satisfy  $\eta \leq \min \left\{\frac{N}{2CKL}, \frac{1}{NKL}\right\}$  where  $K \geq 2$ , let the relaxation coefficient  $\beta \leq \frac{\sqrt{2}}{12}$  and after training  $T$  rounds, the global model  $w^t$  generated by FedInit satisfies:
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| f \left(w ^ {t}\right) \| ^ {2} \leq \frac {2 \left(f \left(w ^ {0}\right) - f \left(w ^ {\star}\right)\right)}{\lambda \eta K T} + \frac {\kappa_ {2} \eta L}{\lambda N} \sigma_ {l} ^ {2} + \frac {3 \kappa_ {1} \eta K L}{\lambda N} G ^ {2}, \tag {9}
$$
where  $\lambda \in (0,1)$ ,  $\kappa_{1} = \frac{1300\beta^{2}}{1 - 72\beta^{2}} + 17$ , and  $\kappa_{2} = \frac{1020\beta^{2}}{1 - 72\beta^{2}} + 13$  are three constants. Further, by selecting the proper learning rate  $\eta = \mathcal{O}\left(\sqrt{\frac{N}{KT}}\right)$  and let  $D = f(w^{0}) - f(w^{\star})$  as the initialization bias, the global model  $w^{t}$  satisfies:
$$
\frac {1}{T} \sum_ {t = 0} ^ {T - 1} \mathbb {E} \| f \left(w ^ {t}\right) \| ^ {2} \leq \mathcal {O} \left(\frac {D + L \left(\sigma_ {l} ^ {2} + K G ^ {2}\right)}{\sqrt {N K T}}\right). \tag {10}
$$
Theorem 1 provides the convergence rate of the FedInit method without the  $PL$ -condition, which could achieve the  $\mathcal{O}(1 / \sqrt{NKT})$  with the linear speedup of  $N\times$ . The dominant term of the training convergence rate is the heterogeneous bias  $G$ , which is  $K\times$  larger than the initialization bias  $D$  and stochastic bias  $\sigma_{l}$ . According to the formulation (10), by ignoring the initialization bias, the best local interval  $K = \mathcal{O}(\sigma_l^2 /G^2)$ . This selection also implies that when  $G$  increases, which means the local heterogeneity increases, the local interval  $K$  is required to decrease appropriately to maintain the same efficiency. More importantly, though FedInit adopts a weighted bias on the initialization state at the beginning of each communication round, the divergence term  $\mathbb{E}\| w_{i,K}^{t - 1} - w^t\|^2$  does not affect the convergence bound whether  $\beta$  is 0 or not. This indicates that the FL paradigm allows a divergence of local clients from the optimization perspective. Proof details are stated in Appendix A.2.3.
Theorem 2 Under Assumptions  $1 \sim 3$  and 5, let participation ratio is  $N / C$  where  $1 < N < C$ , let the learning rate satisfy  $\eta \leq \min \left\{\frac{N}{2CKL}, \frac{1}{NKL}, \frac{1}{\lambda\mu K}\right\}$  where  $K \geq 2$ , let the relaxation coefficient  $\beta \leq \frac{\sqrt{2}}{12}$ , and after training  $T$  rounds, the global model  $w^t$  generated by FedInit satisfies:
$$
\mathbb {E} \left[ f \left(w ^ {T}\right) - f \left(w ^ {\star}\right) \right] \leq e ^ {- \lambda \mu \eta K T} \mathbb {E} \left[ f \left(w ^ {0}\right) - f \left(w ^ {\star}\right) \right] + \frac {3 \kappa_ {1} \eta K L}{2 N \lambda \mu} G ^ {2} + \frac {\kappa_ {2} \eta L}{2 N \lambda \mu} \sigma_ {l} ^ {2}, \tag {11}
$$
where  $\lambda, \kappa_1, \kappa_2$  is defined in Theorem 1. Further, by selecting the proper learning rate  $\eta = \mathcal{O}\left(\frac{\log(\lambda\mu NKT)}{\lambda\mu KT}\right)$  and let  $D = f(w^0) - f(w^\star)$  as the initialization bias, the global model  $w^t$  satisfies:
$$
\mathbb {E} \left[ f \left(w ^ {T}\right) - f \left(w ^ {\star}\right) \right] = \widetilde {\mathcal {O}} \left(\frac {D + L \left(\sigma_ {l} ^ {2} + K G ^ {2}\right)}{N K T}\right). \tag {12}
$$
To bound the  $\mathcal{E}_O$  term, we adopt Assumption 5 and prove that FedInit method could achieve the  $\mathcal{O}(1 / NKT)$  rate where we omit the  $\mathcal{O}(\log (NKT))$  term. It maintains the properties stated in the Theorem 1. Detailed proofs of the convergence bound are stated in Appendix A.2.4.
4.3.2 Generalization Error  $\mathcal{E}_G$
Uniform Stability. One powerful analysis of the generalization error is the uniform stability [11, 21, 50]. It says, for a general proposed method, its generalization error is always lower than the bound
of uniform stability. We assume that there is a new set  $\widetilde{\mathcal{C}}$  where  $\mathcal{C}$  and  $\widetilde{\mathcal{C}}$  differ in at most one data sample on the  $i^{\star}$ -th client. Then we denote the  $w^{T}$  and  $\widetilde{w}^T$  as the generated model after training  $T$  rounds on these two sets, respectively. Thus, we have the following lemma:
Lemma 1 (Uniform Stability. [11]) For the two models  $w^T$  and  $\tilde{w}^T$  generated as introduced above, a general method satisfies  $\epsilon$ -uniformly stability if:
$$
\sup  _ {z _ {j} \sim \left\{\mathcal {D} _ {i} \right\}} \mathbb {E} \left[ f \left(w ^ {T}; z _ {j}\right) - f \left(\widetilde {w} ^ {T}; z _ {j}\right) \right] \leq \epsilon . \tag {13}
$$
Moreover, if a general method satisfies  $\epsilon$ -uniformly stability, then its generalization error satisfies  $\mathcal{E}_G \leq \sup_{z_j \sim \{\mathcal{D}_i\}} \mathbb{E}[f(w^T;z_j) - f(\widetilde{w}^T;z_j)] \leq \epsilon$  [50].
Theorem 3 Under Assumptions 1, 2, 4, and 5, let all conditions above be satisfied, let learning rate  $\eta = \mathcal{O}\left(\frac{1}{KT}\right) = \frac{c}{T}$  where  $c = \frac{\mu_0}{K}$  is a constant, and let  $|\mathcal{S}_i| = S$  as the number of the data samples, by randomly selecting the sample  $z$ , we can bound the uniform stability of our proposed FedInit as:
$$
\begin{array}{l} \mathbb {E} \| f (w ^ {T + 1}; z) - f (\widetilde {w} ^ {T + 1}; z) \| \\ \leq \frac {1}{S - 1} \left[ \frac {2 \left(L _ {G} ^ {2} + S L _ {G} \sigma_ {l}\right) (U T K) ^ {c L}}{L} \right] ^ {\frac {1}{1 + c L}} + (1 + \beta) ^ {\frac {1}{\beta c L}} \left[ \frac {U L T K}{2 \left(L _ {G} ^ {2} + S L _ {G} \sigma_ {l}\right)} \right] ^ {\frac {c L}{1 + c L}} \sum_ {t = 1} ^ {T} \frac {\sqrt {\Delta^ {t}}}{T}, \tag {14} \\ \end{array}
$$
where  $U$  is a constant and  $\Delta^t = \frac{1}{C}\sum_{i\in \mathcal{C}}\mathbb{E}\| w_{i,K}^{t - 1} - w^t\|^2$  is the divergence term at round  $t$ .