Distributed Gradient Descent with Many Local Steps in Overparameterized Models

Thanks for your positive comments and feedback! Please kindly find our responses below.

1. [Linear models] We agree that non-linear models are more popular in real-world applications. But, understanding the linear models is commonly the first step to establishing theoretical foundations. Many phenomenon in machine learning, such as implicit bias, or double-descent, are first understood for linear models. Many theoretical analysis focuses on the linear models first since we can get clean (closed-form) results and non-trivial conclusions in this setting. The implicit bias of non-linear models is still an active area and not well understood. The results on some specific non-linear models are not closed-form (see Line 538-539 in Conclusion and references in Related Work), and it will be cumbersome to analyze the aggregation of these models in distributed setting. In experiments, we also fine-tuned a pretrained neural network beyond linear regression and classification.

2. [Extension to SGD] It is possible to extend the convergence results of our work to stochastic case based on the implicit bias of SGD in linear models [1]. But understanding the behavior of GD is the necessary first step. Extension to local SGD needs refined analysis of implicit bias of SGD, which will be significantly complicated. In experiments, we use SGD to fine-tune the last layer of a pretrained neural network, which yields similar results on the difference between global model and centralized model, leading us to believe that the behavior of SGD will be similar.

3. [Convergence rate and sample complexity] Firstly, in this paper we are studying a different problem unlike the standard federated learning papers that usually analyze convergence rate of some optimization algorithm. We are showing that, even with a large number of local steps, the centralized solution (with all data together) and the distributed solutions are same in the over-parameterized regime, where this should not have to be the case. In the overparameterized setting, there are many points corresponding to zero training loss. Our work characterizes which specific point Local GD will converge to. Since each local problem is run for a large number of local steps (to infinity) to use the implicit bias results, the general convergence rate will not be well-defined in our setting.

Secondly, we cannot define the sample complexity for our problem, as we have no information about the client distributions $\{D_i\}_{i\in [M]}$, apart from the fact that we have sampled $N$ samples from each of them. Even our results (Theorems 1, 2 and 3) depend on only empirical quantities, for instance, the minima of empirical loss of sampled data points $X_c$ and $y_c$. Obtaining sample complexity by appropriate assumptions on client distributions is therefore an interesting direction for future work.

4. [Numerical error] If the numerical error refers to the error brought by finite number of local steps, our theoretical result do not guarantee that Local GD would converges to a specific point. However, in experiments these do not show up as a factor, so we assume Local GD to be stable to these errors.

5. [L and $\eta$] The number of local steps should be large enough to ensure local convergence. In experiments, we set it to be large enough that the local model is stable. The implicit bias of gradient descent towards the minimum-norm solution in linear models does not depend strictly on the learning rate. We only need the learning rate to be sufficiently small to ensure convergence. For example, in linear regression the learning rate $\eta$ is smaller than $\frac{2}{\lambda_{\max}(X^T X)}$, where $\lambda_{\max}(X^T X)$ is the largest eigenvalue. As this value is also the smoothness constant of the objective function, GD with learning rate smaller than this converges. The exact value of $\eta$ does not affect the implicit bias and our conclusions.

6. [Convergence in direction] For binary linear classification, only the direction of trained model matters, since we use $\mathrm{sign}(x^T w)$ to predict the label. You can rescale the model to any magnitude which does not influence the prediction.

7. [Typos] Thanks for your careful review. We will correct the typos in the revised version.

References

[1] Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3051–3059. PMLR, 2019.

Dear Reviewer Muac,

We appreciate your time in reviewing our work and providing feedback. We would like to know whether our responses have addressed your concerns. We would be happy to provide further clarification.

Thanks for your feedback. Please kindly find our responses below.

1. [Linear models] We have stated in the abstract that we analyzed linear models in the paper. We can add the linear models in the title: though many phenomenon in machine learning, such as implicit bias, or double-descent, are first understood for linear models. In experiments, we also fine-tuned a pretrained neural network beyond linear regression and classification. Many theoretical analysis focuses on the linear models first since we can get clean (closed-form) results and non-trivial conclusions in this setting. The implicit bias of non-linear models is still an active area and not well understood. The results on some specific non-linear models are not closed-form (see Line 538-539 in Conclusion and references in Related Work), and it will be cumbersome to analyze the aggregation of these models in distributed setting.

