# Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

arXiv: https://arxiv.org/abs/2604.25467 · HTML: https://mlsanity.com/paper/2604.25467/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 34 references that are in the map.
- First posted on arXiv 28 Apr 2026.

## Builds on (most cited first, 30 of 30)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Federated Learning: Strategies for Improving Communication Efficiency](https://mlsanity.com/paper/1610.05492/) (2016)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Local SGD Converges Fast and Communicates Little](https://mlsanity.com/paper/1805.09767/) (2018)
- [Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training](https://mlsanity.com/paper/1712.01887/) (2017)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Gradient Sparsification for Communication-Efficient Distributed Optimization](https://mlsanity.com/paper/1710.09854/) (2017)
- [Achieving Linear Speedup with Partial Worker Participation in Non-IID Federated Learning](https://mlsanity.com/paper/2101.11203/) (2021)
- [Qsparse-local-SGD: Distributed SGD with Quantization, Sparsification, and Local Computations](https://mlsanity.com/paper/1906.02367/) (2019)
- [Stochastic Controlled Averaging for Federated Learning with Communication Compression](https://mlsanity.com/paper/2308.08165/) (2023)
- [DoubleSqueeze: Parallel Stochastic Gradient Descent with Double-Pass Error-Compensated Compression](https://mlsanity.com/paper/1905.05957/) (2019)
- [Federated Accelerated Stochastic Gradient Descent](https://mlsanity.com/paper/2006.08950/) (2020)
- [An Improved Analysis of Gradient Tracking for Decentralized Machine Learning](https://mlsanity.com/paper/2202.03836/) (2022)
- [LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statistics](https://mlsanity.com/paper/2410.16103/) (2024)
- [Decentralized Gradient Tracking with Local Steps](https://mlsanity.com/paper/2301.01313/) (2023)
- [Momentum Provably Improves Error Feedback!](https://mlsanity.com/paper/2305.15155/) (2023)
- [Heterogeneous Federated Learning](https://mlsanity.com/paper/2008.06767/) (2020)
- [FedMuon: Accelerating Federated Learning with Matrix Orthogonalization](https://mlsanity.com/paper/2510.27403/) (2025)
- [Fast Composite Optimization and Statistical Recovery in Federated Learning](https://mlsanity.com/paper/2207.08204/) (2022)
- [Greedy Low-Rank Gradient Compression for Distributed Learning with Convergence Guarantees](https://mlsanity.com/paper/2507.08784/) (2025)
- [GraVAC: Adaptive Compression for Communication-Efficient Distributed DL Training](https://mlsanity.com/paper/2305.12201/) (2023)
- [Beyond Local Sharpness: Communication-Efficient Global Sharpness-aware Minimization for Federated Learning](https://mlsanity.com/paper/2412.03752/) (2024)
- [Towards Bias Correction of FedAvg over Nonuniform and Time-Varying Communications](https://mlsanity.com/paper/2306.00280/) (2023)
- [A Unified Linear Speedup Analysis of Federated Averaging and Nesterov FedAvg](https://mlsanity.com/paper/2007.05690/) (2020)
- [Non-convex composite federated learning with heterogeneous data](https://mlsanity.com/paper/2502.03958/) (2025)
- [FedCanon: Non-Convex Composite Federated Learning with Efficient Proximal Operation on Heterogeneous Data](https://mlsanity.com/paper/2504.11903/) (2025)
- [An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data](https://mlsanity.com/paper/2509.05213/) (2025)
- [Widening the Network Mitigates the Impact of Data Heterogeneity on FedAvg](https://mlsanity.com/paper/2508.12576/) (2025)

## Built on by (most cited first, 1 of 1)

- [ReasFlow: Assisting Reasoning-Centric Scientific Discovery in Applied Mathematics via a Knowledge-Based Multi-Agent System](https://mlsanity.com/paper/2607.14178/) (2026)

## Abstract (arXiv)

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 widetildeO(1/T+1/√NKT). Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.