<figure><img src="image_1.png" /></figure>

HAL
open science

# A PAC-Bayes bound for deterministic classifiers

Eugenio Clerico, George Deligiannidis, Benjamin Guedj, Arnaud Doucet

► To cite this version:

Eugenio Clerico, George Deligiannidis, Benjamin Guedj, Arnaud Doucet. A PAC-Bayes bound for deterministic classifiers. 2022. hal-03815146

HAL Id: hal-03815146

https://inria.hal.science/hal-03815146

Preprint submitted on 14 Oct 2022

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.
# Generalisation under gradient descent via deterministic PAC-Bayes

Eugenio Clerico, Tyler Farghly, George Deligiannidis, Benjamin Guedj, Arnaud Doucet

► To cite this version:

Eugenio Clerico, Tyler Farghly, George Deligiannidis, Benjamin Guedj, Arnaud Doucet. Generalisation under gradient descent via deterministic PAC-Bayes. ALT 2025 - Association for Learning Theory, Feb 2025, Milan, Italy. hal-03815146v2

HAL Id: hal-03815146

https://inria.hal.science/hal-03815146v2

Submitted on 6 Jul 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# A PAC-Bayes bound for deterministic classifiers

Eugenio Clerico*¹, George Deligiannidis¹, Benjamin Guedj², ³, and Arnaud Doucet¹

¹Department of Statistics, University of Oxford

²Centre for Artificial Intelligence, University College London

³Inria London

## Abstract

We establish a disintegrated PAC-Bayesian bound, for classifiers that are trained via continuous-time (non-stochastic) gradient descent. Contrarily to what is standard in the PAC-Bayesian setting, our result applies to a training algorithm that is deterministic, conditioned on a random initialisation, without requiring any *de-randomisation* step. We provide a broad discussion of the main features of the bound that we propose, and we study analytically and empirically its behaviour on linear models, finding promising results.

## 1 Introduction

Effectively upper-bounding the generalisation error of modern learning algorithms is one of the main open challenges for the statistical learning theory community (Zhang et al., 2016). Originally, properties of the hypothesis space, such as VC dimension and Rademacher complexity (Bousquet et al., 2004; Vapnik, 2000; Shalev-Shwartz and Ben-David, 2014), were used to establish worst-case generalisation bounds, holding uniformly over all the possible algorithms and datasets. However, these results are often vacuous when applied to the current over-parameterised models, the modern perspective focuses on algorithm-and data-dependent bounds (McAllester, 1998; Bousquet and Elisseeff, 2002; Dwork and Roth, 2014; Hardt et al., 2016; Xu and Raginsky, 2017; Clerico et al., 2022b; Lugosi and Neu, 2022). Among the various approaches proposed in the literature, the PAC-Bayesian framework (Guedj, 2019; Alquier, 2021) has obtained promising empirical results for deep networks (Dziugaite and Roy, 2017; Zhou et al., 2019; Pérez-Ortiz et al., 2021; Clerico et al., 2022a).

Typically, the PAC-Bayesian bounds are upper bounds on the expected population loss of stochastic classifiers, holding with high probability on the random draw of the training dataset. This setting, where the model's parameters are required to be intrinsically random, does not match with the more standard scenario of a non-randomised predictor. Although most available empirical PAC-Bayesian studies on neural networks focus on specifically designed stochastic architectures (Dziugaite and Roy, 2017; Zhou et al., 2019; Pérez-Ortiz et al., 2021; Clerico et al., 2022a), there are also *de-randomisation* methods that allow to apply PAC-Bayesian ideas to deterministic classifiers. For instance, one can approximate a deterministic model with a noisy version, whose parameters are stochastically perturbed. Under suitable stability assumptions, a PAC-Bayesian analysis of the stochastic model can bring a generalisation bound for the deterministic one (Neyshabur et al., 2018; Nagarajan and Kolter, 2019; Miyaguchi, 2019). Alternatively, one can average a stochastic algorithm to get a deterministic one. Then, PAC-Bayesian bounds on the loss of this averaged predictor can be obtained by Jensen's inequality (Germain et al., 2009; Letarte et al., 2019; Biggs and Guedj, 2021) (if the loss function is convex), or through margin-based arguments (Langford and Shawe-Taylor, 2002; Banerjee et al., 2020; Biggs and Guedj, 2022). Finally, besides the PAC-Bayesian bounds in expectation, there are disintegrated bounds that hold with high probability with respect to a random realisation of the stochastic algorithm. For instance, in the case of a stochastic network, this means that we have a generalisation bound for the deterministic classifier obtained by drawing all the parameters once (Catoni, 2007; Blanchard and Fleuret, 2007; Rivasplata et al., 2020; Viallard et al., 2021).

<footnote>*Correspondence: clerico@stats.ox.ac.uk</footnote>
# Generalisation under gradient descent via deterministic PAC-Bayes

Eugenio Clerico*†

Universitat Pompeu Fabra, Barcelona

Tyler Farghly*

Department of Statistics, University of Oxford

George Deligiannidis

Department of Statistics, University of Oxford

Benjamin Guedj

AI Centre and Department of Computer Science, UCL London & Inria London

Arnaud Doucet

Department of Statistics, University of Oxford

EUGENIO.CLERICO@GMAIL.COM

FARGHLY@STATS.OX.AC.UK

DELIGIAN@STATS.OX.AC.UK

B.GUEDJ@UCL.AC.UK

DOUCET@STATS.OX.AC.UK

**Editors:** Gautam Kamath and Po-Ling Loh

## Abstract

We establish disintegrated PAC-Bayesian generalisation bounds for models trained with gradient descent methods or continuous gradient flows. Contrary to standard practice in the PAC-Bayesian setting, our result applies to optimisation algorithms that are deterministic, without requiring any *de-randomisation* step. Our bounds are fully computable, depending on the density of the initial distribution and the Hessian of the training objective over the trajectory. We show that our framework can be applied to a variety of iterative optimisation algorithms, including stochastic gradient descent (SGD), momentum-based schemes, and damped Hamiltonian dynamics.

## 1. Introduction

Effectively upper bounding the generalisation error of modern learning algorithms is an open problem of great importance to the statistical learning theory community (Zhang et al., 2016). Originally, properties of the hypothesis space, such as VC dimension and Rademacher complexity (Vapnik, 2000; Bousquet et al., 2004; Shalev-Shwartz and Ben-David, 2014), were used to establish worst-case generalisation bounds, holding uniformly over all possible algorithms and training datasets. However, as these results are often vacuous in over-parameterised settings, the modern perspective focuses on algorithm and data-dependent bounds (McAllester, 1998; Bousquet and Elisseeff, 2002; Hardt et al., 2016; Xu and Raginsky, 2017; Clerico et al., 2022b; Lugosi and Neu, 2022).

Among the various approaches, the PAC-Bayesian framework (Guedj, 2019; Alquier, 2021) has obtained particularly promising empirical results (Dziugaite and Roy, 2017; Zhou et al., 2019; Pérez-Ortiz et al., 2021a,b; Biggs and Guedj, 2022b; Clerico et al., 2022a). Typically, a PAC-Bayes bound is an upper bound on the expected population loss of a stochastic algorithm, holding with high probability on the random draw of the training dataset. This framework gained popularity after yielding non-vacuous empirical bounds in overparameterised regimes, such as modern neural networks (Dziugaite and Roy, 2017; Zhou et al., 2019; Pérez-Ortiz et al., 2021a; Clerico et al., 2022a). Since the standard PAC-Bayesian framework relies on the randomness of the trainable parameters,

<footnote>* The first two authors contributed equally.</footnote>

<footnote>† This work was carried out while the author was affiliated with the Department of Statistics, University of Oxford.</footnote>
When trying to apply the PAC-Bayesian framework to study a given non-stochastic model, the first of the methods proposed above is the most natural. However, in the present work, we show that for this purpose it is also possible to leverage the disintegrated PAC-Bayesian bounds. Our starting point is noticing that even training a deterministic classifier does usually involve some randomness: the standard procedure when training a neural network almost always starts with the random draw of the initial weights from some known simple distribution (Goodfellow et al., 2016). For a model trained via continuous-time (non-stochastic) gradient descent, we show that it is possible to exploit this source of stochasticity and obtain a disintegrated PAC-Bayes bound that holds with high probability on the random training dataset and initialisation (Proposition 2). Evaluating this bound only requires to be able to evaluate the density for the initial and final configurations of the network, and to compute the integral along the trajectory of the training objective's Laplacian. Although the bound necessitates the ideal situation of continuous training dynamics, to our knowledge this is the first PAC-Bayesian result that directly apply to a standard non-stochastic neural networks with a fully deterministic training, without strong assumptions on the architecture and the need of the introduction of any auxiliary stochastic model. Besides, we find our bound remarkably easy-to-prove, clean, and simple, with no strong hypotheses made on the training objective, which is only required to be twice differentiable.

The paper is organised as follows. First, we set up the learning framework and introduce the notations in Section 2. Section 3 gives a quick overview on the disintegrated PAC-Bayesian bounds, while in Section 4 we establish our main result (Proposition 2). Section 5 discusses the main features of our bound, and finally Section 6 presents some analytical and empirical quantitative results.

## 2 Notation and setting

In a standard supervised learning framework, we consider pairs $(x, y) \in \mathcal{X} \times \mathcal{Y}$, where $x$ denotes an example and $y$ its label. For simplicity, we can imagine that there is a mapping $f^* : \mathcal{X} \to \mathcal{Y}$ that associates to each $x$ a unique correct label $y = f^*(x)$. A learning algorithm takes a labelled training dataset $s = \{x_1...x_m\}$, made of $m$ inputs whose labels $\{f^*(x_1)...f^*(x_m)\}$ are known, and outputs a function $f : \mathcal{X} \to \mathcal{Y}$, which ideally is a good approximation of $f^*$. More generally, the algorithm will return an hypothesis $h \in \mathcal{H}$ (e.g., values of a neural network's parameters), which is understood to parameterise a function $f_h : \mathcal{X} \to \mathcal{Y}$. For simplicity, we will always assume that $\mathcal{H} \subseteq \mathbb{R}^N$, for some dimension $N > 0$. We call the algorithm stochastic when its output $h$ is a random variable on $\mathcal{H}$, whose law can depend on $s$.

The performance of each hypothesis can be assessed via a loss function $\ell$, which takes a pair of labels $(y, y')$ and returns a real value, representing how “far apart” $y$ and $y'$ are. In practice, for any hypothesis $h$, one can always compute the empirical loss on the training dataset

$$ \mathcal{L}_s(h) = \frac{1}{m} \sum_{x \in s} \ell(f_h(x), f^*(x)). $$

However, we are essentially interested in how well $h$ can predict the label of an $x$ which does not belong to $s$. Assuming that the population of examples follows a distribution $\mu$, the actual relevant quantity is the population loss

$$ \mathcal{L}_{\mathcal{X}}(h) = \int_{\mathcal{X}} \ell(f_h(x), f^*(x)) d\mu(x). $$

Upper-bounding $\mathcal{L}_{\mathcal{X}}$, having access to $\mathcal{L}_s$ only, is usually referred to as the generalisation problem. Typically, the training dataset is made of $m$ i.i.d. samples from $\mu$ (i.e., $s \sim \mu^m = \mu^{\otimes m}$). We are interested in upper bounds on $\mathcal{L}_{\mathcal{X}}(h)$ (where $h$ is the hypothesis picked by the algorithm) that holds with high probability on the random draw of $s$ (or of $(h, s)$ in the stochastic setting).

We will consider an algorithm whose output is obtained by optimising an objective $C_s : \mathcal{H} \to \mathbb{R}$ via continuous-time gradient descent. Clearly, $C_s$ depends on the training dataset $s$ and, in practice, it can coincide with the empirical loss, although this is not necessarily the case, as one might use a surrogate loss for the training or add some regularisation term. We require that the objective is twice differentiable, and we define the gradient flow $(h_0, t) \mapsto \Phi_t(h_0)$ as the solution of the dynamics

$$ \partial_t \Phi_t(h_0) = -\nabla C_s(\Phi_t(h_0)); \qquad \Phi_0(h_0) = h_0. $$
this type of analysis is typically applied to specifically designed stochastic models. For instance,
in the setting of neural networks, this requires an architecture featuring stochastic weights and bi-
ases, instead of the standard deterministic ones. To extend these ideas to deterministic settings,
de-randomisation techniques are used. One possibility is to leverage stability properties to approx-
imate a model by randomly perturbing its parameters. While this approach has shown promis-
ing results for feed-forward neural networks (Neyshabur et al., 2018; Nagarajan and Kolter, 2019;
Miyaguchi, 2019; Banerjee et al., 2020), it relies on specific architectural assumptions. An alterna-
tive way to tackle the problem provides bounds for the predictor obtained by averaging a stochastic
one, an approach started by Germain et al. (2009). However, this leads to results that only apply
to models with very specific structures: for instance, Letarte et al. (2019), Biggs and Guedj (2021),
and Biggs and Guedj (2022a) obtained bounds for particular deterministic networks with a rather
unusual erf activation function. Finally, besides PAC-Bayes bounds in expectation, there are disinte-
grated results that hold with high probability on a random realisation of the stochastic model (Catoni,
2004, 2007; Blanchard and Fleuret, 2007; Alquier and Biau, 2013; Guedj and Alquier, 2013; Rivasplata et al.,
2020; Viallard et al., 2021). To the best of our knowledge, this last approach has not been applied to
standard non-stochastic algorithms, such as neural networks trained via gradient descent methods.

Here, we consider models trained by gradient descent-type methods and leverage the framework of disintegrated PAC-Bayes bounds. We start by noticing that often training with a deterministic optimisation scheme does still involve some randomness, due to an initialisation that features a random draw of the initial values of the parameter (e.g., most neural networks (Goodfellow et al., 2016)). Our analysis shows that it is possible to exploit this source of noise and obtain disintegrated PAC-Bayes bounds, holding with high probability on the random training dataset and initialisation. To the best of our knowledge, this is the first PAC-Bayesian result that directly applies to standard non-stochastic settings, without strong requirements on the model or the need for any randomness other than the initialisation. Besides, unlike bounds based on de-randomisation, ours apply with only limited assumptions made about the smoothness of the training objective, and can be computed in closed form using information collected along the trajectory of the parameters during training.

We compare our bounds with other known results, including some outside the scope of the
PAC-Bayesian literature. When compared to uniform stability bounds (Elisseeff, 2005; Hardt et al.,
2016; Bousquet et al., 2020), we find that ours have sharper rates with respect to the size of the
training dataset, and grow slower with the number of iterations. With regards to the recently popu-
larised information-theoretic bounds (Xu and Raginsky, 2017; Negrea et al., 2019; Neu et al., 2021;
Clerico et al., 2022b), ours are noticeably easier to compute and are not limited to bounds in ex-
pectation. Evaluating our bound only requires knowledge of the density of the initial distribution
and of the Hessian of the training objective over the optimisation trajectory. The latter captures
the flatness of the optimisation objective along the training path and can be seen to agree with
the notion that flatter minima generalise better (Hochreiter and Schmidhuber, 1997; Keskar et al.,
2017; Izmailov et al., 2018; He et al., 2019; Neu et al., 2021). We also highlight here how this
term relates to the implicit regularisation occurring in algorithms known to result in improved
generalisation (Blanc et al., 2020; Damian et al., 2021). We demonstrate that this framework is
easily extended to almost all iterative schemes, including stochastic variants of gradient descent
(Kiefer and Wolfowitz, 1952), and iterative procedures based on auxiliary variables, like momentum
schemes (Qian, 1999) or damped Hamiltonian dynamics (Hairer et al., 2006; França et al., 2020).
As $h_0$ is fixed before starting the training, we will often omit the explicit dependence on it, and simply write the solution of the above dynamics as $h_t$. Fixed a time horizon $T > 0$, our algorithm’s output is given by $h_T$.

We will focus on the case of a random initialisation, with law $\rho_0$. We define the distribution $\rho_t$ as the push-forward of $\rho_0$ under the gradient flow, that is

$$ \rho_t = \Phi_t^{\#} \rho_0. $$

In this way, if we draw $h_0$ from $\rho_0$, then we automatically get that $h_t \sim \rho_t$. Note that $h_t$ is non-random when conditioned on the initialisation $h_0$. In particular, although the algorithm is in fact stochastic, we can see it as a deterministic model that has only a random initialisation. This is actually the common setting for modern neural networks, whose initial configuration is usually drawn from some simple distribution (e.g., He et al. (2015)).

## 3 Disintegrated PAC-Bayes bounds

The main focus in the development of the PAC-Bayesian theory has been on bounds in expectation, and only recently Rivasplata et al. (2020) and Viallard et al. (2021) have brought back interest in the disintegrated bounds, which actually date back to the works of Catoni (2007) and Blanchard and Fleuret (2007). We refer to Alquier (2021) for a detailed introductory exposition of the PAC-Bayesian framework that also discusses some disintegrated results.

We consider a stochastic model, which takes a sample $s \sim \mu^m$ and returns a random hypothesis $h \sim \rho^s$, where the superscript $s$ that we use in this section stresses explicitly that $\rho$ depends on $s$.<fnref n="1" /> We denote the joint law of $(s, h)$ as $\mu^m * \rho^s$, that is $d(\mu^m * \rho^s)(s, h) = d\mu^m(s)d\rho^s(h)$. As it is standard in the PAC-Bayesian literature, we will call $\rho^s$ the posterior distribution. We denote as $\pi$ a prior distribution on $\mathcal{H}$, whose only requirement is to be data-agnostic, in the sense that it cannot depend on the specific dataset $s$ used for the training. We write $\mu^m \otimes \pi$ for the law of a pair $(s, h)$, where $s \sim \mu^m$ and $h \sim \pi$ are independent.

A disintegrated PAC-Bayesian bound is an upper bound on $\mathcal{L}_\mathcal{X}(h)$, where $h$ is the output of a stochastic algorithm, that holds with high probability over $(s, h) \sim \mu^m * \rho^s$. The next result, from Rivasplata et al. (2020), allows to derive several of these bounds. We provide a short proof for completeness.

**Proposition 1.** Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be an arbitrary measurable function, and define $\xi$ as the expectation of $e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h))}$ under $(s, h) \sim \mu^m \otimes \pi$. For any $\delta \in (0, 1)$, for any posterior distribution $\rho^s$ absolutely continuous with respect to the prior $\pi$, with probability at least $1 - \delta$ on $(s, h) \sim \mu^m * \rho^s$ we have that

$$ \Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h)) \le \log \frac{d\rho^s}{d\pi}(h) + \log \frac{\xi}{\delta}. $$

*Proof.* By Markov’s inequality, for $(s, h) \sim \mu^m * \rho^s$ we have that with probability at least $1 - \delta$

$$ e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h)) - \log \frac{d\rho^s}{d\pi}(h)} \le \frac{1}{\delta} \int_{\mathcal{X}^m \times \mathcal{H}} e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h)) - \log \frac{d\rho^s}{d\pi}(h)} d\mu^m(s) d\rho^s(h). $$

By noticing that for all $s \in \mathcal{X}^m$

$$ \int_{\mathcal{H}} e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h)) - \log \frac{d\rho^s}{d\pi}(h)} d\rho^s(h) = \int_{\mathcal{H}} \frac{e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h))}}{\frac{d\rho^s}{d\pi}(h)} \chi_{\frac{d\rho^s}{d\pi}(h)>0} d\rho^s(h) = \int_{\mathcal{H}} e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_\mathcal{X}(h))} d\pi(h), $$

we easily get the desired result. □

Choosing $\Psi(u, v) = 2m(u - v)^2$ one finds the disintegrated version of a standard PAC-Bayes bound from McAllester (1999). Alternatively, defining $kl(u||v) = u \log \frac{u}{v} + (1 - u) \log \frac{1-u}{1-v}$ (the relative entropy between two binary Bernoulli distributions) and setting $\Psi(u, v) = mkl(u||v)$, we get a tighter bound, which is the disintegrated counterpart of a bound due to Langford and Seeger (2001) and Maurer (2004). These results follow from the fact that in both cases $\xi$ can be upper-bounded by $2\sqrt{m}$, if the loss $\ell$ is bounded in $[0, 1]$ (Bégin et al., 2016).

<footnote n="1">¹To be rigorous, one should actually require that $s \mapsto \rho^s$ is a Markov kernel.</footnote>
## 2. Notation and setting

We consider the standard supervised learning framework, where examples are pairs instance-label $z = (x, y) \in \mathcal{X} \times \mathcal{Y} = \mathcal{Z}$. A learning algorithm takes a training dataset $s = \{z_1, \dots, z_m\}$ of $m$ examples and outputs a function $f: \mathcal{X} \to \mathcal{Y}$. More specifically, we consider algorithms that choose a hypothesis $h \in \mathcal{H}$, which is understood to parameterise a map $f_h: \mathcal{X} \to \mathcal{Y}$ (e.g., $h$ could be the weights of a neural network). We always assume that $\mathcal{H} \subseteq \mathbb{R}^d$, for some dimension $d > 0$. We call the algorithm stochastic when its output $h$ is a random variable on $\mathcal{H}$, whose law can depend on $s$.

Given a loss function $\ell: \mathcal{H} \times \mathcal{Z} \to \mathbb{R}$, we define the empirical loss on a dataset $s$ by

$$ \mathcal{L}_s(h) = \frac{1}{m} \sum_{z \in s} \ell(h, z). $$

However, what often matters is how well $h$ predicts the labels of instances outside of $s$. Assuming that the population of examples follows a distribution $\mu$, the relevant quantity is the population loss,

$$ \mathcal{L}_{\mathcal{Z}}(h) = \int_{\mathcal{Z}} \ell(h, z) d\mu(z). $$

Upper bounding $\mathcal{L}_{\mathcal{Z}}$ only knowing $\mathcal{L}_s$ is the subject of focus in this paper and in the literature on generalisation bounds more broadly. We assume that $s \sim \mu^m = \mu^{\otimes m}$ (i.e., i.i.d. draws from $\mu$). We are interested in upper bounds on $\mathcal{L}_{\mathcal{Z}}(h)$ (with $h$ the hypothesis picked by the algorithm) holding with high probability on the random draw of $s$ (or on $(h, s)$ in the stochastic setting).

Our results are inspired and naturally find their place within the PAC-Bayesian framework. Although the main focus in the PAC-Bayes literature has been on bounds in expectation, Rivasplata et al. (2020) and Viallard et al. (2021) have recently brought back interest in disintegrated bounds, which actually date back to Catoni (2004, 2007) and Blanchard and Fleuret (2007). We refer to Alquier (2021) for an introductory exposition on PAC-Bayes that also discusses a few disintegrated results. PAC-Bayes bounds deal with a stochastic model that, given $s \sim \mu^m$, returns a random hypothesis $h \sim \rho^s$, where the superscript $s$ stresses $\rho$'s dependence on $s$.<fnref n="1" /> We call $\rho^s$ the posterior distribution and denote the joint law of $(s, h)$ as $\mu^m * \rho^s$, i.e., $d(\mu^m * \rho^s)(s, h) = d\mu^m(s)d\rho^s(h)$. A disintegrated PAC-Bayes bound is an upper bound on $\mathcal{L}_{\mathcal{Z}}(h)$ that holds with high probability over $(s, h) \sim \mu^m * \rho^s$. A fundamental ingredient in this framework is the comparison of the posterior $\rho^s$ with a prior distribution $\pi$ on $\mathcal{H}$, only required to be chosen without knowledge of the training dataset $s$. We write $\mu^m \otimes \pi$ for the law of a pair $(s, h)$, where $s \sim \mu^m$ and $h \sim \pi$ are independent.

This work provides generalisation bounds for algorithms whose output is obtained optimising an objective $\mathcal{C}_s: \mathcal{H} \to \mathbb{R}$ via gradient-based descent methods. $\mathcal{C}_s$ can depend on the training dataset $s$ and, in practice, it can coincide with the empirical loss. However, this is not necessarily the case, as one might use a surrogate loss for the training or add some regularising term. In our analysis, we use $h_t$ (or $h_k$) to denote the parameters at time $t$ (or iteration $k$). Similarly, we use $\rho_t$ (or $\rho_k$) to denote their marginal distribution. All the measures that we consider are absolutely continuous with respect to the Lebesgue measure, and we use the same notation to denote their density. The random initialisation is given by $\rho_0$, that we assume to have strictly positive density on the whole $\mathcal{H}$.

<footnote n="1">1. To be rigorous, one should actually require that $s \mapsto \rho^s$ is a Markov kernel.</footnote>
**Corollary 1.** Assume that the loss $\ell$ is bounded in $[0, 1]$. Then, for any $\delta \in (0, 1)$, for any posterior distribution $\rho^s$ absolutely continuous with respect to the prior $\pi$, each of the following bounds hold with probability at least $1 - \delta$ on $(s, h) \sim \mu^m * \rho^s$

$$ \mathcal{L}_{\mathcal{X}}(h) \le \mathcal{L}_s(h) + \sqrt{\frac{\log \frac{d\rho^s}{d\pi}(h) + \log \frac{2\sqrt{m}}{\delta}}{2m}}; $$

$$ \mathcal{L}_{\mathcal{X}}(h) \le \text{kl}^{-1} \left( \mathcal{L}_s(h) \left| \frac{\log \frac{d\rho^s}{d\pi}(h) + \log \frac{2\sqrt{m}}{\delta}}{m} \right. \right), $$

$$ \text{where } \text{kl}^{-1}(u|c) = \sup\{v \in [0, 1] : \text{kl}(u\|v) \le c\}. $$

# 4 Disintegrated bounds under gradient flow dynamics

We state here our main result, a disintegrated bound for an algorithm trained by continuous gradient-descent optimisation. We recall some notation from Section 2. We denote as $\mathcal{C}_s$ a twice differentiable training objective, and we consider a random initial state $h_0 \sim \rho_0$. The dynamics are described by

$$ \partial_t h_t = -\nabla \mathcal{C}_s(h_t), $$

where we omit the explicit dependence of $h_t$ on $s$ and $h_0$. The algorithm's output is given by $h_T$, with $T > 0$ some fixed time horizon. We define $\rho_t$ as the push-forward of $\rho_0$ under the gradient flow, so that $h_t \sim \rho_t$. For simplicity, for all $t$ we assume that $\rho_t$ admits a density with respect to the Lebesgue measure, which we will again denote as $\rho_t$, with a slight abuse of notation.

