Leveraging Flatness to Improve Information-Theoretic Generalization Bounds for SGD
Authors: Ze Peng
Jian Zhang
Yisen Wang
Lei Qi
Yinghuan Shi
Yang Gao
Venue: ICLR 2025 Poster (2025)
Keywords: information theory; implicit bias; deep learning theory; convex optimization; learning theory
Primary area: learning theory
TL;DR: We make information-theoretic bounds better leverage the flatness bias of SGD. We also use our technique developed in this process to address a known limitation of information-theoretic bounds on convex-Lipschitz-Bounded (CLB) problems.

Abstract: Information-theoretic (IT) generalization bounds have been used to study the generalization of learning algorithms. These bounds are intrinsically data- and algorithm-dependent so that one can exploit the properties of data and algorithm to derive tighter bounds. However, we observe that although the flatness bias is crucial for SGD’s generalization, these bounds fail to capture the improved generalization under better flatness and are also numerically loose. This is caused by the inadequate leverage of SGD's flatness bias in existing IT bounds. This paper derives a more flatness-leveraging IT bound for the flatness-favoring SGD. The bound indicates the learned models generalize better if the large-variance directions of the final weight covariance have small local curvatures in the loss landscape. Experiments on deep neural networks show our bound not only correctly reflects the better generalization when flatness is improved, but is also numerically much tighter. This is achieved by a flexible technique called "omniscient trajectory". When applied to Gradient Descent’s minimax excess risk on convex-Lipschitz-Bounded problems, it improves representative IT bounds’ $\Omega(1)$ rates to $O(1/\sqrt{n})$. It also implies a by-pass of memorization-generalization trade-offs. Codes are available at [https://github.com/peng-ze/omniscient-bounds](https://github.com/peng-ze/omniscient-bounds).

