Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum

arXiv:2504.12742 · cs.LG, cs.DC, math.OC · Submitted 2026-08-08 · Read on arXiv

Listen

Radio episode about this paper

Transcript

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

Tom: Next we'll be talking about the paper "Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum".

Jane: The paper was written by Yuan Zhou, Xinli Shi, Xuelong Li, Jiachen Zhong, Guanghui Wen et al. from Southeast University and China Telecom and Purple Mountain Laboratories.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Welcome back to the show, everyone. Today we're digging into a paper with a real mouthful of a title: "Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum." Jane, I'm going to need you to translate that for our listeners.

Jane: Happy to, Tom. Let's break it down piece by piece. "Federated learning" means you've got a bunch of devices—phones, laptops, edge servers—each with its own private data, and they're all trying to train one shared model together without ever sending that raw data to a central place.

Tom: Right, and "decentralized" takes that one step further. There's no central server at all. The devices just talk to their neighbors in a network, like people passing notes in a classroom rather than everything going through the teacher.

Jane: Exactly. And "nonconvex" means the mathematical landscape they're trying to navigate isn't a nice smooth bowl. It's got hills and valleys, which is what you get with neural networks. "Composite" means they're adding a regularizer—something that encourages the model to be sparse or simple—on top of the main loss function.

Tom: And the last two pieces, "gradient tracking" and "momentum," are the clever tricks. Gradient tracking helps each device keep tabs on what the global gradient actually looks like, so no single device's local data pulls the whole model off course. Momentum is like a heavy ball rolling down a hill—it smooths out the jittery steps of stochastic gradient descent.

Jane: That's a great way to put it. And honestly, the fact that they've combined all of these—decentralization, nonconvexity, composite objectives, gradient tracking, and momentum—in one algorithm with solid theoretical guarantees is genuinely impressive.

Tom: I mean, each of those pieces has been studied before, but putting them all together in one framework that actually works? That's the kind of thing that makes a researcher's eyes light up. And the authors—Yuan Zhou, Xinli Shi, Xuelong Li, and their colleagues—they're clearly tackling a real problem that's been nagging at the field.

Jane: Yeah, because in the real world, you rarely get a perfect central server with unlimited bandwidth. You get messy networks, heterogeneous data, and models that need regularization. This paper is saying, "Hey, we can handle all of that at once."

Tom: And that's exactly why we're excited to dig into the details. Stay with us—next we're going to look at what the paper actually claims to achieve.

Summary: Tom: So we've got the title decoded. Now let's talk about what this paper actually delivers. Jane, what's the big picture here?

Jane: The big picture is that they've built an algorithm called DEPOSITUM—which is a bit of a mouthful itself—that handles decentralized federated learning when the local objectives are nonconvex and there's a nonsmooth regularizer in the mix. And they prove it converges.

Tom: Converges meaning it actually finds a good solution, right? Not just in practice, but with mathematical guarantees?

Jane: Precisely. And the guarantee is solid. They show that to reach an expected epsilon-stationary point—basically a point where the gradient is close to zero, meaning you're near a local minimum or saddle point—you need O(one/epsilon squared) iterations. That's the same order as standard stochastic gradient descent, which is reassuring.

Tom: And here's the part that really got me excited. They do this without assuming bounded gradient heterogeneity. That's a big deal because a lot of prior work just assumes the local gradients across clients can't be too different from each other. But in the real world, with truly heterogeneous data, that assumption often doesn't hold.

Jane: Right. And they also avoid assuming mean-squared smoothness, which is another common but restrictive condition. So their analysis is more general, more applicable to real-world scenarios.

Tom: And they get linear speedup with respect to the number of clients. That means if you double the number of clients, you roughly halve the number of iterations needed to reach a given accuracy. That's the scalability property you want in distributed systems.

Jane: Exactly. And they do all this while allowing multiple local updates between communication rounds. That's huge for communication efficiency—you don't want to be sending messages back and forth every single iteration.

Tom: So let me get this straight. They've got an algorithm that works on nonconvex problems, handles nonsmooth regularizers, doesn't rely on restrictive assumptions, achieves linear speedup, and is communication-efficient?

Jane: That's the summary, yes. And the experiments back it up. They tested it on tabular data, image classification, and even fine-tuning a language model. The results are competitive with or better than existing methods.

Tom: I love it when the theory and the experiments line up. But I'm curious about the details of how they actually pulled this off. What's the secret sauce in the algorithm itself?

Jane: That's the perfect segue into our next segment, where we dig into the method. Stick around.

Improvements: Tom: We've established that DEPOSITUM is a big deal. Now let's talk about what it actually improves upon. Jane, what were the gaps in prior work that this paper fills?

