Paper: Communication-Efficient Heterogeneous Federated Learning with Generalized Heavy-Ball Momentum (ICLR-2024)

Reviewer comment (Theory):
In Eq.(5), what is the MVR term? Could the author explain this in detail? Its first part is a sum of $J_i$-step updates, while its second part is a sum of $j$-step updates. I know these two parts come from the additional term from $\theta_{i,j}^t-\theta_i^{t-\tau_i}$ to $\theta^{t-1}-\theta^{t-\tau_i-1}$, but how does it perform as a variance reduction? The $u$ term is the update performed in each local iteration, so the MVR term is the difference between $j$ local updates and $J_i$ local updates where $j\neq J_i$. Please provide the equivalent form in [1] and demonstrate its variance reduction efficiency in the main text.

Author response:
* **W1:** In our initial analysis of the second term in Eq. (5), we drew parallels with the concept presented in [1], perceiving it as a variance reduction mechanism due to observed similarities. We appreciate the Reviewer's keen observation regarding the distinction between these two terms. This insight has led us to a more accurate interpretation and, consequently, a revision in our manuscript. We now clarify that this term in our formulation is intended to penalize the direction of the last updates at round $t - \tau_i$ ​, in contrast to the progressive update of local steps at the current round $t$.
Although at the moment we cannot formally prove that this term acts as variance reduction, we investigated its importance in the ablation study section, specifically in Figure 2, where it is evident that it provides additional speedup w.r.t Local-GHB and reaches better accuracy (see fig. 3 and table 10).

Reviewer comment (Novelty):
The authors claim that they propose the GHB form in Table 1. However, the same updates have been studied in [2] which adopts a multi-step momentum. By setting the specific coefficients as some fixed constants, [2] performs the same updates as the global GHB. Actually, the claimed global GHB in this paper is only a special case of [2]. It also provides some convergence analysis to understand more complicated cases. Therefore, the contribution of this paper seems to only extend the global GHB to the local GHB without any analysis of optimization or generalization, which greatly reduces the novelty of this paper.

Author response:
* **W2:** Thanks for suggesting this work. We were not aware of it and we will discuss it in the paper. We tried to cover the most relevant works in NeurIPS/ICLR/ICML venues. However, [2] seems not to be a peer-reviewed work, as we only found it on arXiv. Please kindly provide us with the correct reference if we are mistaken so that we can appropriately cite it. Please notice that [2] does not perform the same updates of our FedHBM, as we motivate below.  \
**Important differences with [2]:** \
Their $\delta_{t}=-(x_{t} - x_{t-1})/K$ (line 6 of their algorithm 1) takes the difference of models at two subsequent rounds. To perform that operation, the server must send both models to the client, which sacrifices communication efficiency unless all clients participate in each round, which is not a realistic FL scenario. Their clients also need to keep track of this multiple $\delta_{t-j}$, with linear increase (w.r.t. to model parameters) of the needed client storage. This also exponentially increases the number of hyperparameters $\alpha_{j}, \beta_{j}$ to be tuned. In fact, for practicality, authors fix the number of deltas to two, and so four hyper-params to be tuned. For these reasons, we believe that actually [2] is a special case of GHB with hand-tuned $\tau$ (in practice set to two in [2]). The term “multi-step” just refers to the adoption of the Nesterov momentum formulation instead of the common one adopted by us. \
Conversely, the momentum term in FedHBM is computed over a window, whose size is automatically dictated by the client's participation. This requires keeping only one model stored on the client side, and a single momentum factor $\beta$ to be tuned. Our ablation 4.3 extensively tests the effect of the window size $\tau$, highlighting the importance of $tau>1$ and the optimality of the self-tuned $\tau_i$ for communication efficiency. Crucially this is an important difference that sets us apart from [2]. \
Moreover, we show that the FedHBM update rule implicitly contains a momentum variance reduction term. As already mentioned, it gives an additional speedup, especially in the most heterogeneous scenarios. \
**Concerning the experimental validation.** \
    *  **Small scale vs comprehensive evaluation.** The work [2] only validates small-scale computer vision tasks, with only one architecture, which is not specified. Conversely, in our work, we considered both computer vision and NLP tasks, multiple architectures (LeNet, ResNet-20, MobileNetv2), and extended also to large-scale real-world datasets. Such extensive evaluations are important contributions since they show issues not noticed by previous works.
    * **High vs extremely high heterogeneity.** In [2] authors only test under non-extreme heterogeneity (e.g. as low as $\alpha=0.1$), while we showed that in those non-extreme scenarios, even simple baselines (as FedAvgM) are satisfactory (see fig. 1 in our paper).
