Paper: Learning Guarantees for Non-convex Pairwise SGD with Heavy Tails (ICLR-2024)

Reviewer comment (Writing):
The model setup is a bit misleading. In the title and abstract, the authors call the setting heavy tails and heavy-tailed. However, what the authors really study is the sub-Weibull tails (Definition 3.6) that excludes polynomial decays as contrast to many papers the authors cited in the paper. Although there is no unique definition of heavy-tailedness, many readers would assume that you are talking about polynomial decay noise while you actually did not. I suggest you at least mention in the abstract that you are working with sub-Weibull tails.

Author response:
**A1:** Thanks for your constructive comments. In this work, we indeed assume the gradient noises in SGD have sub-Weibull tails. We have modified the related statements in the abstract of our new manuscript to avoid misunderstanding. Other heavy-tailed distributions (such as $\alpha$-stable distributions[1]) will be considered in our future work to further explore the relationship between heavy tails and generalization performance.

[1]T. Nguyen, et al. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. NeurIPS, 2019.

Reviewer comment (Novelty):
The technical novelty and the necessity of working with sub-Weibull distributed gradient noise is not very convincing. Sub-Weibull type distribution can appear in the concentration type inequalities if you want to obtain some high probability guarantees, but this is not what the authors are doing in this paper. The authors simply use very standard \(L^2\) type arguments to study the SGD. In the proofs in the appendix, the authors only need some assumption to bound the 2nd moment, instead of relying on the full definition of the sub-Weibull distribution as described in Definition 3.6. That means all the existing proof techniques in SGD for finite variance setups can be directly used in the paper. If all that is needed is an application of Lemma C.3. with \(p=2\) (a second-moment estimate), why not simply assume that instead of using Definition 3.6? Will all the results still hold? Essentially, Lemma C.3. with \(p=2\) says that if the second-moment depends on the parameter \(\theta\) in a certain way (where \(\theta\) measures the heaviness of the tail), then it will also appear in the final results. The authors are not using the full information of the sub-Weibull distribution, and there are numerous papers about SGD with finite variance. Therefore, including heavy tails in the title and abstract and using that as a selling point is a bit misleading.

Author response:
**A2:** Thanks. For our bounds in expectation, we use the two bounds of second-moment in Lemma C.3, which can be regarded as the dependence of the variance of loss function on the parameter $\theta$ of heavy tails. It should be noted that our current proof can be developed to establish high probability bounds with concentration inequalities, which has been demonstrated at the beginning of “MAIN RESULTS” section. Limited by the response time, we have only provided a relationship between generalization error and on-average model stability, which removes the expectation of Theorem 4.1 (b), and the corresponding high probability generalization bound for Theorem 4.4 in Appendix C.8 of our new manuscript. The orders of these high probability bounds are similar to our previous bounds in expectation. After the end of Rebuttal, we will provide the high probability bounds for all cases.

Reviewer comment (Novelty):
It is unclear whether the same setting has been studied for pointwise SGD. If the answer is yes, the authors should highlight the technical novelty and difficulty of extending the results to the pairwise setting, which does not seem to be very difficult. Moreover, the authors should compare the results with the pointwise setting. On the other hand, if the answer is no, the authors should explain why they do not study the pointwise setting, which is much more common and popular in the literature, and comment on whether similar results can hold for pointwise SGD.

Author response:
**A3:** As far as we know, there is a gap for the non-convex stability-based generalization work under the sub-Weibull gradient noise setting in pointwise learning. [2] utilized uniform convergence tools to study high probability guarantees of non-convex pointwise SGD under the sub-Weibull gradient noise setting. Therefore, we try to analyze non-convex pairwise SGD with algorithmic stability tools. Our analysis can be directly extended to the corresponding pointwise case. In the following, we make comparisons with some stability bounds of non-convex pointwise learning (such as [3,4]) as the following **Vs. Theorem 4.2** and a discussion about the non-convex stability-based generalization work with heavy tails [5,6] in **Vs. Theorems 4.4, 4.6, 4.9**. 

