Paper: Two Facets of SDE Under an Information-Theoretic Lens: Generalization of SGD via Training Trajectories and via Terminal States (ICLR-2024)

Reviewer comment (Theory):
The limitations of the information-theoretic generalization error analysis approach, such as (implicit/explicit) dimension-dependence and time-dependence, are inherent in the bounds of this paper (although mitigating these limitations can be challenging or important future work). The paper would have been better if there had been a part discussing the limitations and future prospects related to these issues.

Author response:
**Response.** Thank you for this valuable suggestion. We have included a section on limitations and future works to discuss the dimension-dependence. Additionally, we provided additional elaboration on the time-dependence towards the end of Section 3, and consider the time-dependence as motivation for introducing terminal-state bounds in Section 4, which, by their asymptotic nature, are time-independent.

Reviewer comment (Theory):
The approximation of $\Lambda_{w*}$ leads to a dependency on the inverse of the learning rate $\eta$ in the bounds provided in Corollaries 4.1 and 4.2, as well as Theorem 4.2. It seems that these bounds become large or even diverge as the learning rate decreases ($\eta_t \rightarrow 0$). A similar problem appears in the sensitivity-based generalization error analysis for SGLD [1]. I guess that the reason for decaying the learning rate only between the 40,000th and 60,000th iterations in the experiments is to prevent the bounds from diverging, as reducing the learning rate too early might cause divergence before the parameters have converged sufficiently.

Author response:
**Response.**  We would like to clarify that Corollaries 4.1-4.2 and Theorem 4.2 present terminal-state-based bounds. The learning rate considered in these bounds corresponds to the one applied in the final iteration. Furthermore, we think that the reverse relationship between the generalization bound and the learning rate is more favorable. Such bounds suggest that a larger learning rate may lead to better generalization, which aligns with various empirical observations.

Reviewer comment (Theory):
In this paper, the authors make the assumption of a loss function that is both differentiable and satisfies the $R$-subGaussian property. Can you provide specific examples of loss functions that meet both of these assumptions, excluding bounded losses?

Author response:
**Response.** At present, we are not aware of a concrete example of this kind of loss beyond the bounded case. To our knowledge, in order to satisfy the sub-Gaussian assumption, a typical choice in the literature is a bounded cross-entropy loss, such as $-\ln(e^{-\ell_{\max}}+(1-2e^{-\ell_{\max}}))f(x)[y]$, where $f(x)[y]$ is the probability assigned to the label $y$ by $f$ for $x$, as given in [3].

Furthermore, despite the unbounded nature of the original cross-entropy loss, we believe that it indeed has a subGaussian distribution under the SGD training. Specifically, by using the prevalent training methodologies, one can typically ensure that the loss curve consistently descends until it reaches convergence, effectively bounding the loss values during training by the initial value. We have not made an effort in rigorously proving the subGaussianality of the loss. Such an effort would require a careful examination of the SGD dynamics and involve significant technicality. 

[3] Dziugaite and Roy. "Data-dependent PAC-Bayes priors via differential privacy." NeurIPS 2018.

Reviewer comment (Theory):
As I comprehend it, the generalization error bounds presented in Theorem 3.1, Corollary 3.1, Theorem 3.2, Theorem 4.1, and so forth, within this paper, explicitly or implicitly rely on the parameter dimension $d$. I'm curious about the behavior of these bounds as $d$ grows ($d \rightarrow \infty$). Especially, in cases where dimensionality dependence is observed both inside and outside the logarithmic expressions, it appears that these bounds tend to diverge unless factors such as gradient noise are adequately small. Could you please share the authors' thoughts on this matter?

Author response:
**Response.** We note that expressions such as $d\log{\frac{Y}{d}}$ may have an $d$-independent upper bound by the inequality $\log(x+1)\leq x$. For example, in Corollary 3.1 of our paper, as long as the gradient norm is upper bounded by a constant, Corollary 3.1 will vanish regardless of the choice of $d$.

