Communication-Efficient Heterogeneous Federated Learning with Generalized Heavy-Ball Momentum

Authors: Riccardo Zaccone, Sai Praneeth Karimireddy, Carlo Masone, Marco Ciccone
Primary area: optimization
Keywords: federated learning, momentum, distributed learning, deep learning

Abstract:
Federated Learning (FL) has emerged as the state-of-the-art approach for learning from decentralized data in privacy-constrained scenarios. However, system and statistical challenges hinder real-world applications, which demand efficient learning from edge devices and robustness to heterogeneity. Despite significant research efforts, existing approaches (i) are not sufficiently robust, (ii) do not perform well in large-scale scenarios, and (iii) are not communication efficient. In this work, we propose a novel _Generalized Heavy-Ball Momentum_ (GHBM), proving that it enjoys an improved theoretical convergence rate w.r.t. existing FL methods based on classical momentum in _partial participation_, without relying on bounded data heterogeneity. Then, we present FedHBM as an adaptive, communication-efficient by-design instance of GHBM. Extensive experimentation on vision and language tasks, in both controlled and realistic large-scale scenarios, confirms our theoretical findings, showing that GHBM substantially improves the state of the art, especially in large scale scenarios with high data heterogeneity and low client participation.

Decision: reject
Reviewer scores: 6, 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1729492376852

thread_type:
official_review

<thread_object>
creator: reviewer_kz3z

<content>
summary:
This paper proposes a multistep inertial momentum method in federated learning to alleviate the biases of the partial participation. It provides the theoretical improvements that the proposed method can achieve the full-participation rate under partial-participation. Experiments also shows its efficiency on different setups.

strengths:
1. The writing is well-done and easy to understand the motivation and design in this paper.

2. The motivation of reducing the gaps caused by partial participation is also a good perspective, which is also a essential academic issue in FL.

weaknesses:
1. The proposed approach is incremental. The use of momentum averaged over $\tau$ rounds, rather than a single round, was introduced based on experimental intuition. Although a correction scheme was applied both locally and globally, this appears to be a simple extension of the FedCM and FedADC approach [1]. 

[1] Ozfatura, Emre, Kerem Ozfatura, and Deniz Gündüz. "Fedadc: Accelerated federated learning with drift control." 2021 IEEE International Symposium on Information Theory (ISIT). IEEE, 2021.

2. The assumption 4.3 is too strong. Periodic participation is feasible, but the disconnect between the experiments and the theory prevents these two parts from supporting each other. Periodic participation implies that the sampling distribution is time-varying and uneven, which simplifies the impact of heterogeneity. The difference between the two is analogous to the theoretical analysis of SGD versus incremental GD. In fact, incremental GD has been proven under certain conditions to surpass the convergence speed of SGD. I believe this is one of the main weaknesses of the paper. The authors mention that under optimal participation rate design, the heterogeneity assumption is unnecessary. In fact, this is one of the theories behind incremental GD eliminating the stochastic biases. This proof does not accurately demonstrate that such properties can be maintained under random selection.

3. Some of the assumptions in the proofs seem to differ. Lemma 4.4 and Corollary 4.5 adopts the heterogeneity assumption but the authors do not claim it clearly. See question (3).

4. Some of the conclusions seem to lack precision. Line 381 comments "since increasing $\tau$ does not increase the overall error due to this term." Does the term $\frac{1}{\tau}\sum_{k=t-\tau+1}^{t}U_{k}$ is proven a constant order? If I missed this proof, please let me know. Currently, there does not seem to be sufficient evidence to indicate the specific nature of this term.

5. The comparisons in Table 1 lack important information. Although Table 1 shows that the proposed method can achieve global participation convergence speed under partial participation, this seems to strictly rely on Assumption 3 and the cyclical setting. I believe it is necessary to clearly state the conditions of this conclusion, rather than allowing readers to mistakenly think that it has surpassed the upper bound of partial participation as it currently stands.  In fact, after reading the entire paper, I believe that with the help of Assumption 3, the theoretical conclusions of FedCM could also reach the same conclusions presented in this paper. Based on the current results, I do not think this theoretical advancement is due to improvements in the algorithm itself.

questions:
1. What is the meaning of $\alpha$ in Figure.1? And, why is the variacne of the gaps larger on iid than non-iid? In principle, the raw gradients of identically distributed datasets should remain highly similar between local clients, resulting in smaller variance. The author could further elaborate on the experimental process and explain why this conclusion leads to counterintuitive results.

2. I have a question about the form of Equation 3. Why is this considered the desired momentum update? Especially for the term $\widetilde{g}^{k}(\theta^{t-1})$, if the state $\theta^{t-1}$ is fixed, i.e. it is irrelated to $k$, then it seems to be only related to $S^k$. Do authors mean by that under a large $\tau$, the union of $S^k\rightarrow S$? I believe the explanation here is overly simplistic, making the approach appear unreasonable.

