Roughness-Informed Federated Learning

Federated Learning (FL) enables collaborative model training across distributed clients while preserving data privacy, yet faces challenges in non-independent and identically distributed (non-IID) settings due to client drift, which impairs convergence. We propose RI-FedAvg, a novel FL algorithm that mitigates client drift by incorporating a Roughness Index (RI)-based regularization term into the local objective, adaptively penalizing updates based on the fluctuations of local loss landscapes. This paper introduces RI-FedAvg, leveraging the RI to quantify the roughness of high-dimensional loss functions, ensuring robust optimization in heterogeneous settings. We provide a rigorous convergence analysis for non-convex objectives, establishing that RI-FedAvg converges to a stationary point under standard assumptions. Extensive experiments on MNIST, CIFAR-10, and CIFAR-100 demonstrate that RI-FedAvg outperforms state-of-the-art baselines, including FedAvg, FedProx, FedDyn, SCAFFOLD, and DP-FedAvg, achieving higher accuracy and faster convergence in non-IID scenarios. Our results highlight RI-FedAvg’s potential to enhance the robustness and efficiency of federated learning in practical, heterogeneous environments.

💡 Research Summary

The paper addresses a central challenge in federated learning (FL): the degradation of convergence and model performance when client data are non‑independent and identically distributed (non‑IID). In such settings, local updates drift away from the global optimum, a phenomenon known as client drift. Existing remedies—FedProx’s uniform proximal term, SCAFFOLD’s control variates, and DP‑FedAvg’s noise injection—either lack adaptivity, incur extra communication or memory overhead, or sacrifice accuracy for privacy.

To overcome these limitations, the authors introduce a novel metric called the Roughness Index (RI). RI quantifies the oscillatory behavior of a client’s loss landscape by projecting the high‑dimensional loss onto a set of random directions, sampling the loss along a small interval, computing the total variation of each one‑dimensional slice, normalizing by the slice’s amplitude, and finally taking the ratio of the standard deviation to the mean across directions. A high RI indicates a highly fluctuating loss surface, which is typically harder to optimize.

RI‑FedAvg incorporates this metric directly into the local objective. For client (k) at round (t) the regularized loss becomes
\