It is natural to suppose that $\rho_0$ is chosen without relying on the specific dataset $s$ used for the training, so that, in particular, we can select $\rho_0$ as our PAC-Bayesian prior. On the other hand, $\rho_T$ depends on $s$ through the objective $\mathcal{C}_s$, and plays the role of the posterior. Sampling the initialisation $h_0$ of the model from $\rho_0$, and then following the deterministic training dynamics up to $T$, we get a final hypothesis $h_T$, which can be seen as a sample from $\rho_T$. In particular, saying that an event holds with probability at least $1 - \delta$ under $(s, h_T) \sim \mu^m * \rho_T$ is equivalent to state that it has probability at least $1 - \delta$ under $(s, h_0) \sim \mu^m \otimes \rho_0$. We can thus use Proposition 1 to obtain a PAC-Bayesian generalisation bound for an algorithm that, once drawn $s$ and $h_0$, is deterministic.

So far, we have a bound that involves the term $\rho_T(h_T)/\rho_0(h_T)$. Typically, the density $\rho_0$ is known, and evaluating $\rho_0(h_T)$ does not pose any problem. On the other hand, computing $\rho_T(h_T)$ is less straightforward. However, thanks to the continuity of the gradient descent dynamics, it is possible to keep track of the evolution of $\rho_t(h_t)$ along the training trajectory via a simple trick, usually referred to as “instantaneous change of variables” in the normalising flow literature (Chen et al., 2018).

**Proposition 2.** Consider the dynamics $\partial_t h_t = -\nabla \mathcal{C}_s(h_t)$, where $\mathcal{C}_s: \mathcal{H} \to \mathbb{R}$ is twice differentiable. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be a measurable function. Fixed $T > 0^2$ and $\delta \in (0, 1)$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_{\mathcal{X}}(h_T)) \le \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_s(h_t) dt + \log \frac{\xi}{\delta}, $$

$$ \text{where } \Delta \text{ denotes the Laplacian with respect to } h \text{ and } \xi = \int_{\mathcal{X}^m \times \mathcal{H}} e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_{\mathcal{X}}(h))} d\mu^m(s) d\rho_0(h). $$

*Proof.* Since sampling $(s, h_T)$ from $\mu^m * \rho_T$ is equivalent to drawing $(s, h_0) \sim \mu^m * \rho_0$ and following the dynamics up to $T$, Proposition 1 yields directly that, with probability at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$,

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_{\mathcal{X}}(h_T)) \le \log \frac{\rho_T(h_T)}{\rho_0(h_T)} + \log \frac{\xi}{\delta}. $$

The gradient flow's continuity equation states that, for all $h \in \mathcal{H}$,

$$ \partial_t \rho_t(h) = \nabla \cdot (\rho_t(h) \nabla \mathcal{C}_s(h)) = \nabla \rho_t(h) \cdot \nabla \mathcal{C}_s(h) + \rho_t(h) \Delta \mathcal{C}_s(h). $$

<footnote n="2">²To be fully rigorous, one should explicitly assume that the solution $t \mapsto h_t$ of the dynamics is well defined on $[0, T]$, for all initial conditions $h_0 \in \mathcal{H}$.</footnote>
### 3. Disintegrated PAC-Bayes for continuous-time gradient flows

We begin by considering the continuous-time dynamics of the gradient flow. While this setting is less realistic than the one considered in the discrete-time analysis to follow, the discussion is considerably cleaner and will help expose some of the primary ideas of the framework we propose. We define the gradient flow $(h_0, t) \mapsto \Phi_t^s(h_0) \in \mathcal{H}$ as the solution to the differential equation

$$ \partial_t \Phi_t^s(h_0) = -\nabla C_s(\Phi_t^s(h_0)); \qquad \Phi_0^s(h_0) = h_0. \qquad (1) $$

We assume that $C_s$ is such that a solution exists until a fixed time horizon $T > 0$, for all $h_0 \in \mathcal{H}$ and training datasets, $s \in \text{supp}(\mu)^{\otimes m}$. Since $h_0$ and $s$ are fixed prior to training, we will simplify the notation, using $h_t$ to denote the solution of (1). Given a random initialisation $\rho_0$ we define $\rho_t$ as the push-forward of $\rho_0$ under the gradient flow:

$$ \rho_t = \Phi_t^{s\#} \rho_0. $$

By sampling the initial parameters $h_0$ from $\rho_0$ and following the flow dynamics up to $T$, we get a hypothesis $h_T$ that is distributed according to $\rho_T$.

We take a PAC-Bayesian approach to deriving generalisation bounds, selecting $\rho_0$ as the prior and $\rho_T$ as the posterior distribution. With this, we obtain PAC-Bayesian generalisation bounds for an algorithm that, once $s$ and $h_0$ are drawn, is deterministic.

**Theorem 1** Consider the dynamics $\partial_t h_t = -\nabla C_s(h_t)$, where $C_s: \mathcal{H} \to \mathbb{R}$ is twice differentiable, and let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be an arbitrary measurable function. Taking $\delta \in (0, 1)$ and $T > 0$ fixed, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, it holds that

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_Z(h_T)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt + \log \frac{\xi}{\delta} \right), \qquad (2) $$

where $\Delta$ denotes the Laplacian with respect to $h$ and $\xi = \int_{\mathbb{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}, \mathcal{L}_Z(h))} d\mu^m(\bar{s})d\rho_0(h)$.

How one chooses $\Psi$ is dependent on the integrability of $e^{m\Psi(\mathcal{L}_{\bar{s}}, \mathcal{L}_Z)}$ and thus, is dependent on properties of the loss function, initial distribution and data distribution. In the following corollaries, we provide two concrete settings in which $\Psi$ can be chosen to obtain more explicit bounds. In Corollary 2, we consider the setting in which the loss function is sub-Gaussian in $h$, setting $\Psi(u, v) = (v - u)/\sqrt{m}$. When $\ell$ is bounded a tighter bound holds and we can obtain faster rates using the approach of Langford and Seeger (2001) and Maurer (2004), taking $\Psi$ to be $m$ times the relative entropy between two Bernoulli distributions $kl(u||v) = u \log \frac{u}{v} + (1 - u) \log \frac{1-u}{1-v}$. We remark that these choices of $\Psi$ are common in the PAC-Bayesian literature (Bégin et al., 2016; Alquier, 2021).

**Corollary 2** Assume that $\ell(h, \cdot)$ is R-sub-Gaussian<fnref n="2" /> for each $h \in \mathcal{H}$. Then, for any $\delta \in (0, 1)$ and $T > 0$, with a probability of at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_Z(h_T) \le \mathcal{L}_s(h_T) + \frac{1}{\sqrt{m}} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt + \log \frac{1}{\delta} + \frac{R^2}{2} \right). $$

<footnote n="2">2. $\ell(h, \cdot)$ is R-sub-Gaussian if for all $\lambda \in \mathbb{R}$ we have $\log \int_Z e^{\lambda \ell(h, z)} d\mu(z) \le \lambda \int_Z \ell(h, z) d\mu(z) + \frac{R^2 \lambda^2}{2}$.</footnote>
It follows that

$$ \partial_t(\rho_t(h_t)) = \partial_t \rho_t(h_t) + \nabla \rho_t(h_t) \cdot \partial_t h_t = \rho_t(h_t) \Delta C_s(h_t), $$

which yields

$$ \log \frac{\rho_T(h_T)}{\rho_0(h_0)} = \int_0^T \Delta C_s(h_t) dt. $$

Hence

$$ \log \frac{\rho_T(h_T)}{\rho_0(h_T)} = \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt, $$

and we conclude.

□

Note that the time horizon $T$ must be chosen a priori, as it cannot depend on the specific $s$ and $h_0$ used for the training. However, through a classical union bound argument, one can slightly loosen the bound while allowing to select the optimal $T$ from a finite pool of possible values. Concretely, it is possible to pick the best $T$ from a set containing $K$ potential choices at the cost of replacing the term $\log \frac{\xi}{\delta}$ with $\log \frac{K\xi}{\delta}$, in the bounds of Propositions 1 and 2.

Clearly, for $\ell \subseteq [0, 1]$, the choices of $\Psi$ that lead to Corollary 1 will again bring explicit bounds. Setting $\Psi(u, v) = 2m(u - v)^2$ yields

$$ \mathcal{L}_{\mathcal{X}}(h_T) \le \mathcal{L}_s(h_T) + \sqrt{\frac{\log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \log \frac{2\sqrt{m}}{\delta} + \int_0^T \Delta C_s(h_t) dt}{2m}}, \quad (1) $$

while from $\Psi(u, v) = mkl(u|v)$ we get

$$ \mathcal{L}_{\mathcal{X}}(h_T) \le \mathrm{kl}^{-1} \left( \mathcal{L}_s(h_T) \left| \frac{\log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \log \frac{2\sqrt{m}}{\delta} + \int_0^T \Delta C_s(h_t) dt}{m} \right| \right), \quad (2) $$

where each of the above two bounds holds with probability higher than $1 - \delta$ on the random initialisation and training dataset.

We remark that Proposition 2 holds also for time-dependent objectives (as long as they do not depend on $\rho_t$), as the gradient flow's continuity equation is still valid. Leveraging this, we can specialise the result for training with (potentially randomly selected) batches. We refer to Section 5.5 for further details. Finally, an interesting alternative scenario is when the continuous-time dynamics are stochastic. For instance, one might want to look at the Langevin's diffusion $dh_t = -\nabla C_s(h_t) + \sigma dB_t$, where $\sigma > 0$ is a fixed coefficient and $B_t$ a Brownian motion. However, in this case we do not end up anymore with an easily tractable exact expression for $\log \frac{\rho_T(h_T)}{\rho_0(h_0)}$ (see Appendix A).

# 5 Analysis and discussion

## 5.1 A few comments on the Laplacian term

An interesting feature of the bound in Proposition 2 is the integral of the Laplacian of the optimisation objective along the training path:

$$ \int_0^T \Delta C_s(h_t) dt. \quad (3) $$

Under the gradient flow dynamics, $\Delta C_s(h_t)$ is keeping track of how the local probability density is locally varying. Indeed, from a PAC-Bayesian point of view, for the bound to be small we need to end up in some point where the *posterior* density is not too high compared to the initial *prior*. If we are following a trajectory characterised by high values of the Laplacian, we will see a sharp increase of the density. Intuitively, we can picture the situation as if we were attracting the nearby paths and bringing further *probability mass* around us. In this case we are likely to end with a final $\rho_T$ that is much larger than the initial $\rho_0$, and hence, potentially, a loose bound.

It is worth noticing that, in the one-dimensional case, (3) takes a particularly simple form, as it only depends on the initial and final “velocities”. To see this, note that $\partial_t C_s'(h_t) = C_s''(h_t) \partial_t h_t = -C_s''(h_t) C_s'(h_t)$,
**Corollary 3** Assume that $\ell$ is bounded in $[0, 1]$ and $m \ge 8$. Then, for any fixed $\delta \in (0, 1)$ and $T > 0$, with a probability of at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_{\mathcal{Z}}(h_T) \le \text{kl}^{-1}\left(\mathcal{L}_s(h_T)\left|\frac{B_T}{m}\right|\right); \quad B_T = \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t)dt + \log \frac{2\sqrt{m}}{\delta}, \quad (3) $$

where we define $\text{kl}^{-1}(u|c) = \sup\{v \in [0, 1] : \text{kl}(u||v) \le c\}$. In particular, it follows that

$$ \mathcal{L}_{\mathcal{Z}}(h_T) \le \mathcal{L}_s(h_T) + \sqrt{\frac{\mathcal{L}_s(h_T)B_T}{2m}} + \frac{B_T}{2m}. \quad (4) $$

We note that as long as $B_T = O(\log m)$, one can expect standard rates of order $O(\sqrt{\log m/m})$. However, an advantage of the kl formulation is that it can yields faster rates if the empirical loss is controlled by $\mathcal{L}_s(h_T) = O(1/m)$ (i.e., when the model is able to fit very well the training data). Indeed, in such case (4) is a fast-rate bound $O(\log m/m)$ (see, e.g., the discussion after Theorem 1 in Tolstikhin and Seldin, 2013 or in Section 2 of Mhammedi et al., 2019). In Appendix B, we derive additional bounds for settings with weaker concentration guarantees, including sub-exponential concentration, by making use of recent results from Casado et al. (2024).

We remark that, in Theorem 1, the time horizon $T$ must be chosen prior to training and cannot depend on $s$ and $h_0$. However, the result can be generalised to allow for the algorithm to pick the best time horizon among a set of $\kappa$ candidates, $\{T_1, \dots, T_\kappa\}$, only suffering an additional penalty of $\log \kappa$ on the right-hand side of (2) (i.e., replacing $\log \frac{\kappa}{\delta}$ with $\log \frac{\kappa \delta}{\delta}$). This follows from an elementary union argument where for each $T_k$, we consider a bound holding with probability at least $1 - \delta/\kappa$. As a final comment, we mention that the RHS of the bound in Theorem 1 diverges as $T \to \infty$, which reflects the fact that PAC-Bayes bounds are vacuous for degenerate posteriors.

**Proof of Theorem 1** The proof is based on two steps. First, we keep track of how the density evolves during training, obtaining an explicit expression for the posterior density. Then, we apply this in combination with a classical Markov's inequality argument.

For the first step, from the continuity equation<fnref n="3" /> of the gradient flow, we have that

$$ \partial_t \rho_t(h) = \nabla \cdot (\rho_t(h) \nabla C_s(h)), \quad \text{for all } h \in \mathcal{H}, \quad (5) $$

and furthermore, we obtain that $\rho_t$ admits a Lebesgue density for all $t \in [0, T]$. From this, we obtain

$$ \partial_t(\rho_t(h_t)) = \partial_t \rho_t(h_t) + \nabla \rho_t(h_t) \cdot \partial_t h_t = \rho_t(h_t) \Delta C_s(h_t), $$

and in particular,

$$ \log \frac{\rho_T(h_T)}{\rho_0(h_0)} = \int_0^T \Delta C_s(h_t)dt. $$

For the second step, it follows from Markov's inequality that with a probability of at least $1 - \delta$ on $(s, h_T) \sim \mu^m * \rho_T$, it holds that

$$ e^{m\Psi(\mathcal{L}_s(h_T), \mathcal{L}_{\mathcal{Z}}(h_T)) - \log \frac{\rho_T(h_T)}{\rho_0(h_T)}} \le \frac{1}{\delta} \int_{\mathcal{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_{\mathcal{Z}}(h)) - \log \frac{\rho_T(h)}{\rho_0(h)}} d\mu^m(\bar{s})d\rho_T(h). $$

<footnote n="3">3. The continuity equation is a key tool in studying the local evolution of the density (see, e.g., Chapter 8 in Ambrosio et al., 2008). It comes from fluid dynamics and expresses the fact that the local rate of change of the density of a fluid equals the negative divergence of the flux. In our case, the flux can be expressed as minus the gradient of the training objective, leading to (5). Note that (5) can also be seen as a simplified version of the classical Fokker-Planck equation, in the absence of the Brownian noise.</footnote>
with $C'_s$ and $C''_s$ denoting the first and second derivatives of $C_s$. It follows that $C''_s(h_t) = -\partial_t \log |C'_s(h_t)|$, which yields

$$ \int_{0}^{T} C''_{s}(h_{t})dt = \log \frac{C'_{s}(h_{0})}{C'_{s}(h_{T})}. $$

On the other hand, things are not that simple in the multidimensional case. With a few algebraic manipulations, we can write

$$ \int_{0}^{T} \Delta C_{s}(h_{t})dt = \log \frac{\|\nabla C_{s}(h_{0})\|}{\|\nabla C_{s}(h_{T})\|} - \int_{h_{[0:T]}} \nabla \cdot \tau(h) \|\delta h\|, \quad (4) $$

where $\tau(h)$ stands for the unit tangent vector to the trajectory in $h$, and $\int_{h_{[0:T]}} \dots \|\delta h\|$ denotes the line integral along the path $h_{[0:T]}$. The last term has a clear interpretation, as it quantifies how much $h_{[0:T]}$ is attracting the nearby trajectories. We refer to Appendix B for the derivation of (4).

Although we are not aware of any generalisation bound that involves the exact term (3), the Laplacian's integral might find a natural interpretation through the connection between flatness in the loss landscape and good generalisation. It dates back to Hochreiter and Schmidhuber (1997) the belief that an algorithm, able to find wide local optima of the loss, performs well on unseen inputs. Re-popularised by Keskar et al. (2017), this idea has been the object of an intense debate, attracting some criticisms as the definitions of flatness are often susceptible to re-parameterisations of the hypothesis space (Dinh et al., 2017; Neyshabur et al., 2017). However, numerous recent works have provided empirical and theoretical support, as well as new insights, to the flat minima's argument (Izmailov et al., 2018; He et al., 2019; Dziugaite and Roy, 2017; Neu et al., 2021). Going back to our setting, if we assume that the objective $C_s$ coincides with the empirical loss, we clearly see that sharper loss's local optima correspond to higher values of the Laplacian term. As a consequence, converging towards flatter minima will typically yield a tighter bound.

## 5.2 Role of the time horizon

Most of the standard PAC-Bayesian bounds (in expectation) contain the Kullback-Leibler divergence between the *prior* and the *posterior*. This quantity is infinity if the *posterior* is a Dirac delta, and this is why these results cannot apply directly to the deterministic setting. Even the disintegrated framework is not exempt from this pathological behaviour, which is reflected by the fact that in Proposition 2 we cannot let the time horizon $T$ to be infinite even if $h_t$ converges, as the *posterior* would be degenerate and typically consist in a sum of Dirac deltas centred on the local optima. Concretely, we can easily see that if an algorithm is converging to a local minimum (where $C_s$ is approximately quadratic), then (1) will asymptotically go as

$$ \mathcal{L}_{\mathcal{X}}(h) - \mathcal{L}_{s}(h) \lesssim O(\sqrt{T}) $$

for large $T$ (see Appendix C for more details). However, this is not the only possible asymptotic behaviour. For instance, we can consider situations where $h_t$ is not converging but diverging. As we will see in Section 6, this is what happens for linear models trained to optimise a linear objective or the cross-entropy.

Although characterising the typical spectrum of the Hessian of the loss during the training evolution is a challenging open problem, recent empirical studies suggest that after an initial phase where the Laplacian can potentially be negative, later stages of the dynamics are characterised by a few dominant positive Hessian's eigenvalues (Sagun et al., 2018; Ghorbani et al., 2019). As a consequence, in most situations the Laplacian's integral will soon become monotonically increasing with $T$. Moreover, we can expect that $\log \frac{\rho_0(h_0)}{\rho_0(h_T)}$ will grow with $T$ as well, since $h_0$ is likely to lie in a region where the *prior* density is high, and the training will typically push $h_t$ farther and farther from it. Overall, in order to achieve non-vacuous bounds, we need to stop the training at a finite time horizon $T$, which will play the role of a trade-off between how well we can fit the training dataset, and how much we can keep the bound small.

A final remark, already mentioned earlier, is that $T$ must be chosen independently of the dataset $s$ and initial state $h_0$ used for the training. However, adding a penalty $\log K$ to the bounds of Propositions 1 and 2 allows us to pick the best time horizon among a set $\{T_1 \dots T_K\}$ of $K$ possible choices. This is the consequence of an elementary union argument, as for each $T_k$ we can consider a bound holding with probability at least $1 - \delta/K$.
Combining this with the fact that, for all $\bar{s} \in \mathbb{Z}^m$,

$$ \int_{\mathcal{H}} \frac{\rho_0(h)}{\rho_T(h)} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_{\mathcal{Z}}(h))} d\rho_T(h) = \int_{\{\rho_T > 0\}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_{\mathcal{Z}}(h))} d\rho_0(h), $$

we obtain

$$ m\Psi(\mathcal{L}_{\bar{s}}(h_T), \mathcal{L}_{\mathcal{Z}}(h_T)) \le \log \frac{\rho_T(h_T)}{\rho_0(h_T)} + \log \frac{\xi}{\delta}, $$

with a probability of at least $1 - \delta$ on $(s, h_T) \sim \mu^m * \rho_T$. Since sampling $(s, h_T)$ from $\mu^m * \rho_T$ is equivalent to drawing $(s, h_0) \sim \mu^m \otimes \rho_0$ and following the dynamics up to $T$, the bound equivalently holds with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$. Using that

$$ \log \frac{\rho_T(h_T)}{\rho_0(h_T)} = \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \log \frac{\rho_T(h_T)}{\rho_0(h_0)} = \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt $$

we obtain the bound in the statement.

<figure><img src="image_2.png" /></figure>

## 4. Discrete time dynamics

In this section, we consider the gradient descent (GD) algorithm,

$$ h_{k+1} = h_k - \eta_k \nabla C_s(h_k), $$

where $\{\eta_k\}_{k=0}^{K-1}$ is the training schedule and the number of iterations, $K \ge 1$, is fixed. We let $\rho_k$ be the law of $h_k$, and assume that $\rho_0$ admits a positive Lebesgue density on the whole $\mathcal{H}$.

The primary obstacle in reproducing the methodology of the previous section is that we can no longer use the gradient flow continuity equation to keep track of the density change along the trajectory. However, the change in density can still be computed exactly as long as we can ensure that the update map is injective and differentiable along the path.

**Theorem 4** For any dataset $s \in \mathbb{Z}^m$, assume there is a Borel set $A_s \subseteq \mathcal{H}$ on which $C_s$ is twice-differentiable and M-smooth, where $\sup_k \eta_k \le 1/(2M)$. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be measurable. Taking $\delta \in (0, 1)$ and $K \in \mathbb{N}$ fixed, suppose that $\{h_k\}_{k=0}^{K-1}$ lies in $A_s$ with a probability of at least $1 - \delta/2$ under $(s, h_0) \sim \mu^m \otimes \rho_0$. Then, with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$, it holds that

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_\mathcal{Z}(h_T)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_k)} - \sum_{k=0}^{K-1} \mathrm{tr} \log (\mathrm{Id} - \eta_k \nabla^2 C_s(h_k)) + \log \frac{2\xi}{\delta} + \delta \right), \quad (6) $$

where $\xi = \int_{\mathbb{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_{\mathcal{Z}}(h))} d\mu^m(\bar{s}) d\rho_0(h)$.

We refer to Appendix A for the proof. We note that the term $-\mathrm{tr}(\log(\mathrm{Id} - \eta_k \nabla^2 C_s(h_k)))$ in (6) can be upper bounded in various ways, using the fact that $C_s$ is smooth around $h_k$.

**Lemma 5** With the notation of Theorem 4, let $h_k \in A_s$. Then

$$ -\mathrm{tr} \log (\mathrm{Id} - \eta_k \nabla^2 C_s(h_k)) \le \eta_k \Delta C_s(h_k) + \eta_k^2 \|\nabla^2 C_s(h_k)\|_{\mathrm{F}}^2 \le \frac{3}{2} \eta_k \|\nabla^2 C_s(h_k)\|_{\mathrm{TR}}, $$

where $\| \cdot \|_{\mathrm{F}}$ is the Frobenius norm and $\| \cdot \|_{\mathrm{TR}}$ the trace norm.<fnref n="4" />

<footnote n="4">4. For a matrix $U$ with singular values $\{\sigma_i\}$, let $\|U\|_{\mathrm{TR}} = \sum_i \sigma_i$ and $\|U\|_{\mathrm{F}} = \sqrt{\sum_i \sigma_i^2} = \mathrm{tr}[UU^\top]^{1/2}$.</footnote>
## 5.3 Scaling with the dimension of the hypothesis space

When dealing with generalisation bounds, a natural question is how they scale with the dimension $N$ of the hypothesis space. In our case, the answer is not easy, as it is hard to say how the different components of bound depend on $N$. For instance, we can expect the Laplacian to scale linearly with $N$, as it is the sum of $N$ terms. If $T = O(1)$ and the Laplacian’s integral dominates the bound, (1) becomes $O(\sqrt{N/m})$, which is a standard behaviour. On the other hand, if the log-density term $\log \frac{\rho_0(h_0)}{\rho_0(h_T)}$ is the main contribution to the bound, the scaling with $N$ might be more complex, as it will depend on the particular choice of the *prior* density.

Interestingly, the optimal time horizon $T$ can scale with $N$ as well, and empirical results on MNIST show that for a linear model we can have $T \sim O(1/N)$, and hence obtain non-vacuous bounds even in a highly over-parameterised setting (see Figure 1). Intuitively, when $N$ is large, the algorithm can quickly find an hypothesis that fit well the training data and is “close” to $h_0$, hence requiring a low training time to be reached. On the other hand, fewer learnable parameters might need a longer training time before reaching low values of the empirical loss. Despite the simple setting of our experiment, we believe that this is a promising result as the bound seems unaffected by the *curse of dimensionality*. However, a more detailed empirical and theoretical analysis is needed in order to gauge how well our bound performs with more complex models, such as deep networks.

<figure><img src="image_2.png" /><figcaption>**Figure 1:** Behaviour of best time horizon $T$ with the model dimension $N$, from experiments on MNIST for linear models trained to optimise a cross-entropy objective. The optimal $T$ was found to go roughly as $1/N$ (cf. dotted line), preventing the bounds (1) and (2) from exploding, even in highly over-parameterised settings. Note that the results reported here were obtained via discretised gradient dynamics (cf. Section 5.4 and Figure 2). We refer to Appendix E.3 for further details on the experiment.</figcaption></figure>

<figure><img src="image_3.png" /><figcaption>**Figure 2:** Relative error $(B_1 - B_2)/B_2$, for a linear model with $10^5$ parameters trained on MNIST via discretised gradient steps on a cross-entropy objective. Here, $B_1$ and $B_2$ are the values of the bound (2) obtained with the time-steps $\delta_{t,1} = 2 \times 10^{-4}$ and $\delta_{t,2} = 2 \times 10^{-6}$ respectively. The small relative error suggests that the corrections to the bound due to the time-discretised dynamics are likely to be negligible. We refer to Appendix E.3 for further details on the experiment.</figcaption></figure>

## 5.4 Discretised dynamics

One of the main limitations of the present work is that the bound of Proposition 2 requires a continuous-time gradient descent. Apart from a few simple settings, where analytic solutions are available, the standard gradient-based optimisation algorithms consider discretised dynamics, whose trajectories are made of jumps and only approximate the continuous flow. As a consequence, in order to leverage our results and get rigorous generalisation bounds for the standard gradient-descent optimisers, one would need to control the errors introduced by the time discretisation via a careful analysis, which is beyond the scope of the present work. However, we can at least conjecture that, for small enough learning rates, the results obtained via the discretised dynamics are a good approximation of the exact solution of the gradient flow. In Figure 2, the relative error between the values of the bound (2), obtained from two different time discretisations, stays small throughout the training, suggesting that the discretised results can be good approximations of the actual continuous dynamics.
From the first inequality, we see that the continuous time bound of Theorem 1 is recovered when one takes the learning rate to 0. From the final inequality, it follows that this term scales at worst as $O(d)$, as the smoothness assumption ensures that $\eta_k \|\nabla^2 C_s(h_k)\|_{TR} \le d/2$. However, it is likely that in many cases, this translates to an overly pessimistic estimate. For instance, we show in Remark 18, in Appendix C.1, that for a simple random feature model one can upper bound $\|\nabla^2 C_s\|_{TR}$ with a term that is of order $O(1)$ for large $d$.

