Paper: Federated Learning Under Second-Order Data Heterogeneity (ICLR-2024)

Reviewer comment (Theory):
The definition of $w_{k+1} \approx \text{argmin } \phi_k(w)$ in Algorithm 1 is very vague; how is $\approx$ quantified? Further, for the theoretical results in Section 3.2, as far as can be understood from the proof of Lemma 3, $w_{k+1}$ has been taken to be exactly $\text{argmin } \phi_k(w)$. The authors have not discussed this point. It is unreasonable to expect that $\phi_k(w)$ can be exactly minimized, and it would be preferable to see a result where $w_{k+1}$ is an approximate minimizer of $\phi_k(w)$, for example, an $\varepsilon$-stationary point of $\phi_k(w)$ as written in Section 3.1.

Author response:
Thank you for providing your feedback on our submission. 
1. It appears to us that the main concern comes from solving the local subproblem, so we answer it here in detail. The definition $w_{k+1}\approx \arg\min_w \phi_k(w)$ requires two things that are used in the proof: 1) descent in the local objective, $\phi_k(w_{k+1}) \le \phi_k(w_k) - \frac{1}{2\eta}\Vert w_{k+1} - w_k\Vert^2$ (used in the proof of Lemma 1), and 2) $\Vert\nabla \phi_k(w_{k+1})\Vert\le \varepsilon$ (used in the proof of Lemma 3). Both conditions can be achieved by applying gradient descent to $\phi_k(w)$. As is known for nonconvex functions, it will take $H=O(\varepsilon^{-2})$ iterations to satisfy both conditions.

Reviewer comment (Theory):
The complexity of minimizing $\phi_k(w)$ to obtain $w_{k+1}$ has been completely ignored in Section 3.2, which is an important aspect that has not been considered. This could have been captured if there were local steps in the proposed algorithm; specifically, it would be more convincing if there were some $H \geq 1$ local steps of (S)GD on $\phi_k(w)$ to approximately minimize it, and the final convergence bound in Theorem 1 would then be a function of both $K$ and $H$. While the emphasis is probably on reducing communication complexity, deriving results assuming exact minimization of local surrogate functions is unrealistic and too strong. It would be preferable if the theoretical results captured the inexactness of the approximate minimizer of $\phi_k(w)$, as this would likely reveal practically important tradeoffs.

Author response:
Thank you for providing your feedback on our submission. 
1. It appears to us that the main concern comes from solving the local subproblem, so we answer it here in detail. The definition $w_{k+1}\approx \arg\min_w \phi_k(w)$ requires two things that are used in the proof: 1) descent in the local objective, $\phi_k(w_{k+1}) \le \phi_k(w_k) - \frac{1}{2\eta}\Vert w_{k+1} - w_k\Vert^2$ (used in the proof of Lemma 1), and 2) $\Vert\nabla \phi_k(w_{k+1})\Vert\le \varepsilon$ (used in the proof of Lemma 3). Both conditions can be achieved by applying gradient descent to $\phi_k(w)$. As is known for nonconvex functions, it will take $H=O(\varepsilon^{-2})$ iterations to satisfy both conditions.

Reviewer comment (Novelty):
Two relevant works, [1] and [2], have not been discussed. Step 5 in Algorithm 1 of this paper is virtually the same as Step 4 in Algorithm 1 (PAGE) of [1], although [1] is for the centralized setting. [2] proposes a federated version of PAGE called FedPAGE, which, unlike SABER, does not need to minimize any surrogate function. Moreover, FedPAGE also attains $\mathcal{O}(\frac{\sqrt{M}}{\varepsilon^{2}})$ complexity. Thus, the algorithms in [1] and [2] appear similar to SABER but achieve the same complexity without the compute-intensive part of minimizing $\phi_k(w)$. A detailed comparison of SABER with (Fed)PAGE, including experiments, is needed.  
[1]: Li, Zhize, et al. "PAGE: A simple and optimal probabilistic gradient estimator for nonconvex optimization." International Conference on Machine Learning. PMLR, 2021.  
[2]: Zhao, Haoyu, Zhize Li, and Peter Richtárik. "FedPAGE: A fast local stochastic gradient method for communication-efficient federated learning." arXiv preprint arXiv:2108.04755 (2021).

