An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data

arXiv:2509.05213 · cs.LG, cs.DC · Submitted 2025-09-05 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data".

Tom: This work introduces FedSub, an efficient subspace algorithm designed for federated learning on heterogeneous data, addressing challenges related to client drift and high communication/computation costs.

Jane: First, who's behind it and why it matters.

Paper summary: Tom: So Jane and I just read this paper today called "An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data," and honestly, it looks like they're tackling some seriously tough problems in federated learning right from the start. The main idea seems to be that they’re trying to handle client drift caused by data heterogeneity while also cutting down on the massive communication and computation costs associated with deep neural networks.

Jane: That sounds incredibly complex, Tom, but based on what I've read from the summary and abstract, it sounds like they propose FedSub as a solution that uses subspace projection to keep things efficient locally and then adds these low-dimensional dual variables to handle the drift issue. It seems like their core claim is achieving this efficiency while still mitigating those client drift problems.

Lu: From a research angle, the idea of restricting local updates to low-dimensional subspaces because model parameters often have low-rank structures during training makes a lot of sense for reducing complexity. It’s leveraging an empirical observation about DNNs to guide the mathematical structure of the algorithm itself.

Meng: I'm thinking practically, if we can reduce communication and computation costs by using these low-dimensional projections, that translates directly into faster training times and lower infrastructure needs for our clients, which is something I really care about as an engineer. But I wonder how stable those low-dimensional subspaces are when the data distributions are really messy.

Lalam: From my perspective as a language model, this approach suggests a cultural shift in how we deploy AI across different environments; if we can make training more efficient and robust to local variations, it means more diverse and high-quality AI models can be trained collaboratively without needing impossibly large centralized resources.

Tom: Exactly! So, to keep us grounded for a second, the paper introduces FedSub as this new algorithm that combines subspace projection for efficiency with low-dimensional dual variables to fight client drift. It’s all about finding a middle ground between having a full-space method and just using FedAvg when things get messy.

Jane: Right, and they set up the global objective by minimizing the average loss across all clients subject to that consensus constraint, which is standard in federated learning. They then introduce this subspace projection matrix P l which is smaller than the full space dimension, say R l times L l where R l is much smaller than M l.

Paper summary: Lu: That restriction to low-dimensional subspaces is a clever way to impose structure on the local updates, relying on the assumption that these parameters or gradients exhibit low-rank structures. It’s essentially saying, "We don't need to process every single parameter dimension at every step."

Meng: But then they have to define how this projection actually works in practice for a client, and that leads into the part about the local update rules. I need to know how much overhead that projection adds computationally before I can judge its real-world viability for a startup environment.

Lalam: The paper mentions that their algorithm is derived from the primal-dual hybrid gradient method but modified into this subspace variant, which suggests they're building on established optimization techniques while adapting them for this specific constraint. This shows how theoretical frameworks can be adapted to solve real-world distributed learning hurdles.

Tom: And that leads directly to the dual variables part, because they use these low-dimensional dual variables specifically to correct client drift caused by non-IID data distributions. That’s their other big mechanism for handling heterogeneity besides the subspace reduction.

Jane: And when we look at the efficiency analysis, they compare FedSub against full-space methods like FedAvg, and they show clear reductions in communication, computation, and memory usage compared to those baseline scenarios. Specifically, the uplink communication cost drops from O(md) to O(rd) when using the low-dimensional subspace variable B k i in Algorithm one.

Lu: That communication reduction is significant because it directly addresses one of the primary bottlenecks in large-scale federated learning setups where bandwidth can be a major constraint. It shows that the theoretical structure translates into tangible savings in data transfer volume.

Meng: I need to see those complexity numbers carefully; O(rd) versus O(md). If 'r' is small, that difference becomes substantial when we scale up to millions of parameters across many clients. How does the computation cost Cg(rd) compare to Cg(md)?

Lalam: The memory analysis also shows a reduction in memory overhead for FedSub, moving it from O(3rd + Mg(rd) + 2rm + md) down to something much smaller when P k = I, which is comparable to the full-space version. This hints at a very compact representation being achievable.

Paper summary: Tom: And it’s not just about efficiency; they also provide a convergence analysis with Theorem one which gives us bounds on the expected error E∇f(x k)two under certain step size conditions. This shows they’ve done their homework on how the learning process actually behaves.

Jane: The experimental validation is also pretty encouraging; they tested this on Logistic Regression and CIFAR-one hundred classification with ResNet, and the results show that even with subspace projections that might compromise accuracy compared to a full-space method, incorporating those dual variables for drift correction significantly outperforms FedAvg in terms of convergence accuracy.