# 5. Examples with closed-form results

To complement the results of the previous sections, we consider toy examples for which more explicit forms for the bound can be obtained.

## 5.1. Random feature model

To start, we investigate a simple feature model trained via continuous-time gradient descent on a simple regression task. We let $\mathcal{X} = S^{p-1}$ be the unit sphere in $\mathbb{R}^p$ and $\mathcal{Y} = [0, 1]$. The goal is to learn a target function $f: \mathcal{X} \to \mathcal{Y}$. We consider the class of mappings $F_h: \mathcal{X} \to \mathbb{R}$ defined as

$$F_h(x) = \frac{1}{\sqrt{d}} h \cdot \Phi(x),^5$$

for $h \in \mathbb{R}^d$, where $\Phi: \mathbb{R}^p \to \mathbb{R}^d$ is a feature map, kept fixed during training. Since we know that $f$ takes values in $[0, 1]$, we let the prediction by the model be $f_h = \max\{0, \min\{1, F_h\}\}$. We consider a quadratic loss $\ell(h, z) = (f_h(x) - y)^2$, and an optimisation objective $C_s(h) = \frac{1}{m} \sum_{z \in s} \hat{\ell}(F_h(x), y)$, with $\hat{\ell}(F, y) = (F - y)^2$. We focus the setting of random feature models (Rahimi and Recht, 2007; Mei and Montanari, 2022), where the features are given by $\Phi(x) = \phi(Wx)$. Here, $\phi$ is a non-linearity acting component-wise, and $W$ is a $d \times p$ matrix whose components are independently drawn from a standard Gaussian distribution. Interestingly, in the overparameterised regime $d \to \infty$ the bound does not become vacuous (see Appendix C.1 for the proof).

**Proposition 6** For the random feature model described above, we have the limit in probability

$$\lim_{d \to \infty} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt \right) \le T \left( \mathbb{E}_{\zeta \sim \mathcal{N}(0,1)} [\phi(\zeta)^2] + 2\sqrt{C_s(h_0)} \right).$$

The next natural question is how the bound scales with the dataset size $m$. In general, one cannot say much. Indeed, if $f$ is just random noise, it is clear that it is not possible to get a small population loss, although for long enough training time $T$ one could potentially learn the training dataset. However, we conjecture that if the target $f$ is reasonably nice, one can get a bound of order $\mathcal{L}_Z(h_T) \lesssim \frac{\log m}{m}$ by selecting $T \sim \log m$. To be more explicit, we consider a simple setting, where $f$ is a linear combination of finitely many spherical harmonics (see Appendix C.1 for more details).

**Proposition 7** Let $f: S^{d-1} \to [0, 1]$ be a linear combination of spherical harmonics of degree at most $D$. Let $\phi$ be given by $\phi(x) = \sum_{k=1}^J (x/e)^j$, with $J \ge D$.<fnref n="6" /> Then, there exists a constant $\lambda > 0$,

<footnote n="5">5. The factor $1/\sqrt{d}$ is a standard choice for this kind of models, as it brings $\|h/\sqrt{d}\| \sim O(1)$ at initialisation, when the components of $h$ are independently initialised $\mathcal{N}(0, 1)$ (e.g., see Ghorbani et al., 2021 for random feature models).</footnote>

<footnote n="6">6. In general, up to a zero-measure set of coefficients, any polynomial activation of degree $J$ would lead to the same result. The specific one picked here is to give an explicit example.</footnote>
Interestingly, Barrett and Dherin (2021) remarked that the solutions of the discretised dynamics $h_{k+1} - h_k = -\epsilon \nabla C_s(h_k)$ follow closely the continuous paths yielded by $\partial_t h_t = -\nabla(C_s(h_t) + \epsilon \nabla C_s(h_t)\|^2)$, with an error after a single step of order $O(\epsilon^3)^3$. Consequently, if we aim at studying the discrete dynamics, we could obtain a more precise approximation of an actual generalisation bound by replacing the term $\Delta C_s$ with $\Delta(C_s(h_t) + \epsilon \nabla C_s(h_t)\|^2)$. If higher accuracy is required, one can add further correction terms to $C_s$, as it is standard in the backward error analysis of Runge-Kutta methods (Hairer et al., 2006).

## 5.5 Training via batches

So far, we have considered the case of a training objective $C_s$ which is fixed throughout the training. However, the bound of Proposition 2 holds also for time-dependent objectives. In particular, we can apply our result to the case in which the training is performed by optimising the empirical loss over different batches of the training dataset. Let us denote as $s_1...s_M$ the (temporarily ordered) sequence of batches used for the training, where a batch needs to appear multiple times if it is used more than once. It is worth noticing that the batches and their order are allowed to be selected randomly, as it is usually the standard stochastic gradient descent optimisation. Then, we can suppose that from $t = 0$ to $t = t_1$ we use $C_{s_1}$ as optimisation objective, from $t = t_1$ to $t = t_2$, we use $C_{s_2}$, and so on until $t_M = T$. We can get a bound as the one of Proposition 2, with the only difference that $\int_0^T \Delta C_s(h_t)dt$ will now be replaced by

$$ \int_{0}^{t_1} \Delta C_{s_1}(h_t)dt + \int_{t_1}^{t_2} \Delta C_{s_2}(h_t)dt + \dots + \int_{t_{M-1}}^{t_M} \Delta C_{s_M}(h_t)dt. \quad (5) $$

## 5.6 Data-dependent priors

For most of PAC-Bayesian results, the prior is required to be chosen independently of the dataset used to evaluate the bound. However, we can use an additional dataset $s'$ (independent of $s$) to learn the prior. This strategy was first proposed by Seeger (2002) and has lead to the tightest currently available empirical PAC-Bayesian bounds on MNIST and CIFAR10 (Pérez-Ortiz et al., 2021; Clerico et al., 2022a).

In our setting we might proceed in a similar way. Till a fixed time $t_0$, we follow the dynamics $\partial_t h_t = -\nabla C_{s'}(h_t)$, and we select $\rho_{t_0}$ as the PAC-Bayesian prior. Then, we take $C_s$ as training objective and we go on as usual, starting from $h_{t_0}$ and following the gradient flow until the time horizon $T$. In order to evaluate the bound, we need to compute $\log \frac{\rho_T(h_T)}{\rho_{t_0}(h_T)}$. Proceeding as in the proof of Proposition 2, we can get

$$ \log \frac{\rho_T(h_T)}{\rho_{t_0}(h_T)} = \log \frac{\rho_{t_0}(h_{t_0})}{\rho_{t_0}(h_T)} + \int_{t_0}^{T} \Delta C_s(h_t)dt = \log \frac{\rho_0(h_{t_0})}{\rho_{t_0}(h_T)} + \int_{0}^{t_0} \Delta C_{s'}(h_t)dt + \int_{t_0}^{T} \Delta C_s(h_t)dt. $$

Now, the only problem is to evaluate $\log \rho_{t_0}(h_T)$, as we do not have a closed form for $\rho_{t_0}$. However, it is possible to exactly compute this term as

$$ \log \rho_{t_0}(h_T) = \log \rho_0(\hat{h}_{T-t_0}) + \int_{0}^{t_0} \Delta C_{s'}(\hat{h}_{T-t_0+t})dt, \quad (6) $$

where $\hat{h}_t$ is the solution of $\partial_t \hat{h}_t = -\nabla C_{s'}(\hat{h}_t)$ satisfying $\hat{h}_T = h_T$. We hence get

$$ \log \frac{\rho_T(h_T)}{\rho_{t_0}(h_T)} = \log \frac{\rho_0(h_{t_0})}{\rho_0(\hat{h}_{T-t_0})} + \int_{0}^{t_0} (\Delta C_{s'}(h_t) - \Delta C_{s'}(\hat{h}_{T-t_0+t}))dt + \int_{t_0}^{T} \Delta C_s(h_t)dt. $$

Plugging this expression in Proposition 1 we readily obtain a bound whose prior $\rho_{t_0}$ is data-dependent. We refer to Appendix D for an explicit statement of the bound and a derivation of (6).

<footnote n="3">³For $\partial_t h_t = -\nabla C_s(h_t)$ the single-step error is of order $O(\epsilon^2)$.</footnote>
such that for $T = \frac{\log m}{4\lambda}$,

$$ \lim_{d \to \infty} \mathcal{L}_Z(h_T) \le O_p \left( \frac{\log m}{m} \right), \quad \text{as } m \to \infty. $$

## 5.2. Wide neural networks

Next, we investigate the first term of the bound, $\log \rho_0(h_0)/\rho_0(h_k)$, by considering the setting of wide feed-forward neural networks. We assume that each layer has the same activation and the hidden layers have the same width $n$. Each of the weights and bias parameters are initialised with a centred Gaussian distribution with variance $\sigma_w^2/n$ and $\sigma_b^2$, respectively. We consider the quadratic loss as the optimisation objective and have the training schedule scale with $n$ such that $\eta_k = \bar{\eta}_k/n$, where $\bar{\eta}_k$ is independent of $n$. Thus, as $n$ grows, the evolution of the network can be approximated by linear dynamics according to the neural tangent kernel (NTK) (Jacot et al., 2018).

Borrowing the analysis of Lee et al. (2020), we obtain two properties of the large $n$ setting: (i) with high probability, $h_0$ (the initial value of all weights and biases) is such that $C_s$ is smooth and bounded in a region around it and (ii) with the same probability, gradient descent stays close to initialisation. With this, we can show that this setting satisfies the assumptions of Theorem 4, and thus, we obtain the generalisation bound in the proposition that follows. We refer to Appendix C.2 for the formal statement and proof, as well as further discussion.

**Proposition 8 (Informal statement)** Under suitable regularity conditions, for any $\delta \in (0, 1)$, there exists $n_{\min}^{\delta} \in \mathbb{N}$ such that whenever $n \ge n_{\min}^{\delta}$, the assumptions of Theorem 4 are satisfied. In particular, with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$,

$$ \Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \le \frac{1}{m} \left( \frac{C}{\sigma_{\omega}^2} A_K(h_0) - \sum_{k=0}^{K-1} \mathrm{tr} \log (\mathrm{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k)) + \log \frac{2\xi}{\delta} + \delta \right), $$

where we define,

$$ A_K(h_0) = \left( \left( \sum_{k=0}^{K-1} \bar{\eta}_k \right) \wedge \lambda_{\min}^{-1} \right) (J_s(h_0) + 1), \quad J_s(h_0)^2 = \frac{1}{m} \sum_{i=1}^{m} \langle h_0, \nabla_h F_{h_0}(x_i) \rangle, $$

$\lambda_{\min}$ is the minimum eigenvalue of the NTK of the finite width network, $\xi$ is defined as in Theorem 4 and $C > 0$ is a constant.

To better understand the nature of this bound, we discuss how both $A_K$, as well as the curvature dependent term could be further controlled in the limit as $n \to \infty$. In Remark 22, we show that a naive estimate can be used to obtain $J_s(h_0) = O_p(\sqrt{n})$ as $n \to \infty$. Furthermore, we argue that a central limit heuristic could be used to show that as $n \to \infty$, $J_s(h_0) = O_p(1)$. However, stating a formal argument in this direction is beyond the scope of this work. The Hessian of the NTK has been analysed in various recent works with a notable contribution from Jacot et al. (2020), who show that the Laplacian of the NTK training objective converges in the limit as $n \to \infty$ to the trace of the NTK and a term that decays over iterations. In Remark 23, we show that their arguments can also be applied to control the Hessian dependent term as $n \to \infty$. A consequence of this analysis is that for large $k$, with high probability, it holds that

$$ \lim_{n \to \infty} - \mathrm{tr} \log (\mathrm{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k)) \le 3\bar{\eta}_k \mathrm{tr}(\Theta)/m. $$
# 6 Linear models

To give an idea of the behaviour of the bounds that we propose, we report here some analytic and empirical results. We consider models linear in the parameters, in the form

$$F_h^i(x) = \sum_j h^{ij} \Phi^j(x),$$

where the $\Phi^j$ are some fixed mappings $\mathbb{R} \to \mathbb{R}$. For binary classification tasks, where $y \in \{-1, +1\}$, we let $F_h$ be a real function and the model's prediction $f_h$ corresponds to the sign of $F_h$. Conversely, for general multi-class problems the predicted label $f_h$ is the index of the largest component of $F_h$.

As it is standard in the classification setting, we can assume that $\ell$ is the 01-loss<fnref n="4" />, so that $\mathcal{L}_s$ and $\mathcal{L}_X$ represent the empirical and population error. Unless otherwise stated, we will always consider an optimisation objective in the form $\mathcal{C}_s(h) = \hat{\mathcal{L}}_s(h) = \frac{1}{m} \sum_{x \in s} \hat{\ell}(F_h(x), f^*(x))$, with $\hat{\ell}$ denoting some twice differentiable surrogate loss function. Applying the chain rule to $\hat{\ell}$, we find that

$$\Delta \hat{\ell} = \sum_i \frac{\partial \hat{\ell}}{\partial F^i} \Delta F^i + \sum_{ii'} \frac{\partial^2 \hat{\ell}}{\partial F^i \partial F^{i'}} \nabla F^i \cdot \nabla F^{i'}$$

where $\Delta$ and $\nabla$ always refer to derivatives with respect to $h$. In the linear models that we are considering $\Delta F = 0$ and $\nabla F^i \cdot \nabla F^{i'} = \delta_{ii'} \|\Phi\|^2$. Thus

$$\Delta \hat{\ell} = \|\Phi\|^2 \Delta_F \hat{\ell}, \qquad (7)$$

with $\Delta_F \hat{\ell} = \sum_i \partial_{F^i}^2 \hat{\ell}$, the Laplacian of $\hat{\ell}$ with respect to $F$.

## 6.1 Linear loss

A first simple example that we can solve analytically is a binary classification problem with

$$\hat{\ell}(F, y) = -Fy.$$

The surrogate loss is linear in $F$, and so (7) yields $\Delta \hat{\ell} = 0$. In particular, the Laplacian's integral in our bounds vanishes. We can easily solve the evolution dynamics and find that

$$h_T = h_0 + \gamma T,$$

where, recalling that $f^*(x)$ is the correct label of the example $x$, we have defined $\gamma = \frac{1}{m} \sum_{x \in s} f^*(x) \Phi(x)$. Assuming that under $\rho_0$ the components of $h$ are all independent centred random variables, with variance $1/N$ ($N$ being the dimension of $h$), we find that $\log \frac{\rho_0(h_0)}{\rho_0(h_T)} = \frac{N}{2}(\|h_T\|^2 - \|h_0\|^2)$. So, (1) becomes

$$\mathcal{L}_X(h_T) - \mathcal{L}_S(h_T) \le \sqrt{\frac{N}{4m}(2Th_0 \cdot \gamma + T^2 \|\gamma\|^2) + \frac{1}{2m} \log \frac{2\sqrt{m}}{\delta}}.$$

Tipically, $\|h_0\| \sim O(1)$, and we expect $\|h_T\| \sim T \|\gamma\|$ to dominate the dynamics. Notice that the bound scales linearly with $T$. This is different from the $O(\sqrt{T})$ behaviour that we had encountered in section 5.2. Indeed, here $h$ is not approaching a local optimum, but diverging at constant speed $\gamma$.

We tested empirically the behaviour of the bounds (1) and (2) with a linear surrogate loss, by training a model where the learnable vector $h$ had 10000 components. The dataset was made of 500 datapoints, coming from Gaussian clusters in $\mathbb{R}^5$ and labelled into two categories. The map $\Phi$ was in the form $x \mapsto \phi(Wx)$, with $\phi$ the ReLU function and $W$ a $(5 \times 10000)$-matrix whose entries were independent draws from a centred Gaussian with variance $1/5$. We refer to Appendix E.1 for further experimental details and results. Figure 3 compares the behaviour of the two bounds with the test errors obtained from a held-out dataset, for 50 different possible values of the time horizon<fnref n="5" /> and for $\delta = 5 \times 10^{-3}$. Note that for large $T$ (1) is linear, as expected.

<footnote n="4">⁴$\ell(y, y') = 0$ if $y = y'$; 1 otherwise.</footnote>

<footnote n="5">⁵The penalty factor $\log 50$ (see end of Section 5.2) is already included in the values reported.</footnote>
The bound given in Proposition 8 scales linearly with the number of iterations, and therefore, the rate of convergence of the training loss. This rate is governed by the spectrum of the NTK and therefore, it depends on the dataset. In the worst case, the convergence rate is controlled by the minimum eigenvalue of the NTK, making it is sufficient to have $\sum_{k=0}^{K-1} \bar{\eta}_k = \Omega(\lambda_{\min}^{-1} \log m)$. The spectrum of the NTK and more specifically, its minimum eigenvalue, has appeared in a variety of analyses of the NTK that have shown its relationship to the generalisation capabilities of the model (Arora et al., 2019) as well as memorisation properties (Montanari and Zhong, 2022).

## 6. Extension to other algorithms

While we have focused on gradient descent, the analysis can be extended to any iterative scheme

$$h_{k+1} = h_k + V_s(h_k; k),$$

where $V_s: \mathcal{H} \times \mathbb{N} \to \mathcal{H}$ can be any iteration-dependent vector field. This leads to the following generalisation of Theorem 4, which differs only in the trajectory dependent sum, where $V_s$ now replaces $-\nabla^2 C_s$. A similar result is found for continuous flows (Theorem 24 in Appendix D).

**Theorem 9** Consider the dynamics $h_{k+1} = h_k + V_s(h_k; k)$. For each dataset $s \in \mathcal{Z}^m$, denote as $A_s$ a Borel set on which $V_s$ is differentiable and M-Lipschitz, with $M \le 1/2$. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be a measurable function. Fix $K \in \mathbb{N}$ and choose $\delta \in (0, 1)$, such that the trajectory $\{h_k\}_{k=0}^{K-1}$ lies in $A_s$ with probability at least $1 - \delta/2$, under $(s, h_0) \sim \mu^m \otimes \rho_0$. Then, with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$,

$$\Psi(\mathcal{L}_s(h_K), \mathcal{L}_{\mathcal{Z}}(h_K)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_K)} - \sum_{k=0}^{K-1} \mathrm{tr} \log \left( \mathrm{Id} + \nabla V_s(h_k; k) \right) + \log \frac{2\xi}{\delta} + \delta \right),$$

where $\nabla V_s$ is the Jacobian of $V_s$ with respect to $h$, and $\xi = \int_{\mathcal{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_{\mathcal{Z}}(h))} d\mu^m(\bar{s})d\rho_0(h)$.

**Stochastic gradient descent** An immediate corollary of the above is that our theory applies to noisy variants of gradient descent with little modification. For example, we can consider a version of gradient descent that only evaluates $C$ on a mini-batch $s_k \subset s$ at each iteration $k \in \mathbb{N}$, by simply setting $V_s(h; k) = -\eta_k \nabla C_{s_k}(h)$. The resulting generalisation bound applies identically for stochastic variants of this scheme, including stochastic gradient descent, where the resulting bound is a function of the instance of the sampled mini-batches. To make things more explicit, we consider a surrogate loss function $\hat{\ell}: \mathcal{Z} \times \mathcal{H} \to \mathbb{R}$ and for a batch $s_k \subset s$, we write $C_{s_k}(h) = \frac{1}{|s_k|} \sum_{z \in s_k} \hat{\ell}(z, h)$. For a sequence of batches $\{s_k\}$ (potentially randomly selected), we consider the dynamics $h_{k+1} = h_k - \eta_k \nabla C_{s_k}(h_k)$. Then, under suitable smoothness assumptions for $\hat{\ell}$, we can derive from Theorem 9 the following bound

$$\Psi(\mathcal{L}_s(h_K), \mathcal{L}_{\mathcal{Z}}(h_K)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_K)} - \sum_{k=0}^{K-1} \mathrm{tr} \log \left( \mathrm{Id} + \eta_k \Delta C_{s_k}(h_k) \right) + \log \frac{2\xi}{\delta} + \delta \right),$$

which holds with probability at least $1 - \delta$ on the randomness of the training dataset $s$, the initialisation $h_0$, and the choice of the batches (see Proposition 25 in Appendix D.1).
<figure><img src="image_4.png" /><figcaption>**Figure 3:** Behaviour of the bounds (1) and (2) for a linear model with 10000 parameters, trained with a linear surrogate loss, for different time horizons $T$. The data came from eight different Gaussian distributions on $\mathbb{R}^5$ that had been assigned binary labels. The training dataset counted 500 datapoints and a held-out test dataset was used to evaluate the test errors reported in the plot. We refer to Appendix E.1 for further details.</figcaption></figure>

<figure><img src="image_5.png" /><figcaption>**Figure 4:** Behaviour of the bounds (1) and (2) for a linear model with 10000 parameters, trained with a quadratic surrogate loss, for different time horizons $T$. The learning task was the same as in Figure 3. Again, we used a training dataset made of 500 datapoints and we evaluated the test errors on a held-out test dataset. The parameters $\alpha$ and $\beta$ in $C_s$ were both set to 1. We refer to Appendix E.2 for further details.</figcaption></figure>

## 6.2 Quadratic loss

Another natural situation where we can solve things explicitly is the case of a quadratic surrogate loss

$$ \hat{\ell}(F, y) = \frac{1}{2}(F - \beta y)^2. $$

Here, $\beta$ is some fixed positive value, and we are still considering a binary classification problem, where the true labels are +1 or -1 and the network's prediction is the sign of $F$. Now, $\Delta_F \hat{\ell} = 1$ and hence (7) yields $\Delta \hat{\ell} = \|\Phi\|^2$. We consider a training objective in the form

$$ C_s(h) = \frac{\alpha \|h\|^2}{2N} + \frac{1}{m} \sum_{x \in s} \hat{\ell}(F_h(x), f^*(x)), $$

with $\alpha \ge 0$ and $N$ the dimension of $h$. The regularising term $\alpha \|h\|^2/(2N)$ is here introduced to computationally stabilise the algorithm.<fnref n="6" /> We easily find that

$$ \int_{0}^{T} \Delta C_{s}(h_{t}) dt = \Gamma T, $$

with $\Gamma = \alpha + \frac{1}{m} \sum_{x \in s} \|\Phi(x)\|^2$. The training dynamics are now governed by $\partial_t h_t = \beta\gamma - \Theta h_t$, where again we have defined $\gamma = \frac{1}{m} \sum_{x \in s} f^*(x)\Phi(x)$, while $\Theta = \frac{\alpha}{N} \text{Id} + \frac{1}{m} \sum_{x \in s} \Phi(x)\Phi(x)^{\top}$ is a $N \times N$ matrix. Hence, assuming that $\Theta$ is invertible, we have

$$ h_T = h_0 + (\text{Id} - e^{-T\Theta})(\Theta^{-1}\beta\gamma - h_0). $$

If again we consider a Gaussian initialisation, (1) becomes

$$ \mathcal{L}_{\mathcal{X}}(h_T) - \mathcal{L}_{s}(h_T) \le \sqrt{\frac{\frac{N}{2}(\|h_0\|^2 - \|h_T\|^2) + \Gamma T}{2m}}. $$

As the objective is quadratic, for large $T$ this bound will behave as $O(\sqrt{T})$, as it happens more generally every time the algorithm is converging to a local minimum of the objective (see Section 5.2).

Figure 4 shows the bounds obtained by optimising a quadratic objective with $\alpha = \beta = 1$. The dataset and the setting were those that we used for the linear loss's training, which we described in the Section 6.1. Again, we used $\delta = 5 \times 10^{-3}$. We refer to Appendix E.2 for further experimental details and results.

<footnote n="6">⁶In the unregularised case $\alpha = 0$, the matrix inversion $\Theta^{-1}$, necessary to compute $h_T$, can be source of significant computational errors. Introducing a factor $\alpha > 0$ makes the algorithm more stable.</footnote>
**Momentum dynamics** We can also use Theorem 9 to consider settings in which auxiliary variables are used to compute the update. We do this by replacing $h_k$ with the pair $(h_k, v_k)$, where $v_k$ denotes the auxiliary variable. Indeed, this setting applies to a wide range of optimisation schemes. Note that in this scenario, the initial density $\rho_0$ must refer to the pair $(h_0, v_0)$.

An example is the momentum dynamics $h_{k+1} = h_k + v_{k+1}$ and $v_{k+1} = \mu_k v_k - \eta_k \nabla C_s(h_k)$, for some momentum schedule $\{\mu_k\}$. In such a case we obtain a high probability bound in the form

$$ \Psi(\mathcal{L}_s(h_K), \mathcal{L}_{\mathcal{Z}}(h_K)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0, v_0)}{\rho_0(h_K, v_K)} - d \sum_{k=0}^{K-1} \log \frac{1}{\mu_k} + \log \frac{2\xi}{\delta} + \delta \right). $$

We refer to Appendix D.2 for further details and discussion.

**Damped Hamiltonian dynamics** For damped Hamiltonian dynamics (França et al., 2020), we can exploit the fact that the joint density of the ‘position-momentum’ pair $(h, v)$ is conserved under the Hamiltonian flow, a property preserved for discrete time-steps via symplectic integrators (Hairer et al., 2006). Using this, we obtain bounds for discrete-time algorithms without any smoothness assumptions on $C_s$, other than twice differentiability. We refer to Appendix E for the details.

# 7. Comparison with the literature

## 7.1. Comparison with other PAC-Bayes bounds

Contrary to many PAC-Bayesian results, our bounds have the remarkable feature of applying to neural networks with deterministic parameters, trained via standard gradient-descent methods. Yet, this is not a complete novelty. A few previous works in the literature propose to study the generalisation of a deterministic network via the PAC-Bayes analysis of a noisy stochastic perturbation of it. This idea was exploited by Neyshabur et al. (2018), and later Biggs and Guedj (2022a), for a L-layer fully connected architecture with 1-Lipschitz homogeneous activation function and 1-Lipschitz loss. Combining margin arguments with PAC-Bayes techniques they found (up to logarithmic factors)

