PaperlayerView exact revision on arXiv

Source-linked AI summary

FedAWA: Adaptive Optimization of Aggregation Weights in Federated Learning Using Client Vectors

Changlong Shi, He Zhao, Bingjie Zhang, Mingyuan Zhou, Dandan Guo, Yi Chang

arXiv:2503.15842v1cs.LG

TL;DR

Data heterogeneity makes it difficult to choose aggregation weights that support effective federated training while preserving data security. FedAWA adapts server-side weights using client vectors without proxy datasets, and the authors report the best overall performance across varied datasets and heterogeneity settings. The conclusion identifies identical client-model architectures as a boundary of applicability.

  • Problem

    Data heterogeneity makes dataset-size weighting suboptimal, while adaptively adjusting aggregation weights without proxy data remains a challenge for privacy-preserving federated learning.

    Source: §1 · pp. 1–2

  • Method

    FedAWA uses client vectors derived from local model changes to adaptively optimize server-side aggregation weights without additional datasets.

    Source: §1 · p. 2

  • Results

    FedAWA achieves the overall best performance across various datasets and heterogeneity settings.

    Source: §5.2 · p. 7

  • Takeaways & Limitations

    Client vectors capture information about local datasets, and FedAWA offers adaptive aggregation without relying on a proxy dataset.

    Source: §1 · p. 2, §6 · p. 8

    The empirical exploration uses the global dataset, although practical federated clients do not share their data.

    Source: §4.2 · p. 4

Abstract

from arXiv · show

Federated Learning (FL) has emerged as a promising framework for distributed machine learning, enabling collaborative model training without sharing local data, thereby preserving privacy and enhancing security. However, data heterogeneity resulting from differences across user behaviors, preferences, and device characteristics poses a significant challenge for federated learning. Most previous works overlook the adjustment of aggregation weights, relying solely on dataset size for weight assignment, which often leads to unstable convergence and reduced model performance. Recently, several studies have sought to refine aggregation strategies by incorporating dataset characteristics and model alignment. However, adaptively adjusting aggregation weights while ensuring data security-without requiring additional proxy data-remains a significant challenge. In this work, we propose Federated learning with Adaptive Weight Aggregation (FedAWA), a novel method that adaptively adjusts aggregation weights based on client vectors during the learning process. The client vector captures the direction of model updates, reflecting local data variations, and is used to optimize the aggregation weight without requiring additional datasets or violating privacy. By assigning higher aggregation weights to local models whose updates align closely with the global optimization direction, FedAWA enhances the stability and generalization of the global model. Extensive experiments under diverse scenarios demonstrate the superiority of our method, providing a promising solution to the challenges of data heterogeneity in federated learning.

1. Introduction

Data heterogeneity makes local optimization directions diverge, challenging stable and effective aggregation. FedAWA uses client vectors to adapt aggregation weights during training without additional datasets.

1. Introduction

  • Heterogeneous client data can produce divergent local optimization directions, making convergence unstable and degrading overall model performance.FedAvg assigns aggregation weights by local dataset size, while privacy-preserving adaptive weighting remains a challenge.Source: §1 · pp. 1–2
  • Client vectors, formed from parameter changes after local training, capture local-data variation and aggregate in a direction closer to desired model updates.The paper investigates their relationship to local data and their alignment with the desired update direction.Source: §1 · p. 2
  • FedAWA dynamically optimizes server-side aggregation weights without additional datasets, assigning higher weights to client vectors aligned with the overall update direction.The method is intended to improve global-model generalization and training stability.Source: §1 · p. 2
  • The work connects model merging and federated learning by using client vectors to optimize aggregation weights in federated learning.FedAWA does not require additional data or transmission of original client data.Source: §1 · p. 2
  • The authors report experiments across diverse scenarios to demonstrate FedAWA's effectiveness.The supplied introduction does not specify the experimental outcomes.Source: §1 · p. 2

2. Related Works

Prior work addresses heterogeneity through client-side drift correction, server-side aggregation changes, and model merging. FedAWA focuses on adaptive server-side weight optimization without proxy data.

Federated Learning

  • Client-side methods such as FedProx and FedDyn [1] adjust local training, while FedAWA focuses on server-side aggregation and can be combined with these methods.The cited client-side approaches address local bias; the passage says they assign aggregation weights based on local dataset size.Source: §2 · p. 2
  • FedDisco [44] uses fixed weights based on dataset size and category-distribution discrepancy, whereas FedLAW [24] learns weights using a proxy dataset.L-DAWA [34] adaptively uses cosine similarity as weights; FedAWA instead optimizes weights and uses client and global model variations to reflect local data.Source: §2 · p. 2

Model Merging

  • Model merging combines trained models in a single step, but federated learning iteratively updates local and global models over multiple communication rounds.The paper identifies this distinction as a challenge for directly applying model-merging aggregation weights in federated learning.Source: §2 · pp. 2–3

3. Background