However, we acknowledge that most of our bounds are inherently dimension-dependent, so there exist some settings of $d$ and $n$, particularly when $d$ grows faster than $n$ (e.g., $d>\mathcal{O}(n^2)$), such that these bounds becomes vacuous. Such limitations exist in nearly all the information-theoretic generalizations and sometimes lead these bounds to non-vanishing in the stochastic convex optimization, as recently explored in [4]. While these limitations are not easily overcome fundamentally, the nature of non-convexity in deep learning may provide alternative potentials. For example, although deep neural networks are often overparameterized, recent research, as highlighted in works such as [5,6,7], reveals that SGD/GD only operates within a subspace of the neural networks, referred to as the *intrinsic dimension*. In this context, the SDE in Eq.(5) or $C_t$ itself can be defined within these invertible subspaces, and our derivations remain valid. Moreover, the theoretical characterization of the *intrinsic dimension* remains an open problem. The further development of this research direction will undoubtedly enhance and contribute to the advancements in our work.

[4] Haghifam et al. "Limitations of Information-Theoretic Generalization Bounds for Gradient Descent Methods in Stochastic Convex Optimization." ALT 2023. 

[5] Li et al. "Measuring the Intrinsic Dimension of Objective Landscapes." ICLR 2018.

[6] Gur-Ari et al. "Gradient descent happens in a tiny subspace." arXiv preprint arXiv:1812.04754 (2018).

[7] Larsen et al. "How many degrees of freedom do we need to train deep networks: a loss landscape perspective." ICLR 2022.

Reviewer comment (Theory):
Is it feasible to find an approach for approximating $\Lambda_{w*}$ that eliminates the reliance on the inverse of $\eta$? Is delaying the timing of reducing the learning rate the sole method when numerically evaluating your bounds? I believe there is some relationship between the generalization (or the desired convergence) and the speed of learning rate decay. What is your perspective on these concerns?

Author response:
**Response.** As mentioned in the previous response, in our opinion, we think that generalization bounds that reverse-depend on the learning rate are more consistent with empirical observations in practice. Additionally, it's worth noting that the learning rate decay strategy does not affect the numerical evaluation of the bounds in Section 4 since they use the final learning rate.

While we agree with the reviewer that the speed of learning rate decay can impact the optimization and generalization of SGD, our terminal-state bounds might be insufficient to provide meaningful insights into this matter at this point. However, it is essential to highlight that our trajectory-based bounds, capable of tracking the training dynamics, hold potential for investigating the effects of the learning rate. Despite not explicitly incorporating the learning rate, studying how learning rate decay affects the structure of gradient noise might yield valuable insights.

Reviewer comment (Evaluation):
When numerically evaluating generalization errors in experiments, I imagine it involves evaluating the difference in accuracy between the test data and the training data. In such cases, even if the training/test accuracy is low, it's possible for the apparent generalization error to be small. Therefore, if possible, reporting the predictive accuracy after training completion would enhance the reliability of the experimental results (I refer to Appendices H and I in [2] for example).

Author response:
**Response.** Thank you for this suggestion and thanks for bringing these experiments to our attention, we will try to include them in the next revision.

Reviewer comment (Writing):
p.4, the paragraph of "Validation of SDE," the 2nd paragraph: an oder 1 strong... --> an order 1 strong...?

Author response:
**Response.** Thanks for your careful reading, we have fixed this.

Reviewer comment (Presentation):
The overall organization can be done better. For instance, there are many different information-theoretic bounds throughout the papers, each separately discussed below the corresponding theorems. It would be very helpful for first-time readers if the paper had a summarizing section that gathers and summarizes all the results and compares them (e.g. when is this more useful/tighter)

Author response:
**Response.** Thank you for this suggestion. We have included Table 1 in the Appendix of the revision, which summarizes all the obtained bounds in this paper. Additionally, we have provided important remarks about these bounds in the table.