Jane: Great question. So there were essentially two camps of prior work. On one side, you had centralized federated composite optimization methods—things like FedMiD, FedDR, FedADMM. They handle the regularizer and the nonconvexity, but they all rely on a central server to aggregate the models.

Tom: And the other camp?

Jane: The other camp was decentralized methods for nonconvex composite problems—things like ProxGT, DEEPSTORM, Prox-DASA-GT. They remove the server, which is great, but they have their own issues. Some of them require batch sizes that depend on the target accuracy, which is impractical. Others require mean-squared smoothness or bounded gradient heterogeneity.

Tom: So DEPOSITUM is saying, "We can do decentralized, we can do composite, we can do nonconvex, and we don't need those restrictive assumptions"?

Jane: Exactly. And there's one more improvement that's really important. Prior decentralized methods typically communicate at every single iteration. DEPOSITUM allows you to do multiple local updates between communication rounds. That's a massive communication saving.

Tom: And the theoretical analysis quantifies that tradeoff, right? They show how the communication period interacts with the stepsize and the network connectivity.

Jane: Yes, and that's one of the more subtle contributions. They have these parameters delta-one and delta-two that capture the contraction margins for the model consensus and the tracking consensus. The analysis shows exactly how increasing the communication period T-zero shrinks those margins and restricts the admissible stepsize.

Tom: So it's not just "here's an algorithm that works"—it's "here's an algorithm, and here's precisely how the knobs you turn affect its behavior."

Jane: That's the kind of depth that makes a paper valuable to practitioners. You're not just blindly tuning hyperparameters; you understand the mechanics.

Tom: And the momentum part—they support both Polyak and Nesterov momentum. That's nice because different problems benefit from different momentum schemes.

Jane: Right, and they show both variants converge with the same guarantees. That flexibility is genuinely useful.

Tom: Alright, I'm sold on the improvements. But I want to get into the nitty-gritty of the first page of the paper. What's the formal setup?

First Page: Tom: So we're looking at the actual first page of "Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum." Jane, walk us through the setup.

Jane: So the problem is formulated as minimizing a global objective that's the average of local objectives across n clients. Each local objective is a smooth, potentially nonconvex function, plus a common regularizer weighted by a parameter beta. That regularizer can be the l1 norm, the Minimax Concave Penalty, the Smoothly Clipped Absolute Deviation, or even an indicator function for constraints.

Tom: And the regularizer is weakly convex, not necessarily convex. That's what makes it challenging, right?

Jane: Exactly. Weak convexity is a generalization of convexity that still allows for things like MCP and SCAD, which are nonconvex but have nice proximal operators. The proximal operator is the workhorse here—it's what lets you handle the nonsmooth part.

Tom: And the network setup—they assume an undirected connected graph. Each client talks only to its neighbors.

Jane: Right. And there's a mixing matrix that encodes the communication weights. The spectral gap of that matrix—how quickly information spreads through the network—plays a key role in the convergence analysis.

Tom: So the paper defines this stationarity measure, s(x, nu-bar), which combines three things: the proximal gradient mapping, the consensus error, and the gradient estimation error. That's how they measure progress.

Jane: Yes, and that's a thoughtful choice because it captures all the ways the algorithm could be "off track." You want the proximal gradient to be small, you want the models to agree across clients, and you want your gradient estimate to be accurate.

Tom: And they prove that this measure goes to zero at a rate of O(one/sqrt(nT)) after a network-dependent transient. That's the linear speedup we talked about.

Jane: And they do it with a mini-batch size that's independent of the total number of iterations. That's important because some prior methods needed batch sizes that grew with the desired accuracy, which is impractical.

Tom: I also noticed they have a nice comparison table in the introduction. It really helps situate their work relative to the field.

Jane: It does. It shows that DEPOSITUM is the only method that simultaneously avoids mean-squared smoothness, avoids bounded gradient heterogeneity, uses a batch size independent of T, allows local updates, and achieves linear speedup. That's a unique combination.

Tom: Alright, so we've covered the setup and the contributions. Let's bring in our guests to get their take on the broader implications.

Conclusion: Tom: We've spent this whole episode on "Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum," and I think it's fair to say this is one of those papers that moves the needle. Jane, what's your final takeaway?

Jane: My takeaway is that this paper gives you a complete package. You get a theoretically sound algorithm, you get guarantees that don't rely on unrealistic assumptions, and you get experimental validation across a range of tasks. That's rare and valuable.

Tom: And Lu, from a research perspective, what excites you most?

Lu: The fact that they've unified so many ideas—gradient tracking, momentum, local updates, weak convexity—into one framework with clean analysis. That's going to be a reference point for future work. People will build on this.

Tom: Meng, what about from an engineering standpoint?