**Vs. Theorem 4.2** [3] developed the uniform stability bound $\mathcal{O}\left((\beta n)^{-1} L^{\frac{2}{\beta c + 1}} T^{\frac{\beta c}{\beta c + 1}}\right)$ under similar conditions, where the order depends on the smoothness parameter $\beta$ and a constant $c$ related to step size $\eta_t$. If $\log T \leq L^{\frac{2}{\beta c + 1} - 1} T^{\frac{\beta c}{\beta c + 1} - \frac{1}{2}}$, the bound of Theorem 4.2 is tighter than theirs. A stability bound $\mathcal{O}\left((nL)^{-1}\sqrt{L+\mathbb{E}_S[v_S^2]}\log T \right)$ [4] was established in the pointwise setting. Although it is $T^{1/2}$-times larger than the bound of Theorem 4.2, [4] made a more stringent limitation to the step size $\eta_t$, i.e., $\eta_t=\frac{c}{(t+2)\log(t+2)}$ with $0<c<1/L$. If we make the same setting, we will get a similar bound with [4].

**Vs. Theorems 4.4, 4.6, 4.9** For the heavy-tailed distributions except for sub-Weibull distributions, e.g., $\alpha$-stable distributions, there are a few papers [5,6] studied the stability-based generalization bounds and made the conclusion that the dependence of generalization bound on the heavy-tailed parameter is not monotonic. Especially, [5] analyzed the dependencies of several constants on heavy-tailed parameter. Our monotonic dependence on heavy-tailed parameter $\theta$ essentially belongs to the dependence of the variance of the gradient for the loss function on heavy tails. Inspired by [5], we will further study the dependencies of other parameters (e.g., smoothness parameter $\beta$) on $\theta$, except for the monotonic dependence in our bounds.

The above statements have been discussed in Appendix C.7.3 of our modified manuscript.

[2]S. Li and Y. Liu. High probability guarantees for nonconvex stochastic gradient descent with heavy tails. ICML, 2022.

[3]M. Hardt, et al. Train faster, generalize better: Stability of stochastic gradient descent. ICML, 2016.

[4]Y. Zhou, et al. Understanding generalization error of SGD in non-convex optimization. Machine Learning, 2022.

[5]A. Raj, et al. Algorithmic stability of heavy-tailed SGD with general loss functions. ICML, 2023b.

[6]A. Raj, et al. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. ALT, 2023a.

Reviewer comment (Theory):
For the main results, it would be better if the authors can provide some discussions on the monotonic (or not) dependence of the bound on \(\theta\), which measures the heaviness of the tail, and provide some insights.

Author response:
**A4:** Thanks. In our work, the sub-Weibull gradient noise assumption is introduced to derive the monotonic dependence of the bound on $\theta$, which is consistent with the papers we mentioned [7]. However, inspired by [8], our dependence on $\theta$ essentially belongs to the dependence of the variance of loss function on $\theta$. Except for the variance, there are also some dependencies of other parameters on $\theta$, such as the smoothness parameter, which need to be developed. These dependencies may be not-monotonic as [8,9]. We have added some related discussions in Appendix C.7.1 of our modified manuscript.

[7]T. Nguyen, et al. First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. NeurIPS, 2019.

[8]A. Raj, et al. Algorithmic stability of heavy-tailed SGD with general loss functions. ICML, 2023b.

[9]A. Raj, et al. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. ALT, 2023a.

Reviewer comment (Novelty):
On page 2, before discussing the development of previous analysis techniques to the heavy-tailed pairwise cases, the authors should cite some works about algorithmic stability and generalization bounds for pointwise SGD with heavy tails from the literature.

Author response:
**A5:** As mentioned in **A3**, there is a gap in the stability-based generalization work for pointwise SGD with sub-Weibull gradient noise. The most related works are [10,11,12]. [10] used uniform convergence tools to derive some high probability generalization bounds of non-convex pointwise SGD with sub-Weibull tails. [11] and [12] analyzed the stability-based generalization performance of pointwise SGD under $\alpha$-stable Levy process. The latter studied the case of non-convex loss function. We have cited these works before talking about developing previous analysis techniques to the heavy-tailed pairwise cases on page 2 in our modified main text.

[10]S. Li and Y. Liu. High probability guarantees for nonconvex stochastic gradient descent with heavy tails. ICML, 2022.

[11]A. Raj, et al. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. ALT, 2023a.

[12]A. Raj, et al. Algorithmic stability of heavy-tailed SGD with general loss functions. ICML, 2023b.

Reviewer comment (Theory):
Assuming the gradient of the loss function is Lipschitz is very reasonable, but assuming both the gradient and the loss function itself are Lipschitz seems quite strong. The authors should add examples and discussions about Assumption 3.7.

