Lazy Aggregation for Heterogeneous Federated Learning

Abstract

Federated learning (FL) is a distributed neural network training paradigm with privacy protection. With the premise of ensuring that local data isn’t leaked, multi-device cooperation trains the model and improves its normalization. Unlike centralized training, FL is susceptible to heterogeneous data, biased gradient estimations hinder convergence of the global model, and traditional sampling techniques cannot apply FL due to privacy constraints. Therefore, this paper proposes a novel FL framework, federated lazy aggregation (FedLA), which reduces aggregation frequency to obtain high-quality gradients and improve robustness in non-IID. To judge the aggregating timings, the change rate of the models’ weight divergence (WDR) is introduced to FL. Furthermore, the collected gradients also facilitate FL walking out of the saddle point without extra communications. The cross-device momentum (CDM) mechanism could significantly improve the upper limit performance of the global model in non-IID. We evaluate the performance of several popular algorithms, including FedLA and FedLA with momentum (FedLAM). The results show that FedLAM achieves the best performance in most scenarios and the performance of the global model can also be improved in IID scenarios.

1. Introduction

With the development of information technology, the amount of data generated by human beings is unprecedented. To mine the potential value of data, a series of emerging information technologies have been promoted (for example, Internet of Things, Big Data, and Data Mining). Nevertheless, it is tough for industry competition, monopoly, and data protection laws to achieve data sharing. Thus, it is necessary to discuss how to break the “data silos” to achieve secure data sharing. In 2016, Google proposed a privacy-constrained distributed neural network training paradigm named FL. Different from centralized training, FL aggregates the models’ weights of devices to integrate data knowledge of all parties without disclosing local data and improves the performance of the model on unknown data. FL realizes the efficient and safe sharing of data knowledge among multiple parties. FL has been widely explored and applied in data-sensitive fields such as medical treatment, personalized recommendation, and monetary risk assessment.

Unfortunately, there are some challenges in FL [1,2]. First of all, the open distributed architecture requires network transmission for synchronous training. In practice, the local training delay is much smaller than network transmission delay, especially in the case of large scale and big model. The network quality directly determines training efficiency. Secondly, the central server cannot access local data and data statistics of the participants, which restricts data preprocessing operations (e.g., data cleaning, deduplication, and enhancement). In addition, different distributions and data sparsity greatly weaken the final effect and efficiency, and even fail to converge, which has become a dark cloud over FL.

In order to resolve the issues of communication and data heterogeneity in FL, researchers have carried out a series of explorations. McManhan et al. [3] proposed the vanilla FL framework federated averaging (FedAvg), whose optimization focuses on increasing local epochs to reduce the synchronization frequency, and just some devices are chosen in each round. In Independent and Identically Distributed (IID), FedAvg per- forms very well and takes the first step of FL from theory to practice. However, FedAvg not only slows down the convergence speed but also greatly reduces the performance of the global model in non-IID. Li et al. [4] analyze theoretically convergence of FedAvg for the strongly convex and smooth problem, different data distributions expand the weight difference, and more epochs further exacerbate gradient offsets. Duan et al. [5] proved that it could also affect FL performance in the case of global unbalanced dataset, and they proposed to use a global shared and balanced dataset to alleviate gradient offsets caused by different distributions. As sharing proportion increased, it performs better, which maybe cause data leakage unfortunately. Zeng et al. [6] proposed a novel FL framework FedGSP, which counteracts perturbations for heterogeneous data by sequential execution. Sequential training is more robust to heterogeneous data than parallel training; however, the efficiency of sequential training is lower. For this reason, the author proposes the inter-cluster grouping (IGP), which is achieved inter-group homogeneity and intra-group heterogeneity. Grouping requires device characteristics, which has security risks.

In this paper, we propose a novel method for FL, which is shown in Figure 1. Federated lazy aggregating (FedLA) is used to address data heterogeneity. On the basis of FedAvg, FedLA allows more devices, which are from successional rounds, to train sequentially. FedLA doesn’t need extra informations about devices. As paper [6] said, finding out heterogeneous device groups is an NP-hard problem. Due to sequential training performs well under both homogeneous and heterogeneous data, just finding out sufficient sampling to aggregate in time. As for the aggregation timing, this paper proposes the change rate of the models’ weight divergence (WDR) to monitor sampling. Furthermore, we propose FedLA with Cross-Device Momentum (CDM) mechanism, FedLA with Momentum (FedLAM), which relieves vertical optimization jitters and helps local models jump out of local optima and release the performance of the global model in FL. In the end, we introduce several benchmark algorithms for comparison. The result shows that FedLAM achieves excellent performance in most scenarios. Compared with FedAvg, FedLAM promotes test accuracy by 3.1% on Mnist and is even higher by 7.3% in class 1. It has a comprehensive increase of 3.7% on Cifar10 and also has a 2.4% improvement on Emnist.