As such, those scenarios do not render the robustness issues of existing algorithms. Conversely, in our work, we extensively tested the robustness of the algorithms in the most difficult scenarios ($\alpha=0$) which are generally not covered by other works. Such a study on the robustness against extreme heterogeneity is indeed one of the contributions of our work. 
    * **Cross-silo and cross-device.** [2] does not take into account cross-silo and cross-device scenarios, and how their algorithm performs in both cases. We took special care about exploring how our solution works under such conditions, as we presented extended empirical evidence to prove that FedHBM works also in a scenario it was not originally designed for. Large-scale experiments in Table 4 prove that a participation ratio as low as 0.5% is enough to see the practical advantages of our method over the state-of-the-art. As these analyses are important to evaluate the versatility of an FL algorithm for real-world scenarios, they should be considered as a strength of our work.
    * **Reproducibility.** Lastly, [2] lacks necessary implementation details, which together with the absence of official code makes results not easily reproducible. Conversely, we devoted much space in the supplementary to provide the reader with a comprehensive description of the experimental setup and hyperparameter search for all the algorithms we tested (see section B.1, especially table 7). We also attach our code, to underline our commitment to full reproducibility. 

For these reasons, the novelty of our work is by no means diminished by [2], nor the contribution we bring to the FL research community, as we provided value in focusing on underexplored cases, showing failure cases of existing SOTAs, unifying existing momentum-based FL works in one framework (GHB), and finally proposing and extensively validating our new FedHBM algorithm.

We hope that our answer addresses your concerns and that you consider revising your rating of our work for an “accept” decision.

Reviewer comment (Experiments):
In the experiments, it indicates that the non-iid dataset split adopts a $\alpha\rightarrow 0$. While in the appendix, it shows $\alpha=0$. Please unify this statement.

Author response:
* **W3:** Thanks for spotting the typo, we fixed it.

Reviewer comment (Experiments):
The results are incomplete. For instance, in Table.2, 3, 4, there are no results of FedCM and FedADC. However, in Figure.3, the curves of FedCM are stated as an ablation study. As the very important baselines of this paper, the comparison with FedCM and FedADC is very necessary, which are also two SOTA methods among the momentum-based methods. Authors should add their performance tests to this paper.

Author response:
* **W4:** In the first part of section 4.3 we discuss how both FedCM and FedADC are equivalent formulations of GHB with fixed $\tau=1$ (and hence, extra communication): this is shown analytically in table 1, but also experimentally in figure 3, where the performance of GHB $\tau=1$ and FedCM are shown to be equivalent. This is the reason why we didn’t include them in the main results, as we specified in the first paragraph of section 4.1. However, since this may not be immediately obvious to the reader, we added the values of FedCM and FedADC in the main tables, as you suggested, to dispel any doubts about the completeness of our experimental validation.

Reviewer comment (Experiments):
The experiments do not adopt the learning rate decay. On page 15 in the appendix, the authors claim that all experiments do not use a learning rate scheduler for simplicity. However, this will significantly reduce the performance of some algorithms and make comparisons unfair. For instance, the ADMM-based method, i.e. FedDyn, requires to be optimized well enough on the local client. Otherwise, dual variables will be updated with very large biases. A similar phenomenon also happens in SCAFFOLD and even FedCM. I think this is the main weakness of the experiment in this paper. All results only reflect the phenomenon under a fixed learning rate. While in the current machine learning, the learning rate decay in non-convex optimization is very important and one of the major concerns. Although changing the learning rate complicates comparisons, this kind of comprehensive study can broadly reflect the performance of the proposed method.

Author response:
* **W5:** The reason why we did not adopt a learning rate decay is because we considered that this would have caused unfair comparison since algorithms have very different convergence speeds. Please notice that this way of conducting the experiments is in line with most existing well-established works in FL [3,4,5,6,7].  However, to demonstrate that our findings hold even when adopting learning rate decay, we added in the supplementary an additional section showing the accuracy curves for the iid and non-iid training of CIFAR-100 with ResNet-20, with full hyperparameter search the learning rate decay and algorithm-specific hyperparameters. For ease of comparison, we also added in Table 10 the results for iid training of CIFAR-10/100 without learning rate decay.

Reviewer comment (Novelty):
The novelty of this paper is limited. As mentioned in the summary, two existing momentum-based techniques are additively combined with the modification to adjust previous iterate time step parameter. In my opinion, to call the proposed formulation generalized heavy-ball momentum is somewhat a far-stretch.