Author response:
**A7:** Some previous work assumed the gradient and the loss function itself are both Lipschitz [15,16]. The Lipschitz w.r.t. the loss function itself is quite strong. Therefore, we attempt to remove this assumption with the sub-Weibull gradient noise assumption, which is one of our main contributions. We have cited some examples and modified the related statements to emphasize this contribution behind Assumption 3.7 in our new manuscript.

[15]M. Hardt, et al. Train faster, generalize better: Stability of stochastic gradient descent. ICML, 2016.

[16]Y. Lei and Y. Ying. Fine-grained analysis of stability and generalization for stochastic gradient descent. ICML, 2020.

Reviewer comment (Writing):
In Assumption 3.8 and Assumption 3.9, the authors should add a line or two explaining what the expectations are taken with respect to.

Author response:
**A8:** Thanks. For Assumption 3.8, $\mathbb{E}$ is the expectation w.r.t. the samples indexes $i_t, j_t$. For Assumption 3.9, the definition of PL condition should be $\|\nabla F_S(w)\|^2 \geq  2\mu \left(F_S(w)- F_S(w(S))\right)$ without expectation. We have made the related modifications in Assumptions 3.8 and 3.9.

Reviewer comment (Theory):
Currently, no proof outline (or sketch) is provided in the main text. This makes it difficult to understand the extent of the novelty in the ideas underlying the proof. From my understanding, the proof steps seem to build heavily upon existing ideas from the literature on nonconvex pairwise SGD based learning (Lei et al, 2021b).

Author response:
Thanks for your constructive comments. Firstly, for the convenience of understanding the extent of the novelty in the ideas underlying the proof, we have provided a proof sketch figure at the beginning of Appendix C.

Next, the main differences between our results and the non-convex stability bound of [1] are listed as follows.

**1)Different stability tools:** We use $\ell_1$ on-average model stability instead of uniform stability of [1].

**2)Different assumption:** We make a sub-Weibull gradient noise assumption to remove Lipschitz condition, which is one of our main contributions. While [1] doesn’t consider it. Besides, we compare the bound with PL condition and the one without PL condition to theoretically analyze the effect of PL condition. 

**3)Minibatch SGD:** Our analysis for SGD is extended to the minibatch case, while [1] doesn’t consider it.

**4)Better expectation bounds:** [1] provided a uniform stability bound $\mathcal{O}\left((\beta n)^{-1}L^2T^{\frac{\beta c}{\beta c+1}}\right)$, where the constant $c=\frac{1}{\mu}$ ($\mu$ is the parameter of PL condition). In general, $\mu$ is typically a very small value (Examples 1 and 2 in [2]) which leads to a large value of $c$. Thus, $T^{\frac{\beta c}{\beta c+1}}$ is closer to $T$ than the dependencies of our bounds on $T$. In other words, our bounds are tighter than [1].

**5)High probability bounds:** Our proof can be developed to establish high probability bounds which have been added in Appendix C.8 of our modified manuscript. The orders of these high probability bounds are similar to our previous bounds in expectation. 

We have emphasized these differences in Appendix C.7.2 of our modified manuscript.

[1]Y. Lei, et al. Generalization guarantee of SGD for pairwise learning. NeurIPS, 2021b.

[2]Y. Lei and Y. Ying. Sharper generalization bounds for learning with gradient-dominated objective functions. ICLR, 2021.

Reviewer comment (Theory):
The results are stated in expectation throughout, which is weaker than the "high probability" results that exist in the literature for similar learning problems that use stability based analysis.

Author response:
Thanks for your constructive comments. Our current proofs are not only used to derive expectation bounds but also establish high probability bounds, which has been demonstrated at the beginning of “MAIN RESULTS” section. Limited by the response time, we have only provided a relationship between generalization error and on-average model stability, which removes the expectation of Theorem 4.1 (b), and the corresponding high probability generalization bound for Theorem 4.4 in Appendix C.8 of our new manuscript.  The orders of these high probability bounds are similar to our previous bounds in expectation. After the end of Rebuttal, we will give the high probability bounds for all cases.

Reviewer comment (Theory):
In Assumption 3.8, is the expectation over all sources of randomness, including $w_t$?