Reviewer comment (Novelty):
Overall, although the results are better than previous bounds (e.g., Neu et al. (2021)), to me, it's a bit unclear exactly what are the technical novelties in allowing for better results. The discussions provided after each theorem helped me understand the context of how to interpret the new result, but I'm a bit confused about whether the results themselves are new or are just improvements of known bounds. Even in the contributions, although the paper provides much explanation on what the bounds mean and what they are saying, the novelty of the results is somewhat unclear (As most of the implications are already somewhat known in deep learning theory literature).

Author response:
**Response.** We would like to illustrate the reason that our bounds are better than Neu et al. (2021). Specifically, they invoke an auxiliary weight process, we let $\mathcal{A}\_{AWP}$ to denote the auxiliary noisy weight process in Neu et al. (2021). In our paper, we use SDE in Eq.(5), denoted as $\mathcal{A}\_{SDE}$. Let $\mathcal{A}\_{SGD}$ denote the original algorithm of SGD, recall $\mathcal{E}$ is the generalization error, then Neu et al. gives the bound as
$$
\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{AWP})-\mathcal{E}(\mathcal{A}\_{AWP})\leq \text{GenBound}(\mathcal{A}\_{AWP})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{AWP})|.
$$

Our paper gives the bound
$$
\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{SDE})-\mathcal{E}(\mathcal{A}\_{SDE})\leq \text{GenBound}(\mathcal{A}\_{SDE})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|.
$$

Empirical evidence, as shown in Figure 1 in our paper and in previous works such as [2,3], indicates that  $|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|$ is nearly negligible, particularly when compared to the scale of the bound in our result. On the other hand, it remains uncertain that whether the residual term $|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{AWP})|$ for an artificial $\mathcal{A}\_{AWP}$ in Neu et al. (2021)  is small. This is a crucial factor contributing to the improvement in our bound over theirs. To be fair,  $\mathcal{A}\_{AWP}$ can be optimized over the Gaussian variance, but this might be challenging beyond the isotropic covariance case studied in Wang and Mao (2022).

In terms of technical novelty, while the developments in $\text{GenBound}(\mathcal{A}\_{SDE})$ and $\text{GenBound}(\mathcal{A}\_{AWP})$ are closely related, both rooted in [4], previous works, including Neu et al. (2021), analyze both data-independent and state-independent isotropic noise covariance. In contrast, our SDE uses a data-dependent and state-dependent anisotropic noise covariance, rendering some crucial steps not directly applicable. Due to the different noise structure, the step of selecting a "good" prior distribution differs significantly. Additionally, we apply a variant of the log-sum inequality to demonstrate that the population gradient noise covariance is superior to the isotropic covariance, a method not invoked before. Furthermore, prior information-theoretic bounds have not considered the posterior distribution as a mixture of stationary Gaussians, as presented in our Section 4.

[2] Wu et al. "On the noisy gradient descent that generalizes as sgd." ICML 2020.

[3] Li et al. "On the validity of modeling sgd with stochastic differential equations (sdes)." NeurIPS 2021.

[4] Pensia et al. "Generalization error bounds for noisy, iterative algorithms." ISIT 2018.

Reviewer comment (Novelty):
Overall, am I correct in saying that the paper proposes analyses that provide a somewhat unifying information-theoretic perspective on the folklore (some of which have been studied extensively) results in deep learning theory?

Author response:
**Response.** While our findings do align or link with certain prior results in the broader landscape of deep learning theory, it's essential to emphasize that our paper introduces novel aspects. For instance, our Theorem 3.2, demonstrating that the alignment between population gradient noise covariance and batch gradient noise covariance can serve as an indicator for generalization, has not been explored before.

Reviewer comment (Experiments):
Can the authors provide the plot of the proposed generalization bounds by somehow estimating the mutual information?

Author response:
**Response.** This is a good suggestion to verify the numerical tightness between the bounds and $I(W;S)$. However, it may not be feasible to estimate $I(W;S)$ in the most relevant cases, given the challenges associated with their high dimensionality and the continuous nature of these random variables. It's worth noting that even with advanced mutual information estimators like MINE, uncertainties about the approximation error could render the comparison potentially meaningless.

Reviewer comment (Experiments):
Can this be extended to other optimizers such as momentum, adam, adagrad...etc?

