PaperlayerView exact revision on arXiv

Source-linked AI summary

Revisiting Weighted Aggregation in Federated Learning with Neural Networks

Zexi Li, Tao Lin, Xinyi Shang, Chao Wu

arXiv:2302.10911v4cs.LG

TL;DR

Federated learning commonly normalizes aggregation weights, but how their total magnitude and relative allocation affect generalization remains a question. This paper studies weight shrinking and client coherence, then proposes FedLAW, which improves global-model generalization across different datasets and models.

  • Problem

    The paper asks how aggregation weights should be assigned to generate a global DNN model with better generalization, challenging the convention that their sum must equal one.

    Source: §1 · p. 1

  • Method

    The authors study aggregation-weight magnitude γ and relative client weights λ, learning them separately on a proxy dataset to examine shrinking and client coherence.

    Source: §1 · pp. 1–2

  • Results

    FedLAW improves global-model generalization by a large margin across different datasets and models.

    Source: §7 · p. 9

  • Takeaways & Limitations

    Aggregation weights need not be normalized to sum to one: smaller sums can induce global weight shrinking that improves generalization.

    Source: §7 · p. 9

    The paper leaves adapting its server-side shrinking approach to minibatch SGD for future work.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

In federated learning (FL), weighted aggregation of local models is conducted to generate a global model, and the aggregation weights are normalized (the sum of weights is 1) and proportional to the local data sizes. In this paper, we revisit the weighted aggregation process and gain new insights into the training dynamics of FL. First, we find that the sum of weights can be smaller than 1, causing global weight shrinking effect (analogous to weight decay) and improving generalization. We explore how the optimal shrinking factor is affected by clients' data heterogeneity and local epochs. Second, we dive into the relative aggregation weights among clients to depict the clients' importance. We develop client coherence to study the learning dynamics and find a critical point that exists. Before entering the critical point, more coherent clients play more essential roles in generalization. Based on the above insights, we propose an effective method for Federated Learning with Learnable Aggregation Weights, named as FedLAW. Extensive experiments verify that our method can improve the generalization of the global model by a large margin on different datasets and models.

1. Introduction

The paper questions the convention that aggregation weights must sum to one, and examines how their total magnitude and relative allocation affect federated training. It connects weight shrinking and client coherence to generalization, then uses those insights to motivate FedLAW.

1. Introduction

  • The authors learn γ and relative client weights λ separately on a global-objective-consistent proxy dataset to study shrinking and client importance.They use these insights to propose FedLAW, reporting improved generalization across different datasets and models.Source: §1 · pp. 1–2
  • Smaller γ can improve generalization through global weight shrinking, which the paper relates to weight-decay regularization.The paper also notes that too-small γ can have negative effects, so an optimal value balances regularization and optimization.Source: §1 · pp. 1–2
  • A larger global-gradient norm calls for stronger shrinking when local epochs increase, clients are more IID, or training is before near-convergence.The paper identifies global-gradient magnitude as determining the optimal shrinking factor.Source: §1 · p. 2
  • Client coherence captures local-gradient similarity and heterogeneity coherence, which describes how sampled-client data distributions align with the global data.The paper reports a positive-to-negative critical point in local gradient coherence; generalization can benefit when it is positive and dominant.Source: §1 · p. 2
  • Reweighting sampled clients to increase heterogeneity coherence can improve training performance.Heterogeneity coherence concerns consistency between the global data and the distribution of a cohort of sampled clients.Source: §1 · p. 2

2. Related Works

Prior work learns aggregation weights for personalization or peer matching, while the paper focuses on generalization and uses learned weights to study federated training dynamics. It also relates its insights to centralized weight decay and gradient coherence.

2. Related Works

  • AUTO-FEDAVG learns weights on institutional medical data for personalization, and L2C learns weights to match peers in decentralized federated learning.Both use normalized aggregation weights, whereas this paper investigates generalization and the dynamics revealed by learned weights.Source: §2 · p. 2
  • FEDDF and FEDBE use ensemble distillation after weighted aggregation to improve global-model generalization, and the paper compares FEDLAW with them.Both distillation approaches require a server proxy dataset.Source: §2 · p. 2
  • The paper relates federated weight shrinking and client coherence to centralized weight decay and gradient coherence.In centralized learning, gradient coherence measures pairwise gradient similarity among samples; the cited work connects high similarity with stronger directional updates and better generalization.Source: §2 · p. 3

3. Preliminary and Problem Setup

The setup describes federated training with local client updates followed by server aggregation. It separates the total aggregation-weight magnitude from the relative weights and studies how client heterogeneity and local epochs affect training dynamics.

