Stability and Sharper Risk Bounds with Convergence Rate $\tilde{O}(1/n^2)$

Thanks for your careful review. Here are the replies.

Weakness 1: As mentioned above, the main weakness of the paper comes from the fact that the analysis requires stronger assumptions than do prior works (eg, PL vs SC, requiring lower bounds for the sample size). Here, I think the authors should revise the language in the abstract ("under the same common assumptions"), which is not correct.

Answer: Thanks for your feedback. We will revise the language in the final version to avoid any potential misunderstanding. We appreciate your valuable feedback.

Weakness 2: Since the paper mainly focuses on the dependence on the sample size $n$, it is unclear how the bounds in terms of problem parameters are improved or degraded. Here a comparison with prior work will be helpful.

Answer: Thank you for your suggestion. We have updated the table to include Lipschitzness $M$, smoothness $\gamma$, and strong convexity parameters/PL condition $\mu$ and compared with existing works. We will use this revised table in the final version, as it presents comparisons with prior studies more clearly and straightforwardly.

Weakness 3: Minor.

Answer: Thanks for noting these typos. We'll fix all minor errors in the final version to the best of our ability.

Question: Can the authors provide a comparison with prior works in terms of the dependence on the problem parameters (Lipschitzness, smoothness, strong convexity parameters)?

Answer: Please refer to the answer for Weakness 2 and the table given above.

We’re happy to address any further questions you may have.

Reference

L. Zhang, T. Yang, and R. Jin. Empirical risk minimization for stochastic convex optimization: $O(1/n)$-and $O(1/n^2)$-type of risk bounds. COLT 2017.

Y. Xu and A. Zeevi. Towards optimal problem dependent generalization error bounds in statistical learning theory. Mathematics of Operations Research 2024.

S. Li and Y. Liu. Improved learning rates for stochastic optimization: Two theoretical viewpoints.

Y. Klochkov, and N. Zhivotovskiy. Stability and deviation optimal risk bounds with convergence rate $O(1/n)$. NeurIPS 2021.

Thank you for your thorough review. Here are the replies.

Weakness 1 & Question 1 & Limitation 1: Expectation is usually taken with respect to the randomness of $A$.