$$ \mathcal{L}_{\mathcal{Z}}(h) \le \mathcal{L}_{s}^{\gamma}(h) + O\left(\sqrt{dn\Gamma/(\gamma^2 m)}\right), \qquad (7) $$

where $n$ is the width of the network, $\mathcal{L}_s^\gamma$ the margin empirical loss with margin $\gamma > 0,7$ and $\Gamma = \sum_{l=1}^L (\|W_l\|_F^2 \prod_{l' \neq l} \|W_{l'}\|)$, with $\{W_l\}_{l=1}^L$ the weights of the network, $\| \cdot \|_F$ denoting the spectral norm, and $\| \cdot \|_F$ the Frobenius norm. One of the main issues of this result is $\Gamma$'s exponential dependence on the depth, due to the product of the norm of the weights. On the other hand, our bounds involve the Hessian term, which are at most of order $d$ (see Lemma 5), and the contribution $\log \frac{\rho_0(h_0)}{\rho_0(h_T)}$. When the weights are independently initialised as $\mathcal{N}(0, 1)$, this last term is upper bounded by $\frac{1}{2} \sum_{l=1}^L \|W_l\|_F^2$. Moreover, as shown in (4), for small enough empirical loss $\mathcal{L}_s = O(1/m)$, our bound can have a fast-rate dependence of $O(1/m)$ on the training dataset size.

Later, building on the ideas from Neyshabur et al. (2018), Nagarajan and Kolter (2019) obtained a bound that does not suffer of the exponential dependence on the depth. However, this comes at the price of inversely scaling with the smallest absolute value of the pre-activations on the training data, leading to vacuous bounds in practice. We also mention that a result similar to that of

<footnote n="7">7. The margin $\gamma$ is a standard measure of the confidence of the network's prediction. We refer to Neyshabur et al. (2018) or Biggs and Guedj (2021) for a definition of the margin loss. Note however that we always have $\mathcal{L}_s^\gamma \ge \mathcal{L}_s$.</footnote>
## 6.3 Cross-entropy loss

For a general classification task, which can possibly have more than two labels, when $f_h = \operatorname{argmax}_i F_h^i$ a common optimisation objective is the cross-entropy loss

$$ \hat{\ell}(F, y) = \log \sum_i e^{F^i} - F^y . $$

In this case we do not have simple analytic dynamics, and we will need to rely on some numerical solution. We have

$$ \partial_t h_t^{ij} = -\frac{1}{m} \sum_{x \in s} \left( \frac{e^{(h_t \Phi(x))^i}}{\sum_k e^{(h_t \Phi(x))^k}} - \delta^{if^*(x)} \right) \Phi^j(x), $$

where $\delta^{if^*(x)}$ denotes the Kronecker delta between the labels $i$ and $f^*(x)$, and

$$ \Delta \hat{\ell} = \|\Phi\|^2 \left( 1 - \frac{\sum_k e^{2F^k}}{(\sum_k e^{F^k})^2} \right) . $$

Contrarily to what happened with the quadratic and linear losses, where the objective's Laplacian stayed constant during the training, here we are in the interesting situation where the Laplacian tends to decrease during the training. Indeed, for $\hat{\ell}(F, y)$ to be small, we need $F^y$ to be much larger than the other components of $F$, so that $\sum_k e^{F^k} \simeq e^{F^y}$. But this also means that the Laplacian is small, since $(\sum_k e^{2F^k})/(\sum_k e^{F^k})^2 \simeq 1$. This is a different behaviour from what happens when the algorithm is converging to a local minimum of the objective, as as now we can expect that the Laplacian's integral will have a sub-linear behaviour in $T$. Indeed, in the present situation, we are not approaching a local optimum. Conversely, $h_t$ is slowly diverging, reaching flatter and flatter regions of $\mathcal{H}$.

We trained linear models of different sizes on MNIST, using the cross-entropy loss as objective. Since the dynamics does not admit a simple analytic solution, we performed the classical gradient-descent algorithm on batches of the training dataset, with finite time-steps. As discussed in Section 5.4, what we obtain in this way are not rigorous generalisation bounds, as we are not following the continuous dynamics. However, we believe that the result that we get are a good approximation of the ones that one would

obtain with the exact solution of the gradient flow (see Figure 2 and the discussion in Section 5.4). The feature mapping $\Phi$ that we used was again the composition of a ReLU function and a linear transformation $x \mapsto Wx$, where the components of $W$ were independently sampled from a normal distribution, with variance $1/784$ (784 being the dimension of the input). We refer to Appendix E.3 for further experimental details.

Table 1 compares the bound (2) with the test error evaluated on a held-out dataset, for different dimensions of the hypothesis space $\mathcal{H}$, and always with $\delta = 5 \times 10^{-3}$. We chose the best bound among those obtained with 25 different values of $T$, at the price of adding a penalty factor $\log 25$ (see Section 5.2).

As we had already stressed in Figure 1 (which refers to the same experiment), note that the bounds are not negatively affected by the high over-parameterisation, essentially because the optimal time horizon $T$ decreases roughly as $1/N$. The tightest bound was achieved for a network of width $10000$ (i.e., $N = 10^5$). For this model, Figure 5 shows the evolution of the bounds (1) and (2) during the training. The values of the empirical error and of the penalty factor ($\log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \log \frac{25\sqrt{m}}{\delta} + \int_0^T \Delta C_s(h_t)dt)/m$ are also reported. Figures 6 compares the contributions of the Laplacian's integral and the log-density term, showing that in this case it is the latter that dominates.

As a final remark, we notice that the bounds of Table 1 are comparable to those reported by Clerico et al. (2021) for a stochastic single-hidden-layer network. Notably, the results therein were obtained by using a PAC-Bayesian bound as optimisation objective, while here we are directly optimising the empirical cross-entropy loss, without any regularisation term.

<table><caption>Table 1: Bounds on MNIST</caption><thead><tr><th>Dimension N</th><th>Bound (2)</th><th>Test Error</th></tr></thead><tbody><tr><td>10<sup>3</sup></td><td>.3251</td><td>.2049</td></tr><tr><td>3 × 10<sup>3</sup></td><td>.3034</td><td>.1967</td></tr><tr><td>10<sup>4</sup></td><td>.2786</td><td>.1680</td></tr><tr><td>3 × 10<sup>4</sup></td><td>.2795</td><td>.1541</td></tr><tr><td>10<sup>5</sup></td><td>.2751</td><td>.1602</td></tr><tr><td>3 × 10<sup>5</sup></td><td>.2812</td><td>.1534</td></tr><tr><td>10<sup>6</sup></td><td>.2757</td><td>.1568</td></tr></tbody></table>
Neyshabur et al. (2018) was previously established by Bartlett et al. (2017), without PAC-Bayes techniques.

The bounds discussed so far only take into account the final output of the algorithm, while our result looks at the evolution of the model during the training. A similar spirit is shared by Miyaguchi (2019), where the author focuses on the continuous time setting and studies the evolution of the generalisation gap under the gradient flow training dynamics. Applying this result to a multilayer network, it is possible to re-derive (7) under slightly weaker assumptions.

Most other PAC-Bayesian results for deterministic models cannot be applied to standard training algorithms, as they require strong, and often unusual, assumptions on the model architecture (e.g., Letarte et al., 2019; Germain et al., 2009; Biggs and Guedj, 2021), or sampling the parameters from the posterior distribution (e.g., Zantedeschi et al., 2021; Viallard et al., 2021; Rivasplata et al., 2020).

Finally, it is worth mentioning the work of Luo et al. (2022), which also studies de-randomised PAC-Bayesian guarantees for GD methods. However, it is hard to directly compare our findings with their results, as their framework and ours differ significantly. Specifically, their analysis focuses on discretised versions of GD and SGDm and their bounds are stated in terms of the gradient discretisation error, and become vacuous as this error approaches zero (standard GD/SGD). Moreover, their bounds require the prior distribution to depend on a subset of the training dataset. Overall, we believe that our results address more conventional versions of GD and SGD, albeit with the trade-off of requiring a smoothness assumption not needed in their work.

## 7.2. Comparison with the stability literature

Another method for obtaining algorithm-dependent generalisation bounds is the framework of uniform stability (Hardt et al., 2016; Pensia et al., 2018; Farghly and Rebeschini, 2021; Raj et al., 2023), proposed by Elisseeff (2005). This approach has recently received attention due to its application to fundamental optimisation methods, such as gradient descent and its stochastic counterpart (Hardt et al., 2016), while other works have leveraged it to obtain bounds in high-probability (Feldman and Vondrak, 2018, 2019; Bousquet et al., 2020). Hardt et al. (2016) considered the SGD training with non-convex training objectives. For $C_s$ is $M$-smooth in $h$ (uniformly on $\mathcal{Z}$), $\ell$ is $L$-Lipschitz and bounded in $[0, 1]$, and the step-size satisfying $\eta_k \le c/(k+1)$, the analysis of Hardt et al. (2016) and Bousquet et al. (2020) leads to the bound (with probability at least $1-\delta$ and for some constant $C$)

$$ \mathcal{L}_{\mathcal{Z}}(h) \le \mathcal{L}_s(h) + C \left( \frac{(K/c)^{\frac{Mc}{Mc+1}} \log(m/\delta)}{Mm} + \sqrt{\frac{\log \delta^{-1}}{m}} \right), \quad (8) $$

To compare with these results, under the same assumptions for GD, Theorem 4 yields (cf. also (4))

$$ \mathcal{L}_{\mathcal{Z}}(h_K) \le \mathcal{L}_s(h_K) + \sqrt{\frac{\mathcal{L}_s(h_K)B_K}{2m}} + \frac{B_K}{2m}, $$

with

$$ B_K = \log \frac{\rho_0(h_0)}{\rho_0(h_K)} + \sum_{k=0}^{K-1} (\eta_k \Delta C_s(h_k) + \eta_k^2 \|\nabla^2 C_s(h_k)\|_F^2) + \log(m/\delta). \quad 8 $$

<footnote n="8">8. Note that since here we have that the smoothness holds on the whose $\mathcal{H}$, we do not need the additional term we can replace $\log(2\delta/m) + \delta$ in (6) with $\log(\delta/m)$.</footnote>
<figure><img src="image_6.png" /></figure>