2. Related Work

At present, there are three main solutions to handle the FL problem in non-IID, including (1) adaptive optimization, (2) model personalization, and (3) data-based optimization.

2.1. Adaptive Optimization

FedAvg excels in IID but is not ideal in non-IID. To this end, there is a series of works to modify FedAvg to adapt to non-IID scenarios. Li et al. [7] eased the extreme point oscillations by dynamically reducing the local learning rate; moreover, a proximal term is added to local loss to limit the gradient offset. In the presence of straggler devices, aggregation stability is greatly improved. In general, it is similar to FedAvg. Shamir et al. [8] used multi-party distributed optimization to coordinate all parties to find the global optimal solution. However, FL cannot fully participate in each round. Karimireddy et al. [9] proposed a control state scheme, Scaffold which uses device-state and server-state control variables to coordinate the differences in training objectives of all parties, which greatly improves the convergence speed and upper bound of the global model, but its transmission overhead is twice as much as FedAvg. Durmus et al. [10] proposed a dynamic regularization method, FedDyn, by dynamically modifying local objectives to make them asymptotically consistent with the global objective and avoids transferring extra data. It should be noted that the method is very sensitive to learning rate.

2.2. Model Personalization

Just as its name implies, this kind of scheme allows device models to be personalized, and each model may be completely different. Although the personalized models are different, they still contain some global shared knowledge. Finn et al. [11] proposed a federated meta-learning framework, MAML which pretrains a meta-model using auxiliary dataset, and then fine-tunes the meta-model with local data. Fallah et al. proposed scheme [12] is to extract the feature extraction layers of pretraining model and fine-tune output layers to implement personalization. Arivazhagan et al. [13] proposed a personalized scheme which splits the model into two parts: base layer and personalized layer, and only upload weights of the base layer for aggregation in each round. It is personalized learning with the lower requirement for datasets. While the personalized model is excellent locally and not often satisfactory. Sattler et al. [14] proposed a fresh multitask framework, clustered federated learning (CFL), which introduced a recursive dual-partition mechanism based on FedAvg to separate heterogeneous devices and combine homogeneous devices to improve generalization and reduce gradient conflicts, but it is computationally inefficiency. Ghosh et al. [15] put forward an iterative federated clustering learning algorithm IFCA similar to EM. First, several group models are randomly initialized, and the group models are distributed to selected devices in each round, and selected devices choose the optimal weight base on local data. Finally, the server aggregates the gradients by group. The FedSem scheme of Xie et al. [16] is based on $l_2$ distance between device and group model weights. The above two schemes can improve grouping efficiency. The former needs to transmit more redundancy models, while the latter requires complex model initialization. Duan et al. [17] proposed a pre-grouping scheme, FedGroup, which uses decomposed cosine similarity instead of Euler distance to avoid the “Dimensional Curse” and improve the efficiency and stability of the cluster algorithms. Although CFL establishes the connections of devices to a certain extent, each group is independent. In reality, the relationships of devices are complex and diverse, and hard grouping will cut off connections between devices. Therefore, Li et al. [18] proposed a federated learning with soft clustering algorithm, FLSC based on IFCA. The devices of each group can overlap, and the updated weights can be shared by multiple groups. Its experiment shows that group models of FLSC have better performance than IFCA.

2.3. Data-Based Optimization

Data sharing [4,19,20] is a simple and effective scheme. FedShare [19] uses a global shared and balanced dataset $G$ to initialize the global model, distributes the global model and partial $G$ in each round, and devices use the distributed and local data to train jointly. Experiments show that only 5% of shared dataset makes the global model improve by 30% on Cifar10. It is difficult to obtain dataset such as this in practice. Data augmentation [21,22] can expand dataset through random transformation, knowledge transfer etc., which can effectively alleviate the impact of data sparsity. However, these methods need to acquire statistical informations, which will undoubtedly break the privacy constraint. Knowledge distillation [23] is a kind of promising solution for FL in non-IID, which transfers the knowledge of teacher model to student models with the help of auxiliary dataset. Lin et al. [24] used the ensemble distillation method to fuse multiple local models (teacher) to generate the next global model (student) and extract the knowledge of local models to ensure privacy, but the efficiency of knowledge distillation is relatively low.

