Paper: Scaff-PD: Communication Efficient Fair and Robust Federated Learning (ICLR-2024)

Reviewer comment (Novelty):
The formulation seems not to be quite new given existing studies. Although the authors show that their formulation is a generalization of existing ones, directly applying the distributionally robust objective (i.e., client weights optimized over a set of candidate distribution) is not new. The $\Phi$ regularizer seems not to add too much new things into the formulation as well. In sum, although the formulation is general, the idea and individual components are not new.

Author response:
Thank you for your suggestions on clarifying the contributions of our work. We agree that the formulation is not new, and we are leveraging formulations developed in the DRO literature. Instead, **our work is trying to tackle the optimization problem of how to solve the DRO problem more efficiently in federated learning**. In below, we would like to reiterate our key theoretical contributions of this work: 
- New provable algorithms for solving DRO in federated learning. In this work, we (1). design new algorithm components to perform gradient correction; (2). analyze and control the drift error on each device such that the proposed algorithm converges to the optimal solution; and (3). The resulting algorithms – SCAFF-PD – are guaranteed to converge to the optimal saddle point solution of DRO problems in federated learning. 
- Faster convergence rates in both strongly-convex-strongly-concave setting and strongly-convex-concave setting. Our analysis shows that SCAFF-PD enjoys faster convergence rates compared to existing algorithms [BDK+2022, HTF+2021] in the federate learning literature, for both strongly-convex-strongly-concave setting and strongly-convex-concave setting.

Reviewer comment (Novelty):
The algorithms do not incorporate many new stuffs. For example, the SCAFFORD bias correction, the extrapolated step for min-max acceleration have been well studied. The proposed algorithm seems to apply to the general federated min-max problem. However, since the DRO-FL problem has its special structure, i.e., $\lambda$ is linear in the first part of the total objective and the regularizer may still have some benign and dedicated structure in terms of $\lambda$. It may be more interesting to explore such architecture to get better results rather than in the worst-case setting of federated minmax case. In sum, the developed algorithms are not existing enough.

Author response:
There seems to be a misunderstanding. **Our proposed algorithm SCAFF-PD is exploring and leveraging the special structure** —$\lambda$ is linear in the first part of the total objective— of the DRO problem. For example, our proposed algorithm **cannot** be directly applied to solve generic min-max optimization problems in federated learning. The update for the dual variable (Eq. (4.1)) relies on the special structure (the coupling term is linear in $\lambda$). Furthermore, we were able to achieve the accelerated convergence rates that are faster than all existing algorithms because we leveraged the special structure of the DRO, **whereas most existing min-max federated learning algorithms [BDK+2022, HTF+2021] are designed for solving generic min-max optimization problems and do not achieve the accelerated rate developed in this work**.


[BDK+2022] A. Beznosikov, P. Dvurechenskii, A. Koloskova, V. Samokhin, S. U. Stich, and A. Gasnikov. Decentralized local stochastic extra-gradient for variational inequalities. Advances in Neural Information Processing Systems, 35:38116–38133, 2022.

[HTF+2021] C. Hou, K. K. Thekumparampil, G. Fanti, and S. Oh. Efficient algorithms for federated saddle point optimization. arXiv preprint arXiv:2102.06333, 2021.

