Optimization and Bayes: A Trade-off for Overparameterized Neural Networks
Authors: Zhengmian Hu
Heng Huang
Venue: NeurIPS 2023 poster (2023)
Keywords: Bayesian learning
Generalization

Abstract: This paper proposes a novel algorithm, Transformative Bayesian Learning (TansBL), which bridges the gap between empirical risk minimization (ERM) and Bayesian learning for neural networks. We compare ERM, which uses gradient descent to optimize, and Bayesian learning with importance sampling for their generalization and computational complexity. We derive the first algorithm-dependent PAC-Bayesian generalization bound for infinitely wide networks based on an exact KL divergence between the trained posterior distribution obtained by infinitesimal step size gradient descent and a Gaussian prior. Moreover, we show how to transform gradient-based optimization into importance sampling by incorporating a weight. While Bayesian learning has better generalization, it suffers from low sampling efficiency. Optimization methods, on the other hand, have good sampling efficiency but poor generalization. Our proposed algorithm TansBL enables a trade-off between generalization and sampling efficiency.

Optimization and Bayes: A Trade-off for Overparameterized Neural Networks
Zhengmian Hu, Heng Huang Department of Computer Science
University of Maryland
College Park, MD 20740
angkeennmaa@@gmiall com, hom..hengguaanghh@@gmaall com
Abstract
This paper proposes novel algorithm, Transformative Bayesian Learning (TransBL), which bridges the gap between empirical risk minimization (ERM) and Bayesian learning for neural networks. We compare ERM, which uses gradient descent optimize, and Bayesian learning with importance sampling for their gen- eralization and computational complexity. We derive the first Algorithm-depende PAC-Bayesian generalization bound for infinitely wide networks based on an exact KL divergence between the trained posterior distribution obtained by infinitesi- mal step size gradient descent and Gaussian prior. Moreover, we show how transform gradient-bart based optimization into importance sampling by incorporating weight. While Bayesian learning has better generalization, suffers from low sampling efficiency. Optimization methods, the other hand, have good sampling efficiency but poor generalization. Our proposed algorithm TransBL enables a trade-off between generalization and sampling efficiency.
Introduction
Deep Neural Networks (DNNs) have achieved remarkable success in machine learning and related applications. It repeatedly outperformed conventional machine learning approaches. resulting in ground-breaking research such as human-level performance in computer vision [1 substantial
progress in natural language processing (10), and mastering the game of Go [65].
The success of DNNs is the result of critical combination of the complexity and generalization. On the one hand, universal approximation theorem 34 guarantees that any continuous function can be approximated arbitrarily well by using deep network. On the other hand, deep architectures, together with large-scale training data and back pack-propagation algorithm, present a good generalization ability towards unseen data.
Although expressive power and generalization ability are both desirable on their own, they are mutually incompatible. Being one of the main contributions of statistical learning theory, the probably approximately correct (PAC) learning [69] allows us to establish an upper bound of generalization gap by capacity of a hypothesis space. Generally speaking, if model is capable of fitting random labels, then it must generalize poorly.
Recent result [79] shows that overparameterizee DNNs do fit random label perfectly. However, it is observed that more steps of stochastic gradient descent (SGD) are needed train neural network to fit random labels. This phenomenon suggests that the learning capability of DNNs increases with the number of training steps. light this observation, the algorithm-depenndet generalization bounds which control the complexity by maximum training steps preferred over uniform convergence bounds.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

In this paper, we follow the PAC-Bayesian approach to derive our algorihhm-depnnde generalization bounds. PAC-Bayesian bounds were first introduced by McAllester [53] and further developed
These bounds apply for stochastic learning algorithms and focus on expected generalization error over a probability on parameter space. The generalization is related to the KI. divergence between the output distribution of a learning algorithm and prior distribution.
in 61
52 13
Typically, deep neural networks contain large number of parameters and are trained by numerous training epochs through some gradient descent updates. The training dynamics of deep neural networks is complicated, and hence direct theoretical analysis on the KL divergence is intractable. Fortunately, such situation can be largely simplified in the infinite width limit [55] 73 42 18
39 43 76 3). Under this limit, the output distribution of stochastic neural networks drawn from the prior is approximated by Gaussian distribution. Moreover, the training dynamics of gradient-base based optimization is governed by the kernel gradient descent, which guarantees that the evolution of network output only depends on the function values themselves. Thus, our first problem to study in this paper is: Q1 In the infinitely wide neural network limit, can we theoretically derive the formula of this KL divergence?
Optimizing the PAC-Bayesian bound of expected error gives rise the Gibbs measure. It is also widely termed as posterior, as it shares the same form as posterior in Bayes' rule, given that error interpreted as likelihood. Drawing samples from this posterior is very hard in practice. Markov Chain Monte Carlo (MCMC) methods have been explored for deep Bayesian learning. However, it suffers from slow convergence and high computational cost on high- dimensional parameter spaces. Non-exact minimization of PAC-Bayesian bound gives rise to variational approximation (VA) which more computationally efficient, but biased due to the difference between variational distribution and true posterior.
Given that MCMC and VA have their own disadvantages in Bayesian learning, our paper investigates the second problem: Q2 - Does there exist a Bayesian learning method with non-diminishing sampling efficiency even for infinite wide neural network?
Finally, we notice that the gradient-bast based optimization is efficient in training DNNs, but comes with larger generalization error. Bayesian learning. on the other hand, optimizes the expected loss but has lower efficiency. Our third question to study is: Q3- Does there exist an interpolation between optimization and Bayesian learning for trade-off between computation efficiency and generalization error?
In this paper, we give positive answers to above three questions and the main contributions of this paper are summarized as follows:
1) We analyze the infinitely wide neural network trained by gradient flow, or equivalently infinitesimal step size gradient descent. We show that the infinite width limit largely simplifies the training dynam- ics, and the KL divergence between the output distribution and Gaussian prior can be formulated as a function of training time and training data.
2) As a byproduct of our analysis on the gencralization and sampling efficiency, we prove that the trace of Hessian for DNN is not diminishing under infinite width limit and depends on the initialization and training. To the best of our knowledge, this dynamics of Hessian trace is new and maybe of independent interest.
3) We show that if the determinant of Jacobian of optimization flow is available, we can compute a weight for each optimized predictor, such that the weighted output distribution is just posterior. The sampling efficiency in infinite width limit is also derived. We call this type of algorithm as Transformative Bayesian Learning (TransBL) because it is transformed from an optimization procedure.
4) We show that modifying the additional weight in Trans gives rise to an interpolation between optimization and Bayesian learning. The behaviour of TransBL. is increasingly similar to optimization when the weight is changed toward being uniform. This interpolation doesn't thus enables flexible trade-off between sampling efficiency and generalization.
alter training dynamics,
2 Background
We first explain the setting and briefly review PAC-Bayesian and Bayesian learning as background. Other related works is discussed in Appendix A

---

Problem Setup Given ipput space X' and label space }/ we assume that labeled training data are drawn independently from an unknown data distribution D over &' x y. The training set denoted where m is the size of the training sample. A predictor is function h: X - IR and the set of predictors is denoted as H. We introduce two losses 1 : IR x 3 - R. where the first one is used to measure the error rate and the second smooth surrogate loss with polynomially bounded second order derivative is used for providing learning signal in gradient-based training. We define expected loss as R(h)) = (((=)))[(((((s))) 2)] and empirical loss r(h) = Σx=1 [(h(sa), z.)) The expected loss is the central quantity which are interested in but cannot be unobserved in general. Therefore, the currently popular approach to statistical learning derive an upper bound of expected loss.
as 5 [(sa= ...)),,..,. m.),
x3 + [o. 1] and
PAC-Bayesian Bounds We will follow the PAC- AAC-aayesian approach to the generalization problem.
Theorem 1 (Theorem 1.2.6. in (I3). For any positive X and E (0, 1), with at least 1 in the probability of training samples, for all distribution g on space HH, we have
Enni [R(h)] s $ BBhong (r(h)] DKI.(gll)) log 3
The KL divergence is (4 (p) E [log [(a) and the auxiliary function is defined
p(æ)
- (v) 1 - cxp(-o)
exp(-av)
S 1-
000
Although there are huge amounts of research on PAC-Bayesian approach, all of them rely on the KL divergence between output distribution of stochastic predictor learned from certain algorithm and prior p. therefore the analysis of this KL divergence is natural question.
Bayesian Learning
We define Gibbs measure as (0) +ंę -
(r(0) p(0) and partition function
as A = Eo Ar])).
The Gibbs measure is related to posterior in Bayesian inference 24]. and minimize right-hand side in theorem [ inducing much tighter bound:
Enn [R(h)] < log (Z5)) S m(1 - exp(
log Z
Despite the fact that expected loss is optimized, Bayesian learning bears significantly more difficulty than popular ciumization-based algorithms. It is noted that the sampling from high dimensional distribution with irregular shape may lead to high fluctuations and low efficiency, yet in typical neural network, the number of parameters to be trained is very large, and loss landscape is complicated. Many sampling techniques have been developed and explored, but the state-of thethee still far from satisfactory. In this paper, we concentrate on theoretical understanding of the Bayesian approach and consider a special class of importance sampling in Section 4 where we can compare the computation efficiency of Bayesian method and optimization approach.
3 Training by Optimization
We now derive the KL divergence under deterministic optimization setting. We assume all predictors are parameterized by some parameters 0 O The expected loss R(h) and empirical loss r(h) could therefore be regarded as functions of parameter 9. An optimization flow is a function f θ - e, which typically relies on the training data to update the model's parameters 0 and minimize the empirical loss r(h). The updating can be achieved through various methods such as single or multiple steps of gradient descent, gradient flow with infinitesimally small step sizes, more sophisticated methods detailed in Appendix We require the Jacobian determinant ( = det ( f()) exists and the optimization flow is bijective, i.e. the inverse f-1 exists everywhere. This is satisfied for gradient flow and gradient descent when step size is smaller than z- More situations where this assumption holds discussed in Appendix B The starting point of optimization is initialized from a prior distribution ()), and we define a corresponding energy V(0) parameter independent constant. The output distribution could be characterized as follows:
log p(0) + C where C is
(f A) = (f)( A), g(0) p(f p() J)-)))
3

---

