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

TL;DR

This work addresses the challenge of data heterogeneity in federated learning by proposing FedAWA, a server-side method that adaptively tunes aggregation weights using client vectors—differences between local and global models—without requiring proxy data. By optimizing the weights to align locally updated directions with a merged global direction, FedAWA enhances stability and generalization of the global model; a layer-wise variant FedAWA-L offers finer control at the cost of higher computation. The approach is validated across CIFAR-10/100 and Tiny-ImageNet with varying degrees of non-IID data, outperforming strong baselines including FedAvg, FedDisco, FedLAW, and L-DAWA, and showing robustness across architectures. The work highlights a practical, privacy-preserving mechanism for improving federated learning under heterogeneous data distributions with potential extensions to heterogeneous-model scenarios.

Abstract

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.

Figures (8)

• Figure 1: The illustration of client vectors. $\tau^t_k$ is the client vector of the $k$-th client, and $\tau^t_g$ represents the global vector obtained by aggregating the client vectors.
• Figure 2: Differences between client local datasets $\mathcal{D}_{1:K}$, client vectors $\tau_{1:K}^t$, and client models $\theta_{1:K}^t$.
• Figure 3: Distance with the ideal update vector $\tau^t_{ideal}$.
• Figure 4: Adjustment of the aggregation weights.
• Figure 5: Top-1 test accuracy (%) on different client numbers, sample rates, and local epochs.