Generalization of noisy SGD under isoperimetry

Dear Reviewer DPmq,

We sincerely thank the reviewer for their time and their review of our paper.

We would like to provide a clarification on the reviewer's summary:

The authors first show that [LSI] holds, an upper bound on the KL divergence can be derived. They further demonstrate that under [dissipativity] the distribution satisfies the LSI, thus leading to the desired bound.

Our goal is to show a time-independent generalization bound in non-convex settings. In section 5, we assume the function is dissipative, this allows us to show a per-iterate LSI constant. With this result we are able to prove that noisy SGD applied on an unbounded non-convex loss generalizes (KL-stability)and that it is also differentially private (Renyi stability). In section 6, we do not assume dissipativity. We only assume that the target function satisfies the LSI. We derive a technique to establish KL stability without using per-iterate LSI, and only using the final target LSI. This enables us to avoid the poor dimension dependence at the cost of introducing some additional constants in the bound.

Weaknesses:

1. Building on prior results: We would like to stress that theorems 6 and 8 are foundational results. The citation for Theorem 8 for instance refers to Chafai's 2004 book. Theorem 12 which uses Chen et al 2021 result is only applicable because of our preceding analysis. We are building on fundamental results to obtain theorems that are relevant, and our results are not immediate derivations.
2. Assumption 15 For our privacy result, we need a bounded sensitivity assumption because the Renyi divergence introduces expectations that are rather difficult to control. Assumption 15 is standard in the privacy literature (see Def 2.10 in [C]) and it holds for the logistic loss over bounded data or any regularized Lipschitz loss.
3. Dimension dependence: As we are in a non-convex setting dimension dependence is unavoidable. We note this at the end of section 5 and propose section 6 to improve this dimension dependence.

We answer your questions in the following.

1. Isoperimetric vs. Log-Sobolev Inequality: We followed the terminology introduced in Vempala and Wibisono [B] where functional inequalities like the LSI and the Poincare inequality are referred to as isoperimetric inequalities. This follows from the close equivalences that exist between isoperimetric inequalities and functional inequalities [A]. We will add clarifications.

2. What is \Tilde{X}\_k' Thank you for spotting the typo; there is no tilde. It should simply be $X_k'$ . It is a modified expectation under $X_k'$ (see Definition 21). This is precisely where the difference between KL and Renyi divergences is most seen. For KL, the modified expectation is simple, for Renyi the modified expectation is more complex requiring stronger assumptions like Assumption 15.

3. S_k: You raise a good point. This quantity measures the gradient difference on two datasets sampled from the same distribution. It is referred to as the 'sensitivity' since it measures how sensitive the model is to changes in the dataset. As the batch size increases, the gradient estimates approach the population gradient thus making $S_k$ smaller. There is an inverse relationship between the batch size and $S_k$.

We thank you for the time you took to review our work. We remain at your disposal for any further clarifications you may need. If we addressed your concerns we kindly ask that you consider raising your score.

[A] Rothaus. "Analytic inequalities, isoperimetric inequalities and logarithmic Sobolev inequalities." Journal of functional analysis(1985)

[B] Vempala, Wibisono. "Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices." Neurips (2019).

[C] Bok, Jinho, Weijie Su, and Jason M. Altschuler. "Shifted Interpolation for Differential Privacy." arXiv preprint arXiv:2403.00278 (2024).

Dear Reviewer DPmq,

We are sincerely grateful that you acknowledged our response. We would like to respond to your remaining reservations in the following.

1. We had highlighted in our introduction to Section 4 that we first establish a simplified proof template on which our contributions in sections 5 and 6 build. Section 4 is setting up the following sections by identifying Theorem 6 as the crucial entrance point for our new results.

2. The central goal of our work is to establish a bound that does not diverge as the iteration count increases. Unlike some existing generalization bounds, our bound is finite and does not scale with iteration counts $k$ which today exceed $2^{32}$ (see for example training durations in Table 6 of diffusion models). Second, it is this dimension dependence that we mention is the very reason that motivated our section 6 where we mitigate for this dependence.

