Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

TL;DR

This paper introduces SSF, a subspace optimization method for federated learning that reduces communication overhead while effectively handling non-IID data, achieving competitive convergence rates.

Contribution

The paper proposes a novel subspace optimization approach for federated learning that minimizes communication costs and manages data heterogeneity more efficiently than existing methods.

Findings

SSF achieves a convergence rate of (1/T+1/ff(NKT)) under standard assumptions.

Experiments demonstrate favorable accuracy-efficiency trade-offs with SSF on heterogeneous data.

SSF outperforms existing heterogeneity correction methods in communication efficiency.

Abstract

Federated learning increasingly operates in a large-model regime where communication, memory, and computation are all scarce. Typically, non-IID client data induce drift that degrades the stability and performance of local training. Existing remedies such as SCAFFOLD introduce heterogeneity-correction mechanisms to address this challenge, but they incur substantial extra communication and memory overhead. This paper proposes a subspace optimization method for federated learning (SSF), which performs heterogeneity-corrected optimization in a low-dimensional subspace using only projected quantities, while preserving full-dimensional control information through a backfill-style update that retains residual components whenever the active subspace changes. Under standard smoothness and bounded-variance assumptions, SSF attains a non-asymptotic rate of order $\widetilde{\mathcal{O}}(1/T+1/\sqrt{NKT})$. Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.