Lu: That experimental finding is interesting because it suggests the trade-off between efficiency gained from the subspace restriction and the accuracy loss is manageable when you have those dual variables helping to keep things stable across heterogeneous data. It confirms their theoretical bounds in practice.

Meng: So, if we look at the bigger picture, this paper suggests that we can build more scalable federated systems that handle real-world data imperfections much better than what was previously possible with standard methods. It moves us closer to deploying AI on devices with limited resources more effectively.

Lalam: From my perspective, the implication is that we can deploy highly customized, robust AI systems across a much wider range of hardware and data sources without needing massive centralized retraining pipelines for every single client variation. This democratization of training power is significant for the future of distributed AI deployment.

Tom: It seems like the authors, Jiaojiao Zhang, Yuqi Xu, and Kun Yuan from Great Bay University and Peking University, have really put together a cohesive framework here in "An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data." The title itself sets up the problem perfectly.

Jane: And they clearly show that by combining subspace projection to reduce overhead with dual variables to manage drift, we get an algorithm that is both resource-friendly and actually performs well compared to simpler methods like FedAvg. It’s a balanced approach.

Lu: The work opens up new avenues for understanding the interplay between low-rank approximations inherent in neural network weights and the statistical challenges posed by non-IID data in a federated setting. It’s a solid piece of research that connects structural properties of models to distributed learning dynamics.

Paper summary: Meng: I'm still focused on the practical deployment implications; if this works well, it means we could deploy more complex AI models across more edge devices without overwhelming the network with parameter synchronization costs. It’s about making distributed training viable for real-world scenarios.

Lalam: The cultural impact I see is that this level of efficiency allows for faster iteration cycles in developing AI applications, meaning new, specialized models can be tested and deployed much quicker across different user groups. This speeds up the whole innovation pipeline.

Tom: So to wrap up the summary of "An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data," it’s FedSub, which uses subspace projection to save resources and dual variables to fix drift, proving efficiency without sacrificing much accuracy.

Jane: It really is a thoughtful approach, Tom, showing how you can tackle multiple challenges simultaneously in a distributed learning framework with targeted mathematical tools.

Lu: This paper provides valuable insights into how we can leverage the inherent structure of models while maintaining the necessary robustness against data skew in federated learning systems.

Meng: For us, it confirms that optimizing the architecture for efficiency is just as important as optimizing for raw accuracy when dealing with large-scale distributed training.

Lalam: I think the biggest cultural implication is showing that sophisticated, privacy-preserving AI training doesn't have to be prohibitively expensive or slow to implement on diverse, real-world data sets.

Tom: That’s a great way to put it; it’s about making advanced AI training accessible and practical for a broader range of applications and users.

Jane: And the convergence analysis provides the necessary mathematical backing to trust these efficiency gains, showing exactly under what conditions the algorithm is expected to work correctly.

Lu: It’s a good example of how theoretical constraints, like low-rank structures, can be effectively integrated into practical distributed learning algorithms to achieve better performance metrics.

Meng: I hope the practical implementation details we'll see in future work allow us to build systems that are truly scalable and deployable, not just theoretically sound.

Lalam: And ultimately, it gives us a blueprint for building more resilient and broadly applicable distributed AI systems that respect data constraints while maximizing collaborative training potential.

Conclusion: Tom: So, we’ve been diving deep into this paper called "An Efficient Subspace Algorithm for Federated Learning on Heterogeneous Data," and now it's time to wrap up how important this stuff really is for everyone listening.

Jane: We’ve seen how FedSub uses subspace projection and dual variables to keep things running smoother when clients have different data, so I think the title perfectly captures that balance between speed and accuracy.

Lu: I see the implication as a way to make large-scale AI training much more accessible because it leverages the mathematical structure of neural networks without needing massive centralized compute for every single client update.

Meng: From my side, what really strikes me is how this tackles the practical issue of communication costs; if we can genuinely reduce that overhead, it means we can deploy more complex models on smaller devices without crippling bandwidth.

Lalam: I think the cultural shift here is toward building AI systems that are inherently more robust to real-world data variation, which means less reliance on perfect, uniform datasets and more reliable deployment across diverse user groups.

Tom: Exactly! The authors of this paper, Jiaojiao Zhang, Yuqi Xu, and Kun Yuan from Great Bay University and Peking University, have delivered a framework that manages these complexities in a very structured way.

Jane: They’ve shown that by combining these two techniques—subspace restriction for efficiency and dual variables for drift correction—they can achieve strong convergence even when the data is messy.

Lu: The research really opens up new avenues for understanding how the low-rank properties of models interact with the statistical challenges of non-IID data in a federated setting, which is super creative stuff.

Meng: It’s fascinating to see that this isn't just theoretical; they have solid convergence analysis and experimental validation on real tasks like CIFAR-one hundred which gives us confidence that this approach works reliably in practice.