3. We appreciate that you find our result interesting. We have rewritten the introduction of section 6 in our revision to clarify the proof template. Moreover, the terms in Theorem 18 are explicitly available in the appendix (see equations 8 and 9 on line 503). The poly term in the Corollary is merely a combination of line 1153 with Lemma 32. We restate them here for your convenience, the poly term in Corollary 20.1 is

   $$\left(2\eta\beta^2L^2 + 2(1+\frac{\eta}{2\beta})\right)\left(c_\pi D_\mathrm{KL}(X_0, \pi) + \frac{\eta}{\beta}(8dL^2 + 2\sigma^2)\right) + 8\left(1+\frac{\eta}{2\beta}\right)\left(c_\pi' D_\mathrm{KL}(X_0', \pi') + \frac{\eta}{\beta}(8dL^2 + 2\sigma^2)\right)$$

   The key of our result is not the exact form of the polynomial dependence above but in the fact that it is polynomial. This is the reason we kept the expansions in the appendix. The takeaway message is the following: at the cost of some additional polynomial terms, we can rely on the LSI of the stationary distribution instead of the per-iterate one. This is the key message (see also our response here for further details).

We thank you for your time and we remain available for further discussions.

Dear Reviewer 4Kqp,

We are grateful for your thorough review and feedback. We are encouraged by the strengths you see in our work. In the following, we would like to answer some of your questions.

(W1) Assumption 15 For our privacy result, we need a bounded sensitivity assumption because the Renyi divergence introduces expectations that are rather difficult to control. Assumption 15 is standard in the privacy literature (see Def 2.10 in [A]) and it holds for the logistic loss over bounded data or any regularized Lipschitz loss.

Questions:

(1) In the second paragraph under contributions. Thank you for spotting this typo.

(2) In Lemma 2, is the constant the same one from Assumption 1? Yes the constant is the same and we have fixed the theorem statement.

(3) Some intuitive explanations about the half-step technique: The expansion is done by the gradient step and the contraction is done by the noise. Conditioned on the current step, the KL divergence measuring the gradient step alone would be infinity as it would correspond to the KL divergence between two Dirac distributions. However, if we split the noise in two and add some Gaussian noise to the gradient step, the KL divergence becomes finite and easy to analyze. This is the main reason why we split the noise. A secondary reason is the application of Lemma 17: we need a small amount of Gaussian noise to smoothe the distribution before analyzing the contraction.

(4) Assumption 14. Indeed you are correct, we borrow the assumption from prior work [B, Assumption 3.1] where there is a $\|z - z'\|$ appearing on the RHS. We then absorbed $\|z - z'\|$ into $\theta$ and $D$ by assuming a bounded data distribution. We will add clarifications.

We thank you for the time you took to review our paper and remain at your disposal for any clarifications you may need.

References

[A] Bok, Jinho, Weijie Su, and Jason M. Altschuler. "Shifted Interpolation for Differential Privacy." arXiv preprint arXiv:2403.00278 (2024).

[B] Zhu, Lingjiong, et al. "Uniform-in-time Wasserstein stability bounds for (noisy) stochastic gradient descent." Advances in Neural Information Processing Systems 36 (2024).

Dear Reviewer aiX5,

We are grateful for your thorough, well organized review and your detailed feedback. In the following, we would like to address the points you raise.

1. bounds in the literature To strengthen our Related Work Section 3, we will add a table listing previous bounds considering iterative noisy gradient schemes in non-convex settings. Our focus in our paper is on bounds that apply specifically to noisy iterative algorithms and not on bounds that are agnostic to the learning algorithm, as is the case for most generalization bounds which only consider the function class.

2. lower bounds Our motivation is to amend a gap in previous bounds which explode as the iteration count increases. Our dissipative bound, unlike previous results, matches the lower bound of (Chourasia et al 2021, Theorem 3), for $q=1$, as strongly convex functions are special cases of dissipative functions. Our bounds in section 6, however, are unlikely to be tight but they are a significant improvement over previous diverging bounds.

3. experimental evaluation We believe we are in a reverse situation where practice has already shown that large training times can perform well (see for instance the $2^{32}$ training iterations of state of the art diffusion models) and generalization bounds that aligned with these observations were absent as most existing bounds diverged with training time. Our work is amending theory to make it match what is already practically observed. Our motivation is to show, theoretically, that there exist non-convex settings where long training runs are not harmful.

Technical comments:

1. O(1/sqrt(n)) vs O(1/n) bounds For information-theoretic bounds the fast rate is indeed achievable but it comes with a trade-off: it becomes analytically difficult to obtain a decay factor to control the number of iterations. With current techniques, to obtain the fast rate one has to accept a diverging bound, see Proposition 15 and Corollary 16 of [D]. Our goal is to characterize algorithms that are run for a large number of iterations. For this reason, we cannot use $1/n$ generalization bounds that only depend on the function space nor can we use information-theoretic bounds which are not tractable when applied to noisy iterative algorithms.

2. Conversion of Rényi DP to DP We are very grateful for this reference, we were unaware of this refinement. We will include the improved conversion in Lemma 3.

3. the dependence on the order of Rényi divergence For the privacy result, indeed, your observation that the decay factor in (Chourasia et al 2021) does not appear to degrade with $q$ is correct. The reason is not the two-step analysis but rather the tighter "Renyi log-Sobolev inequality" they use. Their Lemma 3 has a $\partial R_\alpha/\partial\alpha$ appearing which is the derivative of the Renyi divergence with respect to the order. As Renyi divergence is increasing, this term is positive and can be ignored. By disregarding it, we obtain our result. When including it in the analysis, Chourasia et al obtain the equivalent of our Theorem 6 but with a changing Renyi order per iteration (see their equation 66-67). We chose not to operate with changing Renyi orders as it quickly becomes notationally heavy but their tightening can be applied to our result with a similar PDE argument (their lines 59-61).

4. dimension dependence Going from the strongly convex setting section 5.1 to the non-convex setting incurs a dimension dependence. This worst-case dimension dependence in noisy iterative schemes is expected in non-convex settings (see discussion around 4.3 in [B]). However, we mitigate for this poor dimension dependence in Section 6.

5. a lower bound on how the KL or Rényi divergence changes with iterations seems doable Lower bounds in non-convex settings are notoriously difficult to construct. In strongly convex settings, the Gaussian serves to establish lower bounds, in the non-convex setting, the KL divergence between even the simplest non-convex distribution, a mixture of Gaussians, is not analytically tractable. In the space of noisy iterative algorithms, lower bounds are a major open problem (see page 3 of [C]).

Most of the authors comments are well received. I have raised my score accordingly. Here are the concerns that weren't sufficiently addressed.

"For this reason, we cannot use generalization bounds that only depend on the function space nor can we use information-theoretic bounds which are not tractable when applied to noisy iterative algorithms."

Authors goal for the paper is to "characterize algorithms that are run for a large number of iterations" and the generalization bound presented has the desired $O(1)$ dependence on number of iteration. But since the dependence on the dataset size $n$ suffers (perhaps due to fundamental limitation of the analytical tools used, perhaps not), I believe it is important to acknowledge this gap in the paper.

"Our dissipative bound, unlike previous results, matches the lower bound of (Chourasia et al 2021, Theorem 3), for , as strongly convex functions are special cases of dissipative functions."

It seems to me like this is an incorrect assertion. The LSI constant in Theorem 12 hides a dependence on dimension d that does not appear in the lower bound lower bound of (Chourasia et al 2021, Theorem 3). The paper should acknowledge that the bounds presented under dissipativity might not be tight (at least in dimension d and in dataset size n).

Dear Reviewer J6MA,

We are grateful for your thorough review and your feedback. We would like to address the points you have raised.

W. 1. Assumption 19 and LSI They are the same assumption. We made a typographic choice on our part to have theorems preceded by the assumptions they use. It is for this reason that we restate the LSI assumption formally in Assumption 19. Note that dissipativity is a stronger assumption than Assumption 19 (see Figure 2). So in our analysis, Section 5 operates with the stronger dissipativity assumption and Section 6 operates under the relaxed Assumption 19.

W. 2&3. Section 6. We hear the reviewer's feedback and we have rewritten the introduction of Section 6 to make our contribution clear. Both sections 5 and 6 rely on the template in section 4. Indeed, A central component in the analysis is the step-wise contraction with a factor $\gamma$ in Theorem 6. The existence of a contraction factor is what ensures a finite bound. In section 5, we analyze dissipative functions to show that a uniform LSI exists which gives the desired $\gamma$ decay factor. In Section 6, we notice that instead of a clean contraction in Theorem 6, we can instead have

where the additive terms are bounded. This is what we call approximate contraction: Theorem 18 has a similar contraction as in Theorem 6 but with added additive terms. By substituting Theorem 6 with Theorem 18, we obtain our KL generalization result.

what is the main proof framework. The main technical contribution in section 6 lies in showing Theorem 18. Instead of taking a per-iterate LSI like in theorem 6, Theorem 18 only requires LSI of the target measure. To obtain this result, one needs to change, at some point in the analysis, the distribution with respect to which the KL divergence is being taken: from the distribution of the iterates to the distribution of the the target. To perform this change of distribution, we use Lemma 18. Observe there that an expectation under $Y$ on the left can be switched to an expectation under $X$ on the right. Lemma 18 however only applies to sufficiently smooth functions. This is precisely where Lemma 17 comes in. Lemma 17 shows that Gaussian convolution has a smoothing behavior which makes us able to apply Lemma 18. Combining these results with a continuous time analysis of Gaussian convolution yields Theorem 18.

Writing lines 90, 439, 504: Thank you for spotting these mistakes.

We are grateful for your time and remain available for any further clarifications we could provide.