**Figure 5:** Behaviour of the bounds (1) and (2) at different time horizons *T*, for a linear model with $10^5$ parameters trained with the cross-entropy loss on MNIST. The training dataset counted *m* = 60000 datapoints and a held-out test dataset was used for the test errors reported. The plot also shows the values of the penalty factor ($\log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \log \frac{25\sqrt{m}}{\delta} + \int_0^T \Delta C_s(h_t)dt/m$, where $\delta = 5 \times 10^{-3}$ See Appendix E.3 for more experimental details.

<figure><img src="image_7.png" /></figure>

**Figure 6:** For the same experiment illustrated in Figure 5, comparison of the log density term $\log \frac{\rho_0(h_0)}{\rho_0(h_T)}$ and the Laplacian's integral $\int_0^T \Delta C_s(h_t)dt/m$. The former clearly dominates on the latter. Note that the contribution of $\log \frac{25\sqrt{m}}{\delta} \approx 14.02$ to the bound is negligible.

# 7 Conclusion

We established disintegrated PAC-Bayesian bounds for deterministic algorithms, holding with high probability on the random training dataset and initialisation. We tested our findings empirically on linear models, obtaining promising results that are non-vacuous even for highly over-parameterised settings. It is worth noticing that most of empirical PAC-Bayesian results in the literature are derived by using a PAC-Bayesian bound as the optimisation objective (Dziugaite and Roy, 2017; Pérez-Ortiz et al., 2021; Clerico et al., 2021, 2022a), while our bounds directly apply to the standard (continuous-time) empirical risk optimisation. Yet, our results on MNIST closely match the bounds obtained by Clerico et al. (2021). We defer to future work a more extensive empirical analysis, focusing on data-dependent priors and non-linear models.

A main limitation of our framework is the necessity of continuous-time dynamics. However, it might be possible to obtain rigorous results for the standard discretised gradient-descent optimisation via a careful numerical analysis. Alternatively, a scenario which might be worth exploring is the limit of infinite width, where the NTK dynamics (Jacot et al., 2018) might provide interesting continuous-time exact results. Finally, we mention as possible further developments a deeper study of stochastic training dynamics, where the gradient flow's continuity relation is replaced by the Fokker-Planck equation, and the case of a training objective that depends on the bound itself.

**Acknowledgments:** This work has been partially supported by the UK Engineering and Physical Sciences Research Council (EPSRC) through the grants EP/R513295/1 (DTP scheme) and CoSInES EP/R034710/1. The authors would like to thank Umut Şimşekli, Tyler Farghly, and Patrick Rebeschini for the valuable comments and suggestions.
Using Proposition 25 a similar bound can be obtained for SGD.

As a first point of comparison, we note that our analysis does not require the Lipschitz assumption but only smoothness along the path of GD with high probability. Additionally, our analysis holds in both the stochastic and non-stochastic setting, whereas the technique used by Hardt et al. (2016) fundamentally requires random mini-batches to have bounds that decay with $m$. While the bound of (8) decays at a rate of $m^{-1/2}$, our bound can decays faster with rates $\log(m)/m$ (for $\mathcal{L}_s \sim 1/m$). The fact that $\sum_{k=0}^{K-1} \eta_k \sim c \log K$ suggests that our bound may scale better with $K$, though this would require the $\Delta \mathcal{L}_s(h_k)$ and $\|\nabla^2 \mathcal{L}_s(h_k)\|_F$ terms to not grow too quickly with $k$. In the worst case, the smoothness can be used to upper bound these terms by $3dM/2$.

Lastly, we note that one of the main criticisms to the uniform stability approach is that it is solely related to the algorithm and does not consider specifics of the data or the distribution of the labels, raising doubts on its ability to distinguish whether a model has been trained on true or random labels (Zhang et al., 2016). On the other hand, our bound can depend on the data distribution through the optimisation objective $\mathcal{L}_s$ and its landscape along the training trajectory.

### 7.3. Comparison with information-theoretic bounds

Another popular direction within the literature on generalisation bounds uses ideas from information theory to upper bound the expected generalisation error in terms of the mutual information (Xu and Raginsky, 2017; Russo and Zou, 2019). This has been particularly practical for developing data-dependent bounds for noisy iterative methods, such as stochastic gradient Langevin dynamics and SGD (Mou et al., 2018; Negrea et al., 2019; Neu et al., 2021).

The general approach to this requires controlling the mutual information between the training data and the update of each iterate. Therefore, this technique is restricted to settings where noise is applied at each iteration, and the bounds explode when the amount of noise is reduced. To apply this to GD and SGD, Neu et al. (2021) consider a surrogate model trained by a Gaussian perturbation of these iterates. When the loss is $R$-sub-Gaussian, they bound the expected generalisation error by

$$ \mathbb{E}\mathcal{L}_Z(h_K) \le \mathbb{E}\mathcal{L}_s(h_K) + \left( \frac{R^2}{m} \sum_{k=1}^{K} \eta_k^2 \mathbb{E}A(h_k) \right)^{1/2} + |\mathbb{E}B(h_K)|, $$

where $A(h)$ and $B(h)$ measures the sensitivity of the gradient and the loss function, respectively, to perturbations in the parameters and dataset at $h$. A notable difference between this technique and our method is that this can only provide bounds in expectation. This comes with the downside that the right-hand side can usually not be computed exactly and the expectation of $A$ and $B$ should be approximated using a Monte Carlo average. In contrast, our bound is based on the instance of the optimisation trajectory, it can be computed exactly. However, our bounds have similar dependence on $m$ but worse dependence on $\eta_k$.

## 8. Conclusion

We derive novel high-probability generalisation bounds for models learned via optimisation algorithms such as gradient descent. Contrary to the standard PAC-Bayesian framework, our guarantees apply to models whose only randomness lies in the initialisation without requiring any *de-randomisation* step. To the best of our knowledge, our results are the first to leverage the disinte-
References

P. Alquier. User-friendly introduction to PAC-Bayes bounds. arXiv:2110.11216, 2021.

A. Banerjee, T. Chen, and Y. Zhou. De-randomized PAC-Bayes margin bounds: Applications to non-convex and non-smooth predictors. arXiv:2002.09956, 2020.

D.G.T. Barrett and B. Dherin. Implicit gradient regularization. ICLR, 2021.

F. Biggs and B. Guedj. Differentiable PAC-Bayes objectives with partially aggregated neural networks. Entropy, 23(10), 2021.

F. Biggs and B. Guedj. On margins and derandomisation in PAC-Bayes. AISTATS, 2022.

G. Blanchard and F. Fleuret. Occam's hammer. COLT, 2007.

O. Bousquet and A. Elisseeff. Stability and generalization. Journal of Machine Learning Research, 2, 2002.

O. Bousquet, S. Boucheron, and G. Lugosi. Introduction to Statistical Learning Theory. Springer, 2004.

L. Bégin, P. Germain, F. Laviolette, and J.F. Roy. PAC-Bayesian bounds based on the Rényi divergence. AISTATS, 2016.

O. Catoni. PAC-Bayesian supervised classification: The thermodynamics of statistical learning. IMS Lecture Notes Monograph Series, 2007.

R.T.Q. Chen, Y. Rubanova, J. Bettencourt, and D.K. Duvenaud. Neural ordinary differential equations. NeurIPS, 2018.

E. Clerico, G. Deligiannidis, and A. Doucet. Wide stochastic networks: Gaussian limit and PAC-Bayesian training. arxiv:2106.09798, 2021.

E. Clerico, G. Deligiannidis, and A. Doucet. Conditionally Gaussian PAC-Bayes. AISTATS, 2022a.

E. Clerico, A. Shidani, G. Deligiannidis, and A. Doucet. Chained generalisation bounds. COLT, 2022b.

L. Dinh, R. Pascanu, S. Bengio, and Y. Bengio. Sharp minima can generalize for deep nets. ICML, 2017.

C. Dwork and A. Roth. The algorithmic foundations of differential privacy. Foundations and Trends in Theoretical Computer Science, 9, 2014.

G.K. Dziugaite and D. M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. UAI, 2017.

P. Germain, A. Lacasse, F. Laviolette, and M. Marchand. PAC-Bayesian learning of linear classifiers. ICML, 2009.

B. Ghorbani, S. Krishnan, and Y. Xiao. An investigation into neural net optimization via Hessian eigenvalue density. ICML, 2019.

I. Goodfellow, Y. Bengio, and A. Courville. Deep Learning. MIT Press, 2016.

B. Guedj. A primer on PAC-Bayesian learning. Proceedings of the Second Congress of the French Mathematical Society, 2019.

E. Hairer, C. Lubich, and G. Wanner. Geometric numerical integration. Springer-Verlag, 2006.

M. Hardt, B. Recht, and Y. Singer. Train faster, generalize better: Stability of stochastic gradient descent. ICML, 2016.

H. He, G. Huang, and Y. Yuan. Asymmetric valleys: Beyond sharp and flat local minima. NeurIPS, 2019.

K. He, X. Zhang, S. Ren, and J. Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. ICCV, 2015.
grated PAC-Bayesian framework to analyse such settings. We make this explicit by stating a bound
that holds for wide neural networks trained via gradient descent.

A strength of our bounds is that it assumes little about the model or training procedure. For the
continuous gradient flow dynamics, we require only that the optimisation objective be twice differ-
entiable. For the discrete-time algorithm, we require smoothness and twice differentiability only in
high probability on the trajectory. Additionally, we show that our results can be extended to settings
more general than gradient descent and give explicit bounds for SGD, momentum schemes, and
damped Hamiltonian dynamics. We foresee that this should motivate further work into developing
generalisation bounds for other optimisation algorithms.

A promising direction for future work could be designing computationally efficient methods for
computing these bounds. We would also like to evaluate the tightness of our guarantees and compare
them with other results known in the literature with thorough empirical investigation. Finally, we
believe that our results can be improved by identifying more easily verifiable assumptions to make
the framework more broadly applicable.

Acknowledgments

Eugenio Clerico was partially supported by the UK Engineering and Physical Sciences Research Council (EPSRC) through the grant EP/R513295/1 (DTP scheme) and Arnaud Doucet by EPSRC CoSInES EP/R034710/1. Tyler Farghly was supported by EPSRC EP/T5178 and by the DeepMind scholarship. Benjamin Guedj and Arnaud Doucet acknowledge support of the UK Defence Science and Technology Laboratory (DSTL) and EPSRC grant EP/R013616/1. This work was part of the collaboration between US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative. Benjamin Guedj acknowledges partial support from the French National Agency for Research, grants ANR-18-CE40-0016-01 and ANR-18-CE23-0015-02. The authors would like to thank Jake Fawkes, Shahine Bouabid, Umut Şimşekli, and Patrick Rebeschini for the valuable comments and suggestions.
S. Hochreiter and J. Schmidhuber. Flat minima. Neural Computations, 9, 1997.

P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A.G. Wilson. Averaging weights leads to wider optima and better generalization. UAI, 2018.

A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: Convergence and generalization in neural networks. NeurIPS, 2018.

N.S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P.T.P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.

J. Langford and M. Seeger. Bounds for averaging classifiers. CMU technical report, 2001.

J. Langford and J. Shawe-Taylor. PAC-Bayes & margins. NeurIPS, 2002.

G. Letarte, P. Germain, B. Guedj, and F. Laviolette. Dichotomize and generalize: PAC-Bayesian binary activated deep neural networks. NeurIPS, 2019.

G. Lugosi and G. Neu. Generalization bounds via convex analysis. COLT, 2022.

A. Maurer. A note on the PAC Bayesian theorem. arXiv:0411099, 2004.

D.A. McAllester. Some PAC-Bayesian theorems. COLT, 1998.

D.A. McAllester. PAC-Bayesian model averaging. COLT, 1999.

K. Miyaguchi. PAC-Bayesian transportation bound. arXiv:1905.13435, 2019.

V. Nagarajan and J.Z. Kolter. Deterministic PAC-Bayesian generalization bounds for deep networks via generalizing noise-resilience. ICLR, 2019.

G. Neu, G.K. Dziugaite, M. Haghifam, and D. M. Roy. Information-theoretic generalization bounds for stochastic gradient descent. COLT, 2021.

B. Neyshabur, S. Bhojanapalli, D.A. McAllester, and N. Srebro. Exploring generalization in deep learning. NeurIPS, 2017.

B. Neyshabur, S. Bhojanapalli, and N. Srebro. A PAC-Bayesian approach to spectrally-normalized margin bounds for neural networks. ICLR, 2018.

M. Pérez-Ortiz, O. Risvaplata, J. Shawe-Taylor, and C. Szepesvári. Tighter risk certificates for neural networks. Journal of Machine Learning Research, 22, 2021.

O. Rivasplata, I. Kuzborskij, C. Szepesvári, and J. Shawe-Taylor. PAC-Bayes analysis beyond the usual bounds. NeurIPS, 2020.

L. Sagun, U. Evci, V.U. Güney, Y. Dauphin, and L. Bottou. Empirical analysis of the Hessian of over-parametrized neural networks. ICLR Workshop, 2018.

M. Seeger. PAC-Bayesian generalisation error bounds for Gaussian process classification. Journal of Machine Learning Research, 3, 2002.

S. Shalev-Shwartz and S. Ben-David. Understanding Machine Learning - From Theory to Algorithms. Cambridge University Press, 2014.

Y. Song, J. Sohl-Dickstein, D.P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. arXiv:2011.13456, 2020.

V.N. Vapnik. The Nature of Statistical Learning Theory. Springer, 2000.

P. Viallard, P. Germain, A. Habrard, and E. Morvant. A general framework for the disintegration of PAC-Bayesian bounds. arXiv:2102.08649, 2021.

A. Xu and M. Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. NeurIPS, 2017.

C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. Communications of the ACM, 64, 11 2016.

W. Zhou, V. Veitch, M. Austern, R.P. Adams, and P. Orbanz. Non-vacuous generalization bounds at the ImageNet scale: a PAC-Bayesian compression approach. ICLR, 2019.
## References

P. Alquier. User-friendly introduction to PAC-Bayes bounds. arXiv:2110.11216, 2021.

P. Alquier and G. Biau. Sparse single-index model. *Journal of Machine Learning Research*, 14(1), 2013.

L. Ambrosio, N. Gigli, and G. Savaré. *Gradient Flows in Metric Spaces and in the Space of Probability Measures*. Birkhäuser, 2008.

S. Arora, S. Du, W. Hu, Z. Li, and R. Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. *ICML*, 2019.

A. Banerjee, T. Chen, and Y. Zhou. De-randomized PAC-Bayes margin bounds: Applications to non-convex and non-smooth predictors. arXiv:2002.09956, 2020.

P.L. Bartlett, D.J. Foster, and M.J. Telgarsky. Spectrally-normalized margin bounds for neural networks. *NeurIPS*, 2017.

F. Biggs and B. Guedj. Differentiable PAC-Bayes objectives with partially aggregated neural networks. *Entropy*, 23(10), 2021.

F. Biggs and B. Guedj. On margins and derandomisation in PAC-Bayes. *AISTATS*, 2022a.

F. Biggs and B. Guedj. Non-vacuous generalisation bounds for shallow neural networks. *ICML*, 2022b.

G. Blanc, N. Gupta, G. Valiant, and P. Valiant. Implicit regularization for deep neural networks driven by an Ornstein-Uhlenbeck like process. *COLT*, 2020.

G. Blanchard and F. Fleuret. Occam's hammer. *COLT*, 2007.

S. Boucheron, G. Lugosi, and P. Massart. *Concentration Inequalities - A Nonasymptotic Theory of Independence*. Oxford University Press, 2013.

O. Bousquet and A. Elisseeff. Stability and generalization. *Journal of Machine Learning Research*, 2, 2002.

O. Bousquet, S. Boucheron, and G. Lugosi. *Introduction to Statistical Learning Theory*. Springer, 2004.

O. Bousquet, Y. Klochkov, and N. Zhivotovskiy. Sharper bounds for uniformly stable algorithms. *COLT*, 2020.

L. Bégin, P. Germain, F. Laviolette, and J.F. Roy. PAC-Bayesian bounds based on the Rényi divergence. *AISTATS*, 2016.

I. Casado, L.A. Ortega, A.R. Masegosa, and A. Pérez. PAC-Bayes-Chernoff bounds for unbounded losses. arXiv:2401.01148, 2024.

O. Catoni. *Statistical learning theory and stochastic optimization*. Ecole d'été de probabilités de Saint-Flour XXXI-2001. Springer, 2004.
# Appendix

## A Langevin dynamics

We briefly discuss what happens if we try to use a Langevin dynamics instead of a deterministic gradient method, that is

$$ dh_t = -\nabla C_s(h_t)dt + \sigma dB_t. \qquad (8) $$

Here, $\sigma > 0$ is a diffusion coefficient and $B_t$ denotes the standard Brownian motion. Note that the introduction of noise prevents the bound from diverging at large $T$. Indeed, under suitable regularity conditions for $C_s$, $\rho_t$ will converge to the Gibbs equilibrium distribution

$$ \rho_{\infty}(h) = \frac{1}{Z}e^{-2C_s(h)/\sigma^2}, $$

with $Z$ the normalising constant. As $\rho_{\infty}$ is bounded by $1/Z$, asymptotically $\log \frac{\rho_T(h_T)}{\rho_0(h_T)} \lesssim -\log(Z\rho_0(h_T))$ does not diverge with $T$.

For (9), the Fokker-Planck equation yields

$$ \partial_t \rho_t(h) = \nabla(\rho_t(h)\nabla C_s(h)) + \frac{\sigma^2}{2}\Delta\rho_t(h) = \nabla\rho_t(h) \cdot \nabla C_s(h) + \rho_t(h)\Delta C_s(h) + \frac{\sigma^2}{2}\Delta\rho_t(h). \quad (9) $$

Applying Itô's lemma we get

$$ d(\rho_t(h_t)) = (\rho_t(h_t)\Delta C_s(h_t) + \sigma^2\Delta\rho_t(h_t))dt + \sigma\nabla\rho_t(h_t) \cdot dB_t. $$

Now, writing $u_t = \log \rho_t$, we have

$$ du_t(h_t) = \left( \Delta C_s(h_t) + \frac{\sigma^2}{2} \|\nabla u_t(h_t)\|^2 + \sigma^2 \Delta u_t(h_t) \right) dt + \sigma \nabla u_t(h_t) \cdot dB_t, \quad (10) $$

where we used $\frac{\Delta\rho_t}{\rho_t} = \|\nabla u_t\|^2 + \Delta u_t$. This expression cannot easily be integrated in $t$, as it happened for the deterministic dynamics, since now we have a dependence on $u_t$ (and hence on $\rho_t$) in the RHS. In order to obtain a computable bound, one would probably need to upper-bound the intractable quantities in the RHS, an approach that we defer to future work.

An alternative way to tackle the stochastic dynamics consists in converting the SDE (8) in an ODE inducing the same marginal distributions (an idea already exploited in Song et al. (2020)). In fact, consider the deterministic flow

$$ \partial_t \hat{h}_t = -\nabla \left( C_s(\hat{h}_t) + \frac{\sigma^2}{2} \log \hat{\rho}_t(\hat{h}_t) \right), \quad (11) $$

where we denote as $\hat{\rho}_t$ the density of $\hat{h}_t$ and we impose $\hat{\rho}_0 = \rho_0$. The continuity equation now reads

$$ \partial_t \hat{\rho}_t(h) = \nabla \cdot \left( \hat{\rho}_t(h) \nabla \left( C_s(h) + \frac{\sigma^2}{2} \log \hat{\rho}_t(h) \right) \right) = \nabla(\hat{\rho}_t(h) \nabla C_s(h)) + \frac{\sigma^2}{2} \Delta \hat{\rho}_t(h). \quad (12) $$

Comparing (12) with (9), we immediately see that $\rho_t$ and $\hat{\rho}_t$ undergo the same dynamics and, since $\rho_0 = \hat{\rho}_0$, we have that $\rho_t = \hat{\rho}_t$ for all $t$. In practice, this means that we have two equivalent ways to get $h_T \sim \rho_T$: given an initial point from $\rho_0$, we can follow either the stochastic dynamics (8), or the deterministic ones (11), up to time $T$. Moreover, in order to evaluate the factor $\log \rho_T(h_T)$, needed for our PAC-Bayesian bound, instead of using (10) we can leverage the continuity equation of the deterministic flow to derive

$$ \log \rho_T(h_T) = \log \rho_0(\hat{h}_0) + \int_0^T \Delta \left( C_s(\hat{h}_t) + \frac{\sigma^2}{2} \log \rho_t(\hat{h}_t) \right) dt, $$

where $\hat{h}_t$ denotes the solution of (11) satisfying $\hat{h}_T = h_T$. Although this expression looks easier than (10), its actual computation is typically unfeasible in practice, due to the RHS's dependence on $\log \rho_t$ (both explicit and implicit, through $\hat{h}_t$). In the generative modelling context, score matching and neural networks techniques are used to approximate $\nabla \log \rho_t$, but this introduces approximation errors difficult to quantify (Song et al., 2020).
O. Catoni. PAC-Bayesian supervised classification: The thermodynamics of statistical learning. *IMS Lecture Notes Monograph Series*, 2007.

E. Clerico, G. Deligiannidis, and A. Doucet. Conditionally Gaussian PAC-Bayes. *AISTATS*, 2022a.

E. Clerico, A. Shidani, G. Deligiannidis, and A. Doucet. Chained generalisation bounds. *COLT*, 2022b.

A. Damian, T. Ma, and J.D. Lee. Label noise SGD provably prefers flat global minimizers. *NeurIPS*, 2021.

A. Daniely, R. Frostig, and Y. Singer. Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity. *arXiv:1602.05897*, 2017.

G.K. Dziugaite and D.M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. *UIA*, 2017.

A. Elisseeff. Stability of randomized learning algorithms. *Journal of Machine Learning Research*, 6, 2005.

M. Erdogdu, L. Mackey, and O. Shamir. Global non-convex optimization with discretized diffusions. *NeurIPS*, 2018.

T. Farghly and P. Rebeschini. Time-independent generalization bounds for SGLD in non-convex settings. *NeurIPS*, 2021.

V. Feldman and J. Vondrak. Generalization bounds for uniformly stable algorithms. *NeurIPS*, 2018.

V. Feldman and J. Vondrak. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. *COLT*, 2019.

G. França, J. Sulam, D. Robinson, and R. Vidal. Conformal symplectic and relativistic optimization. *NeurIPS*, 2020.

P. Germain, A. Lacasse, F. Laviolette, and M. Marchand. PAC-Bayesian learning of linear classifiers. *ICML*, 2009.

B. Ghorbani, S. Mei, T. Misiakiewicz, and A. Montanari. Linearized two-layers neural networks in high dimension. *The Annals of Statistics*, 49, 2021.

I. Goodfellow, Y. Bengio, and A. Courville. *Deep Learning*. MIT Press, 2016.

U. Grenander. Stochastic processes and statistical inference. *Arkiv Matematik*, 1(3), 10 1950.

B. Guedj. A primer on PAC-Bayesian learning. *Proceedings of the Second Congress of the French Mathematical Society*, 2019.

B. Guedj and P. Alquier. PAC-Bayesian estimation and prediction in sparse additive models. *Electronic Journal of Statistics*, 7, 2013.

E. Hairer, C. Lubich, and G. Wanner. *Geometric numerical integration*. Springer-Verlag, 2006.
As a final remark, we mention that the dynamics (11) also appear when trying to build an algorithm that directly optimises a PAC-Bayesian bound. Indeed, for $\ell \subseteq [0, 1]$, the choice $\Psi(u, v) = \frac{2}{\sigma^2}(u - v)$ (for an arbitrary $\sigma > 0$) in Proposition 1 leads to a bound in the form

$$ \mathcal{L}_{\mathcal{X}}(h_T) \le \mathcal{L}_s(h_T) + \frac{1}{4m\sigma^2} + \frac{\sigma^2}{2} \left( \log \frac{\rho_T(h_T)}{\rho_0(h_T)} + \log \frac{1}{\delta} \right). $$

Using this bound as optimisation objective is equivalent to follow the dynamics (11), provided that we set $\mathcal{C}_s(h) = \mathcal{L}_s(h) + \frac{1}{4m\sigma^2} + \frac{\sigma^2}{2}(-\log \rho_0(h) + \log \frac{1}{\delta})$. We defer to future work a deeper analysis of this connection.

## B Rewriting the Laplacian's integral

We explicitly derive here (4). First, notice that $\partial_t \log \|\nabla \mathcal{C}_s(h_t)\| = -\frac{\nabla \mathcal{C}_s(h_t)}{\|\nabla \mathcal{C}_s(h_t)\|} \cdot \nabla \|\nabla \mathcal{C}_s(h_t)\|$. Since, for all $h$,

$$ \Delta \mathcal{C}_s(h) = \nabla \cdot \nabla \mathcal{C}_s(h) = \nabla \cdot \left( \frac{\nabla \mathcal{C}_s(h)}{\|\nabla \mathcal{C}_s(h)\|} \right) \|\nabla \mathcal{C}_s(h)\| + \frac{\nabla \mathcal{C}_s(h)}{\|\nabla \mathcal{C}_s(h)\|} \cdot \nabla \|\nabla \mathcal{C}_s(h)\|, $$

we get

$$ \Delta \mathcal{C}_s(h_t) = \nabla \cdot \left( \frac{\nabla \mathcal{C}_s(h_t)}{\|\nabla \mathcal{C}_s(h_t)\|} \right) \|\nabla \mathcal{C}_s(h_t)\| - \partial_t \log \|\nabla \mathcal{C}_s(h_t)\|. $$

We conclude that

$$ \int_{0}^{T} \Delta \mathcal{C}_s(h_t) dt = \log \frac{\|\nabla \mathcal{C}_s(h_0)\|}{\|\nabla \mathcal{C}_s(h_T)\|} + \int_{0}^{T} \nabla \cdot \left( \frac{\nabla \mathcal{C}_s(h_t)}{\|\nabla \mathcal{C}_s(h_t)\|} \right) \|\nabla \mathcal{C}_s(h_t)\| dt. $$

The integral in the RHS is a line-integral along the path $h_{[0,T]}$, as $\|\nabla \mathcal{C}_s(h)\|$ is the norm of the flow's "velocity" in $h$. Moreover, $\tau(h) = -\frac{\nabla \mathcal{C}_s(h)}{\|\nabla \mathcal{C}_s(h)\|}$ is the unit tangent vector to the gradient flow in $h$. We can thus write

$$ \int_{0}^{T} \Delta \mathcal{C}_s(h_t) dt = \log \frac{\|\nabla \mathcal{C}_s(h_0)\|}{\|\nabla \mathcal{C}_s(h_T)\|} - \int_{h_{[0:T]}} \nabla \cdot \tau(h) \|\delta h\|, $$

which is (4).

## C Asymptotic behaviour around a local minimum

We consider the case in which the algorithm is approaching a quadratic local optimum of $\mathcal{C}_s$. We will hence say that for $t$ larger than some $\hat{t}$ we have

$$ \mathcal{C}_s(h_t) \approx \mathcal{C}_s(h^*) + \frac{1}{2}(h_t - h^*)^\top H (h_t - h^*), $$

where $H$ is a positive definite matrix, namely the Hessian of $\mathcal{C}_s$ evaluated at $h^*$. Hence, for $t \ge \hat{t}$ the objective's Laplacian is roughly constant,

$$ \Delta \mathcal{C}_s(h_t) \approx \mathrm{Tr}[H]. $$

For a large enough time horizon $T$, we can asymptotically write $\int_0^T \Delta \mathcal{C}_s(h_t) \approx T \mathrm{Tr}[H]$. Moreover, since the Laplacian's integral is diverging while $\frac{\rho_0(h_0)}{\rho_0(h_T)} \to \frac{\rho_0(h^*)}{\rho_0(h_T)}$, asymptotically we find that (1) takes the form

$$ \mathcal{L}_{\mathcal{X}}(w_T) - \mathcal{L}_s(w_T) \lesssim O \left( \sqrt{\frac{T \mathrm{Tr}[H]}{2m}} \right). $$

## D Data-dependent priors

As we discussed in Section 5.6, if we want to use a prior $\rho_{t_0}$, learnt on a dataset $s'$ (independent of $s$), we need to evaluate $\log \rho_{t_0}(h_T)$. Define $\hat{h}_t$ as the solution of $\partial_t \hat{h}_t = -\nabla \mathcal{C}'_s(\hat{h}_t)$ satisfying $\hat{h}_T = h_T$. We have that

$$ \log \rho_{t_0}(h_T) = \log \rho_{t_0}(\hat{h}_T) = \log \rho_0(\hat{h}_{T-t_0} + t) + \int_0^T \partial_t (\log \rho_t(\hat{h}_{T-t_0+t})) dt. $$
M. Hardt, B. Recht, and Y. Singer. Train faster, generalize better: Stability of stochastic gradient descent. *ICML*, 2016.

H. He, G. Huang, and Y. Yuan. Asymmetric valleys: Beyond sharp and flat local minima. *NeurIPS*, 2019.

S. Hochreiter and J. Schmidhuber. Flat minima. *Neural Computation*, 9, 1997.

P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A.G. Wilson. Averaging weights leads to wider optima and better generalization. *UAI*, 2018.

A. Jacot, F. Gabriel, and C. Hongler. Neural tangent kernel: convergence and generalization in neural networks. *NeurIPS*, 2018.

A. Jacot, F. Gabriel, and C. Hongler. The asymptotic spectrum of the hessian of DNN throughout training. In *International Conference on Learning Representations*, 2020.

N.S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P.T.P. Tang. On large-batch training for deep learning: Generalization gap and sharp minima. *ICLR*, 2017.

J. Kiefer and J. Wolfowitz. Stochastic estimation of the maximum of a regression function. *The Annals of Mathematical Statistics*, 23(3), 1952.

J. Langford and M. Seeger. Bounds for averaging classifiers. *CMU technical report*, 2001.

J. Lee, L. Xiao, S.S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. *Journal of Statistical Mechanics*, 2020.

G. Letarte, P. Germain, B. Guedj, and F. Laviolette. Dichotomize and generalize: PAC-Bayesian binary activated deep neural networks. *NeurIPS*, 2019.

G. Lugosi and G. Neu. Generalization bounds via convex analysis. *COLT*, 2022.

X. Luo, L. Bei, and J. Li. Generalization bounds for gradient methods via discrete and continuous prior. *NeurIPS*, 2022.

A. Maurer. A note on the PAC Bayesian theorem. *arXiv:0411099*, 2004.

D.A. McAllester. Some PAC-Bayesian theorems. *COLT*, 1998.

S. Mei and A. Montanari. The generalization error of random features regression: Precise asymptotics and the double descent curve. *Communications on Pure and Applied Mathematics*, 75(4), 2022.

Z. Mhammedi, P. Grünwald, and B. Guedj. PAC-Bayes un-expected Bernstein inequality. *NeurIPS*, 2019.

K. Miyaguchi. PAC-Bayesian transportation bound. *arXiv:1905.13435*, 2019.

A. Montanari and Y. Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. *The Annals of Statistics*, 50(5):2816–2847, 2022.
Since, for $t \in [0, t_0]$, $\rho_t$ is the push-forward of $\rho_0$ under the gradient dynamics with objective $C_{s'}$, we have that for all $h$

$$ \partial_t \rho_t(h) = \rho_t(h) \Delta C_{s'}(h) + \nabla \rho_t(h) \cdot \nabla C_{s'}(h) $$

and so in particular

$$ \partial_t(\rho_t(\hat{h}_t)) = \partial_t \rho_t(\hat{h}_t) + \nabla C_{s'}(\hat{h}_t) \cdot \partial_t \hat{h}_t = \rho_t(\hat{h}_t) \Delta C_{s'}(\hat{h}_t). $$

We conclude that

$$ \log \rho_{t_0}(h_T) = \log \rho_0(\hat{h}_{T-t_0}) + \int_0^T \Delta C_{s'}(\hat{h}_{T-t_0+t})dt, $$

which is (6). We can hence state the following variant of Proposition 2 for data-dependent priors.

**Proposition 3.** Let $s$ and $s'$ be independent training dataset, with $s \sim \mu^m$. Let $C_s$ and $C_{s'}$ be two twice differentiable training objectives, fix $T > 0$ and $t_0 \in [0, T)$. Consider an algorithm that, for an initial state $h_0$, follows the dynamics $\partial_t h_t = -\nabla C_{s'}(h_t)$ for $t \in [0, t_0]$, and $\partial_t h_t = -\nabla C_s(h_t)$ for $t \in [t_0, T]$, and finally outputs $h_T$. Then, for any measurable function $\Psi: \mathbb{R}^2 \to \mathbb{R}$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_X(h_T)) \le \log \frac{\rho_0(h_{t_0})}{\rho_0(\hat{h}_{T-t_0})} + \int_0^{t_0} (\Delta C_{s'}(h_t) - \Delta C_{s'}(\hat{h}_{T-t_0+t}))dt + \int_{t_0}^T \Delta C_s(h_t)dt + \log \frac{\xi}{\delta}, $$

where $\hat{h}_t$ is the solution of $\partial_t \hat{h}_t = -\nabla C_{s'}(\hat{h}_t)$ with $\hat{h}_T = h_T$, and $\xi = \int_{\mathcal{X}^m \times \mathcal{H}} e^{\Psi(\mathcal{L}_s(h), \mathcal{L}_X(h))} d\mu^m(s) d\rho_{t_0}(h)$.

# E Experimental details

## E.1 Linear loss

For our experiments with the linear surrogate loss (see Section 6.1), we used a binary toy dataset, which was created as follows. First, we sampled eight independent points in $\mathbb{R}^5$, from a standard normal distribution. These eight points were then used as the means of eight Gaussian distributions, symmetric and with variance 0.1. We sampled eight Gaussian clusters of 5000 points from each of this eight distributions. We then assigned the label $-1$ to four clusters, and $+1$ to the remaining four. Finally, all the points were projected on the unit sphere. 500 of them, randomly chosen, were used as training dataset, while the remaining 39500 constituted a held-out test dataset.

The linear model that we considered was a neural network with a single-hidden layer, where the hidden nodes were kept fixed at their initial values and only the outer layer was trained. Explicitly, we had $F(x) = h \cdot \Phi(x)$, with $\Phi(x) = \phi(Wx)$. Here, $W$ denotes a random $5 \times N$ matrix, whose components were independently drawn from a centred normal distribution with variance 1/5, and $\phi$ is the ReLU function $\phi(x) = \max(0, x)$. We chose the distribution $\rho_0$ as a multivariate normal distribution, so that all the components of $h_0$ were independent draws from a Gaussian random variable with variance $1/N$.

We tried different values for the dimension $N$ of $h$: 100, 1000, and 10000. We reported in Figure 3 in the main text the results obtained for $N = 10000$. The parameter $\delta$ was fixed at .005, so the bounds hold with probability higher than .995. Since the values reported are already comprehensive of the factor $\log 50$ due to the union bound, it is possible to select the horizon time $T$ yielding the tightest result. Figure 7 compares the values of (2) for the different sizes $N$.

## E.2 Quadratic loss

For the experiments with the quadratic surrogate loss (Section 6.2), we used the same binary dataset and models that we described in Appendix E.1. The parameters $\alpha$ and $\beta$ that define $C_s$ were both set to 1. Figure 4 in the main text reports the results obtained for $N = 10000$. The value reported are already taking into account the $\log 50$ factor (due to the union bound), and $\delta$ was again chosen to be .005. Figure 8 compares the values of the bound (2) for different dimensions $N$.
W. Mou, L. Wang, X. Zhai, and K. Zheng. Generalization bounds of SGLD for non-convex learning: Two theoretical viewpoints. *COLT*, 2018.

V. Nagarajan and J.Z. Kolter. Deterministic PAC-Bayesian generalization bounds for deep networks via generalizing noise-resilience. *ICLR*, 2019.

R.M. Neal. Bayesian learning for neural networks. Springer Science & Business Media, 118, 1995.

J. Negrea, M. Haghifam, G.K. Dziugaite, A. Khisti, and D.M. Roy. Information-Theoretic generalization bounds for SGLD via Data-Dependent estimates. *NeurIPS*, 2019.

G. Neu, G.K. Dziugaite, M. Haghifam, and D.M. Roy. Information-theoretic generalization bounds for stochastic gradient descent. *COLT*, 2021.

B. Neyshabur, S. Bhojanapalli, and N. Srebro. A PAC-Bayesian approach to spectrally-normalized margin bounds for neural networks. *ICLR*, 2018.

G.A. Pavliotis. *Stochastic Processes and Applications: Diffusion Processes, the Fokker-Planck and Langevin Equations*. Springer, 2014.

A. Pensia, V. Jog, and P.L. Loh. Generalization error bounds for noisy, iterative algorithms. *IEEE International Symposium on Information Theory*, 2018.

M. Pérez-Ortiz, O. Risvaplata, J. Shawe-Taylor, and C. Szepesvári. Tighter risk certificates for neural networks. *Journal of Machine Learning Research*, 22, 2021a.

M. Pérez-Ortiz, O. Rivasplata, B. Guedj, M. Gleeson, J. Zhang, J. Shawe-Taylor, M. Bober, and J. Kittler. Learning PAC-Bayes priors for probabilistic neural networks. *arXiv:2109.10304*, 2021b.

N. Qian. On the momentum term in gradient descent learning algorithms. *Neural Networks*, 12(1), 1999.

M. Raginsky, A. Rakhlin, and M. Telgarsky. Non-convex learning via stochastic gradient Langevin dynamics: a nonasymptotic analysis. *COLT*, 2017.

A. Rahimi and B. Recht. Random features for large-scale kernel machines. *NeurIPS*, 2007.

A. Raj, L. Zhu, M. Gürbüzbalaban, and U. Şimşekli. Algorithmic stability of Heavy-Tailed SGD with general loss functions. *arXiv:2301.11885*, 2023.

O. Rivasplata, I. Kuzborskij, C. Szepesvári, and J. Shawe-Taylor. PAC-Bayes analysis beyond the usual bounds. *NeurIPS*, 2020.

D. Russo and J. Zou. How much does your data exploration overfit? Controlling bias via information usage. *IEEE Transactions on Information Theory*, 66(1), 2019.

S. Shalev-Shwartz and S. Ben-David. *Understanding Machine Learning - From Theory to Algorithms*. Cambridge University Press, 2014.
<figure><img src="image_8.png" /><figcaption>**Figure 7:** For the experiment with the linear surrogate loss, comparison of the evolution of the bound (2) for different sizes *N* of the model. See Section 6.1 and Appendix E.1 for more details.</figcaption></figure>

<figure><img src="image_9.png" /><figcaption>**Figure 8:** For the experiment with the quadratic surrogate loss, comparison of the evolution of the bound (2) for different sizes *N* of the model. See Section 6.2 and Appendix E.2 for more details.</figcaption></figure>

## E.3 Cross-entropy loss

We performed our experiments with the cross-entropy as surrogate loss (see Section 6.3) on the MNIST dataset, with the standard split of 60000 images used for the training and 10000 for the testing of the model. The training was performed on batches of 500 elements. The images were flattened to one-dimensional arrays and renormalised. We used again linear architectures $F(x) = h\Phi(x)$, but now with $h$ being a $10 \times w$ matrix ($w$ denoting the width of the network). $\Phi$ was always in the form $\Phi(x) = \phi(Wx)$. Here $\phi$ acts component-wise as the ReLU function, and $W$ is a $w \times 784$ matrix (784 being the dimension of the input datapoints), whose components were independently drawn from a centred normal distribution with variance $1/784$. The dimension of the hypothesis $h$ was $N = 10w$, as the networks had 10 output nodes. We tested different values of $N$, as reported in Table 1. In all the bounds $\delta = .005$.

Since closed-form solutions for the flow are not available, we used discretised gradient dynamics. For each batch $s_k$ we made an update $h_{t+\delta_t} = h_t + \delta_t \nabla C_{s_k}(h_t)$, for some small $\delta_t$. Recall that, when training by batches, the Laplacian's integral takes the form (5). Here we chose that $t_k = k\delta_t$, so that every time we switched to a new batch after a single time-step. In particular, we get $\int_{t_{k-1}}^{t_k} \Delta C_{s_k}(h_t)dt = \Delta C_{s_k}(h_{t-1})\delta_t$ in (5). Table 2 reports the different values of time discretisations $\delta_t$, used for the models that we trained.

Figure 2 in the main text shows the relative error between the values of (2) obtained for the same model ($N = 10^5$) for two different time discretisations ($\delta_{t,1} = 2 \times 10^{-4}$ and $\delta_{t,2} = 2 \times 10^{-6}$). In both cases we trained for a time-interval of $2 \times 10^{-4}$ before switching batch. This means that for $\delta_{t,1}$ each batch was used for a single time-step, before switching to the next one, while for $\delta_{t,2}$ 100 gradient-steps were performed on each batch before passing to the next one.

The values of (2) that we obtained for the different dimensions $N$ are reported in Table 1, in the main text. The penalty log 25 is already taken into account for the values reported, as for each width the best bound has been chosen among the values obtained with 25 different time horizons. The optimal $T$ reported in Figure 1 are the values of the best time horizon (among the 25 ones that we tried) that we obtained for each model's width. Figures 5 and 6 refer to the model with size $N = 10^5$.

<table><caption>**Table 2:** Time discretisation</caption><thead><tr><td>Dimension N</td><td>Final Time</td><td>δ<sub>t</sub></td></tr></thead><tbody><tr><td>10<sup>3</sup></td><td>60</td><td>2.00 × 10<sup>-2</sup></td></tr><tr><td>3 × 10<sup>3</sup></td><td>20</td><td>6.67 × 10<sup>-3</sup></td></tr><tr><td>10<sup>4</sup></td><td>6</td><td>2.00 × 10<sup>-3</sup></td></tr><tr><td>3 × 10<sup>4</sup></td><td>2</td><td>6.67 × 10<sup>-4</sup></td></tr><tr><td>10<sup>5</sup></td><td>.6</td><td>2.00 × 10<sup>-4</sup></td></tr><tr><td>3 × 10<sup>5</sup></td><td>.2</td><td>6.67 × 10<sup>-5</sup></td></tr><tr><td>10<sup>6</sup></td><td>.06</td><td>2.00 × 10<sup>-5</sup></td></tr></tbody></table>
J. Shawe-Taylor, C.K.I. Williams, N. Cristianini, and J. Kandola. On the eigenspectrum of the gram matrix and the generalization error of kernel-pca. *IEEE Transactions on Information Theory*, 51 (7), 2005.

Y. Song, J. Sohl-Dickstein, D.P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. *ICLR*, 2021.

I.O. Tolstikhin and Y. Seldin. PAC-Bayes-empirical-Bernstein inequality. *NeurIPS*, 2013.

V.N. Vapnik. *The Nature of Statistical Learning Theory*. Springer, 2000.

P. Viallard, P. Germain, A. Habrard, and E. Morvant. A general framework for the disintegration of PAC-Bayesian bounds. *arXiv:2102.08649*, 2021.

A. Xu and M. Raginsky. Information-theoretic analysis of generalization capability of learning algorithms. *NeurIPS*, 2017.

G. Yang. Tensor programs II: Neural tangent kernel for any architecture. *arXiv:2006.14548*, June 2020.

G. Yang and E. Littwin. Tensor programs IIb: Architectural universality of neural tangent kernel training dynamics. *ICML*, 2021.

G. Yang and H. Salman. A fine-grained spectral perspective on neural networks. *ICLR 2020*, 2020.

V. Zantedeschi, P. Viallard, E. Morvant, R. Emonet, A. Habrard, P. Germain, and B. Guedj. Learning stochastic majority votes by minimizing a PAC-Bayes generalization bound. *NeurIPS*, 2021.

C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understanding deep learning requires rethinking generalization. *Communications of the ACM*, 64, 11 2016.

W. Zhou, V. Veitch, M. Austern, R.P. Adams, and P. Orbanz. Non-vacuous generalization bounds at the ImageNet scale: a PAC-Bayesian compression approach. *ICLR*, 2019.
## Appendix A. Omitted proofs of Sections 3 and 4

**Corollary 2** Assume that $\ell(h, \cdot)$ is $R$-sub-Gaussian for each $h \in \mathcal{H}$. Then, for any $\delta \in (0, 1)$ and $T > 0$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_Z(h_T) \le \mathcal{L}_S(h_T) + \frac{1}{\sqrt{m}} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_S(h_t) dt + \log \frac{1}{\delta} + \frac{R^2}{2} \right). $$

**Proof** The results follows from Theorem 1 with $\Psi(u, v) = (v - u)/\sqrt{m}$. In this case, we have that for each fixed $h$

$$ \int_{\mathcal{X}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\mu^m(\bar{s}) = \left( \int_{\mathcal{X}} e^{(\mathcal{L}_Z(h) - \ell(h, z))/\sqrt{m}} d\mu(z) \right)^m \le e^{R^2/2}, $$

by the definition of sub-Gaussianity. So, $\xi \le R^2/2$, and the desired bound follows immediately. ■

**Corollary 3** Assume that $\ell$ is bounded in $[0, 1]$ and $m \ge 8$. Then, for any fixed $\delta \in (0, 1)$ and $T > 0$, with a probability of at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_Z(h_T) \le \mathrm{kl}^{-1} \left( \mathcal{L}_S(h_T) \left| \frac{B_T}{m} \right. \right); \quad B_T = \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_S(h_t) dt + \log \frac{2\sqrt{m}}{\delta}, $$

where we define $\mathrm{kl}^{-1}(u|c) = \sup\{v \in [0, 1] : \mathrm{kl}(u||v) \le c\}$. In particular, it follows that

$$ \mathcal{L}_Z(h_T) \le \mathcal{L}_S(h_T) + \sqrt{\frac{\mathcal{L}_S(h_T) B_T}{2m}} + \frac{B_T}{2m}. $$

**Proof** The bound follows from Theorem 1 with $\Psi(u, v) = \mathrm{kl}(u||v)$, after using the fact that with this choice one has $\xi \le 2\sqrt{m}$ if the loss is bounded in $[0, 1]$ and $m \ge 8$ (see, e.g., Theorem 1 in Maurer, 2004). For the second bound, we use the property,

$$ \mathrm{kl}^{-1}(u|c) \le \min\{u + \sqrt{uc/2} + c/2, u + \sqrt{c/2}\}, $$

(see, for example, Tolstikhin and Seldin 2013 and the discussion and references therein). ■

**Theorem 4** For any dataset $s \in \mathcal{Z}^m$, assume there is a Borel set $A_s \subseteq \mathcal{H}$ on which $\mathcal{C}_s$ is twice-differentiable and $M$-smooth, where $\mathrm{sup}_k \eta_k \le 1/(2M)$. Let $\Psi : \mathbb{R}^2 \to \mathbb{R}$ be measurable. Taking $\delta \in (0, 1)$ and $K \in \mathbb{N}$ fixed, suppose that $\{h_k\}_{k=0}^{K-1}$ lies in $A_s$ with a probability of at least $1 - \delta/2$ under $(s, h_0) \sim \mu^m \otimes \rho_0$. Then, with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$, it holds that

$$ \Psi(\mathcal{L}_S(h_T), \mathcal{L}_Z(h_T)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_k)} - \sum_{k=0}^{K-1} \mathrm{tr} \log \left( \mathrm{Id} - \eta_k \nabla^2 \mathcal{C}_S(h_k) \right) + \log \frac{2\xi}{\delta} + \delta \right), $$

where $\xi = \int_{\mathcal{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\mu^m(\bar{s})d\rho_0(h)$. ■
**Proof** First, for any $s \in \mathbb{Z}^m$, we introduce the notation $G_{\eta}^s$ for the mapping (from $\mathcal{H}$ to $\mathcal{H}$) $h \mapsto h - \eta C_s(h)$. Then, define

$$A_s^0 = \{h_0 \in \mathcal{H} : h_k \in A_s \text{ for all } k \in \{0, \dots, K-1\}\},$$

which is a Borel set thanks to the regularity of the $G_{\eta_k}^s$ on $A_s$.

We start by noticing that for all $k$, the restriction of $G_{\eta_k}^s$ to $A_s$ is injective. Indeed, if $h$ and $h'$ are in $A_s$, we have that

$$\|G_{\eta_k}^s(h) - G_{\eta_k}^s(h')\| \ge \|h - h'\| - \eta_k \|\nabla C_s(h) - \nabla C_s(h')\| \ge \frac{1}{2} \|h - h'\|.$$

For any fixed $s$, if we condition on $h_0 \in A_s^0$, this condition is satisfied for all $k \in \{0, \dots, K-1\}$. Now let $\tilde{\rho}_k^s = \rho_k(\cdot|h_0 \in A_s^0)$ be the law of $h_k$, conditioned on $h_0 \in A_s^0$. If we denote as $\tilde{G}_{s,k}$ the restriction of $G_{\eta_k}^s$ to supp $\tilde{\rho}_k^s$, we have that $\tilde{G}_{s,k}$ is a differentiable bijection $\text{supp} \tilde{\rho}_k^s \leftrightarrow \text{supp} \tilde{\rho}_{k+1}^s$. In particular, we see by induction that this implies that $\tilde{\rho}_{k+1}^s$ admits a Lebesgue density (since $\tilde{\rho}_0^s \ll \rho_0$), and by the change of variable formula

$$\tilde{\rho}_{k+1}^s(h) = \tilde{\rho}_k^s \circ \tilde{G}_{s,k}^{-1}(h) \det \left[ \text{Id} - \eta_k \nabla^2 C_s \circ \tilde{G}_{s,k}^{-1}(h) \right]^{-1}, \quad \forall h \in \text{supp} \tilde{\rho}_{k+1}^s.$$

For $h_k \in \text{supp} \tilde{\rho}_k^s$, since $\tilde{G}_{s,k}(h_k) = h_{k+1}$, we get

$$\tilde{\rho}_{k+1}^s(h_{k+1}) = \tilde{\rho}_k^s(h_k) \det \left[ \text{Id} - \eta_k \nabla^2 C_s(h_k) \right]^{-1}.$$

In particular, for $h_0 \in A_s^0$ we have that $h_k \in \text{supp} \tilde{\rho}_k^s$ by definition of $\tilde{\rho}_k^s$, and so

$$\begin{aligned} \log \frac{\tilde{\rho}_K^s(h_K)}{\tilde{\rho}_0^s(h_0)} &= \sum_{k=0}^{K-1} \log \frac{\tilde{\rho}_{k+1}^s(h_{k+1})}{\tilde{\rho}_k^s(h_k)} \\ &= - \sum_{k=0}^{K-1} \log \det \left[ \text{Id} - \eta_k \nabla^2 C_s(h_k) \right] = - \sum_{k=0}^{K-1} \text{tr} \log \left[ \text{Id} - \eta_k \nabla^2 C_s(h_k) \right]. \end{aligned} \quad (9)$$

where the last equality follows from the Jacobi formula for positive definite matrices, namely $\log \det = \text{tr} \log$.

We now use the same Markov argument as in Theorem 1, with posterior $\tilde{\rho}_K^s$ and prior $\rho_0$. Explicitly, we have that with probability at least $1 - \frac{\delta}{2}$ on $(s, h_K) \sim \mu^m * \tilde{\rho}_K$

$$e^{m\Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) - \log \frac{\tilde{\rho}_K^s(h_K)}{\rho_0(h_K)}} \le \frac{2}{\delta} \int_{\mathcal{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h)) - \log \frac{\tilde{\rho}_K^s(h)}{\rho_0(h)}} d\mu^m(\bar{s}) d\tilde{\rho}_K^s(h).$$

Since for every $\bar{s} \in \mathbb{Z}^m$ we have

$$\begin{aligned} & \int_{\mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h)) - \log \frac{\tilde{\rho}_K^s(h)}{\rho_0(h)}} d\tilde{\rho}_K^s(h) \\ &= \int_{\{\tilde{\rho}_k > 0\}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h)) - \log \frac{\tilde{\rho}_K^s(h)}{\rho_0(h)}} d\tilde{\rho}_K^s(h) \le \int_{\mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\rho_0(h), \end{aligned}$$
we get that

$$ \mu^m * \rho_K \left( m\Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \le \log \frac{\tilde{\rho}_K^s(h_K)}{\rho_0(h_K)} + \log \frac{2\xi}{\delta} \middle| h_0 \in A_s^0 \right) \ge 1 - \delta/2. $$

Now, we note that for any $h_0 \in A_s^0$, the following holds:

$$ \log \frac{\tilde{\rho}_K^s(h_K)}{\rho_0(h_K)} = \log \frac{\tilde{\rho}_K^s(h_K)}{\tilde{\rho}_0^s(h_0)} + \log \frac{\rho_0^s(h_0)}{\rho_0(h_K)} - \log \rho_0(A_s^0), $$

which is further bounded noticing that $-\log(1 - \delta/2) \le \delta$ as $\delta \in (0, 1)$. In particular, using the change of density formula (9) we get that

$$ \begin{aligned} & \mu^m * \rho_K \left( m\Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \right. \\ & \le \log \frac{\rho_0(h_0)}{\rho_0(h_K)} - \sum_{k=0}^{K-1} \operatorname{tr} \log \left[ \operatorname{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k) \right] + \log \frac{2\xi}{\delta} + \delta \middle| h_0 \in A_s^0 \right) \ge 1 - \delta/2. \end{aligned} \quad (10) $$

Now, note that for any event E we have

$$ \begin{aligned} \mu^m * \rho_K(E) &= \int_{\mathbb{Z}^m} (\rho_K(E|h_0 \in A_s^0)\rho_0(A_s^0) + \mu^m * \rho_K(E|h_0 \in A_s^0)(1 - \rho_0(A_s^0))) \, d\mu^m(s) \\ &\le \mu^m * \rho_K(E|h_0 \in A_s^0) + \frac{\delta}{2}, \end{aligned} $$

since $\mu^m \otimes \rho_0(h_0 \notin A_s^0) \le \delta/2$ by hypothesis. Applying this to (10), we get that

$$ m\Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \le \log \frac{\rho_0(h_0)}{\rho_0(h_K)} - \sum_{k=0}^{K-1} \operatorname{tr} \log \left[ \operatorname{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k) \right] + \log \frac{2\xi}{\delta} + \delta, \quad (11) $$

with probability at least $1 - \delta$ on $(s, h_K) \sim \mu^m * \rho_K$. Since sampling from $\rho_K$ for a given $s$ is equivalent to sample from $\rho_0$ and follow the dynamics until the K-th step, we can claim that (11) holds with probability at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$. ■

**Lemma 5** With the notation of Theorem 4, let $h_k \in A_s$. Then

$$ - \operatorname{tr} \log \left( \operatorname{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k) \right) \le \eta_k \Delta \mathcal{C}_s(h_k) + \eta_k^2 \| \nabla^2 \mathcal{C}_s(h_k) \|_{\mathrm{F}}^2 \le \frac{3}{2} \eta_k \| \nabla^2 \mathcal{C}_s(h_k) \|_{\mathrm{TR}}, $$

where $\| \cdot \|_{\mathrm{F}}$ is the Frobenius norm and $\| \cdot \|_{\mathrm{TR}}$ the trace norm.

**Proof** To obtain the first upper bound, denote as $\{\lambda_i(h_k)\}_{i=1}^d$ the spectrum of $\nabla^2 \mathcal{C}_s(h_k)$. Then we have that

$$ - \operatorname{tr} \log \left( \operatorname{Id} - \eta_k \nabla^2 \mathcal{C}_s(h_k) \right) = - \sum_{k=0}^{K-1} \sum_{i=1}^{d} \log(1 - \eta_k \lambda_i(h_k)). $$
Using that $-\log(1-u) \le u(u+1)$ for $|u| \le 1/2$, we obtain that for each $k$

$$ -\sum_{i=1}^{d} \log(1 - \eta_k \lambda_i(h_k)) \le \eta_k \sum_{i=1}^{d} \lambda_i(h_k) + \eta_k^2 \sum_{i=1}^{d} \lambda_i(h_k)^2 = \eta_k \Delta C_s(h_k) + \eta_k^2 \|\nabla^2 C_s(h_k)\|_F^2. $$

For the second upper bound, note that $\Delta C_s(h_k) \le \|\nabla^2 C_s(h_k)\|_{TR}$ and

$$ \eta_k^2 \|\nabla^2 C_s(h_k)\|_F^2 = \eta_k^2 \sum_{k=1}^{d} \lambda_i(h_k)^2 \le \eta_k \sum_{k=1}^{d} |\eta_k \lambda_i(h_k)| |\lambda_i(h_k)| = \frac{1}{2} \eta_k \|\nabla^2 C_s(h_k)\|_{TR}, $$

where we used that $|\eta_k \lambda_i(h_k)| \le \eta_k \|\nabla^2 C_s(h_k)\| \le 1/2$ since $C_s$ is $1/(2\eta_k)$-smooth in $h_k \in A_s$. ■

## Appendix B. A few more explicit bounds

In this section we leverage and extend the approach of Casado et al. (2024) (we show here that it is possible to drop the absolute continuity assumption for the loss required therein) to give a possible way to choose $\Psi$ (in Theorems 1 and 4) so that the term $\xi$ appearing the the bounds can be explicitly controlled. For the sake of conciseness we will state everything in the continuous time framework, as the discrete time counterpart of these results will naturally follow replacing the use of Theorem 1 with Theorem 4.

For $b \in (0, +\infty]$, let $\psi: [0, b) \to [0, +\infty)$ be a convex, continuously differentiable function, such that $\psi(0) = \psi'(0) = 0$. We define

$$ B = \lim_{\lambda \to b} \frac{\psi(\lambda)}{\lambda} $$

and

$$ A = \begin{cases} \lim_{\lambda \to b} (\lambda B - \psi(\lambda)) & \text{if } B < +\infty; \\ +\infty & \text{otherwise.} \end{cases} $$

For $u \in \mathbb{R}$, we define the Fenchel conjugate of $\psi$ as

$$ \psi^*(u) = \sup_{\lambda \in [0, b)} (\lambda u - \psi(\lambda)). $$

The following lemma summarises a few properties of $\psi^*$. Its claims are either immediate to check or classical results for convex functions (see, for instance, Chapter 2 in Boucheron et al. 2013).

**Lemma 10** $\psi^*$ satisfies the following properties:

* $\psi^*(u) = 0$ if $u \le 0$;

* $\psi^*(u) = +\infty$ if $u > B$ and $b = +\infty$;

* $\psi^*(u) = A + b(u - B)$ if $u > B$ and $b < +\infty$;

* the restriction of $\psi^*$ to $[0, B]$ is a continuously differentiable bijection $[0, B] \leftrightarrow [0, A]$.
Moreover, we have that for $v \in [0, A]$

$$ \psi^{*-1}(v) = \inf_{\lambda \in (0, b)} \frac{v + \psi(\lambda)}{\lambda} \in [0, B] $$

and, if $b < +\infty$, $\psi^*$ is invertible on $[0, +\infty)$ and for $v > A$

$$ \psi^{*-1}(v) = \frac{v - A}{b} + B. $$

Given a loss function $\ell$, we will say that $\ell$ has sub-$\psi$ concentration if, for all $h \in \mathcal{H}$, it satisfies

$$ \log \int e^{\lambda(\mathcal{L}_Z(h) - \ell(h, z))} d\mu(z) \le \psi(\lambda), $$

for all $\lambda \in [0, b)$. As an explicit example, we note that for $\psi(\lambda) = R^2\lambda^2/2$ defined on $[0, \infty)$, the sub-$\psi$ concentration follows from R-sub-Gaussianity.

**Lemma 11** If $\ell$ has sub-$\psi$ concentration, then for each $h$ we have that for any $u \in \mathbb{R}$

$$ \mu^m(\mathcal{L}_Z(h) \ge \mathcal{L}_s(h) + u) \le e^{-m\psi^*(u)}. $$

**Proof** For $u < 0$ the result is trivial, as the RHS is 1. For $u \ge 0$, it follows from a simple application of Markov's inequality, since we have that

$$ \begin{aligned} \mu^m(\mathcal{L}_Z(h) \ge \mathcal{L}_s(h) + u) &= \mu^m \left( \sum_{z \in s} (\mathcal{L}_Z(h) - \ell(h, z)) \ge m u \right) \\ &= \inf_{\lambda \in [0, b)} \mu^m \left( e^{\lambda \sum_{z \in s} (\mathcal{L}_Z(h) - \ell(h, z))} \ge e^{m \lambda u} \right) \\ &\le \inf_{\lambda \in [0, b)} e^{-m \lambda u} \prod_{z \in s} \int_{\mathcal{Z}} e^{\lambda (\mathcal{L}_Z(h) - \ell(z, h))} d\mu(z) = e^{-m \psi^*(u)}, \end{aligned} $$

which is what we wanted.

<figure><img src="image_3.png" /></figure>

**Corollary 12** If $\ell$ has sub-$\psi$ concentration, then for each $h$, for any $v \in \mathbb{R}$

$$ \mu^m(m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \ge v) \le \min\{1, e^{-v}\}. $$

**Proof** First, note that the result if trivially true if $v \le 0$, as the RHS is 1. For $v > 0$, first notice that if $b < +\infty$ then $\psi^*$ is non-decreasing and invertible on $[0, +\infty)$ by Lemma 10, and so

$$ \mu^m(m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \ge v) = \mu^m(\mathcal{L}_Z(h) - \mathcal{L}_s(h) \ge \psi^{*-1}(v/m)). $$

In this case the conclusion follows immediately by Lemma 11.

On the other hand, if $b = +\infty$, for $v > Am$, we have that $m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \ge v$ only if $\mathcal{L}_Z(h) - \mathcal{L}_s(h) > B$. However, by Lemma 11 this is a zero-probability event. Indeed, we have that, for any $\varepsilon > 0$, $\mu^m(\mathcal{L}_Z(h) \ge \mathcal{L}_s(h) + B + \varepsilon) \le e^{-m\psi^*(B+\varepsilon)} = 0$, since $\psi^*(B+\varepsilon) = +\infty$ by Lemma 10. So, for $v > Am$ we have

$$ \mu^m(m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \ge v) = 0 \le e^{-v}. $$
We are only left with the case $v \in [0, Am]$ and $b = +\infty$. In such a case, the invertibility of $\psi^*$ on $[0, B]$ is enough and we get that (almost surely)

$$ \{m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \geq v\} = \{\mathcal{L}_Z(h) - \mathcal{L}_s(h) \geq \psi^{*-1}(v/m)\}, $$

and so again

$$ \mu^m(m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \geq v) = \mu^m(\mathcal{L}_Z(h) - \mathcal{L}_s(h) \geq \psi^{*-1}(v/m)). $$

Again, we conclude by Lemma 11, since $\psi^*$ is a bijection $[0, B] \leftrightarrow [0, A]$ by Lemma 10. ■

**Lemma 13** If $\ell$ has sub-$\psi$ concentration, for any $m' < m$, for any fixed $h$, we have

$$ \int_{\mathcal{Z}^m} e^{m'\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h))} d\mu^m(s) \leq \frac{m}{m - m'}. $$

**Proof** We have

$$ \begin{aligned} \int_{\mathcal{Z}^m} e^{m'\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h))} d\mu^m(s) &= \int_0^\infty \mu^m \left( e^{m'\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h))} \ge u \right) du \\ &= \int_0^\infty \mu^m \left( m\psi^*(\mathcal{L}_Z(h) - \mathcal{L}_s(h)) \ge \frac{m}{m'} \log u \right) du \le \int_0^\infty \min(1, e^{-\frac{m}{m'} \log u}) du \\ &= 1 + \int_1^\infty e^{-\frac{m}{m'} \log u} du = \frac{m}{m - m'}, \end{aligned} $$

where we applied Lemma 12 to obtain the inequality. ■

We can now state the following generalisation bound, which follows from Theorem 1.

**Proposition 14** Consider the dynamics $\partial_t h_t = -\nabla C_s(h_t)$, with $C_s: \mathcal{H} \to \mathbb{R}$ twice differentiable. Assume that $\ell$ has sub-$\psi$ concentration. Fix $T > 0$ and $\delta \in (0, 1)$. We have that, with probability at least than $1 - \delta$ on the draw of the pair $(s, h)$ from $\mu^m \otimes \rho_0$,

$$ \mathcal{L}_Z(h_T) \le \mathcal{L}_s(h_t) + \inf_{\lambda \in (0,b)} \left( \frac{1}{\lambda(m-1)} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt + \log \frac{m}{\delta} \right) + \frac{\psi(\lambda)}{\lambda} \right). \quad (12) $$