Author response:
**Response.** It is possible to extend to the momentum case but may require some necessary simplifications as used in [4] and Neu et al. (2021). Extending to Adam would likely require a substantially different SDE approximation, which might not be straightforward. Fortunately, the SDE of Adam has been recently explored in [5], providing a valuable starting point. We have acknowledged these potential extensions in the future work section of our revised paper.

[5] Malladi et al. "On the SDEs and scaling rules for adaptive gradient algorithms." NeurIPS 2022.

Reviewer comment (Theory):
Any connection to the information bottleneck theory of deep learning?

Author response:
**Response.** We believe that the notion of establishing such a connection is likely on the minds of most researchers studying information-theoretic bounds, and it still remains a challenge. Bridging this gap could involve decomposing the original $I(W;Z)$ in information-theoretic bounds into $I(W;X)$ and $I(W;Y|X)$ using the chain rule. While IB theory proposes that increasing $I(W;X)$ is beneficial for reducing training error and controlling $I(W;Y|X)$ is crucial to prevent overfitting, the precise connection with the SGD or SDE process remains uncertain to us.

Reviewer comment (Experiments):
Can the authors do something similar with noisy SGD [1]?

Author response:
**Response.** Thank you for bringing this paper into our attention. Although label noise gradient descent is introduced in Section 2.2, their primary focus is on investigating the continuous model, specifically the label noise stochastic gradient flow. Given this, the analytic methods used in [6] might be more suitable. Additionally, our initial exploration of the discrete model presented in their Section 2.2 has proven challenging to interpret, yielding no useful insights thus far. We will continue to explore this aspect.

[6] Mou et al. "Generalization bounds of sgld for non-convex learning: Two theoretical viewpoints." COLT 2018.

Reviewer comment (Theory):
Corollary C.1 states that the generalization is controlled by the distance from initialization. Does this mean that lazy training is when generalization is the best? Or is it that the lazy training phase is not compatible with the required assumptions?

Author response:
**Response.** Under the precondition that overparameterized models can attain zero training loss, we think our results do indicate that lazy training can lead to good generalization.

Reviewer comment (Writing):
Motivation of this work is not clear. It seems the main motivation is the connections between dynamics of SGD and its associated discrete SDE, but the only evidence in the paper is a plot in the introduction.

Author response:
**Response.** We would like to clarify the motivation behind our work. Analyzing SGD directly through information-theoretic bounds is often intractable due to its discrete nature (i.e. leading to a degenerate distribution in KL divergence), as also mentioned in your comments. However, the SDE approximation, where the updating process is smoothed by Gaussian noise, is amenable to analysis through information-theoretic bounds. Specifically, let $\mathcal{A}\_{SGD}$ denote the original algorithm of SGD, and let $\mathcal{A}\_{SDE}$ be the SDE in Eq.(5),  recall $\mathcal{E}$ is the generalization error, then $\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{SDE})-\mathcal{E}(\mathcal{A}\_{SDE})\leq \text{GenBound}(\mathcal{A}\_{SDE})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|$.

The key observation that $|\mathcal{E}(\mathcal{A}{SGD})-\mathcal{E}(\mathcal{A}{SDE})|$ is nearly negligible, as demonstrated in Figure 1, encourages our focus on analyzing the generalization bound of SDE. We also note that similar empirical evidence has been observed in some previous studies, such as [1,2].

To provide further clarity on our motivation, we have added an additional section in Appendix B.1.

[1] Wu et al. "On the noisy gradient descent that generalizes as sgd." ICML 2020.

[2] Li et al. "On the validity of modeling sgd with stochastic differential equations (sdes)." NeurIPS 2021.

Reviewer comment (Novelty):
In the related work section, the authors miss many prior works and the discussion appears incomplete.

Author response:
**Response.** We have augmented the section on related works in our paper and expanded on the validation of SDE in Appendix B.2. We note that previous theoretical results aim to illustrate that SDE can weakly approximate SGD in the distribution sense, rather than precisely matching any single sample path of SGD. In addition, although we have made an effort to include all related works, there is a chance that we may have overlooked some important prior works. We genuinely welcome the reviewer to identify any such works, and we would be more than willing to incorporate them into our paper.