The KL divergence between output distribution and prior can be derived as follows:
(
(((())
p(0)
DKL (xl(g -Eenyloog
= Bompllog
p(S(0)) Eg llog p(S(0)) log [J)(0)]]-
p())
Based on a physics analog of Helmholtz free energy change in isothermal process, we define energy term and entropy term separately:
( V V()), S log [
The energy term only depends on the prior. If Gaussian prior is used, V corresponds to the change of squared norm of parameters. Although the KL. divergence is positively related to the energy term, and energy term is monotonically increasing with norm of trained parameters, naively compressing parameters to reduce parameter norm doesn't 't lead decrease in KL divergence, because the decrease in energy term is offset by the change in entropy term.
The entropy term determines the gain loss of entropy caused by an optimization procedure. The pure gain means that the optimization flow maps region of parameter space into another region with larger area. For example, consider the simple function f(x) 2x in a one-dimensional setting. The Jacobian of this function, which describes how much the function stretches or compresses space, constant 2. Imagine we start with an input that is uniformly distributed within the interval [0, 1]. After applying the function, the output fz) becomes uniformly distributed over this larger interval [0, 2). Therefore the entropy increase by log(2). However, most optimization procedures pursue minimization of certain objective function which is only small in an extremely small region of parameter space, and thus loss of entropy inevitable. descent algorithm, the loss of entropy is related to the local curvature. We also extend the discussion to algorithms with enriched state space, such as Momentum SGD and Adagrad in Appendix R22
Section 7 we show that for gradient
Finally, we have
DEL (gllp) Eg- (A, A,5()))
The above KL divergence represents an increase in free energy, or the information gain from the training, depending on the point of view. Both energy term and entropy term can be evaluated empirically, though the computational cost depends on the choice of optimization method, network architecture and prior.
4 Transformative Bayesian Learning
In this section, we introduce a class of importance sampling algorithms solving the Bayesian learning problem. These algorithms based on the optimization flow, energy change, and loss of entropy discussed in Section 3
The most simple example of TransBI. algorithm comes from using the output distribution q of an optimization flow f as the proposal distribution. For any function F on parameter space θ we have
Eg [F()]] =Eang F(O) Px = Z> (0)
Eamp F(f(0)) (r(f(o) p(f(()) {(f()))
(2)
1
Ea-p F(f(0))e -x((1(0)))))) (
Zx
The above expectation could be regarded as attaching an additional weight to the stochastic predictor obtained from ordinary optimization, and should use weighted average of the results obtained by different initialization. This process doesn't 't involve any more training of parameter 0 than the optimization procedure that TransBI. is based on, though additional computation might be required for the value of (0) and (())
The unnormalized weight of above importance sampling is wx (0) = e A That means not all results obtained by training treated equally. In particular, the solution with lower empirical loss, lower energy increase, and producing higher entropy is more important than others.
One benefit of TransBL is that we can measure the computation efficiency by comparing it to the training by optimization. A simple way to do this is to calculate the ratio of effective sample size and

---

sample size:
effx= = (Eappl x (0)))
Ex- 1 (0)]
We define the above value as sampling efficiency, and a non-diminishing efficiency at certain
limit indicates that Bayesian learning is most constant factor slower than optimization. The
numerator of the above expression is determined by the partition function because Eg-p 150-p/Wx(0 [0) = Zy. Therefore, only the denominator depends on the optimization procedure, and we expect Eg [w] (0) be as small as possible for efficient Bayesian learning. Notice that this quantity again relies on energy change A (Ø) and entropy change A We will give theoretical analysis of these two values for infinitely wide neural network Section
4.1 An Illustrative Example
For illustration purpose. we consider univariate loss function that presents two distinct global minima with zero loss. One of these characterized sharp minimum, while the other represents a flat minimum. A direct initialization from the prior, followed by training using gradient flow, often results in an ensemble with significant deviation from the posterior. This is because the optimization process fails to recognize the presence of the sharp minimum, while insights from PAC Bayesian indicate that the flat minimum is surrounded by a higher posterior probability density. TransBL method applies small weight to the solution found within the sharp minimum. Consequently, TransBL can adeptly recreate the posterior, as shown in the Figure 6b figures for the illustration are moved to Appendix F.
Due to space limitations, all
5
Connections between Bayesian Learning and Optimization
5.1 A Bayesian Perspective for Optimization
We show that, in order to achieve low expected loss, the optimization flow in deep learning should reshape the initial distribution to good estimation of posterior. A formal argument relies on applying Donsker and Varadhan's 's variational formula [17] obtain the following equation:
Ba [r())) + DKL (gllp) = - x
log (oxpp(xp((>>(()) + X P
(4)
The first term is independent of training procedure and only relies on the definition of prior. The second term measures the KI. divergence between output distribution and Gibbs measure. This also the gap between expected loss bound obtained by some training methods and the optimal one. Notice that the above KI. divergence could also be expressed in terms of weight for TransBI., energy change, and entropy loss:
DKL (gl(gl = log ZA Ea-pl- log wx (0) = log ZA + Bomp (r(()))))(((())))))+++)))) AA,()) (5)
5.2 Efficiency of Optimization Flows for TransBL
In this section, we show that, in order achieve high sampling efficiency, the optimization flow in deep learning is again expected to reshape the initial distribution to a good estimation of posterior. The only difference to the previous section is that the deviation between output distribution and posterior is measured under different divergence.
In order to see this, we define X2 divergence D2*(pllg) = Exwg (1 ((ll)) )
Theorem 2. For sampling efficiency of TransBL whose proposal distribution is q, we have:
DKL log(1++ (px(l))) log effx
(6)
With the above result, we can establish the equivalence between optimal generalization bound and optimal sampling efficiency, i.e. an optimization flow is optimal in the sense minimizes the expected loss upper bound if and only if the sampling efficiency of TransBI. is 1. A more general correspondence between generalization and sampling efficiency desired. However, the complication lies in different divergences used in two quantities. By comparing Eq. (4) and Eq. (6)
5

---