Meng: The communication efficiency is the big win for me. In real deployments, bandwidth is often the bottleneck, not compute. Being able to do multiple local steps between communication rounds, with provable convergence, that's directly applicable.

Tom: And Lalam, you've been quiet. What's the cultural impact here?

Lalam: I think the impact is about democratizing AI training. Decentralized methods like this mean you don't need a massive centralized data center to train good models. Smaller organizations, edge networks, even personal devices can participate. That's a shift toward more distributed ownership of AI.

Jane: That's a beautiful way to put it. And it's not just about the algorithm—it's about who gets to build and control these systems.

Tom: Alright, we've covered the title, the summary, the improvements, the first page, and the broader implications. I think we've given this paper a proper send-off.

Jane: Agreed. It's a strong contribution, and we'll be watching to see what these authors do next.

Tom: Thanks for joining us, everyone. We'll be back with the next paper soon. Until then, keep learning.

Jane: And keep questioning. Goodbye, everyone.

Yuan Zhou, Xinli Shi, Xuelong Li, Jiachen Zhong, Guanghui Wen, Jinde Cao

Southeast University · China Telecom · Purple Mountain Laboratories

cs.LG, cs.DC, math.OC

Submitted: 2026-08-08

Code: https://github.com/Zhouxyz/DEPOSITUM

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 53/100

The gist: The paper addresses the Decentralized Nonconvex Composite Optimization Problem (NCOP), formulated as: `min ϕ(x) = f(x) + βh(x), f(x) ≜ (1/n) Σ i=1 n fi(x)` where each local function `fi` is

Key concepts

Federated Learning
A machine learning approach where multiple devices train a shared model using their private data. The raw data remains on the local device, ensuring privacy while allowing the group to build a single, collective model.
Decentralized Learning
An extension of Federated Learning where there is no central server. Devices communicate directly with neighboring nodes in a network, passing information locally rather than sending all data through one point.
Nonconvex Composite Optimization
A challenging mathematical problem type used in training neural networks. It involves optimizing functions that have complex shapes (nonconvex) and requires adding regularizers to encourage the model to be simple or sparse.

Terminology

Summary

Summary

The paper Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum by Yuan Zhou, Xinli Shi, Xuelong Li, Jiachen Zhong, Guanghui Wen, and Jinde Cao, proposes a new algorithm for decentralized federated learning in settings where local objectives are nonconvex and include a nonsmooth, weakly convex regularizer.

Problem Addressed:

The paper addresses the Decentralized Nonconvex Composite Optimization Problem (NCOP), formulated as:

min ϕ(x) = f(x) + βh(x), f(x) ≜ (1/n) Σ i=1 n fi(x)

where each local function fi is smooth but not necessarily convex, and h is a common weakly convex function that may be nondifferentiable (e.g., l1 norm, MCP, SCAD). The goal is to solve this problem in a decentralized federated learning setting without a central server, where clients communicate only with their neighbors over a connected graph. The paper notes that existing methods for federated composite optimization are mostly centralized, and decentralized methods often require communication at every iteration or rely on restrictive assumptions like bounded gradient heterogeneity or mean-squared smoothness.

Proposed Algorithm: DEPOSITUM

The authors propose the Decentralized fEderated PrOximal Stochastic gradIent tracking with momentUM (DEPOSITUM) algorithm. Key features of DEPOSITUM include:

  • Proximal Updates: It integrates proximal gradient steps to handle the nonsmooth regularizer.

  • Gradient Tracking: It maintains an auxiliary variable y to track the average of local gradients, which corrects for client drift caused by data heterogeneity without requiring a bounded gradient heterogeneity assumption.

  • Momentum: It incorporates both Polyak and Nesterov momentum variants.

  • Local Updates: It allows for multiple local updates between communication rounds (T0), improving communication efficiency.

The algorithm's update rules are:

  • Model update: x t+1 = W t prox h 1/(αβ) x t - αy t

  • Momentum update: For Polyak, ν t+1 = γν t + (1-γ)g t+1; for Nesterov, µ t+1 = γµ t + (1-γ)g t+1 and ν t+1 = γµ t+1 + (1-γ)g t+1.

  • Tracking update: y t+1 = W t (y t + ν t+1 - ν t).

Theoretical Contributions:

The paper provides a rigorous convergence analysis for DEPOSITUM under standard assumptions (smooth nonconvex local functions, weakly convex regularizer, unbiased stochastic gradients with bounded variance, and a doubly stochastic mixing matrix). The main theoretical results are:

  • Iteration Complexity: DEPOSITUM achieves an expected ϵ-stationary point with an iteration complexity of O(1/ϵ2).

  • Linear Speedup: With an appropriate stepsize and momentum schedule (specifically, 1 - γ = √(n/(T+1)) and α = √n/(24L√(T+1))), and a mini-batch size independent of T (e.g., B = ⌈√n⌉), the averaged stationarity measure achieves a rate of O(1/√(nT)) after a network-dependent transient, specifically when T = Ω(n2). This demonstrates linear speedup with respect to the number of clients n.

  • No Restrictive Assumptions: The analysis does not require bounded gradient heterogeneity or mean-squared smoothness assumptions.

  • Trade-off Analysis: The analysis quantifies the trade-off among the communication period T0, the admissible stepsize, and network connectivity (spectral gap 1-λ). Increasing T0 reduces communication frequency but restricts the stepsize and enlarges the network-dependent constant.

Experimental Validation:

The paper validates DEPOSITUM through extensive experiments on five datasets (A9A, MNIST, FMNIST, CIFAR-10, and Dolly-15k for LLM fine-tuning). The experiments cover various models (Linear, MLP, CNN, ResNet-18, and Qwen2.5-0.5B with LoRA), regularizers (l1, MCP, SCAD), data partitions (IID and non-IID via Dirichlet distribution), network topologies (Ring, 2-hop Ring, 3-hop Ring, Complete), communication periods, and client counts. Key findings include:

  • Hyperparameter Effects: The experiments characterize the effects of stepsize α, momentum parameter γ, communication period T0, and regularization weight β, confirming the theoretical trade-offs.

  • Topology Comparison: DEPOSITUM outperforms DEEPSTORM and Prox-DASA-GT in terms of test loss and stability across different graph topologies.

  • Client Count: The experiments confirm the linear speedup trend with respect to the number of clients n.

  • LLM Fine-tuning: DEPOSITUM achieves the lowest response perplexity (PPL) on the Dolly-15k test tasks compared to DEEPSTORM and Prox-DASA-GT in the composite setting, and remains competitive with smooth-only methods like GUT, D-PSGD, and DSGT.

  • Comparison with Centralized Methods: In additional experiments, DEPOSITUM attains the highest mean accuracy in most settings compared to centralized FL methods FedMiD, FedDR, and FedADMM.

Conclusion:

The paper concludes that DEPOSITUM is an effective and theoretically sound method for decentralized nonconvex composite federated learning, offering linear speedup, robustness to data heterogeneity, and communication efficiency. The authors note that extending the guarantees to time-varying or directed networks and imperfect communication remains future work.

Improvements for AI systems

Based on the scientific paper, here are the specific improvements I can make to AI systems and what the improved system can do:

Improvement: Implement the DEPOSITUM algorithm to replace centralized federated learning architectures. The system uses a peer-to-peer communication graph where each client exchanges model updates only with neighbors, eliminating the single-point-of-failure and communication bottleneck of server-based approaches.

What the improved system can do:

  • Train models in environments where a central coordinator is unavailable, such as edge networks, vehicular networks, or military/adversarial settings

  • Maintain training continuity even if individual clients fail or drop out, since there is no central server to crash

  • Scale to large client populations without the server becoming a communication bottleneck

Improvement: Incorporate the algorithm's support for weakly convex regularizers (e.g., l1 norm, MCP, SCAD) and nonconvex loss functions, which are common in real-world applications but poorly handled by standard SGD-based methods.

Improvement: Use the dual momentum schemes (Polyak and Nesterov) combined with gradient tracking to correct for data heterogeneity without requiring bounded gradient heterogeneity assumptions.

Improvement: Implement the periodic communication schedule (T0 local updates between communication phases) to reduce communication frequency while maintaining convergence guarantees.

Improvement: Leverage the theoretical guarantee of O(1/√(nT)) convergence rate, achieving linear speedup with respect to the number of clients n.

Improvement: The algorithm does not require mean-squared smoothness or bounded gradient heterogeneity assumptions, making it applicable to a wider range of problems.

Improvement: Apply the algorithm to decentralized parameter-efficient fine-tuning (e.g., LoRA) of large language models, as demonstrated in the experiments with Qwen2.5-0.5B.

Improvement: Use the tunable regularization weight β to control the trade-off between model sparsity and accuracy, as demonstrated in the experiments.

The improved AI system can:

  1. Train models in fully decentralized networks without any central server

  2. Handle nonconvex, nonsmooth, and heterogeneous data distributions robustly

  3. Achieve faster convergence with momentum while correcting for data drift

  4. Reduce communication costs by 5-10x through local updates

  5. Scale efficiently with client count (linear speedup)

  6. Fine-tune LLMs collaboratively while preserving data privacy

  7. Produce sparse, efficient models with controlled accuracy trade-offs

  8. Operate reliably in edge, vehicular, and privacy-sensitive environments where centralized approaches are infeasible

Sources

Related papers