Reviewer comment (Writing):
The validation of SDE section in the paper includes many technical terms without exactly defining them.

Author response:
**Response.** We have augmented the section on related works in our paper and expanded on the validation of SDE in Appendix B.2. We note that previous theoretical results aim to illustrate that SDE can weakly approximate SGD in the distribution sense, rather than precisely matching any single sample path of SGD. In addition, although we have made an effort to include all related works, there is a chance that we may have overlooked some important prior works. We genuinely welcome the reviewer to identify any such works, and we would be more than willing to incorporate them into our paper.

Reviewer comment (Theory):
Regarding Section 3 of the paper, many parts are unclear. In particular, it is assumed that the transition kernel of SGD is Gaussian around the minima, but it is unclear why this is a valid assumption, since the transition kernel of SGD is just a bunch of delta measures based on the different realizations of the mini-batches.

Author response:
**Response.** We think the reviewer might be referring to Section 4 rather than Section 3, as the Gaussian distribution assumption around the local minima is used in Section 4 (please correct us if we have misunderstood anything). We have provided this background information in Appendix B.3. In the continuous case, this is rooted in a classical result of the Ornstein–Uhlenbeck process, which is a Gaussian process. The stationary solution to its corresponding Fokker–Planck equation yields a Gaussian density. In the discrete case, in the long-time limit, $W_T$ represents the weighted sum of Gaussian noises. It's important to note that in both cases, we require the quadratic approximation of the loss (i.e., Eq. (7)).

Reviewer comment (Novelty):
In general, the motivation of the paper is rather weak.

Author response:
**Response.** We would like to clarify the motivation behind our work. Analyzing SGD directly through information-theoretic bounds is often intractable due to its discrete nature (i.e. leading to a degenerate distribution in KL divergence), as also mentioned in your comments. However, the SDE approximation, where the updating process is smoothed by Gaussian noise, is amenable to analysis through information-theoretic bounds. Specifically, let $\mathcal{A}\_{SGD}$ denote the original algorithm of SGD, and let $\mathcal{A}\_{SDE}$ be the SDE in Eq.(5),  recall $\mathcal{E}$ is the generalization error, then $\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{SDE})-\mathcal{E}(\mathcal{A}\_{SDE})\leq \text{GenBound}(\mathcal{A}\_{SDE})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|$.

The key observation that $|\mathcal{E}(\mathcal{A}{SGD})-\mathcal{E}(\mathcal{A}{SDE})|$ is nearly negligible, as demonstrated in Figure 1, encourages our focus on analyzing the generalization bound of SDE. We also note that similar empirical evidence has been observed in some previous studies, such as [1,2].

To provide further clarity on our motivation, we have added an additional section in Appendix B.1.

[1] Wu et al. "On the noisy gradient descent that generalizes as sgd." ICML 2020.

[2] Li et al. "On the validity of modeling sgd with stochastic differential equations (sdes)." NeurIPS 2021.

Reviewer comment (Theory):
What does it mean that the distribution around local minima is a Gaussian distribution? The assumptions in Section 3 are not well understood.

Author response:
**Response.** We think the reviewer might be referring to Section 4 rather than Section 3, as the Gaussian distribution assumption around the local minima is used in Section 4 (please correct us if we have misunderstood anything). We have provided this background information in Appendix B.3. In the continuous case, this is rooted in a classical result of the Ornstein–Uhlenbeck process, which is a Gaussian process. The stationary solution to its corresponding Fokker–Planck equation yields a Gaussian density. In the discrete case, in the long-time limit, $W_T$ represents the weighted sum of Gaussian noises. It's important to note that in both cases, we require the quadratic approximation of the loss (i.e., Eq. (7)).

Reviewer comment (Writing):
The statement of the theorems can be improved. For instance, it is difficult to understand the assumptions of Theorem 4.1.