3. Eq (8) and (9) adopt the bounded heterogeneity assumption, since there is $G$ terms to show the deminished errors. However, the absence of the $G$ term's influence in the subsequent conclusions. Additionally, are the inferences from Equations (8) and (9) still valid for the later conclusions?

4. Can authors provide the general result of Theorem 4.7? For instance, under any selection on $\tau$ without fixing it as $1/C$. Also, how to recover the $\tau =1$ case? And how does the selection $\tau$ impact the final results?

5. I notice the general rounds are set as $10,000$. What is number of local steps $J$? This seems to be a very exaggerated setup, as the referenced baseline algorithms typically do not exceed 1,000 training rounds on the same task. Though FedCM provides results after 4,000 rounds. Additionally, why Table 3 does not contain the results of FedCM and FedADC, and other methods in Table 2?

6. The selection of $\beta$ seems to be related to $T$ and $J$. So, this term is a parameter that needs to be precisely designed in this algorithm rather than being a constant? Does it imply that different total training lengths $T$ will affect the choice of $\beta$? 

Some typos:
(1) Line 142, $\theta^{t,J_i}$ should be corrected as $\theta_i^{t,J_i}$.

---------------------------------- After rebuttal -----------------------------------------------

I hope the author can include a discussion section in the main text to draw an analogy between the discussion of IGD comparisons and heterogeneity in FL. This will help enhance the paper's impact, as it truly highlights the core source of this technique. Secondly, regarding the author's mention of research on lower bounds, I encourage them to pursue related studies. Research on lower bounds has always been very challenging yet valuable, but it does not need to be included as a contribution in this paper. Finally, I suggest further revising the narrative in the paper to explicitly state the conditions and assumptions for any claimed theoretical advancements to prevent readers from being misled. Regarding the author's statement that other methods cannot achieve certain results due to the lack of guarantees in existing proofs, I understand this writing style—it is sometimes valid. However, in the context of this paper, it is clearly debatable. After all, new assumptions can also be applied to many previous algorithms to derive entirely new conclusions. I have raised my score to 6.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730317949064

thread_type:
official_review

<thread_object>
creator: reviewer_gubz

<content>
summary:
This paper proposes propose a Generalized Heavy-Ball Momentum (GHBM), proving that it enjoys an improved theoretical convergence rate w.r.t. existing FL methods based on classical momentum in partial participation, without relying on bounded data
heterogeneity. This paper also presents FEDHBM as an adaptive, communication-efficient by-design instance of GHBM. Extensive experimentation on vision and language tasks, in both controlled and realistic large-scale scenarios, confirms our theoretical findings, showing that GHBM substantially improves the state of the art, especially in large scale scenarios with high data heterogeneity and low client participation.

strengths:
1. The authors present a novel formulation of momentum called Generalized Heavy-Ball (GHBM) momentum, which extends the classical heavy-ball (Polyak, 1964), and propose variants that are robust to heterogeneity and communication-efficient by design.
2. The theoretical convergence rate of the proposed GHBM is given for non-convex functions, which extend the previous result of Cheng et al. (2024) of classical momentum, showing that GHBM converges under arbitrary heterogeneity even (and most notably) in partial participation.
3. The authors empirically show that existing FL algorithms suffer severe limitations in extreme non-iid scenarios and real-world settings. In contrast, FEDHBM is extremely robust and achieves higher model quality with significantly faster convergence speeds than other client-drift correction methods.

weaknesses:
Although this paper is written very clearly, the Assumption 4.3 (Cyclic Participation) is important and should be further discussed in this paper.

questions:
Although this paper is written very clearly, the Assumption 4.3 (Cyclic Participation) is important and should be further discussed in this paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645954844

thread_type:
official_review

<thread_object>
creator: reviewer_zb5t

<content>
summary:
The paper proposes new federated learning (FL) algorithm variants—GHBM, LocalGHBM, and FedHBM—which generalize previous momentum-based methods by allowing the momentum term to be computed using multiple past iterations, rather than just the most recent one. The authors theoretically prove that the iterates of GHBM, with access to a stochastic gradient oracle, converge to a first-order stationary point in a mean-squared sense for non-convex functions, even with partial client participation and arbitrary client dissimilarity. Experimentally, they demonstrate that for large neural network training, their methods outperform existing FL approaches, especially in scenarios with high client dissimilarity and low client participation.

strengths:
The paper advances previous theoretical results that relied on assumptions of bounded client heterogeneity or full client participation. The authors’ method generalizes the Fed-Avg-M approach from Cheng (2024), achieving the same convergence rate without requiring full client participation. Experimental results are compelling, with the proposed methods showing significant improvements over existing methods used for comparison.