In summary, these works have made improvements for heterogeneous federated learning from the data to model and system architecture. To mitigate the difference in data distribution and deal with the conflict gradient are the main goals. Our work focuses on the former, and continuously sampling multiple devices is to make sampling closer to global data distribution. Data sparsity is also very common on the device side, which makes the local model fall into local optimization. To this end, CDM is introduced. More importantly, the two-point optimization does not introduce extra communication overhead and FL privacy is effectively guaranteed, even reducing the server communication load.

3. Methods

3.1. Problem

FedAvg is the prototype of a series of FL algorithms whose theoretical basis is the federated optimization [25] problem, which is essentially a distributed optimization problem with privacy constraints. The distributed optimization goal of FL is:

$$\underset{w}{min}\left\{F\left(w\right)={\displaystyle \sum }_{k=1}^N{a}_k{L}_k\left(w\right)\right\}$$

(1)

where $N$ is the number of total devices and $p_k$ is the global weight of the k-th device $a_k\ge 0$, ${{\displaystyle \sum }}_{k=1}^N{a}_k=1$. Generally, $a_k=n_k/{{\displaystyle \sum }}_{k=1}^N{n}_k$; $n_k$ is the number of the k-th device samples. In general, $L_k\left(w\right)$ is widely used cross-entropy loss as the loss function of the classification task:

$$L_k\left(w\right)={\displaystyle \sum }_{i=1}^C{p}_k\left(y=i\right)E_{x|y=i}\left[log{f}_i\left(x,w\right)\right]$$

(2)

The local optimization goal is to find the weight $w_k$ with the minimum loss, and FL algorithms look for global optimal weight $w$ with the minimum total loss. To cope with the complex network environments, FedAvg only selects a subset $S_t$ of all devices $\left\{N\right\}$ in each round. Then, they execute training in parallel. Finally, the server aggregates recycled updates to renew the global model weight $w^{t+1}$ whose transfer expression is:

$$w^{t+1}=\frac{1}{{\sum }_{k\in S_T}{n}_k}{w}_k^{t+1}$$

(3)

When the local distributions of devices are consistent with global distribution (IID), there are few distinctions between partial and full training in each round, while the amount of computation and communication is greatly reduced. In reality, local data are sparse and distributed inconsistently, and gradient conflicts are prone to occur. The effect of gradient aggregation is poor [4], and it is easy to form a low-quality model.

3.2. FedLA

When the sampling is sufficient and IID, the gradients are unbiased estimates, and training in parallel can accelerate convergence of the global model. Otherwise, the gradient directions are divergent or even conflicting, which impacts aggregation [26]. It is known from the literature [6] that sequential SGD is more robust to heterogeneous data. An intuitive idea is to make devices train one by one. However, this way is inefficient, especially for large-scale FL. Moreover, sequential training has higher requirements for network connectivity, and too long sequential training would cause knowledge forgetting phenomenon. On the one hand, the efficiency of parallel training is higher in IID. On the other hand, sequential training is non-negligible robustness. We combine the two methods to form a mutual and allow homogeneous devices to train parallelly and heterogeneous devices to train sequentially. In FL settings, the server cannot obtain data distributions on the device-side, grouping by data distribution becomes an NP-Hard like knapsack problem. Therefore, Zeng et al. [6] proposed an IGP approximation to solve the problem, but the operation of acquiring device characteristics may cause data leakage. We consider that it isn’t necessary to judge whether devices are homogeneous or not, and more attention should be paid to how to sample fully. That is to say, the gradients of only one round of devices are often biased in non-IID, which makes aggregation unsatisfactory. Reducing the frequency of the server’s aggregation maybe a better choice.

The accuracy of gradient depends on whether data sampling is sufficient. In non-IID settings, data distribution is inconsistent with global distribution on the device- side, which is likely to cause gradient bias. The cross-device sequential training can sample more samples and gain high-quality gradients. Too long sequential training also leads to knowledge forgetting and is low efficiency. It is necessary to aggregate and cut off sequential chains periodically. To this end, we introduce the change rate of the models’ weight divergence (WDR) to judge the aggregation timing. Set WDR in the t-th round as:

$$WD{R}_t=\left|\frac{W{D}_t-W{D}_{t-1}}{W{D}_t}\right|$$

(4)

$W{D}_t$ is the models’ weight divergence (WD) in the t-th round:

$$W{D}_t=\frac{1}{K}{\displaystyle \sum }_{0\le i<j<K}\Vert w_t^i-w_t^j\Vert$$

(5)