Lalam: I think the biggest vision here is fostering a culture where we prioritize building AI that is inherently resilient and adaptable to the messy reality of diverse data sources, rather than trying to force everything into a rigid standard.

Tom: So, it boils down to FedSub providing a blueprint for resource-efficient and robust distributed training that respects both the hardware constraints and the data diversity in modern AI applications.

Jane: It really is a thoughtful approach because it shows how you can tackle multiple challenges simultaneously in this complex distributed learning space without sacrificing the necessary performance metrics.

Lu: This work provides valuable insights into leveraging inherent model structure while maintaining robustness against data skew, which is a big piece of the puzzle for advanced AI research.

Meng: If we can scale up these efficient methods, it confirms that optimizing the architecture for resource usage is just as critical as optimizing for raw accuracy when dealing with large-scale distributed training.

Lalam: Ultimately, this paper gives us a blueprint for building more resilient and broadly applicable AI systems that respect data constraints while maximizing collaborative training potential across different environments.

Great Bay University · Peking University

cs.LG, cs.DC

Submitted: 2025-09-05

Updated: 2026-09-28

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 72/100

The gist: This work introduces FedSub, an efficient subspace algorithm designed for federated learning on heterogeneous data, addressing challenges related to client drift and high communication/computation

Key concepts

Subspace Projection: This technique restricts the local updates on each client to a much smaller, lower-dimensional subspace of the full parameter space. This is based on the idea that model parameters often have low-rank structures during training, allowing for significant reduction in complexity.
Dual Variables: These are auxiliary variables introduced into the optimization process to specifically counteract 'client drift,' which occurs when clients use non-identical data. They help steer the local updates toward a consensus that mitigates divergence caused by heterogeneity.
Communication Efficiency: FedSub reduces communication overhead by having clients only send a low-dimensional subspace variable (size r x d) to the server, instead of sending the full model parameters (size m x d). This lowers uplink data transfer costs significantly.
Client Drift Mitigation: Client drift happens when different clients train on different data, causing their local models to diverge. FedSub uses low-dimensional dual variables as a correction term in the update rule to actively counteract this divergence and maintain convergence.

Terminology

Federated Learning — 25× in this paper · explained in 14 episodes
Accuracy — 20× in this paper · explained in 1 episode
Data Heterogeneity — 7× in this paper · explained in 3 episodes
Deep Neural Networks — 5× in this paper
Non-IID Data — 5× in this paper · explained in 2 episodes
Fine-Tuning — 3× in this paper · explained in 5 episodes
Adam — 2× in this paper · explained in 2 episodes
Adaptation — 2× in this paper · explained in 1 episode
Convergence Guarantees — 2× in this paper · explained in 3 episodes
Core — 2× in this paper · explained in 2 episodes
Robustness — 2× in this paper · explained in 5 episodes
Coordinate Descent — 1× in this paper
Decentralized Learning — 1× in this paper · explained in 2 episodes
Edge Devices — 1× in this paper · explained in 2 episodes
Federated Learning (FL) — 1× in this paper · explained in 7 episodes
Loss Function — 1× in this paper · explained in 2 episodes
Regularization — 1× in this paper · explained in 3 episodes
Spherical Smoothing — 1× in this paper
Trace — 1× in this paper · explained in 2 episodes

Summary

This work introduces FedSub, an efficient subspace algorithm designed for federated learning on heterogeneous data, addressing challenges related to client drift and high communication/computation costs. The core contribution is utilizing low-dimensional subspace projections to simultaneously reduce communication, computation, and memory overhead while employing low-dimensional dual variables to mitigate client drift caused by non-IID data distributions.

The gist

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.

Problem Formulation and Subspace Restriction

The global objective in federated learning is to minimize the average loss across all clients:

min x = 1/n Σ i f i(x), subject to xi - 1/n Σ i xi = 0, ∀i. To improve efficiency, the paper restricts local updates on each client to low-dimensional subspaces. For each layer l ∈ [L], a subspace projection matrix Pl ∈ R m × rl is introduced, where rl ≪ ml. A full-space variable xl is projected onto a low-dimensional subspace via P l T xl ∈ R rl × dl. This restriction is motivated by the empirical observation that model parameters or gradient matrices in Deep Neural Networks (DNN) often exhibit low-rank structures during training.

Algorithm and Dual Variables for Drift Mitigation

The proposed algorithm is derived from the primal-dual hybrid gradient (PDHG) method [17] and modified into a subspace variant. The use of dual variables enables FedSub to mitigate client drift. Specifically, the update rule for the local variable B k,t+1 i incorporates a correction term:

B k,t+1 i = B k,t i - η g i(B k,t i) + (P k) T Pk-1 1/τ τX-1 t=0 ḡ(Bk-1,t i) − g i(Bk-1,t i), where the term c k i is a correction for client drift. This correction term is essentially the scaled dual variable 1/ητ Λ k i.

Efficiency Analysis: Communication and Computation

The efficiency of FedSub is analyzed in terms of communication, computation, and memory costs compared to full-space methods (like FedAvg when P k = I). The complexity analysis focuses on a single layer.

(a) Communication Efficiency:

In Line 8 of Algorithm 1, client i sends the low-dimensional subspace variable B k i,τ ∈ R r × d to the server. This reduces the uplink communication cost from O(md) to O(rd).

(b) Computation Efficiency:

The total computation complexity of FedSub is O(τ (mrd + Cg(rd)) + 2mrd), where Cg(r d) denotes the computation cost of evaluating a gradient with size r × d. This is significantly reduced compared to the full-space setting, which costs O(τCg(md)).

(c) Memory Efficiency:

The total memory cost for FedSub is O(3rd + Mg(rd) + 2rm + md). When P k = I, the memory cost reduces to O(3md + Mg(md)), which is comparable to or better than the full-space version.

Convergence Analysis and Experimental Validation

The convergence analysis provides bounds on the error term. Theorem 1 establishes that under Assumptions 1-3 and a specific step size condition (η ≤ O 1/θrθ squared mL 2f nτ !), the expected error E∥∇f(x k)∥ squared is bounded. The paper demonstrates effectiveness through experiments on Logistic Regression and CIFAR-100 classification with ResNet. Results show that Our-CD, despite using subspace projections which compromise accuracy compared to P k = I, significantly outperforms FedAvg in terms of convergence accuracy when dual variables are incorporated for correcting client drift. Furthermore, the accuracy of FedSub improves with increasing subspace dimensions r, aligning with the theoretical results.

Conclusion

FedSub successfully addresses data heterogeneity and large-scale model training challenges by combining low-dimensional subspace projections to reduce resource usage and dual variables to mitigate client drift. The theoretical analysis confirms its convergence properties, while experimental results validate its competitive accuracy on classification tasks. Future work includes extending FedSub to adaptive optimizers like Adam.

How it works

The global objective is formulated as minimizing the average loss across clients subject to a consensus constraint (Equation 2).

Improvements for AI systems

Based on the provided scientific paper, here are specific improvements that can be made to AI systems by implementing the FedSub algorithm, along with a description of what those improved systems could achieve:

)

Improve communication efficiency in Federated Learning (FL) for massive Deep Neural Networks (DNNs). By utilizing subspace projection and low-dimensional dual variables, the system can significantly reduce uplink communication costs from being proportional to the full model size to being proportional only to the lower-dimensional subspace dimension.
Reduce computational burden on edge devices during local training iterations. The algorithm achieves this by performing gradient computations within a restricted low-dimensional subspace, leading to a reduction in computation complexity (from O(τCg(md)) towards O(τCg(rd))), which is crucial for resource-constrained edge clients.
Mitigate Client Drift caused by data heterogeneity (non-IID data) without requiring extra communication overhead. The incorporation of low-dimensional dual variables acts as an implicit correction mechanism, ensuring that local updates remain aligned with the global consensus subspace, thereby preventing model divergence caused by client drift.
Achieve faster and more stable convergence in FL settings on heterogeneous data distributions. The theoretical analysis (Theorem 1) provides explicit step-size conditions and convergence rates, allowing practitioners to tune hyperparameters for optimal performance, ensuring that the global model converges efficiently even when clients have non-IID data.
Enable deployment of FL on highly complex models (like large language models or ResNets) by making training feasible despite high dimensionality. The subspace projection allows the system to handle the high-dimensional weight matrices without incurring prohibitive communication and memory bottlenecks associated with full-space updates, effectively enabling collaborative fine-tuning of very large neural networks.

The improved AI systems can specifically:

Compute collaboratively trained models across thousands of distributed devices (e.g., mobile phones, IoT sensors) with minimal data transfer bandwidth usage, as the communication load scales with the subspace dimension rather than the full model dimension.
Train large-scale models (like LLMs or ResNets) in a federated manner where clients have heterogeneous datasets (e.g., different user demographics or task subsets), maintaining high accuracy by actively correcting for client drift through dual variables, leading to more robust and accurate global models than standard FedAvg.
Operate efficiently on edge devices by drastically reducing the memory footprint required during local training steps, allowing for the deployment of sophisticated FL algorithms on hardware with limited memory and processing power.
Maintain high convergence accuracy even when using aggressive compression techniques (like subspace projection), as demonstrated by the experimental results showing that FedSub can achieve accuracy competitive with full-space methods (e.g., achieving approximately 10−7 error in Logistic Regression).

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.