Author response:
In Assumption 3.8, $\mathbb{E}$ is the expectation w.r.t. the samples indexes $i_t, j_t$. $w_t$ isn’t included in the expectation. We have made the related modification in Assumption 3.8.

Reviewer comment (Theory):
In Theorem 4.1 (b), it is not clear to me whether this applies to any learning algorithm A, or is specific to SGD? This is because of the gradient noise assumption made therein which suggests that it is for SGD.

Author response:
Thanks. Due to Assumption 3.8, Theorem 4.1 (b) is specific to SGD. We have modified the related statements in our main text.

Reviewer comment (Novelty):
Have the stability bounds of this paper, i.e. $\tilde{O}(T^{1/2}/n)$ and $\tilde{O}(T^{1/4}/n)$ under an additional PL condition, been established in the pointwise setting under the same assumptions or are they completely new?

Author response:
As far as we know, there is a gap for the non-convex stability-based generalization work under the sub-Weibull gradient noise setting in pointwise learning. [4] utilized uniform convergence tools to study high probability guarantees of non-convex pointwise SGD under the sub-Weibull gradient noise setting. Therefore, we try to analyze non-convex pairwise SGD with algorithmic stability tools. Our results can be directly extended to the case of pointwise learning. In the following, we make comparisons with some stability bounds of non-convex pointwise learning [5,6] in **Vs. Theorem 4.2** and a discussion about the non-convex stability-based generalization work with heavy tails [7,8] in **Vs. Theorems 4.4, 4.6, 4.9**. 

**Vs. Theorem 4.2** [5] developed the uniform stability bound $\mathcal{O}\left((\beta n)^{-1} L^{\frac{2}{\beta c + 1}} T^{\frac{\beta c}{\beta c + 1}}\right)$ under similar conditions, where the order depends on the smoothness parameter $\beta$ and a constant $c$ related to step size $\eta_t$. If $\log T \leq L^{\frac{2}{\beta c + 1} - 1} T^{\frac{\beta c}{\beta c + 1} - \frac{1}{2}}$, the bound of Theorem 4.2 is tighter than theirs. A stability bound $\mathcal{O}\left((nL)^{-1}\sqrt{L+\mathbb{E}_S[v_S^2]}\log T \right)$ [6] was established in the pointwise setting. Although it is $T^{1/2}$-times larger than the bound of Theorem 4.2, [6] made a more stringent limitation to the step size $\eta_t$, i.e., $\eta_t=\frac{c}{(t+2)\log(t+2)}$ with $0<c<1/L$. If we make the same setting, we will get a similar bound with [6].

**Vs. Theorems 4.4, 4.6, 4.9** For the heavy-tailed distributions except for sub-Weibull distributions, e.g., $\alpha$-stable distributions, there are a few papers [7,8] studied the stability-based generalization bounds and made the conclusion that the dependence of generalization bound on the heavy-tailed parameter is not monotonic. Especially, [7] analyzed the dependencies of several constants on heavy-tailed parameter. Our monotonic dependence on heavy-tailed parameter $\theta$ essentially belongs to the dependence of the variance of the gradient for the loss function on heavy tails. Inspired by [7], we will further study the dependencies of other parameters (e.g., smoothness parameter $\beta$) on $\theta$, except for the monotonic dependence in our bounds.

We have added these statements in Appendix C.7.3 of our modified manuscript. 

[4]S. Li and Y. Liu. High probability guarantees for nonconvex stochastic gradient descent with heavy tails. ICML, 2022.

[5]M. Hardt, et al. Train faster, generalize better: Stability of stochastic gradient descent. ICML, 2016.

[6]Y. Zhou, et al. Understanding generalization error of SGD in non-convex optimization. Machine Learning, 2022.

[7]A. Raj, et al. Algorithmic stability of heavy-tailed SGD with general loss functions. ICML, 2023b.

[8]A. Raj, et al. Algorithmic stability of heavy-tailed stochastic gradient descent on least squares. ALT, 2023a.

Reviewer comment (Theory):
What is the intuition behind improving prior dependencies on $T$ to $T^{1/2}$ under smoothness and $T^{1/4}$ under PL and smoothness? Specifically, how does the PL allow a transition from $T^{1/2}$ to $T^{1/4}$?

Author response:
Thanks. The key of the transition from $T^{1/2}$ to $T^{1/4}$ is the upper bound of step size $\eta_1$. 

