Tackling Intertwined Data and Device Heterogeneities in Federated Learning with Unlimited Staleness

Most existing solutions to staleness in AFL are based on weighted aggregation Chen and Jin. (2019); Wang (2022); Chen (2020). For example, [7] suggests that a model update’s weight exponentially decays with its amount of staleness, and some others use different staleness metrics to decide model updates’ weights Wang (2022). Chen (2020) decides these weights based on a feature learning algorithm. These existing solutions, however, are improperly biased towards fast clients, and will hence affect the trained model’s accuracy when data and device heterogeneities in FL are intertwined, because they will miss important knowledge in slow clients’ model updates.

To show this, we conducted experiments using a real-world disaster image dataset Mouzannar et al. (2018), which contains 6k images of 5 disaster classes (e.g., fires and floods) with different levels of damage severity. In FL of 100 clients, we set data heterogeneity as that each client only contain samples in one data class, and set device heterogeneity as a staleness of 100 epochs on 15 clients with images of severe damage. When using this dataset to fine-tune a pre-trained ResNet18 model, results in Figure 1 show that staleness leads to large degradation of model accuracy, and weighted aggregation results in even lower accuracy than direct aggregation, because contributions from images of severe damage on stale clients are reduced by the weights¹¹1In synchronous FL, stale updates will be simply skipped, corresponding to applying zero weights on these updates. Hence, similar performance degradation is also expected for synchronous FL..

On the other hand, if we increase the contributions from stale clients by using larger weights, although the model accuracy on these images of severe damage will improve, the larger weights will amplify the impacts of errors contained in stale model updates and hence affect the model’s overall accuracy in other data classes. Detailed results can be found in Appendix B.

In practical scenarios such as natural disasters, such large or unlimited staleness is common due to interruptions in communication at disaster sites, and the staleness is too large for server to wait for any slow clients. The large performance degradation of weighted aggregation, then, motivates us to instead convert stale model updates to unstale ones.

The only existing work on such conversion, to our best knowledge, uses the first-order Taylor expansion to compensate for errors in stale model updates Zheng et al. (2017). For a stale update $g(w_{t-\tau})$, the estimated unstale update is:

$$g(w_{t})\approx g(w_{t-\tau})+\nabla g(w_{t-\tau})(w_{t}-w_{t-\tau}).$$

(1)

Since the Hessian matrix $\nabla g(w_{t-\tau})$ is difficult to compute for neural networks, it is approximated as

$$\nabla g(w_{t-\tau})\approx\lambda\cdot g(w_{t-\tau})\odot g(w_{t-\tau})$$

(2)

where $\lambda$ is an empirical hyper parameter. However, this method can only applies to small amounts of staleness Zhou et al. (2021); Li et al. (2023a); Tian et al. (2021), in which the high-order terms in the Taylor expansion can be negligible. To verify this, we use the same experiment setting as above and vary the amount of staleness from 0 to 50 epochs. As shown in Table 1, the error caused by high-order terms in Taylor expansion, measured in cosine distance and L1-norm difference with the unstale model updates, both significantly increase when staleness increases. These results motivate us to design techniques that ensure accurate conversion with unlimited staleness.

Our proposed approach builds on the existing techniques of gradient inversion Zhu and Han. (2019b), which recovers the original training data from the gradient of a trained model. Its basic idea is to minimize the difference between the trained model’s gradient and the gradient computed from the recovered data. Denote a batch of training data as $(x,y)$ where $x$ denotes input data and $y$ denotes labels, gradient inversion solves the following optimization problem:

$$(x^{\prime*},y^{\prime*})=\arg\min\nolimits_{(x^{\prime},y^{\prime})}\|\frac{\partial L[(x^{\prime},y^{\prime});w^{t-1}}{\partial w^{t-1}}-g^{t}\|^{2}_{2},$$

(3)

where $(x^{\prime},y^{\prime})$ is the recovered data, $w^{t-1}$ is the trained model, $L[\cdot]$ is model’s loss function, and $g^{t}$ is the gradient calculated with training data and $w^{t-1}$. This problem can be solved using gradient descent to iteratively update $(x^{\prime},y^{\prime})$.

The quality of recovered data relates to the amount of data samples recovered. Recovering a larger dataset will confuse the learned knowledge across different data samples and reduce the quality of recovered data, and existing methods are limited to recovering a small batch ($<$48) of data samples Yin (2021); Geiping (2020); Zhao and Bilen. (2020). This limitation, however, contradicts with the typical size of clients’ datasets in FL, which are usually more than hundreds of samples Wu et al. (2023); Reddi et al. (2020). This limitation indicates that we may utilize gradient inversion to estimate clients’ training data distributions without revealing individual samples of clients’ local data.

To compute $D_{rec}$, we first fix the size of $D_{rec}$ and randomly initialize each data sample and label in $D_{rec}$. Then, we iteratively update $D_{rec}$ by minimizing

$$Disparity[LocalUpdate(w^{t-\tau}_{global};D_{rec}),w^{t-\tau}_{i}],$$

(6)

using gradient descent, where $Disparity[\cdot]$ is a metric to evaluate how much $w^{t-\tau}_{i}$ changes if being retrained using $D_{rec}$. In FL, a client’s model update comprises multiple local training steps instead of a single gradient. Hence, to use gradient inversion in FL, we substitute the single gradient computed from $D_{rec}$ in Eq. (3) with the local training outcome using $D_{rec}$. In this way, since the loss surface in the model’s weight space computed using $D_{rec}$ is similar to that using $D_{i}$, we can expect a similar gradient being computed.

We first visualize it by using MNIST dataset to train LeNet model. Figure 3 shows that, the loss surface computed using $D_{rec}$ is similar to that using $D_{i}$ in the proximity of ($w^{t-\tau}_{global}$), and the computed gradient is very similar, too.

To verify the accuracy of using $\hat{w}^{t}_{i}$ to estimate $w^{t}_{i}$, we compare this estimation with first-order estimation, by computing their discrepancies with the true unstale model update under different amounts of staleness. Results in Figure 4 show that, compared to First-order CompensationZheng et al. (2017), our estimation based on gradient inversion can reduce the estimation error by up to 50%, especially when staleness excessively increases to more than 50 epochs.

Another key issue is how to decide the size of $D_{rec}$. Since gradient inversion is equivalent to data resampling in the original training data’s distribution, a sufficiently large size of $D_{rec}$ is necessary to ensure unbiased data sampling and sufficient minimization of gradient loss through iterations. On the other hand, when the size of $D_{rec}$ is too large, the computational overhead of each iteration would be unnecessarily too high. More details about how to decide the size of $D_{rec}$ are in Appendix D. Further results about our method’s error with various local training programs can also be found in Appendix E.

As shown in Figure 4, the estimation made by gradient inversion also contains errors, because the gradient inversion loss can not be reduced to zero. As the FL training progresses and the global model converges, the difference between the previous and current global models will reduce to 0, and hence the difference between stale and unstale model updates will also reduce, eventually to 0. In this case, in the late stage of FL training, the error in our estimated model update ($\hat{w}_{i}^{t}$) will exceed that of the original stale model update $w^{t-\tau}_{i}$.

To verify this, we conducted experiments by training the LeNet model with the MNIST dataset, and evaluated the average values of $E_{1}(t)=Disparity[\hat{w}^{t}_{i};w^{t}_{i}]$ and $E_{2}(t)=Disparity[w^{t-\tau}_{i};w^{t}_{i}]$ across different clients, using both cosine distance and L1-norm difference as the metric. Results in Figure 5 show that at the final stage of FL training, $E_{2}(t)$ is always larger than $E_{1}(t)$.

In practice, when we make such switch, the model accuracy in training will experience a sudden drop due to the inconsistency of gradients between $\hat{w}^{t}_{i}$ and $w^{t-\tau}_{i}$. To avoid such sudden drop, at time $t+\tau^{\prime}$, instead of immediately switching to using $\hat{w}^{t}_{i}$ in server’s model aggregation, we use a weighted average of $\gamma\hat{w}^{t}_{i}+(1-\gamma)w^{t-\tau}_{i}$ in aggregation, so as to ensure smooth switching. $\gamma$ linearly decays from 1 to 0 within a time window, and the length of this window can be flexibly adjusted to accommodate the optimization of model accuracy. Experiment results in Table 3 show that, when this length is set to 10% of training time before reaching the switching point, the model accuracy is maximized.

Our basic design rationale is to retain the clients’ FL procedure to be unchanged, and offload all the extra computations incurred by gradient inversion to the server. In this way, we can then focus on reducing the server’s computing cost of gradient inversion, which is caused by the large amount of iterations involved, using the following two methods.

First, we reduce complexity of the objective function in gradient inversion by sparsification, which only involve the important gradients with large magnitude into iterations of gradient inversion. Existing work has verified that gradients in mainstream models are highly sparse and only few gradients have large magnitudes Lin et al. (2017). Hence, we use a binary mask $Mask[\cdot]$ to selecting elements in $w^{t-\tau}_{i}$ with the top-$K$ magnitudes and only involve these elements to gradient inversion. As shown in Table 4, by only involving the top 5% of gradients, we can reduce around 80% of computation measured as the number of iterations in gradient inversion, with very slight increase in the error of estimating unstale model updates. Besides, we further explored the impact of such error caused by sparsification on the model accuracy, and results are in Appendix F.

Since in most cases the clients’ local data remains fixed, we do not need to start iterations of gradient inversion every time from a random initialization, but could instead optimize $D_{rec}$ from those calculated in the previous training epochs. Our experiments in Table 5 show that, when the clients’ local data remains fixed, we can further reduce the amount of iterations in gradient inversion by another 43%. Even if such client data is only partially fixed (e.g., changed by 20%), we can still achieve non-negligible reduction of such iterations.

Note that, we only apply gradient inversion to stale model updates containing unique knowledge not present in other model updates. Besides, most FL systems Charles et al. (2021) keep the number of clients in a global round constant. Once such number is sufficient (e.g., 10-50 even for FL with thousands of clients), further increasing such number yields little performance gains but increases overhead and causes catastrophic training failure Ro et al. (2022). Hence, the server’s overhead of gradient inversion, even in large-scale FL systems, will not largely increase. Such scalability is further discussed in Appendix G.

Although we used gradient inversion to estimate local data distributions from stale model updates, in most FL settings, it would be difficult or nearly impossible for the server to recover, either the stale clients’ local data or the labels, from the knowledge about such distributions, especially when applying the sparsification method described before.

However, even under the easiest scenario where client’s dataset only contains one sample and local training is just one-step gradient descent, such recovery will still be unsuccessful.

More specifically, although gradient inversion can recover the majority of data samples’ pixels as shown in Figure 8(a) when no sparsification is applied, the quality of such recovery quickly drops when moderate sparsification is applied, as shown in Figure 8(c) and 8(d). This is because sparsification effectively reduces the scope of knowledge available for gradient inversion to recover data. Results in Table 6 with multiple perceptual image quality metrics, including LPIPS Zhang (2018) and FID Heusel et al. (2017), further verify that such recovered images cannot be recognized in human eyes. Essentially, when 95% sparsification rate is applied, the quality of recovered images is similar to that of random noise. We also assessed the possibility of a neural network classifier (e.g., a ResNet-18 model) to recognize the recovered images. Results in the last row of Table 6 show that with the 95% sparsification rate, the classification accuracy is nearly equivalent to random guessing.

Besides, since our method only modifies the FL operations on the server and keeps other FL steps (e.g., the clients’ local model updates and client-server communication) unchanged, statistical privacy methods, such as differential privacy, can also be applied to local clients in our approach, just like how it applies to vanilla FL. Each client can independently add Gaussian noise to its local model updates, before sending the updates to the server Geyer et al. (2017). Similarly, it can also apply to our privacy protection method, by adding noise to the gradient after sparsification.

Gradient inversion should only be applied to stale clients when data and device heterogeneities are intertwined, i.e., the clients’ local data is unique and unavailable elsewhere. However, to properly decide such uniqueness, the server will need to know the class labels of client’s data, hence impairing the clients’ data privacy. Instead, we decide the data uniqueness by comparing the directions of stale clients’ model updates with the directions of other model updates from unstale clients, and only consider the stale clients’ data as unique if such difference is larger than a given threshold.

We quantify such difference between model updates $w_{i}^{t}$, $w_{j}^{t}$ from client $i$ and $j$ using cosine distance, such that

$$D_{c}(w_{i}^{t},w_{j}^{t})=1-w_{i}^{t}\cdot w_{j}^{t}\biggl{/}\|w_{i}^{t}\|\|w_{j}^{t}\|,$$

(7)

and the threshold is computed as the average of cosine distances between unstale model updates at $t-\tau$:

$$\frac{1}{\|S^{t-\tau}_{unstale}\|^{2}}\sum\nolimits_{j,k\in S^{t-\tau}_{unstale}}[D_{c}(w_{j}^{t-\tau},w_{k}^{t-\tau})]$$

(8)

, where $S^{t-\tau}_{unstale}$ is the set of unstale clients. Since the scale of cosine distance changes during FL training Li et al. (2023b), the average value of cosine distance adds adaptivity to the threshold.

We conducted preliminary experiments to evaluate if the server can accurately detect important model updates from unique client data. In the experiment, we emulate data heterogeneity by assigning each client with data samples from one random class, and results in Table 8 and Figure 9 show that the accuracy quickly grows to $>$90% as training progresses, and the average detection accuracy is 93%.

In all experiments, we consider a FL scenario with 100 clients. Each local model update on a client is trained by 5 epochs using the SGD optimizer, with a learning rate of 0.01 and momentum of 0.5.

We evaluate the FL performance by assessing the trained model’s accuracy in the selected data class being affected by staleness, and evaluate the FL training time in number of epochs. We expect that our approach can either improve the model accuracy, or achieve the similar accuracy with the baselines but use fewer training epochs.

In the fixed data scenario, 3 standard datasets and 1 domain-specific dataset are used in evaluations:

• •

  Using MNIST LeCun and Burges. (2010) and FMNIST Xiao et al. (2017) datasets to train a LeNet model, and data class 5 is affected by staleness;

• •

  Using CIFAR-10 Krizhevsky (2009) dataset to train a ResNet-18 model, data class 2 is affected by staleness;

• •

  Using a disaster image dataset MDI Mouzannar et al. (2018) to fine-tune the ResNet-18 model pre-trained with ImageNet.

The trained model’s accuracies⁴⁴4Compared to centralized training, FL models often exhibit lower accuracy, particularly under high data and device heterogeneity, as also reported in existing studies Morafah et al. (2024). using different FL schemes, with the amount of staleness as 40 epochs, are listed in Table 9. The training progresses of 1st-Order and W-Pred closely resemble that of Unweighted aggregation, suggesting that estimating stale model updates with the Taylor expansion is ineffective under unlimited staleness. Similarly, Weighted aggregation will lead to a biased model with much lower accuracy. In contrast, our gradient inversion based compensation can improve the trained model’s accuracy by at least 4%, compared to the best baseline. Such advantage in model accuracy can be as large as 25% when compared with Weighted aggregation. Besides image data, our method is also applicable to other data modalities such as text and time-series data. Results and discussions on these modalities with large real-world datasets are in Appendix A.

Figure 11 further show the FL training procedure over different epochs, and demonstrated that our method can also improve the progress and stability of training while also achieve higher model accuracy during different stages of FL training. Furthermore, we conducted experiments with different amounts of data and device heterogeneity. Results in Tables 10 and 11 show that⁵⁵5Training times in Tables 10-13 represent the relative training time required to reach convergence of the global model, assuming our method’s training time is 100%., compared with the baselines, our method can generally achieve higher model accuracy or reach the same accuracy with fewer training epochs, especially when the amount of staleness is large or the amount of data heterogeneity is high. We also use other large-scale real-world dataset to conduct experiments and results are in Appendix C.

To continuously vary the global data distributions, we use two public datasets, namely MNIST and SVHN Netzer (2011), which are for the same learning task but with different feature representations as shown in Figure 12. Each client’s local dataset is initialized as the MNIST dataset in the same way as in the fixed data scenario. Afterwards, during training, each client continuously replaces random data samples in its local dataset with new data samples in the SVHN dataset.

Experiment results in Figure 13 show that in such variant data scenario, since clients’ local data distributions continuously change, the FL training will never converge. Hence, the model accuracy achieved by the existing FL schemes exhibited significant fluctuations over time and stayed low ($<$40%). In comparison, our technique can better depict the variant data patterns and hence achieve much higher model accuracy, which is comparable to FL without staleness and 20% higher than those in existing FL schemes.

We also conducted experiments with different amounts of staleness and different rates of data distributions’ variations. We apply different variation rates of clients’ local data distributions by replacing different amounts of such random data samples in the clients’ local datasets in each epoch. To prevent the training from stopping too early when the variation rate is high, we repeatedly varied the data when the variation rate exceeded 1 sample per epoch. Results in Tables 12 and 13 showed that our method outperformed the baselines with different amounts of staleness. Weighted aggregation performs the worst since it leads the model to bias toward other unstable clients and other baselines show similar performance since they cannot compensate such large staleness.