$K$ is the number of selected devices per round. The model weight $W_t^i$ is updated for the selected i-th device after local training in the t-th round. In this way, the server recycles updated weights, it judges the timing of aggregation according to (4) calculated WDR. When the WDR is less than the threshold, the server aggregates and forms a new global model, otherwise, the server would randomly distribute collected and updated weights to the next round of devices. When the model weights are not aggregated, the models’ WD becomes larger and larger, but their WDR becomes smaller and smaller. Because, with the increase in continuous sampling, the updated directions of the models’ weights are gradually consistent with the global optimization direction.

Zhao et al. [19] have theoretically proved that the WD between FedAvg and vanilla SGD is related to the probabilistic distance of the local and global distribution. Since the global data distribution is unknown as usual, the WD focuses on the divergence of the models’ weights, which are to be aggregated in this paper.

$$\begin{array}{c}\Vert w_t^i-w_t^j\Vert \le \left(1+\eta \beta \right)\Vert w_{t-1}^i-w_{t-1}^j\Vert \\ +\eta g_{max}(w_{t-1}^j){\displaystyle {\displaystyle \sum }_{k=1}^C}\Vert p^i\left(y=k\right)-p^j\left(y=k\right)\Vert \end{array}$$

(6)

where $0<\beta ={{\displaystyle \sum }}_k^C{p}^i\left(y=k\right){\lambda }_{x|y=k}$ and $g_{max}(w_{t-1}^j)=\underset{0\le k<C}{max}{\nabla }_w{E}_{x|y=k}[log{f}_k(x, w_{t-1}^j)]$ and $\eta$ is the fixed learning rate (LR) of local SGD in our research. We can conclude that the WD mainly comes from the difference in the initial weights and sampling distributions. The former can be reduced by synchronous aggregation, which is the focus of previous works, but this paper has shifted its attention to the latter. For the aggregation, the result of multi-epochs training with a small LR isn’t different from an epoch training with a lager LR on the device side, and multiple devices with few samples train sequentially which is equivalent to a device training with a large number of samples. Therefore, we can generalize the conclusion of the inequality (6). The sampling range is extended from batch to device.

With sampling, the WD growth caused by distribution difference disappears gradually. Therefore, monitoring WDR is a great way to judge whether sampling is necessary. If WDR is greater than the threshold, the current round of weights won’t be aggregated but forwarded to the next round of selected devices. On the contrary, we consider that the sampling is sufficient. Only when their whole distributions are consistent with the global distribution will the WD growth slow down. On this basis, this paper proposes a federated lazy aggregate (FedLA) as shown in Algorithm 1:

Algorithm 1: Federated Lazy Aggregate (FedLA)
Input:$T, K, E, B, w_0, d_0, \epsilon$
Initialize $n_0^0,n_0^1\dots n_0^{k-1}\leftarrow 0; w_0^0,w_0^1\dots w_0^{K-1}\leftarrow w_0$
for t = 0, 1, 2, …, T − 1 do
	Sample devices $S_t\in \left[N\right], \left|S_t\right|=K\le N$
	Transmit $w_t^0 w_t^1, \dots , w_t^{K-1}$ to selected $K$ devices, respectively
	for each device k $\in {\mathrm{S}}_{\mathrm{t}}$ in parallel do
		$n_k, w_{t+1}^k=ClientUpdate\left(w_t^k, E, B\right)$
		Transmit device $n_k, w_{t+1}^k$ to server
	end for
	for k = 0, 1, 2, …, T − 1 do
		$n_{t+1}^k\leftarrow n_t^k+n_k$
	end for
	$d_{t+1}=WD\left(w_{t+1}^0, w_{t+1}^1, \dots w_{t+1}^{K-1}\right)$
	if$\left|\frac{d_{t+1}- d_t}{d_{t+1}}\right|\le \epsilon$ do
		$w_{t+1}\leftarrow \frac{1}{{{\displaystyle \sum }}_{k=0}^{K-1} n_{t+1}^{k }}{{\displaystyle \sum }}_{k=0}^{K-1} n_{t+1}^{k }{w}_{t+1}^k; w_{t+1}^0,w_{t+1}^1\dots w_{t+1}^{K-1}\leftarrow w_{t+1}$
		$n_{t+1}^0,n_{t+1}^1\dots n_{t+1}^{K-1 }\leftarrow 0; d_{t+1}=0$
	else
		$w_{t+1}\leftarrow w_t$
	end if
end for
Output:$w_T$