we see that x² divergence is used for measuring sampling efficiency. Although an upper bound of KL divergence Dxi. (Pxllg could be obtained, this is not directly comparable the Dxi. (gllpx) used in generalization bound because of asymmetry of KI. divergence. That difference also justifies the interpolation between optimization and Bayesian learning shown in the next section.
Interpolation of Optimization and Bayesian Learning
We recall that the output distribution g for the optimization and the posterior Px for Bayesian learning are connected by an additional weight wx= px( z, w'x (0)) These two distributions have distinct properties. The distribution bears worse generalization, but is easy to sample from given an established optimization oracle. The posterior PA, on the other hand, is better shaped, but hard to sample from. Therefore, natural question is whether we can find a distribution PA as intermediate and interpolation between them.
Given our formulation of TransBL., it turned out to be quite easy to construct such an interpolation by modifying the weight. For function Ug R + R* where 3 is the parameter used for interpolation, we can define modified weight as up (w> (0)). The interpolation distribution is
(0) ue (w> (f))(())))))))
z
x
Note that the superscript doesn't mean power but means the interpolation is controlled by B. The expectation a function F(0) on p could be computed similarly as Eq. (2)
1 Eg-p! ((())F([( (0))
E 8app [F(O)]
7
We note Eq. still holds when we change q into px and the gap between the expected loss bound obtained by px and the posterior is:
(o(ww(())
(rAll))) log
Bamp log
z
7
wx(0)
Notice that the above identity degenerates into Eq. (5) if we use v-(-) = which corresponds to the original optimization.
The sampling efficiency of r is again defined the ratio of effective sample size and sample size: eff. (Eo=plus(wx Eg (u (0) and enjoys relation similar Theorem
DKI (x log (x log effy.
6.1 Trade-off between Generalization and Sampling Efficiency
According to Eq. (5) we can see that the generalization is highly affected by extremely small (f). small weight indicates that the optimization produces certain parameters too often. This issue could be addressed by assigning a small weight for these parameters in importance sampling. On the other hand, the sampling efficiency of Bayesian learning is very sensitive large value of (x)),, according to Eqq (3) The high weight appears on the regime where output probability is low but posterior density is high. That essentially means the optimization oracle is not effective in exploring certain region, and the only way to bypass this bottleneck of Bayesian learning is changing the optimization oracle to produce distribution with heavier tail than posterior. Due the exponential dependence on energy change in weight, the distribution of weight is typically heavy-tailed, therefore the efficiency is low. Fortunately, this issue could be largely alleviated with interpolation. Since generalization bound is insensitive to large value of weight, it is possible to apply appropriate modification to weights such that the sampling efficiency is improved and generalization doesn't 't deteriorate too much.
Searching for the Pareto optimality among various trade-off methods is not concern in this paper, although obviously this is important issue and is subject for future work. In this paper, we only
ww w S
consider simple weight clipping method:w:(w)
1 otherwise

---

At the limit of B + 0. p degenerates to output distribution of optimization. At the limit of 5 - oc,
Pi converges to posterior px. A lower bound of sampling efficiency eft 2 (Z /B)2 could be
established. More importantly, weight clipping always achieves smaller generalization bounds than optimization, and higher sampling efficiency than Bayesian learning.
Theorem 3. For parameters o0), DRL (v> (1xx) and efty are both monotonically decreasing functions of B.
Overparameterized Neural Network
The energy change V and entropy change S() play pivot role in previous sections. All quantities of interest, including KL divergence in the PAC-Bayesian bounds, weights in TransBL and sampling efficiency, depend these two terms. Therefore, it is essential to understand the dynamies of these quantities. In this section, we will explore some simple infinitely wide neural network. In particular, this analysis will provide insight for further developments in the deep Bayesian learning and PAC- BByessin bounds.
7.1 Network Definition
We consider a feedforward network with fully connected layers. The parameter 0 vec(W1, .W....,, Wa) vector of flattened weights. For a specific input Sa. the forward propaga- tion of the network is defined as follows:
20 sa) = Sa, y:(0, sa) c, (W,,)-1(0,sa))
s.)) (v)((sss))
The output 30 sa) is scalar. The point-wise function (3)) is activation function with bounded first to fourth order derivative. The hidden layer width at i-th layer is n, such that y:(0, sa), (0,00, s.) E - W; R= are trainable weights. The coefficient c; = 1/ used to ensure the output lies in a proper scale at infinite width limit, and is widely adopted in previous works 43 51 70 40
39
3 18 19]. The network is initialized with Gaussian prior, such
that (W.),, ~ (0,
We introduce the following notation to simplify the result. For a multi-variable scalar function f(X), af(X) x is row vector when X vector. X is a matrix, then (Xxf(X)) ax;,; and x share the same shape as X For vulti-variable vector-value function F(X), we define (Vx ((X)) (X.))
ax,
The network trained by gradient flow or infinitesimal step size gradient descent. Let output vector be (Y (0))), = Wa 8.) and total loss function as L(Y) >a (Y zo), the gradient flow is defined as:
dt ((t) = (TY)=---YYX
7.2 Energy Term
The energy for the prior defined as V(0) = Ed= Wll. We define y2 (0, sa) = Tr(V w. (0, and the dynamics of V (0((t) is:
(0,
d v(O(1) VY (0())) (0(1)) dt
(8) We find that Va (0, sa) can be computed by group of auxiliary vector as Eq. (27) detailed in the appendix, and could be regarded as the tangent vector at output space by propagating a tangent vector W: from i-th layer.
For deriving the dynamics of (Y)) = v/ll (0, sa), ) we define -0) (V) and it follows
d
6) = - -
(8)
di
We note that the definition of (d) follows Y y(d), therefore ((4), as a gradient Gram matrix, just neural tangent kernel (NTK) which has been studied in 39 3 [8
7

---

7.3 Entropy Term
We write down the gradient flow in an abstract form d (t) g(((t)) and denote the solution as (t)) = ftt (()) 8(0) J (( This gradient flow is special example of optimization flow in Section Here the entropy change is S,) = log J (t..) (f). According to Liouville formula, have d
dt (A,55)()))) -divg(O(t)) (r()) Given that the gradient has form (( <Vy (O)) VoY (0))², the entropy change could be formulated follows:
d dt (4,5(0(0)) = Tr V. VV
Σ
The content in the first trace term in Eq. (9) called Gauss-Newton matrix, which only represents curvature information of loss 1" The matrix Ha VoYa is the Hessian of ya(0, sa). It is known that spectral norm of Hessian vanishes for wide neural network [47]. However, we find that the trace is not constant and depends initialization and training. For all 1 S < S d. 1...., we introduce auxiliary vectors (g(t.)), = - (1))a 1I and 7kk = (V (Vo,ua)o"((s) Notice that we left out parameters (f, 8a) for concision.
Theorem 4. Let @ be element-wise product, we have
Tr(Ha, = (V, (d((o))
1 Σ 1=1
The above formula is a multi-variable version of (fa-1 ... f1)" = Ejzi(fa- o ... f ft ((fs-10..... fi)) g(t.)) plays the role of squared gradient and is a kind of diagonal NTK in middle layers.
We define (T()
7 (2) (0, (0, s.) and matrices (1))(0) (Var= Then we have
for all 1 S
d (0()) = - tt (00)))XX))))))))))
(10)
d
7.4 Infinite Width Limit
The dynamics in previous two sections can be dramatically simplified in infinite width limit.
Theorem S. When hidden layers width n = Rl....., Rd-1 approaches infinity, the random vectors (0(0)) and (()) (0(0)) at initialization converges law to a multivariate Gaussian distribution with mean and covariance matrix Σ. Moreover, the following quantities converge in probability to constant and don't 't change during finite training time:
( (0()) (6) Φ (0(1) + 410), (0(1,, (=(is) )a
(11)
The specific formulas of Σ, ( {(t), and (b.d) are shown in Appendix D
The last line of above limits is for all I,..., ny. therefore theorem could be simplified to:
(r Ea (())))
(EEW)
1
TH( =
Theorem demonstrates a function space picture, where the dynamics of y (O(L)) and r (O())
only depend on these values themselves but not internal dynamics of the network. The auxiliary
vectors ()) (0()) and (6) (0(t)) could be obtained by solving ODEs ( and (T)) äe V and $ $ can be integrated according to Eqs. (7) and (9) Therefore, the KL divergence is accessible as an expectation in Eg. (1) Based on this KL divergence, can obtain the final result of the PAC Bayes
bound formulated as a function of training time and training data. The complete formula of this PAC
Bayes bound is shown in Appendix F.4
In order to get a sense of the how the generalization bound and sampling efficiency are influenced by training time, show the asymptotic behaviour under mild assumption on loss function.

---

Funtion
V(o(0) V1e00)
0.0
Space
Spece
Hossiam Ttaso
Pazameter Space
0.10-
-0.1
-0.2
0.00
10
10
Training Time
Traming
Time
Trainin Time
Figure 2: Comparison of dynamics in parameter space and function space.
Corollary 6. If there exist C1 C2 such that a (v. 1) Ty
CD
L (v.)).))) S C2 then
DRa(gel (p) S O(1) O(1)&² effx 2 ex1)exp( e ( for any finite 1, where ((1) represents positive constant irrelevant to I, 9t is the output distribution of an infinitely wide network after training for length of time, and effx the sampling efficiency of importance sampling with qu as proposal distribution.
Experiments
We first illustrate that Hessian trace doesn't 't vanish for overparameterized network and our analysis induces efficient estimation of this value. Next, verify theoretical finding by comparing the dynamics of an overparameterized network in function space and parameter space. Finally, we demonstrate the interpolation of sampling and optimization.
8.1 Non-Diminishing Hessian Trace and Efficient Estimation
We consider a three hidden layers feedforward network with 2048 as hidden layer width and tanh as nonlinear function.
Dur Asprotacrion Eundunoo
Homin Tooc
The synthetic dataset is constructed circle. We choose 20 values uniformly
unit
distributed between 0 to and we let 2- dimensional input to the network be s(T) [cos(T) sin(r)].
2.5
Fleddan Tince
For each input s(T), we are concerned about Hessian trace Tr(V Voya(O, s(T))). Fast Hes- sian trace estimation by itself challenging problem, due to the high dimensional parame- ters and complex structure of the network. We
Figure 1: Hessian trace for one randomly initial- ized network (left) and the probability density of Hessian trace initialization (right).
follow previous works 36 78 and adopt Hutchinson's 's method to estimate the ground truth value of Hessian trace.
On the other hand, our analysis on the dynamics of overparameterized network reveals that certain factors which Hessian trace depends on are insensitive to initialization and training, therefore could be estimated by fixed value to reduce computation. More details are shown in Appendix F.3 We apply the Hutchinson's 's method with 1000 independent random Rademacher vectors. The 30 interval is plotted in the Figure [ It is shown that our approximation aligns with ground truth well. Based on this approximation, the distribution of Hessian is also calculated in Figure]
8.2 Comparing the Dynamics in Parameter Space and Function Space
We use the same toy model as the previous section and set target values to be t(r) sin(33) and loss be mean squared error (MSE). For dynamics parameter space, we run SGD with finite step size 0.01 and
size 1. For dynamics in function space, solve Egs. (7) to (10) with fixed matrices (6), ( and E(.)) The result is plotted in Figure where for SGD, the training time is defined as step size times iteration number. The ground truth of Hessian trace on one input again
estimated by Hutchinson's 's method. We can see from all three kinds of outputs that the function space

---

dynamics produces similar output SGD. The error between them is attributed to discretization error for finite step size and finite width fluctuation.
8.3 Interpolation of Optimization and Bayesian Learning
We consider one-shot learning on Fashion-MNIST 73 We randomly select two classes for binary classification and select one sample for each class as training dataset. We use a single hidden layer network with width being 1024 and softplus activation. We use loss [[y, t) = 1/ + exp(y()) and surrogate loss 1 (v, t) = log(1 + exp(-yf)) for gradient descent. For Gibbs measure, we fix 180. The entropy change approximately evaluated by integrating Eq. (9) with finite step size and fixed (()))
We train 10 independent network and the results are shown in Figure The shadow region represents 3a interval. We first notice that during the training, the expectation of TransBL prediction is stable, yet the variance decreases, indicating the distribution of ensemble is becoming closer to Gibbs measure. This could be verified in middle figure in Figure 3 as the distribution of weights concentrates toward the partition function log X x= 43.6.
1.0
Training
Tratted
0.800
Sampling yutuleury
0.825
0.8
0.00
0.7
0.6
GD Thuin Test Trandtt
10"1
0.773
0.770
0.5
10
Training Time
Ing//r)
log.
Figure 3: Test accuracy for ensemble trained by GD and TransBL with clipping (left). Distribution of weights in TransBL (middle). Trade-off between sampling efficiency and test accuracy (right).
for weight
We also show a clear trade-off between sampling efficiency and test accuracy for TransBL with weight clipping, where the parameter 3 plays pivotal role in balancing these aspects. When 3 is set high, the model leans towards a Bayesian posterior framework. Although this setting noticeably reduces the number of effective samples obtained, observed an enhancement in the model's 's generalization ability, as depicted in the right figure of Figure 3 Conversely, with a smaller /3, the model's 's hehavior tends resemble that of typical deep ensemble, indicating high sampling efficiency but worse generalization.
Conclusion
We study the generalization and sampling efficiency of Bayesian learning with special attention to overparameterized network. We show that both KL divergence, which governs generalization in PAC-Bayesian theory, and x" divergence, which determines sampling efficiency for importance sampling, are related to the change energy and entropy loss during training. The dynamics of these two quantities in DNNs are studied, leading to function space picture in infinite width limit. Our study also contributes to the understanding of DNNs Hessian trace, due to its involvement in entropy loss. By considering importance sampling with output distribution from gradient-basee optimization as proposal distribution, we bridge the gap between optimization and Bayesian learning problems, and provide flexible interpolation for accuracy- computation trade-off.
Acknowledgement
This work was partially supported by NSF IIS 1838627, 1837956, 1956002, 221149 2412, CNS 2213701, CCF 2217003, DBI 2225775.
10

---

References
[1] Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring. ICML, pages 1591-1598. PMLR, 2012.
[2] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in Neural Information Processing Systems, volume 32, 2019.
[3] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. parameterization. ICML, pages 242-252 PMLR, 2019.
convergence theory for deep learning via over-
[4] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In ICML, pages 254-263. PMLR, 2018.
[5] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In ICML. pages 322-332. PMLR, 2019.
[6] Haim Avron and Sivan Toledo. Randomized algorithms for estimating the trace of an implicit symmetric positive semi-definite matrix. Journal of the ACM (JACM), 58(2)1134,, 2011.
[7] Peter Bartlett, Dylan J Foster, and Matus Telgarsky. Spectrally-normalize margin bounds for neural networks. arXiv preprint 17X.v:1706.08498, 2017.
[8] Jeremy Bernstein and Yisong Yue. Computing the information content of trained neural networks. arXiv preprint arXiv:2103.01045, 2021.
[9] Charles Blundell, Julien Coreebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In ICML, pages 1613-1622. PMLR, 2015.
[10] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Language models are few sho- leamerr. In Advances in Neural Information Processing Systems, volume 33, pages 1877-1901,, 2020.
[11] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information P Processing Systems, 32:10836-10846. 2019.
[12] Yuan Cao and Quanquan Gu. Generalization error bounds of gradient descent for lcarning over- parameterized deep relu networks. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 34, pages 3349-3356, 2020.
[13] O Catoni. Pac-hayesian supervised classification: The thermodynamics of statistical learning. institute of mathematical statistics lecture notes monnogaphh series 56. IMS, Beachwood, OH. MR2483528, 5544465, 2007.
[14] Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In ICML, pages 1683- 1691. PMLR, 2014.
[15] Eugenio Clerico, George Deligiannidis, and Arnaud Doucet. Conditional gaussian pac-bayes. arXiv preprint 1xXx:21100111800 2021.
[16] Amnk S Dalalyan. Thcoretical guarantees for approximate sampling from smooth and log-coneave densities. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 79(3):651-67) 6666. 2017.
[17] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov pracess expectations for large time- iii. Communications on pure and applied Mathematics, 39(4):389-466 1976.
[18] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai, Gradient descent finds global minima of deep neural networks. In ICML,, pages 1675-1685. PMLR, 2019.
[19] David Duvenaud, Dougal Maclaurin, and Ryan Adams. Early stopping as nonparametric variational inference. In Artificial Intelligence and Statistics, pages 1070-1077. PMLR, 2016.
[20] Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations, 2020.

---

[21] Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In UAI 2017. AUAI Press, 2017.
[22] Cong Fang, Hanze Dong, and Tong Zhang. Mathematical models of overparameteriied neural networks. Proceedings of the IEEE, 10965)6633703, 2021.
[23] MA Ganaie, Minghui Hu, et al. Ensemble deep learning: A review. arXiv preprint arXiv:2104.02395 2021.
[24] Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julien. Pac-bayesian theory meets bayesian inference. Advances Neural Information Processing Systems, volume 29, 2016.
[25] Alison L Gibbs and Francis Edward Su. On choosing and bounding probability metrics. International statistical review, 70(3):41-433, 2002.
[26] Peter W Glynn and Chang-han Rhee. Exact estimation for markov chain equilibrium expectations. Journal of Applied Probability, 51(A):377-38, 2014.
[27] Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size- independent sample complexity of neural networks. Conference On Learning Theory, pages 297-299. PMLR, 2018.
[28] Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems, 24, 2011.
[29] Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations, 2020.
[30] Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In ICML, pages
PMLR, 2016.
[31] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
[32] Jeremy Heng and Pierre E Jacob. Unbiased hamiltonian monte carlo with couplings. Biometrika, 106(2): 287-302, 2019.
[33] Mathew D Hoffmman David M Blei, Chong Wang. and John Paisley. Stochastic variational inference. Journal of Machine Learning Research, 14(5), 2013.
[34] Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks, 255)-359-36, 1989.
[35] Zhengmian Hu and Heng Huang. On the random conjugate kernel and neural tangent kernel. In Interna- tional Conference on Machine Learning, pages 4359-4368 PMLR, 2021.
[36] Michael F Hutchinson. stochastic estimator of the trace the influence matrix for laplacian smoothing splines. Communications in Statisties-Simulation and Computation, 180)::059--176, 1989.
[37] Pavel Izmailov, Sharad Vikram, Matthew D Hoffmann, and Andrew Gordon Gordon Wilson. What are bayesian neural network posterions really like? In ICML, volume 139, pages 4629 4640. PMLR, 2021.
[38] Pierre E Jacob, John O' Leary, and Yves F Atchadé. Unbiased markov chain monte carlo methods with couplings. Journal of the Royal Statistical Society: Series B (Statisticul Methodology). 82(3):543-600, 2020.
[39] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kemel:
Convergence and generalization in neural networks. In Advances in Neural Information Processing Syxtems, volume 31, 2018.
[40] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint 1xXi:17100011966 2017.
[41] Pitas Konstantinos, Mike Davies, and Pierre Vandergheynst. Pac-bayesian margin bounds for convolutional neural networks-technical report. arXiv preprint arXiv:1801.00171 2017.
[42] Jaehoon Lee, Jascha bick-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks gaussian processes. In ICLR, 2018
12

---

[43] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz. Yasaman Bahri, Roman Novak. Jaseha Sohl-Dicksteins and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. Advances in neural information processing systems, 12:8572-8583, 2019.
[44] Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent only converges minimizers. In Conference on learning theory, pages 1246-1257. PMIRR, 2016.
[45] Xingguo Li, Junwei Lu, Zhaoran Wang. Jarvis Haupt, and Tuo Zhao, On tighter generalization bound for deep neural networks: Cnns, resnets, and beyond. arXiv preprint arXiv: (00:::806003000 1806.0515,, 2018.
[46] Etai Littwin. Tomer Galanti, and Lior Wolf. On random kernels of residual architectures. In Uncertainty in Artificial Intelligence, pages 897-907. PMLR, 2021.
[47] Chaoyue Liu, Libin Zhu, and Mikhail Belkin. On the linearity of large non- non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems, 33, 2020.
[48] Stephan M, Matthew D. Hoffman, and David M. Blei. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research, 18:1-35, 2017.
[49] David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computation, 4 (3):488447) 472, 1992.
[50] Dougal Maclaurin, David Duvenaud, and Ryan P Adams. Early stopping is nonparametric variational inference. arXiv preprint arXiv:150400134 2015.
[51] Alexander G de G Matthews, Mark Rowland, Jiri Hron, Richard E Turner, and Zoubin Ghahramani. Gaussian process behaviour in wide deep neural networks. arXiv preprint XXi::8804112777 2018.
[52] Andreas Maurer. note on the pac bayesian thcorem. arXiv preprint cs/04/ 1099, 2004.
[53] David A McAllester. Some pac-bayesian theorems. Machine Learning, 37(3):35--333, 1999.
[54] Chris Mingard, Guillermo Valle-Péres Joar Skalse, and Ard A Louis. Is sgd bayesian sampler? well, almost. Journal of Machine Learning Research, 22, 2021.
[55] Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
[56] Kirill Neklyudov and Max Welling. Orbital meme. In International Conference on Artificial Intelligence and Statistics, pages 5790-5814, PMLR, 2022.
[57] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of over-paramettrization in generalization of neural networks. In ICLR, 2018
[58] Theodore Papamarkou. Jacob Hinkle, M Todd Young, and David Womble. Challenges in markov chain monte carlo for bayesian neural networks. arXiv preprint arXiv:1910.06532 2019.
[59] Daniel Park, Jascha Sohl- Dickstein, Quoe Le, and Samuel Smith. The effect of network width on stochastic gradient descent and generalization: an empirical study. In ICML, pages 5042-5051. PMLR, 2019.
[60]
Konstantinos Pitas. Dissecting non-vacuous generalization bounds based on the mean-field approximation. In International Conference on Machine Learning, pages 7739-7749. PMLR, 2020.
[61] Matthias Seeger. Pac-hayesian generalisation error bounds for gaussian process classification. Journal of machine learning research, 3(Oct):233-22))) 226. 2002.
[62] Daniel Seita, Xinlei Pan, Haoyu Chen, and John Canny. efficient minibatch acceptance test for metrooolii-hasttnnss. In Proceedings of the TwentySeventh International Joint Conference oncArtificial Intelligence, pages
7 2018.
[63] Mariia Seleznova and Gitta Kutyniok. Neural tangent kernel beyond the infinite-widie limit: Effects of depth and initialization. In International Conference on Machine Learning, pages -9920. PMLR, 2022.
[64] Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems, 32, 2019.
[65] David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieserr Ioannis Antonoglou. Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489,, 2016.
13

