An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data

Jiaojiao Zhang; Yizhao Fan; Yuqi Xu; Kun Yuan

cs.LGSep 5, 2025

An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data

Authors: Jiaojiao Zhang, Yizhao Fan, Yuqi Xu, Kun Yuan

Organizations: Intelligent Computing Research Center, Great Bay University, Dongguan, China · Center for Machine Learning Research, Peking University, Beijing, China

Abstract

This work addresses the key challenges of applying federated learning to large-scale deep neural networks, particularly the issue of client drift due to data heterogeneity across clients and the high costs of communication, computation, and memory. We propose FedSub, an efficient subspace algorithm for federated learning on heterogeneous data. Specifically, FedSub utilizes subspace projection to guarantee local updates of each client within low-dimensional subspaces, thereby reducing communication, computation, and memory costs. Additionally, it incorporates low-dimensional dual variables to mitigate client drift. We provide convergence analysis that reveals the impact of key factors such as step size and subspace projection matrices on convergence. Experimental results demonstrate its efficiency.

Figures & tables

	FedSub (Algorithm 1 )	FedSub- $P^{k}=\mathrm{I}$	FedAvg	Scaffold
communication (uplink)	$rd$	$md$	$md$	$2md$
computation	$\tau(mrd+C_{g}(rd))+2mrd$	$\tau C_{g}(md)$	$\tau C_{g}(md)$	$\tau C_{g}(md)$
memory	$3rd+M_{g}(rd)+2rm+md$	$3md+M_{g}(md)$	$md+M_{g}(md)$	$3md+M_{g}(md)$

TABLE I : Assume a single layer.

\mathcal{O}

-notation is omitted for simplicity.

C_{g}(rd)

denotes the computation cost of evaluating a gradient with size

r\times d

M_{g}(rd)

denotes the memory cost of computing a gradient of dimension

r\times d

including intermediate quantities such as activations. As shown in Remark 1 , although FedSub-

P^{k}=\mathrm{I}

is mathematically equivalent to the well-known Scaffold [ 5 ] , Scaffold requires each client to send correction terms to the server, which increases communication overhead.

B_{i}^{k,t+1}=B_{i}^{k,t}-\eta\left(\frac{r}{m}(P^{k})^{T}\hat{\nabla}f_{i}(x^{k}+P^{k}B_{i}^{k,t})+\frac{\Lambda_{i}^{k}}{\eta\tau}\right)

Algorithm 1 Proposed Algorithm (FedSub)

Model	LLaMA-60M			LLaMA-350M			LLaMA-1B
Model	$r=$ 64	$r=$ 128	$r=$ 256	$r=$ 64	$r=$ 128	$r=$ 256	$r=$ 64	$r=$ 128	$r=$ 256
FedSub-CD	$\underset{\text{(2.72M)}}{\text{35.50M}}$	$\underset{\text{(5.44M)}}{\text{38.22M}}$	$\underset{\text{(10.88M)}}{\text{43.66M}}$	$\underset{\text{(16.27M)}}{\text{81.86M}}$	$\underset{\text{(32.54M)}}{\text{98.12M}}$	$\underset{\text{(65.08M)}}{\text{130.66M}}$	$\underset{\text{(32.50M)}}{\text{163.68M}}$	$\underset{\text{(65.01M)}}{\text{196.18M}}$	$\underset{\text{(130.02M)}}{\text{261.19M}}$
FedSub- $P^{k}=\mathrm{I}$	58.07M (25.30M)			367.97M (302.38M)			1339.08M (1207.91M)
FedAvg [ 1 ]	58.07M (58.07M)			367.97M (367.97M)			1339.08M (1339.08M)
FlexLoRA [ 12 ]	$\underset{\text{(5.00M)}}{\text{37.77M}}$	$\underset{\text{(9.99M)}}{\text{42.77M}}$	$\underset{\text{(19.99M)}}{\text{52.77M}}$	$\underset{\text{(29.91M)}}{\text{95.50M}}$	$\underset{\text{(59.82M)}}{\text{125.40M}}$	$\underset{\text{(119.64M)}}{\text{185.22M}}$	$\underset{\text{(59.77M)}}{\text{190.94M}}$	$\underset{\text{(119.53M)}}{\text{250.71M}}$	$\underset{\text{(239.07M)}}{\text{370.24M}}$
FedSVD [ 9 ]	$\underset{\text{(2.72M)}}{\text{35.50M}}$	$\underset{\text{(5.44M)}}{\text{38.22M}}$	$\underset{\text{(10.88M)}}{\text{43.66M}}$	$\underset{\text{(16.27M)}}{\text{81.86M}}$	$\underset{\text{(32.54M)}}{\text{98.12M}}$	$\underset{\text{(65.08M)}}{\text{130.66M}}$	$\underset{\text{(32.50M)}}{\text{163.68M}}$	$\underset{\text{(65.01M)}}{\text{196.18M}}$	$\underset{\text{(130.02M)}}{\text{261.19M}}$