Where $E$ represents the number of local training epochs, $B$ represents the batch size for loading data, and $w^k$ represents the weight of k-th device which is initialized to the received weight in the t-th round. $T$ represents the number of total rounds; $w_0$ represents initial global weight; $d_0$ indicates the initial WD value, generally assign 0; $\epsilon$ is the threshold of WDR; and $n_t^k$ indicates the number of cumulative samples in the t-th round. FedLA adopts the same local solver as FedAvg, which optimizes the received weight by mini-batch SGD, and then the updated weights are sent back to server. Unlike FedAvg, aggregation doesn’t occur every round, but is determined according to WDR. The impact of heterogeneous data could be effectively mitigated through aggregation with intervals. In addition to adding WDR (computational complexity $\mathrm{O}({\mathrm{K}}^2\mathrm{d}$), where $d$ is the number of parameters in the model) and postponing aggregation, no other optimizations are introduced. Thus, FedLA has the same convergence rate $\mathrm{O}(1/\mathrm{T}$) as FedAvg [4] for strongly convex and smooth problems in non-IID. In practice, only classifier weights are needed for judgment, and the weights of feature layers are transmitted when aggregation is possible which ensures efficiency and reduces redundant aggregation. Moreover, it is compatible with other optimization and privacy strategies.

3.3. FedLAM

FedLA eliminates redundant aggregations and allows more devices to train sequentially, enabling more data to be sampled and alleviating data sparsity (e.g., few samples, class imbalance, and poor quality). Continuous training sequentially without aggregation occurs more jitters than vanilla FL on vertical optimization. It is more urgent for momentum mechanism, especially for cross-devices. It can correct current gradients by previous gradients and raise the utility of historical gradients. Firstly, we assign $K$ groups of gradient momentums on the server side, the initial value is zero gradient $\widehat{0}$ and its transfer equation is:

$$m_t^k\leftarrow \mu m_{t-1}^k+\mathsf{\Delta }{w}_t^k$$

(7)

where $\mu$ is the momentum coefficient, the bigger coefficient means that historical gradients are dominant and have a longer survival period; $m_t^k$ is the momentum of the k-th group in the t-th round; and $\mathsf{\Delta }{w}_t^k$ indicates the k-th updated gradient in the t-th round. Therefore, we propose the FedLA with momentum (FedLAM) that adds the CDM mechanism. Algorithm 2 is as follows:

Algorithm 2: Federated Lazy Aggregate with Momentum (FedLAM)
Input:$T, K, E, B, w_0, d_0, \epsilon , \mathit{\mu }$
Initialize ${\mathit{m}}_{\mathbf{0}}^{\mathbf{0}},{\mathit{m}}_{\mathbf{0}}^{\mathbf{1}}\dots {\mathit{m}}_{\mathbf{0}}^{\mathit{K}-\mathbf{1}}\leftarrow \widehat{\mathbf{0}}; n_0^0,n_0^1\dots n_0^{K-1}\leftarrow 0; w_0^0,w_0^1\dots w_0^{K-1}\leftarrow w_0$
for t = 0, 1, 2, …, T − 1 do
	Sample devices $S_t\in \left[N\right],\left|S_t\right|=K\le N$
	Transmit $w_t^0 w_t^1, \dots , w_t^{K-1}$ to selected $K$ devices, respectively
	for each device $k\in S_t$ in parallel do
		$n_k, \mathbf{\Delta }{\mathit{w}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}=ClientUpdate\left(w_t^k, E, B\right)$
		Transmit device $n_k, \mathbf{\Delta }{\mathit{w}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}$ to server
	end for
	for k = 0, 1, 2, …, T − 1 do
		$n_{t+1}^k\leftarrow n_t^k+n_k$
		${\mathit{m}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}\leftarrow \mathit{\mu }{\mathit{m}}_{\mathit{t}}^{\mathit{k}}+\mathbf{\Delta }{\mathit{w}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}$
		$w_{t+1}^k\leftarrow w_t^k+{\mathit{m}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}$
	end for
	$d_{t+1}=WD\left(w_{t+1}^0, w_{t+1}^1, \dots w_{t+1}^{K-1}\right)$
	if$\left|\frac{d_{t+1} - d_t}{d_{t+1}}\right|\le \epsilon$ do
		$w_{t+1}\leftarrow \frac{1}{{{\displaystyle \sum }}_{k=0}^{K-1} n_{t+1}^{k }}{{\displaystyle \sum }}_{k=0}^{K-1} n_{t+1}^{k }{w}_{t+1}^k; w_{t+1}^0,w_{t+1}^1\dots w_{t+1}^{K-1}\leftarrow w_{t+1}$
		${\mathit{m}}_{\mathit{t}+\mathbf{1}}\leftarrow \frac{\mathbf{1}}{{{\displaystyle \sum }}_{\mathit{k}=\mathbf{0}}^{\mathit{K}-\mathbf{1}} {\mathit{n}}_{\mathit{t}+\mathbf{1}}^{\mathit{k} }}{{\displaystyle \sum }}_{\mathit{k}=\mathbf{0}}^{\mathit{K}-\mathbf{1}} {\mathit{n}}_{\mathit{t}+\mathbf{1}}^{\mathit{k} }{\mathit{m}}_{\mathit{t}+\mathbf{1}}^{\mathit{k}}; {\mathit{m}}_{\mathbf{0}}^{\mathbf{0}},{\mathit{m}}_{\mathbf{0}}^{\mathbf{1}}\dots {\mathit{m}}_{\mathbf{0}}^{\mathit{K}-\mathbf{1}}\leftarrow {\mathit{m}}_{\mathit{t}+\mathbf{1}}$(*)
		$n_{t+1}^0,n_{t+1}^1\dots n_{t+1}^{K-1 }\leftarrow 0; d_{t+1}=0$
	else
		$w_{t+1}\leftarrow w_t$
	end if