3. Preliminary and Problem Setup

  • Client heterogeneity is synthesized with Dirichlet sampling controlled by α, where smaller α means more NonIID data.The study treats α and local epochs E as key factors affecting federated training dynamics.Source: §3 · p. 3
  • Clients train locally from the global model using SGD, then the server samples clients and aggregates their local models or pseudo-gradients.The number of local epochs is E; the setup says larger E improves communication efficiency but makes updates more asynchronous.Source: §3 · p. 3
  • Vanilla FedAvg uses normalized aggregation weights proportional to client data sizes, with the global learning rate set to 1.The weights are denoted μ.Source: §3 · p. 3
  • The paper assumes aggregation weights need not sum to one and separates their l1 norm γ from relative client weights λ.It defines γ = ||μ||1 and λi = μi / ||μ||1.Source: §3 · p. 3
  • For studying γ and λ, the authors use a server proxy dataset matching the global objective, while separately validating FedLAW with tiny or biased proxies.Unless stated otherwise, the analysis in sections 4 and 5 uses CIFAR-10 and SimpleCNN.Source: §3 · p. 3

4. Global Weight Shrinking

The section finds that γ below 1 can regularize global aggregation, with an optimal value balancing regularization and optimization. It then studies adaptive γ through proxy-data learning and relates its behavior to training dynamics and generalization.

4.2. Adaptive Global Weight Shrinking and Training Dynamics

  • Adaptive GWS improves on FedAvg in both IID and NonIID settings, with greater benefit when clients perform fewer local epochs.Figure 2 reports the test-accuracy gains; the passage gives no numeric values.Source: Figure 2 · p. 4, §4.2 · p. 4

4.1. Global Weight Shrinking and Its Impacts on Optimization

  • When γ < 1, global weight shrinking regularizes the model; its strength increases as γ falls, and an optimal γ balances regularization against optimization.Moving either below or above the optimal value may reduce performance.Source: §4.1 · pp. 3–4

4.2. Adaptive Global Weight Shrinking and Training Dynamics

  • Adaptive GWS raises γ as the global gradient decays, keeping the gradient-term ratio near 19; stronger shrinking accompanies larger gradients, more local epochs, more IID data, and pre-convergence.Near convergence, smaller regularization is needed.Source: Figure 3 · p. 4, §4.2 · p. 4
  • Adaptive GWS leaves the DNN function similar under rescaling, yields nearly twice as many parameters close to zero as FedAvg, and has flatter curvature and smaller loss.The authors relate the flatter loss landscape to better generalization.Source: §4.2 · pp. 4–5, Figure 4 · p. 5
  • Adaptive GWS is compatible with local weight decay and can further improve performance; unlike local weight decay, it is hyperparameter-free and adapts γ.As local weight decay strengthens, the learned γ increases, weakening global shrinking.Source: §4.2 · p. 5

5. Client coherence

Client coherence clarifies why clients contribute differently across training: more coherent clients matter before a critical point, while adaptive global weight shrinking can sustain positive coherence afterward. Reweighting can also improve the distribution consistency of sampled cohorts with the global data.

5.1. Basic Concept and Formulation

  • Client coherence extends gradient coherence from mini-batch samples to FL clients; local gradient coherence measures similarity among clients’ local gradients or updates.The paper relates greater gradient similarity to stronger global gradients and improved generalization.Source: §5.1 · pp. 5–6
  • Attentive LAW learns relative aggregation weights on a proxy dataset; optimal weights can depart from data-size weighting when equal-sized clients differ in heterogeneity.The method fixes γ while learning λ to study clients’ roles in training dynamics.Source: §5.2 · p. 6
  • Before the critical point, attentive LAW assigns larger weights to clients with more balanced data; early-stopped attentive LAW performs comparably after that point.Figure 5 shows that balanced clients have higher coherence and receive nearly all weights in the first two rounds.Source: Figure 5 · p. 6, §5.2 · p. 7
  • After the critical point, adaptive GWS maintains positive local gradient coherence and further accuracy gains, unlike FedAvg’s negative coherence and marginal gains.Figure 6 compares their training dynamics.Source: §5.2 · p. 7, Figure 6 · p. 7
  • With partial participation, reweighting improves heterogeneity coherence by balancing sampled clients’ class distributions, and attentive LAW gains substantially over FedAvg.Attentive LAW with SWA performs better by seeking a more generalized minimum in the aggregation-weight hyperplane.Source: §5.2 · p. 7, Figure 7 · p. 7

6. FedLAW

FedLAW combines adaptive global weight shrinking with learnable relative client weights. It outperforms baselines across reported datasets, settings, client counts, and architectures, and remains effective under proxy-data variation and corrupted clients.

6.1. Method

  • FedLAW combines adaptive global weight shrinking and attentive learnable aggregation weights to optimize γ and λ; its optional SWA variant uses two-stage optimization.The two stages first optimize γ with λ fixed, then optimize λ with the learned γ fixed.Source: §6.1 · p. 8

