SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xuming An $^{1}$  Li Shen $^{2*}$  Han Hu $^{1*}$  Yong Luo $^{3}$
$^{1}$  School of Information and Electronics, Beijing Institute of Technology, China  
 $^{2}$  JD Explore Academy, China  $^{3}$  School of Computer Science, Wuhan University, China  
{anxuming,hhu}@bit.edu.cn, mathshenli@gmail.com, luoyong@whu.edu.cn
Federated Learning (FL) is an emerging collaborative machine learning framework where multiple clients train the global model without sharing their own datasets. In FL, the model inconsistency caused by the local data heterogeneity across clients results in the near-orthogonality of client updates, which leads to the global update norm reduction and slows down the convergence. Most previous works focus on eliminating the difference of parameters (or gradients) between the local and global models, which may fail to reflect the model inconsistency due to the complex structure of the machine learning model and the Euclidean space's limitation in meaningful geometric representations. In this paper, we propose FedMRUR by adopting the manifold model fusion scheme and a new global optimizer to alleviate the negative impacts. Concretely, FedMRUR adopts a hyperbolic graph manifold regularizer enforcing the representations of the data in the local and global models are close to each other in a low-dimensional subspace. Because the machine learning model has the graph structure, the distance in hyperbolic space can reflect the model bias better than the Euclidean distance. In this way, FedMRUR exploits the manifold structures of the representations to significantly reduce the model inconsistency. FedMRUR also aggregates the client updates norms as the global update norm, which can appropriately enlarge each client's contribution to the global update, thereby mitigating the norm reduction introduced by the near-orthogonality of client updates. Furthermore, we theoretically prove that our algorithm can achieve a linear speedup property  $\mathcal{O}\left(\frac{1}{\sqrt{SKT}}\right)$  for non-convex setting under partial client participation, where  $S$  is the participated clients number,  $K$  is the local interval and  $T$  is the total number of communication rounds. Experiments demonstrate that FedMRUR can achieve a new state-of-the-art (SOTA) accuracy with less communication.
FL is a collaborative distributed framework where multiple clients jointly train the model with their private datasets [27, 28]. To protect privacy, each client is unable to access the other dataset [2]. A centralized server receives the parameters or gradients from the clients and updates the global model[46]. Due to the limited communication resource, only part of the clients is involved in the collaborative learning process and train the local model in multiple intervals with their own datasets within one communication round [23]. Due to the data heterogeneity, clients' partial participation and multiple local training yield severe model inconsistency, which leads to the divergences between the directions of the local updates from the clients and thus reduces the magnitude of global updates [17]. Therefore, the model inconsistency is the major source of performance degradation in FL [40, 15].
*Corresponding authors: Li Shen and Han Hu
So far, numerous works have focused on the issues of model inconsistency to improve the performance of FL. Many of them [20, 16, 42, 1] utilize the parameter (or gradient) difference between the local and global model to assist the local training. By incorporating the global model information into local training, the bias between the local and global objectives can be diminished at some level. However, the parameter (or gradient) difference may fail to characterize the model bias due to the complex structure of modern machine learning model and the Euclidean space has limitations in providing powerful and meaningful geometric representations [10]. Meanwhile, incorporating the difference introduces extra high computation and communication costs because of the high-dimensional model parameter, which is common in the modern machine learning area[25]. Some other works [22, 39, 44, 24] exploit the permutation invariance property of the neurons in the neural networks to align and aggregate the model parameters for handling the model inconsistency issues, but the extra computation required for neuron alignment may slow down the speed of FL. In addition, Charles et al. [3] demonstrate that after multiple rounds, the similarities between the client updates approach zero and the local update direction are almost orthogonal to each other in FL. If the server aggregates the local updates, each client's contribution is little, which reduces the global update step. Therefore, we need to reduce the model inconsistency and compensate for the global norm reduction introduced by the near-orthogonality of client updates.
In order to alleviate the model inconsistency and compensate for the global reduction, we propose a practical and novel algorithm, dubbed as FedMRUR (Federated learning with Manifold Regularization and Normalized Update Reaggregation). FedMRUR adopts two techniques to achieve SOTA performance. i) Firstly, FedMRUR adopts the hyperbolic graph fusion technique to reduce the model inconsistency between the client and server within local training. The intuition is that adding the manifold regularizer to the loss function to constrain the divergence between the local and global models. Unlike the Euclidean space, the hyperbolic space is a manifold structure with the constant negative curvature, which has the ability to produce minimal distortion embedding[8] with the low storage constraints[30] for graph data. And the neural network, the most prevail machine learning model, has a graph structure[34, 24], we map the representations to the hyperbolic space and compute their distance to indicate the model bias precisely. Considering the numerical stability[6], we select the Lorentz model to describe the hyperbolic space and the squared Lorentzian distance[19] to indicate the representations' proximity. By adopting the hyperbolic graph fusion technique, FedMRUR can constrain model inconsistency efficiently. ii) Secondly, FedMRUR aggregates the client's local updates in a novel normalized way to alleviate the global norm reduction. In the normalized aggregation scheme, the server aggregates the local update norms as the global norm and normalizes the sum of the local updates as the global direction. Compared with directly aggregating local updates, the new aggregation scheme enables each customer's contribution to be raised from its projection on the global direction to its own size. As a result, the size of the global update becomes larger and compensates for the norm reduction introduced by model inconsistency, which improves the convergence and generalization performance of the FL framework.
Theoretically, we prove that the proposed FedMRUR can achieve the convergence rate of  $\mathcal{O}\left(\frac{1}{\sqrt{SKT}}\right)$  on the non-convex and L-smooth objective functions with heterogeneous datasets. Extensive experiments on CIFAR-10/100 and TinyImagenet show that our proposed FedMRUR algorithm achieves faster convergence speed and higher test accuracy in training deep neural networks for FL than several baselines including FedAvg, FedProx, SCAFFOLD, FedCM, FedExp, and MoFedSAM. We also study the impact on the performance of adopting the manifold regularization scheme and normalized aggregation scheme. In summary, the main contributions are as follows:
- We propose a novel and practical FL algorithm, FedMRUR, which adopts the hyperbolic graph fusion technique to effectively reduce the model inconsistency introduced by data heterogeneity, and a normalized aggregation scheme to compensate the global norm reduction due to the near-orthogonality of client updates, which achieves fast convergence and generalizes better.
- We provide the upper bound of the convergence rate under the smooth and non-convex cases and prove that FedMRUR has a linear speedup property  $\mathcal{O}\left(\frac{1}{\sqrt{SKT}}\right)$ .
- We conduct extensive numerical studies on the CIFAR-10/100 and TinyImagenet dataset to verify the performance of FedMRUR, which outperforms several classical baselines on different data heterogeneity.
McMahan et al. [27] propose the FL framework and the well-known algorithm, FedAvg, which has been proved to achieve a linear speedup property [43]. Within the FL framework, clients train local models and the server aggregates them to update the global model. Due to the heterogeneity among the local dataset, there are two issues deteriorating the performance: the model biases across the local solutions at the clients [20] and the similarity between the client updates (which is also known as the near-orthogonality of client updates) [3], which needs a new aggregation scheme to solve. In this work, we focus on alleviating these two challenges to improve the convergence of the FL algorithms.
Model consistency. So far, numerous methods focus on dealing with the issue of model inconsistency in the FL framework. Li et al. [20] propose the FedProx algorithm utilizing the parameter difference between the local and global model as a prox-correction term to constrain the model bias during local training. Similar to [20], during local training, the dynamic regularizer in FedDyn [1] also utilizes the parameter difference to force the local solutions approaching the global solution. FedSMOO [36] utilizes a dynamic regularizer to make sure that the local optima approach the global objective. Karimireddy et al. [16] and Haddadpour et al. [9] mitigate the model inconsistency by tracking the gradient difference between the local and global side. Xu et al. [42] and Qu et al. [31] utilize a client-level momentum term incorporating global gradients to enhance the local training process. Sun et al. [37] estimates the global aggregation offset in the previous round and corrects the local drift through a momentum-like term to mitigate local over-fitting. Liu et al. [26] incorporate the weighted global gradient estimations as the inertial correction terms guiding the local training to enhance the model consistency. Charles et al. [4] demonstrate that the local learning rate decay scheme can achieve a balance between the model inconsistency and the convergence rate. Tan et al. [38] show that the local learning rate decay scheme is unable to reduce the model inconsistency when clients communicate with the server in an asynchronous way. Most methods alleviate the model inconsistency across the clients by making use of the parameter (or gradient) difference between the local and global model.
Aggregation scheme. There are numerous aggregation schemes applied on the server side for improving performance. Some works utilize classical optimization methods, such as SGD with momentum [45], and adaptive SGD [5], to design the new global optimizer for FL. For instance, FedAvgM [13, 35] and STEM [17] update the global model by combining the aggregated local updates and a momentum term. Reddi et al. [32] propose a federated optimization framework, where the server performs the adaptive SGD algorithms to update the global model. FedNova [41] normalizes the local updates and then aggregates them to eliminate the data and device heterogeneity. In addition, the permutation invariance property of the neurons in the neural networks is also applied for improving robustness to data heterogeneity. FedFTG [48] applies the data-free knowledge distillation method to fine-tune the global model in the server. FedMA [39] adopts the Bayesian optimization method to align and average the neurons in a layer-wise manner for a better global solution. Li et al. [22] propose Position-Aware Neurons (PANs) coupling neurons with their positions to align the neurons more precisely. Liu et al. [24] adopt the graph matching technique to perform model aggregation, which requires a large number of extra computing resources in the server. Many deep model fusion methods [21] are also applied in the research field of FL, such as model ensemble [47] and CVAE [12]. The aforementioned algorithms utilize the local parameters or gradients directly without considering the near-orthogonality of client updates, which may deteriorate the convergence performance of the FL framework.
The proposed method FedMRUR adopts the hyperbolic graph fusion technique to reduce the model inconsistency and a normalized update aggregation scheme to mitigate the norm reduction of the global update. Compared with the previous works, we utilize the squared Lorentzian distance of the features in the local and global model as the regularization term. This term can more precisely measure the model bias in the low-dimensional subspace. For the update aggregation at the server, FedMRUR averages the local updates norm as the global update norm, which achieves to alleviate the norm reduction introduced by the near-orthogonality of the client updates.
In this section, we first formally describe the problem step for FL and then introduce the FedMRUR and the two novel hyperbolic graph fusion and normalized aggregation techniques in FedMRUR.
Figure 1: The workflow of FedMRUR. Once the global parameter  $x_0$  is received, the client initializes the local model  $x_p$  and starts hyperbolic graph fusion. In the hyperbolic graph fusion, the client first takes the local and global model to get their representations and maps them into the hyperbolic space. Then, the client use their distances in the hyperbolic space as a regularizer to constrain model divergence. Next, the client performs local training and uploads the updates to the server. The server adopts the normalized scheme to aggregate the local updates and performs the global model update.
We consider collaboratively solving the stochastic non-convex optimization problem with  $P$  clients :
$$
\min  _ {w} f (w) := \frac {1}{P} \sum_ {p \in \mathcal {P}} f _ {p} (w), \text {w i t h} f _ {p} (w) := \mathbb {E} _ {z \sim D _ {p}} [ l (w, z) ], \tag {1}
$$
where  $w$  is the machine learning model parameter and  $z$  is a data sample following the specified distribution  $D_{p}$  in client  $p$ ; meanwhile  $l(w,z)$  represents the model loss function evaluated by data  $z$  with parameter  $w$ .  $f_{p}(w)$  and  $f(w)$  indicate the local and global loss function, respectively. The loss function  $l(w,z)$ ,  $f_{p}(w)$ , and  $f(w)$  are non-convex due to the complex machine learning model. The heterogeneity of distribution  $D_{p}$  causes the model inconsistency across the clients, which may degrade the performance of the FL framework.
Notations. We define some notations to describe the proposed method conveniently.  $\| \cdot \|$  denotes the spectral norm for a real symmetric matrix or  $L_{2}$  norm for a vector.  $\langle \cdot ,\cdot \rangle$  denotes the inner product of two vectors. For any nature  $a,b,a\wedge b$  and  $a\vee b$  denote min  $\{a,b\}$  and max  $\{a,b\}$ , respectively. The notation  $O(\cdot)$ ,  $\Theta (\cdot)$ , and  $\Omega (\cdot)$  are utilized to hide only absolute constants that don't depend on any problem parameter.
3.2 FedMRUR Algorithm
In this part, we describe our proposed FedMRUR algorithm (see Figure 1 and Algorithm 1) to mitigate the negative impacts of model inconsistency and improve performance. We add a manifold regularization term on the objective function to alleviate the model inconsistency. To eliminate the near-orthogonality of client updates, we design a new method to aggregate the local updates from the clients. Within one communication round, the server first broadcast the global model to the participating clients. During local training, the client takes the sampled data into the local and received global model and gets the representations. Then the client maps the representations into the hyperbolic space and computes their distance, which is used to measure the divergence between the local and global models. Next, the client adopts the distance as a manifold regular to constrain the model bias, achieving model fusion in the hyperbolic graph. After local training, the client uploads its local update to the server. The server aggregates the local update norms as the global update step
Algorithm 1 FedMRUR
1: Input: initial parameter  $w^0$ ; momentum coefficient  $\alpha$ ; perturbation radius  $\rho$ ; local interval  $K$ ; communication rounds  $T$ ; set of selected clients  $S_t$ ; global and local learning rate  $\eta_g, \eta_l$ .
2: Output: Global parameter  $w^t, \forall t \in T$ .
3: Initialization: Initialize  $\triangle^0 = 0$  and  $w^0$  as the global parameter at the server.
4: For  $t = 0,1,\ldots ,T - 1$  do:
5: The server broadcasts parameter  $w^t$  and global update  $\triangle^t$  to the selected clients  $S_t$ .
6: For client  $p \in S_t$  in parallel do:
7: client  $p$  initialize the local parameter as  $w_{p}^{t,0} = w^{t}$ .
8: For  $k = 0, \dots, K - 1$  do:
9:  $\widetilde{w}_p^{t,k} = w_p^{t,k} + \rho \frac{\nabla F_p(w_p^{t,k})}{\|\nabla F_p(w_p^{t,k})\|}$ .
10:  $v_{i}^{t,k + 1} = \alpha \nabla F_{p}(\tilde{w}_{p}^{t,k}) + (1 - \alpha)\triangle^{t}.$
11:  $w_{i}^{t,k + 1} = w_{i}^{t,k} - \eta_{l}v_{i}^{t,k + 1}$
12: End for.
13:  $\triangle_p^t = w_p^{t,K} - w_p^{t,0}$
14: End for
15: Aggregate  $\triangle^{t+1} = \frac{\sum_{p \in S_t} \|\triangle_p^t\|}{|S_t||\sum_{p \in S_t} \triangle_p^t||} \sum_{i \in S_t} \triangle_p^t$ .
16: Update global parameter  $w^{t + 1} = w^t - \eta_g \triangle^{t + 1}$ .
17: End for.
and normalizes the sum of the local updates as the global update direction. Utilizing the normalized aggregation scheme, the server can update the model with a larger step and improve the convergence.
Hyperbolic Graph Fusion. In FL, the Euclidean distances between parameters[11, 20] (or gradients[16, 42]) between the client and the server is utilized to correct the local training for alleviating the model inconsistency. However, the Euclidean distance between the model parameters can't correctly reflect the variation in functionality due to the complex structure of the modern machine learning model. The model inconsistency across the clients is still large, which impairs the performance of the FL framework. Since the most prevail machine learning model, neural network has a graph structure and the hyperbolic space exhibits minimal distortion in describing data with graph structure, the client maps the representations of the local and global model into the hyperbolic shallow space[29] and uses the squared Lorentzian distance[19] between the representations to measure the model inconsistency.
To eliminate the model inconsistency effectively, we adopt the hyperbolic graph fusion technique, adding the distance of representations in the hyperbolic space as a regularization term to the loss function. Then the original problem (1) can be reformulated as:
$$
\min  _ {w _ {0}} F \left(w _ {0}\right) = \frac {1}{P} \sum_ {p} \left[ f _ {p} \left(w _ {p}\right) + \gamma * R \left(w _ {p}, w _ {g}\right) \right], \quad s. t. w _ {g} = \frac {1}{P} \sum_ {p} w _ {p} \tag {2}
$$
where  $R(w_{p},w_{g})$  is the hyperbolic graph fusion regularization term, defined as:
$$
R \left(w _ {p}, w _ {g}\right) = \exp \left(\left\| L _ {p} - L _ {g} \right\| _ {\mathcal {L}} ^ {2} / \sigma\right), \left\| L _ {p} - L _ {g} \right\| _ {\mathcal {L}} ^ {2} = - 2 \beta - 2 \langle L _ {p}, L _ {g} \rangle_ {\mathcal {L}}. \tag {3}
$$
In (2) and (3),  $L_{p}$  and  $L_{g}$  are the mapped Lorentzian vectors corresponding to  $Z_{p}$  and  $Z_{g}$ , the representations from local model  $w_{p}$  and global model  $w_{g}$ .  $\gamma$  and  $\sigma$  are parameters to tune the impact of the model divergence on training process.  $\beta$  is the parameter of the Lorentz model and  $\langle x,y\rangle_{\mathcal{L}}$  denotes the Lorentzian scalar product defined as:
$$
\langle x, y \rangle_ {\mathcal {L}} = - x _ {0} \cdot y _ {0} + \sum_ {i = 1} ^ {d} x _ {i} \cdot y _ {i}, \tag {4}
$$
where  $x$  and  $y$  are  $d + 1$  dimensional mapped Lorentzian vectors. The new problem (2) can be divided into each client and client  $p$  uses its local optimizer to solve the following sub-problem:
$$
\min  _ {w _ {p}} F _ {p} \left(w _ {p}\right) = f _ {p} \left(w _ {p}\right) + \gamma * R \left(w _ {p}, w _ {g}\right). \tag {5}
$$
The regularization term  $R(w_{p},w_{g})$  has two benefits for local training: (1) It mitigates the local over-fitting by constraining the local representation to be closer to the global representation in the hyperbolic space (Lorentzian model); (2) It adopts representation distances in a low-dimensional hyperbolic space to measure model deviation, which can be more precisely and save computation.
(a) Vanilla Aggregation.
(b) Normalized Aggregation.
Figure 2: A toy schematic to compare the naive aggregation and normalized aggregation of the local updates, where the number of clients is 2 and the local intervals are set as 1. The solid line indicates the client's local update  $\triangle_{i}$ ,  $\theta_{i}$  is the angle between the local update and the global update, and the dotted line represents the clients' contribution on the global update. The red lines are the aggregated global update  $\triangle$ . The main difference is  $\|\triangle\|$ , the norm of the global update. When adopting the naive aggregation method, the global norm  $\|\triangle\| = \sum_{i} \|\triangle_{i}\| \cos \theta_{i}$ . When adopting the normalized aggregation method, the global norm  $\|\triangle\| = \sum_{i} \|\triangle_{i}\|$ . We can see that the norm of global update in the direct aggregation is less than the norm in the normalized aggregation, due to the fact that  $\cos \theta \leq 1$ .
Normalized Aggregation of Local Updates. According to [3], after a number of communication rounds, the cosine similarities of the local updates across clients are almost zero, which comes from the model(gradient) inconsistency between the server and client sides. In the vanilla aggregation way, the server takes the mean of local updates from participated clients as the global gradient. As shown in Figure 2(a), client  $i$  makes  $\| \triangle_i\| \cos \theta_i$  contribution on the aggregation result. When the data heterogeneous is significant, the cosine similarities across the local updates are small. Correspondingly, the cosine similarities  $\cos \theta_{i}$  between the clients and the server narrows down, and the global update norm shrinks which slows down the convergence. To alleviate the negative impact of near-orthogonal local updates, we propose a new normalized aggregation method to compute the global update. The direction of the global update can be acquired by normalizing the sum of the local updates and the result is the same as the one obtained by the vanilla way. For the norm, the server computes it by taking the average of norms of the received local updates. As shown in Figure 2(b), with the proposed normalized aggregation method, the client  $i$ 's contribution on the global update increases from  $\| \triangle_i\| \cos \theta_i$  to  $\| \triangle_i\|$ . Accordingly, the norm of the global update  $\| \triangle \|$  grows and accelerates the convergence.