Author response:
* **W1:** First, let us restate the novelty of our approach. The generalization we claim is based on the fact that, while standard momentum and its declinations in FL use the updates that occurred in the last round, we consider a larger window (GHB) across multiple rounds. While this modification results in a simple implementation, we showed that enlarging this window to a length $\tau>1$ is crucial to developing the robustness against extreme heterogeneity (see fig. 3). On top of this idea, we developed communication-efficient variants of GHB, namely Local-GHB and our proposed FedHBM, which self-control the window length according to client participation and this is optimal both in terms of window length and communication costs (see again Fig 3). 
As an additional point, we frame existing momentum FL algorithms inside this generalization, and show how they are special cases (see Table 1): in this sense, our work brings some order to the literature on momentum in FL, which is part of our contribution.

Reviewer comment (Theory):
Some theoretical analysis not provided. Although averaging has been applied, it is not clear why having this gap or "window" is desirable in general and particularly in the federated learning setting. Motivation/intuition needs to be given.

Author response:
* **W2:** We acknowledge that a formal proof could strengthen our paper. However, we also argue that theoretical convergence rates may only render an incomplete picture, whereas extensive experiments may reveal more than an elaborate yet very specific proof. As a matter of fact, in our experiments, empirical evidence demonstrated that in extreme cases state-of-art algorithms, despite being supported by formal guarantees, happen to actually fail (e.g. see Table 4). \
**Motivation behind a window $\tau > 1$.** \
The main intuition behind our method is that the trajectory of the server updates over a window $\tau>1$ provides a better estimate for the momentum term in a federated setting. Intuitively, as $\tau$ increases,  the momentum term increasingly incorporates information from a broader range of clients. A key observation is that when $\tau$ equals the average period length (e.g. $\tau = \frac{1}{C}$), under uniform client sampling, the momentum term contains the information on the global distribution and hence is optimal.  We experimentally verified this hypothesis, demonstrating its validity in practice as we showed by purposely varying $\tau$ in Figure 3. We revised the manuscript to better describe these details. \
**Is it always beneficial?** \
Yes. According to our experiments, **in FL this is always beneficial** since it improves even in the iid scenario, where the simple FedAvgM is often a strong baseline. However, its effect is most noticeable in extreme non-iid scenarios where computing the momentum over a larger window mitigates the effects of heterogeneity and partial participation (see Fig. 2 and 3)..

Reviewer comment (Reproducibility):
Not clear if the results are reproducible since code is not given.

Author response:
* **W3:** We take reproducibility into great consideration since we provided a precise algorithmic description in Algorithm 1 and accurately reported the datasets, the split procedures, the hyperparameters, and the framework used (see section B and in particular table 7). We did not attach the code to the original submission, but following your suggestion we are attaching it now, we hope this may dispel your concerns.

Reviewer comment (Novelty):
Missing some related work:  
* Xin, R. and Khan, U. Distributed heavy-ball: a generalization and acceleration of first-order methods with grading tracking. 2018.  
* Das, R. et al. Faster non-convex federated learning via global and local momentum. 2022.  
* Kim, G. et al. Communication-Efficient Federated Learning with Acceleration of Global Momentum. 2022.

Author response:
* **W4:** Thanks for reporting to us these works we were not aware of, we will take the time to go into details and we will include them in our revision of the related works.

Reviewer comment (Theory):
The main weakness of the paper lies on the absence of theoretical results.

Author response:
* **W1:** We acknowledge that a formal proof could strengthen our paper. Certainly, a theoretical proof of convergence is one way to support the claims of a method, but it is not the only way. Here we chose to have an empirical approach, where we test extensively FedHBM and existing methods in a wide range of tasks, extreme scenarios, and real-world cases, as we are pleased you appreciated. By doing so we provide evidence that our method consistently outperforms existing methods which are supported by articulated formal proofs [1,2]. We also argue that theoretical convergence rates may only render an incomplete picture, whereas extensive experiments may reveal more than an elaborate yet very specific proof. As a matter of fact, in our experiments, empirical evidence demonstrated that in extreme cases state-of-art algorithms, despite being supported by formal guarantees, happen to actually fail (e.g. the case of FedDyn diverging already known in the literature, and SCAFFOLD and MimeMom in our table 4).

Reviewer comment (Novelty):
The proposed algorithm is a rather simple generalization of the heavy-ball momentum and as a result the novelty is limited.

