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
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- Pk=I
FedAvg
Scaffold
communication (uplink)
rd
md
md
2md
computation
τ(mrd+Cg(rd))+2mrd
τCg(md)
τCg(md)
τCg(md)
memory
3rd+Mg(rd)+2rm+md
3md+Mg(md)
md+Mg(md)
3md+Mg(md)
TABLE I : Assume a single layer. O -notation is omitted for simplicity. Cg(rd) denotes the computation cost of evaluating a gradient with size r×d . Mg(rd) denotes the memory cost of computing a gradient of dimension r×d including intermediate quantities such as activations. As shown in Remark 1 , although FedSub- Pk=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.
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- Pk=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- Pk=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τ ) for LLaMA-1B fine-tuning.
Tasks
Batch Size
Sequence Length
r
τ
Learing Rate η
Generation Freq
Pk
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 Pk 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 Pk 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.
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 O(1/T+1/NKT). Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.
Federated learning enables a population of clients to collaboratively train machine learning models without exchanging their raw data, but standard algorithms such as FedAvg suffer from slow convergence and high communication and memory costs in heterogeneous, resource-constrained environments. We introduce FedSLoP, a federated optimization algorithm that combines stochastic low-rank subspace projections of gradients, thereby reducing the dimension of communicated and stored updates while preserving optimization progress. On the theoretical side, we develop a detailed nonconvex convergence analysis under standard smoothness and bounded-variance assumptions, showing that FedSLoP is guaranteed to converge to a first-order stationary point at a rate of O(1/NT). On the empirical side, we conduct extensive experiments on federated MNIST classification with heterogeneous data partitions, showing that FedSLoP substantially reduces communication volume and client-side memory while achieving competitive or better accuracy compared with FedAvg and representative sparse or low-rank baselines. Together, our results demonstrate that random subspace momentum methods such as FedSLoP provide a principled and effective approach to communication- and memory-efficient federated learning. Codes are available at: https://github.com/pkumelon/FedSLoP.git.
Federated learning (FL) is a communication-efficient distributed learning paradigm. However, client drift remains one of the most critical challenges, hindering the efficient training of a global model. In this study, we propose a novel latent information sharing scheme that directly mitigates data heterogeneity across clients. Our theoretical and empirical results show that sharing a small amount of hidden-layer activations significantly improves training efficiency while preserving convergence guarantees and data privacy. Furthermore, we compare our method with existing FL approaches designed to address client drift, including FedProx, SCAFFOLD, FedPVR, FedProto, and SplitFed, and demonstrate superior model accuracy under a fixed round budget without incurring excessive communication overhead. Overall, this work presents a promising new knowledge aggregation scheme and provides a comprehensive analysis of the impact of activation sharing on federated optimization.