end for
Output:$w_T$

It is noteworthy that the selected devices send their gradients instead of weights, and the weights of each group are updated by corrected gradients which enhance the horizontal optimization and offset the vertical optimization jitters. In non-IID settings, the local models fall easily into the local optimum [7], which is further passed to the global model through their gradients [27]. The step of momentum aggregation is optional in Algorithm 2(*). It is suggested to train early and close it later to expand the search optimization space, which contributes to the walking out of the saddle point. CDM could help FL alleviate gradient disappearance and take the global performance to a higher level. Compared with [28], our scheme doesn’t pass momentum on the device, which undoubtedly reduces the traffic between the device side and server side.

4. Evaluation

In this section, we will evaluate the actual performance of FedLA and FedLAM. First, we implement the above and several benchmark algorithms based on PyTorch and Ray frameworks. Three public datasets are selected for testing. In addition, we set up IID and two kinds of non-IID scenarios (label distribution skew, [5] Dirichlet [29]) to simulate actual environments in FL.

4.1. Experiment Setup

To compare different benchmark algorithms simply and fairly, this section selects three kinds of image classification datasets, Mnist [30], Emnist [31], and Cifar10 [32], and simulates IID and non-IID scenarios by partitioning the datasets. All experimental datasets are divided into 100 copies; each device holds one, and the number of selected devices K = 10. It should be noted that the local training doesn’t use momentum, weight decay options of SGD optimizer and gradient clipping technologies, and the constant learning rate is used by local solver. The datasets, partitions and models used in our experiments are shown in Table 1.

4.1.1. Datasets and Models

• Mnist, a widely used image classification task, contains a training dataset of 60,000 handwritten digit pictures and a test dataset of 10,000. Its elements are 28 × 28 pixel black and white pictures with 10 classification labels. Since the dataset is relatively simple and pure, and complex models (e.g., CNN, RNN) are little discriminated, for comparison, only single-layer MLP is used in this section.
• Emnist (Extend Mnist) expands 402,000 numbers and 411,000 26 letters samples on the basis of Mnist. Due to the huge amount of dataset, six types of splits (e.g., by class, by merge, letters etc.) are introduced. To distinguish from Mnist, the letters split is used, and a two-layer simple CNN is used for the test model.
• Cifar10 contains 60,000 32 × 32 pixel color images in 10 categories, and its classes are completely mutually exclusive (e.g., car and airplane), while the samples of the same category are quite different (e.g., car and truck). The difficulty of classification tasks has undoubtedly increased significantly, so we choose a three-layer convolutional layer and a two-layer fully connected layer CNN.

4.1.2. Data Partitions

As shown in Figure 2, data distributions of the clients are under three federated data partitions.

• IID: The dataset is randomly shuffled and divided into several subsets of the same size.
• Label distribution skew: Zhao et al. [5] proposed a more demanding non-IID setting, the outstanding feature of which is that each subset only holds a few classes of samples.
• Dirichlet: Wang et al. [29] proposed a partition scheme based on the Dirichlet distribution, where each subset is also class-imbalanced with different amounts of samples. It is closer to a partition scheme than the real FL environments.

Figure 2. Schematic diagram of class distribution of data partitions.

Figure 2. Schematic diagram of class distribution of data partitions.

Table 1. Summary of federated datasets and models.

Dataset	Samples	Classes	Devices	Partitions	Models
MNIST	69,035	10	100	IID Class 1/3 Dir-0.1/0.3	MLP	101,770
CIFAR10	60,000				CNN	553,514
EMNIST (letters)	145,600	26		Class 3/9	SimpleCNN	28,938