6.2. Experiments

  • FedLAW outperforms baselines across datasets and models in IID and NonIID settings, scales to more clients and architectures, and performs well with partial participation.The reported comparisons include Tables 3–5 and Figure 8.Source: Table 4 · p. 8, Table 5 · p. 8, §6.2 · p. 9
  • Server-side baselines are sensitive to proxy-set size and can overfit when it is too small or too large, while FedLAW works with an extremely tiny proxy set and benefits from larger sets.Figure 8 compares participation ratios and proxy-dataset sizes.Source: §6.2 · p. 9, Figure 8 · p. 9
  • FedLAW remains effective when proxy and client data distributions differ and is robust to corrupted clients, matching ensemble-distillation methods using the same proxy dataset.Corruption is generated by swapping two labels in clients’ local training data.Source: §6.2 · p. 9

7. Conclusion

The paper revisits weighted aggregation in neural-network federated learning, identifying global weight shrinking and client coherence as insights into training dynamics. It proposes FedLAW and reports improved global-model generalization across datasets and models.

Appendix

  • Appendix experiments cover distribution shifts, corrupted clients, and additional fixed-γ results across architectures and NonIID settings.The distribution-shift tables vary whether client data or proxy data are long-tailed.Source: Table 6 · p. 9, Table 7 · p. 9, Table 8 · p. 9

7. Conclusion

  • FedLAW builds on global weight shrinking and two forms of client coherence, and experiments report substantially improved global-model generalization across datasets and models.The paper challenges the convention that aggregation weights must sum to 1.Source: §7 · p. 9

A.1. Model Aggregation in Federated Learning

  • The appendix provides additional experiments, FedLAW details, and setup information; its aggregation review focuses on convex combinations of clients’ local models.It contrasts prior layer-wise alignment approaches with the paper’s focus on conventional weighted aggregation.Source: Appendix · p. 13, §A.1 · p. 13

A.2. Generalization and Training Dynamics of Neural Networks

Prior work links neural-network loss-landscape curvature and early training dynamics to generalization, while heterogeneous federated data can make performance sensitive to tuning. The paper uses hyperparameter-optimization tools to examine federated training dynamics, rather than presenting aggregation-weight optimization as its main novelty.

A.2. Generalization and Training Dynamics of Neural Networks

  • Smaller top Hessian eigenvalues indicate flatter curvature, which prior work associates with better generalization.Prior approaches seek flatter minima by implicitly regularizing the Hessian.Source: §A.2 · p. 13
  • Prior work identifies critical early training periods in both neural networks and federated learning, while the paper examines client heterogeneity through local gradient coherence.The cited federated-learning study found that adding data later did not recover better accuracy after reducing data early.Source: §A.2 · p. 13

A.3. Federated Hyperparameter Optimization

  • Heterogeneous client data can lead to divergent local models and performance collapse, making careful hyperparameter tuning important.The passage names gradient-descent and reinforcement-learning approaches for hyperparameter optimization.Source: §A.3 · p. 13
  • The paper uses hyperparameter-optimization tools to study federated training dynamics; appendix tables cover fixed-γ architecture settings and adaptive GWS at different global learning rates.The authors state that optimizing aggregation weights is not their main novelty.Source: §A.3 · p. 14, Table 9 · p. 14, Table 10 · p. 14

B.1. Global Weight Shrinking

Adaptive global weight shrinking improves performance across heterogeneous settings, with the learned shrinking factor varying with data heterogeneity and gradient scale. The accompanying analyses associate adaptive shrinking with flatter loss landscapes than FedAvg.

B.1. Global Weight Shrinking

  • Adaptive GWS boosts performance across tested NonIID settings, though its benefit is smaller in the extremely NonIID setting (α = 0.1).The learned γ generally decreases as data become more IID, except for an outlier at α = 10; the authors relate this to increasing global-gradient norms and stronger regularization.Source: §B.1 · p. 14, Figure 9 · p. 15
  • Adaptive GWS produces smaller model parameters and flatter loss curvature during training, while FedAvg tends to generate sharper global models.The passage uses the top Hessian eigenvalue and the ratio of the top 1 and top 5 eigenvalues as flatness indicators.Source: §B.1 · p. 15, Figure 10 · p. 15
  • The ratio of the global gradient to the regularization pseudo-gradient across the reported experiments is approximately Gaussian, with mean around 20.5.The distribution aggregates experiments from Figures 2 and 9.Source: §B.1 · p. 15

B.2. Client Coherence

The client-coherence analysis finds that optimal aggregation weights vary with client heterogeneity, so dataset size alone does not capture client contribution. The appendix also reports FedLAW’s learning curves and convergence behavior.

