REVIEW 2 major objections 5 minor 35 references
Non-convex composite federated learning with heterogeneous data
T0 review · 2 major / 5 minor · reviewed 2026-08-09 · deepseek-v4-flash
Pith's one-line read Sending the pre-proximal vector to the server recovers the true average gradient and removes client drift in composite federated learning.
desk verdict A solid non-convex composite FL paper with one repairable lemma inconsistency and an honest Bg residual; worth full peer review. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The argument is carried by three linked devices. The first is the two-model bookkeeping: only the pre-proximal local model is transmitted, so the server sees the sum of gradients before the regularizer's proximal operator acts on them. The second is the correction term $c^{r+1}_i = \frac{1}{\eta_g\eta\tau}(P_{\tilde{\eta}}(x^r)-x^{r+1}) - \frac{1}{\tau}\sum_{t=0}^{\tau-1}\nabla f_i(z^r_{i,t};B^r_{i,t})$, which each client reconstructs from the broadcast global model; its average over clients is zero, which is exactly what keeps the transmitted average undistorted. The third is the auxiliary function $\Omega^r = F(P_{\tilde{\eta}}(x^r))-F^\star + \frac{1}{n\tilde{\eta}}\|\Lambda^r - \bar{\Lambda}^r\|^2$, whose first term tracks global suboptimality and whose second term bounds the local-drift error; Lemma A.1 and Lemma A.2 close the recursion on this drift term. The parameter $(t+1)\eta$ inside the local proximal operator is chosen so that the post-proximal local iterates approximate a centralized proximal gradient step on the global model, which is what makes the drift bound small enough to close.
What would settle it
A direct calculation settles the internal consistency: the proof of Lemma A.1 invokes $8(t+1)\eta^2 L^2 \leq 1/(8\tau)$, which requires $\tilde{\eta} \leq \eta_g/(8L)$, stricter than the lemma's stated $\tilde{\eta}\leq \eta_g/\sqrt{20L}$; the theorem's condition (13) with $\eta_g\ge 1.5$ and $\tilde{\eta}\le 1/(10L)$ implies the stricter bound, so the main theorems survive but the lemma statement is imprecise. Separately, run Algorithm 1 with full local gradients on a non-convex composite problem where $g$ is the indicator function of a box, so Assumption 3.1 fails and $B_g=\infty$, and check whether $\|G(P_{\tilde{\eta}}(x^r))\|^2$ goes to machine zero; exact convergence would show the $B_g$ residual is an artifact of the proof.
Extended reading notes
Core claim
The load-bearing claim is that each client keeps two local models at every step, the pre-proximal variable $\hat{z}^r_{i,t}$ and the post-proximal variable $z^r_{i,t}=P_{(t+1)\eta}(\hat{z}^r_{i,t})$, evaluates stochastic gradients at the post-proximal model, and after $\tau$ updates sends only $\hat{z}^r_{i,\tau}$ to the server. Because the drift-correction terms $c^r_i$ average to zero across clients, the server's update becomes $x^{r+1}=P_{\tilde{\eta}}(x^r)-\tilde{\eta}\,\frac{1}{n\tau}\sum_{t=0}^{\tau-1}\sum_{i=1}^n\nabla f_i(z^r_{i,t};B^r_{i,t})$, which is a centralized proximal gradient step affected only by stochastic sampling, with no distortion from the nonlinear proximal map. This is the decoupling: the server can run a proximal step on the global model using the true average local gradient. The convergence analysis measures progress with the gradient mapping $G(P_{\tilde{\eta}}(x^r))=\frac{1}{\tilde{\eta}}(P_{\tilde{\eta}}(x^r)-\tilde{x}^{r+1})\in \nabla f(P_{\tilde{\eta}}(x^r))+\partial g(\tilde{x}^{r+1})$, and proves $\frac{1}{R}\sum_{r=1}^R \mathbb{E}\|G(P_{\tilde{\eta}}(x^r))\|^2 \leq \frac{\mathbb{E}[\Omega^1]}{0.3\tilde{\eta}R} + \frac{20\sigma^2}{n\tau b} + \frac{187L^2\tilde{\eta}^2 B_g^2}{\eta_g^2}$ for general non-convex problems, and $\mathbb{E}[\Omega^{R+1}] \leq (1-\frac{\mu\tilde{\eta}}{3})^R \mathbb{E}[\Omega^1] + \frac{18\sigma^2}{\mu n\tau b} + \frac{168L^2\tilde{\eta}^2 B_g^2}{\mu\eta_g^2}$ under the proximal PL inequality. These are stated as Theorem 3.5 and Theorem 3.6, with residuals set by stochastic variance and by the squared subgradient bound of the regularizer.
Load-bearing premise
The regularizer $g$ must have globally bounded subgradients (norm at most $B_g$), which fails for constraint-set indicator functions and makes the stated residual error infinite.
Editorial extensions
If this is right
- A single $d$-dimensional vector per client per communication round is enough to remove client drift for non-convex composite objectives; the correction term is reconstructed locally from the broadcast global model, so no control-variate communication is needed.
- Communication frequency can be reduced through local updates while the method keeps a per-round convergence guarantee; the variance term $\sigma^2/(n\tau b)$ shrinks with larger mini-batches and vanishes when full local gradients are used.
- Under the proximal PL inequality the composite loss converges linearly to a neighborhood whose radius is set by the stochastic variance and by $B_g^2$, the squared subgradient bound of the regularizer.
- If the regularizer has unbounded subgradients, such as an indicator function of a constraint set, the stated residual bounds are vacuous; the paper explicitly notes in Remark 3.7 that it cannot remove the $B_g$ term in the non-convex case.
- Full-gradient experiments on sparse logistic regression show exact convergence with no visible $B_g$ residual, suggesting that the theoretical residual is conservative rather than a fundamental obstruction.
Reading between the lines
- The pre-proximal transmission idea may extend to other nonlinear client-side maps, such as variance-reduced proximal operators or adaptive-regularizer steps, whenever the server needs the average of a pre-image rather than a post-image.
- The $B_g^2$ residual likely comes from the two-sided subgradient bound used in Lemma A.1; a drift bound that exploits one-sided or structure-dependent properties of $g$ could remove or shrink it for specific regularizers.
- The proof of Lemma A.1 requires the stricter condition $\tilde{\eta} \leq \eta_g/(8L)$ at the step where it uses $8(t+1)\eta^2 L^2 \leq 1/(8\tau)$, while the lemma states only $\tilde{\eta} \leq \eta_g/\sqrt{20L}$; the theorem's global condition (13) appears to imply the stricter bound when $\eta_g\ge 1.5$, so the main theorems survive, but the lemma as written is imprecise.
- Because the correction term uses the previous round's global gradient information, the drift correction is inherently one round stale; a testable variant with momentum or variance reduction over older gradients could shrink the residual at no extra communication cost.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper addresses the composite federated learning problem (1) with a non-convex smooth loss f and a convex, possibly non-smooth regularizer g. It proposes Algorithm 1, in which each client maintains both a pre-proximal and a post-proximal local model, performs τ local updates using mini-batch stochastic gradients plus a client-drift correction term, and transmits only the pre-proximal model to the server. The server then updates a pre-proximal global model so that the server-side recursion mimics centralized proximal gradient descent. Under Assumptions 3.1–3.4 and step-size condition (13), the paper proves an average bound on the gradient mapping norm for general non-convex f (Theorem 3.5) and a linear convergence bound for the function value under the proximal PL inequality (Theorem 3.6), both with explicit residual errors depending on the stochastic variance and on the subgradient bound Bg. Numerical experiments on sparse logistic regression and CNN training on MNIST support the algorithmic claims.
Significance. If the results hold, Algorithm 1 is a meaningful step for composite federated learning: it decouples proximal evaluation from communication, requires only one d-dimensional vector per client per communication round, does not assume bounded data heterogeneity, and handles non-smooth regularizers. The convergence analysis is non-trivial and the appendix contains a full Lyapunov-style proof with explicit lemmas. The paper is also candid about its limitations, notably the residual term depending on Bg. The main concerns are a step-size condition mismatch in the appendix and the restrictive bounded-subgradient assumption on g; both are local and fixable, and the central algorithmic idea appears sound.
major comments (2)
- [Appendix A.4, Lemma A.1 and Lemma A.2] Lemma A.1 is stated under the condition η̃ ≤ ηg/√(20L), but the proof of (A.20) uses the inequality 8(t+1)η²L² ≤ 1/(8τ), which is equivalent to η̃ ≤ ηg/(8L). For typical values of L this is a stricter condition than the one stated, so the lemma as written is not fully proven. Lemma A.2 has the same issue and also contains a typographical garble in its stated condition ("η̃ ≤ ηg√(20L)"). The step-size condition (13) in Theorems 3.5 and 3.6 does imply the stricter bound, so the main results are not invalidated, but the lemma statements must be corrected and Lemma A.2's hypothesis should match the proof requirements.
- [Assumption 3.1 and Theorems 3.5–3.6] Assumption 3.1 requires the subgradient of g to be uniformly bounded by Bg. For many canonical composite problems, g is the indicator function of a convex constraint set, for which subgradients on the boundary are unbounded, making Bg infinite and the residual terms 187L²η̃²Bg²/ηg² (Theorem 3.5) and 168L²η̃²Bg²/(μηg²) (Theorem 3.6) vacuous. The authors acknowledge this limitation in Remark 3.7 and the conclusion, but the abstract and introduction present the result as a general treatment of non-convex composite FL. I recommend adding an explicit qualification in the abstract and a short discussion of settings where Bg is finite (e.g., ℓ1, ℓ∞, ReLU) and of possible avenues for removing the Bg residual under additional assumptions.
minor comments (5)
- [Equation (2) and Appendix A.1] In the compact form (2), the two sums in the correction term are typeset identically, although one is the averaged gradient across clients and the other is the per-client gradient; please add the overbar or a textual clarification to avoid confusion.
- [Section 4.2] The text reads "The remaining 3,0000 samples" and should be "The remaining 30,000 samples".
- [Appendix A.4, Lemma A.2] The stated condition in Lemma A.2 appears as "η̃ ≤ ηg√(20L)" but should presumably be "η̃ ≤ ηg/√(20L)", and in any case must be updated to the stricter condition used in the proof.
- [Notation and introduction] The phrase "sends only a single d-dimensional vector per communication round" is not fully precise because the server also broadcasts a d-dimensional vector; the per-client communication cost is correct, but the wording in the abstract could be clarified.
- [Reference [22]] The proof of Lemma A.1 cites "[22, Corollary C.1]" for an inequality on the variance of averaged stochastic gradients; please verify that the cited result indeed appears in reference [22] and that the reference is complete.
Circularity Check
No significant circularity: the non-convex and proximal-PL convergence guarantees are derived from explicitly stated assumptions via a Lyapunov argument, and the only self-citation ([33]) is used for a separate strongly-convex theorem, not for the paper's central claims.
full rationale
The paper's central results, Theorems 3.5 and 3.6, are derived from Assumptions 3.1-3.4 using a constructed auxiliary function Omega^r, the drift bound in Lemma A.1, the recursion in Lemma A.2, and the server-side descent inequality based on Fact 1. None of these steps assumes the conclusion: the residual terms O(sigma^2/(n*tau*b)) and O(L^2*eta_tilde^2*B_g^2/eta_g^2) arise algebraically from the stochastic variance bound and from the bounded-subgradient constant B_g in Assumption 3.1, not from fitting any parameter to the output. The algorithm and the strongly-convex theorem from the authors' conference paper [33] are cited in Remark 3.7, but that citation is not load-bearing for the non-convex or proximal-PL proofs; those proofs are self-contained once the stated assumptions are accepted. The paper explicitly acknowledges that the (t+1)eta parameter choice is heuristic and that the B_g-dependent residual cannot be eliminated in the non-convex case; these are honest limitations, not circular reductions. The apparent mismatch in Lemma A.1 between the stated step-size condition and the stricter condition used inside its proof is a proof-technical imprecision, not a case of defining a prediction in terms of the input. Overall, the derivation chain is self-contained and the claims are not forced by definition or by self-citation.
Assumptions & free parameters
assumptions (5)
- domain assumption The subgradient of the convex regularizer g is uniformly bounded: ∥∇g(x)∥ ≤ Bg for all x.
- domain assumption Each client loss fi is L-smooth.
- domain assumption The loss is either generally non-convex or satisfies the proximal PL inequality (3.3).
- domain assumption Local mini-batch gradients are unbiased with variance bounded by σ²/b.
- standard math Proximal inequality Fact 1 from Reddi et al. [10] and nonexpansiveness of proximal maps.
Cite this review
Pith. "Pith review of Non-convex composite federated learning with heterogeneous data." pith.science (2026). https://pith.science/paper/TPS2ZUUS
@misc{pith2026250203958,
author = {Pith},
title = {Pith review of: Non-convex composite federated learning with heterogeneous data},
year = {2026},
howpublished = {\url{https://pith.science/paper/TPS2ZUUS}},
note = {Machine review of arXiv:2502.03958}
}
read the original abstract
We propose an innovative algorithm for non-convex composite federated learning that decouples the proximal operator evaluation and the communication between server and clients. Moreover, each client uses local updates to communicate less frequently with the server, sends only a single d-dimensional vector per communication round, and overcomes issues with client drift. In the analysis, challenges arise from the use of decoupling strategies and local updates in the algorithm, as well as from the non-convex and non-smooth nature of the problem. We establish sublinear and linear convergence to a bounded residual error under general non-convexity and the proximal Polyak-Lojasiewicz inequality, respectively. In the numerical experiments, we demonstrate the superiority of our algorithm over state-of-the-art methods on both synthetic and real datasets.
Figures
Reference graph
Works this paper leans on
-
[1]
Mahmoud M Badr, Mohamed Mahmoud, Yuguang Fang, Mohammed Abdulaal, Abdulah J Aljohani, Waleed Alasmary, and Mohamed I Ibrahem. Privacy-preserving and communication-efficient energy prediction scheme based on federated learning for smart grids. IEEE Internet of Things Journal, 10(9):7719–7736, 2023
work page 2023
-
[2]
Fast composite optimization and statistical recovery in federated learning
Yajie Bao, Michael Crawshaw, Shan Luo, and Mingrui Liu. Fast composite optimization and statistical recovery in federated learning. In International Conference on Machine Learning, pages 1508–1536, 2022. 14
work page 2022
-
[3]
A fast iterative shrinkage- thresholding algorithm for linear inverse problems
Amir Beck and Marc Teboulle. A fast iterative shrinkage- thresholding algorithm for linear inverse problems. SIAM Journal on Imaging Sciences , 2(1):183–202, 2009
2009
-
[4]
Gianluca Bianchin, Jorge I Poveda, and Emiliano Dall’Anese. Online optimization of switched LTI systems using continuous-time and hybrid accelerated gradient flows. Automatica, 146:110579, 2022
work page 2022
-
[5]
Sparse linear regression from perturbed data
Sophie M Fosson, Vito Cerone, and Diego Regruto. Sparse linear regression from perturbed data. Automatica, 122:109284, 2020
work page 2020
-
[6]
Distributed networked real-time learning
Alfredo Garcia, Luochao Wang, Jeff Huang, and Lingzhou Hong. Distributed networked real-time learning. IEEE Transactions on Control of Network Systems , 8(1):28–38, 2020
work page 2020
-
[7]
Mini- batch stochastic approximation methods for nonconvex stochastic composite optimization
Saeed Ghadimi, Guanghui Lan, and Hongchao Zhang. Mini- batch stochastic approximation methods for nonconvex stochastic composite optimization. Mathematical Programming, 155(1):267–305, 2016
2016
-
[8]
A strictly contractive Peaceman–Rachford splitting method for convex programming
Bingsheng He, Han Liu, Zhaoran Wang, and Xiaoming Yuan. A strictly contractive Peaceman–Rachford splitting method for convex programming. SIAM Journal on Optimization , 24(3):1011–1040, 2014
work page 2014
Show all 35 references
-
[9]
Distributed networked learning with correlated data
Lingzhou Hong, Alfredo Garcia, and Ceyhun Eksin. Distributed networked learning with correlated data. Automatica, 137:110134, 2022
2022
-
[10]
Proximal stochastic methods for nonsmooth nonconvex finite-sum optimization
Sashank J Reddi, Suvrit Sra, Barnabas Poczos, and Alexander J Smola. Proximal stochastic methods for nonsmooth nonconvex finite-sum optimization. Advances in Neural Information Processing Systems , 29, 2016
2016
-
[11]
Linear convergence of gradient and proximal-gradient methods under the Polyak- Lojasiewicz condition
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal-gradient methods under the Polyak- Lojasiewicz condition. In Machine Learning and Knowledge Discovery in Databases, pages 795– 811, 2016
2016
-
[12]
Mime: Mimicking centralized stochastic algorithms in federated learning
Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri, Sashank J Reddi, Sebastian U Stich, and Ananda Theertha Suresh. Mime: Mimicking centralized stochastic algorithms in federated learning. arXiv preprint arXiv:2008.03606, 2020
2008 arXiv
-
[13]
Scaffold: Stochastic controlled averaging for federated learning
Sai Praneeth Karimireddy, Satyen Kale, Mehryar Mohri, Sashank Reddi, Sebastian Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning. In International Conference on Machine Learning, pages 5132–5143, 2020
2020
-
[14]
Understanding notions of stationarity in nonsmooth optimization: A guided tour of various constructions of subdifferential for nonsmooth functions
Jiajin Li, Anthony Man-Cho So, and Wing-Kin Ma. Understanding notions of stationarity in nonsmooth optimization: A guided tour of various constructions of subdifferential for nonsmooth functions. IEEE Signal Processing Magazine, 37(5):18–31, 2020
2020
-
[15]
Federated optimization in heterogeneous networks
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems , 2:429–450, 2020
2020
-
[16]
On the convergence of FedAvg on non-iid data
Xiang Li, Kaixuan Huang, Wenhao Yang, Shusen Wang, and Zhihua Zhang. On the convergence of FedAvg on non-iid data. In International Conference on Learning Representations, 2019
2019
-
[17]
Private non-convex federated learning without a trusted server
Andrew Lowy, Ali Ghafelebashi, and Meisam Razaviyayn. Private non-convex federated learning without a trusted server. In International Conference on Artificial Intelligence and Statistics , pages 5749–5786, 2023
2023
-
[18]
Error bounds and convergence analysis of feasible descent methods: a general approach
Zhi-Quan Luo and Paul Tseng. Error bounds and convergence analysis of feasible descent methods: a general approach. Annals of Operations Research , 46(1):157–178, 1993
1993
-
[19]
Communication- efficient learning of deep networks from decentralized data
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient learning of deep networks from decentralized data. In Artificial Intelligence and Statistics , pages 1273–1282, 2017
2017
-
[20]
Gradient free cooperative seeking of a moving source
Elad Michael, Chris Manzie, Tony A Wood, Daniel Zelazo, and Iman Shames. Gradient free cooperative seeking of a moving source. Automatica, 152:110948, 2023
2023
-
[21]
Primal-dual subgradient methods for convex problems
Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical Programming, 120(1):221–259, 2009
2009
-
[22]
Differentially private federated learning on heterogeneous data
Maxence Noble, Aur´ elien Bellet, and Aymeric Dieuleveut. Differentially private federated learning on heterogeneous data. In International Conference on Artificial Intelligence and Statistics , pages 10110–10145, 2022
2022
-
[23]
FedSplit: An algorithmic framework for fast federated optimization
Reese Pathak and Martin J Wainwright. FedSplit: An algorithmic framework for fast federated optimization. Advances in Neural Information Processing Systems , 33:7057–7066, 2020
2020
-
[24]
Local SGD converges fast and communicates little
Sebastian U Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2018
2018
-
[25]
FedDR–randomized Douglas-Rachford splitting algorithms for nonconvex federated composite optimization
Quoc Tran Dinh, Nhan H Pham, Dzung Phan, and Lam Nguyen. FedDR–randomized Douglas-Rachford splitting algorithms for nonconvex federated composite optimization. Advances in Neural Information Processing Systems , 34:30326–30338, 2021
2021
-
[26]
FedADMM: A federated primal-dual algorithm allowing partial participation
Han Wang, Siddartha Marella, and James Anderson. FedADMM: A federated primal-dual algorithm allowing partial participation. In 2022 IEEE 61st Conference on Decision and Control (CDC) , pages 287–294, 2022
2022
-
[27]
Decentralized nonconvex optimization with guaranteed privacy and accuracy
Yongqiang Wang and Tamer Ba¸ sar. Decentralized nonconvex optimization with guaranteed privacy and accuracy. Automatica, 150:110858, 2023
2023
-
[28]
On stochastic gradient and subgradient methods with adaptive steplength sequences
Farzad Yousefian, Angelia Nedi´ c, and Uday V Shanbhag. On stochastic gradient and subgradient methods with adaptive steplength sequences. Automatica, 48(1):56–67, 2012
2012
-
[29]
Federated composite optimization
Honglin Yuan, Manzil Zaheer, and Sashank Reddi. Federated composite optimization. In International Conference on Machine Learning, pages 12253–12266, 2021
2021
-
[30]
On convergence of FedProx: Local dissimilarity invariant bounds, non-smoothness and beyond
Xiaotong Yuan and Ping Li. On convergence of FedProx: Local dissimilarity invariant bounds, non-smoothness and beyond. Advances in Neural Information Processing Systems, 35:10752–10765, 2022
2022
-
[31]
A family of inexact SQA methods for non-smooth convex minimization with provable convergence guarantees based on the Luo–Tseng error bound property
Man-Chung Yue, Zirui Zhou, and Anthony Man-Cho So. A family of inexact SQA methods for non-smooth convex minimization with provable convergence guarantees based on the Luo–Tseng error bound property. Mathematical Programming, 174(1):327–358, 2019
2019
-
[32]
Secure and decentralized federated learning framework with non-iid data based on blockchain
Feng Zhang, Yongjing Zhang, Shan Ji, and Zhaoyang Han. Secure and decentralized federated learning framework with non-iid data based on blockchain. Heliyon, 10(5), 2024
2024
-
[33]
Composite federated learning with heterogeneous data
Jiaojiao Zhang, Jiang Hu, and Mikael Johansson. Composite federated learning with heterogeneous data. arXiv preprint arXiv:2309.01795, 2023
2023 arXiv
-
[34]
Fedaudio: A federated learning benchmark for audio tasks
Tuo Zhang, Tiantian Feng, Samiul Alam, Sunwoo Lee, Mi Zhang, Shrikanth S Narayanan, and Salman Avestimehr. Fedaudio: A federated learning benchmark for audio tasks. In IEEE International Conference on Acoustics, Speech and Signal Processing, pages 1–5, 2023
2023
-
[35]
FedPD: A federated learning framework with adaptivity to non-iid data
Xinwei Zhang, Mingyi Hong, Sairaj Dhople, Wotao Yin, and Yang Liu. FedPD: A federated learning framework with adaptivity to non-iid data. IEEE Transactions on Signal Processing, 69:6055–6070, 2021. 15
2021
Reviewed August 9, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.