FedVG: Gradient-Guided Aggregation for Enhanced Federated Learning

TL;DR

FedVG is a novel gradient-based federated aggregation framework that leverages a global validation set to guide the optimization process, thereby enabling more informed and adaptive federated aggregation.

Abstract

Federated Learning (FL) enables collaborative model training across multiple clients without sharing their private data. However, data heterogeneity across clients leads to client drift, which degrades the overall generalization performance of the model. This effect is further compounded by overemphasis on poorly performing clients. To address this problem, we propose FedVG, a novel gradient-based federated aggregation framework that leverages a global validation set to guide the optimization process. Such a global validation set can be established using readily available public datasets, ensuring accessibility and consistency across clients without compromising privacy. In contrast to conventional approaches that prioritize client dataset volume, FedVG assesses the generalization ability of client models by measuring the magnitude of validation gradients across layers. Specifically, we compute layerwise gradient norms to derive a client-specific score that reflects how much each client needs to adjust for improved generalization on the global validation set, thereby enabling more informed and adaptive federated aggregation. Extensive experiments on both natural and medical image benchmarking datasets, across diverse model architectures, demonstrate that FedVG consistently improves performance, particularly in highly heterogeneous settings. Moreover, FedVG is modular and can be seamlessly integrated with various state-of-the-art FL algorithms, often further improving their results. Our code is available at https://github.com/alinadevkota/FedVG.

Figures (16)

• Figure 1: Comparison of federated learning methods under high data heterogeneity ($\alpha = 0.05$). FedVG consistently achieves the highest or near-highest accuracy across all settings.
• Figure 2: FedVG Framework. Locally trained client models, which converge to varying points, are sent to the server. A global validation set is used to compute their validation gradients, obtain client scores $s_k$, and form a global model by weighted sum. Individual client updates are represented as dotted lines in the loss landscape on the right, while aggregated models are solid.
• Figure 3: FL algorithm performance as $\alpha \rightarrow{} 0$. Shaded areas indicate standard deviations. Additional results for TinyImageNet and DermaMNIST datasets on the ResNet-50 model are provided in the Appendix.
• Figure 4: FL algorithm performance on ViT models as $\alpha \rightarrow{} 0$. Shaded areas show standard deviations.
• Figure 5: FedVG integration performance across CIFAR-10, OrganAMNIST, and COVID19 as $\alpha \rightarrow{} 0$. Dashed lines denote the base methods, and solid lines their FedVG-enhanced counterparts. Shaded gray areas show the standard deviation of base methods, while colored regions reflect the variability of enhanced methods.