weaknesses:
The main weakness of the paper’s contribution is the similarity to (Cheng 2024). The theoretical novelty is limited as the authors admit the analysis is based on (Cheng 2024). Further, this paper compares against Fed-Avg-M from (Cheng 2024), but makes no mention of the SCAFFOLD-M algorithm from the same paper which has convergence guarantees in the desired setting of unbounded client heterogeneity and partial client participation. SCAFFOLD-M incurs a slightly worse convergence rate but does not require cyclic client participation. The line in the introduction section that reads“existing momentum-based FL methods still theoretically rely on bounded heterogeneity in partial participation…” is thus misleading. The absence of discussion of SCAFFOLD-M in terms of theoretical guarantees and experimental performance is a major weakness, especially since (Cheng 2024) showed that SCAFFOLD-M performed much stronger than Fed-Avg-M under severe client heterogeneity.

Although the experimental results are strong relative to the other methods chosen for comparison, the discussion on why the proposed methods significantly outperform existing methods is unsatisfying. It is mentioned throughout the paper that existing methods perform poorly in the “large-scale” setting and the proposed methods improve in this area, but this difference does not seem to relate to the theoretical convergence results established in the paper.

questions:
Can you compare your theoretical and experimental results to SCAFFOLD-M? It seems to be the most natural algorithmic competitor to the proposed methods. Also, can you offer an explanation as to why the proposed methods outperform existing methods in the “large-scale” setting especially? This point is made many times but does not seem to be supported by the theoretical discussion. Are the proposed methods just converging faster than other methods, or is there something about using the generalized momentum method that biases the networks to learn better quality solutions?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721621198

thread_type:
official_review

<thread_object>
creator: reviewer_f1jy

<content>
summary:
This paper introduces a novel Generalized Heavy-Ball Momentum (GHBM) technique for federated optimization with partial participation, enhancing convergence without relying on the bounded heterogeneity assumption. The proposed method is both theoretically and experimentally shown to improve convergence rates. Empirical results across various federated learning benchmarks demonstrate benefits over previous approaches.

strengths:
1. The paper tackles the important problem of federated learning optimization in heterogeneous settings and under partial participation.

2. The proposed technique builds on the classical momentum mechanism in a straightforward manner, making it promising for practical implementation in federated learning systems.

3. The authors conduct diverse experiments showing the advantages of GHBM (and variations) across multiple realistic computer vision and language tasks.

weaknesses:
**W1.** The **literature overview** in Subsection 3.2 omits some recent works, such as ScaffNew [1] and related papers, that achieve accelerated communication complexity in heterogeneous settings through control variates similar to SCAFFOLD. In addtion, there are recent papers on federated optimization [2, 3], which leverage momentum as a local correction term to benefit from second-order similarity (such as MIME and CE-LSGD).

**W2.** On **presentation and clarity.** Many equation links, particularly those referencing the appendix (e.g., Eq. (122)), are not functional, making the paper cumbersome to follow. 
The symbol $G$ is used without definition in Lemma 4.4.
Referring to notations introduced in the appendix, such as equations (11) and (12), makes understanding Lemma 4.6 challenging without flipping back and forth between sections.

A clear, intuitive explanation of Assumption 4.3 would be beneficial to the reader.

**W3.** The local version (LocalGHBM) may not be well-suited for cross-device settings, as it necessitates storing local states similarly to SCAFFOLD, which has been identified as a limitation by Reddi et al. (2021).


The following work has already been published:

> Konstantin Mishchenko, Eduard Gorbunov, Martin Takác, and Peter Richtárik. Distributed learning with compressed gradient differences, 2019.
___

[1] Mishchenko, Konstantin, et al. “Proxskip: Yes! local gradient steps provably lead to communication acceleration! finally!.” International Conference on Machine Learning. PMLR, 2022.

[2] Mishchenko, Konstantin, et al. “Federated learning under second-order data heterogeneity.” https://openreview.net/forum?id=jkhVrIllKg, 2023.

[3] Karagulyan, Avetik, et al. “SPAM: Stochastic Proximal Point Method with Momentum Variance Reduction for Non-convex Cross-Device Federated Learning.” arXiv preprint, 2024.

questions:
1. What is actually the problem being tackled in this paper: cross-device or cross-silo?

2. The results in Figure 1 are somewhat counterintuitive, showing higher variation for the “iid” case. Could the authors clarify this observation?

3. The work by Patel et al., (2022) suggested an algorithm CE-LSGD and provided some lower bounds, which seem to be circumvented by this submission. Please explain how is it possible? Please also provide a comparison to that paper.

4. How important is the cyclic participation Assumption 4.3 for achieving the main theoretical contribution of the paper?

5. Do the convergence rates benefit from introducing local steps, as seen in ScaffNew? From equation (122), it seems that the local step size $\eta_l$ must decrease inversely with the number of local steps $J$.