Federated learning iteratively aggregates client-trained models at a central server without sharing local data. The background contrasts dataset-size weighting with the task-vector perspective on model merging.

Federated Learning

  • Federated learning has K clients with private datasets and a central server that seeks to learn a global model without data sharing.In each communication round, the server broadcasts its model, clients train locally, and their models are returned for aggregation.Source: §3 · p. 3
  • FedAvg weights each client's model by its share of the total local dataset size during server aggregation.The background identifies this weighting as suboptimal under high data heterogeneity and frames adaptive, proxy-free weighting as the goal.Source: §3 · p. 3

Task Arithmetic

  • Task arithmetic defines a task vector as the difference between fine-tuned and pretrained parameters, then scales and sums vectors to form merged parameters.The cited passage contrasts this data-free, single-step merging with federated learning's iterative rounds.Source: §3 · p. 3
  • Figure 1 depicts client vectors and labels the aggregated client vectors' result as the global vector.The caption identifies the global vector as the vector obtained by aggregating client vectors.Source: Figure 1 · p. 3

4. Method

FedAWA uses client-update vectors to adapt server-side aggregation weights without additional datasets. Its design is motivated by evidence that client vectors reflect local-data differences and that their aggregation can guide weight adjustment.

4.1. Motivation

  • The method replaces task vectors with client vectors, defined as the difference between global and locally trained model parameters, to represent client-specific data variation.The paper hypothesizes that these vectors encode local-data information useful for optimizing aggregation weights while preserving privacy.Source: §4.1 · pp. 3–4

4.2. Empirical Observations

  • Client-vector differences closely resemble local-dataset differences, while distances between model parameters do not effectively reflect dataset relationships.The authors use this observation to motivate client vectors as a guide for aggregation.Source: §4.2 · p. 4
  • The paper treats the size-weighted merged vector as a proxy for overall data distribution and compares it with an ideal vector derived from training on the global dataset.The ideal vector represents the update direction of a model trained directly on the union of clients’ datasets.Source: §4.2 · p. 4

4.3. FedAWA

  • FedAWA weights clients by how closely their update directions align with the merged global vector, with a regularizer to limit deviation from the previous global model.The weights are optimized on the server and constrained to sum to one; the method uses 1 - cosine similarity as its distance function.Source: §4.3 · pp. 4–5
  • FedAWA-L learns separate aggregation weights for each model layer, enabling finer-grained adjustments than a single set of client weights.The method uses only information already available in the basic federated learning algorithm, without requiring additional data.Source: §4.3 · pp. 5–6

4.4. Discussions

  • FedAWA is presented as a server-side module that can complement client-side adjustments, while fixed server-side weights can serve as optimization starting values.The discussion also emphasizes that FedAWA does not require proxy data and uses model parameters and updates already accessible to the server.Source: §4.4 · p. 6

5. Experiments

Experiments compare FedAWA with seven baselines across image-classification datasets and non-IID settings. FedAWA performs best overall, improves several existing algorithms, and shows benefits across architectures and aggregation-weight analyses.

5.1. Experiment Setup

  • The experiments use CIFAR-10, CIFAR-100, and Tiny-ImageNet with shared model architectures per dataset, Dirichlet-partitioned client data, and seven representative baselines.Smaller Dirichlet α values represent greater non-IID heterogeneity.Source: §5.1 · pp. 6–7

5.2. Main Results

  • FedAWA achieves the overall best test accuracy across the evaluated datasets and heterogeneity settings, while FedAWA-L performs better at higher computational cost.The authors describe FedAWA as a more balanced performance–computation trade-off.Source: §5.2 · p. 7
  • Across varying data heterogeneity, FedAWA consistently improves FedAvg, FedProx, and FedDisco, including when FedDisco initializes its aggregation weights.Figure 6 summarizes these combinations across datasets and heterogeneity settings.Source: §5.2 · p. 7, Figure 6 · p. 7
  • FedAWA reduces aggregation execution time relative to L-DAWA and FedLAW, but takes slightly longer than FedAvg.The timing comparison uses ResNet20 on CIFAR-10; FedAWA-L incurs additional overhead.Source: §5.2 · p. 7, Table 2 · p. 7

5.3. Ablation Studies

  • FedAWA performs effectively across CNN, ResNet, WRN, DenseNet, and ViT architectures, including as network depth and width increase.Table 3 reports the comparisons across model architectures.Source: §5.3 · p. 7, Table 3 · p. 7
  • FedAWA achieves higher accuracy than FedAWA-COS, supporting iterative weight optimization; FedAWA-COS also outperforms L-DAWA across the compared datasets and heterogeneity levels.FedAWA-COS sets weights directly from client-vector/global-vector similarity, while L-DAWA uses model-parameter similarity.Source: §5.3 · p. 8, Table 4 · p. 8
  • FedAWA’s similarity to the dataset-vector ideal rises and converges during training, whereas FedAvg and FedDisco remain unchanged and L-DAWA and FedLAW fluctuate.Figure 7 compares aggregation-weight similarity with a dataset vector treated as ideal weights.Source: §5.3 · p. 8, Figure 7 · p. 8
  • Removing the regularization term decreases performance, and cosine-distance regularization generally performs slightly better than Euclidean-distance regularization.Table 5 compares FedAWA without regularization and with the two distance-based regularizers.Source: §5.3 · p. 8, Table 5 · p. 8