2. [Number of large local steps in practice] The statement that FedAvg with a large number of local steps works well in practice is motivated by the recent work on the training of LLM [1,2]. FedAvg with 500 local steps are compared to centralized training baselines (see Figure 2 and Table 2 in [1], Figure 3 and Table 1 in [2]). The results show the FedAvg can obtain the same or slightly better performance than centralized baseline. They also test it on non-i.i.d. data distribution, which shows the FedAvg with a large number of local steps converges slower on non-i.i.d. data distribution but gets the same final generalization as i.i.d data distribution (see Figure 5 in [1]). In the revised version we specify this statement in the distributed training of LLM. Our work provides theoretical insight into why FedAvg (Local-GD) works well in practice (in linear models).

Also, there are other existing works to argue FedAvg converges very well in real-world tasks even with heterogeneous data [3]. It also points out that the negative theoretical results about FedAvg are mysteriously inconsistent with practical observations. Our work is also an attempt to challenge the conventional wisdom that FedAvg degrades a lot with a large number of local steps in heterogeneous setting.

3. [Extension to SGD] It is possible to extend the convergence results of our work to stochastic case based on the implicit bias of SGD in linear models [4]. But understanding the behavior of GD is the necessary first step. Extension to local SGD needs refined analysis of implicit bias of SGD, which will be significantly complicated. In experiments, we use SGD to fine-tune the last layer of a pretrained neural network, which yields similar results on the difference between global model and centralized model, leading us to believe that the behavior of SGD will be similar.

4. [Dirichlet distribution] We have added experiments on linear classification with Dirichlet distribution. Note that it is a binary classification problem and we use Dirichlet distribution to unbalance the positive and negative samples. Please see Appendix B in the revised version for the results. The experimental results show that global model obtained from Local-GD can converge to the centralized model in direction with large dimension, which is similar to our reported results on the manuscript.

References

[1] Arthur Douillard, Qixuan Feng, Andrei A Rusu, Rachita Chhaparia, Yani Donchev, Adhiguna Kuncoro, Marc’Aurelio Ranzato, Arthur Szlam, and Jiajun Shen. Diloco: Distributed low-communication training of language models. arXiv preprint arXiv:2311.08105, 2023.

[2] Sami Jaghouar, Jack Min Ong, and Johannes Hagemann. Opendiloco: An open-source framework for globally distributed low-communication training. arXiv preprint arXiv:2407.07852, 2024.

[3] Jianyu Wang, Rudrajit Das, Gauri Joshi, Satyen Kale, Zheng Xu, and Tong Zhang. On the unreasonable effectiveness of federated averaging with heterogeneous data. Transactions on Machine Learning Research, 2024

[4] Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3051–3059. PMLR, 2019.

1. [Linear Models] We have mentioned in the abstract that we deal with linear models for regression and classification. As the reviewer correctly points out, network fine-tuning test is a real-world application of theory of linear models by considering the input to final layer as the "input feature". We will revise our submission to remove the line "beyond linear regression and classification" and emphasize that it is an application of linear regression and classification.

2. [Large number of local steps in practice] The reference [3] requires the local objective functions to be strongly convex, see Lemma 1 in [3] , which in our case would correspond to the empirical covariance matrices on each client, $X_i X_i^\top$, to be invertible. For our overparameterized case, $d \geq mn$, each covariance matrix is rank-deficient and is thus not invertible. Therefore, the theory developed in [3] on the benefit of local steps cannot be applied to our setting. Note that the intuition in [3] about local directions cancelling each other at optimum is not required in the overparameterized case, as long as a centralized solution exists. As for the results in heterogeneous data sets with Dirichlet distribution in Figure 2, the theory in [3] is not applicable due to overparametrization, however, our theoretical results (Theorems 2 and 3) are applicable and explain the results.

Thanks for your positive comments. Please kindly find our responses below.

1. [Convergence Rate] Note that, we are addressing a different problem than convergence rate of a distributed optimization algorithm (unlike most federated learning papers). We are showing that, even with a large number of local steps, the centralized solution (with all data together) and the distributed solutions are $same$ in the over-parameterized regime, where this should not have to be the case. The work you suggested and more work as other reviewers mentioned exploit the property of zero training loss to refine the conventional convergence rate of Local SGD, but do not consider which specific point it will converge to. In overparameterized setting, there are many points corresponding to zero training loss. Our work answers this question by analyzing the implicit bias in distributed setting, which is completely different from the existing work on overparameterized FL. We have included the reference provided, along with some others, in the related works section.

2. [Extension to SGD] It is possible to extend the convergence results of our work to stochastic case based on the implicit bias of SGD in linear models [1]. But understanding the behavior of GD is the necessary first step. Extension to local SGD needs refined analysis of implicit bias of SGD, which will be significantly complicated. In experiments, we use SGD to fine-tune the last layer of a pretrained neural network, which yields similar results on the difference between global model and centralized model, leading us to believe that the behavior of SGD will be similar.