Reviewer comment (Theory):
Theory is not strong. The analysis is only conducted for strongly-convex-(strongly)-concave settings, while most practical examples work under the nonconvex setting. In addition, there are quite a few works on general federated nonconvex-concave/PL/strong-concave setting (see https://arxiv.org/abs/2302.04249 and related works therein). It would be great to discuss why only strongly-convex-(strongly)-concave settings are studied here. Again, the analysis is mainly developed for general case without taking the DRO min-max structure into account.

Author response:
We would like to re-iterate our key contributions in this work. In the strongly-convex-(strongly)-concave (SC-SC or SC-C) setting, existing algorithms cannot match the convergence rate of the centralized one and achieve the accelerated convergence rates, and our work resolves this open problem. We believe such theoretical contributions are valuable additions to the existing federated min-max optimization literature. 

**Why study strongly-convex-(strongly)-concave settings?** As shown in previous work [YWK+2022], one can convexify the non-convex problem and solve the reformulated convex problem, which leads to improved performance compared to previous methods. Therefore, we believe solving the SC-SC and SC-C efficiently makes theoretical and empirical contributions to the literature. 

On the other hand, we strongly agree it would be interesting to investigate how to deal with the non-convex setting in future work, and it is out of the scope of this work. Please refer to **A2** for discussion on leveraging the special structure of DRO studied in our work.

Reviewer comment (Evaluation):
Experiments can be made more convincing. The improvements in Table 1 are not that significant. In some cases, the worst-20% results are improved with too much loss in average accuracy. For example, for $\alpha=0.1$, worst-20% increases from 29.5 to 29.78, but the average accuracy drops from 46.11 to 41.23.

Author response:
Thank you for your suggestion. We will provide additional experiments on SCAFF-PD in our camera-ready version.

We thank you again for your thoughtful review and comments. Please let us know if you have other questions or comments.

Reviewer comment (Writing):
What is $\bar{\tau}$ in Condition 5.1 and how do you set $\gamma_0$?

Author response:
**A1**: We set $\bar{\tau}$ and $\gamma_0$ such that they satisfy the condition described in (B.35), for example, $\bar{\tau} = O(1/L_{xx})$ and $\sigma = O(L_{xx}/L^2_{\lambda x})$. $\bar{\tau}$ and $\gamma_0$ can be interpreted as the initial values for series {$\{\tau_r, \gamma_r\}$}.

Reviewer comment (Theory):
Lemma B.2 is wrong and the proof is also wrong, which leads to the soundness of Theorem B.6 and Theorem 5.1. Specifically, to prove  
\[
t_r\left(\frac{1}{\tau_r}+ \mu_{\boldsymbol{x}}\right) \geq \frac{t_{r+1}}{\tau_{r+1}},
\]  
we should not start from rewriting this equation, which already assumes that this inequality is correct. The procedure I did is shown below:  
\[
t_r\left(\frac{1}{\tau_r}+ \mu_{\boldsymbol{x}}\right) =\frac{t_r}{\tau_r} (1+ \mu_{\boldsymbol{x}}{\tau_r}) = \frac{t_r}{\tau_r} \frac{\gamma_{r+1}}{\gamma_r}  = \frac{\sigma_r}{\sigma_0\tau_r} \frac{\gamma_{r+1}}{\gamma_r}  = \frac{\sigma_r}{\sigma_0\tau_r} \frac{\sigma_{r+1}/\tau_{r+1}}{\sigma_r/\tau_r}   = \frac{\sigma_{r+1}/\tau_{r+1}}{\sigma_0} =  \frac{t_{r+1}}{\tau_{r+1}}.
\]  
That means there is no greater than or equal to relationship, the two sides always equal to each other. Meanwhile, $\frac{t_r}{\sigma_r} = \frac{t_{r+1}}{\sigma_{r+1}}$.

Author response:
**A2**: We would like to clarify that, as you mentioned the two terms equal to each other, the ‘greater or equal to’ result still holds, therefore the argument and proof of Lemma B.2 is correct. We will revise the proof of Lemma B.2 to improve the presentation clarity.

Reviewer comment (Theory):
From the assumptions and the definition of all symbols, I recognized that $\tau_r = J\eta_l\eta_g$, which indicates that $\tau_{r+1} = J\eta_l\eta_g$, too, if I did not miss anything. That means, for equation B.38, with $\frac{t_r}{t_r+1} = \theta_r$, we must have $V_{r+1}\leq \theta_{r+1}Z_{r+1}$, which however, is based on the condition that $\theta_{r+1}>1$. However, $\theta_{r+1} = \frac{\sigma_r}{\sigma_{r+1}} = \frac{\gamma_r\tau_r}{\gamma_{r+1}\tau_{r+1}} = \frac{\tau_r}{(1+ \mu_{\boldsymbol{x}}\tau_r)\tau_{r+1}} = \frac{1}{1+ \mu_{\boldsymbol{x}}{\tau_r}}<1$, which contradicts the assumption, which leads to me doubt about the soundness of the theorems. Please correct me if I am wrong.

Author response:
**A3**: We would like to clarify that $\tau_r$ is a parameter that may change over the communication round, which is defined in Lemma B.2. The local learning rate $\eta_{\ell}$ is a parameter that depends on $\tau_r$, where $\eta_{\ell}=\tau_r / (J \eta_g)$ (defined in Lemma A.3). $\eta_g$ is the global learning rate and $J$ is the number of local steps, and both $J$ and $\eta_{g}$ are constants. When $\tau_r$ changes over communication rounds, the value of $\eta_{\ell}$ will change accordingly. Therefore, $\tau_r/\tau_{r+1}\neq 1$ and $\frac{\tau_r}{(1+\mu_x \tau_r)\tau_{r+1}} \neq \frac{1}{1+\mu_x \tau_r}$, which means the last equality in **Q3** does not hold. Thus, the value of $\theta_{r+1}$ does not contradict with the assumption. 

**To summary, the proof of the Theorem is correct.** We will clarify the definition of the local step $\eta_{\ell}$ in our revision.

Reviewer comment (Evaluation):
The experiments on synthetic and real datasets seem to contradict each other. For example, Figure 2 shows that larger $\rho$ leads to higher convergence rate while Figure 3 as well as in the analysis says “Meanwhile, the experimental results suggest that smaller $\rho$ leads to faster convergence w.r.t. worst-20% accuracy for our algorithm.” Which one is correct and why larger/smaller $\rho$ leads to faster convergence?

Author response:
**A4**: There seems to be a misunderstanding. Figure 2 and Figure 3 are visualizing two different quantities: 
- **The y-axis Figure 2 represents the distance to the optimal solution.** For the experiments in Figure 2, we consider the synthetic data setup and study the convergence of different algorithms --- whether they converge to the unique optimal saddle point solution as well as the convergence rate of each algorithm.
- **The y-axis Figure 3 represents training/test accuracy of the model.** For the experiments in Figure 3, we study the performance of different algorithms on real data, and we evaluate the training/test accuracy instead of the distance to the optimal solution.

*Both experiments are correct.* We apologize for the confusion. Thank you for your suggestion, and we will provide the experimental results of the distance to optimal solutions for real data (the ones shown in Figure 3) in our camera-ready version.

Reviewer comment (Theory):
In equation (3.1), what is the range of each $\lambda_i$? Is there any constraints on all $\lambda_i$'s, for example, does $\sum_i \lambda_i =1$?

Author response:
**A5**: The $\lambda$ vector is constrained in the set $\Lambda$, for example, the $\Lambda=\Delta^{N}=${$\lambda:\sum_i \lambda_i = 1, 0\leq \lambda_i \leq 1$}. The $\lambda$ vector is always contained in the simplex $\Delta^{N}$, and our formulation is flexible and also allows adding additional constraints to $\lambda$, for example, refer to the constraints described in the equation before 'Finally, we can recover the Q-FL' sentence on Page 4, i.e., $\Lambda=$ {$\lambda \in \Delta, \lambda_i \leq 1/(\alpha N)$}.

Reviewer comment (Writing):
What is the definition of $\Delta$ in equation (3.3)?

Author response:
**A5**: The $\lambda$ vector is constrained in the set $\Lambda$, for example, the $\Lambda=\Delta^{N}=${$\lambda:\sum_i \lambda_i = 1, 0\leq \lambda_i \leq 1$}. The $\lambda$ vector is always contained in the simplex $\Delta^{N}$, and our formulation is flexible and also allows adding additional constraints to $\lambda$, for example, refer to the constraints described in the equation before 'Finally, we can recover the Q-FL' sentence on Page 4, i.e., $\Lambda=$ {$\lambda \in \Delta, \lambda_i \leq 1/(\alpha N)$}.

Reviewer comment (Writing):
What is the definition of $D(\boldsymbol{\lambda}, \boldsymbol{\lambda}^r)$ in equation (4.2)?

Author response:
**A6**: $D(\lambda, \lambda^{r}) = \| \lambda - \lambda^{r} \|^{2}$, which is described in Lemma A.4. We will highlight this in our notation part in the main body.

Reviewer comment (Theory):
In Corollary 5.2, what is the definition of $\epsilon$?

Author response:
**A7**: $\varepsilon$ is the upper bound of the square of the distance between the $R$-th round iterate $x^{R}$ and the optimal solution $x^{\star}$, i.e., $\| x^{R} - x^{\star} \|^{2}\leq \varepsilon$.

Reviewer comment (Writing):
What is the definition of $L_{xx}$ in equation (5.3)?

Author response:
**A8**: $L_{xx}$ is the smoothness parameter for the primal variable $x$, and the smoothness condition is defined in Eq (3.4). Thank you for pointing this out, we have added the $L_{xx}$ condition to Theorem 5.1 and Theorem 5.5 in our revision.

Reviewer comment (Writing):
In the Introduction part, you use $T$ to indicate the convergence while in the theorems presented, you use $R$. Shouldn’t those be consistent?

Author response:
**A9**: Thank you for your suggestion. We used $T$ and $R$ to denote the number of communication rounds. We have changed $T$ to $R$ in our introduction section.

Reviewer comment (Experiments):
For Figure 2, how is the local update iteration $J$ chosen?

Author response:
**A10**: In Figure 2, we chose $J=100$, which is described in Sec 6.1 – ‘For both algorithms, we set the number of local steps to be 100 and select the algorithm parameters through grid search’. We also studied the effect of local update iterations $J$ in Figure 5 (b) in the appendix.

Reviewer comment (Novelty):
The algorithm design and theoretical analysis rely on SCAFFOLD, encompassing the hypothesis and proof framework. This extensive reliance on prior work may potentially diminish the originality and contribution of the proposed method in this paper.

Author response:
Thank you for your suggestions on highlighting our technical contributions and the novelty of this work. We summarize the technical challenges as follows: the gradient correction used in SCAFFOLD only considers minimization problems; we need to design new gradient correction components in SCAFF-PD, which are **significantly different from the ones in SCAFFOLD, for solving the minimax problems in DRO**. The key novelty of our paper is to tackle the above technical challenges by (1). designing new algorithm components to perform gradient correction for min-max DRO problems; and (2). analyze and control the drift error on each device such that the proposed algorithm converges to the optimal solution. To this end, our analysis shows that SCAFF-PD enjoys faster convergence rates compared to existing algorithms for solving DRO problems in the federate learning literature.

Reviewer comment (Theory):
A notable issue arises in the algorithm design, as it necessitates two times of communications with nodes at each round, transmitting distinct content. This introduces a huge communication overhead. Additionally, contradictory to the federated context where nodes may join or leave the network at any time, the proposed algorithm must consistently maintain a stable participation at every round. This operational requirement may pose challenges to the practical applicability of the proposed algorithm.

Author response:
First of all, our goal of this work is to develop a faster federated optimization algorithm to reduce the number of communication rounds and leverage the local computational resource on each device. We agree that our algorithm requires two rounds of communication for updating the model parameters. However, compared to all existing algorithms in the literature, our proposed algorithms achieve the fastest convergence rate (w.r.t. communication) in terms of the condition numbers of the problem. Thus, when the condition number (i.e., $L^2_{\lambda x}/(\mu_{x}\mu_{\lambda})$) becomes larger, our algorithm achieves provably faster convergence than existing ones. Furthermore, through our experimental evaluations, we found *our algorithm achieves faster convergence performance on real-world tasks/datasets* **when using the same communication rounds as existing algorithms**. 

Regarding the setting where each node could join or leave the network at any time, this is definitely an interesting problem to investigate. However, this is out of the scope of this work. Even for the simpler case we studied in this work, previous existing algorithms cannot even match the convergence rate of the centralized algorithms, and our work resolves this open problem. It would be interesting to investigate how to deal with this challenging situation in future work.

Reviewer comment (Theory):
The first thing is related to the paragraph about choosing $\psi$ and $\Lambda$. From the convergence analysis, $\Lambda$ is a bounded set. However, there is no discussion about this in the main part.

Author response:
Because in the dual update step, Eq. (4.1), the center (the *server*) is solving the constrained subproblem (in the second row of Eq. (4.1)), where $\lambda \in \Lambda$. In our proof, we applied the fact that $\lambda^{r+1}$ is the optimal solution for the constrained subproblem. Thank you for your suggestion, we will add more discussion about the constraint in the main part.

Reviewer comment (Reproducibility):
In the main part there is no expression for local stepsize.

Author response:
The local step $\eta_{\ell}$ is defined in Lemma B.2. The local learning rate $\eta_{\ell}$ depends on $\tau_r$, where $\eta_{\ell}=\tau_r / (J \eta_g)$, where $\eta_g$ is the global learning rate and $J$ is the number of local steps, and both $J$ and $\eta_{g}$ are constants. We will clarify the definition of the local step $\eta_{\ell}$ in our revision.

Reviewer comment (Theory):
The formulation of Theorem 5.5 is not full. There is no word about the smoothness of function $f$.

Author response:
Thank you for your suggestion, we have included the smoothness conditions in Theorem 5.5 in our revision.

Reviewer comment (Experiments):
In the experiments, there is no comparison between the performances of SCAFF-PD and ProxSkip.

Author response:
Thank you for your suggestion. We will include and discuss the comparison between SCAFF-PD and ProxSkip in the main body of our camera-ready version.

Reviewer comment (Evaluation):
Could you explain whether your proposed method has a speed up related to number of clients and number of local steps, which is observed for SCAFFOLD?

Author response:
This is again a great observation. We answer this question in two parts: 
- **Speedup wrt clients $N$.** Our setting and proof allows any $\Lambda \subseteq \Delta^{N}$ i.e. any subset of the simplex. In particular, suppose we set $\Lambda = e_1$  i.e. all the weight is placed only on client 1 with $\lambda_1=1$ and $\lambda_j=0$ for all other clients $j\neq 1$. In this case, our DRO formulation (1.1) reduces to $\min_{x} f_1(x)$. This is effectively becomes a single client optimization, and we should not expect any speedup w.r.t. $N$. More generally, the DRO problem can potentially result in different clients being weighed drastically differently, with the worst case putting all weight on a single client. Thus, we cannot have a speedup wrt the number of clients $N$.

- **Speedup wrt local iterations $J$.** Our result with some minor modifications does in fact show a speedup wrt the number of local iterations $J$, similar to SCAFFOLD. We had chosen to focus mainly on the dependence of $\varepsilon$ in our proofs and in comparison with prior works. This is because all prior work for the DRO problem only had a $1/\varepsilon^2$ rate, compared to our much faster $1/\varepsilon$ (matching SCAFFOLD). However, the reviewer is absolutely right that for a fine-grained comparison to SCAFFOLD, we need to focus on the constants $J$ and $N$ as well.
To show a speedup wrt $J$, we can make one of the two following modifications: 

  - **Option (i)** The computation of $c_i = g_i(x)$ in SCAFF-PD needs to be done over a batch-size $J$ every round. This would increase the gradient complexity $\times 2$ but does not change the theoretical rates. However, this results in $c_i$ and $c$ having a much smaller variance of $\zeta^{2}/J$, and can be used to improve Lemma A.2 to derive $\mathcal{E}_{r} \leq \frac{12\tau^2}{\eta_g^2} \mathbb{E}[\| \nabla_x \Phi (x^{r}, \lambda^{r+1})\|^{2}] + \frac{15\tau^2}{\eta_g^2}(1+\chi)\frac{\zeta^{2}}{\color{red}J}.$ This improvement wrt $J$ can then be propagated in the rest of the proof where Lemma A.2 is used - e.g. in Eqn. (B.17). Doing this gives us the required speedup w.r.t. $J$ in all our results, including Corollary 5.6.

  - **Option (ii)** Alternatively, we can use the gradients computed in the past round to compute . This is exactly the same as *option (ii) in the SCAFFOLD algorithm* --- see Eqn. (5) in [KKM+2020]. This would be more efficient since we do not need to compute additional stochastic gradients for $c_i$, however the proof is even more complex and messy. 

We chose the simpler approach in this work since the focus of our work was the min-max component of the DRO. We would be happy to extend our analysis for this version of SCAFF-PD if the reviewer thinks this would significantly improve our work.

Reviewer comment (Theory):
Please, add explanation for derivation of eq. (B.59) from eq.(B.57).

Author response:
We were using the fact that $\Delta_r \leq 0$ in the first inequality of Theorem B.9. We will add this explanation after Eq. (B.59) for better clarity.