Published as a conference paper at ICLR 2025
LEVERAGING FLATNESS TO IMPROVE INFORMATION- THEORETIC GENERALIZATION BOUNDS FOR SGD
Ze Peng', Jian Zhang', Yisen Wang", Lei Qi Yinghuan Shit", Yang Gao
State Key Laboratory for Novel Software Technology. Nanjing University
State Key Lab of General Artificial Intelligence,
School of Intelligence Science and Technology, Peking University
School of Computer Science and Engineering, Southeast University
(pengze, zhangji i an736 0773669688 smaiil nju.. edu en, yisen. weag..awg@@ku.ee.. edu. cn, gilei@seu. edu..edu.e.o en, {syh, en edu.en
ABSTRACT
Information-theoreti (IT) generalization bounds have been used to study the generalization of learning algorithms. These bounds are intrinsically data- and algorithn-dependen so that one can exploit the properties of data and algorithm to derive tighter bounds. However, we observe that although the flatness bias is crucial for SGD's generalization, these bounds fail to capture the improved gen- eralization under better flatness and are also numerically loose. This is caused by the inadequate leverage of SGD's flatness bias in existing IT bounds. This paper derives a more flatness-leveraging IT bound for the flatness-favoring SGD. The bound indicates the learned models generalize better if the large- variiner direc- tions of the final weight covariance have small local curvatures in the loss land- scape. Experiments on deep neural networks show bound not only correctly reflects the better generalization when flatness is improved, but is also numeri- cally much tighter. This achieved by flexible technique called "omniscient trajectory".. When applied to Gradient Descent's 's minimax excess risk on convex- Sunschitz-Boundee problems, it improves representative IT bounds' 22(1) rates 0(1/v/n It also implies by-pass of eeemorizationngeeeraaiaaa trade-offs.
INTRODUCTION
wver-parameterizeed deep models trained by Stochastic Gradient Descent (SGD) are observed to generalize well, contradicting classic statistical learning theories that over-parameterization Ieads overfitting (Vapnik et al., 1998; Bartlett & Mendelson, 2002; Zhang, 2002). The drawback of these theories is that they are too general and unable to leverage the specific properties of learning algorithms and data (Nagarajan & Kolter, 2019). Therefore, modern learning theories have turned to data- and algorithm-dependent bounds that leverage the properties of data and popular algorithms (e.g., the limited hypothesis subset reached by SGD, various norms matrix weights, oowrrankness and sparsity of parameters or hidden representations, etc.) to derive tighter bounds specific to them (Brutzkus et al., 2018; Allen-Zhu et al., 2019; Arora et al., 2019; Cao & Gu, 2019; Neyshabur al., 2019; Pesme et al., 2021; Muthukumar & Sulam, 2023; Alquier, 2023).
Recently, generalization bounds have been developed using information thhooretic measures (Hell- ström et al., 2024), because these measures are defined with the data distribution and the conditional distributions of the output hypothesis given the training data and are naturally data- and algorithm- dependent. Representative examples include the PAC-Bayesian bounds (McAllester, 1999) and the bounds using mutual information (MI) between the training data and the output of the algorithm (Russo & Zou, 2020; Xu & Raginsky, 2017). Thanks to the dependence, PAC- BAC-Bayesian approaches have led to the first non-vacuous numerical generalization bound for deep networks (Dziugaite & Roy, 2017), later tightened (Pérez-Ortiz al., 2021) and scaled up (Zhou al., 2019; Lotfi al., 2022: 2024). MI bounds, the focus of this paper, can be seen as PAC-Bayesian bounds with optimal priors (Alquier, 2023). Tighter variants of the MI bound have been developed, e,g., its chaining
"Codes available at https: / thub. com ppegm ze/ omn sci ant-bo

---

Published as a conference paper at ICLR 2025
Givertization creloue
Eetirrtte proceitation
b-00006
b-2.00
E-0.05
Io-0.02
Boand
(Cross Ensnupyy
13x18
(Cross Ennmpp))
12x00
Badd
Band
Baskikk
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure 1: Wang & Mao (2022)'s bound for ResNet- 18 on CIFAR-10 under varied flatness.
variants (Asadi et al., 2018), its individual-saman sample variants (Bu et al., 2020), the conditional MI framework (CMI) (Steinke & Zakynthinou, 2020), the evaluated variants (Harutyumyan et al., 2021; Hellström & Durisi, 2022; Wang & Mao, 2023), and bounds using other measures like rate-distortion (Sefidgaran et al., 2022) and Wasserstein distance (Wang et al., 2019). By bounding the measures for specific algorithms, they have been applied to SGLD (Pensia et al., 2018; Negrea et al., 2019; Wang et al., 2021b; Futami & Fujisawa, 2023), discretized SDE (Wang & Mao, 2024b), and SGD (Neu et al., 2021; Wang & Mao, 2022).
Another line of research on algorithmic properties finds SGD favors flat minima (Hochreiter & Schmidhuber, 1997; Keskar et al., 2017). Flat minima are minima at wide and flat basins of the loss landscape, and they are robust to loss landscape changes between the training and testing set. Con- sequently, flatness has been used to understand and improve the generalization of SGD-trained deep models (Achille & Soatto, 2018; Jiang et al., 2019; Foret et al., 2020; Cha et al., 2021; Zhao et al., 2022), where it is formulated with the Hessians ils, H,, of empirical and population losses (Achille & Soatto, 2018; Orvieto et al., 2022). Therefore, as algorithn-dependent bounds, information- theoretic bounds for SGD should leverage the flatness as important algorithmic property-
However, existing information ceemmatioon-hhoreiic bounds for SGD do not adequately leverage the flatness bias. By controlling flatness through varying learning rate and batch size (Jastrzębski et al., 2018; He et al., 2019; Wu et al., 2022), we empirically observe how the generalization under varying flatness is captured by Wang & Mao (2022)'s 's bound (which is tighter than Neu et al. (2021)'s)
2R na² 2
o²
[Generalization Error] S inf a²0
V [gcl
T
ErE Ms H
(E(T)
Flainess Term
Trajectory
with 9. being the update at step t. As shown in Figure la, batch size decreases, the actual general- ization error decreases while Wang & Mao (2022)'s bound increases. That is, their result misaligns with the true generalization under varied flatness. The bound consists of two terms: the trajectory term and the flatness term. Figures 1b and Ic show the flatness term can capture the generalization to some extent while the trajectory term cannot, causing the misalignment. Unlike the flatness term that depends on Hessians, which are the explicit formulation of flatness, the trajectory term depends on gradient variance that is an implicit measure of flatness and requires extra conditions (e.g., near- zero losses) to approximate the Hessian traces (Zhu et al., 2019; Martens, 2020; Feng et al., 2023). Therefore, we postulate it is the trajectory term's implicit dependence on flatness that causes the mis- alignment, and we intend to make the whole bound fully explicitly depend on and leverage flatness. To this end, Proposition of Neu et al. (2021) (see Proposition 2) a good start, which involves finer-graines optimization and depends more algorithmic properties, potentially including flat- ness. It is proved by a technique of auxiliary trajectory that is constructed by randomly perturbing the SGD trajectory with independent Gaussian noises of covariance Σ. The bound optimizes Σ as parameter and takes the form
[Generalization Error] < Ew [fW [W;
L.e., mixture of an optimization over Σ and an expectation over output weight W. By constructing from the Hessians of empirical and population losses, the bound will fully explicitly depend on the flatness. However, find the bound still has two drawbacks: Firstly, it is suboptimal because the optimization is outside the expectation and cannot depend on or leverage the specific properties of W instances. Secondly, accurate estimation is crucial for evaluating, comparing bounds and helping
2

---

Published as a conference paper at ICLR 2025
algorithm design. Yet, the bound is not easy to estimate: It is similar to optimizing the "population risk" of "sample" W over "hypothesis" Σ, which requires sampling some W and minimizing the "empirical risk" (ERM). Since only few W can be sampled due to the computation and data cost of training deep models, the "ERM" is prone to "overfitting". Le., negative bias in estimation.
Both issues can be addressed by moving the optimization inside the expectation (i.e., bounds like Ew [infy f(W;)) so that the optimization depends on W: The inside optimization can lever- age the specific properties (e.g. population Hessian instance) of W, leading to a tighter bound; moreover, the empirical mean & E infz. /(W;;) of the inside (instance-overfittingg opti- mization is an unbiased estimator and overfitiigg is no longer source of negative bias in es- timations. To implement this interchange, we find it is the independence of the Gaussian aux- iliary perturbation that makes E the same when conditioned on any W instance, i.e., D cannot adapt to W instances and unable to leverage their specific properties. Therefore, we extend the auxiliary trajectory technique by building it from auxiliary perturbation that is no longer inde- pendent but at least depends on W. This is equivalent to moving the optimization inside (check inls(w) E [JWW; E(W))] E finfy: f(W; E)). Pushing this insight further, we make the pertur- bation depend on all random variables (e.g., the training data) in the SGD training process (e.g,, to leverage the instance empirical Hessian that depends on training data), leading to the "omniscient trajectory"?, which leads to tighter bounds if well optimized.
The technique yields a tighter generalization error bound for SGD that explicitly and fully depends on the flatness of the SGD output instance. Intuitively, it indicates the algorithm generalizes well when the output weights are flat and the flatness aligns with the covariance of output weights. Here, the covariance of output weights is computed after independently training multiple models, and the alignment means that the variance low along sharp directions and high along flat directions, as illustrated in Figures 2b and 2c. As discussed in Section 3.4, our alignment better leverages the flatness compared to similar notions of alignments (Wang & Mao, 2024b; Wang et al., 202 The bound is evaluated on ResNet-18 18 trained by CIFAR-10. When varying batch size, it aligns well with the actual generalization error, indicating better exploitation of flatness. Moreover, our bound is numerically tight by being only few percentages looser than the truth across hyperparameters.
Our omniscient trajectory technique has a simple nature, namely, interchanging the order of the expectation and optimization. Thanks the simplicity, our technique can be flexibly combined with many existing techniques. Furthermore, recent works (Livni, 2024; Haghifam et al., 2023; Attias al., 2024) have highlighted the limitations of anformation-theoreti bounds for having an ((1) lower-bound for GD or any accurate learners on some convex- Liieec--pppchhz-BBooue Bounded (CLB) prob- lems. It can also be seen as an information-aceute aaccuorcy trade-off highlighting the complex relationship between memorization and generalization. We find our simple and flexible technique yields an o(1//n) minimax rate for GD on CLB problems. Therefore, despite being simple, it provides asymptotic improvements and addresses significant limitation of existing iifrmatiion-hhoreee generalization theory. It also implies a by-pass of the trade-off: although accurate learners them- selves memorize lot, they are quite close to some oracle learners that memorize little.
Our contributions are summarized as follows: 1) We derive an iiformatioo-theoretii generalization bound for SGD that better leverages its flatness bias and is numerically tighter; 2) our bound shows how the direction of flatness affects generalization; 3) we introduce a flexible omniscient trajectory technique that also 4) yields an O(1/vn) iitoomation-theoretfe bound for GD on CLB problems.
2 PRELIMINARY
We first introduce basic notations. To present existing information-theoret bounds for SGD and discuss important insights behind them for our (re)use in Section 2.3, we first introduce the algorithm of interest in Section 2.1 and the formulation and properties of flatness in Section 2.2. For k E N+, let [k]] {1, 2, ,,.,,...Kkk k... For sequence ag, .i,..., let atr == (a:), and let a_; denote the rest after excluding a; For vector x E R and matrix A E kRk let |||||| == vxr xr and, with an abuse,
'The above anisotropic Gaussian noises with covariance K are only an example of auxiliary trajectory and Theorem generalizes the technique to more general auxiliary trajectories. Therefore, note the technique is about neither anisotropic noises (Neu et al., 2021; Wang & Mao, 2024b) nor general noises (Sefidgaran et al., 2022), but the auxiliary trajectory's 's extra dependence all randomness and the instance-level optimization.
3

---

Published as a conference paper at ICLR 2025
||=ll = Ar. Random variables, realizations, and domains are denoted by capital, lowerease, and calligraphic letters, respectively, For example, Z is a random sample, Z is the sample space, and 2 is realization. Let a over Z be the sample distribution. Let X' denote an LLD. copy of random variable X. Let d be the number of model parameters and W C TR be the hypothesis space. If not specified otherwise, assume W TRd. Let S (71,..., 7n,) fe" be the training set of LLD. samples. A stochastic learning algorithm is formulated as conditional distribution Pwis.
Let & : W Z + R be loss function. The ultimate goal of the learning algorithm is to opti- mize the population risk Lu(w; (w, = Ez~μ [f(we, Z)] over w E W. Since p. is not fully accessi- ble, one uses empirical risk minimization (ERM) by sampling a training set S E z" from μ and optimizing the empirical risk defined by (, E) = E/ l(w, 21) |s|_ The difference in expec- tation over training sets and algorithmic randomness is the (expected) generalization error (gap) gen(u", Pwis: E) == E(s.W) (W, Ls (W, e) ,where μ²" Pwis denotes the joint distribution determined by μº and Pwis. We may omit the loss function when it is clear from the context. Let the loss difference under transnational perturbation E IRd and Gaussian perturbation & with covariance Σ E Rdddd be A (w,s) s) = Bg-N(1,E) [L.(w + + E) - (w)) We write o² instead of Σ in the superscript if Σ = I²1, and omit ry or they are zero.
2.1 ITERATIVE STOCHASTIC ALGORITHMS
To facilitate analysis, we assume an abstract form of iterative stochastic algorithms to hide unneces- sary details and improve generality. We assume the algorithm first prepares independent random variable V for internal randomness, then starts from an independent initial hypothesis Wo E W and updates the hypothesis iteratively for T N steps by W, = W gs S. (W.-1. S, V, Wo.-2) E W for E [7], and finally outputs W == Wr. Here, is a deterministic function. The algorithm speci- fies random process Wo.T, referred to as the original trajectory. We may omit yr dependence on (W S. V, Won-2) to save space. SGD with batch size b can be obtained by first generating the indices B1. E ([n]")' for each mini-batch, then saving them in V and finally letting Ir compute the gradients in the mini-bater defined by B,. Lastly, with access to past weights Woa-2 and all randomnesses V. It can recover past gradients and covers momentum or Adam, etc. As a result, our theoretical results can be directly applied to these algorithms.
2.2 FLATNESS
The flatness at E W is formulated by the Hessians ls(w) and Hy(w) of the empirical and population losses, respectively. Empirically, the flatness of the empirical loss is highly anisotropic for deep models: after SGD training, most empirical Hessian eigenvectors have small eigenvalues while only the rest few have large eigenvalues (Sagun et al., 2018; Papyan, 2019). We distinguish these cigenvectors by "flat" versus "sharp" directions (Jastrzębski et al., 2019; Wu et al., 2022). Perturbations in the weight space along the sharp directions cause large loss changes, while those along the flat directions cause only slight loss changes.
2.3 INPORMATION-THOORE
BOUNDS AND APPLICATION ON SGD
A random variable X is R-sub-Gaussian if E [EXXX-IXX]) S eA² R²/2 for any À E R. Let [(A; B) be the mutual information (MI) between a pair of random variables (A, B) (Cover & Thomas, 2006). The MI between the training set and the output weight can bound the generalization error:
Lemma (Xu Raginsky (2017)) Assume f(w,) .) is R-sub-Gaussian on P for any w E W. The generalization error of Pwis on μ is bounded by gen(p", Pwis) S (2RT S)/nn.
MI is hard to compute for SGD and one must bound it to apply Lemma 1. However, directly bound- ing MI SGD is not enough because the MI of SGD can be infinite (Hellström et al., 2024), leading to a vacuous generalization bound. Neu et al. (2021) propose auxiliary trajectories to address this problem. Auxiliary trajectories the MI between the training set and the output weight. If well designed, the larger the perturbation is, the more the MI decreases. It is often easier to bound the MI of the auxiliary trajectory, which finally
perturbed versions of the original trajectory. aiming to decrease
4

---

Published as a conference paper at ICLR 2025
becomes the "trajectory term" in the bound. Given the auxiliary trajectory. to transform the bound of it into bound for the original trajectory, one must pay a penalty term the loss difference between their output weights. The penalty comprises of several instances of A terms, which approxi-
mated by Taylor expansion to (1, 8) TE ( (w) + (7 + E) H.(w) x (7 + f)/2 As result, the penalty is related the flatness: It will be very large if the perturbation (7 + E has large projections onto the top eigenvectors of H,, i.e., the sharp directions, while it is safe to have large projections onto the flat directions. Trading off between the trajectory and penalty terms,
have Insight for designing auxiliary perturbation and trajectories.
Insight The ideal perturbation for the auxiliary trajectory should have large projections onto flat directions reduce MI and the trajectory term while maintaining small projections onto sharp directions to keep the penalty term small.
We now present the existing tories." Neu et al. (2021) propose isotropic Gaussian as the perturbation for the auxiliary trajectory:
theoretic bounds for SGD with the help of auxiliary trajec-
Definition 1 (SGLD-Like Trajectory) The SGLD-like trajectory Worr of any trajectory Worr is
Wo Wo. W, = W + (W; - Wi-1) +No.
where No- ((0,of1) is an independent isotropic Gaussian noise with > 0.
After building the SGLD-like auxiliary trajectory of the original trajectory, Neu al. (2021) exploit the properties of Gaussian noises to bound the MI of the SGLD-like output weight. Based on the same SGLD-like trajectory, Wang & Mao (2022) improve the technique for bounding MI, providing numerically tighter result in Proposition 1.
Proposition 1 (Theorem 2 of Wang & Mao (2022)) SGD's generalization error is bounded by
R
gen Pw[s)
E 9 E
E - S')
Penalty for the D-LD- trojectory (Flainers Term)
bound (Trajectory Term)
where the expectations are over the randonmess of training set sampling, initialization, and V.
Since the trajectory term and the flatness term correspond to the MI bound and the penalty of the SGLD-like trajectory, respectively, these two terminologies will be used interchangeably. The SGLD-like trajectory decreases MI at the cost of the penalty term. The flatness is exploited to ensure the penalty is not very large. However, as shown in Figure/ 1, Proposition I's trajectory term does not adequately exploit the flatness. This drawback is partially due to the isotropic Gaussian that adds perturbations the same strength along all directions, violating Insight 1. To exploit the anisotropy of flatness, one can use non-isotropic Gaussian noises in Proposition 2.
Proposition 2 (Proposition 8 of Neu et (2021)) SGD's generalization error is bounded by
Ri2 E
gen(u", (aol"",Wiss) I inf EES,
W E BE
where S4 C paxd is the set of symmetric positive definite matrices.
Here, Σ is the covariance of the Gaussian. By setting, for example, - [1g
the noise has
large variances along the directions with small curvatures, better exploiting the flatness. Neverthe- less, we find Proposition still has room for improvement, as elaborated in Section 3.1.
They are slightly modified to have simpler, similar and more comparable forms to each other. These modifications only make the bounds tighter, and will not cause unfair comparisons.
5

---

Published as a conference paper at ICLR 2025
Los Cantour
Cutpur
Original
We
gen(Wrr Wy gen(Wr) Wy)
Orniscient
We
W
W2
SGLD-like
We
gen(Wr)
(a) The relationships between trajectories.
(b) Aligned.
(c) Not aligned.
Figure 2: The relationship between the original, omniscient, and SGLD-like trajectories, and illus- trative examples of of alignment and misalignment between flatness and output weight covariance. The two trajectories have decoupled roles: the omniscient trajectory optimizes the whole bound while the SGLD-like trajectory bounds the MI for the omniscient trajectory.
3 PROPOSED OMNISCIENT TRAJECTORY
3.1 ANALYSIS ON EXISTING RESULTS' DRAWBACKS
To see the inefficiencies of Proposition 2, we recall a generic principle:
Principle Being specific enhances optimization: min Ex [f(X; 2 Ex [mina f(X; A)]
Through the lens of this principle, Proposition 2 is closer to LHS of Principle since E to be opti- mized is shared among all instances in the expectation, leading to trade-offs between these instances. As result, it is suboptimal compared to a fully specific and dependent one.
Moreover, Proposition 2 is prone to overfitting when estimated in numerical studies. Since the dis- tribution of output weights is not fully accessible, one must sample some of them before optimizing over Σ. This procedure identical to ERM, leading to potential overfitting to the sampled weights. Sampling output weights involves training multiple models on multiple training sets. The computa- tional costs and data requirements of training deep models limit the number of samples. In contrast, S+ is a ((d")-dimension)) manifold. Therefore, the overfitting is severe at least in the classic view and the estimated bound has a severe negative bias. Although the overfitting may be benign as it is in deep learning, a geteraaenaiiinnnnnnnn theory is needed to ensure it, which is too complex and may even induce a meta-metta-vvrrftiing problem with more parameters. To our knowledge, Proposition 2 lacks numerical results, possibly due to the large parameter count and the overfitting.
Some examples of the two drawbacks can be found in Appendix B.1. As indicated by Principle 1, both issues be addressed by moving the optimization inside the expectation. Regarding the suboptimality, Principle 1 explicitly shows better tightness. Regarding overfitting, the RHS of Prin- ciple already the result of overfitting to instance trajectories, eliminating overfitting as bias source. Furthermore, depending on more random variables, such as the training set, makes it easier access and exploit the empirical flatness. Motivated by these benefits, we propose fully dependent auxiliary trajectories in the form of functions of all random variables in the training process. With knowledge of all random variables, the trajectories are termed "omniscient (auxiliary) trajectories".
3.2 OMNISCIENT TRAJECTORY
The omniscient trajectory that depends on all random variables in the training is defined as follows:
Definition (Omniscient Trajectory) The omniscient trajectory Worr of (8, VI.T, Worr) spec- ified by omniscient perturbation Agı- is given by
Wo Wo. = W (W, Wo. - (n VI.T,, Wo.r)),
where Agt is deterministic function. Let r, == Σt= Ag. V. Vir, Wo.),, then W. = W.+L,
To bound the generalization error of the original SGD trajectory, transform the generalization error of the original output weight into that of the omniscient output weight at the cost of a "penalty term" for the translational perturbation Ty. To bound generalization error and MI of the omniscient

---

Published as a conference paper at ICLR 2025
output weight, construct an SGLD-like auxiliary trajectory Wo:7 based on the omniscient trajec- tory and transform the generalization of the omniscient trajectory into that of the SGLD-like trajec- tory, at the cost of another penalty, i.e., the flatness term. We then bound the generalization error of the SGLD-like output weight using the MI I (WW5)) The relationships between the trajectories and their generalization errors are summarized in Figure 2a. To bound the MI, we extend the key Lemma 4 of Wang & Mao (2022) that bounds the mutual information using variances in Lemma B.1 so that it can be used for the double-layered auxiliary trajectories despite its heavy dependence to S. Putting these problem transformations and bounds together, we obtain a generalization bound Theorem using trajectory statistics similar to Proposition but modified by the omniscient perturbation.
Theorem
Assume R-sub-Gaassien Gaubsiaatsittttt for For any Agur and any o1r E (R²")", we have
R2
T
gen(k", (en(""PW[[s) S
11g. - E [ Agh
IP]
+ ( (W,,)) - Ar..(W. S S')
n
Penalty for the omni. trajectory (Penalty Term)
MI bound (Irajectory term)
A of
(Wr - S) - of A + +1r
E
Penalty for the SGLD-like trajectory (Flatness Term)
Setting Age- = 0 recovers Proposition 1. We further optimize it to tighten the bound.
3.3 OPTIMIZING THE OMNISCIENT TRAJECTORY AND EXPLOITING FLATNESS
To avoid optimizing over 7 elements Aga:, we first simplify Theorem 1 with Age = (g: E [gee) 1 - (gr- - E [9-1)) + AAG. As result, for any deterministic function AG of (S, V, g1.T, Wor)
T
and > O, we have Corollary B.1 that can be informally stated as (with AW, = W. Wo)
R² nT
|IAW, - - [AWr] + AGI*] + +E A (WT, S) - AX (Wr, s*)
gen
Appendix B.4 verifies such Ag. is optimal under the constraint IT AG, and the generality is not harmed by this simplification. Setting AG = O recovers the isotropic version of Proposition 24.
We then optimize over AG. It can decrease the trajectory term by canceling AWr E (AWr)- However, large AG would increase the penalty term, leading to a trade-off. Fortunately, near flat minima, most directions are flat, and the penalty term is insensitive to perturbations along them. Following Insight I, we confine AG to align with the flat directions. This is done by approximating the penalty terms to the second order, where Hessians emerge, and solving an optimization problem formed by the output weights and the Hessians (see Appendix B.6). Theorem 2 presents the result, where the three expectations correspond to the penalty, flatness, and trajectory terms, respectively.
Theorem 2 Assume -(w,-) is R-sub-Gaussian on P for any w W, ((.,2) and (.(()) are thirdly continuously differentiable w.r.t w for any 2 E Z, and there is a constant D > O such that D [|AW]] bounds the residuals in the third-order Taylor expansions of ((:,2) and L,(-) for any E Z. Then for any >>0 trading- off between the trajectory and penalty terms, we have
Tanlu",Pwiss S E EEAAs - AAT, 3 2
R2 E [tr (AH(WT + AG)]
E [| - - (1+ HunN+HHpe/22)----1 (aoen - -[AW+]) - (2XC1 + Hpen) -1. 27 +r, (1)
12
where AH(w)= == As(w) Hw). Haat is chosen from (H())) and Ass(W)), Apen is cho- sen frofromm (W( and (WT), and J is chosen from ((&(WW)) - (WT) and (s(W)), (AAC(AW--E[AW-l))
and r = Old is the residual in a second-order approximation. See Eq. (B.6) in Appendix B.7 for the form of o..
C == po{( (r(Baa)))" '0, AG = (2XCI+Afpm))
UUsing non-isotropic Gaussian noises for the SGL.D-like trajectory fully recovers Proposition 2.
7

---

Published as a conference paper at ICLR 2025
3.4 DISCUSSION AND COMPARISON
To see the dependence on the flatness more clearly, we approximate the result of Theorem 2. To simplify the approximation, only for this subsection, we assume J == V v (WT), and the model is sufficiently trained so that (Lr(W)) = O. We also assume À is large enough so that 2XC' surpasses H pen'ss 's top singular value approximate T - (T / Be noted that the above assumptions are used only in this subsection, i.e., Theorem 2 and experi- ments in Section 4 do not require them. Results without some of these assumption can be found in Appendix B.8. Approximating the penalty term to the second order leads to the following result:
H /ppe/22C.
gen 2
E
AATT - AE [Wr]l] (ww.(laa((.))--HHM.... + VL (
4x2C2
232
Corresponding Penalty Term
3
2R VV
E l|tr (AH(WT+ + AG)J] E | -
E [AWT]l f
+ IVVC
4x2c2
ex-C2
Corresponding Flatness Term
Corresponding to Trajectory Term
where| |.| replaces the eigenvalues of matrix with their absolute values. The details can be found in Appendix B.8. It can be seen that all terms except the two population gradient norms depend on Hessians. Particularly, the "norms" of deviation AWT - E [WWr] are defined by the Hessians. As a result, the components of the deviation along flat directions contribute little to the bound. The flatter the minima are, the more and flatter the flat directions there are, and more components of the deviations contribute little to the bound. Moreover, the flatness is aligned with the covariance, smaller bound. As a result,
most large deviations can be found near the "flat subspace", leading generalization is better minima are flat and the flatness aligns with the covariance.
This focus alignment is similar to the work by Wang & Mao (2024b) and Wang al. (2021a). These alignments are compared in Appendix B.8.1 in detail. Briefly, the main difference is that our alignment directly relies on flatness. Our technique is also similar to the rate-distortion bound (Sefidgaran et al., 2022) as both involve weight-dependent perturbations. However, our omniscient trajectory depends on more random variables, such as the training data that is erucial for leveraging the empirical Hessian. Besides, the existing chaining (Asadi et al., 2018) and evaluated (Harutyun- yan et al., 2021; Hellström & Durisi, 2022: Wang Mao, 2023) bounds, which can leverage the similarity between adjacent hypotheses, are potentially useful because the flatness of a minimum reflects similarity with neighboring weights. However, chaining requires partitioning the hypoth- esis space, which is difficult for deep neural networks; the evaluated bounds directly rely on model losses, obscuring insights in the language of weights and flatness.
4 EXPERIMENTAL STUDY
In this section, we experimentally show how our bound captures the true generalization error/gap measured by Cross Entropy (0-1 loss does not have Hessians and is motivationally incompatible) compared to the existing bounds. We vary the hyperparameter and train 6 independent ResNet- 18 models on CIFAR-10 at each hyperparameter. To ensure uub-Gaussianity, capped cross-entropy (CE) is used in testing and bound estimation, while vanilla CE is used for training for efficient train- ing. We estimate Theorem 2 with À E {1,10 109}. Hnat = Ap = (Wr) to make estimation easier and vLs (W,,) VL (W,,) for numerical tightnes as detailed in Appendix C.2. Es- timating the bounds requires splitting datasets: The 6 models are trained by 6 random splits of the training set. Terms involving population statistics (e.g.. the population Hessians in the flatness terms and the population gradient in the penalty term) are approximated to the second order and estimated on validation sets. But for existing bounds, we assume C (W) < (5_,((0, [C,, (W + E)] (Wang & Mao, 2022) and avoid the population Hessian. The true generalization error is estimated on separate test set. See Appendix for more details and the results with population Hessians.
We first evaluate whether our bound can better capture the generalization under varying flatness. As shown in Figure 3, the bound correctly captures the tendency of generalization error under varied batch size and learning rate. Specifically, the tendency of the trajectory term w.r.t. the batch size
8

---

Published as a conference paper at ICLR 2025
Im-6001
I--0
b-0.1
I--8
Extimatidl Gen Gip
(Cross Emnoogy)
12.10
(Cros Entrops)
1.2
Gen Gup
L1416
14'
0.4
300
380
306
230
520
0.5
200
20
300
500
530
09
100
200
530
420
502
206
460
4-10"
(a) Bound
(b) Trajectory Term
(c) Penalty Term
(d) Flatness Term
Figure 3: Thm. 2( = 109) on CIFAR-10 and ResNet-18 with varied learning rate and batch size.
Wang Mao (221))
oo (2021)
Then Oe
DeH05)
Generllaatio
Copfbood
(Cross Bintos)
Tra
Lacinde aed
Rlah
Dost: Sae scllnng
Figure 4: Numerical results for ResNet-18 on CIFAR-10 under varied training data usage, label noise level, width, depth, and weight decay. The isotropic version of Proposition is used.
corrected. The penalty term for the omniscient trajectory correlates well with the generalization error w.E.f both learning rate and batch size, thanks to the Hessians in Eq. (2). The flatness term also generally has the correct tendency w.r.f. both hyperparameters. Unfortunately, the tendency of the trajectory term w.r.. learning rate is still incorrect. As shown in Section 3.4, the trajectory term essentially the product between Hessians and the variance of the last-step weight. We conjecture is the increased variance when learning rate increases that overpowers the improved flatness and makes the tendency uncorrected. In contrast, the variance less sensitive to batch size. Detailed discussion can be found in Appendix C.4.1. Nevertheless, after multiplied together, they contribute little to the bound and the uncorrected tendency does not affect the bound very much.
Figure 4 shows how our bound captures the generalization under varied gemeralization-critte hy- perparameters. can be seen that our bound with large (i.e., putting more weights on reducing the trajectory term) can well capture the generalization under these variations, while the some existing bounds fail to capture the improved generalization under stronger weight decay.
Results for MLP on MNIST can be found in Appendix C.4. According Figures and 4 and results in the appendix, our bounds are numerically much tighter. The main improvement comes from the trajectory term, result of both Proposition and the optimized omniscient trajectory leveraging flatness. See Appendix C.4.1 for detailed discussions. In Appendix C.6, we evaluate our bounds under weight scaling, where the generalization is still well captured.
The trajectory term and the flatness almost diminish after multiplied together, and the bound mainly comprises of the penalty term. Theorem 2 and Eq. (2), the term heavily relies on the population gradients. We consider this reliance to be the major limitation of our bound.
EXTENSIONS
Our technique so flexible that it can be combined with existing techniques, such the individual- sample technique and CMI framework. We can also combine it with the data-dependent prior tech- nique of Negrea et al. (2019) for SGLD to yield an SGD bound. The combined results are listed in Appendix D. Notably, the omniscient trajectories in these results depend all new random vari- ables introduced by the variants. As a result, the omniscient trajectory can leverage the features of the variants and optimize more specifically.
Our technique can also address the known limitations of representative existing information- theoretic bounds
Nonvex-Lipschizz-Be Bounded (CLB) problems in stochastic convex optimization (SCO). An SCO problem is a triplet (W, E), where W is convex and f(+,2) 2) is convex given any E Z. The CLB subclass Ct,D of SCO further requires W is closed and bounded with diameter

---

Published as a conference paper at ICLR 2025
D, while 2) is L-Lipschitz given any Z. The goal is to minimize the excess (population) risk C under an unknown sample distribution, which can be bounded by the
Ew [c,) excess optimization error and the generalization error. The minimax excess risk reflects the worst- case generalization: inf Pwise. ()))) Rw- (W") - wnw (u("),, where A is a family of algorithms and M1 is the set of all distributions over Z.
An O(LD/VI) minimax rate has been obtained for Gradient Descent (GD) using uniform stabil- ity. However, Haghifam et al. (2023) have shown there exist CLB instances where PAC-Bayesian bounds, the vanilla (C)MI bounds. their Gaussian pertturbd variants (Propositions and fall in this category if adapted to CLB settings), and representative variants all have an (Q(1) lower-bound, meaning these representative bounds cannot explain the generalization of GD on CLB problems. Attias et al. (2024) have shown any accurate CLB learner with low excess generalization risk must have high CMI, í.e., the CMI-accuracy trade-off. Recently, Wang & Mao (2024a) have addressed the counterexample by Haghifam et al. (2023) and all CLB problems by augmenting CMI bounds with stability. However, a lot of complicated CMIs are involved and due to this complexity it is hard intuitively interpret their relationship with the CMIlaccuracy trade-off. Therefore, we try our technique to see whether it leads to a similar but simpler result. Combining individual-sample bounds with omniscient trajectory, we easily resemble the stability argument:
Theorem (Informal) For CLB problem (W, 2,6) and distribution j over Z, we have
gen S
inf AG
(DDV2T (WT + AG 2 2 + E [[AA(AA AG (Wr, 2.) A A
(Jor stable algovithms)
21
(for GD)
E[ E [Wr Z-I]Il] (with 2 steps size v.)
8T.2
n
This bound recovers that Bassily et al. (2020) up to a constant. Following similar steps of Haghi- fam et al. (2023), one can recover the best-achievable O(LD/ TT) excess risk bound. Theorem 3 and its proof essentially state the original trajectory has a vanishing distance (measured by loss dif- ferences) to the omniscient trajectory of zero individual MI. Therefore, Theorem 3 provides ternative to the trade-off on GD or stable algorithms: although the accurate learners have high CMI, they are quite close auxiliary (oracle) learners with low MI.. Extending this intuition for more general accurate learners in Theorem 4 of Appendix D.5, we explore to replace CMI with a new in- formation measure induced by our technique and break through the aafoomatiionaaccuracy trade-off: this cetrraliaationbbbundie orem requires vanishing excess optimization error and is still partial and preliminary. Lastly, when losses are smooth and convex, we also recover stability-based (1//n) rate for SGD, as detailed in Appendix D.6.
al-
new information measure vanishes together with risks. However, The-
CONCLUSION AND LIMITATION
In this paper, we address the inadequate leverage of flatness in existing information-theores gener- alization bounds for SGD. By continuing the "being more specific" trend in generalization theory, propose the omniscient trajectory that can be optimized depending on all random variables in the training process to better leverage the flatness. Our bound shows that an algorithm generalizes better if its output flatness aligns with its output covariance. When applied to deep neural networks, bound aligns well with the generalization under varied flatness and is numerically tighter. When ap- plied to CLB problems for GD, the technique yields an ((1 minimax rate, addressing a recently highlighted limitation of information-theoree generalization theory.
However, our bound relies on population gradients and Hessians. Although tolerable for theories (e.g., Neu et al. (2021) and Wang & Mao (2022; 2024b) also rely on population statistics), this problem prevents the bound from being a part of self-certified algorithms (Pérez-Ortiz al., 2021). Moreover, the most iments (see Eq. (D.5) and Figures 3c and C.4c). We conjecture these two issues may be addressed bounding the generalization of higher-order statistics.
components vanish in both the CLB bound and the exper-
in future works by More detailed discussions on the limitations can be found in Appendix E.
The auxiliary oracle leamners E. [Wr]] 2_ does obey the MI-accuracy trade-off due its access to μ..
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
All proofs and full versions of informally stated results can be found in the appendices for the corresponding sections. Details of experiments can be found in Appendix C. The codes for the ex- periments can found in the supplementary material or https: /ttl::///gtthaee hgghhtb.thuhgggg com/ www.pengccpppppp///eee/ece...e.e peng-zel ommiscient-bounde
ACKNOWLEDGMENTS
This work was supported by the NSFC Project (62192783, 92370129, 62222604, 62206052, 62376010), China Postdoctoral Science Foundation (2024M750424), Fundamental Research Funds for the Central Universities (020214380120, 0202 283801201,, State Key Laboratory Fund (ZZKT2024A14) 14), the Postdoctoral Fellowship Program of CPSF (GZC20240252), Jiangsu Funding Program for Excellent Postdoctoral Talent (2024ZB242), Jiangsu Science and Technology Major Project (BG2024031), and Beijing Nova Program (20230484344, 20240484642).
REFERENCES
Alessandro Achille and Stefano Soatto. Emergence of invariance and disentanglement in deep rep- resentations. Journal of Machine Learning Research, 19(50):1-34, 2018.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameter- ized neural networks, going beyond two layers. In Advances Neural Information Processing Systems, volume 32, 2019.
Pierre Alquier. friendly introduction to PAC-Bayes bounds, 2023. URL httt::/arxxii org/abs 11216.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In Proceedings of the 36th International Conference on Machine Learning, pp. 477-502. PMLR, 2019.
Amir Asadi, Emmanuel Abbe, and Sergio Verdu. Chaining mutual information and tightening gen- eralization bounds. In Advances in Neural Information Processing Systems, volume 31, 2018.
Idan Attias, Gintare Karolina Dziugaite, Mahdi Haghifam, Roi Livni, and Daniel M. Roy. Informa- tion complexity Stochastic Convex Optimization: Applications to generalization and memo- rization, 2024. URL http: http://aaavv..g.aaaa. www.///gg/bb/eeeeeeee 402. 09327. arXiv:2402.09327 [cs].
Peter I. Bartlett and Shahar Mendelson. Rademacher and Gaussian complexities Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):463-482, 2002.
Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of Stochastic Gradient Descent on nonsmooth convex losses. Advances Neural Information Processing Systems, volume 33, pp. 4381-4391, 2020.
Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev- Shwlar-Shwwrtz. SGD learns over- parameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representutions, 2018.
Yuheng Bu, Shaofeng Zou, and Venugopal V. Veeravalli. Tightening mutual information based bounds on generalization error. IEEE Journal on Selected Areas in Information Theory, ((1): 121-130, 2020.
Yuan Cao and Quanquan Gu. Generalization bounds of Stochastic Gradient Descent for wide and deep neural networks. Advances in Neural Information Processing Systems, volume 32, 2019.
Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park. SWAD: Domain generalization by seeking flat minima. In Advances in Neural Information Processing Systems, 2021.
Thomas M. Cover and Joy A. Thomas. Elements of Information Theory 2nd Edition ( (Wiley Series in Telecommunications and Signal Processing). Wiley-Interscience, 2006. ISBN 0471241954.
11

---

Published as a conference paper at ICLR 2025
Imre Csiszár and János Körmer. Information theory: Coding theorems for discrete memoryless systems. Cambridge University Press, 2011.
Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau.
How to URL
compute hessian-vector products?
In ICLR Blogposts 2024, 2024.
https: /ttprs///ilrrbbre woweelr-bloooooooo www.bblloopssssooooooooooo.sssss #inverse hessein http://sss...itbbb.o//////////oa hww.b///o///..chhhphhhh..
prroducts-itvss--e cetn--optimizttiiii
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the 33rd Annual Conference on Uncertainty Artificial Intelligence (UAI), 2017.
Yu Feng, Wei Zhang, and Yuhai Tu. wetivt duality in ferd-forward neural networks re- veals two co-determinants for generalization. Nature Machine Intelligence, 5(8):908-911) 2023. Publisher: Nature Publishing Group UK London.
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference Learning Represen- tations. ICLR, 2020.
Futoshi Futami and Masahiro Fujisawa. Time- -ime-independent information-theoree generalization bounds for SGLD. In Advances in Neural Information Processing Systems, 2023.
Mahdi Haghifam, Borja Gálvee Ragnar Thobaben, Mikael Skoglund, Daniel M. Roy, and Gintare Karolina Dziugaite. Limitations of information--heoretiee generalization bounds for Gradient Descent methods in Stochastic Convex Optimization. In Proceedings The 34th Inter- national Conference on Algorithmic Learning Theory, pp. 663-706. PMLR, 2023.
Hrayr Harutyunyan, Maxim Raginsky, Greg Ver Steeg, and Aram Galstyan. information-theored generalization bounds for black-box learning algorithms. In Advances Neural Information Processing Systems, 2021.
Fengxiang He, Tongliang Liu, and Dacheng Tao. Control batch size and learning rate to generalize well: Theoretical and empirical evidence. In Advances in Neural Information Processing Systems, volume 32, 2019.
Fredrik Hellström and Giuseppe Durisi. A new family of generalization bounds using samplewise evaluated CMI. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Fredrik Hellström, Giuseppe Durisi, Benjamin Guedj, and Maxim Raginsky. Generalization bounds: Perspectives from information theory and PAC-Bayes, 2024. URL httt:://axiiv.ot
04381.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Stanislaw Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Ben- gio, and Amos Storkey. Three factors influencing minima in SGD, 2018. URL http: //arxiv.org/aattttiee 1111/.11111 04623.
Stanislaw Jastrzębski, Zachary Kenton, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey. On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations. ICLR, 2019.
Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where to find them, 2019. URL 1912. 02178.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On batch-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations. ICLR, 2017.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEFF, (6(11):2278-22) 1998.
12

---

Published as a conference paper at ICLR 2025
Yunwen Lei and Yiming Ying. Fine-grained analysis of stability and generalization for stochastic gradient descent. In International Conference on Machine Learning, pp. 5809-581111 5819. PMLR, 2020.
Roi Livni. Information theoretic lower bounds for information theoretic upper bounds. Advances in Neural Information Processing Systems, 36, 2024.
Sanae Lotfi, Marc Anton Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum, and An- drew Gordon Wilson. PAC-Bayes compression bounds tight that they explain generaliza- tion. In Advances in Neural Information Processing Systems, 2022.
Sanae Lotfi, Mare Finzi, Yilun Kuang, Tim G. J. Rudner, Micah Goldblum, and Andrew Gordon Wilson. Non-vacuous generalization bounds for large language models, 2024. URL https: /ttp:v/..rr//ab///2..1.2220 //aaaxivvrrg///aeeee///// btarg/abs/233/0
James Martens. New insights and perspectives on the Natural Gradient Method. Journal of Machine Learning Research, 21(16)):1-76, 2020.
David A. McAllester. PAC-Bayesian model averaging. In Proceedings of the Twelfth Annual Con- ference on Computational learning theory, pp. 164-170, 1999.
Ramchandran Muthukumar and Jeremias Sulam. Sparsity-aware generalization theory for deep neural networks. In Proceedings of Machine Learning Research, volume 195, pp. 5311-5342. ML Research Press, 2023.
Vaishnavh Nagarajan and J. Zico Kolter. Uniform convergence may be unable explain gener- alization in deep learning. In Advances in Neural Information Processing Systems, volume 32, 2019.
Jeffrey Negrea, Mahdi Haghifam, Gintare Karolina Dziugaite, Ashish Khisti, and Daniel M Roy. Information-theoretie generalization bounds for SGLD via data-dependent estimates. In Advances in Neural Information Processing Systems, volume 32, 2019.
Gergely Neu, Gintare Karolina Dziugaite, Mahdi Haghifam, and Daniel M. Roy. Information- theoretic generalization bounds for Stochastic Gradient Descent. In Proceedings of Thirty Fowrth Conference on Learning Theory, pp. 3526-35244 PMLR, 2021.
Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of over-parameerrizaion in generalization of neural networks. 7th International Conference on Learning Representations, ICLR, 2019.
Orabona.
SGD converges (even in unbounded do- https: tttpss//parameterfreet www.eoocc///eeeee 2tt//2200/////// of -aaee------s-d-ooorrr-e---e converges -eeerrgeeeeee-e---eeeee inniin--uunnn wnwounduouuddeeeeeeee cowpouunddddddmaiis///.oo domains/.
Last
iterate URL
mains), 2020.
Antonio Orvieto, Hans Kersting, Frank Proske, Francis Bach, and Aurelien Lucchi. Anticorrelated noise injection for improved generalization. Proceedings of the 39th International Conference on Machine Learning, pp. 77094-17116. PMLR, 2022.
Vardan Papyan. The full spectrum of deepnet hessians at scale: Dynamics with SGD training and sample size, 2019. URL http:// aaxvvv.arrivvva/a org/ abs 1811.070000
Ankit Pensia, Varun Jog, and Po-Ling Loh. Generalization error bounds for noisy, iterative gorithms. In 2018 IEEE International Symposium Information Theory, pp. 546-550. IEEE, 2018.
Scott Pesme, Loucas Pillaud-Vivien, and Nicolas Flammarion. Implicit bias of SGD for diagonal linear networks: a provable benefit of stochasticity. Advances in Neural Information Processing Systems, volume 35, pp. 2011--92230. 2021.
María Pérez-Ortiz, Omar Rivasplata, John Shawe- Shawe-Taylor, and Csaba Szepesvári. Tighter risk certifi- cates for neural networks. Journal of Machine Learning Research, 22((27):1-40 2021.
13

---

Published as a conference paper at ICLR 2025
Daniel Russo and James Zou. How much does your data exploration overfit? Controlling bias via information usage. IEEE Transactions on Information Theory, 66(1):302 -233222,,, 2020.
Levent Sagun, Utku Evci, V. Ugür Güney, Yann Dauphin, and Léon Bottou. Empirical analysis of the hessian of over-par parametrized neural networks. 6th International Conference on Learning Representations, ICLR 2018
Workshop Track Proceedings. ICLR, 2018
Milad Sefidgaran, Amin Gohari, Gaël Richard, and Umut Simsekli. Rate-distortion theoretic gener- alization bounds for stochastic learning algorithms. In Proceedings of Thirty Fifth Conference on Learning Theory, pp. 4416-4463. PMLR, 2022.
Thomas Steinke and Lydia Zakynthinou. Reasoning about generalization via conditional mutual information, 2020. URL https: peattps://axxe http//aavvorrg//e/e///// wwwpor//////o///eeeeeeee /2011 09122.
Vladimir Naumovich Vapnik, Vlamimir Vapnik, al. Statistical learning theory. Wiley New York, 1998.
Bohan Wang, Huishuai Zhang. Jieyu Zhang. Qi Meng, Wei Chen, and Tie-Yan Liu. Optimizing infommation-theorrtical generalization bound via anisotropic noise of SGLD. In Advances in Neural Information Processing Systems, volume 34, pp. 26080-26090, 2021a.
Hao Wang, Mario Diaz, Jose Candido S. Santos Filho, and Flavio P. Calmon. An information- theoretic view of generalization via Wasserstein distance. In 2019 IEEE International Symposium on Information Theory, pp. 577-581, IEEE, 2019.
Hao Wang, Yizhe Huang, Rui Gao, and Flavio Calmon. Analyzing the generalization capability of SGLD using properties of Gaussian channels. In Advances in Neural Information Processing Systems, volume 34, pp. 24222-24234, 2021b.
Ziqiao Wang and Yongyi Mao. On the generalization of models trained with SGD: Information- theoretic bounds and implications. In International Conference on Learning Representations. ICLR, 2022.
Ziqiao Wang and Yongyi Mao. Tighter information-theorete generalization bounds from super- samples. In Proceedings of the 40th International Conference on Machine Learning, pp. 36111- 36137. PMLR, 2023.
Ziqiao Wang and Yongyi Mao. Sample-conditioned hypothesis stability sharpens information- theoretic generalization bounds. In Advances Neural Information Processing Systems, vol- ume 36, 2024a.
Ziqiao Wang and Yongyi Mao. Two facets of SDE under an information-theoretic lens: Gener- alization of SGD via training trajectories and via terminal states. In The 40th Conference Uncertainty in Artificial Intelligence, 2024b.
Lei Wu, Mingze Wang, and Weijie Su. The alignment property of SGD noise and how it helps select flat minima: A stability analysis. In Advances in Neural Information Processing Systems, volume 35, pp. 4680-4693, 2022.
Aolin Xu and Maxim Raginsky, Information-theorete analysis of generalization capability of learn- ing algorithms. In Advances in Neural Information Processing Systems, volume 30, 2017.
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael Mahoney, PyHessian: Neural Networks Through the Lens of the Hessian. In 2020 IEEE International Conference on Big Data, pp. 581- 590. 2020.
Tong Zhang. Covering number bounds of certain regularized linear function classes. Journal of Machine Learning Research, 2(Mar):527-55) 2002.
Yang Zhao, Hao Zhang, and Xiuyuan Hu. Penalizing gradient norm for efficiently improving gen- eralization deep learning. In Proceedings of the 39th International Conference on Machine Learning, volume 162, pp. 26982-269999 26992. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2025
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz. non-vacuous gen- eralization bounds at the imagenet scale: a PAC-bayesian compression approach. In International Conference Learning Representations. ICLR, 2019.
Zhanxing Zhu, Jingfeng Wu, Bing Yu, Lei Wu, and Jinwen Ma. The anisotropic noise in Stochastic Gradient Descent: Its behavior of escaping from sharp minima and regularization effects. In International Conference on Machine Learning, pp. 7654-7663, PMLR, 2019.
15

---

Published as a conference paper at ICLR 2025
A TECHNICAL LEMMAS
Lemma A.1 X, E Rd are independent random variables, then we have E [|| E [V]]]] - E [ - -II]-
Proof By the convexity of the L2 norm and the independence between X and Y, we have
| Y =xx By ||| Ylll 2 Ex Ey [X Y] Ex [|] [Y]]].
Corollary A.1 If X2 E ga are LLD.. copies of X, then we have E E|| E [X]ll] [[|X] - X2ll]-
Lemma A.2 If X1,X2 E Rd are L.D. copies of X. then we have E ||X E [x]]i?] S
E
||X
X2ll²j [|21|
Proof
|X - X211a E
=E
X 11
2B
X X₂
+ E
|221
=E
| 2
2B [
E [X2] + EE [|
=E
||
2E
X
E
[
+
|X||]
2
||x
E
[X]]1]]
Lemma A.3 Let f W Ryo be non-negative convex function. Let W W be random variable, then we have
E [[S[B s(W)]] < 2E [S(W)].
Proof
E [I [W] - f(W)]] < +
=E (E [W]) + [S(W)]
(f20) (convexity of f)
<2E [WW)
B PROOFS AND DETAILS FOR SEC. 3 PROPOSED OMNISCIENT TRAJECTORY
B.1 DETAILS FOR SECTION 3.1
In this section, we present illustrative examples of suboptimality and overfitting of Proposition mentioned in Section 3.1.
B. SUBOPTIMALITY
Different (S, W) has different empirical and population Hessians with different sharp directions. For instance, assume U1 is the only sharp direction for ($1, wi) and vy for (s2, Wz). Then the non- specific Σ would apply small noises in both directions. if and v2 are nearly orthogonal, the noises along va and in insufficient for (#1; wi) and (82, w2), respectively.
16

---

Published as a conference paper at ICLR 2025
B.1.2 OVERRITIING
When k output weights are sampled, they lie in a k-dimensional subspace. As a result, noises along the remaining d - k directions will not decrease the trajectory term. The optimal induces infinitely small noises along the remaining d & directions decrease the penalty terms. Therefore, the optimal Σ will have rank at most when samples are used. If one additional weight is sampled and it does not reside in the k-dimensional subspace, E cannot reduce the variance along the (k- 1)- th direction. As a result, the bound optimized for k samples is suboptimal for the + samples, indicating overfitting.
B.2 LEMMA B.I
Lemma B.1 extends Lemma of Wang & Mao (2022) that bounds the mutual information using gradient variances.
Lemma B.1 Let X,Y Rd., A E ER and O E Rdd be arbitrary random variables. Let N (0, 1) be a di-dimensional Gaussian noise that is independent of (X,Y, K,Y,A,,. A,O). Then for any a > O, for any deterministic gd functiion f of (X.Y.A) we have for any function Qt solely of (Y,O),
1 [II(X,Y.A)-
(U(X,Y,A)+oN) (Y,0) S
(B.1)
2n
and for any function Q solely Y
_ -
1
(U(X,Y,A)+
(B.2)
2.
Proof
We prove Eq. (B.1) first. Let 0g E Re denote k-dimensional zero vector. Define X'
[XT 02 0 Y' = [YT Oa, OTJ A'= = [O], A 0 and f" +di+ds x
Rd. Rd and Q Rd Rd by
f'(0, u) = fuidd Uiddy va, +1ad +dg), 2'(v) = 2(Wi:d Ud, +d tdaatad)
where 022 E Rd is the zero vector. By assumption, (X,Y", A') is independent of N. Consequently, by Lemma 4 of Wang Mao (2022), have
Y) S d (Y)
E [I + - YY) + 11]
log
do2
[IIS'(Y" + A,,') 2Y
20
By construction, have fY + A = A) and 12 !(Y, Combined with the definition of Y", they lead to Eq. (B.2) of Lemma B.I.
To obtain Eq. (B.2), let O" be the constant 0g, and S (y, o²)) = 0(y). Then by Eq. (B.1), have
I(f(X, ()) +o;; =
1
<
F.
20
1
[If(X,Y, A) -
E
20
which is exactly Eq. (B.2).
17

---

Published as a conference paper at ICLR 2025
B.3 PROOF OF THEOREM 1
By the technique of a change of trajectory, we first transform the generalization error of the original trajectory into that of the omniscient and SGLD-like trajectories at the cost of penalty terms:
gen(u", Pwis)
=E
[w) [sw)
(W))W) (
(Wr) (W(W))
(WT) (Wrr)
(WT) (W(
+ (WWr) (s(WT)
E (W)())) -
[a (Wr, Ar, (Wr,s')) S')]
= (WT) (W(r)
AΣ of (WT + ST, af (Wr + ST)
+E (WT,, S) Ar, ('),
(B.3)
With the penalty terms matching those in Theorem the goal reduces to bounding the generalization error of Wr. To this end, we apply Lemma that bounds the error by mutual information and then decompose the mutual information using the chain rule. Since ((w,)) is R-sub-Gaussing by assumption, we have
E ([((W))-Ls(((r))) - gen(ji, Psr 5) sV V I(Wr;S),
2R
where
I(Wr; SI(Wo:r:))
(Data-processing inequality)
=I(Wo: S) + Σ I(WaS Woa-1)
(Chain rule of MI)
T
S(W:S) 1o.-1)
(Wo Wo is independent of S)
Σ I(W-1 (ge Ag.) N;SS Wot-1).
(Definition 1)
To bound the stepwise conditional mutual information, apply Lemma B.I with X S, Y Wor-1. A (Wo.T, V) and
f(X, A) = W,-1 (w.,W,-1,,, V, Wo.-2) - (S, S, V, gI. Wo.T)). Since (X,Y,A) A) is a function of (S, V, 91.r, Wor)), which independent with the SGLD noise N₂ ~ (0,0021) I21), by Lemma B.1, we have
I(W:5| Wo.-1) =[(W-11 - (gt Agi) N;: Wo.-1)
W (g: Ag.) ((Wot-1) [|]
20
for any function Ω of Y Wot-1- To make things simpler by removing W.-1 and minimize the expected squared norm, we set Q (2( = ú1-1 E [gt] and obtain
I(W;S) S Wo.-1) 20
[1g. E [oe] Ag
(B.4)
Plugging Eq. (B.4) back to Eq. (B.3) leads to Theorem 1.
Remark B.1
The bound can be further tightened by setting Q( - [ E [Agr]
However, this eventually leads to an extra E [AG] in the trajectory terms of Corollary B.I. This extra term creates a "feedback" between instances that hinders optimization: when optimizing AG, we must consider all instances of AG(S, V. g1.,, Worr) even if we are optimizing against a single instance of trajectory to form Theorem 2. For simplicity, only set ((ma.t11) W E [gt].
18

---

Published as a conference paper at ICLR 2025
B.4 SELECTION OF Agi GIVEN FT
In the result of Theorem 1, the penalty and flatness terms only rely on rr among omniscient-related variables. Therefore, we can first fix IT and optimize Agi.T.
Assume for this subsection that rr is fixed as AG, function of (S, V.V.V Irr,, Wo.r), and that Anır is constrained to satisfy AG = Σ Agr. Penalty and flatness terms are now fixed since they only rely on I' and we only need to optimize the trajectory term
arg min Ag.r: A9
R
T Σ t=1
B
E [ocl
- Agell)
over AgI.T. We then set o, o, throw the square root and coefficients away, and move the summa- tion inside the expectation, arriving at
T
arg min
E
Σ
19:
E [or]
Ag: |?
Ag
Ag
Set i1 == 9t E [gt]. wisely:
Since the constraints are instance-wise, we can also optimize Ag: instance-
T
Σ
|p, Agel = =T T Ee-opry [IIe- Agill) 2T Emuur [ve- - A
arg min
Agrr:: Ag. AG
t=1
=T Ewor [ - - AG
where inequality due to the convexity of the squared norm and takes equality if and only if Ag is constant (W.F.. t), say A. Then we have Σ, Ag, = (E, 14) - T A AG, which implies
Ag: = gt
E [e E E [gg+) AG..
B.5 COROLLARY B.
Corollary B.1 Assume E(w.,) Gaub-iaa on μ for any w E W. For any deterministic
function AG of (S, 91.T, Wor) and any >> O, we have
IAW - E [AW,]] + AGI?
R2 no²
gen(je, PWis) S
A (Wr, - (Wr, (W) +E
T
where AWr = WT Wo the change before and after the training. If we give up further optimization, i.e., AG O, then have
gen(je", Pwis) Sv
nTo2 VAA + [AT (Wr, AT (Wr,
R2
where V [X] = |X
[ denotes variance.
Proof We prove Corollary B. by instantiating Ag. of Theorem
Since Agt depend on gI.T, we can use "; == - E' [gn] and D = Ee Pt when constructing Agt- Let AG be deterministic function of (S, V. 91.T, Wo.r) and a > be constant. Then
Agi = (v22 D) + AGG is indeed deterministic function of (S, Virr Wo.)).
With I'T E, Ag -. AG AG and o, constant a, by Theorem obtain
H2
gen(u"" Sou(u",Pwis) S,
E
11 (ur - D) - 7 AG
E (WTAAG(WT:5S) Asc(Wr,S")
n
E
A To²
(Wy + AG, T (W₂ + AG, S*)
19

---

Published as a conference paper at ICLR 2025
where
E V1
(ur D)
1 AG
Σ E
1
9+
E
(g-1)
AG
=T
1
E
Σ 9t
E
Σ
9t
AG
T
Since Wr Wo Ee St. have E.Ot
-
A WW and
Σ Vt (v 17 1
AG
- [
+AG]?
Putting everything together leads to Corollary B.I.
B.6 DETAILS FOR SECTION 3.3
Firstly, need explicit form of the penalty terms to balance the tradeoff. Therefore, the penalty terms are approximated to the second order in the surrogate optimization, where the Hessians at the output weight show up. The approximation is the same as Appendix B.7. Secondly, assume the empirical and population Hessians in the penalty term for the SGLD-like trajectory does not change too much, namely Hs + AG) HWT + AG) z HWr) Hy (Wr). ). to simplify the surrogate optimization. These simplifications lead the following surrogate optimization target (be noted that we have not replacing terms with Haat Heen or J):
min AG
3
3
R
E
tr
(Wr) - H,((W)))] [IAWT E [AW + AG| 1 AG (As(WT) A - H.T)
n
E AG (
where (s) is the Hessian of the empirical loss at the trajectory terminal and (w) is that of the population loss. If we optimize the above surrogate target, we will rely on testing/validation sets to estimate the population gradient and Hessian. Therefore, one may want to partially avoid this reliance much as possible by simply ignoring them. On the other hand, optimiz- ing against validation sets leads to better numerical tightness. To accommodate both needs, we leave it option by letting Amat E {AH(W)), (ss(W)) ,Mpon (WT)W))) Hs(Wr) H((W))
and J E ((r(()))) (Wr)), (Ls(W)) and replacing the empirical-populatinn differences in the above surrogate losses. Another difficulty that, one has to consider other independent runs when optimizing AG for one trajectory. To further simplify optimization, we modify the surrogate optimization target so that only one trajectory is considered:
min R2
3
[Haat)
|AW,- - E [AW7] + AGI AG J + AH
Hpe
where AG is function of the trajectory. To obtain a convex optimization problem, we partially remove the cubic root, giving
1 |AGT Apen AG,
min AC ||AWT - B[AWT] AGI AGG J AG
where C ==
Re
(Faat))
Bound parameter À here is used to compensate the "weight"
change after twisting the target: When the terms under the cubic root become smaller, the derivative of the cubic root increases fast; but after the cubic root is removed, the derivative of the norm becomes smaller when the term becomes smaller. Therefore, we should use large À to emphasize the reduction trajectory term. The convex optimization problem has closed form solution
AG (2XCII Hpea)
(2XC(AWZ
[AWr]) J).
(B.5)
20

---

Published as a conference paper at ICLR 2025
Since all variables in Eq. (B.5) is function of weights, empirical and population gradients and Hessians, which are functions of weights, the training set and the population distribution, Eq. (B.5) indeed specifies function of (S, V. VI.T, Wo.r). This omniscient trajectory is used in Theorem 2.
B.7 PROOF OF THEOREM 2
Based on Corollary B.1, Theorem considers the specific omniscient trajectory defined by AG -(2CCI +
[AWr]) + J) as Eq. (B.5). We first optimize a given the specific AG, for which we approximate the penalty for the SGLD-like trajectory to the third order. Specifically, for the empirical part we have
A (WT + AG,S)
[[s 2s(W 22 AG]]
(00
Be < VLs(WT+ + AG) + 1 7/5 fts(WT +
a.jiiiEEss + iSiSk
Σ
T ur [Hs(W] + +G)) + Eg [no]
2
where To is the residual at the third order and the coefficients aij.k are independent &. By the assumption of Theorem 2, we have Iral S D || and [E& [roll < DE [||||] Dd(d+2)o' One can obtain similar results for the population loss difference, and combining the both we have
= oT tr (Hs(Wr AG) Hp + AG)) +r,
E (WT AG, S) A RT (WT + AG,
2
where |r| S Dd(d + 2) = O1.
Similarly Theorem 2 of Wang & Mao (2022), optimize 0 to balance
R E
lAW,
B|AWs] AGII?
n
and
TTT
tr (As(WT+ AG) HWT + AG) o2B
2
For A, B> 0, A/o o² takes the minimum 3(A/2)2/3 B1/3
(//2B)
G.
1/3
nt (I + pen AW,) (2XCI + Heen) - 12
n
73/2 tr (AH AG))
(B.6)
The proof is finished by setting a this value and putting everything together.
B.8 DETAILS FOR SECTION 3.4
We take J v(Ls(WT) (Wrr) and sufficiently large À 0 so that [x, (Hpen/ /2X) < 1.
We assume the training sufficient so that V(s(Wr) as O. We also assume that the initial weight Wo is fixed.
21

---

Published as a conference paper at ICLR 2025
By applying || y/1² 2 -|/1=|1 + 2 Ilyll" and approximating the penalty for the omniscient trajec- tory to the second order, we obtain
gen(u" PwIs)
3
21²
[IAW, B (W V.L(() (2ACI + pp (WWW)))
EE +
p0mm 1)2 (2xcr1+++)))))))
2
n
+
E
+ E
(VLs(W))) (. (2XCI + Hpen) )
+ E
|AWT - E ( |I |1 Blod) - poou) + -A0
By selection of A, have (2XCI+ Henn) 1 1/2XC and I- (I HI+ 22C))11
Hpen)
(Hppn/2CC)) ppnn /2C. These approximations lead to
gen(n" PW(s)
21
E [Ir +G]]]] B |
- IE [AWT]] (iw.lllip....eee /4XR/ /111//42202
(VL)(W')) (AW))
(AW: - BAre-EJAWWE B(AWs))
+
(WT)) (WT)) T/22AC
A [AWT)]liHH + || J AA///A//
2R
E ||t (W AG) |AWr [AWT]l [AW.IllR....uue + ||l11 /4X202
n
E [JT (AWT - E [AW?]]
E [1J1I /2XC|
+ [IAWT [AWG + || [MPAa//ae))]]
The second term E (W,)
(AWr)) can also be transformed to Ressian-induced norms by exploiting the symmetry of the terminal deviation AW, E [WWT]- To this end, ap- proximate the population gradient (.,,,W,)) at the mean terminal weight [rr by VL,,, (Wr) Ro
LL [Wr]) + H (E [Wr])(WT- E [Wr]). Then we have
E
J (AWr - E [a)]))
=E
(E (W7)) V (E, + AH(E (W) EE(T)) - E[AWT))
||W E E[AW+]1
Plugging this back leads to
gen (n" Pwis)
3 2R
E (A + A
/Xx0= |AWr- [AWT]] BA ||/11 /4X202
221
|AWTT - E ]I
[IJI /2XC| + E AW - [AWr]] AHI + |
The above results apply without the assumption that vis(W,) ~ 0. Using the assumption V (s leads to Eq. (2).
22

---

Published as a conference paper at ICLR 2025
B.8.1 COMPARISONS WITH SIMILAR ALIGNMENTS WITH LOCAL GEOMETRY
In Section 3.4, Theorem 2 is shown to connect with an aligmment between terminal weight devia- tions/covariance and Hessians. This is similar to the work by Wang & Mao (2024b) and Wang al. (2021a) on the alignment and the more fine-grained structure of local geometry.
Wang & Mao (2024b) provides two types of generalization bounds for discretized SDE: the trajectory-state bound based on the statistics along the full trajectories and terminal-state bound based on the statisties after the last update, which are separately compared below.
The trajectory-state bound focuses on the alignment between the population gradient noise co- variance (GNC) E, (W, = Vz'zpl [VE(W, Z')] that depends on W, and the empirical GNC S), b(n Vo-Um V 2:)] that depends on (W,, S). Therefore, this alignment and ours are both aligmment between weight/gradien distribution and the local geometry. They also have the same dependence: the weight/gradient statisties both Wi-1-dependens and the local geometric properties are both (W,-1, -)-eependent. Differences mainly lie in the type of geometric properties used and how directly they reflect flatness. Wang & Mao (2024b)'s trajectory-state bound uses the inverted GNC C to capture local geometry. which is an indirect measure of flatness. By comparing Wang & Mao (2024b)'s Figure 2 and 3. one can see that their bound does not correlate correctly with Hessian traces even along training steps. In contrast, we directly use Hessians, allow- ing us better leverage the flatness as required by our motivation. Their choice of C, is a result of the discretized SDE algorithm instead of bound optimization (E, is the ingredient corresponding to bound optimization). Therefore, it is determined by the algorithm and it is difficult to switch to Hessians or other statistics. Another result of the above determining is that C, can only reflect the flatness at step instead of the flatness of terminal weights. In contrast, our alignment only involves the flatness at terminal weights. According to the empirical results from Wang & Mao (2024b)'s Figure 3 and Jastrzębski et al. (2019)'s Figure 2, the eigenvalues of Hessians during the middle of training much larger than those at terminal. Therefore, it is better only use terminal flatness, as it is in our alignment. Regarding Wang & Mao (2024b)'s 's terminal-state bound, both this bound and our bound involve terminal states. Still, their bound uses the inverted weight covariance instead of Hessians to reflect local geometries.
Wang et al. (2021a) optimize the noise in SGLD, with new information-theorett generalization bound as surrogate for real risks. They show the square root of expected GNC is greedily optimal under their bound and use this noise to improve the optimization and generalization of SGLD and closes the gap with SGD. Their results indicate the importance of direction and alignment of noises. If we put their algorithms and bounds together, then the result will be very similar with Wang & Mao (2024b)'s. Therefore, our alignment has similar differences with Wang et al. (2021a) as with Wang & Mao (2024b).
Moreover, our omniscient trajectory can be seen as a surrogate algorithm, and optimizing it is very similar to designing new algorithms with better generalization. As result, additional similarities and differences can be found on the goal of introducing local geometry. In Wang et al. (2021a)'s work, before Theorem 1, the only use of local geometry is in Constraint I on not rising empirical risks. Therefore, it is used to guide how to increasing noises without sacrificing the empirical risks. In our bound, as discussed in Insight 1, we use local geometry to guide how to pull terminal weights together/closer without changing the losses too much. Therefore, local geometry in alignments has similar role between Wang et al. (2021a)'s results and ours of guiding information reduction without sacrificing losses. Regarding differences, Wang et al. (2021a) decrease information measures by adding noises, while we do so by mapping things together or closer.
C DETAILS FOR SEC. 4 EXPERIMENTAL STUDY
Codes are available in the supplementary material or at wwtssciint-ouuudne
wttpcuuu.oom///eggee ze/
C.I HESSIAN-RELATED DETAILS
In experiments, all Hessian traces computed using PyHessian (Yao et al., 2020).
23

---

Published as a conference paper at ICLR 2025
The optimized omniscient trajectory in Theorem 2 requires products between besssian-related inverse matrix (I + penn/2CC) //X22= and vectors, say, Since Hessians of deep models are too large to store and compute with, it is ubiquitous to directly compute the inverse-vector product (iVP). For iVP, following Dagréou et al. (2024), we conjugate gradient method that only requires the product (I + Hppe/22XC) of the original matrix with an arbitrary vector e, which is supplied with PyHessian's is hv-product. More specifically, we approximate such that (1 + Heen /2AC
to obtain û. We run conjugate gradient method until the error than 1% of ||ull", or the maximum iteration 20 is reached.
[ Hpen/2AC)J
is less
C.2 BOUND ESTIMATION DETAILS
This section presents the details in estimating Eq. (1) of Theorem 2. For the randomness in weights, we train 1: > 1 models independently to obtain weights Wi.kk. To ensure the training data for each model is LLD. sampled, we assume the whole training set is sampled in an I.I.D manner, randomly
partition the whole training set into subsets { i=1 and use St to train W For fair comparison, then make unbiased positively biased estimates for expectations in Eq. (I) separately and then compute the bound using these estimates. Being concave, the cubic root will bring negative bias. Nevertheless, this negative bias also be found in baseline existing bounds. Moreover, in experiments, the expectation product under the cubic root is extremely small. Therefore, this negative bias will not influence the overall result too much and we leave it as it is for simplicity.
When optimizing the bound to obtain AG, we use a validation set 5" to estimate J.
For the penalty terms, iee,, E. [Aa SW Aar(Wr, S')] and E [tr (A (H(WT + AG))], unbiasedly estimate them with the test set. The same test set is used for all weights to fully use the data. The loss differences are directly computed by modifying the parameters and forwarding the testing data S" instead of approximating them.
For the flatness-optimized trajectory term, the estimate is more complicated, since we must also estimate E [WWr] within the estimate and handle the matrix inverses and the population gradients in J.
To estimate E []W,]] in expectations, one should use samples of Wr that are independent of W However, another draw of l weights results in further partitioning on the training set, and the data used for each weight is much less. To fully use the data, similarly to cross validation, we estimate the inner expectation by AW, = k-1
Σ Evei
AWY, which independent of W
To avoid the population Hessian in the inverses and make estimation easier, we select ueen (Wr). In this way, the inverses solely depend on the training data and the terminal weight, and their interaction with the testing set is only linear, which is much more friendly to expectations. Define E' = -- (1 + maaen//XC'- and F = (2XC1 Hppe)_1, which only depend on (S, WT). Initial experiments on CIFAR-10 indicate that it is necessary to include the population gradients for numerical tightness, and we set J == (LT) - (Wrr). Under these specific settings, we will estimate
E ||IEAAW E(AW,)) F(VLs(W+) - V.,)( |²]
(C.1)
Replacing the expectation with the empirical mean and the random variables with L.D.. copies, we obtain estimator of the trajectory term
- - - (W()))
k
24

---

Published as a conference paper at ICLR 2025
which is positively biased because
VCs 2)
E
E
k
E F| - -
- V / 115 w7
s w
k
|| (E{FAAAS- - ( S5t wi
Bst. W
| - E
(E (wi) 5 wi))
||WW -|E - -V
St st,W w
=E B(AWT - E -
(W
where the second step is because x|
Σ E [X] [
Among the above estimations, the most loose step comes from the last estimation, i.e., the estimation on the trajectory term, because other expectations are estimated unbiasedly and this expectation is positively biased. In principle, we can estimate it without bias. To see this, set A == FAWr (W))) expand Eq. (C.1) as
F(VLs
E
B(AWT )) - (VVL (WT)
E
||
x E IAWII
E
[||||]] [|A|]]]
2E
E BAW T
ET A
+
| FIExEAAVE (
=E
2E
|WW]
E [ETA]
+
TEXEIAWW,
which can be term by term estimated without bias. However, since E involves Hessian inversion, this unbiased estimator requires more Hessian iVPs (iHVP): Computing A requires iHVPs, while ET A and E x E [AWr] require another two iHVPs. As result, the unbiased estimator requires 4k iHVPs. In contrast, the positively biased estimate only requires 2k iHVPs. Since IVPs, especially iHVPs, are extremely time-consuming, we still use the positively biased estimator in the experiments.
C.3 TRAINING DETAILS
For both MNIST and CIFAR-10, 10, train 1 = 6 independent models from the same randomly cho- sen initial weight, which means each model receives 10, 000 training samples. We use 2-layer MLPs for MNIST Wang & Mao (2022) while ResNet-18 for CIFAR-10. Since bounded loss is natu- rally sub-Gaussian (Xu & Raginsky, 2017), we employ cross-entropy (CE) losses capped at 12 log c for both testing and bound estimation, where e represents the class number. Here, log corresponds to the expected gross-entropy at random initialization, which are heuristically the maximum of the average loss during the training. However, we still use vanilla CE losses for training. As Wang & Mao (2022), we use SGD with momentum of 0.9. The use of momentum does to violate the assumptions in Section 2.1 because each update can access history weights to recover gradients and compute momentum. ResNet-18 18 models are trained for 200 epochs while 2-layer MLP models are trained for 500 epochs as Wang & Mao (2022). We start from a base hyperparameter, where the learning rate is 0.01, the batch size is 60, and dropout weight decay is used. For the 2-layer MLP, the base hidden width 512. To enhance the generalization the networks on CIFAR- 10 for evaluating the bounds, we random horizontal flip and random resized crop. This of data
25

---

Published as a conference paper at ICLR 2025
augmentations does not violate the assumptions in Section 2.1, they can be achieved by passing in V the random seeds for the augmentations and letting Ir augment the samples using these seeds before computing the gradients.
For thnggl--hyerpprammtee variations in Figures 4 and C.5, we vary that single hyperparameter while keeping others the same as the base. For Figures 3 and C.4, we vary the learning rate and the batch size in a grid-search manner while keeping others the same as the base. When the learning rate is too small too large, or the batch size is too large, the model even cannot fit the training data well. Therefore, we exclude the training records whose final weight has a training accuracy less than 95%.
Models are trained on 12 NVidia RTX4090D GPUs for 2 day with auto mixed precision (BF16) and t torch. comp ile () to save memory and increase parallclization. To ensure consistency of model and training details with previous results, our codes are modified based on Wang & Mao (2022) with bound estimation modules totally re-implemented.
We must remark that the test accuracy of our models on CIFAR-10 is relatively low, because each model is trained by a one-sixth subset of CIFAR-10.
C.4 MORE EXPERIMENTAL RESULTS
C.4.1 MORE RESULTS FOR RESNET- 18 ON CIFAR-10 AND MORE DISCUSSIONS
Here, we display more results of existing bounds for ResNet-18 on CIFAR-10. Figure C.I displays the result of the isotropic Proposition 2 of Neu et al. (2021).
From Figure C. I, we find that the trajectory term of isotropic Proposition is insensitive to batch size. which is improvement over the incorrect tendency of Proposition 1. However, this im- provement is only partial compared to our trajectory term that decreases as the batch size decreases. Thanks to the flatness term that has the correct tendency w.E.f batch size under large learning rates, the bound also has the correct tendency w.t. batch size learning rate is large enough. Compared to this result, our bound can even better capture the generalization under all learning rates by having a more aligned trajectory term.
Now we foeus on the tendency w.r.l. learning rate. Our bound fails to fully correct the wrong tendency of the trajectory term of Propositions and 2 w.rll learning rate. Nevertheless, the wrong tendency seems partially corrected because the lines corresponding to the trajectory terms under different learning rates are well separated in Figures 1b and C.Ib of the existing bounds, while some of those from our bounds twist together Figure 3b. However, this partial correction is not enough to invert the wrong tendency, unlike the tendency w.e.. batch size. We conjecture it because the trajectory term is not solely determined by flainess, but a "product" between the Hessians and the output weight variance (see Eq. (2)). Increasing the learning rate increases the step size, which then increases the variance of the output weight. This variance increase is empirically confirmed by Figure C.1b: we display the trajectory term with o~ excluded in this figure. As result, what is displayed in this figure is proportional to the output weight variance. It clearly shows that the variance increases the learning rate increases. Therefore, there is a competition between the increased variance and the improved flatness. Unfortunately, it seems the increased variance overpowers the improved flatness in this process and our bound fails to invert the wrong tendency w.r.t learning rate. On the other hand, decreasing the batch size seems to also increase the variance and we must explain why the tendency w.r.t. batch size is fixed. The reasons is that, the variance is, in fact, very insensitive to batch size: From Figure C.1b, we can see that the variance is almost fixed when batch size changes. As a result, the trajectory term is dominated by the improved flatness when the batch size decreases.
Regarding the tightness improvement (mainly due to the trajectory term), Proposition 2 has a much smaller trajectory term than Proposition 1 (~ 10-2 v.s. ~ 109). It mainly happens because gradients of SGD are noisy and gradients from different steps often partially cancel. In Proposition 1, the canceled parts are still counted in the trajectory term while in Proposition 2 where gradients are summed together before taking norms. The components canceled in the trajectory indeed cancel in the trajectory term and do not contribute to the bound. This improvement is pushed further by the omniscient trajectory and the optimization under flatness (~~ 10-2 2 v.s. ~ 10 8). Since Theorem can be seen as the combination of Proposition 2 and the omniscient trajectory (see Corollary B. and
26

---

Published as a conference paper at ICLR 2025
Generalizztion
Estimated generalization
0006
Ir-0.001
1--0.02
13x10
Gen Gup Bound (Crose Limropy
12
(Cross Entropy)
Exiimatea Gen, Gap
11
11"10"
10
0.02
10
100
200
303
503
400
202
603
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure C.I: the isotropic version of Proposition for ResNet-18 18 on CIFAR-10 under varied flatness.
Generalization bound
Estimated generalization
000000001
0..05
1r-0.001
1m00.005
Ir-0.01
1m00.02
1m-0,04
64
4x10"
1.5
Gen Gup Bound (Closs Lmmroyy)
03
2x10"
(Cross Entrppy)
Exiimated Gen, Gap
10"
0.
0.0
200
202
400
600
2003
400
500
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure C.2: Wang & Mao (2022)'s bound for 2-layer MIP on MNIST under varied flatness.
how connects to Proposition by setting AG = 0), the improved tightness of our bound compared to Proposition 1 is two-folded, done by both Proposition and the omniscient trajectory.
C.4.2 RESULTS FOR 2-LAYER MLP ON MNIST
The numerical results MNIST (LeCun al., 1998) with the 2-layer MIP displayed in Fig- ures C.2 to C.5.
C.5 EXPERIMENTS OF EXISTING BOUNDS WITH POPULATION HESSTAN
In Figures 1, and C.3, we have displayed results of existing bounds without using population Hessians. That is, we computing the bounds whose (approximated) flatness term is
a 2T tr (Wr))
2
Generalization error bound
Estimated generalization error
1r-0.001
Ir-0.0055
1r-0.001
In-0.005
Ir-0.01
0.02
1--0.04
64
0.04
104
4x10
Gen Gup Bound (Cross Ennopyy)
Escimited Gen. Gan
(Cioss Camppy
02
16"
01
6x10
200
650
600
250
620
200
450
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure C.3: the isotropic version Proposition 2 for 2-layer MLP on MNIST under varied flatness.
27

---

Published as a conference paper at ICLR 2025
Generallazioo bound
Estimated generalizzzion
1r0000001
Ir 0L0035
Ir00.001
k-00005
1=0.01
1r-002
I--0.4
Gen. Gip Bound
(Cros Encopp)
0.3
0.3
1,00
1=-1
9.3-
03
10'
2-100
1-10"
(Cros (CrospEntrupp Entropp)
Eatimneed Gen. Gap
0.1
6.1
200
250
460
200
266
500
Baach
(a) Bound
(b) Trajectory Term
(c) Punishment Term
(d) Flatness Term
Figure C.4: Numerical results of Theorem 2 with varied learning rate and batch size.
= 10 on MNIST and a 2-layer MLP with
Weng (200
Pry. of (203))
The 3 Q-110
The 500-10000
The Q0-1-888
Gen CapyBound (Cross Entrgy)
10
10"
0.4
250
500
755
1920
0.00
0.01
0.50
Label moiic level
Hiddin loycr wills
Weight docny
Figure C.5: Numerical results for a 2-layer MLP on MNIST under varied label noise level, width, and weight decay. The isotropic version of Neu al. (2021)'s Proposition is used.
instead of T tr (Hs(W))) - tr (Wr)) In contrast, our bound depends heavily on pop- ulation statistics like the mean gradient and Hessians. One may doubts whether this comparison is unfair due to our extra dependence on population Hessians. This concern is amplified by the fact that the population Hessian potentially improves the numerical tightness of existing bounds: If one assumes as Wang & Mao (2022) that (W,, S Eg-N((o.22) (WT + E)], one is essentially assuming Hμ. (Wr) is positive semi-definite, whose trace non-negative. Propositions and 2, the non-negative traces are subtracted from the bound, which is tighter than the case where it is not subtracted.
To address this concern, we also allow the existing bounds to depend on the population statistics, especially the population Hessian, and evaluate them in experiments. The results are surprising and confusing: the population Hessians indeed have positive traces but their traces are orders of magni- tude larger than the traces of the empirical Hessians. As a result, the flatness term becomes negative. If one optimize the bound parameter a, one would arrive at -o: bound of the generalization error, which is obviously incorrect. We conjecture this is due to the second order approximation used in the expectation, which becomes inaccurate for population loss changes under the SGL DDklike noise of variance o2TI. This is also a drawback of existing bounds to require the large noises to de- crease the impact of the trajectory term in the a-optimized bound. In contrast, our bound has an orders-of magnitude smaller trajectory term and requires orders of magnitude smaller noises. This supported by empirical results in Table where the value of (average) o² (a, is the optimal values of in >, of of Proposition 1, the a in Σ o"TI of isotropic Proposition 2, or the in o"T in Corollary B.1 Theorem 2) in different experiments are presented. Be noted that although of values themselves seem already small for all experiments, the actual noises in the flatness term has covari- ance TTT where 7 is the number of training steps (~ 10tt for CIFAR10) 10) and identity matrix / has a size equal to the number of parameters (~ 10th for ResNets). As a result, the magnitude of a3 have significant impact on the second-order approximation in the flatness terms and the second-order approximation is much more accurate for evaluating our bound.
To alleviate this problem, we take absolute values of the flatness terms to make them positive in the main results of Wang & Mao (2022) and Neu et al. (2021). That is, we now estimate the bound with flatness term replaced by
0 2T
E Ir (fs((W)) ir (H,)
2
28

---

Published as a conference paper at ICLR 2025
Lr. and bas
label noise
width
weight decay
Prop. w/o pop. Hess.
15381.622 627 10-9
16372.131
10-9
19904.060 10-9
17674.105 10-9
Prop. w/o pop. Hess.
1434.559 10
1278.293
10-
1660.805 x 10
1341.889
10
Prop. w/ pop. Hess.
6659.479 x 10-9
1572.998.v
10
9175.748 x 10
9149 648 10-9
Prop. w/ pop. Hess.
676.226 x 10-9
545.839 10
767.169 x 10
643.772 10-9
Thm. 2 (A= 1)
291.674 x 10
696.630 x 10"
303.843 x 10
307.434 x 10-9
Thm. 2 (A 109)
4.354 x 10-9
6.269 x 10
3.866 x 10
3.167 10-9
Thm. (A = 10")
0.193 x 10-9
0.204 10"
0.244 x 10-9
0.167 10-9
Table 1: The average value of o across bounds and CIFAR-10 experiments varying different hy- perparameters. It can be seen that our bound requires orders of magnitude smaller noises for the SGLD-like trajectory's is Gaussian noises when is suitably selected. As a result, the distortion due to the second-order approximation has less impact to our bound.
Generalization ence bound
Estimsted genernlizaionn ernor
I--0 0006
I--0 001
1--0
In-O
I--0.02
Cap (Cross Entrops)
(Crov
TrresEotorr Entoopy)
115
A5511 13"
380
130
266
30a
400
50
260
420
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure C.6: Wang & Mao (2022)'s bound with population Hessians for ResNet-18 on CIFAR-10 under varied flatness.
The population Hessian is estimated on new validation set. After this modification, we display the numerical results for Propositions and 2 in Figures C.6 and C.7.
Since the modification only involves the flatness term, the trajectory term is almost the same those in Figures lb and C.Ib of the existing bounds without the population Hessian. Therefore, we mainly focus on the flatness term.
Since the subtracted population Hessian traces are orders of magnitude larger than the empirical Hessian, the flatness terms with absolute value are dominated by the population Hessian traces and are orders of magnitude larger than those in Figures Ic and C.Ic. As result, the existing bounds even become looser if fed with validation sets.
When it comes to the tendency, the flatness terms' tendencies w.e.t batch size also become wrong when batch size is large or learning rate is small. As a result, Proposition 2's (Neu al., 2021) partial improvement of trajectory term w.r.f. batch size can not compensate the wrong tendency of the flatness term, and the whole bound scales incorrectly w.r.t batch size.
Generalization ence bound
Esticated geteer scterrlizatiio eraa
1m-0.0005
Ir-0,001
k 0.005
Im-0.01
0.002
Gen Bound (Cross Enirspy)
(Coss Estrupy)
0.34
11 If
0.00
19*
30
Bacc
den
(a) Bound
(b) Trajectory Term
(c) Flatness Term
Figure C.7: The isotropic version of Proposition 2 with population Hessians for ResNet-18 on
CIFAR-10 10 under varied flatness.
29

---

Published as a conference paper at ICLR 2025
Wazz and Maa (2221)
Pop, & Nez ctaa (2321)
Tim. 3 (-110-00))
Then Qe1Leett)
Thx (ee1
Oenerdizaaion coee
[on, Caydhard (hamet
Dxpo
Hagasdoce
1.5
128
Figure C.8: Numerical results of existing bounds with population Hessians for ResNet-18 on CIFAR-10 under varied training data usage, label noise level, width, depth, and weight decay. The tendency w.r.f. weight scaling is also added. However, the results are taken from (unbiased) MLP on MNIST because ResNets are more complicated models and takes time to develop weight scaling scheme that does not essentially change the prediction. The isotropic version of Proposition 2 (Neu al. (2021)'s Prop. 8) is used. a model has channel scaling factor c, means the model has times the number of channels each layer compared to a standard ResNet-18.
Weg (60
Pry. af (203))
The (6-11-000
The 3 1-110
The Sev-1Deaa
Thes 10-1-0--
Gien. ClayBound
(Cross Fintropy)
10'
1.00
10"
0.10
0.00
-0.05
2 Weight scaling factor
Weight scaling factor
(a) Generalization gap on all samples. (b) Gen. gap on correct classification.
Figure C.9: Cross Entropy generalization gap and bounds under weight scaling of MLPs MNIST.
Figure C.9a displays the generalization gap on all testing samples. Figure C.9b displays the gener- alization gap on only the correctly classified testing samples. Since Figure C.9b involves negative values, we use log-scale above 10 and linear scale below 10-4.
The correct tendency of the flatness term without population Hessians and the wrong tendency of that with population Hessian indicate population Hessians have wrong tendencies when batch size large or the learning rate is small. This somehow points to the notion the generalization and overfitting of higher order statistics, Figures Ic and C. yet the population Hessian has the wrong tendency. Small batch size and large learning rate seem also helpful for this higher-order generalization.
the empirical Hessian has the correct tendency shown in
As shown Figure C.8, when weight decay is used, the existing bounds with population Hessian better capture the tendency of generalization under regularization. However, when the hyper- parameters in Figure C.8 are varied, the existing bounds with population Hessians are times looser than those without the population Hessians.
To sum up, the existing bounds cannot fully exploit the population statistics even we allow them access the validation set. Therefore, the comparison between our bound and the existing bounds is still fair, at least for existing bounds in its current form estimated after the conventional (Wang & Mao, 2022) second-order approximation.
C.6 EXPERIMENTS ON WEIGHT SCALING
To test whether flatness-related generalization bound captures the generalization under weight scaling, we train homogenous two-layer MLPs, whose output layer has no bias, on MNIST and scale the weight at the end the last step of the training. The results displayed in Figure C.9a, where generalization gap is computed for clipped Cross Entropy losses. In this case, our bounds with large À can well capture the generalization.
However, the results in Figure C.9a out of expectation. One would expect the generalization im- proves as the weight is scaled up because most MNIST samples would be correctly classified and the
30

---

Published as a conference paper at ICLR 2025
generalization risk measured by Cross Entropy would be decreased by increasing the "confidence". However, in Figure C.9a, Cross Entropy increases when scaling up the weights. We conjecture is because split the training set into k = 6 subsets, the generalization of these MLPs are in fact harmed and about 4% of testing samples are misclassified. As result, scaling the weight also in- creases the generalization risk the misclassified testing samples. The increase on misclassified lesting samples are arguably much faster than those on correctly classified samples, because weight scaling on misclassification pushes the loss to infinity while weight scaling on correct classification pushes the loss to zero. As result, the generalization gap increases in Figure C.9a.
To see how our bound captures the improved generalization on fully correctly classified scenarios, we filter the testing sets and only keep the correctly classified samples, simulating the cases where all samples are correctly classified. We implement this simulation using the following loss:
aarg max logits, = y min CrossEntropy(logitt, y), 12 log c).
Clipped Cross Entropy
Filtering the correct samples
where e is the number of classes It is bounded, thus sub-Gaussian and can be rigorously used in the existing and our bounds, whose results can be found in Figure C.9b. In Figure C.9b, when weights are scaled down, the training loss increases and increases faster (maybe the faster increase related to overfitting), decreasing the gap. However, when weights are sufficiently scaled up, the generalization gap indeed decreases with the decreases of scaling (maybe because now the training loss is almost zero and its decrease slows down). In the latter phase, the bound from Wang & Mao (2022) and Neu et al. (2021) all have weak or wrong dependency w.r.f. weight scaling. In contrast, our bound still well captures the improved generalization Our bounds even captures the strange negative generalization error. By inspecting the raw data, we find the negative value mainly comes from the empirical part of the penalty term, indicating that the negativeness of the bound value is not mainly due the dependence on the validation set.
both cases, if is suitably selected.
To sum up, our bounds can well capture the generalization under weight scaling in our experiments.
D PROOFS AND DETAILS FOR SEC. 5 EXTENSIONS
In this section, we present and prove the combined results of the omniscient trajectory wth existing results to demonstrate the flexibility of the technique. The proofs generally follow those of the existing results. But note that the omniscient and the SGLD-like trajectories may depend on random variables more than the training set S. Since the existing results often assume the trajectory solely depends on S. we must re-prove some of them instead of directly applying them to handle the extra dependence. Sometimes we must rearrange the order of the steps of the original proofs to insert the omniscient trajectory. Setting the omniscient perturbation to zero will recover the existing results, which we will not bother to restate.
D.1 COMBINATION WITH THE INDIVIDUAL- SAMPLE TECHNIQUE
Corollary D.1 Assume E(w,)) :) is Resub-Gaussian on i for any w E W. Recall n is the number of training samples. Then for any family { (Ai'r)" 1 of omniscient perturbations and any TiT e
(R>O)T,
T we have
gen(u" Pw s) <
R
n
21(W;;
Z,)
E A
(Wr, Z,) - (Wr.s") s*]
71 1=1
i=1
R
1 n
A E, of (Wr, 2,)
A o? (Wrr
21W 2z
E
13
n
n
T
E
19t -
E lg.
1_ - Aoill)]
0
Σ, (We, 2,)
- & (D.))
E
1=1
31

---

Published as a conference paper at ICLR 2025
where Agt is a deterministic functioo of (5, 9I.T,
g1.T, Wo.,, Z-i), and r == -Er= Ag rú:=
Remark D.1 The omniscient trajectory additionally depends on and -- the two featuring ran- dom variables of the individual-sance sample bounds.
Proof The crux of the individual-sample technique is first extract the summation in the empirical risk and change its order with the expectation, before other things like using the MI bound or building the auxiliary trajectories:
gen(n". " Pwis)
E
Pp (WT)
Σ
(WTr 2s)
(WT) - e(Wr,
72 1=1
1=1
1 n Σ
E (W77)
E(W7, 2.)
1 + n
Σ
Ary. (W1, Z,) Ary (Wr, S*).
where Wi = W + r is the i-th omniscient trajectory. can be seen that the omniscient penalty term has matched that in Eq. (D.1). Thus, will then bound the individual-sample generalization error of the omniscient trajectories. The process is similar to the proof of Theorem 1, Lee. adding Gaussian noises to obtain the SGLD trajectories:
n
E (W)))) e(W, 20]
E
(W7) E(W7, 2ii)
n1
of (WT + + 2,) AE, of (Wr + rri $5]
where Wi- = Wi W Ný. and N WN 12) is independent Gaussian noise. Again we throw away the already matched penalty terms and focus on bounding the individual-sample generalization error of the SGLD-lik trajectory. To this end, generalization error of an algorithm that takes one sample (but stealthily samples n - 1 samples as algorithm's 's internal randomness) and apply Lemma 1:
each term in the error as the
- E(Wg z.) -
1 Σ
" gen(n R
22 PRziz,
9)
--
2R 11
n
1(W;:Z))
1
i=1
R
21(W;; 7; 2. Z=:).
T1
We proceed by a similar arguments as in Theorem to bound the mutual information: for any i [n], we have
T
T(W: 2 2-) <T(W;: 2: 2-1)+ + Σ 7(W;; 7 W. oit-1, 2-1)
Σ I(Wi 1- (g. - Agt) + +N;: Z, W 0it-1 Z)
By instantiating X Zi, Y Wi Wit-- A (Wo.T, V. V-)-,), O Z-, and
f(x' r A Wi (9:(We11, (Z Z,- Z, Z Z4i..... Zn), V. Won-2)
- -Ag V, 91.7, Wo.,, Z-i)
and applying Eq. (B.1) in Lemma B.1, for any deterministic function L of solely (W)..-122 Z-i); respectively, we have
r(W; 1 (g: Agt) No: Zi
Wi Wo1-1:
1
E
W (9- Agi) 0:1-1: 2=:i)(?]]
20,
32

---

Published as a conference paper at ICLR 2025
Setting D2 (w) 0::00:11::00 2-1) أ E 19 Z- z-1] leads
I(W[ (gi - Ag;) N Z 1. - 220?
1
E
[19
E 19:
Z-1]
=oill
2
After putting everything together, the corollary is proved.
D.2 COMBINATION WITH THE CMI FRAMEWORK
Conditional MI (CMI) framework (Steinke & 7akynthinou, 2020) is developed to fundamentally solve the potential unboundedness of MI on some algorithms. CMI framework considers the fol- lowing random process: Firstly, a supersample Š Zim.0:1 E g of 2n samples are sampled as effective discrete "sample space". Then, n independent indices U1:= E {0, 1)" are sampled as the discrete "sample". The training is constructed by S (Z,,,, using the supersam- ple and indices. CMI then bounds the generalization by how much the output reveals the indices given the supersample. We combine the omniscient trajectory with CMI framework and give the following result:
Corollary D.2 Assume f(w.-) is Sau--Gaussian on y for any w E W. Then for any omniscient trajectory Agur that additionally depends on (S, U), and any 01. (R>0)² T we have
gen(n", (en(u,,Ww(s) 52
2R I(Wr;U U ()
A o (Wr, Z,)
E ot (W,, sr) Ar
π
R²
99: 1²
2
E
9t
E
96
S
E. of 17
(Wr, Z,) - (
E
where Agt is a deterministic functioo of (S, V, 9I., Wor,, S.U).
Remark D.2 The omniscient perturbation additionally depends on (3,U), which are the featuring random variables of CMI bounds.
Proof After paying the penalties for changes of trajectories, ization error on the SGLD-like trajectory. With the omniscient trajectory depending on (S,U). the distribution of the SGLD-like trajectory is determined by 1 We will apply the following re-proved CMI bound that handles the extra dependence on (5,U) of Wr:
focus on bounding the general-
E
(W/r) - Ls(Wr)] = 21F [(W())) - [W
3
E
(WT) (s(WT) |/] = =2Eg gen Pu, Pir)ess) 3
2R²
<2B3
| |
(Lemma 1)
2R²
18= 51] = =2 / /2R ([WTUUSS
<2
Eg
1(Wr;U U 15)
n
where the first step is because
E
(WWT) (s(WT))
1 "
=E
CZ (Wr) + L EZ (Wr)
Σ
Ezm
==1
E
E (Wr) Ls((VS))
1
E
[Ca(WW)) - Ls(Wrr)
2
33

---

Published as a conference paper at ICLR 2025
The following steps are similar the proof of Theorem and Corollary D.I. including application of the chain rule to the conditional MI and an application Lemma B.1, where
==S,Y= Wo:t1 (Worr, V, V,S.U) and ? Wt gt
D.3 COMBINATION WITH NEGREA ET AL. (2019))))' DATA-DEPENDENT PRIOR
To state Corollary D.3, we need extra notations. For an array a1d of length and a list U [ of indices, define dg = (av,) )= For two lists of non-repeated indices U E N and U² E N let U U2 be the list of indices that can be found in U but not in U2 ordered in Ut. Let U n be their intersection ordered as in UL.
Corollary D.3 Assume ((w.)) is R-sub-Gaussian on μ for any w E W. Assume the following specific form for the algorithm:
The randomness
BI.T, where Br E [n]] specifies the samples in the b-sized batch at step with no repeated entries;
The update function can be written as
9: WW
Won-2) V, V L.S, (W--1);
where V stands for samplewisely clipped gradient operator (i.e., compute gradients for each sample, clip and then average).
Let U E [n]m be the result of m uniform (over the remaining indices) sample indices without replacement in [n], independent previous random variables. Define
& = 10 Bel (FLsEiro (We-1) -b b
VLs, (W.-1)
to be the gradient incoherence (see it as the difference between the true update at step t and prediction of the update using only samples in U). For omniscient trajectories defined by bounded Agyr that additionally depends on U, then we have
gen(u", PW]s) S
R
16
A
+E
E. Ar
(W
(Wr,
A (
Corollary D.3 is application of the following more general result:
Corollary D.4 Assume f(w,-) -) is Sabb-Gaussian on μ for any W. Assume samples have interactions in updates in the following sense:
The randomness in the update is V BAT where B, [n]* specifies the samples in the b-sized butch at step with no repeated entries;
The update function can be written
Je 1, Wo.-2) = Pe (W,,,, Woa-2),
where Pt W. - z x W + W is defined by
|s| 0,
Pr Pr(W--1,3, 2or-2) =
fe(wi-1, 81, wor |s] > 0,
and ft W x 2x 1 + W is deterministic
Also, let U [m] be the result of m uniform (over the remaining indices) sam- ples without replacement in [n], independent of previous random variables. Lastly, that given SU, V. Wot-1), then Pe Pe Wo.t22)PP((:-- Sv. Wo.-2) and
34

---

Published as a conference paper at ICLR 2025
Pe SBnU, Woa-2) all have bounded second moments (e.g., by samplewise gradient clip- ping).
Define
buB (p SBAU, Wox-2) Pc 11,,, St, Woa-2)) & ==
b
to the gradient incoherence. For omniscient trajectories defined by AgeT that additionally de- pend on U, if Ag also satisfies the above conditional bounded second moments (e.g., bounded in value), then we have
gen(u", Pwis) $
R²
T
1 - A
+E A
aF (WI, St) - AFT of (Wr, S')
Remark D.3 The omniscient trajectory additionally depends on U and then & the two featuring random variables of Negrea et al. 22 SGLD bound.
Proof this proof, we combine Theorem 2.4 and Theorem 3.1 of Negrea et al. (2019) with our technique. As Negrea al. (2019), let St == S(1.......) Due extra dependence on (U, V) of the omniscient trajectory, the SGLD-like trajectory is determined by ((U, violating the assumptions of Negrea et al. (2019)'s Theorem 2.4. Therefore we must re-prove this data-dependent prior bound, during which we insert the auxiliary trajectories:
gen(u",Pw[s)) E (u(W)) Lsg (Wr)
=E (WiT)) Lsg (WT) E E. (W,, A
where the first equality is because given (S, Wr), Lsg (Wr) is an unbiased estimator on the em- pirical loss. With the penalty term matching with the statement of the corollary, we focus on the generalization eπror of the SGLD-like trajectory, which can be bounded by
(WVT) - Esg (Wr) = EU,so, [ [<n((W>) - Lig (Wr) Sv, v]]
-EV,So.V gen(ju" nm .w W..WUSSS
<EU.So.. V
21² I(WT: Si U = U,SU = Sv, V = -v- (Lemma 1)
By the "golden formula" of MI (Eq.(8.7) Csiszár & Körner (2011)) stating I(X;Y) = Ey [DKL (Pxly ll Qx)]. for any data-dependent prior Qw (u, Su, v) over WT+1 under the Definition 2.1 of Negrea et al. (2019), we have
Ey [DKL (PxiY || Px)] S
[Ca(WT) Log (Wr)
V
2R n m
(DKL (PWor |l Qwa SU, ri]]
$ Ly,su,
E 1,,
Now turn bound the KL divergence given any realization (u, Sus Sh of the conditioned variables. We restrict the data-dependent prior to have the same distribution as the real distribution for Wo. Le., Pwo u.s.. = Pwo Q Qw: (u, Su, v), By Proposition 2.6 of Negrea et al. (2019), we have
DKL (P (Piar uo Qw,, (u, Sa, v)) s DKL (PWorte.ousss Qwi (u, Sus suy u))
T
[DKL (PiilWorr.eeeeee. ll Qwylwarn (u, Sus su, b))]
E Eoon- Pwnc in
(D.2)
35

---

Published as a conference paper at ICLR 2025
We fix a condition c (u, Su; sq. U, W0:--1) and turn to bound the KL divergences in the above expectation. We also instantiate the data-dependent prior to the prior designed by Negrea et al. (2019) for SGLD, which uses samples in St predict the update using B, so that the error between the predicted and the true update eventually becomes the trajectory term. The gradients and updates depend the model weight and so do the predictions. Therefore, the data-dependent prior must somehow access the (t - 1)-th weight of the original trajectory to make predictions. In Negrea al. (2019)'s proof for SGLD, the data-dependent prior technique is directly applied the original trajectory and thus the (t - 1)-th original weight can be directly found in the condition c. However, in our case where the technique is applied to the auxiliary SGLD-like trajectory, we only have the latest step of the SGLD-liike trajectory in the condition, which is not the original SGD weight but corresponds many 1 or W on the SGD the omniscient trajectory, given that the SGLD-like trajectory is constructed by adding Gaussian noises to the omniscient trajectory. Directly predicting and comparing gradients at the SGLD-like trajectory will make the final result related to the SGLD-like trajectory. Switching back using local smoothness along the trajectory will result in "local gradient sensitivity" term as Neu et al. (2021) that accumulates very fast along the training process (Wang& & Mao, 2022). The ideal adapted proof would be somehow "tracing back" to the SGD weight W,-1 that leads to the SGLD weight in the condition, making predictions based on the SGD weight Wi-1, and comparing it with the update gr (W) on the SGD trajectory, followed by averaging over the distribution of Wi-1 given the SGLD-like trajectory weight. This intuition suggests a coupling between the true and the predicted update through W which reminds us of a lemma by Neu al. (2021):
Lemma D.1 (Lemma 4 of Neu et al. (2021)) Let X and Y be random variables taking values in Rdd with bounded second moments and let a > O. Letting ~ o0,,a2I) be independent of (X, Y), the KL divergence between the distributions of X + and Y + e bounded
DKL. (Px+c II Pr+c) S 20 [ |IX -11
20
where the expectation is taken over any joint distribution with any coupling between X and Y as long as the marginals are still Px, Py, respectively.
The arbitrariness of the coupling allows a specific one through W.-,. Therefore, we turn to construct the coupling of the data-dependent prior.
Since the coupling involves multiple random variables, it must start from a joint distribution. For be the joint distribution over V given by the Markov chain Mor + (M, M,) + (Me, M,), where Morr is sampled from PWo (u. -1 and
any condition (u, 8a; s v, (i):-1), let Q Mo Mft. Md,
M, : 11 Ma.-2) - Agi(s, u, UI.T, Mo.T, u)), PMMMMMM, 8u: Mon-2)) + b - uel Pe (M, 8u Mor-2);
|uel
Mr -
(D.3)
b
M, -M, + Nr, = Nt, + +Nz.
where j; = jn is the samples in j contained in the current batch, N; ~ (O0 o) is independent of other random variables in the random process defining Q. By assumption, terms defining M, and Mr have bounded second moments, satisfying the assumption of Lemma D.I. It can be easily verified that the marginal distribution Q Qit,le Pw.le- Moreover, all used samples in Eq. (D.3) are contained in Sy- As result, the marginal distribution of M; is independent of Si- Therefore, by setting the data-dependent prior W (u, Su, v) Q (igg (M....ss)) y wo. 1,), M. where (-) is a function that outputs some (conditionally) supported realization of St. we have Q @wiWarr (u, Su, u) u) = Q Mege With the help of this coupling and Lemma D.I. 1 can bound the
36

---

Published as a conference paper at ICLR 2025
KL divergence in Eq. (D.2) by
DKL (Pizie Qwiiwoo-, (u, Su, su,
DKL (Qiaale Qaasiee
=DkI (Rin+tife ll QM((N//)) S 20, ( NT) 20gle || xirl
(Lemma D.I)
1 BQ ?
[ Mo:-22 - (ge (8, v. M1:T; Mo.T, u)
- -(1a Pe Su.: Mo - luel b
Pe(Me-1. Sur Mose-2) )|²
Eq. (D.3)
[o((W 2orc2) U. Wo.r.
Ep
20
[ue] pt (W2 Suce Wor-2) b - Wel (W Su, W Wort
)
b
E ||E. Agell [c]
?
where the penultimate step is because, by construction, we have Qo.rle PWo.rle
After plugging everything together, interchanging the order of the expectation and the square root by the latter's concavity, and putting expectations together, the corollary is proved.
D.4 GENERALIZATION OF STABLE ALGORITIMMS AND GD ON CLB PROBLEMS
Theorem Assume n, d N+, L O and rj > T E N+ Assume the algorithm starts from fixed initialization Wo = E W and the whole trajectory remains within W, i.e., Wor E WT+1, For any CLB problem (W, 2,0) E) CLDD and any data distribution μ E (2) we have the following dula-dependent and -agnostic bounds that recover stability bounds:
gen(u", Pwis)
LD (W + AG; Z, 2
S inf AG
Σ
+
E (Wr,
2.) Asc: (Wr,
1=1
2L
[ WT E W Z-:]l]-
n
If the algorithm is a GD algorithm using projected subgradients of step size and step count T, we have
gen(u", Pwis) < 8L2 8T.2T
n
Therefore, we have the following worst-case generalization error bound for CLB and GD:
sup gen(k"" "ennk" < 8L2 SL2Tm n
sup (W,Z,fjECZ.D (2
Proof Let (W, CL.D be an SCO problem and μ Mi(2) be a distribution of samples. Therefore, the sample loss f, the empirical loss Ls(+) and the population loss Ls(-) are L-Lipschitz w.E.t the weight and W bounded with diameter and convex. Let Wor be the trajectory given by GD, where W, W.
The proof essentially resembles the uniform stability argument of Bassily et al. (2020) through the superior expressivity added to the MI bounds by the omniscient trajectory. Since uniform stability
37

---

Published as a conference paper at ICLR 2025
considers replacing one sample in the training set, we need to focus one sample instead of the entire training set. To keep other samples "unchanged" in the replacement, we also need to condition on other samples. Therefore, we will start from the "individual sample" technique (Bu et al., 2020). This observation motivates us to use Corollary D.I. However, losses in CLB problems are not sub- Gaussian in general. Therefore, (2023) that replaces sub-Gaussianity with Lipschitzness and boundedness. By repeating the proof of Corollary D.I but with Theorem 11 of Haghifam et al. (2023) instead of Lemma 1, we obtain
need variants similar Theorem 11 of Haghifam et al.
n LD
T
gen(u", (en(s, -
Σ
21
WT +
Agk: Z; Z-1
n
1=1
E Ary (Wr, 2.) - -ry (r, 8*)]
(D.4)
where {Agi,] 1= is family of omniscient perturbations, which moves the terminal within W, i.e. WT + 17 E W.
Remark D.4 The SGLD-like trajectory is not used in this proof. This is because the SGLD- trajectory is used only to bound the MI of the omniscient trajectory. As one will see later, the omniscient trajectory in this proof has very special form (constant) with trivial MI bound (0). As a result, an MI bound through the DGLD-like trajectory is no longer needed.
This proof proceeds by bounding the penalty terms in Eq. (D.4) with the help of the inherited Lips- chitzness of É in the empirical and the population loss:
Dry (Wr, Z+) = Lz (WT + r2) Ez (Wr)
s L ||
E Ar, (Wr, s') = E [WT + r - Cw(()) -,,r) - LE ||
Therefore, we have the following bound for the penalty terms:
n
2L n
Σ
Ar. (Wr, Z,) Ar,, 5
Σ E ||]]
n
Plugging this bound Eq. (D.4) leads to
LD
2L
gen(y" PWis)) S
I(WT + P Z, Z_:) +
E
1=1
By setting Agt 91- E Jot
Z-i). we have Wr + ríJ E W Z_1]. By convexity of W and
that Wy E W, we have W, + E W. Notably, given Z_ Wq + D is constant regardless of Z₁- Therefore, we have T(Wr + T4; Z z_ = 0 and
2L
gen(/", " PwIs) S
n Σ |+
(D.5)
Note that E Wr E [ Z-IIl Z_,] is the expected distance to the center given Z_,, By Corollary A. this value is smaller than E - WWW Z_i). where W/ is the terminal weight trained by the same Z_, and an independently sampled i-th sample.
Now we turn to GD. By Theorem 3.2 of Bassily et al. (2020), given Z- | WT W of GD has data-agnostic upperbound 4LVTTy + 4LTn Plugging this bound, we obtain
gen(n", Pw]s) ([l(u, _ 2L
Bz_. E []]WT - E [Wr Z-IIl Z_1]]
i=1
2L
Ez_; E []] W 2-11
8L2Tm
<8L2VTT +
n
38

---

Published as a conference paper at ICLR 2025
Remark D.5 Although adding the omniscient trajectory can address the limitation of representative information-theorette bounds CLB problems, currently do not know whether it is exactly our technique that makes it happen. This is because, between the Gaussian Gaussian-perturbed individaa-ssapp (C)MI bound considered by Haghifam et al. (2023) and the omniscient bound, there exist bounds derived by non-isotropic Gaussian perturbations, bounds by non- non--aassian but general independent perturbations, and bounds by general weight-dependent perturbations fudgaran et al., 2022)). It is possible some of these bounds can already audrress the limitation but we have not obtained any positive or negative results for them before the submission deadline. Nev- ertheless, our technique makes the proof extremely simple and without it the proof will be at least much longer. For example, will be harder make the (conditional) MI diminish without knowing Z-
Caate-Distootion bounds (Se-
D.5 EXTENSION TO G-LEEARNERS ON CLB PROBLEMS
Attias et al. (2024) have found that any e-learners have CMI at least Q((/e) on CLB problems, indicating CMI-accuracy trade-off. Here, learning algorithm is e-learner is for CLB problem (W, 2, e) if fir every data distribution μ over 2, the excess generalization risk of the algorithm is at most
Our Theorem 3 and its proof have given an intuitive alternative to the trade-off: although GD or stable learners have large CMI themselves, they quite close learners with low CMI. However, Theorem 3 assumes GD or stable algorithms, yet Attias et al. (2024)'s trade-off covers more gen- eral algorithms. Therefore, we explore whether our technique and alternative can extend more sospectation-ee-canneee in Theorem 4. It states our technique can indeed extends to more e-learners under some assumptions. It also states if one sees the omniscient trajectory augmented MI as new information measure, then the
trade-off is penetrated because both can vanish as n + oc. However, the result still partial and preliminary, because we only covers e-learners that are also O( )optimizezs,,, i.e., they are "well-behaved" in the sense that its excess optimization error is not too large compared the excess generalization error e. Nevertheless, we believe this assumption is rather gentle for well- well-behaved learners used practically, e.g., deep models.
> 0 and let Theorem 4 Assume d N+, L.D>
E CL.D be CLB problem. Assume for
any w E W and z E 2, we have E(w, 2) E [-ID, [-L,D++LD) +LD). If not, shift the loss functions. This shifting does not affect the excess generalization risks the excess optimization errors, which this theorem mainly assume on.
Let O be function of sample number n. Let (An : 2" W) be a family of expectation- learners for (W, 2, E). That is, for every sufficiently large n, for every distribution over 2, one has Esnyy (An (S)) infwew Lu(w)= Se
Assume } is also an -ottonite That is, An has an excess optimization error satisfying Esnya (S))A()) inf [s(w) & O(e) for any sufficiently large n.
Assume n N+ is large enough so that e bounds the excess generalization risk and O(e) bounds the excess optimization error. Then, we have
LD "
gen(u", Pw]s) <
inf (AG*) )
2W + AGG; Z_.)
E D|AA (W, AaGs (W,
n
71
i=1
LD
<0(c) +0
where W == (S) is the output of the algorithm, and (AG}) ) Jim is a family of (one-step) omni- scient perturbations that additionally depend on i and Z
Remark D.6 We assume the algorithu {An} is deterministic but the extension to random ones is straightforward.
39

---

Published as a conference paper at ICLR 2025
Proof After repeating the initial steps of the proof and obtaining Eq. (D.4), have
n LD
gen("", ;W]s)
21 (W + + 2-1)
n
1
+
E A AC (W, - (W,
1=1
LD
21 (W + AG"; 2,
Z
Cs( + AG") - Ls(
72
72
E ( +G*) [w))]
72
After applying similar omniscient perturbation AG Theorem 3, we have
E [W
Z_1] W as in the proof of
LD
gen(k", Pwis)
2T ( 2-1:2
LS( W z_1) - [W
1=1
+
E E1| (E W Z_1l) - (W)II
E
Ls( W 2-1)) [s(W)
1
E | [E W z-1) (W)I)
n
(D.6)
We first bound the population loss difference by applying Lemma A.3 to the difference with convex f(w) (w( Lu("") ) 2 O, where w" is the weight with optimal generalization risk:
W Z-1)) [_(w")) (C,(WW) - Lx(w")) z-:Il [ (E W Z-1)) (W)II =Ez_, E ( E
<2Ez-, E [f( Z_ = 2(T (C), Lu(w²)) s 2c.
Now we bound the empirical loss difference. Let w the empirical risk minimizer of training set The main difficulty is that different W corresponds to different S, and one cannot find an f to be (.(+) Ls(wr, and Ls Ls (wr.) ) the same time. Fortunately, with the individual technique, we can put weights corresponding to Z-; together. Most of their training set is the same, while the only different sample only contributes 1/n of the loss, which vanishes n co. Therefore, have
E [SS 2-:)) - - =Ez_ [ [B[ 2-1) Lss()) [2:]]
where
Cs(E [W| Z=.)) [s(W)) [=]] < n
=n-1B E [Z (E [W] 2=1) (W) [24]
E
Lz, (E [ Z_1l) Lz, (w) Z_n]
n
(Lipschiteness and boundedness of CLB problems)
n2
13 [z. (18 [W 2_+1) (2) |2-] +
LD
40

---

Published as a conference paper at ICLR 2025
Applying Lemma A.3 with convex f(w) == Lz (w)
Lz .(09 ,) 2 0 leads to
1 Ls( IW z-:l) - [s(W) [--]
n 23 [f(W) Z-1] +
LD
n
n-1 1
n
2B [z. ,W) - Lzx (w2 .) Z-s + LD
n
E (W) (17,)) Z-i 20)
11
22
[Ls(W) (m-1 Lz. (wž. / (z( Z-- +
1
Lz,
LD
n
n
11
Taking expectation over =-; leads to
E [[sSBB[W z--)) - [s()
s 21E (S((W)) -2 n (n--B
=2"-Ayx E [ (ws..
1 [
LD (D.7)
n
where Sn-1 pm-1. We need relate m=EE [Ss_ (wB...) ) + f [2xx²) to E L (wxz
We can bound L [Zx(W)) by LD again by the assumption that losses are bounded in [-LD, LDD, leaving n-1 [..., (w3, (w3,_.) ) and E [[sww)) which somehow forms the sta- bility of the empirical risk minimizer. Since (projected) GD with fixed initialization has been proved to be stable and has been proved to approximate the empirical minimizer, we use (projected) GD to bridge the empirical minimizers.
According to Orabona (2020) and (1) of Haghifam et al. (2023), on CLB problems, (projected) GD algorithm GD, = - W with step size r7 and step count T and a fixed initialization has approximate the empirical
D² + (log7+2)m12
an excess optimization error 2n/ minimizers using GDs with errors
As result,
D2
(log Ta-1 + 2) 2)
E [$S. (WS._1)) E [C.. (GD I.T, (Sn-1))]
21mn-1n-1
2
E [s(W3)) EE [CS. (GD (8))
D²
(log T, + +2mm,L2
<
2n,TTn
for any (nm 1 Tn-1) and any (Jm: TT). We assign m)-1 = 1/: Ta-1 Tn and select them suitably as in Haghifam et al. (2023)'s (3),,), which bounds the approximation errors by (LLD) vn) at the same time:
E [s.. ( (w....) )} - E [R....
(Sn- 1))
50
LD
Vn
LD
E [S(WW)) Ls., (GD, GD. (S))
50
Vn
which is the removal-based stability of GD. To this end, for each 8n-1, we construct an artificial sample 25,-1 such that
We then need to relate E Ls._, (GD, r. (Sn-1)) and E [.s (GD (5))
[ss
f(w. f(w.2s_.)) 2S...) = (w),
and denote $ == Sn-1 U (2s. }.. By construction, we have Egt - (w) = Ls. (w) for any w E W. iee. the optimizations using S,-1 and S happen on the same loss landscape. Since GD relies on (sub-jgradients and thus only relies on loss landscape, have GD, D T, (Sn-1)
41

---

Published as a conference paper at ICLR 2025
m. (1) If pair Sn-1 and Sn, then S and Sp only differs by sample, allowing us to apply the replacement based uniform stability for GD on CLB from Bassily et al. (2020):
Cs._.. - (GD, Dn (Sn )) - E Ls, (GD, t/ 2, (S))
1 E Ls.... (GD mr (S 1)) Ls. (GD, Pr. T. (5))
[Sa-1]
< Es..
[E [||D Mn. (S GD nn in (5) S
= L Es.-1
[ [[]GD To (S GD no. To (5) [
0 122 L2 Tulm + L2Tn
With the same selection of (Tx, Th), the above difference can be bounded by O As result, Eq. (D.7) can be bounded by
LD Vn
CB [W
z-+)) [W)
22 E [W))
n n- L [[S(wU3) n
LD
LD vn.
-
+
n
[W)]
[ +
LD
L.D vn
<2
+
n
LD
2 0(c) +0
+
Plugging everything back to Eq. (D.6) finishes the proof.
D.6 EXTENSION TO SGD AND SMOOTH LOSSES
In this subsection, we apply our technique to SGD under smooth losses. We will prove an omniscient information-theone theoretic bound and then show it recovers some existing stability- based bounds.
To prove the omniscient bound, we need basic form of information ieformatioon-thoortiic bounds like Lemma and Theorem of Haghifam al. (2023). This done by results from Lemma D.2 Corol- lary D.5. After that, we make the basic bound omniscient in Theorem 5. Finally, we bound the omniscient bound by the stability- bound Proposition
Lemma D.2 (2-aasserstain Distance Generalization Bound under Smoothness) Assume for any sample E Z, f( is non-negative, differentiable in Hd and 3-smooth, i.e., for any w, w/ E Rd
[V((w,)) VE((w')) =)]| S 8 -w'l-
Then we have the following 2- wwwssrstaann-aasseee based informatton-thhoreti (indiiitattiii bound:
gen(u", PwIs) 3 [Rs(w]] +
B+ 2n
Σ
Ez. We
(PW1,,
where >> 0 is constant, and W2 (..-) denotes 2- Wasserstein distance, the optimal transport under squared L2 norm.
Proof Let constant 7 > Given any index and any instance of the i-th training sample 2;s let n. be the coupling that approximates the 2-Wasserstein distance Wa (Pwiz,=e, Pw) between
42

---

Published as a conference paper at ICLR 2025
Pwizz=e and Pw by error most >> 0. Then for any > 0. have
gen(p", Pw]S) B(w.))).w (e(W", Z,) E(W, Z,)l
Ezz Baw.wywng [e(W', 2,) -E(W, z,)]
B(w,W' (W' W) Ve(W, Z,) (W"WW)VVEWW, 222 I lW'- w(]) -W---
8
Ez:
We then follow Lei & Ying (2020) to handle the inner product as in their Appendix B: (W' - -W) Ve(W, 2.) < I W|I |V&WW, ZJll
W' wil 1
[VE(W,
2
2
Thanks to the self-bounding property of positive smooth functions (Lemma A.I of Lei Ying (2020)), we have |/V(W, 1 = 2,3 ((W, Z;) and
B e(w., Z.). (W" - wW) V((W, Z,) = | - -WI1=
7
Plugging this back leads to gen(n", PwIss) S
3 _e( L(W. Z,) +
3 ++7
I -w1|]
Ez,
72
2
1=1
s
[W])
18+7
a Σ
Ez,
W2 (Pwiz,: Pw) + +
2n
By arbitrariness of E > O, we have
gen(u", Pw]s)
= [[s(W]]
+
8 +
Σ
Ez.
Ww (w]2,,Pw)]
7
2n
i=1
Lemma D.3 For any random variables (X, Y, Z) such that Y is independent of Z, have Ey (w (Pxjr, Px)] < Ez [Ey~Pvi [W2 (Pxix.2,
Proof Let Th. be the coupling that approximates the 2-Wasserstein distance W2 (Pxlg.2r Px(m)
between Px|y. and Pxly by an error at most c > 0. Then for any > O, we have
Ez [EY [W2 (PxIr,2, Px]z)]] (x,x)xxxl22]]
-Ey Ez [W2 (PXIY.2, Px [x12)]]
(Independence between Y, 2)
BBy Ez E(X,X')~mt2
||x x'11
=Ey E(z.x,X)) v
| x11²]]
Eyy [W2 (PxIY, Px)]
The lemma follows the arbitrariness of 6 > 0.
Lemma D.2 is very similar to the fact that I(X; Y) < I(X;Y if is independent of Z. Follow- ing this similarity, write (conditional) expected 2- Wasserstain distances similar (conditional) MI, or equivalently, replace the KL-divergence in MI with the Wasserstain distance to compare the prior and posterior:
43

---

Published as a conference paper at ICLR 2025
Definition 3 For any random variables (X, Z), ler
In,, (X;Y) EEy W2 (Pxir, Tw (X; 2) F= Py [W2 (PxiY,2; Px|z)]]
Lemma D.4 X and are independent given Z, then T (X;Y 2) = 0.
leads to the following corollary:
Corollary D.5 (Stability-Style 2- Z-Wasserstain Generalization Bound) Under the same assump- tions as Lemma D.2, we have
I LB [[s(W)]
3 7
gen(u", PW]s)
+
Ina (W; Z;)
7
2n
iml
8 S γ
[W))
3 + 7
+
Tw (W; Zz Z_;)
2n
Proof Z; is independent of Z
Theorem (Omniscient 2-Wasserstain Bound under Smoothness) Assume is non-negative, differentiable and B-smooth. Let 7/1 72 > B. Let {AG' 1 weight) perturbations, each of which additionally depends on Z4, Then we have
be a family of omniscient (output-
gen(n", R Pw]s): I
31
[&(W + AG Z:)] +
+ 71
Σ
Twi(W + AG Zz Z-i)
2n
+ n
AG (W,
Dac (W,
and
B
3+1 71
n
gen(u"; Pw]s) 1
[E(W AGG Z:)] +
Ina (W + AG"; Zi Z
2n
72
2 R [s)
3 + 72
Σ |AG|)²
(D.8)
7/2
Proof This proof is very similar to the proof of Theorem 3. By repeating the proof of Corollary D. but with Corollary D.5 instead of Lemma 1, we obtain the first inequality in the theorem statement:
115 Σ E [e(W + AG", AG) 71 n
B B 7 2n
n
gen(n", (ann) S
Iw (W + AG'; Z, Z_,)
E [A (W,2) AAG (W, 2')
The penalty terms can be bounded the same way as Lemma D.2:
Asc (W,=)l < (AG') VV(W, =)|
8
AAi 2
2
< AG' [[( 2] + 3
||
2
22 ||A ||2
8 [(W,>)
_
||
2
72
2
((W, 2) +
AG
72
2
44

---

Published as a conference paper at ICLR 2025
Therefore, have
E [&(W + AG", A.)] + 7
gen("",Pwi)) <
Σ
Iw (W + AG'; Zi Z-i)
m1
2n.
1=1
1=1
+ 72
E [W)) + E (C)((W)) + (3
72)
1 n
ΣE
n
|AG||]]
Population loss E (W)] appears at the right of the inequality. To move it to the left, we pair it with a virtual empirical loss term and moving the consequent gen()"" Pwis) to the left:
B
+ m
gen(n", Pwis) _
[eW + AG", 2i)] +
Iw (W + AG'; Zn Z-:)
71
2n
i=1
1=1
+
8
2
[s] + gen("", ;
8 + 72
Σ E 1=1
n
| |?]
72
n2
(1 3 _ gen(u", Pwis) _
8 1
[&( +AG Zi)] + B 71
Iw, (W
AG" Z-i)
72
71 n
2n
23 +
[[s(W))
B + 72
|AG|S
E
E
Y/2
Restricting 72 > 3 allows us to divide the inequality by the inequality:
f 72
without changing the direction of
8
3 +
gen(u", Pwis) (on(s,,w((s)
E
e(W +
AG",
22)]
+
Iw,(W ( + AG"; Z, Z-;)
1
11n
2n
i=
1=1
2,3
[W))
3 + 72
E [|AC)())
+
EE
7/2
Now that we have proved the omniscient bound for smooth losses, we turn to recovering some existing stability-based bounds.
Proposition 3 Under the same setting as Theorem 5, we have
gen(u" Pwis) S
inf
[RHS Eq.. (D.8)
< inf 28 E [sW)]]
3 + 72 +
stability
72 1
72
2
where
stability = -EZ.,V
1
W W/| [
Z
(rephrased in our notation) is exactly the Eg on-average model stability in Desinition 4 of Lei & Ying (2020).
Remark D.7 Proposition 3 recovers the relationship between stability and generalization in Lei & Ying (2020)'s 's Theorem 2(b) up to constants.
Proof After setting AG WE [W| Z-,, V] as a function of Z_, and V. we have W +AG E. [ Z-i, V), which is function of Z,-independent 7-; and V. As result, W AG' is
45

---

Published as a conference paper at ICLR 2025
independent of Z and (W + AG"; Z Z-4) 0 according Lemma D.4. Therefore, we have
gen(u", Pwis) v 2m,
inf
RHS of Eq. (D.8)
23 [W]) + 8 72
n Σ RR [||G)]*]
$ inf
[e(W + AG 2.)] +
2>0,7>0
1
B
7/2
7/2
2,3
[W)) +
+ 72
E W 2 i, V W- w(1]]
< inf 2>
E
(71 + +00)
72
n
With a closer look, one can find on-average model stability (Lei & Ying, 2020) term at RHS:
| (W Z_4, -1- -
1
B ..
[W' E
Z-,, V1- Z= v]]
n
n2
1
R Ez
| w11 Z-,, r1]
1 2
1 n
A(S.V) A(S m1
E
Es,s' v
2n
E2 un average model stability in Delinition 4 of Lei & Ying (2020)
where the second step follows Lemma A.2, A(., denotes the SGD when the random seed e and (())
means replacing the i-th sample of S with the i-th sample from S
Now that we have recovered stability arguments, we can directly borrow stability of SGD to derive excess risk bounds. The following results are based on the on-average model stability bound derived by Lei & Ying (2020).
Proposition 4 Assume the loss is non-negative, convex and B-smooth. Assume the training algo- rithan is projected SGD that starts from fixed Wn = we E and runs T steps with non-increasing T+1 step sizes {m), ) such that π). S 1/28. Then for any y > G, we have the following excess risk
bound:
E [C Cu(w""]]
72 fu(ww))
28
1+ 1 T/n 40y 722)
1/1
122
+2
M141 Σ
1-1
Eniticulu")
Σnn+ nh
n
7/2 -3
T Σ
+ 12 +
(1/2+ 871) || 1 + 2,3
Σ
1141 (u(l(u)))))))
M++1
72 - B
1=0
T=U
where Wace is the accumulated weight
Σ Met Itti W₁
Wraco
Σt=o l1t11
Remark D.8 In separable settings, i.e., when Cu(w*) O, the excess risk bound simplifies
E (Wacc)))
(ww))
+ T/n 4 + 72) (T1 [I]= + 7/2 +
(11/2 11 11))
n 7/2 -
(12 - B) ΣT= Dr+1
After setting 7). = 17 5 1/2 and reparameterizing 722 = k:3 for 1: >>> 1, we have
E. (Wrcc)
(uw"))
[4ck:3'1]
n 42k 1/2 + 81)
T
k: ||w0 llw²112)
1: + 1
Tn
46

---

Published as a conference paper at ICLR 2025
By minimizing over T, we can obtain the following:
E. (Wrece - L.p (u)()))
[Ack35n
2
V4ck32 (1/2 + 8n)
( (k=1 - || |100)13))
72
n
Now set 1/28 to obtain the following:
[ (W -
inf E>1 n
23
(ek (2kk)
(K-1 ||0 O(3 | 1 /n),
which indicates an (n)/n) sample complexity for smooth, convex and separable settings. This result recovers the Theorem 5 and the OL1/n) rate in Lei & Ying (2020) up constants.
Proof This proof is adapted from Appendix C.2 of Lei & Ying (2020). The excess risk can be decomposed into (excess) optimization error and generalization error. The optimization error bound is directly borrowed from Lei & Ying (2020). The generalization error is bounded by combining the recovered stability bound Proposition and the stability of SGD from Lei Ying (2020).
Let w" be the weight that achieves the optimal population loss.
D.6.1 EXCESS OPTIMIZATION ERROR
According to Lemma A.)(c) Lei and 1) I 1/2L and non-increasing, then for any constant w and constant s, one has
Lei & Ying (2020), if the loss is non-negative, convex and B-smooth,
Σ 1 (W+) L. (w) S s] < (1/2 + 3m) ||0| + 23 T=0
Σ 1
T=U
After setting to w" and taking expectation training sets, we have
Σ Mr+ [Ss(W.) -(mi)---ga Ls(w")) s (1/2 + Bm) | ' 11² + 23
Σ th Cy
(D.9)
Since the excess training error bound only given after summing over steps, one has to sum the generalization error bound over steps well.
D.6.2 STABILITY AND GENERALIZATION ERROR
Theorem 3 of Lei & Ying (2020) states that if the loss is non-negative, convex and (-smooth, and SGD has step size T)e S 2//T, I., the for any p > 0 one has
8(1 + 1/2), Σ 11 (1+p/n)
stability
// [s(W))]
where e,
stability is the & on-average model stability at step t:
stability =E V
" n
1
[Il: [Iw: w. [
Let 72 > 0 be a constant. Plugging this stability bound into Proposition 3 leads to gen(H", Pwr1s)
1
233
E (W.))
8 72 stability +
B/72
72
2 62
1
23
E (Wa)
3 + 72 8(1 + 1/p)
12 E [Es(W.))
8/72
72
n
1
233
E [[s(
4(1 + 1Bp)(B 11 p/n)
n E [ss(W))
3/722
72
47

---

Published as a conference paper at ICLR 2025
There empirical losses the trajectory at RHS of the above inequality. They can be bounded by Eq. (A.5) of Lei & Ying (2020), which states given any training set and any constant i E W, one has
:-1 Σ 1E [.(W))
S s S m | + 2
Σ T=0
M 1 C (1))
Setting to wn and taking expectation over training sets, we have
t-1
Σ 1 E [[s(W))] s m 110011 + Σ
We t Lp("))
r=0
Plugging it back leads to
Senlu",P,,tt S 1 8/72
1
23 1 [[s(W))] 401+1/D)C +22) pp/n)'11 72
Th
[| 12 +2 2 1 (p
1-1
n
As in Lei Ying (2020), one can choose p = n/T have (1 + p/n)(-1 S (1 p/n)T-1=
(1 + 1/T)T-1 As a result, we have
gen(n"; Pw.js) S
4(1 T/n) + B2) 2B E [Ws) 3/72
1
m ||| + 1h (n
))
72
n
T=0
To align the generalization error bounds with the weighted summation form of the optimization error, we weight them by 7)+1 and sum over steps the above inequality:
T
7 gen(u", Pwwlss
T
1
2,3
E [ss)
+ 4(1 + T/n)(3 + +72)3e
( 1| +2 ?2 2
<
1
8/12
72
1=0
-=U
2,3
T
7 E [cs(Wa)
+ 4(1 (/n)(B + +2)
Σ
7h2 |||l"=12 + Σ
Rht1 Cn( )
8/72
7/2
11-
t=0
t=0
T=0
To get rid of the empirical losses Et Ne+ (W:).. we apply Eq. (D.9) again and obtain
T
7 gen(u", Pw.1s)
1=0
T
23
(1/2 + Bm)) 1| + 23
Σ
M+ (ww") )+0 + Σ
2
M++1 Cy("")
8/72
7/
1=0
t=0
4(1 + T/n)(3 + 72)33
2-1
Th 1h Lulw"") (ulur))
Σ
7/+1 ?)1 | 2 + Σ
+
11
48

---

Published as a conference paper at ICLR 2025
After obtaining the above generalization error bound, the bound for excess risks can be obtained by summing up it and the optimization error bound Eq. (D.9):
T
7 E (W;) - [u("")]
1=0
T
23
5
(1/2 + 3m1) |1="|1=-20 Σ
1 (u(w") +
2
Σ
nh 1 ku(()))
1 8/72 ?2
t=
t=0
4(1 + TT/)((
32
2-1
Σ Me+1
||00012 112 + Σ
Wh In+=)=(("))
-//92
n
+ +1/22+ + 9m))
(m)]|[""1 ||00|1 112 1/2+ 23 Σ t=0
n M
28 T
1 1+ T/n 3(3 722
T
1-1
Th+1 (p(()))
Σ
1110|112 +2 r=0
2 Σ C Cu(ww)
1h+1
72 8
72 -
t=0
t=
+ 12 + B
(1/2+ B11)
||w0| + 23
Σ 1 ((ti(u((())
72- - B
1=0
The stated inequality can be obtained by dividing Et= 76+11 and applying the Jensen's 's inequality to the convex Ls()).
E DISCUSSION ON THE LIMITATION
The major limitation of our bound is that our bound still relies on population gradients and Hessians. This limitation harms the applicability of our bound to self-certified algorithms (Pérez-Ortiz. et al., 2021).
However, the limitation is not unique jectory technique. The most essential step of this technique is to switch from the original trajectory to the auxiliary trajectory with better properties. However, back to the original trajectory by adding their differences into the bound. In this process, the loss differences are used to measure such differences, resulting in the population loss difference and pop- ulation statistics. As result, previous representative works all have explicit reliance on population statistics. See S" in Propositions and 2. This reliance can be alleviated through some assumptions like (Wr)) < (0.-) (WT + <)] of Wang & Mao (2022). Nevertheless, one still must verify this assumptions on the population set to rigorously apply them, especially when the model is under-fitted or the generalization is bad so that the output weight is far from local minima of the population loss. To sum up, existing representative results based on auxiliary trajectory must rely on population statistics at least implicitly.
bound, but is inherent limitation of the auxiliary tra-
must relate the auxiliary trajectory
In terms of dependence to population/valiidtion statistics, we also optimize the omniscient trajectory using validation statistics, which is heavier dependence. This may forms unfair comparison with existing bounds. However, we have tried to make the comparison fair by allowing the existing bounds to rely on the validation statistics (see Appendix C.5). Even with full access validation sets, the existing bounds cannot exploit them and are still much numerically looser than ours. Lastly, the results of the existing and our bounds be seen not only competitors, but also different trade-offs between the dependence on validation set and bound tightness. Our bound demonstrates how tight bound can be if one allows heavy dependence on validation sets, while previous works show the looseness when one controls the access to validation sets. Future works can start from these two extremes to achieve better trade-off or even break it.
49


Reviews:
[{"review_id": "yJ365BRKY1", "reviewer": "Reviewer_CXNU", "summary": "This paper studies improved information-theoretic generalization bounds for SGD by more effectively incorporating the concept of flatness. Specifically, the authors point out limitations in previous information-theoretic generalization bounds, which fail to align well with some generalization phenomena observed in deep learning, such as the tendency for flatter minima to generalize better. The authors then derive new information-theoretic generalization bounds by constructing an \"omniscient trajectory\", which allows for a better capture of the flatness of SGD solutions in the loss landscape. Empirical results demonstrate that these bounds improve upon previous ones. Additionally, the authors extend their analysis by applying techniques from previous works, such as the individual technique, to mitigate limitations of information-theoretic bounds in stochastic convex optimization (SCO) problems.", "questions": "1. How did you estimate the trace of the Hessian (or other matrices) and the inverse of the matrix in your experiments?\n\n2. In Eq. (1), it seems the expression could potentially be simplified by setting $\\xi\\sim\\mathcal{N}(\\gamma,\\Sigma)$, is this valid?\n\n3. In Eq. (3), what specifically does $H$ represent in $I+H/2\\lambda C$? Also, the notation for $H(W)$ is unclear here; you state that \"$H(w)$ is defined as $\\Delta H(w)$ or $\\hat{H}_S(w)$\", yet in Appendix B.5, it seems that $H(w)$ should correspond to $\\Delta H(w)$. Furthermore, regarding the definition for $C$, does $H(W_T)$ refer to the population Hessian? The notations about Hessian become somewhat confusing due to their inconsistency.\n\n4. In Line 356, the phrase \"such $\\Delta g_t$ is optimal under ...\" appears. Where is the justification for this optimality provided in the text?\n\n4. In the proof of Theorem 3, could you clarify why the CMI term in Eq. (D.4) takes the form $I(W_T+\\sum_{t=1}^T\\Delta g_t^i;Z_i|Z_{-i})$? I may have overlooked something, but it seems that the construction might no longer require the Gaussian noises.\n\n5. Is there potential to extend your analysis to SGD with momentum? Neu et al. (2021) give a method for extending their analysis to SGD with momentum, though they acknowledge limitations in demonstrating any advantage for SGD with momentum over vanilla SGD using their method.\n\n\n**I enjoyed reading this paper, and I would be happy to increase my score if the authors could adequately address my concerns.**", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper introduces novel information-theoretic generalization bounds for SGD that improve upon previous results, particularly when SGD finds flat minima.\n\n2. Constructing the auxiliary weight process twice, specifically through the omniscient trajectory, is a notable technical contribution in this field.\n\n3. The results provide an explanation for the learnability of some SCO problems, where many previous information-theoretic bounds do not apply (e.g., non-vanishing bounds).\n\n4. The paper is clearly written and easy to follow.", "weaknesses": "1. Some related works, whether cited or uncited in this paper, require further discussion in relation to these results. For example, [R2] also discusses limitations in [R1] and suggests that, if the SGD process can be well-approximated by SDEs (motivated by empirical observations), the auxiliary process could be the SDE approximation of SGD, with perturbed Gaussian noise depending on the training data and current state (i.e. depending on $S$ and $W_{t-1}$ at step $t$). Moreover, while not directly related to the generalization of SGD, [R3], which is missing in this paper, investigates the optimal covariance matrix of Gaussian noise in SGLD. Since these works also consider more complex geometric structures for anisotropic Gaussian noise, and emphasize the importance of alignment between flatness and the covariance of weight distributions for generalization, it would be beneficial to discuss whether the results here are comparable or related to them in any sense, such as [R2, Theorem 3.2].\n\n[R1] Ziqiao Wang, and Yongyi Mao. \"On the Generalization of Models Trained with SGD: Information-Theoretic Bounds and Implications.\" International Conference on Learning Representations 2022.\n\n[R2] Ziqiao Wang, and Yongyi Mao. \"Two Facets of SDE Under an Information-Theoretic Lens: Generalization of SGD via Training Trajectories and via Terminal States.\" The Conference on Uncertainty in Artificial Intelligence 2024.\n\n[R3] Bohan Wang, et al. \"Optimizing information-theoretical generalization bound via anisotropic noise of SGLD.\" Advances in Neural Information Processing Systems 2021.\n\n2. In Section 3.4, assuming that $\\nabla\\hat{\\mathcal{L}}_S(W_T)\\approx0$ could limit the applicability of your results in modern deep learning. For example, in some settings, especially for large-scale neural networks, models can perform well without vanishing gradient norms, as noted in [R4]. Can the results here explain generalization in such cases?\n\n[R4] Jingzhao Zhang, et al. \"Neural network weights do not converge to stationary points: An invariant measure perspective.\" International Conference on Machine Learning 2022.\n\n3. In the proof of Theorem 1, particularly in Eq. (B.4), why not let $\\Omega(\\tilde{w}\\_{0:t-1})=\\tilde{w}\\_{t-1}-\\mathbb{E}[g\\_t]-\\mathbb{E}[\\Delta g\\_t]$? This choice seems valid if the expectation is taken over a carefully chosen distribution. Such a choice would result in an additional term, $-\\mathbb{E}[\\Delta G]$ in the trajectory term of Corollary B.1. Could you clarify whether this alternative choice would obtain a tighter bound?\n\n4. In Theorem 3, while I do believe the contribution to information-theoretic generalization bounds here is novel, it seems similar in insight to [R5]: namely, if the learning algorithm is stable, information-theoretic generalization bounds can explain the learnability of these CLB problems. The main distinction in Theorem 3 is that the stability parameter is explicitly expressed in terms of $T$ and $\\eta$, as in Bassily et al. (2020), while [R5] simply uses the stability parameter. In other words, both this work and [R5] explain the generalization of stable algorithms (e.g., GD with smooth or non-smooth convex Lipschitz loss) but may not be able to overcome the limitations highlighted in CLB counterexamples in [R6].\n\n[R5] Ziqiao Wang, and Yongyi Mao. \"Sample-conditioned hypothesis stability sharpens information-theoretic generalization bounds.\" Advances in Neural Information Processing Systems 2023.\n\n[R6] Idan Attias, et al. \"Information Complexity of Stochastic Convex Optimization: Applications to Generalization, Memorization, and Tracing.\" International Conference on Machine Learning 2024.\n\nMinor comments:\n\n1. Some citation errors should be corrected: Wang \\& Mao (2021) was published in ICLR 2022, not NeurIPS 2021. The evaluated variant of CMI bounds is first introduced in Section 6.2 of the arXiv version of Steinke \\& Zakynthinou (2020). Additionally, Hafez-Kolahi et al. (2020) did not use evaluated variants of CMI, so the sentence \"the evaluated variants (Hafez-Kolahi et al., 2020; Harutyunyan et al., 2021; Wang \\& Mao, 2023b)\" is incorrect. A missing reference here is [R7], which applies individual e-CMI bounds to obtain new generalization bounds. Accordingly, the correct sentence should be \"the evaluated variants (Harutyunyan et al., 2021; Hellström & Durisi, 2022; Wang & Mao, 2023b)\" to reflect $f$-CMI, e-CMI and ld-CMI.\n\n[R7] Fredrik Hellström, and Giuseppe Durisi. \"A new family of generalization bounds using samplewise evaluated CMI.\" Advances in Neural Information Processing Systems 2022.\n\n2. The Hessian notation is inconsistent: in Section 2.2, $\\hat{H}\\_w$ and $H\\_w$ are used to denote the Hessians at $w$, but in Line 218, $\\hat{H}\\_s$ is used, and in Theorem 2, $\\hat{H}\\_s(w)$ and $H\\_\\mu(w)$ are used. Please clarify these notations for consistency.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C5kBLeE9ar", "reviewer": "Reviewer_6UMN", "summary": "This paper introduces a new approach to developing Information-theoretic Generalization Bounds for Stochastic Gradient Descent (SGD). The core idea is to incorporate an \"omniscient\" auxiliary trajectory into the proof, which better leverages the flatness of the loss landscape. Previous information-theoretic bounds struggled to align with generalization performance as batch size varied. This paper addresses that limitation, offering bounds that correlate more accurately with generalization across different batch sizes. Additionally, prior bounds for Gradient Descent (GD) on convex-Lipschitz-Bounded (CLB) problems could not achieve the optimal minimax rate of $O(\\frac{1}{\\sqrt{n}})$, but this paper’s approach successfully reaches this rate.", "questions": "1) Do you think it would be possible to simplify the approach by using only one auxiliary trajectory (the omniscient one) in the proof?\n2) Can you provide more clarifications about the point 2) in the weaknesses above?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "-The contribution is novel and significant: The \"omniscient\" auxiliary trajectory is original and allows to overcome 2 limitations of previous information-theoretic generalization bounds: \n1) correlation between the bounds and observed generalization behavior under different batch sizes\n2) $O(\\frac{1}{\\sqrt{n}})$ rate for Gradient Descent on convex-Lipschitz-Bounded problems.\n\n-The paper effectively motivates the problem and provides a useful high-level overview of the main ideas in the introduction.\n\n-The paper includes rigorous proofs for the results, along with experiments that support the theory. However, I have not verified the mathematical details of the proofs.", "weaknesses": "1) The bound requires an additional penalty term to be introduced. This additional term could  worsen the bound in some applications if the benefits to the trajectory term are not high enough. Also, the proof involves two auxiliary trajectories (omniscient and SGLD-like) instead of one, complicating the proof and the bound.\n\n2) In the introduction, it is postulated that the incorrect behavior of the previous bound in regards to varying batch sizes comes from the trajectory term not capturing the flatness more directly. However, in the experiments, the tendency of the new trajectory term w.r.t. learning rate is still incorrect. The overall bound seems to correlate correctly with generalization performance because the additional penalty term behave correctly and because this additional penalty term dominates the bound. It is therefore not clear to me if the paper provides the correct explanation in the introduction about what is ''really'' solving the problem.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cpPKQIKQm5", "reviewer": "Reviewer_Bn1c", "summary": "This work belongs to the line of information-theoretic generalization bound and improves upon the results of [Neu et al., 2021] and [Wang and Mao, 2023]. In particular, it utilizes geometry-aware perturbation and establishes a novel generalization upper bound that better captures the influence of flatness-related terms and outperforms [Wang and Mao, 2023] by more accurately predicting the empirical observations. Moreover, the authors devise a more powerful auxiliary sequence called \"omniscient trajectory\" to bound the generalization error of the trajectory of (stochastic) gradien-based methods. As an application of the novel result, the authors derive the optimal minimax rate for GD under the CLB setting.", "questions": "The authors apply their result to CLB problem and obtain the optimal rate. I would be particularly interested if the author could show how their bound will imply a sample rate when the stochastic function is smooth and realiable, which is provably to have a better rate like O(1/n^{2/3}) or O(1/n) than the CLB setting from the recent upper or lower bound result. This might be unrelated to the major topic but I would greatly appreciate if the authors could answer this.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Overall, this is a solid and well-written paper. The motivation and intuition behind the proposed methodology are well-explained. It provides a refinement for the information-theoretic generalization bound by exploiting the geometric properties of the objective function and by devising a more powerful auxiliary sequence. This allows to better reflect empirical pbeservations and recovers the minimax rate when applying to CLB-based bounds. In this regard, this paper takes a step further towards the understanding of how local geometry might affect the generalization.", "weaknesses": "Despite a rather interesting paper, some concerns arise across different aspects.\n\nFirst, the major concern comes from the novelty. The paper should be regarded as a refinement over [Wang and Mao, 2021] and the generalization bound is not new and decisively better compared to the prior results. Also, the methodology is also not completely original. For example, the use of anisotropic Gaussian perturbation also appears in previous works like [Neu et al., 2021], as mentioned by the authors themselves.\n\nSecond, it seems that the applicability of this result slightly suffers from the form of the iterative update in Sec 2.1 and ingredients like the population Hessian. For instance, the iterative update (it would be great if the authors provided a label for the update) might rule out the use of past gradient information and hence exclude many algorithms. It would be great if the author could provide some intuition to show that these ingredients are indeed necessary.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TnazTkE8Bc", "reviewer": "Reviewer_aqMJ", "summary": "This work introduces a novel information theoretic bound. As is common for bounds of this type, the starting point is a well-known result that states that the generalization error can be bounded in term of the mutual information (MI) between the training set $S$ and the learning algorithm’s output $W$. For SGD algorithms, under Gaussian weight priors and Gaussian noise injection on the final iterate, this boils down to a generalization bound that consists of a flatness term and a trajectory term (see e.g. Neu et al. and Wang et al.). \n\nSome existing MI bounds of this form have a crucial defect: they increase in the batch size even though small batch SGD generally converges to flatter minima due to larger gradient noise. This defect is purported to be caused by the hidden flatness-dependence of the trajectory term in those MI bounds, though this wrong tendency can be addressed with i.i.d. per-iterate auxiliary noise injections (see Neu et al.). \n\nThe main contribution of this paper is to improve upon existing MI bounds such as Neu et al. and Wang and Mao via a novel technique. More specifically, the authors of this work conduct a more fine grained analysis wherein, rather than choosing the optimal noise covariance based on the expected trajectory $W$, one chooses the best noise covariance for the specific instances of $W$, and then calculates the expected error bound over $W$. By doing so, one can a) sidestep an overfitting issue with estimating the best noise covariance based on (expensive/scarce samples of) $W$; and b) get a numerically tighter bound by virtue of exchanging the $\\inf\\mathbb E$ with the $\\mathbb E \\inf$ operator.\n\nThis approach, based on the construction of an \"omniscient trajectory\", is general enough so that it can be applied to other existing methods. Moreover, it is shown to be empirically tighter than existing, comparable bounds, and it is shown to be theoretically optimal (up to constants) for SGD on Convex-Lipschitz-Bounded (CLB) problems for a specific choice of training iterations and step size.", "questions": "- There exist PAC-Bayes based generalization bounds that are non-vacuous even on ImageNet, while yours seems to be vacuous already for CIFAR10. How do these bounds compare conceptually to your type of bound? Perhaps they have in some sense an “unfair advantage”?\n- What is the general tendency of your bound to changes in depth, width, training set size, and scaling of the full weight vector? If the tendency is in the wrong direction, is this shared by other bounds based on MI? I am asking because the bound being numerically tighter does not necessarily rule out the fact that it now scales in the wrong directions for some parameters.\n- I suggest that you already introduce some of the notation around the equation in line 82. For example, it might not be clear to some readers unfamiliar with these results that $g_t$ denotes the gradients.\n- Could you please expand on the difficulties of evaluating the bound in practise, specifically on which estimation steps you consider the most loose ones? How does the hardness/accuracy of estimation compare to that of the bounds of Neu et al. and Wang and Mao, for example in terms of how the training set has to be partitioned to carry out all the required estimations in a sound manner, and the trade-offs induced by the specific choices of such partitions? \n- I would suggest that you include some of the most important aspects of empirical evaluation in the main part of the paper, since this might be considered more important to some readers than some of the more technical/theoretical details.\n- Related to the above question: if there are indeed particularly loose estimation steps, do you perhaps have some high level ideas on how to address these inadequacies in future work?\n- I see in Appendix C.1 that for your experiments, you evaluate the penalty terms over the test set, which should make the bound incomparable to some of the other competing generalization bounds which work only on the training set. If you disagree with this statement, could you please elaborate? In my view, a fair comparison would require that you set aside a part of the training set for such estimation procedures.\n\n============================================================================\n\nTypos/minor comments:\n- Line 34: perhaps \"that\" instead of \"because\"\n- Line 932 \"Hessians\" -> \"Hessian\"\n\n=============================================================================\n\nReferences:\n\n- Gergely Neu, Gintare Karolina Dziugaite, Mahdi Haghifam, and Daniel M. Roy. Information- theoretic generalization bounds for Stochastic Gradient Descent. In Proceedings of Thirty Fourth Conference on Learning Theory, pp. 3526–3545. PMLR, July 2021. ISSN: 2640-3498.\n- Ziqiao Wang and Yongyi Mao. On the generalization of models trained with SGD: Information- theoretic bounds and implications. In Advances in Neural Information Processing Systems, October 2021.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed approach is sound and very effective in achieving the goal of better modelling the effect of flatness on generalization. For example, it achieves a decrease of the trajectory term on the order of $10^6$ (multiplicative) in the case of CIFAR-10 over the one proposed by Neu et al.. This is a valuable contribution.\n- The paper is overall well-written.\n- Like the bound of Neu et al., the proposed bound has the right scaling w.r.t. the batch size, while also being numerically tighter.\n- In fact, the bound is numerically tight in a very simple setting (MNIST + 2-layer MLP). \n- The omniscient-trajectory technique is to a degree general purpose and can be applied to a whole class of similar bounds.", "weaknesses": "- While the bounds is quite tight for MNIST, it seems that it is numerically vacuous (i.e. larger than 1 for the 0-1 classification accuracy) already for CIFAR10. So while it is numerically tighter than comparable bounds, it does not provide a major qualitative improvement over bounds such as the one by Neu et al.. \n- In my view, even if one lets go of requiring tightness of a generalization bound, the next most important goal is to have the right tendency/correlation w.r.t. to the parameters of the learning setup. However, it is my impression that this bound does not “fix” any of such qualitative defects of existing comparable bounds, whereas, for example Neu et al. \"fixes\" the scaling w.r.t. batch-size and other flatness bounds improve upon prior ones by enforcing invariance w.r.t. rescaling the parameter vector.\n- Empirically, applying Theorem 2 is difficult and requires multiple assumptions/heuristics, e.g.:\n    - Crucially, multiple population means have to be estimated from i.i.d. samples of $W$, which is difficult. Given a training set of fixed size $N$, one needs to partition it by averaging of $k$ subsets of size $N/k$, yielding a bound for SGD on training sets of size $N/k$. If one heuristically were to apply that obtained bound to the output of SGD over the full training set $S$, it must be quite loose unless $S$ is highly redundant, which is typically not the case.\n    - The bound has a hyperparameter $\\lambda$ that regulates a trade-off between different terms in the bound and at the same time must upper bound the spectral norm of the Hessians\n    - The bound assumes exactly zero training error at the final output. Thus it does not apply in the case of early stopping, and the application to approximately zero training error outputs requires a small leap of faith.\n    - In the case of unbounded losses, one must clip the loss function in order to guarantee the sub-Gaussian concentration required for the analysis.\n- The result gives a bound on the expected generalization error. This is qualitatively weaker than a high probability upper bound. Arguably, this is a rather minor weakness since deep learning models tend to be exceptionally stable w.r.t. different initializations and draws of the SGD batches.\n- These kinds of bounds apply to the generalization error in terms of the expected surrogate loss, rather than the $0-1$ which is the real measure of interest in classification. I acknowledge that this is essentially unavoidable here in order to link generalization to the optimization trajectories, which are obtained from minimization of the surrogate loss. However, it should be noted that this makes the analysis inherently more loose since the surrogate risk is always an upper bound on the $0-1$ risk (i.e. the test accuracy).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)