For Theorem 4.4, we set $\eta_1 \leq \frac{1}{2\beta}$ with the reason that the term $\|\nabla F_S(w_t)\|^2$ can be removed without any cost, which can be found in the second inequality of Appendix C.4.

For Theorem 4.6, the upper bound of $\eta_1$ is set to be tighter than Theorem 4.4. In this case, if we directly remove $\|\nabla F_S(w_t)\|^2$, the bound is equal to Theorem 4.4, which is meaningless. So we introduce the additional assumption, PL condition, to decompose $\|\nabla F_S(w_t)\|^2$. The bound of Theorem 4.6 demonstrates that the tighter upper bound of $\eta_1$ combined with PL condition can lead to the tighter stability bound. The above statement has been added to improve readability in Appendix C.7.1 of our modified manuscript.

Reviewer comment (Theory):
What is the dependence of the bounds in Theorems and Corollaries 4.6 to 4.11 on $\mu$? Without any hidden dependencies, it seems that one can let $\mu \to 0$ to prove the same results without the PL condition. Similarly, it seems like the dependence on $\beta$ is hidden in most statements which might be useful to highlight.

Author response:
**1)Dependence on $\mu$:**
The dependence on $\mu$ is indeed hidden. After the modifications on Theorems and Corollaries 4.6 to 4.11, this dependence $1-\prod\limits_{i=1}^t \left(1-\frac{1}{2} \mu \eta_i\right)$ has been uncovered in our modified manuscript. With this dependence, we can not let $\mu \rightarrow 0$ to prove the same results without the PL condition due to the tighter upper bound of $\eta_1\leq \frac{1}{4\beta}$. **A3** provides the consideration about the upper bounds of step size $\eta_1$.

**2)Dependence on $\beta$:** 
All dependencies on smoothness parameter $\beta$ are hidden in our results. In our modified manuscript, we have also uncovered them. For all results from Theorem 4.2 to Theorem 4.11, the dependencies are $\beta^{-1}$, which are similar to [9,10]. We have added the above discussions in Appendix C.7.1 of our new manuscript.

[9]W. Shen, et al. Stability and optimization error of stochastic gradient descent for pairwise learning. arXiv, 2019.

[10]Y. Lei, et al. Generalization guarantee of SGD for pairwise learning. NeurIPS, 2021b.

Reviewer comment (Theory):
The bounds of this paper are in expectation, while many similar bounds in the literature are stated with high probability. It might be useful to add a discussion on the possibility of establishing high probability bounds, especially how such bounds would interact with the heavy-tailed noise of SGD.

Author response:
Thanks for your constructive comments. Our current proofs of the bounds in expectation can be developed to establish high probability bounds, which has been demonstrated at the beginning of “MAIN RESULTS” section. Limited by the response time, we have only provided a relationship between generalization error and on-average model stability, which removes the expectation of Theorem 4.1 (b), and the corresponding high probability generalization bound for Theorem 4.4 in Appendix C.8 of our new manuscript. The orders of these high probability bounds are similar to our previous bounds in expectation. After the end of Rebuttal, we will give the high probability bounds for all cases.

Reviewer comment (Theory):
It seems that a term $|\mathbb{E}[F_S(w(S))] - F(w^*)|$ is missing from the RHS of Equation (4).

Author response:
The term $|\mathbb{E}[F_S(w(S))] - F(w^*)| = 0$ due to the expectation w.r.t. the training set $S$, which has been emphasized before Equation (4) to make it clearer.

Reviewer comment (Theory):
Why is SGD initialized at zero in Definition 3.1? Is this a fundamental limitation or is it possible to handle arbitrary initialization?

Author response:
Since our proof can handle arbitrary initialization, we choose a simple initialization $w_1=0$ for convenience.

Reviewer comment (Theory):
Perhaps there could be a discussion on why Definition 3.5 is called $\ell_1$ on-average stability even though the $\ell_2$ norm is used in the definition.

Author response:
Thanks. Our pairwise $\ell_1$ on-average model stability is defined following Definition 4 [11]. It should be noted that $\ell_1$ means the exponent of the term $\left\|A(S_{i,j}) - A(S)\right\|$ is 1, which is not related to the $\ell_2$ norm.

[11]Y. Lei and Y. Ying. Fine-grained analysis of stability and generalization for stochastic gradient descent. ICML, 2020.