Author response:
* **W2:**  We agree that our algorithm is actually simple to implement, and we consider it as a strength. In fact, more often than not, very complex and articulated methods require careful tuning, which may be case-specific, and are prone to introducing coding errors. On the contrary, we think that _simplicity is the ultimate sophistication_ [Leonardo Da Vinci]. 
We also believe that simplicity should not be considered as an indication of a lack of novelty.
To the best of our knowledge, we are the first to propose a generalization of the heavy-ball momentum computed over a period $\tau>1$, and we showed that it has a significant impact on performance in FL (figure 3). Within this framework (GHB), FedHBM allows to capture **locally** (i.e., on the client side) global information that encompasses more than one round, without communication overhead. 
We analyzed its practical advances under a wide range of scenarios, to support our claims of FedHBM being superior to the state-of-the-art FL methods. These results, as well as the practical demonstration of the failure of well-established FL algorithms in extreme non-iid cases, are themselves a novelty that adds to the knowledge base of the community, and points to new problems to be considered.

Reviewer comment (Experiments):
In the cross-device setting the algorithm requires a 'proper' initial model to achieve the required improvement.

Author response:
* **W3:** Results in Table 4 use the formulation proposed in algorithm 1, and nonetheless our method is in most cases the best, so the practical variation is not strictly necessary to achieve good performance. Our point is that, although our algorithm is primarily designed for cross-silo scenarios, it can still be used for cross-device scenarios, as the variation we proposed can help obtain a speedup earlier (see Figure 4).

Reviewer comment (Evaluation):
In table 2 it seems that FEDDYN achieves higher target accuracy (90%) in the seemingly more challenging cross-device setting. I would appreciate it if the authors could elaborate on that.

Author response:
* **W4:** Table 2 reports the number of rounds to reach a target accuracy w.r.t the accuracy of centralized training (reported in table 6), hence the lower the number the faster the algorithm. In a cross-device scenario, FedDyn reaches 90% of the accuracy reached by centralized training in 7600 rounds, which is worse than our FedHBM, which needs only 6510 rounds to reach the same result.
Namely, table 2 reports the speed while Tables 3 and 4 report the final performance of the algorithms. Together, these results show that our algorithm is the most communication efficient and that it yields the best model quality.

Reviewer comment (Presentation):
In the plots of Figure 3 and Figure 4 (a) it is hard to distinguish between some of the methods. I would recommend to either utilize different colors or increase the scale.

Author response:
* **W5:** Thank you for your suggestion, we updated the manuscript accordingly.

Reviewer comment (Writing):
Page 3, paragraph 3, in "..existing algorithms can be express as special.." replace express by expressed.

Author response:
* **Minor issues:** Thanks for spotting the typos, we fixed them. On page 6, par. 4, $k$ simply stands for $\times1000$. As we understand it can generate confusion, we replaced it with $\alpha=10.000$
We hope that our answer addresses your concerns and that our revision will contribute positively to the evaluation of our work.

Reviewer comment (Writing):
Page 6, paragraph 4, in "$\alpha = 10k$" define $k$.

Author response:
* **Minor issues:** Thanks for spotting the typos, we fixed them. On page 6, par. 4, $k$ simply stands for $\times1000$. As we understand it can generate confusion, we replaced it with $\alpha=10.000$
We hope that our answer addresses your concerns and that our revision will contribute positively to the evaluation of our work.

Reviewer comment (Presentation):
In table 5  for $C\approx 0.5$ FEDAVG's performance is in bold which appears to be a typo.

Author response:
* **Minor issues:** Thanks for spotting the typos, we fixed them. On page 6, par. 4, $k$ simply stands for $\times1000$. As we understand it can generate confusion, we replaced it with $\alpha=10.000$
We hope that our answer addresses your concerns and that our revision will contribute positively to the evaluation of our work.

Reviewer comment (Reproducibility):
The paper does not provide any insight on how to set the values of the step sizes \eta and \beta. In addition, the authors do not mention how they set it for their experiments. This is a major limitation to practical use of the algorithm.

Author response:
* **W1-W2-Q1:** For clarity, let us restate the hyperparameters involved in our method. $\eta$ is the server learning rate, $\eta_l$ is the client learning rate, $\hat{\beta_i}$ is our momentum factor, defined as $\frac{\beta C}{J_i}$, where $J_i$ is the number of local steps and $C$ is the participation ratio ($0<C\leq1$). The hyperparameters to be tuned are $\eta$, $\eta_l$, and $\beta$.
For each method, the hyperparameters have been searched via a grid search. As we take into great consideration reproducibility, in Table 7 we reported our extensive hyperparameter search for all the algorithms, indicating in bold the best values.

Reviewer comment (Experiments):
In relation with aforementioned issue, the authors do not discuss how other hyperparameters were set for competitors, thus comparisons are not easy to interpret (was their hyperparameter optimization for all algoritms ?)