**Proof** We apply Theorem 1 with the choice

$$ \Psi(u, v) = \frac{m-1}{m} \psi^*(v-u). $$

This leads to the high probability bound

$$ \psi^*(\mathcal{L}_Z(h_T) - \mathcal{L}_s(h_T)) \le \frac{1}{m-1} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt + \log \frac{\xi}{\delta} \right), $$

with

$$ \xi = \int_{\mathcal{Z}^m \times \mathcal{H}} e^{(m-1)\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\mu^m(\bar{s}) d\rho_0(h). $$
Applying Corollary 12 with $m' = m - 1$ we get that $\xi \le \log m$, and the conclusion follows from inverting $\psi^*$ as in Lemma 10. ■

As a consequence of the result above, we can get a bound for sub-exponential losses. In particular, we recall that $\ell(h, \cdot)$ is $(R, \alpha)$-sub-exponential if for any $\lambda \in (-\alpha^{-1}, \alpha^{-1})$

$$ \log \int_{\mathcal{Z}} e^{\lambda \ell(h,z)} d\mu(z) \le \frac{R^2 \lambda^2}{2}. $$

**Corollary 15** Assume that, for each $h \in \mathcal{H}$, $\ell(h, \cdot)$ is $(R, \alpha)$-sub-exponential. Then, for any $\delta \in (0, 1)$ and $T > 0$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_{\mathcal{Z}}(h_T) \le \mathcal{L}_s(h_T) + \inf_{\lambda \in (0, \alpha^{-1})} \left( \frac{1}{\lambda(m-1)} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_s(h_t) dt + \log \frac{m}{\delta} \right) + \frac{R^2 \lambda}{2} \right). $$

**Proof** We have that $\ell$ has $\psi$-concentration, with $\psi(\lambda) = R^2\lambda^2/2$ and $b = 1/\alpha$. Hence, the result follows directly from Proposition 14. ■

**Corollary 16** Assume that, for each $h \in \mathcal{H}$, $\ell(h, \cdot)$ is $R$-sub-Gaussian. Then, for any $\delta \in (0, 1)$ and $T > 0$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \mathcal{L}_{\mathcal{Z}}(h_T) \le \mathcal{L}_s(h_T) + R \sqrt{\frac{2 \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_s(h_t) dt + \log \frac{m}{\delta} \right)}{(m-1)}}. $$

**Proof** As $\ell(h, \cdot)$ has $\psi$-concentration, with $\psi(\lambda) = R^2\lambda^2/2$ and $b = \infty$, we can apply Proposition 14. Setting

$$ \lambda = \sqrt{\frac{2 \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta \mathcal{C}_s(h_t) dt + \log \frac{m}{\delta} \right)}{R^2(m-1)}} $$

in (12) we obtain the desired result. ■

We note that in the $R$-sub-Gaussian case, setting $\lambda = 1$ in (12) leads to the bound of Corollary 2, up to replacing $m$ with $m - 1$ and to neglect a term $\log m$. As a consequence, for large $m$ one can expect the bound of Corollary 16 to be tighter than the one in Corollary 2.

# Appendix C. Examples

## C.1. Random feature model

We recall from Section 5.1 that $\mathcal{X} = S^{p-1}$ and $\mathcal{Y} = [0, 1]$. The goal is to learn a target function $f: \mathcal{X} \to \mathcal{Y}$, with $F_h(x) = \frac{1}{\sqrt{d}}h \cdot \Phi(x)$ and $f_h = \min\{0, \max\{1, F_h\}\}$. We consider a quadratic loss $\ell(h, z) = (f_h(x) - y)^2$ and an optimisation objective $\mathcal{C}_s(h) = \frac{1}{m} \sum_{z \in s} \hat{\ell}(F_h(x), y)$, with $\hat{\ell}(F, y) = (F - y)^2$. The features are given by $\Phi(x) = \phi(Wx)$, where $\phi$ is a non-linearity acting component-wise, and $W$ is a $d \times p$ matrix whose components are independently drawn from a standard Gaussian distribution.
**Proposition 6** For the random feature model described above, we have the limit in probability

$$ \lim_{d \to \infty} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} + \int_0^T \Delta C_s(h_t) dt \right) \le T \left( \mathbb{E}_{\zeta \sim \mathcal{N}(0,1)} [\phi(\zeta)^2] + 2\sqrt{C_s(h_0)} \right). $$

**Proof** A simple calculation shows that $\Delta C_s(h) = \frac{2}{m} \sum_{z \in s} \frac{1}{d} \|\Phi(x)\|^2$. As $d \to \infty$ we get

$$ \frac{1}{d} \|\Phi(x)\|^2 = \frac{1}{d} \sum_{i=1}^{d} \phi(W_i \cdot x)^2 \to \mathbb{E}[\phi(W_1 \cdot x)^2] = \mathbb{E}_{\zeta \sim \mathcal{N}(0,1)} [\phi(\zeta)^2], $$