Author response:
Thank you for providing your feedback on our submission. 
2. Thank you for providing additional references. We have cited SARAH, Spider, and EF21 for the use of recursive updates, but we see that PAGE uses the same mechanism too. As for the complexity of SABER and FedPAGE, the key difference is that the latter operates under $L$-smoothness of the problem. As discussed in prior work on second-order heterogeneity, for instance by Khaled & Jin (2023), $L$-smoothness implies second-order heterogeneity with constant $\delta = 2L$, but the reverse is not true. In other words, second-order heterogeneity is strictly more general and $\delta$ can be orders of magnitude smaller than $L$. Moreover, as far as we can see, if we apply PAGE to the local subproblem of SABER without regularization, then we immediately recover FedPAGE as a special case.
3. We will update the experiments for the camera-ready version to add FedPAGE as a baseline as the reviewer requested. We tried to do it before the end of the discussion period, but couldn't due to the short discussion period and since the FedPAGE paper doesn't provide code.

Reviewer comment (Experiments):
There are concerns regarding the experimental results. In Section 4.1, $p=0.5$ and 50 out of 100 clients are used in line 8 of Algorithm 2 for CIFAR-10, meaning 50% of the clients are used in every other round; it is unclear if SABER’s benefits are due to this. For FEMNIST, the total number of clients is not stated (the authors should provide this), but it is mentioned that 100 clients are used in line 8 of Algorithm 2 again with $p = 0.5$. The value $p = 0.5$ seems too large, and the algorithm would be practically useful only if it works with much smaller values of $p$ such as 0.1. An ablation study showing the effect of varying $p$ would be useful. It is also unclear if the hyperparameters were tuned for the other algorithms. Additionally, as mentioned in Point 4, there should be empirical comparisons with FedPAGE given its similarity to SABER.

Author response:
Thank you for providing your feedback on our submission. 
3. We will update the experiments for the camera-ready version to add FedPAGE as a baseline as the reviewer requested. We tried to do it before the end of the discussion period, but couldn't due to the short discussion period and since the FedPAGE paper doesn't provide code. 
4. Regarding the hyperparameter search, we tuned the hyperparameters for all methods, not just ours.  
R2: 3. We agree that using a smaller $p$ is of interest for making the method practical but we didn't have enough time to present an ablation study as we had to work on the rebuttals and participate in paper discussions as reviewers too.

Reviewer comment (Theory):
Solution of the subproblem. The subproblem can be non convex. In this case GD for the subproblem would require $\epsilon^{-2}$ iterations to find $\epsilon$-stationary point. That might significantly slow down the convergence of the whole procedure. Moreover, it is assumed that the subproblem could be solved exactly. In my point of view, this not only makes the method impractical, but also simplifies the theoretical analysis.

Author response:
As the reviewer correctly pointed out, gradient descent would require $\epsilon^{-2}$ iterations to find the $\epsilon$-stationary point of the local objective. This is, however, expected since employing GD makes it a method from the first-order class, for which there exist a lot of lower bounds. Woodworth et al. (2023) studied a method similar to ours in non-federated context, and the complexity improvement was in having the smoothness constant replaced by the heterogeneity constant, while in terms of $\epsilon$, it remained the same.

Reviewer comment (Theory):
Abstract, first bulletpoint on page 3, page 6. I think it should be $O(\delta \epsilon^{-2} \sqrt{M})$ instead of $O(\delta \epsilon^{2} \sqrt{M})$.

Author response:
Thanks for pointing out the typo in the abstract, we also found similar typos in the text, and fixed all of them.

