FedA4: Federated Learning with Anti-Bias Aggregation and TrAjectory-Based Adaptation

Non-Independent and Identically Distributed (Non-IID) data pose a fundamental challenge in Federated Learning (FL). It usually causes a severe client drift issue (various client model update directions) and thus, degrades the global model performance. Existing methods typically address this by assigning appropriate weights to client models or optimizing model update directions. However, these methods overlook client model update trends. They focus solely on the final client models to be aggregated at the server at each communication round, ignoring model optimization trajectories, which may contain richer information to aid model convergence. To address this issue, we propose FedA4, a novel FL framework with Anti-bias Aggregation and trAjectory-based Adaptation, which leverages clients’ optimization trajectories, rather than only their final model snapshots. For anti-bias aggregation, by observing a phenomenon termed model collapse, where biased clients tend to predict any input data as the dominant classes in their own datasets, we quantify the class dominance and analyze the level of client drift for each client. We evaluate a prediction entropy, namely concentration, so as to assign an optimal weight to each client at each training round. To further mitigate the negative effect of clients with high levels of client drift (biased clients), we then develop a gradient adaptation mechanism termed trajectory-based adaptation, which analyzes clients’ trajectories to correct each client’s contribution to the aggregated global model. Extensive experiments on CIFAR-10, CIFAR-100, STL-10, and Fashion-MNIST demonstrate that FedA4 significantly outperforms state-of-the-art baselines, particularly in scenarios with extreme data heterogeneity (high level of Non-IID).