Author response:
**Response.** We have revised the statement to make it be more reader-friendly. Please do let us know if any further improvements are needed or if there are specific aspects that remain challenging to follow.

Reviewer comment (Novelty):
An important related work is Wang, Bohan, et al. "Optimizing information-theoretical generalization bound via anisotropic noise of SGLD." Advances in Neural Information Processing Systems 34 (2021): 26080-26090. In this paper, the authors also consider the problem of obtaining generalization bounds for SGD plus anisotropic noise. What are the technical differences between the results in this paper and Wang et al.?

Author response:
**Response.** We have included a discussion on the comparison with Wang et al. (2021) in Appendix D.7 (refer to the discussions after Lemma D.4). This work is indeed closely related, and it was an oversight on our part to merely mention it without elaborating on the differences. 

To clarify, their paper studies the algorithm of SGD with anisotropic noise, while our SDE analysis focuses on GD with anisotropic noise. This means that the discrete gradient noise arising from mini-batch sampling still exists in their analyzed algorithm, whereas the gradient noise is fully modeled as Gaussian in our Section 3. Moreover, Wang et al. (2021) uses matrix analysis tools to optimize the prior distribution. A significant distinction lies in their optimization analysis, which relies on the assumption that the trace of gradient noise covariance remains unchanged during training (see **Constriant 1** in their paper). Additionally, their final optimal posterior covariance is derived based on the assumption that the posterior distribution of $W$ is invariant to the data index, see Assumption 1 in their paper (to be fair, we also uses this assumption to give an additional result in Appendix, i.e. Theorem G.1). In contrast, our Section 3 avoids making these assumptions and demonstrates the superiority of population gradient noise covariance (GNC) in Lemma 3.2, by invoking a variant of the log-sum inequality. In summary, our proof is simpler and more straightforward, while Wang et al. (2021) makes a stronger claim about the optimality of population GNC based on their additional assumptions.

Reviewer comment (Novelty):
Some inappropriate comparisons and missing discussions: The bounds in this paper are established for noisy GD, but comparisons to [1] and [2], which are results on SGD analyzed through an auxiliary sequence, are made. It is strange to say a bound on a different object is tighter. More important are the missing discussions of other data dependent bounds established for schemes very similar to the studied noisy GD. For example, the very paper this work builds on has a result on SGLD where the batch noise covariance appears [3, theorem 3.1], comparisons with such results would be more appropriate. Especially since theorem 3.1 in this work is a straightforward application of the method in [3].

Author response:
**Response.** The comparison with [1,2] is based on the observation that the SDE approximation exhibits nearly the same generalization error as SGD. Specifically, [1,2] invoke an auxiliary weight process, denoted as $\mathcal{A}\_{AWP}$, to obtain the generalization bound. In our paper, we use the SDE in Eq.(5), denoted as $\mathcal{A}\_{SDE}$. Let $\mathcal{A}\_{SGD}$ be the original algorithm of SGD. Recall that $\mathcal{E}$ is the generalization error, [1,2] provide the bound as follows:
$$
\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{AWP})-\mathcal{E}(\mathcal{A}\_{AWP})\leq \text{GenBound}(\mathcal{A}\_{AWP})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{AWP})|.
$$

Our paper presents the bound as:
$$
\mathcal{E}(\mathcal{A}\_{SGD})=\mathcal{E}(\mathcal{A}\_{SGD})+\mathcal{E}(\mathcal{A}\_{SDE})-\mathcal{E}(\mathcal{A}\_{SDE})\leq \text{GenBound}(\mathcal{A}\_{SDE})+|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|.
$$

In the experiments of this paper, given that $|\mathcal{E}(\mathcal{A}\_{SGD})-\mathcal{E}(\mathcal{A}\_{SDE})|$ is nearly negligible (particularly when compared with the scale of the GenBound) as shown in Figure 1 (and also in [4,5]), we thus directly compare $\text{GenBound}(\mathcal{A}\_{SDE})$ with the bound in [1,2].

We have included these discussions in Appendix B.1 in the revision.