Reviewer comment (Experiments):
The method is claimed to be stateless by design, but what the authors mean in practice is not clearly defined nor tested in the experiments.

Author response:
Thank you for asking about the stateless nature, it is an important feature of our method and it seems we forgot to properly define it when first mentioning it. Concretely, we assume that the clients do not carry states from round to round. For example, they may randomly drop out and never come back again.

Reviewer comment (Evaluation):
Lack of discussion on what happens in the case of node errors, communication drop-outs, and/or stragglers.

Author response:
Issues such as node errors, communication drop-outs, and stragglers that the reviewer mentions are indeed common in federated learning. The proximal term (regularization) used in the subproblem of SABER was originally motivated in FedProx to mitigate the issue of stragglers. For example, when one client does much more work than the others, the update magnitude is still limited thanks to the regularization. We will add a  discussion on the consequences of using regularization in the paper.

Reviewer comment (Presentation):
Have the authors tampered with the template? I found that I could not search using my accessibility tools and no text was selectable - is there any particular reason for doing this? I find that people that use specialty equipment might struggle with this...

Author response:
We used the standard ICLR template but we split the pdf into the main body and the supplementary and this seems to have affected the text. We will split the pdf differently to fix the issue.

Reviewer comment (Reproducibility):
Why the code was not included in the paper submission? I understand that this can be made public after review but I think it is an essential for proper review and it can be attached as supplementary material not disclosed publicly.

Author response:
We apologize for not including the code in the submission, we were working hard on the manuscript and didn’t find enough time to clean the code to be included.

Reviewer comment (Writing):
Clarify what the authors mean by "stateless by design" as per weaknesses above.

Author response:
Thank you for asking about the stateless nature, it is an important feature of our method and it seems we forgot to properly define it when first mentioning it. Concretely, we assume that the clients do not carry states from round to round. For example, they may randomly drop out and never come back again.

Reviewer comment (Reproducibility):
What happens if a client drops during the computation? Is that handled?

Author response:
If a client drops during computation, the fact that our theory works under second-order heterogeneity implies that there should be no bias as long as we can sometimes get an update from that client to maintain a meaningful estimate of the global direction $v_k$. In other words, the penalty for the client dropping is that the update would be less useful, but it will still be unbiased.

Reviewer comment (Experiments):
Does the framework assume equal participation of clients? What happens in the case that there are clients that "dominate" the training? Does this affect the end result (perhaps, due to the uniform sampling used...?)?

Author response:
If some clients dominate the training in terms of how much work they can do, this issue is already handled using the regularization term. If you mean that they dominate in terms of participation, then the required change would be to sample them less frequently. We can add a discussion on how the theory works in this case.

Reviewer comment (Novelty):
The novelty of the paper seems limited. The local objective function constructed by the author is an intuitive combination of FedProx and SCAFFOLD. The similar idea can be found in Lin et al. (2023)

Author response:
As the reviewer pointed out, the local objective construction is intuitive and is similar to SCAFFOLD and FedProx, but we believe this is a positive feature. Not only it connects naturally to the existing literature, it also lets us intuitively explain the meaning of every term in the subproblem. Note that constructing the subproblem itself was never supposed to be the main contribution, it’s rather in the analysis of the method. Moreover, the analysis of our paper doesn’t have anything in common neither with SCAFFOLD nor with FedProx.

Reviewer comment (Novelty):
The algorithm design of SABER is not surprising to the reviewer in the sense that the proposed SABER algorithm uses the similar idea of SVRG to deal with data heterogeneity.

Author response:
Even if there are similarities with SVRG, our analysis is not based on it, so we are not sure why it is seen as a weakness by the reviewer. Furthermore, unlike SVRG, our algorithm does not store an anchoring point that is reset in the outer loop.

Reviewer comment (Writing):
The writing of the paper needs to be improved. In section 3.3, the result for μ-PL objectives should be formulated as a formal theorem.

Author response:
We are happy to rewrite Section 3.3 to improve the structure.