Answer: This is a typo. Our Definition 2, consistent with the definition of uniformly-stable in gradients in (Lei. 2023), also needs to take expectation over the randomness of the algorithm $A$. For the ERM and PGD algorithms, since they are deterministic, this does not affect our results. For the proof regarding SGD, we need to add an expectation with respect to the randomness of the algorithm (since the noise introduced by the algorithm's inherent randomness is typically assumed to be independent of the dataset $S$). Our conclusion remains unchanged, meaning for Theorem 3, for any $\delta \in (0,1)$, when $n \geq \frac{16\gamma^2 \log\frac{6}{\delta}}{\mu^2}$, with probability at least $1-\delta$, we have:

\mathbb{E}_A [F(w_{T+1})] -F(w^*) = O \left(\frac{ \left\lceil \log_2 (n) \right\rceil^2 \log(T) \log^5(1/\delta)}{T}\right) + O \left(\frac{ \left\lceil \log_2 (n) \right\rceil^2 \log^2(1/\delta)}{n^2} + \frac{ F ( w^{\ast} ) \log (1 / \delta ) }{n} \right).

Thank you for your thorough review. Here are the replies.

Weakness 1 & Question 1: Is it possible to decrease the order of $\log\frac{1}{\delta}$ terms in Theorems 6 and 7?

Answer: After careful consideration, we acknowledge that the current analysis does not readily support a further reduction in the logarithmic order of these terms. While we recognize the theoretical importance of optimizing such bounds, our present framework relies on concentration inequalities and union-bound arguments that inherently introduce this logarithmic dependency. We agree that this is an important direction for future work, and we are actively exploring alternative techniques that might mitigate these terms.

Weakness 2: Aside from the use of PL condition and the consequential order improvement in convergence rate, the technique to achieve them is based on Klochkov and 147 Zhivotovskiy [19]. Nevertheless, I believe this is minor compared to the $\tilde{O}(1/n^2)$ rate.

Answer: Thank you for acknowledging our work. Although our proof approach was inspired by theirs, we encountered significant challenges. The existing method is fundamentally built upon a one-dimensional framework, while ours operates in a high-dimensional regime. Their technique cannot be generalized to our case.

Weakness 3: Minor/Typos.

Answer: Thanks for noting these typos. We'll fix all minor errors in the final version to the best of our ability.

Weakness 4: The sentence in L267-L269 is quite confusing: "However, the generalization analysis method proposed by Klochkov and Zhivotovskiy [19] does not use smoothness assumption, which only derive ... under strongly convex and smooth assumptions."

Answer: This statement suggests that Klochkov and Zhivotovskiy [19] did not employ the smoothness assumption. Instead, under the same number of iterations, their use of the PGD algorithm can only achieve $O(1/n)$ rates. Thanks for your advice. We will thoroughly polish the paper in the final version.

Question 2: The following questions may be not necessary. It would be greatly appreciated if the authors could satisfy our curiosity...

Answer: The core idea of our approach rests on performing an expansion at first-order stationary point where the gradient vanishes. The point where we perform this expansion is itself sufficiently small, enabling higher-order generalization. However, if we attempt an analogous expansion involving higher-order terms like the Hessian, the Hessian at that point is not guaranteed to be sufficiently small, preventing us from achieving tighter generalization bounds.

We’re happy to address any further questions you may have.

Reference

Yegor Klochkov, and Nikita Zhivotovskiy. Stability and deviation optimal risk bounds with convergence rate $O(1/n)$. NeurIPS 2021.

Thanks for your review. Here are the replies.

Strength 1: The generation bounds were developed before, but the one in this paper is dimension free.

Answer: We disagree with your claim presented in this paper that "generalization bounds were developed before". High-probability stability-based excess risk upper bounds achieving $O(1/n^2)$ have never been presented before; the best existing result only achieves $O(1/n)$.

Furthermore, the stability-based generalization error bound for the gradients we derived is also tightest known. By utilizing the vector Bernstein inequality, the constant coefficient preceding the $\sqrt{1/n}$ term now depends not on the maximum gradient norm $M$, but on a coefficient related to the variance. Critically, this conclusion also holds (is applicable) in non-convex settings.

Weakness 1: The technical novelty is somewhat limited, with proofs heavily relying on appropriate combinations or refinements of existing techniques (e.g., Zhang & Zhou 2019, Klochkov & Zhivotovskiy 2021, Fan & Lei 2024).

Answer: We respectfully disagree with the assessment of limited novelty. We summarize the difference with the mentioned literature.

1. (Zhang & Zhou, 2019) studied optimization problems, whereas we study generalization problems. I guess you intended to cite the paper (Zhang et al., 2017) here.

2. Totally different techniques to (Zhang et al. 2017): They use the uniform convergence techniques to get generalization bounds which dependent on dimension $d$. As the uniform convergence method inevitably introduce a dimension-dependent bound, we improve this with a totally different techniques, the gradient stability, and get our dimension-free $O(1 / n^{2})$ bound.

3. From function stability to gradient stability. Compared to (Klochkov & Zhivotovskiy 2021), they prove their high probability generalization bound via function stability (Definition 6 in (Bousquet and Elisseeff, 2002)), which typically implies only a $O(1/n)$ bound. We improve this by introducing gradient stability (Definition 2 in this paper), which leads to the $O(1/n^2)$ generalization bound demonstrated in this paper.

4. We sharpen the fundamental gradient stability bound via a novel Bernstein's inequality (Theorem 2): In (Fan & Lei 2024), they built the gradient stability (we show this in this paper in Theorem 1), similar to our Theorem 2. But the critical difference is that constants in $1/\sqrt{n}$ is improved from $O\left(M \sqrt{\log{(e/\delta)}} \right)$ to $O \left( \sqrt{\mathbb{E}\_{Z} \left[ || \nabla f(A(S);Z) ||^2 \right] \log{(1/\delta)}} \right)$, since $M$ is the maximal gradient norm. Notably, the proved dependence $\mathbb{E}_{Z} \left[||\nabla f(A(S); Z) ||_2^{2} \right]$ via novel Bernstein's inequality is interpreted as the squared gradient norm of the loss functions on test data, under the optimized model parameters $A(S)$. Thus, it is supposed to be small when the algorithm is good, compared with constant $M$.

Weakness 2: In the strongly convex/ PL setting, if $\mu$ is small, the bound is not good as it has $1/\mu^3$.

Answer:

1. The $\mu$ is not necessary to be small: We only require the population risk $F$ to satisfies PL condition (weaker than the PL condition for empirical risk $F_S$) in Theorem 3, and we verify it is $2$ under a linear regression example in Remark 5.

2. The dependence on $\mu$ is standard: Existing works have also dependence on $\mu$ as presented in the following table.

Weakness 3: The assumption $n\geq\frac{16\gamma^2\log(6/\delta)}{\mu^2}$ and $F(w^*)<O(1/n)$ might not be practical in the strongly convex applications.

Answer:

1. For the condition on $n$, this is a standard condition to expect the machine learning model can generalize: In general, when consider generalization bound, the sample size is expected goes to infinity. Our assumption only requires the sample size is larger than a constant size, if the condition is not satisfied, none of existing literature obtain a small generalization bound.

2. In many cases, the optimal population risk $F(w^{\ast})$ can be extremely small, and it is a standard assumption in related works to reach $O(1/n^2)$ bounds, e.g., Remark 2 in (Zhang et al. 2017), Theorem 3 in (Xu and Zeevi, 2024), Corollary 5 in (Lei and Ying, 2020) and Theorem 22 in (Li and Liu, 2021) and so on. To see this more clear, we check the following mean estimation example.

If we denote the observed data as $X_1, X_2, \ldots, X_n$ drawn from an unknown distribution, we assume that this distribution is bounded (for example, any random variable $X$ that follows this distribution satisfies $a \leq X \leq b$).

To estimate the mean of the random variable, our objective function is defined as $f(m; X) = (X - m)^2$, which yields the usual least squares estimator. Using the definition of M-estimators, our goal becomes:

We can easily verify that $f(m; X)$ satisfies the following properties: it is $4(b-a)$-Lipschitz, $2$-strongly convex, and $2$-smooth with respect to $m$. According to Theorem 4, we know that the mean estimated by this method, which minimizes the least squares loss (or the second central moment), will converge to

as the sample size $n$ increases, where $\mathbb{V}[X]$ is the variance of the distribution. It’s worth noting that since this example pertains specifically to estimating the mean of a random variable, we have no additional parameters involved. Therefore, in this case, $F(w^{\ast} )$ in Theorem 5, which represents the variance of the distribution, can be very small if $\mathbb{V}[X]$ is small.

Question: Comment: The authors should specify what you mean by PGD. I thought it was GD but I have to find the original reference to know it was projected GD.

Answer: At line 31, we explained that PGD stands for Projected Gradient Descent upon its first appearance. At the beginning of Section 4.2, we also provided a detailed description of the Projected Gradient Descent iterative algorithm. We appreciate your feedback and will optimize this in the final version to enhance the paper's readability.

We’re happy to address any further questions you may have.

Reference

L. Zhang and Z.-H. Zhou. Stochastic approximation of smooth and strongly convex functions: Beyond the $O(1/t)$ convergence rate. COLT 2019.

L. Zhang, T. Yang, and R. Jin. Empirical risk minimization for stochastic convex optimization: $O(1/n)$-and $O(1/n^2)$-type of risk bounds. COLT 2017.

O. Bousquet and A. Elisseeff. Stability and generalization. JMLR 2002.

Y. Xu and A. Zeevi. Towards optimal problem dependent generalization error bounds in statistical learning theory. Mathematics of Operations Research 2024.

S. Li and Y. Liu. Improved learning rates for stochastic optimization: Two theoretical viewpoints. 2021

Y. Klochkov, and N. Zhivotovskiy. Stability and deviation optimal risk bounds with convergence rate $O(1/n)$. NeurIPS 2021.

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