TABLE II : NLP fine-tuning: This table shows the number of trainable parameters and the communicated parameters (in parentheses). FedSub-CD significantly reduces both total trainable parameters (e.g., LLaMA-1B with

r=256

drops to 261.19M vs FedAvg/FedSub-

P^{k}=\mathrm{I}

at 1339.08M) and communicated parameters (e.g., LLaMA-350M with

r=256

drops to 65.08M vs FedAvg at 367.97M).

Algorithm	Communication Volume (MB)
FedSub-CD64	9.673
FedSub-CD128	19.345
FedSub-CD256	38.690
FedSub- $P^{k}=\mathrm{I}$	392.028
FedAvg	494.326

TABLE III : Average communication volume per training step (i.e., the total communication volume divided by the total number of training steps

K\tau

) for LLaMA-1B fine-tuning.

Tasks	Batch Size	Sequence Length	$r$	$\tau$	Learing Rate $\eta$	Generation Freq	$P^{k}$
ResNet training on CIFAR-10/100	32	-	3	10	0.1	5	CD
ResNet (partial participation)	32	-	3	10	0.1	50	CD
Pre-training	64	1024	50%	10	0.001	5	SVD/CD
Fine-tuning	16	512	70%	5	0.0001	10	CD
Fine-tuning (runtime)	16	512	$\{64,128,256\}$	5	0.0001	10	CD
Fine-tuning (model scaling)	16	512	70%	5	0.0001	10	CD

TABLE IV : Hyperparameter settings for ResNet training, pre-training, and fine-tuning: To further reduce the computation overhead caused by updating the subspace projection matrix

P^{k}

and to enhance training stability, we adopt a periodic update strategy. The update interval is controlled by the parameter Generation Freq ; for example, Generation Freq = 10 means that a new

P^{k}

is randomly sampled every 10 communication rounds. For the subspace dimension

r

, two approaches are used: one is to set it as a fixed constant (e.g.,

r=3

); the other is to define it as a proportion of the original dimension. For instance, 70% means setting

r

to 70% of the original dimension.

Explore similar work

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data
Apr 28, 2026Shuchen Zhu, Zhengyang Huang, Yuqi Xu +1Heterogeneous Federated Learning Federated Learning
Peking University · Beihang University
FedSLoP: Memory-Efficient Federated Learning with Low-Rank Gradient Projection
Apr 27, 2026Yutong He, Zhengyang Huang, Jiahe Geng +1Federated Learning Fedavg
Peking University · Beihang University
Latent Information Sharing for Accelerating Federated Learning
Oct 1, 2026Seungjun Lee, Ensieh Khazaei, Dimitrios Hatzinakos +2Federated Learning Model Activations
Inha University · University of Toronto · University of Birmingham

Apr 28, 2026, cs.LGJ/K move · Enter open · S save

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

Shuchen Zhu, Zhengyang Huang, Yuqi Xu, Peijin Li

Peking University · Beihang University

Heterogeneous Federated Learning Federated Learning Heterogeneity Subspace

Open paper AskPDF 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

\widetilde{\mathcal{O}}(1/T+1/\sqrt{NKT})

. Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.

Abstract

Explore similar work

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

FedSLoP: Memory-Efficient Federated Learning with Low-Rank Gradient Projection

Latent Information Sharing for Accelerating Federated Learning