6. How does GHBM interact with server momentum, which has been shown to be essential for achieving good convergence rates?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733310252708

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global message after authors-reviewers discussion period

comment:
We thank the reviewers for all the suggestions and comments, in this global message we want to highlight the main points of our rebuttal and discussion, and summarize what changes have been made to the original manuscript.

---
# Contribution and Significance
The main contribution of this work is the GHBM algorithm, **a novel, principled momentum-based FL algorithm, with provable convergence guarantees under arbitrary heterogeneity in (cyclic) partial participation and significant empirical improvement of performance in various scenarios and heterogeneity levels.**

In particular, we clarified that **GHBM is the first algorithm that can use only momentum to converge under arbitrary heterogeneity in cyclic partial participation**. This is our theoretical contribution, as the **best known previous work on (classical) momentum in FL proves the same result only in full participation**.

**This result follows from the key novelty in the design of GHBM**: the momentum is updated with the average of the last $\tau$ pseudo-gradients such that, when $\tau=\frac{1}{C}, in cyclic participation the update is similar to the one classical momentum has in full participation. This is validated by additional experiments provided in section C.9, where **GHBM in partial participation is shown to behave similarly to FedCM in full participation**.

We believe this is a **significant contribution to the field**, as it expands the understanding of the effectiveness of momentum techniques in FL.

---
# Changes on the manuscript

### On the use of cyclic participation assumption.
We clarified that **the use of cyclic participation assumption is legitimate for GHBM** because, while it is needed for proving our theoretical result (as explicitly indicated in Table 1), in our analysis the convergence under unbounded heterogeneity comes explicitly from the algorithmic structure of GHBM (i.e. setting $\tau=\frac{1}{C}$ is necessary).
While we leave as an open problem proving a lower bound for FedCM, its best known analysis in [1] cannot prove a similar result in a context of partial participation. Moreover, **all the experimental results we present strongly indicate that we should expect FedCM to be inevitably affected by heterogeneity in partial participation**. In the final revision **we will make it clear that it is still open to formally prove.**

### Comparison with other methods.
Our work mainly compares with state-of-art FL algorithms that use only momentum to deal with heterogeneity, as the point of our paper is proving that (our generalized) momentum can extend strong convergence guarantees in cyclic participation.

As is, we initially left out of our study SCAFFOLD-M [1], which relies on SCAFFOLD controls variates to solve the issues of heterogeneity. Following the suggestion of Reviewer Zb5t, in the last revision **we added both theoretical and experimental comparison with SCAFFOLD-M**, showing that in difficult scenarios it inherits the empirical weaknesses of SCAFFOLD.

### Practical implications
In our work, we directly demonstrate the effectiveness of GHBM in **large-scale real-world scenarios**, by conducting extensive experimentation on challenging settings characterized by a **large number of clients** (on the order of $10^4$-$10^5$) and **very low participation** (e.g. $C<1\\%$). We clarified that, even if the theoretical guarantees are provided in cyclic participation, **all the experiments are carried under the common random uniform sampling**, to guarantee a fair comparison with other algorithms that do not require that assumption.
Together with strong empirical performance, **the simplicity of GHBM makes it  preferable to more complex alternatives** and further corroborates its value for practical application in FL systems.

**Other modifications.**
Other modifications include the correction of the caption of Figure 1, where we mistakenly wrote “iid” for the case of $\alpha=0$ and “non-iid” for the case of $\alpha=10.000$, as it is obviously the opposite. We want to clarify that **the presented results align with the reviewers’ expectations.**
As other modifications, we fixed all the non-functioning equation links and improved readability by reporting the definitions of the involved quantities lemmas 4.4 and 4.6.

---
# Discussion outcome
We thank the reviewers for engaging in discussion with us. Their valuable feedback has allowed us to enhance the overall quality of our paper.
Based on the most recent interactions with each reviewer, we believe we have adequately addressed all outstanding concerns that initially led to some negative scores.

**In light of our efforts to address these concerns, including providing new experiments, clarifications, and discussions of related work, we kindly encourage the reviewers to reconsider the significance and potential impact of our submission in the next discussion phase.**

[1] Cheng et al., Momentum Benefits Non-IID Federated Learning Simply and Provably, ICLR 2024
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper shows the benefits of computing momentum based on longer history in the FL setting, under heterogeneity and cyclic client participation. Overall, the results are mostly clearly presented and the assumptions clearly stated, especially after the discussion phase. The concerns from the reviewers primarily stem from the similarities with certain previous analysis and limited expiation of why the empirical results are better than closely related methods. The authors have provided detailed responses clarifying certain aspects, however some concerns remained. The concerns regarding the participation assumptions were discussed at length and while it can considered somewhat strong, it is arguably close to where sota analysis is.