6. Conclusion

The conclusion presents FedAWA as adaptively optimizing aggregation weights without proxy data, with experiments reporting performance across settings. Its stated scope is limited to clients with identical model architectures.

6. Conclusion

  • FedAWA adaptively optimizes aggregation weights without a proxy dataset, and experiments show strong performance across varied settings.The authors report that client vectors capture local-dataset information and connect this observation to the method.Source: §6 · p. 8
  • FedAWA currently applies only when client model architectures are identical; accommodating heterogeneous architectures is identified as future work.The conclusion notes that this limitation is shared by many existing FL methods.Source: §6 · p. 8

A. Experiment details

This appendix section introduces the details of the experimental setup, environment, datasets, and model architectures used in the paper.

A. Experiment details

A.1. Client Vector and Local Data.

In a 12-client CIFAR-10 experiment with deliberately distinct class subsets, client-vector distances reflected local-data distribution similarities. Distances between full model parameters did not distinguish those relationships as effectively.

A.1. Client Vector and Local Data.

  • With 12 clients and equal-sized CIFAR-10 datasets, the experiment assigned clients different class subsets to create an extreme heterogeneity scenario.Clients 1–4 held the first five classes, clients 5–8 the last five, and clients 9–12 all classes.Source: §A.1 · p. 11
  • Client-vector distances resembled local-data distribution distances: clients with similar class distributions had smaller differences than clients with divergent distributions.The authors interpret this as evidence that client vectors capture relevant local-data information.Source: §A.1 · p. 11
  • Distances between full model parameters were relatively similar, making them less effective than client vectors for distinguishing local-dataset relationships.Both comparisons use 1 - cosine similarity, normalized to a range of 0 to 2.Source: §A.1 · p. 11

A.2. Aggregation Weights.

The section estimates how closely each client’s data distribution matches the global distribution, then uses those similarities to assess aggregation weights. FedAWA’s aggregation weights are compared with a dataset-derived reference vector.

A.2. Aggregation Weights.

  • The experiment measures local–global dataset similarity using ResNet20 features and Optimal Transport distance, then converts the distance to a similarity score.The local-data partition follows Appendix A.1, and the global dataset is the union of all local datasets.Source: §A.2 · p. 11
  • Similarity(P, Q) = 1 / (1 + d_OT(P, Q)) converts optimal transport distance into a similarity score.Source: §A.2 · p. 11
  • The resulting fixed k-dimensional dataset vector treats each local dataset’s similarity to the global dataset as its ideal aggregation weight.The authors compare methods’ aggregation weights with this vector using cosine similarity; Figure 7 is reported to demonstrate FedAWA’s effectiveness.Source: §A.2 · p. 11

A.3. Datasets

The appendix describes image-data partitions, further evaluations across heterogeneous settings, experimental hyperparameters, and model architectures. FedAWA is reported to improve performance in complex heterogeneous scenarios and to outperform baselines in text classification.

A.3. Datasets

  • The experiments use Dirichlet-sampled non-IID partitions, with smaller α producing greater non-IIDness and α = 100 treated as IID.Clients can differ in both class distributions and dataset sizes; Figure 8 illustrates distributions for some datasets, while CIFAR-100 and Tiny-ImageNet distributions are not shown.Source: §A.3 · pp. 11–12, Figure 8 · p. 12

A.4. Experiment

  • Additional evaluations combine FedAWA with CReFF for long-tail data and FPL for multi-domain data, comparing them with methods tailored to those settings.The comparisons include CReFF and CLIP2FL for long-tail data, and FPL and FedProto for multi-domain data.Source: §A.4 · pp. 12–13
  • FedAWA consistently improves performance in complex heterogeneous settings and outperforms baselines on AG News and Sogou News text classification.The evaluations cover long-tail and multi-domain settings as well as text classification; the corresponding tables are Tables 6–8.Source: Table 6 · p. 12, Table 7 · p. 12, Table 8 · p. 12

A.5. Hyperparameters

  • The default setup uses 20 clients, full participation, one local epoch, an initial learning rate of 0.08, and a 0.99-per-round learning-rate decay.The experiments also use local weight decay of 5e-4 and client-side SGD with momentum 0.9.Source: §A.5 · p. 13

A.6. Models

  • Models include ResNet20 for CIFAR-10 and CIFAR-100, ResNet18 for Tiny-ImageNet, and a CNN with three convolutional and two fully connected layers.The appendix also describes ResNet, WRN, DenseNet, and ViT architectures; WRN56-4 uses four times as many filters per layer.Source: §A.6 · p. 13
Loading 2503.15842v1…