Table 1. Summary of federated datasets and models.

Dataset	Samples	Classes	Devices	Partitions	Models
MNIST	69,035	10	100	IID Class 1/3 Dir-0.1/0.3	MLP	101,770
CIFAR10	60,000				CNN	553,514
EMNIST (letters)	145,600	26		Class 3/9	SimpleCNN	28,938

4.1.3. Baselines

• FedAvg: The vanilla FL Algorithm [3].
• FedProx: A popular FL adaptive optimization scheme that limits excessive offset of the device-side model by adding a proximal term [7].
• FedDyn: A dynamic regularization method for FL, and local objectives are dynamically updated to ensure that local optimums are asymptotically consistent with the global optimum [10].
• CFL: A clustering federated learning framework based on recursive bi-partition that uses cosine similarity to separate the devices with gradient conflicts and form their respective cluster center models [14].

4.1.4. Evaluation Metric

In order to comprehensively evaluate the FL algorithms’ performances, we take top-1 classification accuracy of the devices’ test dataset as the main indicator. Note that all original training and test samples are merged into a new dataset which is divided into 100 copies, including 80% training dataset and 20% test dataset, which is to ensure that the training and test dataset is the same distribution on the devices. As it involves multi-party evaluation and local data are inaccessible, we use weighted accuracy on the test dataset of selected devices as verification accuracy and we take the weighted accuracy of all devices as test accuracy. There are several cluster center models in CFL, and it cannot be directly compared with other algorithms. To facilitate comparison, except for the weighted accuracies of cluster models, the pseudo-global model aggregated by all cluster models is introduced [17]. Note that this paper selects the benchmark algorithms without introducing extra synchronized parameters, so their communication traffics are at the same level.

4.2. Effects of the Proposed Algorithm

4.2.1. Performance

There are several points to explain: (1) The highest score of testing is taken after training with 300 rounds on Mnist and Emnist or 500 rounds on Cifar10. (2) The learning rates of the experiments are unified 0.002. (3) Other hyperparameters are as follows: FedProx: offset penalty coefficient u = 0.01; FedDyn: regularity coefficient α = 0.01; CFL: mean-norm threshold eps1 = 0.035, max-norm threshold eps2 = 0.5; FedLA: WDR threshold $\mathsf{\epsilon }$ = 0.02; and FedLAM: historical gradient momentum coefficient $\mu$ = 0.5. Differences from the above will be noted. We found that the FedLA and FedLAM algorithms perform well in most scenarios, as shown in

Table 2. Comparisons of FedAvg, FedProx, FedDyn, CFL, FedLA, and FedLAM on Mnist, Emnist, and Cifar10 (lr = 0.01) in IID and four kinds of non-IID settings. Local epoch E = 5, batch size B = 32.

		FedAvg	FedProx	FedDyn	FedLA	FedLAM	CFL
Mnist	IID	91.5	91.5	93.7	91.4	95.3	${91.6}_1$91.6
	Class 1	83.3	83.2	85.8	87.6	90.6	${\mathbf{93.4}}_{\mathbf{2}}$70.8
	Class 3	90.0	90.0	88.9	89.8	92.4	${90.0}_1$ 90.0
	Dir 0.1	90.0	90.0	88.7	88.8	91.5	${91.2}_2$88.8
	Dir 0.3	89.8	89.8	89.1	90.3	90.4	${89.8}_1$ 89.8
Emnist	IID	89.4	89.7	91.4	89.9	89.4	${89.9}_1$ 89.9
	Class 3	79.1	79.6	84.9	79.3	85.3	${79.2}_1$ 79.2
	Class 9	87.6	87.4	90.0	85.7	89.1	${87.6}_1$ 87.6
	Dir 0.1	84.3	84.9	88.5	84.2	87.8	${85.1}_1$ 85.1
	Dir 0.3	89.1	89.3	90.7	89.1	89.7	${89.1}_1$ 89.1
Cifar10	IID	70.8	69.6	71.5	70.4	76.3	${70.9}_1$ 70.9
	Class 1	27.0	27.0	41.8	${28.0}_{0.03}$	${30.0}_{0.07,0.2}$	${\mathbf{100.0}}_{\mathbf{10}}$ 17.1
	Class 3	65.3	66.2	64.1	67.3	71.1	${66.0}_1$ 66.0
	Dir 0.1	62.5	62.1	59.3	60.4	64.6	${62.7}_1$ 62.7
	Dir 0.3	66.0	66.1	65.6	68.1	71.5	${66.5}_1$ 66.5

