Accelerating Heterogeneous Federated Learning with Closed-form Classifiers

TL;DR

This paper introduces Fed3R, a closed-form Ridge Regression classifier for federated learning that is robust to data heterogeneity, significantly reduces resource use, and improves training stability in cross-device scenarios.

Contribution

The paper presents Fed3R, a novel closed-form classifier for federated learning that addresses client heterogeneity and enhances efficiency and stability.

Findings

Fed3R is immune to data heterogeneity and client sampling order.

Fed3R requires up to 100 times fewer resources than competitors.

Using Fed3R as initialization stabilizes training and improves feature discrimination.

Abstract

Federated Learning (FL) methods often struggle in highly statistically heterogeneous settings. Indeed, non-IID data distributions cause client drift and biased local solutions, particularly pronounced in the final classification layer, negatively impacting convergence speed and accuracy. To address this issue, we introduce Federated Recursive Ridge Regression (Fed3R). Our method fits a Ridge Regression classifier computed in closed form leveraging pre-trained features. Fed3R is immune to statistical heterogeneity and is invariant to the sampling order of the clients. Therefore, it proves particularly effective in cross-device scenarios. Furthermore, it is fast and efficient in terms of communication and computation costs, requiring up to two orders of magnitude fewer resources than the competitors. Finally, we propose to leverage the Fed3R parameters as an initialization for a softmax classifier and subsequently fine-tune the model using any FL algorithm (Fed3R with Fine-Tuning, Fed3R+FT). Our findings also indicate that maintaining a fixed classifier aids in stabilizing the training and learning more discriminative features in cross-device settings. Official website: https://fed-3r.github.io/.