[4] Wu et al. "On the noisy gradient descent that generalizes as sgd." ICML 2020.

[5] Li et al. "On the validity of modeling sgd with stochastic differential equations (sdes)." NeurIPS 2021.


>- More important are the missing discussions of other data dependent bounds ....

**Response.** Note that our bounds in Thm 3.1-3.2 rely on Lemma 2.1, which, as clarified in [3], uses a distribution-dependent prior rather than a data-dependent one. Consequently, our developments align more closely with [6,1] rather than [3], and we do compare with [6,1] in Corollary 3.1 for the case of GLD.

The previous SGLD bounds in [3] and several other prior works apply specifically to noisy SGD with a diagonal Gaussian covariance matrix. In this scenario, the noises are both data-independent and state-independent (i.e. independent of the current weights), so they can not directly be used to compare with our bounds. The appearance of "population gradient noise covariance" in [3] is a result of their construction of gradient incoherence (refer to Lemma D.3 in our revision). In summary, it arises in their bound due to mean differences between the Gaussian posterior and the data-dependent Gaussian prior, while in our bound, it emerges from both mean and covariance differences.

In response to your suggestions, we also made an effort to incorporate the analysis from [3] into our work. Given that the noise term is now data-dependent and state-dependent, certain critical steps in [3] cannot be directly applied. Consequently, we use a variant of the bound from [3], and even in this case, additional assumptions are necessary to obtain the analytic form, as explained in Section D.7 for further details. 

[6] Pensia et al. "Generalization error bounds for noisy, iterative algorithms." ISIT 2018.

Reviewer comment (Theory):
Interpretability of the bounds: The choice of priors would ideally yield meaningful quantities. Here, aside from obtaining a dependence on trajectory wise gradient norms, a result already known for generic noisy iterative schemes, the paper proves Theorem 3.2 where the alignment between the batch estimate covariance and the population gradient covariance appears. It is difficult to understand why this quantity is interesting as it is not well discussed. Moreover, there should be more discussion on the positive definite assumption because, in practice, for neural networks, there are many more weights than data points, thus making \(C_t\) invertible only for large \(n\).

Author response:
**Response.** Consider $n\gg b$. Recall the alignment term $\frac{\Sigma^\mu_tC_t^{-1}}{b}={\Sigma^\mu_t}{\Sigma_t}^{-1}$ where $\Sigma^\mu_t$ is the population gradient noise covariance (GNC) matrix, and $\Sigma_t$ is the GNC matrix for a specific training sample $S$. If these two matrices exhibit perfect alignment with each other *on average*, it implies that SGD (or its SDE approximation) is effectively performed in the true distribution of data, i.e., $w_t=w_{t-1}-\eta\nabla\mathbb{E}\_{Z}{\ell(w,Z)}+\eta\sqrt{\Sigma^\mu_t}N_t$. In other words, the perfect alignment of these two matrices indicates that SGD is insensitive to the randomness of $S$. Recall the key quantity in information-theoretic bounds, $I(W;S)$, which also measures the dependence of $W$ with the randomness of $S$, the term $\Sigma^{\mu}_{t}\Sigma^{-1}_t$ conveys a similar intuition in this context.

Reviewer comment (Theory):
A mix of informal and formal claims: Throughout the paper, even in the proofs, and especially in section 4, the paper mixes informal and formal results making it quite difficult to track what the result is stating. Approximate equalities \(\approx\) and \(=\) are used interchangeably and results that hold in the limit are taken to be valid for a finite T. This is especially felt in theorem 4.1 and lemma 4.1 and its corollaries. The writing and structure of the paper could really be improved.

Author response:
**Response.** Thank you for pointing out this. We have revisited these parts and clarified that the majority of the results presented in Section 4 are asymptotic in nature. Please inform us if you have any further concerns or if additional improvements are needed in the wording.