where $W_i$ denotes the $i$-th raw of $W$. Hence,

$$ \int_0^T \Delta C_s(h_t) dt = T \mathbb{E}_{\zeta \sim \mathcal{N}(0,1)} [\phi(\zeta)^2]. $$

On the other hand, recalling that under $\rho_0$ all the components of $h$ are independently distributed as standard Gaussians, we get

$$ \log \frac{\rho_0(h_0)}{\rho_0(h_T)} = \frac{1}{2} (\|h_T\|^2 - \|h_0\|^2) = \int_0^T h_t \cdot \partial_t h_t dt = - \int_0^T h_t \cdot \nabla C_s(h_t) dt. $$

Our choice of quadratic learning objective yields

$$ \begin{aligned} -h_t \cdot \nabla C_s(h_t) &= \frac{2}{m} \sum_{z \in s} (y - F_{h_t}(x)) \frac{\Phi(x) \cdot h_t}{\sqrt{d}} = \frac{2}{m} \sum_{z \in s} (y - F_{h_t}(x)) F_{h_t}(x) \\ &= -2C_s(h_t) + \frac{2}{m} \sum_{z \in s} y(y - F_{h_t}(x)) \le \frac{2}{m} \sum_{z \in s} |F_{h_t}(x) - y| \le 2\sqrt{C_s(h_t)}, \end{aligned} $$

where we used that $|y| = |f(x)| \le 1$ and Cauchy-Schwarz inequality. The training dynamics ensure that $C_s(h_t) \le C_s(h_0)$. ■

**Proposition 7** Let $f: S^{d-1} \to [0, 1]$ be a linear combination of spherical harmonics of degree at most $D$. Let $\phi$ be given by $\phi(x) = \sum_{k=1}^J (x/e)^k$, with $J \ge D$. Then, in the overparameterised regime $d \to \infty$, there is a $\lambda > 0$, independent of $m$ and on the training dataset $s$, such that for $T = \frac{\log m}{4\lambda}$ we have $\mathcal{L}_Z(h_T) \le O(\frac{\log m}{m})$, asymptotically for large $m$.

**Proof** Here we assume that $d$ is large enough with respect to $m$, so that, up to negligible correction we can approximate $\frac{1}{d}\Phi(x) \cdot \Phi(x')$ by its overparameterised limit ($d \to \infty$), given by (Daniely et al., 2017)

$$ \Theta(x, x') = \mathbb{E}_{(\xi, \xi') \sim \mathcal{N}(0, \Sigma)} [\phi(\xi) \phi(\xi')]. $$

By Lemma 17, we have that

$$ \Theta(x, x') = \sum_{j=1}^{J} \sum_{l=1}^{N_j} \lambda_j \psi_{l,j}(x) \psi_{l,j}(x'), $$
where $\psi_{l,j}$'s denote the spherical harmonics of degree $l$, and $N_l$ is the number of spherical harmonics of degree $l$.

Let us define the $m \times m$ renormalised Gram matrix $\hat{\Theta} = \frac{1}{m}(\Theta(x, x'))_{(z,z') \in s \times s}$. If we denote as $V_{l,j}$ the $m$-vector $(\psi_{l,j}(x))_{x \in s}$, we have that

$$ \hat{\Theta} = \frac{1}{m} \sum_{j=1}^{J} \sum_{l=1}^{N_j} \lambda_j V_{l,j} V_{l,j}^{\top}. $$

In particular, $\hat{\Theta}$ has rank at most $\sum_{j=1}^{J} N_j$, and its eigendecomposition can be written as

$$ \hat{\Theta} = \sum_{j=1}^{J} \sum_{l=1}^{N_j} \hat{\lambda}_{l,j} \hat{V}_{l,j} \hat{V}_{l,j}^{\top} $$

for some orthonormal vectors $\hat{V}$'s. It is known that for large $m$, the eigenvalues of the Gram matrix tend to those of the kernel. More precisely, from Shawe-Taylor et al. (2005), we have that with probability higher than $1 - \tilde{\delta}$, uniformly on $j$ and $l$

$$ \hat{\lambda}_{l,j} \ge \lambda_j - O\left(\sqrt{\frac{\log \sum_{j=1}^{J} N_j - \log \tilde{\delta}}{m}}\right). $$

We can choose a $\lambda > 0$ such that $\lambda_j \ge 2\lambda$, for all $j$. Then, fixed $\tilde{\delta}$, if $m$ is large enough we have that for all $j$ and all $l$, $\lambda_{l,j} \ge \lambda > 0$ with high probability. In particular, when this is true, we have that $\hat{\Theta}$ has rank $\sum_{j=1}^{J} N_j$, and so all the $V_{l,j}$ are independent.

Let $\Delta_t$ be the $m$-vector $(F_{h_t}(x) - y)_{z \in s}$. The dynamics now reads

$$ \partial_t \Delta_t = -2 \hat{\Theta} \Delta_t, $$

and so we get

$$ \Delta_t = e^{-2t\hat{\Theta}} \Delta_0. $$

Now, let us show that $\Delta_0$ lies in the range of $\hat{\Theta}$. Indeed, we have that the network's output at initialisation behaves as a Gaussian process (Neal, 1995). More precisely (Grenander, 1950), we have that

$$ F_{h_0}(x) = \sum_{j=1}^{J} \sum_{l=1}^{N_j} \sqrt{\lambda_j \zeta_{l,j} \psi_{l,j}(x)}, $$

where the $\zeta$'s are i.i.d. draws from a standard Gaussian. In particular, we know that $F_{h_0}(x)$ lies in the span of the spherical harmonics of degree at most $J$. Since we have that $f$ also lies in this span by assumption, we can write $\Delta_0$ as a linear combination of the $V_{l,j}$'s. Since the $V_{l,j}$'s are linearly independent, we have that each of them is orthogonal to the ker of $\hat{\Theta}$. Indeed, let $w \in \ker \hat{\Theta}$. We have

$$ 0 = \hat{\Theta} w = \frac{1}{m} \sum_{j=1}^{J} \sum_{l=1}^{N_j} \lambda_j (w \cdot V_{l,j}) V_{l,j}. $$
Since the $V_{l,j}$'s are linearly independent, and all the $\lambda_j$ are non-zero, we have that $(w \cdot V_{l,j})$ must be zero for all $l$ and all $j$, which means that each $V_{l,j}$ is orthogonal to the ker of $\hat{\Theta}$, and hence is in its range. From this it follows that $\Delta_0$ is in the range of $\hat{\Theta}$.

Now, we can write

$$C_s(h_t) = \frac{1}{m} \|\Delta_t\|^2 \le e^{-4\lambda t} \frac{1}{m} \|\Delta_0\|^2 = e^{-4\lambda t} C_s(h_0),$$

where we used the fact that $\Delta_0$ lies in the range of $\hat{\Theta}$, where the smallest eigenvalue of $\hat{\Theta}$ is lower-bounded by $\lambda$. Choosing $T = \frac{\log m}{4\lambda}$, we get that

$$C_s(h_T) \le \frac{C_s(h_0)}{m}.$$

Now, fixed any $\hat{\delta}$, we can find a constant $C$ such that with probability higher than $1 - \hat{\delta}$,

$$\sup_{x \in \mathcal{X}} (F_{h_0}(x) - f(x)) \le U.$$

In particular, we have that

$$\mathcal{L}_s(h_T) \le C_s(h_T) = O(1/m),$$

with probability higher than $1 - \tilde{\delta} - \hat{\delta}$. From Proposition 6 and (4) it follows that, with probability higher than $1 - \delta - \tilde{\delta} - \hat{\delta}$

$$\mathcal{L}_Z(h_T) = O\left(\frac{\log m}{m}\right),$$

which is what we wanted to show.

<figure><img src="image_4.png" /></figure>

**Lemma 17** Let

$$\phi(x) = \sum_{j=0}^{J} \gamma^j x^j$$

where $\gamma \in \mathbb{R} \setminus \mathbb{Q}$. The we have that the range of $\Theta$ coincides with the span of the spherical harmonics of degree at most $J$.

**Proof** Since $\phi$ is a polynomial of order $J$, we have that it can be written as a linear combination of the first $J + 1$ Hermite polynomials (up to order $J$). We hence have

$$\phi(x) = \sum_{j=0}^{J} \alpha_j P_j(x).$$

Note that all the Hermite polynomials have rational coefficients. Since $\gamma$ is irrational, $\gamma^j$ cannot be written as a finite linear combination of powers $\gamma^{j'}$ (with $j' \ne j$) with rational coefficient. In particular, we have that

$$\alpha_J = \hat{\alpha}_{J,J} \gamma^J;$$

$$\alpha_{J-1} = \hat{\alpha}_{J-1,J} \gamma^J + \hat{\alpha}_{J-1,J-1} \gamma^{J-1};$$

...

$$\alpha_j = \sum_{J \ge k \ge j} \hat{\alpha}_{j,k} \gamma^k,$$
for some rational non-null coefficients $\hat{\alpha}$'s. From Daniely et al. (2017), we know that we can write
$\Theta(x \cdot x') = R(x \cdot x')$, where $R$ is a polynomial given by

$$
R(u) = \sum_{j=0}^{J} \alpha_{j}^{2} u^{j}.
$$

We can write

$$
\alpha_J^2 = \tilde{\alpha}_{J,2J} \gamma^{2J};
$$

$$
\alpha_{J-1}^{2} = \tilde{\alpha}_{J-1,2J} \gamma^{2J} + \tilde{\alpha}_{J-1,2J-1} \gamma^{2J-1} + \tilde{\alpha}_{J-1,2J-2} \gamma^{2J-2};
$$

$\dots$

$$
\alpha_j^2 = \sum_{2J \ge k \ge 2j} \tilde{\alpha}_{j,k} \gamma^k,
$$

where again the coefficients $\tilde{\alpha}$'s are rational numbers. Note moreover that $\tilde{\alpha}_{j,2j}$ is always non-zero.
Rewriting $R$ as a combination of Gegenbauer polynomials of index $\frac{p-2}{2}$ (see, e.g., Yang and Salman,
2020), we have

$$
R(u) = \sum_{j=0}^{J} \beta_j Q_j(u),
$$

where $Q_j = C_j^{\frac{p-2}{2}}$ is the Gegenbauer polynomial of degree $j$. These polynomial also have only
rational coefficients. Since none of the $\alpha_j^2$'s can be written as a linear combination (with ratio-
nal coefficients) of the others, it must be that none of the $\beta_j$'s is zero. From Theorem H.12 in
Yang and Salman (2020), it follows that $\Theta$ has $J$ non-zero eigenvalues, given by

$$
\lambda_j = \frac{p-2}{p+2j-2} \beta_j
$$

and that the eigenspace associated to $\lambda_j$ is of the spherical harmonics of degree $t$. ■

**Remark 18 (Upper bounding the Laplacian for the cross entropy loss)** We remark here that we used the square loss for the training of the model because it leads to the simple closed form expression for the log density ratio in Proposition 6. If we had considered a multiclass classification task, a more natural loss function would have been the cross entropy loss. Although in such a case it is not clear if a simple upper bound can always be stated for the log density ratio, it is the case that the Laplacian term can be upper bounded easily.

We let

$$
F(x) = \frac{1}{\sqrt{d}} h \Phi(x),
$$

where $h$ is now a $q \times d$ learnable matrix and $\Phi: \mathbb{R}^p \rightarrow \mathbb{R}^d$ is a fixed feature map, chosen randomly at initialisation.

In general, for a generic model with a twice differentiable $\hat{l}$ one has that

$$
\Delta \hat{l} = \sum_i \frac{\partial \hat{l}}{\partial F^i} \Delta F^i + \sum_{ii'} \frac{\partial^2 \hat{l}}{\partial F^i \partial F^{i'}} \nabla F^i \cdot \nabla F^{i'},
$$
which follows easily from the chain rule. Here the model is linear in $h$ and so we are simply left with

$$ \Delta \hat{\ell} = \frac{1}{d} \|\Phi\|^2 \Delta_F \hat{\ell}, $$

where $\Delta_F \hat{\ell} = \sum_i \partial_{F^i}^2 \hat{\ell}$.

When $\hat{\ell}$ is the cross entropy loss, we get

$$ \Delta_F \hat{\ell} = \sum_i \partial_{F^i}^2 \hat{\ell} = \sum_i \left( \frac{e^{F^i}}{\sum_j e^{F^j}} - \frac{e^{2F^i}}{\left(\sum_j e^{F^j}\right)^2} \right) = 1 - \frac{\sum_i e^{2F^i}}{\left(\sum_i e^{F^i}\right)^2} \le 1, $$

and so

$$ \Delta C_s(h) = \frac{1}{m} \sum_{z \in s} \frac{1}{d} \|\Phi(x)\|^2 \left( 1 - \frac{\sum_i e^{2F^i(x)}}{\left(\sum_i e^{F^i(x)}\right)^2} \right) \le \frac{1}{m} \sum_{z \in s} \frac{1}{d} \|\Phi(x)\|^2. $$

As a final remark, we note that if we consider discretised dynamics, since the objective is a convex function of the parameters, we have that

$$ \|\nabla^2 C_s(h)\|_{\text{TR}} = \Delta C_s(h) \le \frac{1}{m} \sum_{z \in s} \frac{1}{d} \|\Phi(x)\|^2. $$

This quantity stays finite as $d \to \infty$ (with arbitrarily high probability), showing that the naive bound $\eta_k \|\nabla^2 C_s(h_k)\|_{\text{TR}} = O(d)$ stated at end of Section 4 can be very loose in practice.

## C.2. Wide Neural Networks

We consider a fully connected neural network $F_h : \mathbb{R}^{n_0} \to \mathbb{R}$, for some $n_0 \in \mathbb{N}$. For simplicity, we require that each hidden layer has the same width, $n \in \mathbb{N}$, and the same activation function, $\phi : \mathbb{R} \to \mathbb{R}$. We assume that the inputs are coming from a compact set $\mathcal{X} \subset \mathbb{R}^{n_0}$.

The network output is determined by an input $x^0 \in \mathcal{X}$, weights $\{W^l\}_{l=2}^L \subset \mathbb{R}^{n \times n}$, $W^1 \in \mathbb{R}^{n_0 \times n}$ and $W^{L+1} \in \mathbb{R}^{n \times 1}$, and biases $\{b^l\}_{l=1}^L \subset \mathbb{R}^n$ and $b^{L+1} \in \mathbb{R}$. We use $h$ to denote the vector of all weights and biases. The network's output is $F_h(x^0) = f_h^{L+1}(x^0)$ where we define the functions $f_h^l$ by,

$$ f_h^0(x) = x, \quad f_h^{l+1}(x) = \phi(f_h^l(x))W^{l+1} + b^{l+1}, \quad \text{for } l = 0, \dots, L, $$

where $\phi$ is applied component-wise. We consider a dataset $s = (x_i, y_i)_{i=1}^m \in \mathcal{Z}^m$ sampled from the measure $\mu^m$. We consider a square loss objective in the form

$$ C_s(h) = \frac{1}{m} \sum_{i=1}^{m} (F_h(x_i) - y_i)^2. $$

We consider a Gaussian initialisation $\rho_0$, where the all of the parameters are independently drawn as

$$ W_{ij}^l \sim \mathcal{N}\left(0, \frac{\sigma_w^2}{n}\right), \quad b_i^l \sim \mathcal{N}(0, \sigma_b^2), $$

for some positive $\sigma_w$ and $\sigma_b$.
If the training step-size is scaled appropriately, the large width limit is known to reduce to the neural tangent kernel (NTK) dynamics (Jacot et al., 2018). Given $x, x' \in \mathcal{X}$, the value of the NTK $\Theta(x, x') \in \mathbb{R}$ is given by the limit (in probability) as $n \to \infty$ of the quantity

$$ \hat{\Theta}(x, x') = \frac{1}{n} \langle \nabla F_{h_0}(x), \nabla F_{h_0}(x') \rangle, $$

with $h_0 \sim \rho_0$.

We borrow the analysis of Lee et al. (2020), who use this fact to study the convergence on the finite width NN under training with gradient descent. To leverage ideas from this analysis, we must make the following additional assumptions:

1. The analytic NTK $\Theta$ is full-rank with minimum and maximum eigenvalues satisfying $0 < \lambda_{\min} \le \lambda_{\max} < \infty$.

2. $\mathcal{Z} = \mathcal{X} \times \mathcal{Y}$ is compact and the data distribution $\mu$ has no atoms.

3. The activation function $\phi$ has Lipschitz continuous and bounded gradients.

Throughout this section, we use the notation $\eta_* = 2(\lambda_{\min} + \lambda_{\max})^{-1}$. We define the functions $(g_h^l)_{l=1}^{L+1}$ by $g_h^l(x) = g_h^{l+1}(\phi(x)W^{l+1} + b^{l+1})$ for $l = 0, ..., L$ such that $F_h = g_h^l \circ f_h^l$ for each $l = 0, ..., L+1$. We also define $\delta_h^l(x) = (\nabla g_h^l) \circ f_h^l(x)$.

**Lemma 19 (Lemma 1, Lee et al. (2020))** Suppose assumptions 1-3 are satisfied, then there exists a constant $R > 0$ such that, for any fixed $C > 0$ and $\delta \in (0, 1)$, for $n$ sufficiently large we can find a convex subset $A(C, \delta, n) \subseteq \mathcal{H}$, such that $\rho_0(A(C, \delta, n)) \ge 1 - \delta/2$ and