Author response:
* **W1-W2-Q1:** For clarity, let us restate the hyperparameters involved in our method. $\eta$ is the server learning rate, $\eta_l$ is the client learning rate, $\hat{\beta_i}$ is our momentum factor, defined as $\frac{\beta C}{J_i}$, where $J_i$ is the number of local steps and $C$ is the participation ratio ($0<C\leq1$). The hyperparameters to be tuned are $\eta$, $\eta_l$, and $\beta$.
For each method, the hyperparameters have been searched via a grid search. As we take into great consideration reproducibility, in Table 7 we reported our extensive hyperparameter search for all the algorithms, indicating in bold the best values.

Reviewer comment (Evaluation):
Overall, the rationale for the choice of the results presented (nb of rounds to reach accuracy level, final model accuracy) is not completely clear. For instance, the computational cost per round is not discussed, so it is difficult to understand the implication of these results. In addition, the non-iid setting is not clear. The authors mention using a Dirichlet distribution, which suggests drawing vectors of probabilities to assign samples to clients, but I don't see how this would lead to heterogeneity (only size imbalance between clients).

Author response:
* **W3-Q2:** We assumed the computational cost per round to be equal for all algorithms, as usually, communication cost is the bottleneck. By doing so, we acted conservatively, as Mime, MimeMom, and MimeLiteMom actually require multiple forward and backward passes on the local dataset, and so require much more computations (we refer to [1] for additional details on their algorithm). All the other algorithms (including ours) have similar computational costs since they just require addition/subtraction to model parameters, which are not expensive. As a matter of fact, this is reflected in the wall-clock training time of our simulations, as indicated in section B.2, under the paragraph “Implementation details and practicality of experiments”.  
* **W3-Q4:** Regarding the non-iid setting, we followed the practice in [2], which is common in FL works. In practice, it consists of drawing class probabilities for each client using a Dirichlet distribution with parameter $\alpha$. The lower the value of $\alpha$, the higher the class imbalance we introduce into local datasets. We acknowledge your suggestion, and we added some more words in the supplementary B.1 to better explain the procedure in [2].

Reviewer comment (Reproducibility):
How were step sizes \eta and \beta set in the experiments ? How about similar hyperparameters of competitors ?

Author response:
* **W1-W2-Q1:** For clarity, let us restate the hyperparameters involved in our method. $\eta$ is the server learning rate, $\eta_l$ is the client learning rate, $\hat{\beta_i}$ is our momentum factor, defined as $\frac{\beta C}{J_i}$, where $J_i$ is the number of local steps and $C$ is the participation ratio ($0<C\leq1$). The hyperparameters to be tuned are $\eta$, $\eta_l$, and $\beta$.
For each method, the hyperparameters have been searched via a grid search. As we take into great consideration reproducibility, in Table 7 we reported our extensive hyperparameter search for all the algorithms, indicating in bold the best values.

Reviewer comment (Evaluation):
Is the computational cost per round similar between algorithms ?

Author response:
* **W3-Q2:** We assumed the computational cost per round to be equal for all algorithms, as usually, communication cost is the bottleneck. By doing so, we acted conservatively, as Mime, MimeMom, and MimeLiteMom actually require multiple forward and backward passes on the local dataset, and so require much more computations (we refer to [1] for additional details on their algorithm). All the other algorithms (including ours) have similar computational costs since they just require addition/subtraction to model parameters, which are not expensive. As a matter of fact, this is reflected in the wall-clock training time of our simulations, as indicated in section B.2, under the paragraph “Implementation details and practicality of experiments”.

Reviewer comment (Experiments):
What does final mean ? Are all algorithms stopped after 10000 rounds ?

Author response:
* **Q3:** _final model quality_ means the top-1 accuracy over the last 100 rounds of training, averaged over 5 independent runs. The number of rounds depends on the dataset and the model used, and it is equal to all the algorithms. For example, for the CIFAR-10 dataset and LeNet, the round budget is fixed at 10.000 rounds (for more details, see Table 8)

Reviewer comment (Experiments):
Could you provide more insight on how you designed the non iid setting ? How do you split the data ? How do you use the Dirichlet distribution ? Is it used to draw probabilities of assignment to each client ? Then, how do you assign samples to clients ? It should depend on the sample value, and I don't see it in your explanations.

Author response:
* **W3-Q4:** Regarding the non-iid setting, we followed the practice in [2], which is common in FL works. In practice, it consists of drawing class probabilities for each client using a Dirichlet distribution with parameter $\alpha$. The lower the value of $\alpha$, the higher the class imbalance we introduce into local datasets. We acknowledge your suggestion, and we added some more words in the supplementary B.1 to better explain the procedure in [2].