Reviewer comment (Theory):
Tenuous links with practice and problems with solely analyzing generalization: The paper makes an effort to link its bounds with practice but is often unconvincing. For instance, concepts such as the "edge of stability" are discussed along with experiments simply because a step size and largest singular value appeared together in lemma 4.1 after two obscure assumptions such as commutation of a Hessian and a covariance matrix and an ambiguous \(\approx\). For the trajectory based bounds, it is important to note that the obtained bounds are time dependent, meaning that a single pass over the data can quickly yield a vacuous result so this should be stated well before these bounds are used to try to justify practice. Putting aside the previous points, it is clear from Corollary 4.2 and its interpretation that solely analyzing generalization bounds can quickly lose meaning: yes, an algorithm that stays close to initialization generalizes better but so does an algorithm that simply ignores the training data entirely. Over analyzing generalization bounds without thinking of the utility of the algorithm can lead to strange interpretations.

Author response:
**Response.** Regarding the "edge of stability", here we give an intuitive explanation to the failure mode of $I(W;S)$ in the deterministic case (i.e. $I(W;S)\to\infty$) based on the interplay between the step size and the largest singular value. This is a by-product of our exploration and could be an avenue for further investigation. It's important to note that we do not aim to make overstated claims, and we have rephrased the relevant sections to prevent any potential confusion.

>-  For the trajectory based bounds, it is important to note ... 

**Response.** We have provided more discussions about the time-dependent nature of trajetory-based bounds in the revision (refer to the highlighted text before Section 4). Moreover, we leverage this nature as motivation to introduce terminal-state-based bounds in Section 4. Essentially, if the preference is for time-independent results, we provide asymptotic results that are free from time dependence. In addition, we acknowledge the significance of adopting time-independent trajectory-based bounds, which, similar to Section 4, may require some additional assumptions, for example, smoothness, dissipativity, and regularity of the initial distribution used in a very recent study by [10]. It's important to note that empirically estimating certain parameters in time-independent bounds of [10] poses challenges.

[10] Futami et al. "Time-Independent Information-Theoretic Generalization Bounds for SGLD." NeurIPS 2023.

>- Putting aside the previous points ... .

**Response.**  We agree that discussing generalization is meaningful only under the precondition that the algorithm has the satisfactory performance on training data. Our discussions following Corollary 4.2 have been revised accordingly.

Reviewer comment (Theory):
Could you please formally restate the result showing that the noisy GD scheme you analyze converges to a mixture of Gaussians centered at local minima? It is only briefly discussed with a mention to a result by Mandt et al.

Author response:
**Response.** We have included this background information in Appendix B.3 of the revision. In the continuous case, this is from a classical result of the Ornstein–Uhlenbeck process, which is a Gaussian process. The stationary solution to its corresponding Fokker–Planck equation has a Gaussian density. While in the discrete case, in the long-time limit, $W_T$ represents the weighted sum of Gaussian noises. Note that in both cases, we need the quadratic approximation of loss (i.e. Eq.(7)).

Reviewer comment (Theory):
In lemma 4.1 are you assuming that there is only a single minimum per sample \(S\)?

Author response:
**Response.** We do not assume this. Lemma 4.1 is developed for a given $s$ and a given local minimal $w^*$ but $s$ itself could have more than one $w^*$. That is why we use  the phrase " when $w$ is close to any local minimal $w^*$" in the statement.

Reviewer comment (Theory):
The results Corollary 4.2 and Theorem 4.2 appear to be algorithm independent as long as the algorithm converges to a distribution over local minima. Is this correct?

Author response:
**Response.** Corollary 4.2 and Theorem 4.2 are still algorithm-dependent. Note that that the expectations in these bounds are computed over the randomness of the local minimum $W_S^*$ for the training sample $S$, and the distribution $P_{W_S^*}$ is algorithm-dependent. For instance, it is well-known that numerous local minima exist, but SGD tends to exhibit implicit biases in selecting specific minima,  (e.g., flat minimal/low norm solution/ low rank solution), namely $P_{W_S^*}$ is non-uniform for SGD. As the exploration of the implicit biases of SGD is an ongoing research direction, gaining a better understanding of these biases may aid in more accurately characterizing $P_{W_S^*}$ in the future.