---

[66] Taiji Suzuki. Generalization bound of globally optimal non-convex neural network training: Transportation map estimation by infinite dimensional langevin dynamics. Advances in Neural Information Processing Systems, 33:19224-19231 2020.
[67] Gerald Teschl. Ordinary differential equations and dynamical systems, volume 140. American Mathemati- Soc., 2012.
[68] Achille Thin, Yazid Jamati El Idrissi, Sylvain Le Corff, Charles Ollion, Eric Moulines, Arnaud Doucet, Alain Durmus, and Christian Robert. Neo: Non equilibrium sampling on the orbits of deterministic transform. Advances in Neural Information Processing Systems, 34, 2021.
[69] Leslie Valiant. theory of the learnable. Communications of the ACM, 27(11::1134-142 1984.
[70] Twan Van Laarhoven. 1.2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.00 05350. 2017.
[71] Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge University Press, 2019.
[72] Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In ICMI., pages 681-688. Citeseer, 2011.
[73] Christopher KI Williams. Computing with infinite networks. Advances in neural information processing systems, pages 295-301, 1997.
[74] Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of general- ization. Advances in Neural Information Prucessing Systems, volume 33, pages 4697-4708, 2020.
[75] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for henchmarking machine learning algorithms. arXiv preprint arXiv:/708.07747, 2017.
[76] Greg Yang. Sealing limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural langent kernel derivation. arXiv preprint 1XXx:::902.04700 2019.
[77] Greg Yang. Wide feedforward recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems, 32, 2019.
[78] Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael Mahoney. Pyhessian: Neural networks through the lens of hessian. 2020 IEEE International Conference on Big Data (Big Data), pages 581-590. IEEE, 2020.
[79] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep lcarning (still) requires rethinking generalization. Communications of the ACM, 64(3):107-115, 2021.
[80] Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient meme for bayesian deep learning. In ICLR, 2020.
[81] Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz. Nan-vacuous general- ization bounds at the imagenet scale: a PAC-bbaesiaa compression approach. In ICLR, 2019.
14

---

A Related Works
Generalization of Overparameterized Network Generalization for neural network has been studied based on stability [2] norm [7
27 45,, compression 4 and special function class 25.2 21 81 achieves non-vacuous PAC- Bayesian generalization bounds based compression and perturbation robustness. Clerico et al. [ derives PAC-Bayes bound and learning method for training a Gaussian process. A realizable PAC- -ay-ssae bound is also derived in [ with sampling
realizable networks instead of training by gradient descent. Mean field theory 22 be applied to obtain generalization bound.
60
66
could also
Early Stopping Restricting training time is a popular trick in deep learning. Related theoretical explanation includes uniform stability approach [30], which related stability generalization bounds with training time, and a variational inference approach [19), which bound marginal likelihood with entropy. Compared to these works, Corollary 6 although only applies for infinitely wide network, provides two new points of view. First we relates KL. divergence with training time, thus provide a PAC-Bayes explanation for early stop. Second, we control the worst-case sampling efficiency with training time, thus provide a novel computational argument.
Bayesian Deep Learning comprehensive introduction of Bayesian Deep Learning and its relation to generalization is provided in [74]. It is known that sampling from the Gibbs measure is costly
(37]. Bulk of work on Bayesian deep learning based on variational approximation [49] 33 28 9 has
been developed. Despite the theoretical difficulty of exact Bayesian deep learning, various heuristic
methods including SGD [48] 54], deep ensemble [23), and cyclic step size [80] have been shown to
be successful in practice.
Hessian of Overparameterized Network The convergence of various stochastic Hessian trace estimation methods are analyzed in (6) There has been work focusing on empirical evaluation of Hessian of neural network (78). Existing result in [47] shows that spectral norm of Hessian vanish in infinite width limit. Our paper contributes this fields in two aspect. First, we show that Hessian
trace doesn't 't vanish in infinite width limit. Second, we propose novel Hessian trace estimation method in Appendix F3
which is much more efficient than commonly used Hutchinson's 's method.
MCMC Markov Chain Monte Carlo (MCMC) utilize a kernel that leaves Bayesian posterior as stationary distribution. One source of bias of MCMC is finite number of iterations. The mixing time of MCMC suffer from curse of dimension, making it impractical run MCMC till converge. Several orbit MCMC (68] (5)] has been explored, seeking to incorporate non- -oneequilibrumm map into the transition kernel reduce mixing time. They typically use damped Hamiltonian system or momentum gradient descent to bypass the entropy loss due to network curvature, thus suffer from increased divergence as discussed in Appendix F22 Telescopic sum argument
(26) with coupling technique [38 32 have also been proposed to tackle this problem. Another source of bias is inaccurate
implementation transfer kernel. The discretization error of various Langevin dynamics method has
been analyzed [6] 64). The bias from noisy loss function [62] and noisy gradient [72] 14] could be reduced but could not vanish. Papamarkou et al. (88) summarize more challenges of MCMC that leads to lack of convergence.
B Discussion on the strength of assumption
We required f to be a bijective and J, exist everywhere. This assumption is not as restrictive it seems.
1. For gradient descent, f(0) = 1VI. If C², V is L-Lipschitz and h S 1/L, then f a diffeomorphism (Proposition 4.5 [44), thus satisfying our assumption. In optimization literature, the above sufficient condition is mild and standard assumption to ensure the convergence of GD.
2. For gradient flow, _ VI. Our assumption immediately follows existence and unique- ness of ODE solution. For example, if I C2 and VI is L-LLipschtzz then (0) is defined for all f and 0 (Corollary 2.6 [67] and is diffcomorphism (Theorem 2.10 667).
15

---

3. Momentum SGD and Adagrad always satisfy our assumption on enriched state space (Appendix F.2
4. Same assumption was also used in previous research (50 Moreover, many previous works use assumptions that imply our assumption as discussed in 1.
In Section we only analyze the infinitesimal step-size. The main reason for this choice is that allows the training dynamics neural networks to be described by an ordinary differential equation (ODE), which simplifies the analysis. While it's possible that our approach could still be viable with discrete step sizes, would make the theoretical analysis much more challenging. Practically, only discrete step sizes are implementable in experiments. However, our findings show that the theoretical predictions based on infinitesimal steps align closely with the experimental outcomes, as demonstrated in Figure 2
One strong assumption imposed by this paper is an infinite width for the neural network. This assumption is vital for the feasibility of our theoretical analysis in Section 7 and is a common
stance in numerous theoretical studies of neural networks (55] 7 42 [T infinite-widhh width perspective simplifies the dynamics and analysis but, admittedly, deviates from the practical. finite-width architectures typically used in real-world applications. Several researchers 29 20 35, corrections" Delving into this adds complexity to the theoretical analysis, and hence it's left potential topic for future research. Moreover, we find that the dynamics in function space, derived from the theory of infinitely wide networks, closely resembles the behavior of actual finite width networks, as can be seen in Figure 2
18 39
4 76
3 The
46 63 have tried to understand the impact of finite width, known as -intte-width
C Additional Proofs
Proof of Theorem
First we show that sampling efficiency determined by the x² divergence.
Eg [
(0)
1
(-
p(f(0)) 27
(PA(0) 27 ((0)
Eana
Es~
Dxa(P)l))) +1
effx
Z
Z
q(f(O))
(12)
The inequality in Eq. (6) is from Theorem in (25]
Proof of Theorem 3
Based on some algebra and calculus, we have
d 2 dB
B I (w) 2 =3]] x (0)) < B)) s o
DKL (xx = [Ropplog Dog wx(0)
For sampling efficiency, we have
1 ds eft,
d log
T @ (x)(wx(0)))
Eo
2Eoup [12x((0) B)]
Eg-pl (w)))()))))))))))))))))))))))))) 1) 2 0
Proof of Eq. (27) We discard the parameter 8a without loss of generality.
Vw, C,T V Ma
(13)
Vj 2 i, Vo, =c V マバ (9)
=c W J+1 V diag(o)(y))))) (
(14)
(1)
16

---

We finish the proof by combining above results.
Tr Tr(Vw,BaW.) = - V v.
TT(V Bacc 1
Tr V V.yaal (
(15)
Tr
V Vun Au-133d ....... Vn Mi+ 12/
(8)
Proof of Theorem
Tw:yd=G4th V y,
=(G va v. vi+1)
(16)
((C diag(d'k (yd-1)) c. W4+1 (iag(o
Only d(() for j. < d depends on W; in the above formula.
V WWWe,, = (c)2)-1)
Ma diag (u
(17)
1
V (Wa)a.. V (W.)). Nd =(C.T;-1)
V Val (o 2 v (We).. 2 V V()).u)))
(18)
d-1
=(q )
,, Mal Va) (n.), 2 (()).))))
d-1
V (W,)a.. V(W.)..e Dd
Σ
V. Va) (o) o
(19)
11
Tr
V ((W.)., V(Wi).,old
V 2, (a) c((d))
(20)
Σ (V Ve, ((())))
1=1
D Analysis of Infinitely Wide Network
We first give the analytic form limits in Theorem
The covariant matrix Σ is defined recursively. We first define series of Gaussian random variables
(i) (s.), 7k (sw), where variables with different subscript ks are independent. We also let ()
(8)
3/ (s.). The covariance of (1) (sw) and (i) (st) are denoted (x)), (sx)]. Similarly,
(d)
17

---

