Roughness-Informed Federated Learning

TL;DR

This paper introduces RI-FedAvg, a federated learning algorithm that uses a Roughness Index to adaptively regularize local updates, improving convergence and accuracy in non-IID, heterogeneous environments.

Contribution

The paper presents RI-FedAvg, a novel FL algorithm that incorporates a Roughness Index-based regularization to mitigate client drift and enhance robustness in non-IID settings.

Findings

Outperforms state-of-the-art FL baselines in accuracy.

Achieves faster convergence in non-IID scenarios.

Provides theoretical convergence guarantees for non-convex objectives.

Abstract

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.