Generalization of noisy SGD in unbounded non-convex settings

We are sincerely grateful for your time, your very detailed review and feedback. We will add a discussion of the limitation you mention.

Thank you for spotting the typos, here are a few details to answer some of the remarks:

• Definition A.1: In integral form the expectation is the following:

$$\Tilde{\mathbb{E}}_{k, q}[h({X'_k})] = \mathbb{E}_{X'_k}[\phi_q(X'_k)h(X'_k)] = \int h(x_k')\phi_q(x'_k) p_{X_k'}(x_k')dx_k'$$

When $q=1$ it cancels out the density of p_X_k'. We hope this clarifies the definition. We will add this in the text. With this clarification, there is no tilde in theorem 4.1, it should just be $X'_k$ in both terms. We will remove it, it was a typo. It is an integral under $X'_k$ modified by $\phi_q$.

• Lemma 5.4: Thank you.
• Theorem 5.5: There should be a $d$ we will fix it.
• Theorem 5.5: We will add the result of Chen et al for completeness so the $\chi^2$ computation is easy to follow. The constant $31/32$ comes from factors $2$ from naive applications of Jensen. It can be tightened as described in appendix C. We will add explicitly the tightening argument.

Questions:

• We will add Li et al. The only reason it wasn't included initially was because their discrete result tightens the Lipschitz dependence but has the same iteration count dependence as result already present in the table. Their secondary analysis was carried out for the continuous version of SGLD and an approximation argument was used to obtain the discrete result. This introduced a need for a vanishing stepsizes which we are trying to relax and it applies to the smaller class of bounded + $\ell_2$ functions.
• Indeed you are correct, in case II of their result they establish a decay factor for the slightly smaller class of Lipschitz + $\ell_2$ functions. This corresponds to a subset of strongly dissipative functions discussed in our section 5.2. Our work relaxes this to simply dissipative. As their result is close so we will highlight it further. We loosen the structural requirements on the function further in section 6 where we show that we only need the LSI to hold without explicit structural requirements like Lipschitz + $\ell_2$.
• The continuous case for bounded + $\ell_2$ has been analyzed before in Li et al but remains open for generally dissipative functions as far as we know. The continuous analysis is interesting as it can help to establish optimality indepedently of discretization error but the we believe the non-vanishing discretization error is the central difficulty of the analysis and so we keep our focus on the discrete setting. Indeed, if we could obtain exact samples of the Gibbs distribution, the generalization bound would be straightforward.

We thank you again for your review and we remain at your disposal for any further discussion.

We sincerely thank you for your time and thorough review of our work. We would like to address some of the points you have raised.

Sampling directly from the Gibbs distribution: You are correct in stating that an exact sample from the Gibbs distribution would attain the fast rate. However, there is no exact sampling algorithm available for non-convex unbounded losses. Approximate samplers either need a vanishing step size or need to have an additive step size dependent term in the generalization bound. This is precisely the technical hurdle we try to overcome, we want a bound that decays to zero when $n$ goes to infinity without needing to set the stepsize as $1/n$. We will add this discussion along with the relevant reference. Removing conditioning would be very interesting and will likely lead to better bounds but we believe it will add complexity in tracking the LSI constant.

reasonable parameter regimes: There is indeed a possible trade-off between composition bounds and our bound in section 5 if the dimension is very large and the iteration count is low in comparison. The dimension dependence in section 5 comes from the fact that we only use dissipativity to establish a worst case LSI constant (which introduces the unfavorable dimension dependence). In section 6 we show that if the LSI is known to not have a good dimension dependence (for example if the mimima are connected), then the exponential dependence can be avoided to yield the polynomial bound in section 6.

Initialization in Corollary 6.6 You are correct that the initial points cannot be diracs in Corollary 6.6. We can warm start the algorithms and initialize them as Gaussians at stationary points of the losses, just like in the sampling literature, to ensure a finite $O(d)$ initial KL divergence.

We are grateful for your time and remain available for any further discussion.

We are sincerely grateful for your thorough review and detailed feedback. We would like to address some of the areas of improvement that have been identified.

Extension of the bound:

You correctly observe that some generalization bounds only hold for an evaluation loss (like the 0-1 loss for example) but the algorithm is run on a different differentiable surrogate. Our work applies to a differentiable surrogate. Often since surrogates are chosen to be upper bounds of the evaluation loss, our result still has implications on the evaluation loss in the following way

loss(test) < surrogate(test) = surrogate(train) + [surrogate(test) - surrogate(train)]

where the final term is the generalization gap we control. Just like Futami and Fujisawa's, our result allows us to control the gap in cases where the surrogate loss and the evaluation are exactly the same. Technically our result bounds the mutual information( denoted $I(W_T, S)$ in Futami and Fujisawa) so we can directly use Theorem 7 from Futami and Fujisawa (also in [B]) and our improved bound will still hold in that setting. We had chosen to omit this component for simplicity but we will add this result as it relaxes the requirement on the surrogate loss: we can include sub-exponential losses and not just sub-gaussian ones.

Connections to gradient variance:

To establish this connection, notice that our stability term is the norm of a difference of gradients and this difference can be expanded into the sum of gradient norms since $\|a - b\|_2^2 \leq 2\|a\|_2^2 + 2\|b\|_2^2$. With this relaxation, our bound can be expressed as a function the gradient norms encountered on the trajectory. We will add a discussion of on using sensitivity term $S_k$ versus gradient norms along the trajectory. The work of Jiang et al shows the difficulty of aligning theoretical generalization bounds with practice and it is important to mention.

Example of loss verifying our assumptions:

We agree that to we should showcase the relaxed requirements by giving a concrete example. We will add the following. Consider a two-layer network with weights $x = [x_2, X_1]$ defined as $NN(x, z) = x_2^\top \sigma(X_1z)$ for a data point $z \in \mathcal{Z}$ and the differentiable logistic sigmoid function $\sigma$. We assume the datapoints are bounded (for example pixels in $[0,1]$ for images). We define a learning task through a Lipschitz loss over the outputs $o = NN(x, z)$ of the network:

$$Lipschitz(NN(x, z))$$

With weight decay (i.e $\ell_2$ regularization), this non-convex loss is dissipative (see section 4 of [C]). It satisfies assumption 5.7 because, for any two $z_1, z_2$,

$$\nabla F(x, z_1) - \nabla F(x, z_2) = Lipschitz'(o_1) \nabla_x NN(x, z_1) - Lipschitz'(o_2) \nabla_x NN(x, z_2)$$

The gradient difference of the neural network is given by the vector

$$\begin{bmatrix} \sigma(X_1z_1) - \sigma(X_1z_2), \\ (x_2\odot \sigma'(X_1z_1))z_1^\top - (x_2\odot \sigma'(X_1z_1))z_2^\top) \end{bmatrix}$$

By virtue of the sigmoid being lispchitz and our bounded data assumption on $z_1, z_1$, this gradient difference can be bounded by the norm $\theta\|x\|_2$ + constants where $\|x\| = \sqrt{\|x_1\|_2^2 +\|X_1\|_2^2}$. We will add this example in the appendix, along with [A]'s result capturing sub-exponential losses like the example above. The Assumption 5.7 is stating that if only the data points differ, then the gradient difference only grows at most linearly with respect to the parameter.

Interpretability of Bounds:

We see our work as an attempt at amending a theoretical gap so it can capture the practically relevant setting of very long training runs on non-convex losses. In essence we are trying to make theory catch up to practice. Our result shows that the presence of noise can compensate for a high iteration count. Concretely, for a fixed choice of $\eta$, we can see that choosing a low $\beta$ (i.e adding more noise) is advantageous for generalization, the trade-off however is that a lower $\beta$ can harm convergence of the training loss. We also see that will add a discussion on the trade-offs between generalization and utility.

We are again grateful for your time and we remain at your disposal to discuss these points further.

---

[A] Futami, Futoshi, and Masahiro Fujisawa. "Time-independent information-theoretic generalization bounds for SGLD." Advances in Neural Information Processing Systems 36 (2023): 8173-8185.

[B] Bu, S. Zou, and V. V. Veeravalli. Tightening mutual information-based bounds on general- ization error. IEEE Journal on Selected Areas in Information Theory, 1(1):121–130, 2020

[C] Raginsky, Maxim, Alexander Rakhlin, and Matus Telgarsky. "Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis." Conference on Learning Theory. PMLR, 2017.

We thank you for your time and review. We would like to address some of your concerns.

We have provided a clear application of our bound in our response to mwzT. Our goal is to have theory catch up to practice where it is already common to train on non-convex losses for several thousand iterations.

Moreover, the form of our result closely matches the one obtained in strongly convex settings. In other words, there are no unecessary terms in our results unlike prior work. However, since there are almost no lower bounds in non-convex settings for sampling algorithms (or noisy iterative schemes), it is difficult to prove optimality of our upper-bound.

We remain at your disposal for further discussion.