SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Our proposed FedMRUR is characterized by Figure 1 and the detailed training process is summarized in Algorithm 1. Concretely, firstly, the server broadcasts the global parameter and updates it to the selected clients  $S_{t}$ . At the initialization stage of local training, client  $i$  utilizes the manifold regularization scheme to construct its own local loss function  $F_{i}$  with the received global parameter  $w^{t}$ . Then, client  $i$  adopts the Sharpness Aware Minimization (SAM) [7] optimizer to compute the gradient  $\widetilde{g}_{i}$  with data sampled randomly. The local updater  $v_{i}$  consists of the stochastic gradient  $\alpha \widetilde{g}_{i}$  and the momentum term  $(1 - \alpha)\triangle^{t}$ , the received global update from the last round. Client  $i$  applies  $v_{i}$  to perform multiple SGD and uploads the accumulated local update  $\triangle_{i}^{t}$  to the server. The server takes two steps to construct the global update: 1) aggregating and normalizing the accumulated local updates from the participated clients  $S_{t}$  as the direction of the global update; 2) averaging the norms of accumulated local updates as the norm of the global update. Finally, the server utilizes the constructed global update to perform one step SGD and get a new global parameter.
Remark 1. FedMRUR is on the top of MoFedSAM [31] due to its excellent performance and our method can also be integrated with other federated learning methods, including FedExp, FedCM, SCAFFOLD, FedDYN, etc., to improve the performance.
In this section, we provide the theoretical analysis of our proposed FedMRUR for general non-convex FL setting. Due to space limitations, the detailed proofs are placed in Appendix. Before introducing the convergence results, we first state some commonly used assumptions as follows.
Assumption 1.  $f_{p}(x)$  is  $L$ -smooth and  $R(x,x_0)$  is  $r$ -smooth with fixed  $x_0$  for all client  $p$ , i.e.,
$$
\left\| \nabla f _ {p} (a) - \nabla f _ {p} (b) \right\| \leq L \| a - b \|, \quad \left\| \nabla R (a, x _ {0}) - \nabla R (b, x _ {0}) \right\| \leq r \| a - b \|.
$$
Assumption 2. The stochastic gradient  $g_{p}^{t,k}$  with the randomly sampled data on the local client  $p$  is an unbiased estimator of  $\nabla F_{p}(x_{p}^{t,k})$  with bounded variance, i.e.,
$$
E [ g _ {p} ^ {t, k} ] = \nabla F _ {p} (x _ {p} ^ {t, k}), E \left\| g _ {p} ^ {t, k} - \nabla F _ {p} (x _ {p} ^ {t, k}) \right\| ^ {2} \leq \sigma_ {l} ^ {2}
$$
Assumption 3. The dissimilarity of the dataset among the local clients is bounded by the local and global gradients, i.e.,
$$
E \left\| \nabla F _ {p} (x) - \nabla F (x) \right\| ^ {2} \leq \sigma_ {g} ^ {2} \tag {6}
$$
Assumption 1 guarantees the gradient Lipschitz continuity for the objective function and regularizer term. Assumption 2 guarantees the stochastic gradient is bounded by zero mean and constant variance. Assumption 3 gives the heterogeneity bound for the non-iid dataset across clients. All the above assumptions are widely used in many classical studies [1, 43, 33, 42, 14, 16], and our convergence analysis depends on them to study the properties of the proposed method.
Proof sketch. To explore the essential insights of the proposed FedMRUR, we first bound the client drift over all clients within the t-th communication round. Next, we characterize the global parameter moving within a communication round, which is similar to the one in centralized machine learning algorithms with momentum acceleration technology. Then, the upper bound for the global update  $\triangle_t$  is provided. Lastly, we use  $\| \nabla F(x_t)\|$  the global gradient norm as the metric of the convergence analysis of FedMRUR. The next theorem characterizes the convergence rate for FedMRUR.
Theorem 1. Let all the assumptions hold and with partial client participation. If  $\eta_{l} \leq \frac{1}{\sqrt{30}\alpha KL}$ ,  $\eta_{g} \leq \frac{S}{2\alpha L(S - 1)}$  satisfying  $\frac{3}{4} - \frac{2(1 - \alpha L)}{KN} - 70(1 - \alpha)K^{2}(L + r)^{2}\eta_{l}^{2} - \frac{90\alpha(L + r)^{3}\eta_{g}\eta_{l}^{2}}{S} - \frac{3\alpha(L + r)\eta_{g}}{2S}$ , then for all  $K \geq 0$  and  $T \geq 1$ , we have:
$$
\frac {1}{\sum_ {t = 1} ^ {T} d _ {t}} \sum_ {t = 1} ^ {T} \mathbb {E} \left\| \nabla F \left(w ^ {t}\right) \right\| ^ {2} d _ {t} \leq \frac {F ^ {0} - F ^ {*}}{C \alpha \eta_ {g} \sum_ {t = 1} ^ {T} d _ {t}} + \Phi , \tag {7}
$$
$$
\begin{array}{l} \Phi = \frac {1}{C} \left[ 1 0 \alpha^ {2} (L + r) ^ {4} \eta_ {l} ^ {2} \rho^ {2} \sigma_ {l} ^ {2} + 3 5 \alpha^ {2} K (L + r) ^ {2} \eta_ {l} ^ {2} 3 (\sigma_ {g} ^ {2} + 6 (L + r) ^ {2} \rho^ {2}) + 2 8 \alpha^ {2} K ^ {3} (L + r) ^ {6} \eta_ {l} ^ {4} \rho^ {2} \right] \\ + 2 K ^ {2} L ^ {4} \eta_ {l} ^ {2} \rho^ {2} + \frac {\alpha (L + r) ^ {3} \eta_ {g} ^ {2} \rho^ {2}}{2 K S} \sigma_ {l} ^ {2} + \frac {\alpha (L + r) \eta_ {g}}{K ^ {2} S N} (3 0 N K ^ {2} (L + r) ^ {4} \eta_ {l} ^ {2} \rho^ {2} \sigma_ {l} ^ {2}) \\ + 2 7 0 N K ^ {3} (L + r) ^ {2} \eta_ {l} ^ {2} \sigma_ {g} ^ {2} + 5 4 0 N K ^ {2} (L + r) ^ {4} \eta_ {l} ^ {2} \rho^ {2} + 7 2 K ^ {4} (L + r) ^ {6} \eta_ {l} ^ {4} \rho^ {2} \\ \left. \left. + 6 N K ^ {4} (L + r) ^ {2} \eta_ {l} ^ {2} \rho^ {2} + 4 N K ^ {2} \sigma_ {g} ^ {2} + 3 N K ^ {2} (L + r) ^ {2} \rho^ {2}\right) \right]. \\ \end{array}
$$
and  $d_{t} = \frac{\sum_{i}\left\|\triangle_{i}^{t}\right\|}{\left\|\sum_{i}\triangle_{i}^{t}\right\|} \geq 1$ . Specifically, we set  $\eta_{g} = \Theta \left(\frac{\sqrt{SK}}{\sqrt{T}}\right)$  and  $\eta_{l} = \Theta \left(\frac{1}{\sqrt{ST}K(L + r)}\right)$ , the convergence rate of the FedMRUR under partial client participation can be bounded as:
$$
\sum_ {t = 1} ^ {T} E \| \nabla F (x _ {t}) \| ^ {2} = O \left(\frac {1}{\sqrt {S K T}}\right) + O \left(\frac {\sqrt {K}}{S T}\right) + O \left(\frac {1}{\sqrt {K T}}\right). \tag {8}
$$
Remark 2. Compared with the inequality  $\frac{F^0 - F^*}{C\alpha\eta_gT} + \Phi$  of Theorem D.7 in MoFedSAM paper[31], the second constant term in (7) is same and the first term is less than the first term in MoFedSAM paper, which validates FedMRUR achieves faster convergence than MoFedSAM.
Remark 3. From (8), we can find that when  $T$  is large enough, the dominant term  $O\left(\frac{1}{\sqrt{SKT}}\right)$  in the bound achieves a linear speedup property with respect to the number of clients. It means that to achieve  $\epsilon$ -precision, there are  $O\left(\frac{1}{SK\epsilon^2}\right)$  communication rounds required at least for non-convex and  $L$ -smooth objective functions.
(a) Test accuracy on CIFAR-100 in the non-iid  $(\mu = 0.3, \mu = 0.6, n = 10$  and  $n = 20)$  settings.
(b) Test accuracy on TinyImageNet in the non-iid  $(\mu = 0.3, \mu = 0.6, n = 40$  and  $n = 80)$  settings.
Figure 3: Test accuracy w.r.t. communication rounds of our proposed method and other approaches. Each method performs in 1600 communication rounds. To compare them fairly, the basic optimizers are trained with the same hyperparameters.
In this section, we validate the effectiveness of the proposed FedMRUR algorithm using the experimental results on CIFAR-10/100 [18] and TinyImageNet [18]. We demonstrate that FedMRUR outperforms the vanilla FL baselines under heterogeneous settings. We also present that both manifold regularization and the proposed normalized update aggregation can improve the performance of SGD in FL. The experiments of CIFAR-10 are placed in the Appendix.
Datasets. We compare the performance of FL algorithms on CIFAR-10/100 and TinyImageNet datasets with 100 clients. The CIFAR-10 dataset consists of  $50K$  training images and  $10K$  testing images. All the images are with  $32 \times 32$  resolution belonging to 10 categories. In the CIFAR-100 dataset, there are 100 categories of images with the same format as CIFAR-10. TinyImageNet includes 200 categories of  $100K$  training images and  $10K$  testing images, whose resolutions are  $64 \times 64$ . For non-iid dataset partitioning over clients, we use Pathological- $n$  (abbreviated as  $\mathrm{Path}(n)$ ) and Dirichlet- $\mu$  (abbreviated as  $\mathrm{Dir}(\mu)$ ) sampling as [13], where the coefficient  $n$  is the number of data categories on each client and  $\mu$  measures the heterogeneity. In the experiments, we select the Dirichlet coefficient  $\mu$  from  $\{0.3, 0.6\}$  for all datasets and set the number of categories coefficient  $n$  from  $\{3, 6\}$  on CIFAR-10,  $\{10, 20\}$  on CIFAR-100 and  $\{40, 80\}$  on TinyImageNet.
Implementation Details. For all algorithms on all datasets, following [1, 42], the local and global learning rates are set as 0.1 and 1.0, the learning rate decay is set as 0.998 per communication round and the weight decay is set as  $5 \times 10^{-4}$ . ResNet-18 together with group normalization is adopted as the backbone to train the model. The clients' settings for different tasks are summarized in Table 1. Other optimizer hyperparameters are as follow:  $\rho = 0.5$  for SAM,  $\alpha = 0.1$  for client momentum,  $\gamma = 0.005$ ,  $\sigma = 10000.0$  and  $\beta = 1$  for manifold regularization.
Table 1: The experiments settings for different tasks.
Task	num of clients	participated ratio	batch size	local epoch
CIFAR	200	0.05	50	3
Tiny	500	0.02	20	2
Baselines. To compare the performances fairly, the random seeds are fixed. We compare the proposed FedMRUR with several competitive benchmarks: FedAvg [43], the most widely used baseline,
Table 2: Test accuracy (%) on CIFAR-100& TinyImagenet datasets in both  $\operatorname{Dir}\left( \mu \right)$  and  $\operatorname{Path}\left( n\right)$  ) distributions.
Algorithm	CIFAR-100	TinyImagenet
Dir(μ))	Path(n)	Dir(μ))	Path(n)
μ = 0.6	μ = 0.3	n = 20	n = 10	μ = 0.6	μ = 0.3	n = 80	n = 40
FedAvg	39.87	39.50	38.47	36.67	30.78	30.64	31.62	31.18
FedExp	44.51	44.26	43.58	41.00	33.49	32.68	33.65	33.39
FedProx	39.89	39.86	38.82	37.15	30.93	31.05	32.09	31.77
SCAFFOLD	47.51	46.47	46.23	42.45	37.14	36.22	37.48	35.32
FedCM	51.01	50.93	50.58	50.03	41.37	40.21	40.93	40.46
MoFedSAM	52.96	52.81	52.32	51.87	42.36	42.29	42.52	41.58
FedMRUR	55.81	55.49	55.21	53.69	45.54	45.41	45.42	45.71
Table 3: Convergence speed comparison on CIFAR100& TinyImageNet datasets. "Acc." represents the target test accuracy on the dataset. "∞" means that the algorithm is unable to achieve the target accuracy on the dataset.
Datasets	CIFAR-100	TinyImageNet
Algorithms	Acc.	Dir(μ)	Path(n)	Acc.	Dir(μ)	Path(n)
0.6	0.3	20	10	0.6	0.3	80	40
FedAvg	38%	513	494	655	∞	30%	972	1078	1002	1176
FedExp	715	782	795	1076	1255	1362	1327	1439
FedProx	480	488	638	∞	1043	1030	1163	1615
SCAFFOLD	301	322	389	585	785	850	766	967
FedCM	120	126	157	255	342	401	366	474
MoFedSAM	154	146	211	300	436	447	415	460
Our	157	179	223	341	473	517	470	570
FedAvg	42%	∞	∞	∞	∞	35%	∞	∞	∞	∞
FedExp	985	1144	1132	1382	∞	∞	∞	∞
FedProx	∞	∞	∞	∞	∞	∞	∞	∞
SCAFFOLD	406	449	558	998	1289	1444	1206	2064
FedCM	173	193	260	527	599	735	674	879
MoFedSAM	197	192	260	392	598	624	583	685
Our	192	230	266	424	671	707	653	788
FedAvg	45%	∞	∞	∞	∞	40%	∞	∞	∞	∞
FedExp	∞	∞	∞	∞	∞	∞	∞	∞
FedProx	∞	∞	∞	∞	∞	∞	∞	∞
SCAFFOLD	521	616	784	∞	∞	∞	∞	∞
FedCM	276	353	470	842	1451	2173	1587	2186
MoFedSAM	243	278	400	575	950	1106	959	1162
Our	241	263	338	484	948	1050	953	1069
firstly applies local multiple training and partial participation for FL framework; SCAFFOLD [16] utilizes the SVRG method to mitigate the client drift issue; FedProx [20] uses a proximal operator to tackle data heterogeneity; FedCM [42] incorporates the client-momentum term in local training to maintain the model consistency among clients; Based on FedCM, MoFedSAM [31] improves the generalization performance with local SAM [7] optimizer; FedExp [14] determines the server step size adaptively based on the local updates to achieve faster convergence.
5.2 Evaluation Results
Figure 3 and Table 2 demonstrate the performance of ResNet-18 trained using multiple algorithms on CIFAR-100 and TinyImageNet datasets under four heterogeneous settings. We plot the test accuracy of the algorithms for a simple image classification task in the figure. We can observe that: our proposed FedMRUR performs well with good stability and effectively alleviates the negative impact of the model inconsistency. Specifically, on the CIFAR100 dataset, FedMRUR achieves  $55.49\%$  on the Dirichlet-0.3 setups, which is  $5.07\%$  higher than the second-best test performance. FedMRUR effectively reduces the model inconsistency and the enlarged global update improves the speed of convergence.
Table 3 depicts the convergence speed of multiple algorithms. From [13], a larger  $\mu$  indicates less data heterogeneity across clients. We can observe that: 1) our proposed FedMRUR achieves the fastest convergence speed at most of the time, especially when the data heterogeneity is large. This
validates that FedMRUR can speed up iteration; 2) when the statistical heterogeneity is large, the proposed FedMRUR accelerates the convergence more effectively.
Impact of partial participation. Figures 4(a) and 4(b) depict the optimization performance of the proposed FedMRUR with different client participation rates on CIFAR-100, where the dataset splitting method is Dirichlet sampling with coefficient  $\mu = 0.3$  and the client participation ratios are chosen from 0.02 to 0.2. From this figure, we can observe that the client participation rate (PR) has a positive impact on the convergence speed, but the impact on test accuracy is little. Therefore, our method can work well under low PR settings especially when the communication resource is limited.
(c)
Figure 4: (a). Training loss w.r.t different client participation ratios; (b). Test accuracy w.r.t different client participation ratios. (c). Test accuracy with different  $\gamma$ . (d). Train loss with different  $\gamma$ . The performance of FedMRUR with different parameters on the CIFAR-100 dataset.
Hyperparameters Sensitivity. In Figures 4(c) and 4(d), we compare the performance of the proposed FedMRUR with different hyper-parameters on the CIFAR-100 dataset. From the results, we can see that our algorithm achieves similar test accuracy and training loss under different  $\gamma$  within a certain range ( $\gamma \in [0.005, 0.5]$ ) and this indicates the proposed FedMRUR is insensitive to the hyperparameter  $\gamma$ . The hyperparameter  $\gamma$  represents the method the punishment on the model bias.
Impact of Each Component. Table 4 demonstrates the impact of each component of FedM-RUR on the test accuracy for CIFAR-100 dataset on the Dirichlet-0.3 setups. For convenience, we abbreviate normalized aggregation as "normalized(N)" and hyperbolic graph fusion as "hyperbolic(H)", respectively. From the results, we can find that both the normalized local update
Table 4: Test accuracy % on CIFAR-100 datasets about without each ingredients of FedMRUR.
Algorithm	normalized(N)	hyperbolic(H)	Acc.
MoFedSAM	-	-	52.81
FedMRUR-N	✓	-	54.27
FedMRUR-H	-	✓	53.57
FedMRUR	✓	✓	55.49
aggregation scheme and the hyperbolic graph fusion can improve performance. This table validates that our algorithm design and theoretical analysis are correct and effective.
In this work, we propose a novel and practical federated method, dubbed FedMRUR which applies the hyperbolic graph fusion technique to alleviate the model inconsistency in the local training stage and utilizes normalized updates aggregation scheme to compensate for the global norm reduction due to the near-orthogonality of the local updates. We provide the theoretical analysis to guarantee its convergence and prove that FedMRUR achieves a linear-speedup property of  $O\left(\frac{1}{\sqrt{SKT}}\right)$ . We also conduct extensive experiments to validate the significant improvement and efficiency of our proposed FedMRUR, which is consistent with the properties of our analysis. This work inspires the FL framework design to focus on exploiting the manifold structure of the learning models.
Limitations&Broader Impacts. Our work focuses on the theory of federated optimization and proposes a novel FL algorithm. During the local training, the representations of the global model must be stored locally, which may bring extra pressure on the client. This will help us in inspiration for new algorithms. Since FL has wide applications in machine learning, Internet of Things, and UAV networks, our work may be useful in these areas.
Acknowledgements. This work is supported by National Key Research and Development Program of China under SQ2021YFC3300128, and National Natural Science Foundation of China under Grant 61971457. Thanks for the support from CENI-HEFEI and Laboratory for Future Networks in University of Science and Technology of China.