we can define 70) (s))) and Ex (sa), (s)))
21 21 2 /
= 21 (5a), 7 = 5b/no
Ex+1ly (ob)) (
(1) (s)))
(x
(sb) = El 1/ (s.))
(86)]
(8/
5 (sn) - E[JA (wu (3a): (8.)) (s)), (s)))
(21)
T.,. (0)) =
i>
ol (VK))b
isk
1> k
Ju Th
(UV))
i ko (vu)) < k:
In the infinite width limit, the covariance initialization of (O(O)) and (0(0)) is Ealy gld (sx)] Similarly, the covariance initialization of (Y (8())) and (T ( (0(0))) is (s y ()) (ss)), the covariance at initialization of and (((0)) is 7 (sa)),
Eah ( (sa),
(st)]-
Other matrices limits are
(0(()a.s j=1
=
Σ
E Blo Elok(tul(ss))o(((us)))) Eglu( u(sb)
(22)
k=)
I (lo(((((ss))(((((())) (ol.lss))r(duuaa)lleee)()))1)))) v(s))l.
j=1 kmda
(23)
(5(1) Ja.b )a.b
Σ
II
Tloo (Mod (sa)
j=i-1
b=j
Bio," ((s="())(s))))(((((ss)))) II Blck(tl).)))/((((s)))
[g(((sa).
(24)
j=l k=max(++1.))
(25)
[ (E()))
I Blok (on (sa)) E((v(sa), (((u(sa).((s)))))
(26)
Before embarking the proof, we first give intuition of above results, We first notice that Vd (0, sa) can defined recursively:
Vj Visis Vj (1) sa) (1,(0,
(27)
J <d (0, s.) Cj+1 S diag(o (0,,ss)))) (00 s0,
18

---

We can give a similar definition by forward propagation for 7/ 3 (0, (0,, (v, sa)), which covers the case when j S i:
Vj S i S d, Dj
(6) (0, sa) = 1y(0, 8a):
7+1 (0, 8o) =c++1 1.+i (0,)),,
(6)
(28)
Vi vijjed d, 7it1 (0, s.) = c diag(o (0, su sal))) (0, (0,
(t) Therefore, both Vd and 7 (1) can be regarded output from other "neural network". We compares these forward pass with original neural network as follows:
Vi S j d, Bj+1 (0, sa) ( (U, sa)).
(29)
Both v and Ra propagate in a same way as original network until i-th layer. After that, Va propagate by linear transformation, with matrix diag(o sa))) being controlled by original
network. %
(6) uses different activation function o, which is different than original network with og
in i-th layer, and propagate by linear transformation in a similar way as v
(i) When we consider the yd, Vd and Va output from a multiple branch neural network, the E in Eq. (21) is just conjugate kernel, and (()), ((, in Eq. (22) are fragment of NTK of the larger network.
Proof of Theorem We first prove the convergence to Gaussian distribution. It follows the proof of convergence of θ ( ((.)) Finally we prove that these values don't 't change during training.
We note that in the forward propagation as shown in Eqs. (27) to (29) only need perform two kind of operations: matrix vector multiplication and map vector into matrix by diag, The latter operation could be rewrite into element-wise product which non-linear dlement-wise vector operation, egg. diag(o', (O, (w)))) (0, 8a) de))=;(( (3, sa)) @ (0) (0,
Therefore, the calculation of ya. U2 and 7 (t) readily satisfies the standard of Tensor Program which is firstly introduced in [77] and further developed in (76). Since we require the continuous polynomially bounded third order derivative for a, the general result (66 for Tensor Program could applies, and the convergence in law of outputs to Gaussian distribution and almost surely convergence of NTK at initialization is justified. What left the computation of the covariance and NTK at the infinite width limit.
The variance could be computed recursively. Given that the covariance matrix of each previous layer is established, and the convergence to Gaussian distribution of each element in previous layer, it is easy to verify the covariance matrix for next layer follows Eq. (21)
We next compute the NTK limit. Note that (d)) as gradients inner product, is exactly the same NTK Moreover, (=(+d)), )a is just the diagonal NTK if we consider sub neural network which coincide with original network in previous j layers, but use one single
that has been studied in 39 3 18
hidden unit (y²)a at j-th layer as output. Therefore, we only need take care about ( and ((6) for i < d.
= (s ( Voyu(ss)
(30)
(2) Tr V w, Vd (so) ( Vw,ua(ss))
Vw, vy (s.) = (
(sa)+
(V)
2 (sa) c (v)((((s)))
(s(s))
k=max(s.))
(31)
The first term is easy to calculate and shares similar form with classical NTK after computing the trace Tr ( (Vw, Va (1)
(sa.² V w. Ma(s))
V ,, Ud (4) (s.) caagdl
... C++ W (Ing() Vy.ya(s)) (32)
19

---

Combining above results, we have
(01 Jah
V
(33)
==1
U.S. The limit of first term is lim V (s(e)) "o Ik= (o) In order to calculate the second term (a ( we that Bw- (ww") and (Wy)a is a standard normal random vector, then we have
(50))1(5))2--16666 = El(e, (W5)) (sa)) ( (3 = (sa)s (W(C)(W)))
so the limit is Eyly (sa). ((b)))
All other terms
comes
from second
order
derivative result.
of non-linear We
unit
Ok and doesn't show
up
in classical NTK
note
in
Tr
Tr((c13511 (V. o (c I. (st) Va (sa)) ((V))) va(su)"
Uuk ( (sa) o (f (V.)
o VE (sa)
0(1) converge to a fixed value, in V (ab V yk T Va V Va ya (8)) diag of (s.)) V (s,) (s.)
(5a x--
@V Ha(a)) V (M V (sa) converges to fixed values times identity matrix, and V Ma(B))))) Vaay also converges to fixed values
times identity matrix. The rest terms converge to expectation which contains y and o' (s)())
( (8a). of (sa))
(2(1)) )a,6)aa--VV -Vooo (Vo' (s.)) (s.)) Voba((s))
d - Tr Vw, (0) (sa) (V, + E Tr VW, (s.) ( Vw, j=it1
(34)
Note that (6) (0, 8z) = (V z: Vo )o)) (26 (0, and when j S i the value y/s (0, 8.) depends on We. Therefore we conduct separate discussion. When j > i, the limit value could be derived in way similar (A()) For <
(1) (4) (8a) =Cf Vw, 7d
Vo, (8a)
d-1
(35)
Cjal - ( (a,,)) o ( o (x((((s)))) V1,))s
k=j
The summations are again from the second order derivative of non non-linear unit and the limit could be derived in a similar way to t(0) We focus on first term is somewhat different since it depends on third order derivative of non-linear unit.
V us (sa) = caWa dWW
... Gi+1 W diag(o""
Cj+1 W (U;((())
(36)
(6) V 2d (sa) differs from V Vy (sa) and V (s() in one term o. (sx()))) and that fact is reflected Eq. (25)
Finally we prove that ( (KJ) are constant during training. Since these values are all
inner product of gradients, just need to show gradients Vw,ya(f, sa). Vw (0) (0)(
(0, sa) and (0, sa) doesn't It change too much. We write down the formula of these gradient without the Vw, 7d (0)
20

---

summation terms which could be controlled similarly.
Vw,ya(O, s.) ",
V (0, 001--100-111___)000,,,,,,, va(0,
Vj >
V Vw,yd (0, 8a) 909--10)-10,__)\\,aaa""0
Vj Si
Vw, Wa,nd (0, sa) (0]_1(0)-
(0,
(0, sa V Na(0, sa)
Vj > i+t
(37)
Vw, (8) (0, 5a) (100=1()--100,8_) Hel 80)
Vj = --=++1
Vw,nd (2) (0, 8.) -05 wi Ma (0, 8u (G+ WH1)x diag(o; (y: (0, 8a )))) v, Ve sa)
Vj <i++1
We rely following Local Lipschitz property prove the stability of gradients during training.
Lemma (Local Lipschitz Condition). Let ))-1 (0) to denote any following values:
o (0-1(U)-1(0,5.))
(1-1(U)1((0,s.))
s.):
(38)
(0,,))
(1) 7j (0, 8a),
(0-1(05-100,s)))
there is a constant K > O such that for every R > O, with probability increasing to at limit of n + 0o, or equivalently, with high probability over random initialization (w.h.p.o.ri.) of 0, for any 0,, 02 E B00, R), where B(0, R) = (0' |10 0|| < R). the following holds
[|u; - -, (02)l < K||0; 0211.
(39)
[1,)( & Kvn.
Moreover, let vy (0) to denote any of following values:
V y, Valla, sa), V yct Ve (A, (0, Sa (c++ Wi+1) diag(o," (y, s,)),
(40)
there is constant K > 0 such that for every R > O, w.h.p.o.ri. of 0, for any 0,, O2 E B(0, R),
1
(1ı)(0)) 1z (02)l <
K|10: 0211.
Vn
(41)
[|,((0))| < K.
By using the upper bound in Lemma 7 have that for any R > 0 and 0 E B(0(0), R), the
gradient norm Tr((Vw, (O, 8a ) ya(0, = ( (0, | (0 (0, sa) || < K is bounded w.h.p.o.ri.. For any finite training time, f could be selected large enough ensure 0(t) E (()()), R) during training. The norm bound of difference in Lemma also ensures that the change of gradient norm is bounded by 1 {l)((t ((0) in B(0(0), R) w.h.p.o.ri. which is
vn
diminishing at the limit n +
Proof of Lemma 7 operations: matrix multiplication bounded derivative and element-wine wise multiplication E = @(() where the element-wise non-linear function á is bounded and Lipschitz.
All vectors By in Lemma 7 could be computed sequentially by three kinds of
We, element-wise non-linearity o(C) with having
For the first kind operation, given that W follows standard normal distribution, we have the also normally distributed and enjoy following high probability bound.
21

---

Eq. (2.19) of [71] For n dimensional standard Gaussian random vector Ç ~~ ~~ (0, I,), for any t
Pr +IIS| 2 n =n+t) S 2e 2e
(42)
Therefore we have that for Ç = WE, ||l - KVVlell for some K with high probability.
For element-wise non-lineed non-linearity- o({) where is a standard Gaussian random vector with i.i.d. coordinates, we derive the Chernoff bound follows.
Pr(Ilo(()I 2 KVū) = Pr Σ
(G))
K²n)
inf Pr
exp 1 1 (GG)3
2 2exp(1X2n)
120
Elexp(t
Ei= ((()))))
S inf 120
elKon
(43)
e e =tk"n
EJ exp E
t20
i=1
< inf 120
E (tcxp(((0(()++000))))))))
Se- (1))33++(())))
Therefore, we have ||=|| = | < K yn w.h.p.o.ri. for large enough K.
For element-wiss wise multiplication o(() Ç², since Φ bounded, ||=|| - O(1)l(() all the time.
Recursively apply above three operations gives the bound of norm ||,/||-
In order to control the norm of difference, we note that for matrix multiplication:
||11 W2222l W Wi + W182 W2eell
< W (E1 =2)]l + |[W W2 l|
(44)
We =2 + || -22ll ||2||
The first term could be controlled with high probability as [E1 =2)1l - [Wo(s1 =2)l - II(W1 Wo)(Ei =2) < K VVlles #2ll + R2]le1 - =2 I (LL) 52ll for large enough n, and Wo is the parameter at initialization, R is the maximum distance between W1 and W.,
For element-wiss non-linearity o(C), since o is Lipschitz, we have the difference ||eı - -2ll lØ((i) ((G2)) s O(1)lI - C2ll
(45)
For clement-wise multiplication = <(() @ Ç', since Φ is bounded and Lipschitz, we have the difference
||s e2/l =1(((()) @ C1 ((G2) &ll
- - G1I ((((2)) (Ci -G-I 50(1)]((11-9) - 0
(46)
- &ill
<0(1)]l<> - 22l +(1)II &lll
The first term is reduced by noticing &i has ((1) elements.
Recursively apply above three operations gives the bound of norm of difference.
Proof of Corollary 6
According Theorem 5
the infinite width limit, we have
d (0(1))
(0)
dt
(47)
d
Σ p
((0()()))
1 (V
di
22

---

We use (1) to denote constants that is irrelevant to t. Given that Vy L are bounded, we have
(O())
WM
((())
+ +(1)
(48)
Σ
(O(t))
(0(0))
+ +(1)
follows that the Hessian is also bounded.
IM
Tr Tr(H,(0(0)) 00
T (((O))
(49)
The time derivative and finite time change of energy and entropy is
d
v(o(t)) & 0(1)
Σ Y 0 (0(0))
O())
dt
(.V((())) = (V(((t) --((0-(()))))))))) - 0(1)
Σ 10 (0(0))
t o(1)
(50)
di (A.)(((0)) s 011)+000
(0(0)
+ +11I
IME
A < ((1)
(0(O)) + +
The KL divergence is DKL(u:lD)) =E V -, S(0(())
(8) (010)) + o(1)
(000 + o(1)+ o
<B O(1)
0(1)
Σ
Σ 0
(51)
20(1) +(1)
The last inequality is because E y (0(0)) follows Gaussian distribution and the expectation of norm is finite.
The sampling efficiency could also be controlled.
F[eep(- (exp(-2\r(O(t) - 24.V(0(0)) + +22,5((00))
eff,
z
$ (lxxp(-2A,V/C)) + 22,5(((0))
zx
1+ +ooy o(ye)
<0(1)E exp 0(1) Σ y (O(O))
1+ O(1)
(((0))
<0(1) exp (xx((((0)++(()))))))))) exp (0(1)
Y (e(o))
|
(0(O))
50(1) (x)()))x((( (11))++((1))) 0(1)t²
(52)
The first inequality comes from r(0) 2 0. The last inequality follows from the fact that, for random vectors with tail similar to log-norma distribution, the expectation exists. More specifically, let a, (Σ.) is eigenvalues of covariance matrix E, have (ex.vooee///x(((llt/t s
II, (E)) (exp(clly:])) II, cxp(c) (E)((),
E Additional Experiment Results
In Figure we show the correlation between KI. divergence and generalization gap during training. We also show how weight distribution evolves.
23

---

Dienggooe
0..88
Deccol
100
Laser Gap
122
Traiinn
Toining
Figure 4: The KL divergence, training and testing loss gap, and time dependence of weight distribution of experiment in Section 8.3
F Additional Discussion on Algorithms
F. An illustrative example of TransBI.
(a) Prior
(b) Loss
(c) Posterior
Figure 5: Bayesian learning process
transformed by adding a weight
(a) Gradient flow
TransBL
Figure 6: Illsutration of TransBL
In the presented motivating example, explore univariate loss function that presents two distinct global minima with zero loss. One of these characterized as sharp minimum, while the other represents flat minimum. If the function were to be randomly initialized and then optimized, it might converge to either the two localities. However, insights from PAC Bayesian indicate that the flat minimum is surrounded by higher posterior probability density. A direct initialization from the prior, followed by training using gradient flow, often results significant deviation of the ensemble from the posterior. This is primarily because the optimization process fails to recognize the presence of sharp minimum. The intuitive approach of the TransBL method is to apply a small weight to the solution found within the sharp minimum. Consequently, TransBL can adeptly recreate the posterior, as depicted in the Figure 6b
Furthermore, we demonstrate the interpolation between the ensemble distribution obtained from optimization and Bayesian posterior through weight clipping in Section 6. Figure 77a with 3 0.2, the curve leans more towards the optimization result, revealing a broader spread and

---

less-defined peaks. This shows scenario where optimization has stronger influence than Bayesian learning. Yet with 5. the distribution almost akin what one would expect from a Bayesian posterior. In essence, the parameter / serves as tuning knob, allowing us to interpolate and traverse the spectrum from pure optimization-driver results to outcomes heavily influenced by Bayesian learning. This interpolation mechanism offers a flexible approach to merge the strengths of both methodologies.
(a) 8 0.2
(b) 8
(c) 8 = 5
Figure 7: Interpolation between optimization and Bayesian learning
F.2 Optimization Algorithms
In the main paper, we analyze the change of energy AV and entropy AS for gradient descent (GD). However, with finite step size, GD might not be reversible. Another difficulty lies in the computation of entropy change. As shown in [Bq. (9) the entropy change is related to the curvature of loss and Hessian trace of neural network, where the second part is hard to compute exactly.
We note that this problem could be alleviated for algorithms with enriched state space, but at price of lower sampling efficiency.
Momentum SGD For parameter and momentum v. the Momentum SGD with fraction y, step size h and gradient oracle can be formulated into following two steps:
7 gZ))
(53)
x x
By defining y = [x v). Momentum SGD turns into a transform T④y) composed by two sub transforms T(y) L2 o La (g) that correspond to two lines separately.
The determinant of Jacobian for L1 is del(77) = = where d is dimension of parameters. The determinant of Jacobian for L2 is 1, therefore, the determinant of Jacobian for Momentum SGD is yd
Adagrad is:
For parameter æ, first-order momentum v and second-order momentum m. the Adagrad
m f- m + g(x)²
yu uu (1 7)
g(x) vm
(54)
+ x - hw
The square, square root and division in above formula are all element-wise.
Similar to Momentum SGD, the whole transform of Adagrad could be decomposed into three sub transform. The first and third step are both volume- volume-preservine preserving, and second step compress momentum by y in each dimension. Therefore the determinant of Jacobian for Adagrad is yd
Both above two algorithms enjoy simple form of Jacobian determinant. However, the sampling efficiency degenerates since x" divergence increases when we consider joint probability distribution of parameter and auxiliary variables, like momentum.
25

---

F.3 Approximate DNN Hessian Trace by Forward Propagation
We first recall that due Theorem we have
d-1 (r(H)(0(t)) Σ 2v,,ul va)
di(u))) E(()))
Our motivation is based the stability of (((.)) We have shown in Theorem that in the infinite width limit, g(k-)) converge to vector with same elements and is stable during training. For the wide finite width network, this property is largely preserved.
If replace g(b-)) with its limit value, obtain following formula:
d-1
(r(H)(O(0)) Σ va((()))))
(E(.)))))))
In order to calculate the above value efficiently, we notice that above summation can be regarded as propagation of tangent vector.
first change the network definiton into
10 8a) 8
(0, sa) = c, -,,, 80)s
Vi = l....dd d
(0, s.) (3)((0,, sa)) + b,,
Vi I....,d-- 1.
Notice that compared to original network, add vectors b,, and when b, is set 0, the output is same as original network.
It can be easily seen that VayBe of (O,(t) sa)) (E(,))))))))) be the tangent vector for b; and let it propagates along the forward pass, the result tangent vector in output space is just an estimation of (r())(
Va, therefore we could just let
E4 Final PAC Bayes bound
+ log
Eg [R()]] 4 Eang. [r(0]] +
A
DKL (lp) E 1(x) " nı ~x \ St]
Vo 0, Su
Y, y(
V, VL(Y) y(
(55)
d S. Tr (V - VL(Y.)
ΣΣ p(0) o
dl
/=1 1=1
d y(1 VL((Y)
di
(T()) VLYY)
dt
The closed form solution of the KI. divergence is hard to obtain, but the numerical solution can be computed efficiently on computer (as shown in Figure 4 by solving the above ODE.
26


Reviews:
[{"review_id": "UZNecfbTyv", "reviewer": "Reviewer_1HBf", "summary": "The paper derives an algorithm-dependent PAC-Bayesian generalization bound for infinitely wide neural networks that is based on the KL divergence between posterior obtained by gradient flow and a Gaussian prior. The work shows how to transform optimization into importance sampling through interpolation, coined TransBL, to trade-off generalization and sample efficiency. Also, it provides a proof of non-diminishing Hessian trace under certain conditions.", "questions": "a) There seems to be a slight drop in accuracy for TransBL in Figure 3.a. Is it expected that this is due to noise?\n\nb) The work seems to rely on gradient-flow with infinitesimal step size, whereas regular neural networks are typically trained with finite step sizes. To what extent can we expect similar results still?\n\nc) What happens empirically for different $\\beta$?\n\n", "limitations": "- experimental setup could be more clear\n\nThe experimental set-up is not very clear to me. This makes it harder to judge how well the empirical results align with the theoretical derivations.\n\n- discussion on current limitations\n\nThe paper compares learning in practice with theoretical results. It would be helpful if the paper could also assess in what cases we would expect current results not to hold anymore.\n\n- typos\n\nIn the abstract, the method is called 'TansBL' (twice), whereas the paper uses 'TransBL'.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is clearly structured and the questions it aims to answer are well motivated.\n\n\n\n- Interesting function space perspective of neural network training dynamics through analysis of the KL divergence under infinite width limit assumptions.\n- Bridges gap between optimization and Bayesian learning problems to allow interpolation between computational efficiency and predictive accuracy.\n- Additional results of Hessian trace that could be independently useful.", "weaknesses": "- Experimental setup sometimes hard to follow\n\nIt was sometimes hard to follow the setup of the experimental section (maybe dedicate a separate section in the appendices for details). Having a better picture of the details (especially how inference for TransBL is performed) would allow reproducibility but also spot details that could cause misalignment between theory and empirical results.\n\n- Limited experiments\n\nThe results are promising and already clearly show that the theory is already well-aligned with learning in practice. But, it would be interesting to extend this to a slightly wider set of problems or configurations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We deeply appreciate your recognition of the clarity, motivation, and contributions of our work.\n\n> It was sometimes hard to follow the setup of the experimental section (maybe dedicate a separate section in the appendices for details). Having a better picture of the details (especially how inference for TransBL is performed) would allow reproducibility but also spot details that could cause misalignment between theory and empirical results.\n\nThank you for the suggestion. We will incorporate a more detailed description of the experimental setup in the appendices, providing specifics like the ensemble sampling used and an explanation for the shading in Fig 3.a.\n\nTo address reproducibility concerns, we've provided the code with seeds in the supplementary material to ensure consistent replication of our results.\n\n> it would be interesting to extend this to a slightly wider set of problems or configurations\n\nWe appreciate your recommendation. One of the main challenges in expanding the experiments is our focus on the distribution obtained from optimization, rather than sampling few solutions as done in many deep learning studies. For instance, to produce Fig 3.b, we ran over 100,000 optimization processes. This complexity makes it challenging to quickly expand to other configurations. However, we do recognize the importance and potential of broader applicability, and we will explore this in future work.\n\n> There seems to be a slight drop in accuracy for TransBL in Figure 3.a. Is it expected that this is due to noise?\n\nWe believe this is likely due to statistical noise. As indicated by the broad shading (which represents a 3σ interval), the variance is quite high in the early training stages. At this phase, due to the initial parameter distribution not aligning with the Gibbs distribution, the effective sample size is relatively small, resulting in a larger variance. Consequently, the estimated mean can be noisy.\n\n> The work seems to rely on gradient-flow with infinitesimal step size, whereas regular neural networks are typically trained with finite step sizes. To what extent can we expect similar results still?\n\nThank you for raising this important point. We opted for the infinitesimal step size primarily to simplify the analysis, as it allows us to describe neural network training dynamics with an ODE. For finite step sizes, we believe the core ideas could still be applicable with discrete step sizes, although the theoretical analysis might be more challenging. Moreover, in our experiment with actual neural networks, we indeed use discrete step sizes and have observed that the empirical results closely align with our theoretical findings, as shown in Fig 2.\n\n>  What happens empirically for different $\\beta$?\n\nEmpirically, β serves as a control between generalization and sample efficiency. Larger values tend to exhibit behavior closer to a Bayesian posterior with a rapid decrease in effective sample size but improved generalization, as illustrated in Fig 3.c. Smaller values, on the other hand, show behavior closer to standard deep ensembles.\n\n> experimental setup could be more clear\n\nWe appreciate your suggestion. Due to space limitations, we had to omit certain details from the main paper, like the aforementioned experimental parameters and figure explanations. However, in light of your feedback, we'll ensure that they're included in the appendices to enhance the clarity.\n\n> It would be helpful if the paper could also assess in what cases we would expect current results not to hold anymore.\n\nThis is a pertinent point. Our current analysis heavily relies on the overparameterization. Real-world neural networks often do not operate in this regime, creating a discrepancy between the theoretical entropy estimation and reality. However, if an estimate for this can be obtained, our analysis in Sections 3-6 remains valid, and methods like TransBL and weight clipping can still be applied.\n\n> In the abstract, the method is called 'TansBL' (twice), whereas the paper uses 'TransBL'.\n\nThank you for spotting this oversight. We apologize for the inconsistency and will correct this typo.\n\nAgain, we appreciate your feedback, which will undoubtedly strengthen our manuscript."}, {"review_id": "5XpJQ1sKpU", "reviewer": "Reviewer_sSya", "summary": "The paper introduces a novel learning algorithm, Transformative Bayesian Learning (TransBL), which aims to navigate the balance between empirical risk minimization (ERM) and Bayesian learning in the context of overparameterized neural networks. The authors establish a PAC-Bayesian bound for these networks and provide a theoretical discourse on the relationship between generalization and sampling efficiency. Their primary goal is to explore an intermediary distribution that bridges the gap between ERM and Bayesian learning.The paper presents a sound theoretical argument which I found to be interesting and with potentially moderate impact. \n\nThe main issue currently I see is the presentation of their results. For example, the paper currently lacks certain buildup and intuitive explanations and some clarifications before giving formal treatment. Besides that, the order and flow of the paper currently uses a lot of ambiguities  until they are cleared up much later. I've added several of these ambiguities. Perhaps the authors assumed the reader to be very intimately familiar with this exact research question who will \"mentally fill in the gaps\", but I don't think this should be the default assumption.  I've faced some these ambiguities as I read the manuscript, and added them in the \"weaknesses\" section. Due to my lack of expertise on some related literature, I defer the judgement about novelty of these ideas to other reviewers. ", "questions": "\n- in section 3 the term output distribution $q$ and optimisation flow $f:\\Theta\\to\\Theta$ are introduced but are left somewhat loosely defined. My impression from reading section 3 multiple times is the following: $p(\\theta)$ is the prior, namely just the weight and bias distributions of the model. The flow $f:\\Theta\\to\\Theta$ corresponds to a single step of the gradient-based optimisation procedure (based on its domain & range sets, and specially its name), such that $f(\\theta)$ corresponds to the next step of the optimisation procedure. And $q(\\theta)$ is the probability of the parameters after one step of the optimization process. This interpretation make an intuitive sense: if we assume the gradient updates $\\theta_{t+1} = \\theta_t + \\epsilon \\nabla f(\\theta)$ for some infinitesimal $\\epsilon$ then the fact that volume of $|\\nabla f(\\theta)|$ appearing as a multiplicative term $q(\\theta)$. I understand that I may be incorrect in these interpretation, since an alternative interpretation is that $q$ and $f$ refer to the \"final\" distribution over parameters (i.e., many steps of the gradient descent). Can authors clarify these points? \n- The definitions given between lines 129-130 given for energy and entropy seem rather interesting but somewhat ambiguous. My question is two folds: 1)  it would be good to be clear about whether these are the contributions of this paper? It would be helpful for readers to know which parts are and are not the contributions of this work. 2) What is the intuition behind calling the log-determinant term \"entropy\"? It would be helpful if authors could provide a concrete example for explaining these concepts to familiarize the readers. \n- In line 155 the authors proclaim that the equation (2) doesn't \"involve any more training. \n- Later in the text after line 261 the authors compute $d \\theta(t)/d t$ and refer to it as \"gradient flow\". Is this gradient flow the same as \"f\" introduced earlier in section 3? \n", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Sound theoretical analysis: The paper provides a theoretical analysis on the relationship between ERM and Bayesian learning, specifically in the context of overparameterized neural networks, that may be novel. \n\n- The introduction and justification of TransBL is innovative and backed by comprehensive derivations. The concept of transforming gradient-based optimization into importance sampling could also be interesting. \n\n- Interpolation Mechanism: The introduction of an interpolation mechanism by modifying weights presents an innovative solution to the trade-off between computation efficiency and generalization error.\n\n- Usage of PAC-Bayesian Bounds: The derivation of algorithm-dependent generalization bounds through the use of PAC-Bayesian theory and infinitely wide neural networks adds depth to the theoretical contribution. \n\n- Analysis of Energy and Entropy Change: The authors provide a comprehensive understanding of the dynamics of energy and entropy changes, which play a pivotal role in their algorithm. This exploration of these dynamics could pave the way for further developments in deep Bayesian learning and PAC-Bayesian bounds.", "weaknesses": "\n\n- While the presented theoretical results are relatively well stated, there is a lack of high level intuition and build up that make the paper rather hard to understand for anyone who's not expert on all the related topics.  This can be alleviated by using one or few relevant examples (toy examples where everyone can grasp and convey some key points, as well as by giving some more intuitive and high level explanation of the ideas, before giving a formal description. \n-  It would be helpful for the authors to make some parts of the paper notations and variables clarified so that a reader can follow the terms without the need to know the cited literature or read the appendix.  I've added several questions in the \"Questions\" section to raise some of these points. As another example, in section 7.3 there seems to be a clarification of what $f$ entails by the equation $\\theta(t)= f(t,\\theta(0))$ . Some earlier mention/clarification on this would have helped me/possibly other readers. In line with previous comment, more concrete explanations, as opposed to merely abstract terms, will help clarify the ideas. \n\n- The assumptions such as infinitely wide networks could limit the applicability of the results. Discussion on these assumptions and how they might affect the applicability of their method to practical settings with finite width will be helpful to the reader. While it is understandable that the limit of infinite width is necessary for the theoretical derivations is necessary, perhaps the authors can add more discussions on the approximation errors that this implies for practice , ie. the discrepancy between infinite for theory and finite width used in practice. \n\n- Lack of sufficient empirical validations and baselines: While the paper can be viewed as a mere theoretical contribution, the central proposition of the paper, that they bridge the gap between ERM and Bayesian learning,  is a claim of empirical nature. In other words, more empirical evidence can substantially strengthen the validity of the work and broaden its applicability. While the authors do compare their method to ERM and Bayesian learning, there is a distinct lack of comparison with other more recent methods aiming to achieve similar objectives. These comparisons will be helpful in informing the reader about the practical implications of these study and broaden the paper's impact.  If the authors do not believe further experiments are necessary, perhaps they can be more upfront about their contribution to be of a theoretical nature. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive feedback. \n> Strengths\n\nWe are grateful for your recognition of the strengths of our paper, especially our theoretical contributions in areas including the interpolation mechanism, usage of PAC-Bayesian bounds, and analysis of energy and entropy change.\n> While the presented theoretical results are relatively well stated, there is a lack of high level intuition and build up that make the paper rather hard to understand for anyone who's not expert on all the related topics. This can be alleviated by using one or few relevant examples (toy examples where everyone can grasp and convey some key points, as well as by giving some more intuitive and high level explanation of the ideas, before giving a formal description.\n\nThank you for this valuable suggestion about providing intuitive explanations alongside our theoretical findings. We plan to include a motivating example, which contains a single-variable loss function with two global minima, showing how our TransBL approach intuitively assigning a smaller weights to the sharp minima due to entropy change. This example will be further discussed with references to PAC Bayesian methods to give readers a clearer understanding.\n> As another example, in section 7.3 there seems to be a clarification of what $f$ entails by the equation  $\\theta(t)= f(t,\\theta(0))$. Some earlier mention/clarification on this would have helped me/possibly other readers.\n\nThank you for pointing this out. We will ensure that the function $f$ is introduced and explained earlier in the manuscript.\n\n> While it is understandable that the limit of infinite width is necessary for the theoretical derivations is necessary, perhaps the authors can add more discussions on the approximation errors that this implies for practice , ie. the discrepancy between infinite for theory and finite width used in practice.\n\nWe concur that infinite width is a strong assumption. Such assumption made the theoretical analysis in Section 7 feasible, and it is prevalent in many theoretical discussions, including the prior work mentioned in lines 43-44. There are some works attempting to analyze finite-width corrections. However, these often complicate the theoretical analysis, and we hope that reviewer could understand our choice to leave such discussions for future work. Nevertheless, we appreciate your suggestion and will enhance our discussion on the necessity and limitations of the infinite width assumption.\n\n> Lack of sufficient empirical validations and baselines\n\nThank you for raising this important point. Our main contribution is indeed theoretical, bridging the ERM and Bayesian learning through a parameter $\\beta$. This paradigm is a departure from previous works, making direct comparisons challenging in some aspects.\n\nExisting variational approximation methods can also be viewed as a compromise between ERM and Bayesian learning. However, our method significantly differs from such variational approximation. As demonstrated in our experiments, especially as shown in Fig 3.a, we achieve continuous interpolation. On the other hand, variational approximation's bias depends on the expressiveness of the function class, which means theoretical guarantee of connecting ERM and Bayesian learning and achieving a continuous control like in Fig 3.a are not possible. This makes direct comparisons with these variational approximation difficult.\n\nApart from our main theoretical contribution, we offer a practical contribution. We've developed a novel method for estimating entropy change, as detailed in Sec 8.1. Our method proves to be notably more efficient than prior approaches based on the Hutchinson method.\n\n> And $q(\\theta)$ is the probability of the parameters after one step of the optimization process\n\n> alternative interpretation is that $q$ and $f$ refer to the \"final\" distribution over parameters (i.e., many steps of the gradient descent). Can authors clarify these points?\n\nYou are correct in the explanation of $q(\\theta)$. Both \"one step\" and \"many steps\" interpretations are valid and do not affect the analyses within sections 3-6. \"f\" can represent any invertible function with an existing Jacobian determinant, encompassing single-step gradient descent (GD), multi-step GD, and gradient flow. Only in Section 7 do we focus on gradient flow to discuss the specific forms of energy and entropy change.\n\n> it would be good to be clear about whether these are the contributions of this paper? It would be helpful for readers to know which parts are and are not the contributions of this work\n\nThe energy and entropy decomposition of KL divergence in lines 129-130 is based on a straightforward analogy. We don't regard this definition itself as a contribution of this work. Our four main contributions are detailed from lines 68-84.\n\n> What is the intuition behind calling the log-determinant term \"entropy\"? It would be helpful if authors could provide a concrete example for explaining these concepts to familiarize the readers.\n\nConsider a one-dimensional function f(x)=2x. The Jacobian here is 2. If input x is uniformly distributed within [0,1], the output f(x) will be uniformly distributed within [0,2]. The entropy increases by ln(2), which corresponds to the log-determinant. We'll include this example in our revisions to make our definitions more intuitive.\n\n> In line 155 the authors proclaim that the equation (2) doesn't \"involve any more training.\n\nYes, equation (2) doesn't involve any additional training of the parameter $\\theta$ beyond the optimization procedure in optimization flow f.\n\n> Is this gradient flow the same as \"f\" introduced earlier in section 3?\n\nYes, the gradient flow discussed in line 261 is the same as the abstract \"f\" discussed in sections 3-6. The analyses in Sections 3-6 deal with \"f\" in an abstract, general sense. Only in Section 7 do we narrow our analysis to the gradient flow.\n\nThank you again for your time and insightful comments."}, {"review_id": "6j9AT37wtu", "reviewer": "Reviewer_5yu1", "summary": "This paper proposes a framework to bridge the gap between ERM and Bayesian learning problems.\nThe authors derived the algorithm-dependent PAC-Bayesian generalization bound for infinitely wide networks where the KL divergence between the posterior distribution obtained by infinitesimal step size gradient descent and a Gaussian prior.\nSince direct theoretical analysis on the KL divergence is intractable, they used the well-used simplification that the output distribution of  neural networks drawn from the prior can be approximated by a Gaussian distribution in the infinite width limit.\nFrom the analysis, they provided an interpolation method for accuracy-computation trade-off.\nIn addtion, as a byproduct, this paper analyze the dynamics of the Hessian trace.", "questions": "Q.1 Would you write explicitly what the final resulting of the novel PAC Bayes bound is obtained?\nSince most of the analysis is in the dynamics of the KL divergence, the paper is gradually being analyzed for KL.\nHowever,  makes it difficult to understand the final results obtained.\nIn particular, I would like to know the final result of the PAC Bayes bound formulated as a function of training time and training data.\n\nQ.2 I would like to know if it is possible to compare the results of Gaussian processes with those in the limit of infinitely wide neural networks.\nA neural network with infinite width can be regarded as a Gaussian process, and the PacBayes bound of the Gaussian process can be obtained by Seeger by using the Radon-Nidodim derivative.\nVery naively , it seems that the PacBayes bound can be obtained by plugging the NTK derived from the infinite width of neural network into the PacBayes bound of the Gaussian process.\nIf this Gaussian process based analysis can be an answer of Q1 in the paper, I would like to know what differences can be observed compared to this case.\n", "limitations": "The structure of the paper is difficult to understand.\nThe theoriy that are ultimately obtained is unclear compared to the exisitng PAC Bayes bound.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper presents new insights into the relationship between ERM and Bayesian learning.\nThe most interesting aspects of this paper are the analysis of the KL divergence using variable transformations and its relation to the change in Helmholtz free energy in isothermal processes.\nMoreover, the dynamics of the Hessian trace, which is typically relate to flatness in loss landscape, is analyzed.\n", "weaknesses": "Although the motivation is interesting and the derivations of the equations are interesting, it is difficult to understand what novel PacBayes bound was finally obtained.\nIt is also difficult to understand how they differ from the previous PacBayes bound and what kind of findings are obtained. \nMany of these problems seem to be presentation issues rather than technical issues.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our work and for your insightful comments and questions. We appreciate your recognition of the novel insights and contributions of our work in bridging the gap between ERM and Bayesian learning.\n\n> Would you write explicitly what the final resulting of the novel PAC Bayes bound is obtained? ...  I would like to know the final result of the PAC Bayes bound formulated as a function of training time and training data.\n\nThe final result of the PAC Bayes bound includes several parts, a main bound and a decomposition of KL divergence, and  dynamics of energy and entropy change. The main PAC-Bayes bound is:\n\n$\\mathbb{E}_{\\theta\\sim q_t}[R(\\theta)]=\\Phi^{-1}\\_{\\frac{\\lambda}{m}}(\\mathbb{E}\\_{\\theta\\sim q\\_t}[r(\\theta)]+\\frac{D\\_{KL}(q\\_t||p)+\\log\\frac{1}{\\delta}}{\\lambda})$\n\nwhere the KL divergence $D_{KL}(q_t||p)$ is given by:\n\n$D\\_{KL}(q\\_t||p)=\\mathbb{E}\\_{\\\\{(Y^{(i)}\\_0,\\Gamma^{(i)}\\_0)\\\\}_{i=1}^d\\sim\\Sigma} [V\\_t-S\\_t] $\n\nThe dynamics of $V_t$ and $S_t$ are given by the differential equations:\n\n$\\begin{align*}\nV\\_0&=0,S\\_0=0\\\\\\\\\n\\frac{d}{dt}V_t &= -\\nabla \\mathcal{L}(Y_t)\\sum_{i=1}^d Y^{(i)}_t\\\\\\\\\n\\frac{d}{dt}S_t &= -\\text{Tr}(\\nabla^2\\mathcal{L}(Y_t)\\Theta^{(d)})-\\nabla\\mathcal{L}(Y_t)\\sum\\_{j=1}^d\\sum\\_{i=1}^j \\Gamma^{(i)}\\odot\\Xi^{(i,j)}\\\\\\\\\n\\frac{d}{dt}Y_t&=-\\Theta^{(i)}\\nabla \\mathcal{L}(Y_t)^\\top\\\\\\\\\n\\frac{d}{dt}\\Gamma_t&=-\\Phi^{(i)}\\nabla \\mathcal{L}(Y_t)^\\top\n\\end{align*}$\n\nThe closed form solution of the KL divergence is hard to obtain, but the numerical solution can be computed efficiently on a computer (as shown in Fig 4 of our paper) by solving the above ODE.\n\n> If this Gaussian process based analysis can be an answer of Q1 in the paper, I would like to know what differences can be observed compared to this case.\n\nThank you for this thought-provoking question.\n\nYou're correct in stating that a neural network with infinite width can be regarded as a Gaussian process, but this relationship only holds at initialization. This correspondence breaks down once the network is trained via gradient descent, leading to a distribution that differs from a Gaussian process. While one can indeed construct a Gaussian process using the conjugate kernel $\\Sigma$ (as opposed to the NTK $\\Theta$) to answer questions about the Bayesian learning of neural network, it cannot directly answer questions about the generalization of neural networks trained by gradient descent, hence it cannot be used to answer the Q1.\n\nThank you once again for your valuable feedback and insightful questions. We hope  this response have addressed your concerns and questions."}, {"review_id": "8wP2BL8Rk9", "reviewer": "Reviewer_TNvC", "summary": "The paper provides a perspective interpolating optimization and Bayesian inference. The gradient based optimization procedure gives an output distribution $q(x)$ and the Bayesian inference procedure gives a posterior $p_\\lambda(x)$, the two distributions are connect by the weight $w_\\lambda$. Then the paper proposes to find an interpolation of gradient based optimization and Bayesian inference by using a different weight $w_\\lambda^\\beta$.", "questions": "1. Some of the main theoretical result still follows from the basic assumption of infinite width limit of fully connected neural networks, which is a commonly used approach in theoretical analysis of deep neural networks. However, it is well known that in practice this assumption is generally not true, so I am curious to know what can be improved to generalize the theoretical setting to finite width neural networks.\n\n2. Why the paper considers specifically a simple clipping weight, rather than the actual power function $w_\\lambda^\\beta$, does the proof require that specific form of weight to hold true? \n", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The paper is very well written, even for a reader with not much background in optimization, I found the paper very easy to follow and the main results clearly presented. \n2. Although I am not familiar with the literature of optimization, but I feel the interpolation of gradient based optimization and bayesian inference is an important contribution to the community.\n3. The paper proposes the entropy term and the energy term as two major terms governing the weight, and hence measures the sampling efficiency and generalization error bound. For finite fully connected layers, the analytic form of energy and entropy is provided, and also empirically studied. The paper provides a promising perspective for future analysis of other types of deep neural networks like CNNs and transformers.\n", "weaknesses": "Only some minor ones.\n\nMinor weakness:\n1. line 275, equation 9 is too long\n2. line 265, Equation (27) is not properly referenced.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your constructive feedback and recognition of our work. Your comments have helped us to improve our manuscript.\n\n> Weaknesses\n\nThank you for the suggestion. We will improve the presentation of the mentioned equations for clarity.\n\n> Some of the main theoretical result still follows from the basic assumption of infinite width limit of fully connected neural networks, which is a commonly used approach in theoretical analysis of deep neural networks. However, it is well known that in practice this assumption is generally not true, so I am curious to know what can be improved to generalize the theoretical setting to finite width neural networks.\n\nIndeed, this is a critical question. The assumption of infinite width made the theoretical analysis in Section 7 feasible. However, such assumption is not always valid, and this is a common challenge in many works analyzing infinitely wide networks. To address this, some researchers [1,2,3] have started to explore the extension of infinite width network analysis to finite widths. We believe that an important direction for future work would be to examine whether similar techniques can be extended to our setup, to provide estimation error for the entropy term and energy term under finite width networks.\n\n> Why the paper considers specifically a simple clipping weight, rather than the actual power function $w_\\lambda^\\beta$, does the proof require that specific form of weight to hold true?\n\nThis is a great insight. In our paper, we considered weight clipping due to the monotonicity guarantee provided in Theorem 3. However, this is not the only possible choice. Following your recommendation, we found that the power function $v_\\beta(w_\\lambda) = w_\\lambda^\\beta$ for $\\beta \\in [0,1]$ can also creates an interpolation between optimization and Bayesian learning. When $\\beta = 0$, the modified weights are constant at 1, which turns the distribution into a deep ensemble. When $\\beta = 1$, the weights remain unchanged, preserving the Gibbs distribution. This power function also maintains a monotonicity property similar to that in Theorem 3. Therefore, it is indeed feasible, and we appreciate your suggestion.\n\nWe express our gratitude once again for your thoughtful comments and questions. Your insightful suggestions have inspired us to explore further.\n\n[1] Hanin, Boris, and Mihai Nica. \"Finite depth and width corrections to the neural tangent kernel.\" arXiv preprint arXiv:1909.05989 (2019).\n\n[2] Littwin, Etai, Tomer Galanti, and Lior Wolf. \"On random kernels of residual architectures.\" Uncertainty in Artificial Intelligence. PMLR, 2021.\n\n[3] Dyer, Ethan, and Guy Gur-Ari. \"Asymptotics of wide networks from feynman diagrams.\" arXiv preprint arXiv:1909.11304 (2019)."}, {"review_id": "k7ebWvXX9w", "reviewer": "Reviewer_4jd4", "summary": "The paper presents a novel approach that bridges the gap between optimization and Bayesian learning problems. It introduces an interpolation technique that connects optimization and Bayesian learning. Additionally, the author proposes a generalization bound for infinite wide neural networks and establishes a connection between generalization and sampling efficiency.", "questions": "- Please elaborate more on the sentence in lines 178-179: “This is also the gap between expected loss bound obtained by some training methods and the optimal one.”\n- For binary classification, there are different surrogate loss functions. Please provide some examples of surrogate loss function in multi-classification.\n", "limitations": "-    The concept of importance sampling is not new, and the author proposes importance sampling algorithms to solve Bayesian learning problems. However, the paper does not explicitly discuss the variance issues associated with importance sampling, which can be observed in the experiment results (Figure 3, left).\n\n-    The analysis for over-parametrized neural networks may not be applicable to networks with activation functions like ReLU or other functions with unbounded first derivatives.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "The primary strength of this work lies in its successful formulation of a generalization bound for infinite wide neural networks. The author also provides insightful insights into the relationship between generalization and sampling efficiency, shedding light on an important aspect of the problem.", "weaknesses": "**Technical comments:**\n\n- It would be beneficial to elucidate the motivation behind the sampling efficiency, denoted as $\\text{eff}_{\\lambda}$ in Equation (3).\n- In the related works section, the author should consider mentioning alternative approaches to over-parametrized neural networks, such as random features and mean-field regimes, as suggested in [1].\n- Corollary 6 is based on a surrogate loss function. Are there any assumptions on the main loss function necessary to derive an upper bound on expected loss?\n\n**Presentation comments:**\n\n- In the problem setup section (Line 88), the input and label spaces are denoted as $\\mathcal{X}$ and $\\mathcal{Y}$, respectively. However, in Line 90, the training dataset is denoted as $(s,t)$. For consistency, it would be advisable to use $(x,y)$ to represent each data sample throughout the manuscript.\n\n- In the equation following Line 127, the right-hand side of the second equality represents an expectation with respect to $\\theta \\sim p(f(\\theta))$. However, the current notation is unclear and could be misinterpreted as $\\theta \\sim p(\\theta)$. This should be clarified.\n\n- It is recommended to provide a reference for Donsker and Varadhan, which is used to derive Equation (4).\n\n- Please provide definitions for effective sample size and sample size before Equation (3).\n\n- Please provide a clear definition of generalization and generalization bound within the context of this paper. Is the bound on true risk?\n\n- The theorems should be self-contained. In Line 295, the author mentions $\\alpha$ without prior explanation in Theorem 5.\n\n**Minor**\n\n- The author employs $t_a$ and $t$ to represent label and time, respectively. It would be preferable to use alternative symbols to avoid confusion.\n\n- Inline equations in Lines 637 and 641 should be displayed as outlined equations for improved readability.\n\n**References:**\n\n-       [1]: Fang, Cong, Hanze Dong, and Tong Zhang. \"Mathematical models of overparameterized neural networks.\" Proceedings of the IEEE 109.5 (2021): 683-703.\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed feedback. We appreciate the time and effort you've spent in reviewing our work. Below, we've addressed your points in the order presented.\n\n> It would be beneficial to elucidate the motivation behind the sampling efficiency, denoted as $\\text{eff}_{\\lambda}$ in Equation (3).\n\nThanks for the clear question. We are willing to explain sampling efficiency in more detail.\n\nWe've defined the sampling efficiency as $(\\mathbb{E}[w])^2/\\mathbb{E}[w^2]$, which is the ratio of effective sample size $n_e^*=\\frac{n(\\mathbb{E} [w])^2}{\\mathbb{E} [w^2]}$ to the actual sample size $n$. The effective sample size is a key concept in importance sampling, controlling the variance of the final estimator.\n\nTo provide a brief overview of the concept, we now provide an elementary derivation of the form of sampling efficiency and effective sample size in the importance sampling context. In importance sampling, $w$ represents the importance weights. Given $n$ samples, each of which is independently sampled as $Z_i$, and the importance weights are $w_i$. The final estimate is $S=\\frac{\\sum_{i=1}^n w_i Z_i}{\\sum_{i=1}^n w_i}$. If the variance of $Z_i$ itself is $\\sigma^2$, the variance of $S$ is $\\frac{\\sum_{i=1}^n w_i^2 \\sigma^2}{(\\sum_{i=1}^n w_i)^2}$. We know that the variance of the mean of $n_e$ number of independent identically distributed variables is $\\frac{\\sigma^2}{n_e}$. Thus, the final variance in importance sampling is **equivalent** to the variance when sampling $n_e=\\frac{(\\sum_{i=1}^n w_i)^2}{(\\sum_{i=1}^n w_i^2}$ independent samples. Therefore, the effective sample size is defined as $n_e=\\frac{(\\sum_{i=1}^n w_i)^2}{(\\sum_{i=1}^n w_i^2}=\\frac{n\\overline{w}^2}{\\overline{w^2}}$. Since $w_i$ is also a random variable, we can derive the population version of $n_e$ by changing the sum to the expectation. Thus, we have $n_e^*=\\frac{n(\\mathbb{E} [w])^2}{\\mathbb{E} [w^2]}$. The effective sample size is generally less than the true sample size, and this decrease factor is the sampling efficiency, which leads to the definition in Eq. (3).\n\nWe hope this clarifies your question. More details can be found in Owen, Art B. \"Monte Carlo theory, methods and examples.\" (2013). (Section 9.3)\n\n> the author should consider mentioning alternative approaches to over-parametrized neural networks, such as random features and mean-field regimes, as suggested in [1].\n\nThank you for the suggestion. We will add references to these alternative approaches in the related works section.\n\n> Corollary 6 is based on a surrogate loss function. Are there any assumptions on the main loss function necessary to derive an upper bound on expected loss?\n\nThe only requirement we have for the main loss function is that its output is in $[0,1]$, as stated in line 91.\n\n> For consistency, it would be advisable to use $(x,y)$ to represent each data sample throughout the manuscript.\n\nThank you for your suggestion. We understand the value of consistent notation. However, in Section 7, 'x' and 'y' are used to represent activations in neural networks. To prevent confusion, we opted to use 's' and 't' to denote inputs and outputs in the manuscript.\n\n> In the equation following Line 127, the right-hand side of the second equality represents an expectation with respect to $\\theta \\sim p(f(\\theta))$. However, the current notation is unclear and could be misinterpreted as $\\theta \\sim p(\\theta)$. This should be clarified.\n\nIt appears that your understanding of this equation is incorrect.\n\nWe confirm that we intended to write $\\theta \\sim p(\\theta)$, not $\\theta \\sim p(f(\\theta))$. In fact, the latter notation would not make sense as $p(f(\\theta))$ cannot be regarded as a distribution. If you could explain why you believe it should be $\\theta \\sim p(f(\\theta))$, perhaps we can better assist in clarifying this misunderstanding.\n\n> It is recommended to provide a reference for Donsker and Varadhan, which is used to derive Equation (4).\n\nWe'll add the following citation:\n\"Donsker, Monroe D., and SR Srinivasa Varadhan. \"Asymptotic evaluation of certain Markov process expectations for large time—III.\" Communications on pure and applied Mathematics 29.4 (1976): 389-461.\"\n\n> Please provide definitions for effective sample size and sample size before Equation (3).\n\nWe will add definition of effective sample size in the manuscript. For more details about effective sample size , olease kindly refer to our explanation in the beginning of the response.\n\n> Please provide a clear definition of generalization and generalization bound within the context of this paper. Is the bound on true risk?\n\nGeneralization refers to the phenomenon that minimizing the empirical loss on a training dataset leads to reduced  expected loss on new, unseen data. Generalization bound is just an upper bound of the difference between the expected loss and the empirical loss. For a precise definition, please refer to lines 88-97.\n\nWe use $R(h)$ to denote true risk, specifically, the expected prediction error of a model on the unknown real data distribution. So, yes. The upper bound is on true risk.\n\n> In Line 295, the author mentions $\\alpha$ without prior explanation in Theorem 5.\n\nWe agree with your observation. We will clarify this by stating \"for all $\\alpha=1,\\dots,j$\" in the Theorem.\n\n> It would be preferable to use alternative symbols to avoid confusion.\n\nThank you for your suggestion. We will change symbol $t$ to $z$ for notation of label.\n\n> Inline equations in Lines 637 and 641 should be displayed as outlined equations for improved readability.\n\nThank you for pointing this out. We will make the necessary adjustments to improve readability."}]

Author rebuttal:

Meta-review: This meta review is based on the reviews, the authors rebuttal and the discussions with the reviewers, discussions with the SAC, and ultimately my own judgement on the paper. There was a consensus that the paper contributes sound and interesting contributions. I feel this work deserves to be featured at NeurIPS and will attract interest from the community. I would like to personally invite the authors to carefully revise their manuscript to take into account the remarks and suggestions made by reviewers. Congratulations!
Decision: Accept (poster)