Two Facets of SDE Under an Information-Theoretic Lens: Generalization of SGD via Training Trajectories and via Terminal States

We thank you sincerely for your valuable comments to our paper. Our responses follow.

• The overall organization can be done better. For instance, ....

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.

• Overall, although the results are better than previous bounds (e.g., Neu et al. (2021)), to me, it's a bit unclear ...

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

Our paper gives the bound

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.

• Overall, am I correct in saying that the paper proposes analyses that provide a somewhat unifying ...?

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.

• Can the authors provide the plot of ...?

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.

We thank you sincerely for your comments to our paper. Our responses follow.

• The limitations of the information-theoretic generalization error analysis approach, such as (implicit/explicit) dimension-dependence and time-dependence ...

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.

• 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 ...

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.

• 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?

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.

• 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$ ...?

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.

Thank you for the constructive discussion, corrections, and responses to my concerns. Although the presentation of the terminal state results remains an issue as Reviewer DrRU said, I remain of the opinion that this paper is an important result in the area of IT-based analysis in SGD. Therefore, I will keep my score. I wish the authors the best of luck.

Sincerely,

Reviewer 4Y6d

We thank you sincerely for your constructive comments. Our responses follow.

• Some inappropriate comparisons and missing discussions: ... It is strange to say a bound on a different object is tighter.

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:

Our paper presents the bound as:

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.

• Interpretability of the bounds: The choice of priors ... It is difficult to understand why this quantity is interesting as it is not well discussed.

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.

We really appreciate the valuable discussions and constructive suggestions provided by the reviewer. These inputs have greatly contributed to improving our paper. Here, we response to two points raised by the reviewer.

• ... I understand that those error terms can only be quantified for specific settings, like bounded or lipschitz losses. So why not compare the bounds in that setting before stating one is tighter than the other.

We attempted to make such a comparison in the case where the loss is $\beta$-smooth. While it is easy to obtain an upper bound for the residual term of $\mathcal{A}\_{\rm AWP}$, namely $\beta d\sum_{t=1}^T\sigma_t^2$ when using an isotropic Gaussian with variance $\sigma_t^2$ in the auxiliary weight process (and $d$ can be replaced by the trace of the sum of the noise covariance matrices in the anisotropic case), we still have some challenges in the analysis of the SDE process. This difficulty stems from the need for a similar analysis as the one used for the previous weak approximation results. Notably, the auxiliary weight process involves the same gradient signal as SGD for updating weights, simplifying the analysis of its cumulative noises. We acknowledge that, by the end of the Reviewer-Author discussion period, we may not be able to provide a concrete result. Nevertheless, we sincerely thank you for this suggestion, and we will keep it in mind.

• However, I still believe that the terminal state results are not properly presented and can erroneously lead the reader into believing stronger results are shown.

We understand the reviewer's concern regarding the importance of avoiding any potential misinterpretation or overselling of results. To further avoid any potential misleading, we have made two additional modifications in the revised paper:

1. We have added additional discussions in the contributions part of the introduction (cf. the last red texts on Page 2), we also state here for convience:

Comparing to the first family of bounds (i.e., trajectories-based bounds), the second family of bounds directly bound the generalization error via the terminal state, which avoids summing over training steps; these bounds can be tighter when the steady-state estimates are accurate. On the other hand, not relying on the steady-state estimates and the approximating assumptions they base upon is arguably an advantage of the first family.

2. We have also clarified at the beginning of Section 4 that terminal state results depend on critical assumptions/approximations (including moving the quadratic loss to the first paragraph in Section 4):

To overcome the explicit time-dependence present in the bounds discussed in Section 3, one must introduce additional assumptions, with these assumptions being the inherent cost. For example, an important approximation used in this section is the quadratic approximation of the loss...

We remain committed to refining the presentation of Section 4 and appreciate the opportunity to improve the clarity. Thank you again for your continued guidance.

• 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+ anisotropic noise. What are the technical differences between the results in this paper and Wang et al?

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.

We thank you sincerely for your valuable feedback on our paper. Our responses follow.

• Motivation of this work is not clear for me ....

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.

• In the other related work, the authors miss many prior work and also the discussion in the related work section seems in-complete. Validation of SDE section in the paper also include lots of technical terms without exactly defining them.

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.

• Regarding Section 3 of the paper, many parts are not clear to me. ...
• What does it mean that the distribution around local minima be Gaussian distribution? I do not understand the assumptions in Section 3.

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)).

• The statement of the theorems can be improved. For instance, ...

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.

Dear Reviewers,

While we anticipated more interaction during the Reviewer-Author discussion period, we understand that the workload might be potentially heavy for the reviewers this year, and some may need to focus on their own rebuttals. Given the approaching end of the discussion, and the possibility of not being able to update the paper again, we provide a brief overview of two minor modifications in the latest revised version (in addition to the previous changes), resulting from further discussions with Reviewer E5SG:

1. We have incorporated additional discussions into the contributions parts of the introduction (refer to the last red-texted portion on Page 2). This addition aims to highlight the strengths and weaknesses of trajectories-based bounds and terminal-state-based bounds.

2. We have clarified at the beginning of Section 4 that terminal state results depend on some important assumptions/approximations, serving as a trade-off to eliminate time-dependence.

Best,