$$ \begin{align*} & \|\nabla F_h(x) - \nabla F_{h'}(x)\| \le \sqrt{n} R \|h - h'\|, \quad \|\nabla F_h(x)\| \le \sqrt{n} R, \quad \text{and} \\ & \|\delta_h^l(x) - \delta_{h'}^l(x)\| \le R_1 \|h - h'\|, \quad \|f_h^l(x) - f_{h'}^l(x)\| \le R_1 \sqrt{n}, \\ & \|\delta_h^l(x)\| \le R_1, \quad \|f_h^l(x)\| \le R_1 \sqrt{n}, \quad \text{for } l = 0, \dots, L, \end{align*} $$

for all $h, h' \in B(h_0, Cn^{-1/2})$, $h_0 \in A(C, \delta, n)$ and $x \in \mathcal{X}$.

Here we use the notation $B(h, r)$ to denote the ball about $h$ of radius $r$. As usual, the notation $\nabla$ denote derivatives with respect to the parameters. The result is not stated so explicitly in the paper of Lee et al. (2020), but can be deduced easily by following the proof therein. The convexity of the set $A(C, \delta, n)$ follows from its construction by upper bounding the operator norms of the weight matrices. Similarly, we state more formally another result from this work.

**Lemma 20 (Theorem G.1, Lee et al. (2020))** Suppose assumptions 1-3 are satisfied and $s \in \text{supp}(\mu^m)$. Then there exists constants $C, R_0 > 0$ such that for sufficiently large $n$, whenever $\sup_k \eta_k < \eta_*/n$ and $h_0 \in A(C, \delta, n)$,

$$ \sum_{j=1}^{k} \|h_{j+1} - h_j\| \le Cn^{-1/2}, \quad C_s(h_k) \le \exp\left(-\frac{\lambda_{\min}}{3} \sum_{k'=0}^{k-1} \bar{\eta}_{k'}\right) R_0, \quad \text{for each } k \in \mathbb{N}. $$
To control the change in magnitude of the parameters we define an additional function:

$$J_s(h)^2 = \frac{1}{m} \sum_{i=1}^{m} |\langle h, \nabla_h F_h(x_i) \rangle|^2. \qquad (13)$$

In Remark 22, we discuss how this can be controlled.

**Lemma 21** Suppose assumptions 1-3 are satisfied and $s \in \text{supp}(\mu^m)$. Then there exists constants $C, R_0 > 0$ such that for sufficiently large $n$, whenever $\sup_k \bar{\eta}_k < \eta_\star$ and $h_0 \in A(C, \delta, n)$,

$$\|h_K\|^2 - \|h_0\|^2 \le \frac{2R_3}{\lambda_{\min}n}(J_s(h_0) + R_2).$$

**Proof** From the mean value theorem, it follows that for each $k$ there exists a vector

$$\tilde{h}_k \in \{\alpha h_k + (1 - \alpha)h_{k+1} : \alpha \in [0, 1]\}$$

such that the following holds:

$$\begin{aligned}
\|h_{k+1}\|^2 - \|h_k\|^2 &= 2\tilde{h}_k^T(h_{k+1} - h_k) \\
&= -2\eta_k \langle \tilde{h}_k, \nabla C_s(h_k) \rangle \\
&\le -2\eta_k \langle h_k, \nabla C_s(h_k) \rangle + 2\eta_k \|\tilde{h}_k - h_k\| \|\nabla C_s(h_k)\|.
\end{aligned} \qquad (14)$$

To bound the first term, we first recognise that for any $l = 1, \dots, L+1$, we have $\nabla_{W_{ij}^l} F_h = (\partial_{W_{ij}^l} g_h^l) \circ f_h^l \cdot \phi(f^{l-1})$ and $\nabla_{b_i^l} F_h = (\partial_{b_i^l} g_h^l) \circ f_h^l$. Thus, we obtain,

$$\langle h, \nabla C_s(h) \rangle = \frac{2}{m} \sum_{i=1}^{m} \left( (F_h(x_i) - y_i) \sum_{l=1}^{L} \langle \delta_h^l(x_i), f_h^l(x_i) \rangle \right).$$

With this, we more easily control the change in $\langle h_k, \nabla C_s(h_k) \rangle$ over the iterations. From Lemma 19, it follows that,

$$\left| \langle \delta_{h_k}^l(x_i), f_{h_k}^l(x_i) \rangle - \langle \delta_{h_0}^l(x_i), f_{h_0}^l(x_i) \rangle \right| \\ \le \|\delta_{h_k}^l(x_i) - \delta_{h_0}^l(x_i)\| \|f_{h_k}^l(x_i)\| + \|\delta_{h_0}^l(x_i)\| \|f_{h_k}^l(x_i) - f_{h_0}^l(x_i)\| \le 2R_1 C.$$

Thus, from the Cauchy-Schwarz inequality we obtain that

$$\begin{aligned}
\langle h_k, \nabla C_s(h_k) \rangle &\le 2C_s(h_k)^{1/2} \sqrt{\frac{1}{m} \sum_{i=1}^{m} \left( \sum_{l=1}^{L} \langle \delta_{h_k}^l(x_i), f_{h_k}^l(x_i) \rangle \right)^2} \\
&\le 2C_s(h_k)^{1/2} \left( \sqrt{\frac{1}{m} \sum_{i=1}^{m} \left( \sum_{l=1}^{L} \langle \delta_{h_0}^l(x_i), f_{h_0}^l(x_i) \rangle \right)^2} + 2R_1 C L \right) \\
&= 2C_s(h_k)^{1/2} (J_s(h_0) + 2R_1 C L).
\end{aligned}$$
To bound the second term of (14), we use both Lemma 19 and Lemma 20 to obtain

$$ \begin{aligned} \|\tilde{h}_k - h_k\| \|\nabla C_s(h_k)\| &\le 2\|h_{k+1} - h_k\| C_s(h_k)^{1/2} \sqrt{\frac{1}{m} \sum_{i=1}^m \|\nabla F_{h_k}(x_i)\|^2} \\ &\le 2C_s(h_k)^{1/2}CR. \end{aligned} $$

With this, we obtain the bound,

$$ \|h_{k+1}\|^2 - \|h_k\|^2 \le 4\eta_k C_s(h_k)^{1/2} (J^* + CR + 2R_1CL). \quad (15) $$

Using Lemma 20, we first obtain the simple bound,

$$ \sum_{k=0}^{K-1} \eta_k C_s(h_k)^{1/2} \le n^{-1} R_0^{1/2} \sum_{k=0}^{K} \bar{\eta}_k. \quad (16) $$

When $K$ is large, we use the contractions of the training loss to obtain a sharper bound:

$$ \begin{aligned} \sum_{k=0}^{K-1} \eta_k C_s(h_k)^{1/2} &\le n^{-1} \sum_{k=0}^{K-1} \bar{\eta}_k \exp \left( - \frac{\lambda_{\min}}{6} \sum_{k'=0}^{k-1} \bar{\eta}_{k'} \right) R_0^{1/2} \\ &\le n^{-1} \exp \left( \frac{\lambda_{\min}}{6} \bar{\eta}_{\max} \right) R_0^{1/2} \int_0^{\infty} \exp \left( - \frac{\lambda_{\min}}{6} t \right) dt \\ &\le \frac{6 R_0^{1/2}}{\lambda_{\min} n} \exp \left( \frac{\lambda_{\min}}{6} \bar{\eta}_{\max} \right). \end{aligned} \quad (17) $$

Thus, to obtain the final bound, we use the identity $\|h_K\|^2 - \|h_0\|^2 = \sum_{k=0}^{K-1} (\|h_{k+1}\|^2 - \|h_k\|^2)$, substituting (15) along with the bounds in (16) and (17). ■

Now we state Proposition 8 more formally.

**Proposition 8 (Rigorous statement)** Suppose assumptions 1-3 are satisfied, then there exists positive constants $C, R, R_0, R_2$, and $R_3$ such that, for any $\delta \in (0, 1)$ and sufficiently large $n$, whenever $\sup_k \bar{\eta}_k < \min\{\eta_*, (2R(R + R_0^{1/2}))^{-1}\}$, with a probability of at least $1 - \delta$ on $(s, h_0) \sim \mu^m \otimes \rho_0$

$$ m\Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \le \frac{R_3}{\sigma_w^2 \lambda_{\min}} (J_s(h_0) + R_2) - \sum_{k=0}^{K-1} \text{tr} \log (\text{Id} - \eta_k \nabla^2 C_s(h_k)) + \log \frac{2\xi}{\delta} + \delta, $$

$$ \text{where } \xi = \int_{\mathbb{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_s(h), \mathcal{L}_Z(h))} d\mu^m(s) d\rho_0(h). $$

**Proof** Let $C$ be the constant given in Lemma 20 and suppose $n$ is sufficiently large so that its conclusion is satisfied. We construct the set $A_s$ as the set containing all points visited by gradient flows starting in $A(C, \delta, n)$. By construction, the entire trajectory $\{h_k\}_{k=0}^K$ lies in this set with probability at least $1 - \delta$. To show that $C_s$ is smooth on this set, we note that for any $h, h' \in$
$A_s(C, \delta, n),$

$$
\begin{align*}
& \|\nabla C_s(h) - \nabla C_s(h')\| \\
& \le \frac{1}{m} \sum_{i=1}^{m} (|F_h(x_i) - y_i| \|\nabla F_h(x_i) - \nabla F_{h'}(x_i)\| + \|\nabla_h F_{h'}(x_i)\| |F_h(x_i) - F_{h'}(x_i)|) \\
& \le \frac{1}{m} \sum_{i=1}^{m} (|F_h(x_i) - y_i| \|\nabla F_h(x_i) - \nabla F_{h'}(x_i)\| + \|\nabla F_{h'}(x_i)\| \|\nabla F_{\hat{h}}(x_i)\| \|h - h'\|),
\end{align*}
$$

for some $\hat{h} \in \{\tau h + (1 - \tau)h' : \tau \in [0, 1]\}$.

From Lemma 19, it follows that $\|\nabla F_{h'}(x_i)\|_{\mathrm{F}} \le \sqrt{n}R$. In fact, this holds for all parameter values of distance $Cn^{-1/2}$ from $A(C, \delta, n)$. Since this is a convex set which both $h$ and $h'$ belong to, $\hat{h}$ must belong to this set also, and so $\|\nabla F_{\hat{h}}(x_i)\|_{\mathrm{F}} \le \sqrt{n}R$. Additionally, we apply Lemma 19 as well as the Cauchy-Schwarz inequality to deduce,

$$
\frac{1}{m} \sum_{i=1}^{m} |F_h(x_i) - y_i| \|\nabla F_h(x_i) - \nabla F_{h'}(x_i)\| \le \sqrt{n} R C_s(h)^{1/2} \|h - h'\| \le \sqrt{n} R R_0^{1/2} \|h - h'\|.
$$

From this, we deduce that

$$
\|\nabla C_s(h) - \nabla C_s(\tilde{h})\| \le (R_0^{1/2} R \sqrt{n} + R^2 n) \|h - h'\| \le R(R_0^{1/2} + R) n \|h - h'\|,
$$

which means that the optimisation objective is M-smooth on $A_s(C, \delta, n)$, with $M = R(R_0^{1/2} + R)n$.

Since, by hypothesis, we consider a schedule such that $\sup_k \eta_k \le 1/(2M)$, Theorem 4 applies.
Moreover, we have

$$
\log \frac{\rho_0(h_0)}{\rho_0(h_k)} \le \frac{n}{2\sigma_\omega^2} \left| \|h_k\|^2 - \|h_0\|^2 \right| \le \frac{n}{2\sigma_\omega^2} \|h_k - h_0\| (\|h_k - h_0\| + 2\|h_0\|).
$$

Thus, using Lemma 21,

$$
\log \frac{\rho_0(h_0)}{\rho_0(h_k)} \le \frac{R_3}{\sigma_\omega^2 \lambda_{\min}} (J_s(h_0) + R_2).
$$

from which, the statement follows.

<figure><img src="image_5.png" /></figure>

As noted in Lee et al. (2020) and Yang and Littwin (2021), the NTK analysis of wide neural networks can be performed in settings where the hidden layers have different widths and on networks with different architectures. Therefore, we should expect a similar argument to that given above can be reproduced in these settings.

**Remark 22 (Discussion on J)** As noted in the proof of Lemma 21, the object J defined in (13) can be written in the following form:

$$
J_s(h)^2 = \frac{1}{m} \sum_{i=1}^{m} \left( \sum_{l=1}^{L} \langle \delta_h^l(x_i), f_h^l(x_i) \rangle \right)^2 .
$$

Thus, from Lemma 19, we readily obtain the bound, $|\langle \delta_{h_0}^l(x_i), f_{h_0}^l(x_i) \rangle| \le R_1^2 \sqrt{n}$ for all $h_0 \in A(C, \delta, n)$ and thus,

$$
J_s(h_0) \le LR_1^2 \sqrt{n}, \quad \text{for all } h_0 \in A(C, \delta, n). \tag{18}
$$
However, by borrowing arguments used in the analysis of the limiting form of the NTK, one could be convinced that this estimate is pessimistic. We provide an outline for the reasoning here – a formal analysis would require a careful investigation of the neural network at initialisation and is beyond the scope of this work. In the analysis of the limiting form of the NTK, the approach taken typically has the vectors computed in the forward pass, such as $f_{h_0}^l$, and the backwards pass, such as $\delta_{h_0}^l$, approximated by a zero mean Gaussian distribution. A critical theoretical tool used in this literature is the gradient independence assumption, which requires that the transpose of any weight matrix is independent to the forward pass in the limit as $n \to \infty$ (see, for example, Yang (2020)). This property is used to show that $\sqrt{n}\delta_{h_0}^l$ converges weakly to a zero mean Gaussian with diagonal covariance (see Yang and Littwin (2021)). Under this assumption, it also follows that the covariance in the vector $\sqrt{n}\delta_{h_0}^l(x_i) \odot f_{h_0}^l(x_i)$ is diagonal. Thus, the inner product could be approximated as follows:

$$ \begin{aligned} |\langle \delta_{h_0}^l(x_i), f_{h_0}^l(x_i) \rangle|^2 &= \frac{1}{n} \left( \sum_{i=1}^n \sqrt{n} (\delta_{h_0}^l(x_i))_i (f_{h_0}^l(x_i))_i \right)^2 \\ &\approx \frac{1}{n} \sum_{i=1}^n \left( \sqrt{n} (\delta_{h_0}^l(x_i))_i (f_{h_0}^l(x_i))_i \right)^2. \end{aligned} $$

Since the entries of $f_{h_0}^l(x_i)$ and $\sqrt{n}\delta_{h_0}^l(x_i)|f_{h_0}^l(x_i)$ are identically distributed, we should thus expect that the sum converges with $n$ increasing. However, proving this formally would require an analysis of the rate with which the covariance decays due to gradient independence.

**Remark 23 (Discussion on the curvature term)** In Jacot et al. (2020), the authors study the Hessian of wide neural networks, showing that the Hessian takes the decomposition $\frac{1}{n}\nabla C_s(h) = I(h) + S(h)$ where $I$ is a positive semi-definite matrix that has the same non-zero eigenvalues as the neural tangent kernel and $S$ is a matrix with a trace given by

$$ \operatorname{tr}(S(h)) = \frac{1}{m} \sum_{i=1}^{m} \Delta_h F_h(x_i) (F_h(x_i) - y_i). $$

They analyse these matrices in the limit as $n \to \infty$, showing that for any $r > 2$, $\operatorname{tr}(S^r)$ vanishes. They also show that the matrices become orthogonal and thus, for any $r \in \mathbb{N}$,

$$ \operatorname{tr}((I(h_k) + S(h_k))^r) \to \operatorname{tr}(I_*^r) + \operatorname{tr}(S_*(h_k)^r), \quad \text{as } n \to \infty. $$

With this, it follows from the power series of the matrix logarithm, that with $n \to \infty$ we obtain the limit,

$$ \operatorname{tr} \log (\operatorname{Id} - \eta \nabla^2 C_s(h_k)) \to \operatorname{tr} \log (\operatorname{Id} - \bar{\eta} I_*) + \bar{\eta} \operatorname{tr}(S_*(h_k)) - \bar{\eta}^2 \operatorname{tr}(S_*(h_k)^2). $$

Using the fact that $\operatorname{tr}(S_*(h_k)^2)$ is positive and that $I_*$ is positive semi-definite, we can control this further using Lemma 5:

$$ \operatorname{tr} \log (\operatorname{Id} - \bar{\eta} I_*) + \bar{\eta} \operatorname{tr}(S_*(k)) - \bar{\eta}^2 \operatorname{tr}(S_*(k)^2) \leq \bar{\eta} \frac{3 \operatorname{tr}(\Theta)}{2m} + \bar{\eta} \operatorname{tr}(S_*(k)). $$

They further analyse the matrix $S_*$ in the setting where the gradient flow is used. They show that the trace decays to zero over the iterations and at initialisation, its distribution is captured by a product of joint Gaussian matrices.
## Appendix D. The analysis for general iterative methods

The proof of Theorem 9 follows immediately from the proof technique of Theorem 4 as soon as $-\nabla C_s$ is replaced by the vector field. Similarly, we obtain the following result for the continuous-time flow dynamics.

**Theorem 24** Consider the dynamics $\partial_t h_t = V_s(h_t; t)$, with $V_s: \mathcal{H} \times \mathbb{R}^+ \to \mathcal{H}$ differentiable in $h$. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be a measurable function. For any $\delta \in (0, 1)$ and $T > 0$, with probability at least $1 - \delta$ on the random draw $(s, h_0) \sim \mu^m \otimes \rho_0$, we have

$$ \Psi(\mathcal{L}_s(h_T), \mathcal{L}_Z(h_T)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_T)} - \int_0^T \nabla \cdot V_s(h_t; t) dt + \log \frac{\xi}{\delta} \right), $$

where $\nabla$ refers to derivatives with respect to $h$ and $\xi = \int_{\mathcal{Z}^m \times \mathcal{H}} e^{m\Psi(\mathcal{L}_s(h), \mathcal{L}_Z(h))} d\mu^m(s) d\rho_0(h)$.

### D.1. Mini-batches

We can consider here a version of gradient descent that only evaluates the training objective on a mini-batch $s_k \subset s$ at each iteration $k \in \mathbb{N}$, as already discussed in Section 6. Note that the choice of the mini-batch can be random and doesn't need to be known a priori. Our result will always apply on the specific realisation of the sequence of batches used for the training.

We consider an objective in the form

$$ \mathcal{C}_s(h) = \frac{1}{m} \sum_{z \in s} \hat{\ell}(z, h), $$

for some surrogate loss function $\hat{\ell}: \mathcal{Z} \times \mathcal{H} \to \mathbb{R}$. For a batch $s_k \subset s$ of $m_k$ elements, we write

$$ \mathcal{C}_{s_k}(h) = \frac{1}{m_k} \sum_{z \in s_k} \hat{\ell}(z, h). $$

**Proposition 25** For a sequence of batches $\{s_k\}$, consider the dynamics $h_{k+1} = h_k - \eta_k \mathcal{C}_{s_k}(h_k)$. We denote as $\tilde{\mu}^m$ the law of $(s, \{s_k\})$, which takes into account the potential randomness in the choice of the batches. For each $k$, let $A_s$ be a Borel where $\hat{\ell}(z, \cdot)$ is twice differentiable and M-smooth for every $z$ in $s$, with $\max_k \eta_k \le 1/(2M)$. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be a measurable function. Fix $K \in \mathbb{N}$ and let $\delta \in (0, 1)$, such that the trajectory $\{h_k\}_{k=0}^{K-1}$ lies in $A_s$, with probability at least $1 - \delta/2$ on the randomness of the training dataset $s$, the initialisation $h_0$, and the choice of the batches. With probability at least $1 - \delta$ on the same randomness, we have

$$ \Psi(\mathcal{L}_s(h_K), \mathcal{L}_Z(h_K)) \le \frac{1}{m} \left( \log \frac{\rho_0(h_0)}{\rho_0(h_K)} - \sum_{k=0}^{K-1} \mathrm{tr} \log \left( \mathrm{Id} + \eta_k \nabla^2 \mathcal{C}_{s_k}(h_k) \right) + \log \frac{2\xi}{\delta} + \delta \right), $$

with $\xi = \int_{\mathcal{Z} \times \mathcal{H}} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\mu^m(\bar{s}) d\rho_0(h)$.
## D.2. Momentum

We can consider the use of auxiliary variables: instead of having just $h_k$, we take the pair of processes $(h_k, v_k)$. If the updates of these processes are of the form

$$ \begin{pmatrix} h_{k+1} \\ v_{k+1} \end{pmatrix} = \begin{pmatrix} h_k \\ v_k \end{pmatrix} + V_s(h_k, v_k; k), $$

for some iteration-dependent vector field $V_s$, then the usual analysis applies immediately. For example, let us consider the momentum scheme

$$ h_{k+1} = h_k + v_{k+1}, \quad v_{k+1} = \mu_k v_k - \eta_k \nabla C_s(h_k), $$

where $\mu_k \in [0, 1]$ is the momentum schedule.

This corresponds to the vector field

$$ V_s(h, v; k) = \begin{pmatrix} \mu_k v - \eta_k \nabla C_s(h) \\ (\mu_k - 1)v - \eta_k \nabla C_s(h) \end{pmatrix}, $$

whose Jacobian reads

$$ \nabla V_s(h, v; k) = \begin{pmatrix} -\eta_k \nabla^2 C_s(h) & \mu_k \operatorname{Id} \\ -\eta_k \nabla^2 C_s(h) & (\mu_k - 1) \operatorname{Id} \end{pmatrix}. $$

We can easily compute

$$ \det(\operatorname{Id} + \nabla V_s(h, v; k)) = \det \left( \mu_k (\operatorname{Id} - \eta_k \nabla^2 C_s(h)) + \mu_k \eta_k \nabla^2 C_s(h) \right) = \det(\mu_k \operatorname{Id}) = \mu_k^d, $$

and so

$$ -\sum_{k=0}^{K-1} \operatorname{tr} \log \left( \operatorname{Id} + \nabla V_s(h_k, v_k; k) \right) = -\sum_{k=0}^{K-1} \log \det \left( \operatorname{Id} + \nabla V_s(h_k, v_k; k) \right) = d \sum_{k=0}^{K-1} \log \frac{1}{\mu_k}. $$

If we consider the schedule $\mu_k = 1 - \alpha(k + 1)^{-1}$, for some fixed $\alpha < 1$, then we obtain that this sum scales with $\alpha d \log K$, and this is made dimension independent by choosing $\alpha \sim 1/d$. Curiously, when $\mu_k \equiv 1$, the sum vanishes. As a final remark, note that one must initialise the pair $(h_0, v_0)$ by drawing it from a fixed distribution $\rho_0$, that we assume to have full support on $\mathcal{H} \times \mathbb{R}^d$. This excludes the case of a deterministic initial value for the velocity.

## Appendix E. Discretised damped Hamiltonian dynamics

Here, we consider a Hamiltonian approach, and hence we introduce $d$ additional variables $v$, representing the velocities (momenta) of the parameters $h$. The idea is to exploit the fact that the joint density of the pair $(h, v) \in \mathcal{H}^2$ is conserved under the Hamiltonian flow, a property that is preserved for discrete time-steps by suitable symplectic integrators (Hairer et al., 2006). In order to solve an optimisation problem, we can alternate conservative Hamiltonian steps with dissipative ones, which only involve $v$ and entail an exactly computable change in density.

Let us make things more concrete. For the rest of this section, we denote as $\rho_k$ the joint density of $(h_k, v_k)$. We consider an increasing differentiable map $\psi: \mathbb{R} \to \mathbb{R}$, such that $\psi(0) = 0$, and
we fix $\eta > 0$. We denote as $\Psi_\eta(v)$ the value at $t = \eta$ of the solution of $\partial_t \tilde{v}_t = -\psi(\tilde{v}_t)$ satisfying $\tilde{v}_0 = v_k$, where with a slight abuse of notation we are here implying that $\psi$ is acting component-wise on $\tilde{v}_t \in \mathbb{R}^d$. From $(h_k, v_k)$, to evaluate $(h_{k+1}, v_{k+1})$ we first proceed with a dissipative step:

$$h_{k+1/2} = h_k; \qquad v_{k+1/2} = \Psi_\eta(v_k).$$

Since this step involves the exact solution of a continuous-time gradient descent evolution, we can appeal to the usual continuity arguments to show that

$$\log \frac{\rho_{k+1/2}(h_{k+1/2}, v_{k+1/2})}{\rho_k(h_k, v_k)} = \sum_{i=1}^{d} \log \frac{\psi(v_k^i)}{\psi(v_{k+1/2}^i)}, \qquad (19)$$

with $v^i$ denoting the $i$-th component of $v$. Indeed, we see that for each component of $\tilde{v}_t$

$$\psi'(\tilde{v}_t^i) = \frac{\partial_t(\psi(\tilde{v}_t^i))}{\partial_t \tilde{v}_t^i} = -\frac{\partial_t(\psi(\tilde{v}_t^i))}{\psi(\tilde{v}_t^i)} = -\partial_t(\log \psi(\tilde{v}_t^i)),$$

and so

$$\int_{0}^{\eta} \psi'(\tilde{v}_{t}^{i})dt = \log \frac{\psi(\tilde{v}_{0}^{i})}{\psi(\tilde{v}_{\eta}^{i})} = \log \frac{\psi(v_{k}^{i})}{\psi(v_{k+1/2}^{i})},$$

from which (19) follows.

After this dissipative step, we apply a symplectic Hamiltonian integrator, such as

$$h_{k+1} = h_{k+1/2} + \eta v_{k+1}; \qquad v_{k+1} = v_{k+1/2} - \eta \nabla_h C_s(h_k).$$

This step conserves the density:

$$\rho_{k+1}(h_{k+1}, v_{k+1}) = \rho_{k+1/2}(h_{k+1/2}, v_{k+1/2}).$$

Indeed, we are applying to $(h_k, v_k)$ the tranformation

$$\begin{pmatrix} h \\ v \end{pmatrix} \mapsto \begin{pmatrix} h + \eta v - \eta^2 \nabla C_s(h) \\ v - \eta \nabla C_s(h) \end{pmatrix},$$

whose Jacobian $\begin{pmatrix} 1 - \eta^2 \Delta C_s(h) & \eta \\ -\eta \Delta C_s(h) & 1 \end{pmatrix}$ has determinant 1 for all $h$ and $v$.

Following the above dynamics for $K$ steps and applying the usual Markov argument we obtain the following result.

**Proposition 26** Consider the damped Hamiltonian dynamics described above, with $C_s$ twice differentiable on the whole $\mathcal{H}$. Let $\Psi: \mathbb{R}^2 \to \mathbb{R}$ be a measurable function. Fix $K \in \mathbb{N}$ and let $\delta \in (0, 1)$. With probability at least $1 - \delta$ on the random draw $(s, h_0, v_0) \sim \mu^m \otimes \rho_0$, we have

$$\Psi(\mathcal{L}_Z(h_K), \mathcal{L}_S(h_K)) \le \frac{1}{m} \left( \log \frac{\xi}{\delta} + \log \frac{\rho_0(h_0, v_0)}{\rho_0(h_K, v_K)} + \sum_{k=0}^{K-1} \sum_{i=1}^{d} \log \frac{\psi(v_k^i)}{\psi(v_{k+1/2}^i)} \right), \quad (20)$$

with $\xi = \int_{\mathbb{Z}^m \times \mathcal{H}^2} e^{m\Psi(\mathcal{L}_{\bar{s}}(h), \mathcal{L}_Z(h))} d\mu^m(\bar{s})d\rho_0(h, v)$.
As a concrete example, we can choose $\psi(v) = \varepsilon|v|^p v$, for $p \ge 0$ and $\varepsilon > 0$, where $|\cdot|$ denotes the absolute value computed component-wise. The case $p = 0$ corresponds to the standard conformal damped Hamiltonian dynamics (França et al., 2020), and yields to

$$h_{k+1} = h_k + \eta v_{k+1}; \qquad v_{k+1} = e^{-\varepsilon\eta} v_k - \eta \nabla_h C_s(h_k),$$

with a density that increases exponentially as

$$\log \frac{\rho_{k+1}(h_{k+1}, v_{k+1})}{\rho_k(h_k, v_k)} = d\varepsilon\eta.$$

Note that this last term goes linearly with the dimension of the hypothesis space, a behaviour that is likely to bring poor bounds in over-parameterised settings. To avoid this, one can choose $p > 0$ and get

$$h_{k+1} = h_k + \eta v_{k+1}; \qquad v_{k+1} = \frac{v_k}{(1 + p\varepsilon\eta|v_k|^p)^{1/p}} - \eta \nabla_h C_s(h_k),$$

and

$$\log \frac{\rho_{k+1}(h_{k+1}, v_{k+1})}{\log \rho_k(h_k, v_k)} = \left(1 + \frac{1}{p}\right) \sum_{i=1}^{d} \log (1 + p\varepsilon\eta|v_k^i|^p).$$

With this choice, if the components of $v$ are smaller than 1 (e.g., when they are sampled from a Gaussian with small variance) the last term in the RHS of (20) will likely have a better behaviour with $d$. However, this improvement might come at the price of a larger $\log \frac{\rho_0(h_0, v_0)}{\rho_0(h_K, v_K)}$, due to the fact that less dissipative dynamics allow the model to explore a wider region of the hypothesis space, potentially ending up with a final state $(h_K, v_K)$ with a $\rho_0$ density much lower than $\rho_0(h_0, v_0)$. It is unclear so far what is the optimal choice of $\psi$ that can lead to tightest bounds. We leave the investigation of this open problem as future work.

## Appendix F. A comment on the Laplacian’s integral

An interesting feature of the bound of Theorem 1 is the integral of the Laplacian of the optimisation objective along the training path. Under the gradient flow dynamics, $\Delta C_s(h_t)$ is keeping track of how the probability density is locally varying. Indeed, from a PAC-Bayesian perspective, for the bound to be small we need to end up in some point where the posterior density is not too high compared to the initial one. If we follow a trajectory characterised by a large Laplacian, we see a sharp increase of the density. Intuitively, we can picture this situation as if we were attracting the nearby paths and bringing further probability mass around us. In such a case, the final $\rho_T$ is likely to be much larger than the initial $\rho_0$ and lead to a loose bound. Note that we can rewrite

$$\int_{0}^{T} \Delta C_s(h_t) dt = \log \frac{\|\nabla C_s(h_0)\|}{\|\nabla C_s(h_T)\|} - \int_{h_{[0:T]}} \nabla \cdot \tau(h) \|\delta h\|, \qquad (21)$$

(see below for a derivation). Here $\tau(h) = -\frac{\nabla C_s(h)}{\|\nabla C_s(h)\|}$ is the unit tangent vector to the trajectory in $h$, and $\int_{h_{[0:T]}} \dots \|\delta h\|$ denotes the line integral along the path $h_{[0:T]}$. The last term has a clear meaning, as it quantifies how much $h_{[0:T]}$ is attracting the nearby trajectories.
To obtain (21), first notice that $\partial_t \log \|\nabla C_s(h_t)\| = - \frac{\nabla C_s(h_t)}{\|\nabla C_s(h_t)\|} \cdot \nabla \|\nabla C_s(h_t)\|$. Since, for all $h$,

$$\Delta C_s(h) = \nabla \cdot \nabla C_s(h) = \nabla \cdot \left( \frac{\nabla C_s(h)}{\|\nabla C_s(h)\|} \right) \|\nabla C_s(h)\| + \frac{\nabla C_s(h)}{\|\nabla C_s(h)\|} \cdot \nabla \|\nabla C_s(h)\|,$$

we get

$$\Delta C_s(h_t) = \nabla \cdot \left( \frac{\nabla C_s(h_t)}{\|\nabla C_s(h_t)\|} \right) \|\nabla C_s(h_t)\| - \partial_t \log \|\nabla C_s(h_t)\|.$$

We conclude that

$$\int_{0}^{T} \Delta C_s(h_t) dt = \log \frac{\|\nabla C_s(h_0)\|}{\|\nabla C_s(h_T)\|} + \int_{0}^{T} \nabla \cdot \left( \frac{\nabla C_s(h_t)}{\|\nabla C_s(h_t)\|} \right) \|\nabla C_s(h_t)\| dt.$$

The integral in the RHS is a line-integral along the path $h_{[0,T]}$, as $\|\nabla C_s(h)\|$ is the norm of the flow's "velocity" in $h$. Moreover, $\tau(h) = -\frac{\nabla C_s(h)}{\|\nabla C_s(h)\|}$ is the unit tangent vector to the gradient flow in $h$. We can thus write

$$\int_{0}^{T} \Delta C_s(h_t) dt = \log \frac{\|\nabla C_s(h_0)\|}{\|\nabla C_s(h_T)\|} - \int_{h_{[0:T]}} \nabla \cdot \tau(h) \|\delta h\|.$$

## Appendix G. Implicit regularisation in first-order methods

In this section, we consider some modifications of gradient descent that are known to lead to better generalisation and we relate this improvement to the bound given in Theorem 4.

The role of noise in both optimisation and generalisation is a topic that has received considerable attention recently. The Langevin dynamics, as well as its discrete-time approximations, have proved to be a popular model for investigating this (Raginsky et al., 2017; Mou et al., 2018; Erdogdu et al., 2018; Farghly and Rebeschini, 2021). The Langevin diffusion is most often described by

$$d h_t = - \nabla C_s(h_t) dt + \sigma dB_t, \qquad (22)$$

where $B_t$ is a Brownian motion. Its generalisation properties have been studied using uniform stability and information-theoretic bounds (see Sections 7.2 and 7.3). The associated Fokker-Planck equation yields that the evolution of the marginal density is identical to that of the deterministic flow

$$\partial_t \hat{h}_t = -\nabla \left( C_s(\hat{h}_t) + \frac{\sigma^2}{2} \log \hat{\rho}_t(\hat{h}_t) \right); \quad \hat{h}_0 \sim \rho_0, \qquad (23)$$

where we denote as $\hat{\rho}_t$ the density of $\hat{h}_t$ (Pavliotis, 2014). This idea was recently exploited in Song et al. (2021) in the context of generative modelling. We note that Theorem 1, with the choice of $\Psi(u, v) = \frac{2}{m\sigma^2}(v-u)$ (with $\sigma > 0$) and with a bounded loss $\ell \subseteq [0, 1]$, leads to the bound

$$\mathcal{L}_{\mathcal{Z}}(h_T) \leq \mathcal{L}_s(h_T) + \frac{1}{4m\sigma^2} + \frac{\sigma^2}{2} \left( \log \frac{\rho_T(h_T)}{\rho_0(h_T)} + \log \frac{1}{\delta} \right).$$

Using this as optimisation objective is equivalent to following (23) with $C_s = \mathcal{L}_s - \frac{\sigma^2}{2} \log \rho_0$.

On another note, applying uniformly distributed label noise during training can improve generalisation (Blanc et al., 2020; Damian et al., 2021). Damian et al. (2021) established that in certain
settings this is characterised by the implicit regularisation term $-\frac{\sigma^2}{2B} \text{tr} \log(1 - \frac{\eta}{2} \nabla^2 C_s)$ that is added to the optimisation objective. Here, $\sigma$ is the scale of the label noise and $B$ the batch size. Up to a scaling factor, this is precisely the term that is summed over in the bound of Theorem 4.

However, we note that these algorithms are not optimising our PAC-Bayes bounds. Indeed, if one were to try to use the bound as an optimisation objective, our generalisation guarantee would change (as it would involve the Hessian of this new objective). We defer to future work for deeper analysis of how our bounds relate to the generalisation properties of these algorithm.