B.2. Client Coherence

  • The appendix describes gradient diversity as contrasting with gradient coherence and notes that prior FL similarity work differs in perspective and focus.The cited discussion distinguishes bounded gradient dissimilarity and personalization from the paper’s focus on generalization dynamics and a critical point.Source: §B.2 · pp. 15–16
  • The reported r values across Figures 2 and 9 are approximated by a Gaussian distribution with mean around 20.5.Here r is the ratio of the global gradient to the regularization pseudo-gradient.Source: Figure 11 · p. 16
  • The three-client CIFAR-10 analysis finds FedAvg weights can be biased from the optimum, and optimal weights vary with clients’ heterogeneity.The optimal-weight distribution is described as Gaussian on the aggregation-weight plane, with its mean shifting toward more inter-heterogeneous directions.Source: §B.2 · p. 16, Figure 12 · p. 17
  • In more heterogeneous environments, the loss landscape is sharper, so bias from optimal aggregation weights is associated with a larger generalization drop.The analysis concludes that client contributions should account for heterogeneity as well as dataset size.Source: §B.2 · pp. 16–17
  • In the hybrid NonIID setting, the first 10 clients have class-balanced data and the last 10 have class-imbalanced data; the correlation analysis compares size and heterogeneity.The figure visualizes the hybrid setting with α1 = 10 and α2 = 0.1.Source: §B.2 · p. 17, Figure 13 · p. 18
  • Dataset size is more dominant in optimized aggregation weights when local epochs are large; otherwise, heterogeneity is more dominant.The combined size-and-heterogeneity indicator is more dominant when the two cases are mixed; Table 11 analyzes accumulated weights.Source: §B.2 · p. 18, Table 11 · p. 18

B.3. Learning curves of FedLAW and baselines

  • FedLAW surpasses baseline algorithms in most reported learning-curve cases and avoids late-training overfitting; γ converges in about 30 server epochs and λ in 80.The curves cover FashionMNIST, CIFAR-10, and CIFAR-100.Source: §B.3 · p. 18

C. Additional Details of FedLAW

FedLAW optimizes server-side aggregation weights under nonnegativity and unit-sum constraints, using an exponential parameterization in its experiments.

C. Additional Details of FedLAW

  • Server-side aggregation weights are constrained to be nonnegative and sum to one, with quadratic and exponential base functions offered to enforce those constraints.Source: C. Additional Details of FedLAW · p. 18
  • The variable x determines the aggregation weights, and its gradients are used to update them; experiments use the exponential function, which performs similarly overall to Softmax.Source: C. Additional Details of FedLAW · p. 18
  • Figure 14 visualizes server training for FedLAW on CIFAR10 with 20 clients, 3 local epochs, NonIID α = 1.0, and ResNet20.Source: Figure 14 · p. 19

D.1. Environment.

  • Experiments ran with Python 3.8.5, Pytorch 1.12.0, and four Quadro RTX 8000 GPUs.Source: §D.1 · p. 19

D.2. Data

The appendix specifies the data partitions and proxy sets, model configurations, repeatability controls, evaluation protocol, and training hyperparameters used in the experiments.

D.2. Data

  • NonIID client data are generated by Dirichlet sampling, with different client class distributions and dataset sizes; selected figures use a hybrid split mixing balanced and imbalanced clients.No data augmentation is used.Source: §D.2 · p. 19
  • For Table 3, the class-balanced proxy set uses 10 test samples per class, while the remaining test data evaluate global models.This gives 100 proxy samples for FashionMNIST and CIFAR-10 and 1,000 for CIFAR-100; other experiments use CIFAR-10 proxy sets of 100 or 1,000 samples.Source: §D.2 · p. 19

D.3. Model

  • CIFAR-10 and CIFAR-100 use a three-convolution, two-fully-connected-layer SimpleCNN, while FashionMNIST uses a three-layer MLP; ResNet and DenseNet follow Li et al. (2018).WRN56_4 denotes Wide-ResNet56-4, with four times as many filters per layer; Figure 15 plots FashionMNIST test-accuracy curves.Source: §D.3 · pp. 19–20, Figure 15 · p. 20

D.4. Randomness

  • Experiments were repeated with seeds 8, 9, and 10, using identical initial models and, for partial participation, the same sampled-client sequences across algorithms.Source: §D.4 · p. 20

D.5. Evaluation

  • Global models are evaluated on each dataset's test set over 200 rounds, with final accuracy averaged across the last 10 rounds; Figure 16 plots CIFAR-10 test-accuracy curves.Training indicators are averaged over rounds 90–110.Source: §D.5 · p. 20, Figure 16 · p. 21

D.6. Hyperparameter

  • Client SGD uses momentum 0.9, and learning rates decay by 0.99 each round; Figure 17 plots CIFAR-100 test-accuracy curves.Initial learning rates are 0.08 for CIFAR-10 and FashionMNIST and 0.01 for CIFAR-100; local weight decay is also used.Source: §D.6 · p. 21, Figure 17 · p. 22
Loading 2302.10911v4…