, compared with the baseline algorithms.

On Mnist, FedLAM is ahead of other algorithms. Compared with FedAvg, it leads by 3.8% in IID, and leads comprehensive by 2.9% in non-IID, especially by 7.3% in class 1. For the more difficult task, Cifar10, it increases by 5.5% in IID and there is an overall increase of 3.2% in non-IID. In most cases, FedProx is very close to FedAvg. The method of adding a proximal term doesn’t overcome gradient conflicts caused by heterogeneous data but only reduces excessive gradient offset. As for the straggler devices mentioned in Section 5.2 [7], we will not explore this case in this article. FedDyn performs well on Emnist by aligning the local targets with the global target in real-time through an attached dynamic regularizer. As shown in Figure 3, FedDyn has a fast convergence speed in the early stage, but the accuracy of the global model is stagnant in the later stage, and its loss fluctuates significantly. The main reason is that it doesn’t adapt to the constant learning rate which needs to be adjusted flexibly. CFL separates the devices with gradient conflicts to build a flesh model and utilizes multiple models to adapt to devices with different distributions to improve comprehensive accuracy. For Cifar10 in class 1, the performances of other algorithms are not satisfactory, and CFL uses recursive dual partition to separate the different classes of devices and bring test accuracies to 100%. We also observed that the pseudo-global model has a poor effect, which verifies the great differences between classes. Although FedLA mitigates the gradient conflict caused by insufficient sampling in a way, it is limited by the aggregation method and it’s inability to fuse gradients well in Cifar10 (class 1). Gradient projection [33] technology solves the impact of opposing gradients. In non-IID scenarios, FedLA has a better performance than FedAvg. It reduces redundant aggregation and finds parallel groups of devices which are consistent with global distribution by computing WDR and overcomes the impact of data sparsity on a single device. Furthermore, the collected gradients could be used to correct vertical oscillation and walk out of the steep fall. FedLAM supplemented by CDM has the best performance in most cases, which is an exciting optimization.

4.2.2. WDR and Cross-Device Momentum

In order to explore the effectiveness of WDR and cross-device momentum, we added fixed-interval aggregations and different proportions of momentum based on Mnist (class 1).

As shown in Figure 4a,c, we sampled 2, 5, 10, and 20 devices for sequential training before aggregation. We found that sampling five devices is enough to obtain unbiased gradients in this experiment. Oversampling will not further improve the performance of the global model and reduces learning efficiency, or even has negative effects (knowledge forgetting, resource waste). According to inequality 3, WD mainly comes from the difference in the initial weights and sampling distribution. We prefer to sacrifice part of the former for the latter in non-IID. However, the number of sampling rounds is unknown. Fortunately, it could be observed that the growth rate of WD slowed down significantly after five rounds. Apart from the difference caused by initial weight, sampling distribution gradually tends to be consistent, and the second term of the inequality (6) approaches 0. This is why this paper proposes that WDR is used to find sufficient sampling in a timely manner. It can be used to adjust aggregated intervals dynamically, which makes the performance curve of the global model smoother and accelerates convergence.

In Figure 4b,d, we explored the influence of proportions of historical gradients in FL. In early training, a higher momentum will hinder the convergence of the global model, but it can bring FL performance to a higher level in the later stage. The gradients of devices will gradually disappear with FL training; nonetheless, the main reason is that the device model falls into the local optimum. Whether increasing sampled devices or the proportions of momentum, their gradient qualities are enhanced remarkably, which is undoubtedly helpful to the promotion of the global model. Inspired by the experiment, the momentum mechanism with dynamic adjustment deserves further exploration in the future. To sum up, FedLA determines aggregation timing by monitoring WDR, flexibly handles complex FL environments, and CDM mechanism improves FL performance in both IID and non-IID.

5. Conclusions

This paper studies the training problem with heterogeneous data in FL. Due to sparse data and different distributions of devices in FL, sampling only one round of devices is not enough and creates a poor model. Therefore, this paper proposes a novel FL framework, FedLA, which allows us to sample more devices by putting off aggregation, making aggregation more robust. We also note that local optimization is easier to fall into the saddle point. Thus, the cross-device momentum mechanism is added to FedLA to further release the performance of the training model in FL. Compared with benchmark algorithms (e.g., FedAvg, FedProx, FedDyn, and CFL), FedLAM has the best performance in most scenarios. In the future, we plan to conduct the theoretical analysis of FedLA in detail, and study more advanced aggregation and sampling judgment strategies. Moreover, the dynamic scheduling strategies of learning rate and momentum are introduced to further accelerate the convergence speed of FedLA.