References

[1] Mor Shpigel Nacson, Nathan Srebro, and Daniel Soudry. Stochastic gradient descent on separable data: Exact convergence with a fixed learning rate. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3051–3059. PMLR, 2019.

Thanks for your feedback! Please kindly find our point-to-point responses below.

1. [Question Regarding Eq. (8)-(9)] The centralized model is trained with the collection of all local datasets. Therefore the centralized solution $w_c$ should satisfy the constraint $X_c w_c = y_c$, where $X_c = [X_1^T,\dots, X_M^T]^T \in \mathbb{R}^{MN \times d}$ and $y_c = [y_1^T,\dots, y_M^T]^T \in \mathbb{R}^{MN \times 1}$ are just concatenations of $\{X_i\}_ {i=1}^M$ and $\{y_i\}_{i=1}^M$. Thus it is equivalent to satisfy the constraints $X_i w_c = y_i, \forall i\in [M]$. We simply substitute $X_i w_c = y_i$ in the first line of (4) to get (8).

2. [Local-GD] We do not call this algorithm Fed-Avg, or local-SGD. Local GD as an algorithm has been defined in Algorithm 1. We agree that for the classification problem the loss function includes a weakly regularized term, but we still use a gradient descent with local steps for the algorithm.

3. [Lemma 3] Thanks for your careful review! Yes, the local feasible set is a polyhedral cone consisting of many hyperplanes. It is unbounded since we consider the overparameterized setting ($d \gg N$). It turns out that, for lemma 3 we do not need the boundedness condition; From Theorem 1 and Proposition 8 in [1], we can directly obtain Lemma 3 without the bounded-ness condition in our paper. The requirement is that the local feasible sets are closed, convex sets. Note that the polyhedral cones are closed and convex. Thus Lemma 3 (with refined condition, see revised version) exactly matches our problem in distributed linear classification, and Theorem 2 in our paper also holds.

4. [Theorem 2 does not show exact convergence] Yes, that is why a modified algorithm is proposed (see Theorem 3).

5. [Data heterogeneity] Your proposed setting is just a special case of our setting. What you suggested can be seen as a special case of (2), where ground truth models $\{w_i^*\}_{i=1}^M$ can be exactly same and serve as the global model. The distribution of noise $z_i$ across compute nodes can be different. This setting is an easier case and all of our results are applicable to this special case. In general, there are different kinds of data heterogeneity in federated learning literature, such as label heterogeneity and feature heterogeneity, see Section 2.1 in [2]. In our setting of linear models, the ground truth models that generate local datasets can be different. This is a relatively strong setting as different compute nodes can have different labels even without noise for the same feature $x$.

6. [Works in Overparameterized FL] Thanks for introducing us to this line of work! We mainly focus on the introduction of implicit bias in distributed setting. We have included these works in our discussion on related works. The common point of these works is that they plug in the property that FedAvg can converge to zero training loss into the convergence analysis, and analyze the convergence rate. However, they do not guarantee which point FedAvg can converge to. Our work is different from these works as: 1. We analyze which point the Local-GD can converge to, which is a more elementary problem before obtaining the convergence rate; 2. We use implicit bias as a technical tool to analyze the overparameterized FL.

7. [Conditions on the number of local steps] Yes, the conclusion that the number of local steps should be no larger than $O(\sqrt{T})$ applies to strongly convex and smooth loss functions[3,4]. We specify this in the revised version. Note that [3] requires $O(\sqrt{T})$ local steps for iid client data in this setting, while Theorem 1 in [4] shows that if the loss function is also Lipschitz, for non-iid clients the number of local steps must not exceed $\Omega(\sqrt{T})$.

8. [Algorithm 3 and 4] We have removed Algorithm 3 and 4 in the revised version as you suggested. The linear regression and classification are covered with the same Local-GD (Algorithm 1 and 2) with different loss functions.

References

[1] Patrick L Combettes. Inconsistent signal feasibility problems: Least-squares solutions in a product space. IEEE Transactions on Signal Processing, 42(11):2955–2966, 1994.

[2] Mang Ye, Xiuwen Fang, Bo Du, Pong C. Yuen, and Dacheng Tao. 2023. Heterogeneous Federated Learning: State-of-the-art and Research Challenges. ACM Comput. Surv. 56, 3, Article 79.

[3] Sebastian U Stich. Local SGD converges fast and communicates little. In International Conference on Learning Representations, 2019.

[4] 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, 2020.

Dear Reviewer zEMo,

We appreciate your time in reviewing our work and providing feedback. We would like to know whether our responses have addressed your concerns. We would be happy to provide further clarification.