Title: Optimization and Bayes: A Trade-off for Overparameterized Neural Networks
Venue: NeurIPS 2023
Accepted: yes
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

Abstract

This paper proposes a 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 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 TransBL enables a trade-off between generalization and sampling efficiency.

1 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 [31], substantial progress in natural language processing [10], and mastering the game of Go [65].

The success of DNNs is the result of a 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 a deep network. On the other hand, deep architectures, together with large-scale training data and back-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 a model is capable of fitting random labels, then it must generalize poorly.

Recent result [79] shows that overparameterized DNNs do fit random label perfectly. However, it is observed that more steps of stochastic gradient descent (SGD) are needed to train a neural network to fit random labels. This phenomenon suggests that the learning capability of DNNs increases with the number of training steps. In light of this observation, the algorithm-dependent generalization bounds which control the complexity by maximum training steps are preferred over uniform convergence bounds.

In this paper, we follow the PAC-Bayesian approach to derive our algorithm-dependent generalization bounds. PAC-Bayesian bounds were first introduced by McAllester [53] and further developed in [61, 52, 13]. 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 KL divergence between the output distribution of a learning algorithm and a prior distribution.

Typically, deep neural networks contain a 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, 51, 18, 39, 43, 76, 3]. Under this limit, the output distribution of stochastic neural networks drawn from the prior is approximated by a Gaussian distribution. Moreover, the training dynamics of gradient-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 to the Gibbs measure. It is also widely termed as posterior, as it shares the same form as posterior in Bayes' rule, given that error is 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 is 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-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) Weanalyze 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 dynamics, 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 generalization 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 TransBL 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 alter training dynamics, thus enables flexible trade-off between sampling efficiency and generalization.

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 input space X and label space Y , we assume that labeled training data are drawn independently from an unknown data distribution D over X × Y . The training set is denoted as S train = { ( s a , z a ) | 1 , . . . , m } , where m is the size of the training sample. A predictor is a function h : X → R and the set of predictors is denoted as H . We introduce two losses l : R ×Y → [0 , 1] and l s : R ×Y → 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 ) = E ( s,z ) ∼ D [ l ( h ( s ) , z )] and empirical loss as r ( h ) = 1 m ∑ m a =1 l ( h ( s a ) , z a ) . The expected loss is the central quantity which we are interested in but cannot be unobserved in general. Therefore, the currently popular approach to statistical learning is to derive an upper bound of expected loss.

PAC-Bayesian Bounds We will follow the PAC-Bayesian approach to the generalization problem. Theorem 1 (Theorem 1.2.6. in [13]) . For any positive λ and δ ∈ (0 , 1) , with at least 1 -δ in the probability of training samples, for all distribution q on space H , we have

The KL divergence is D KL ( q ∥ p ) = E x ∼ q [log q ( x ) p ( x ) ] and the auxiliary function is defined as

Although there are huge amounts of research on PAC-Bayesian approach, all of them rely on the KL divergence between output distribution q of stochastic predictor learned from certain algorithm and a prior p , therefore the analysis of this KL divergence is a natural question.

Bayesian Learning We define Gibbs measure as p λ ( θ ) = 1 Z λ e -λr ( θ ) p ( θ ) and partition function as Z λ = E θ ∼ p [ e -λr ( θ ) ] . The Gibbs measure is related to posterior in Bayesian inference [24], and minimize right-hand side in theorem 1, inducing a much tighter bound:

Despite the fact that expected loss is optimized, Bayesian learning bears significantly more difficulty than popular optimization-based algorithms. It is noted that the sampling from a 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-the-art is 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 θ ∈ Θ . The expected loss R ( h ) and empirical loss r ( h ) could therefore be regarded as functions of parameter θ . An optimization flow is a function f : Θ → Θ , which typically relies on the training data to update the model's parameters θ 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, or more sophisticated methods detailed in Appendix B. We require the Jacobian determinant J f ( θ ) = det( ∇ f ( θ )) exists and the optimization flow is a bijective, i.e. the inverse f -1 exists everywhere. This is satisfied for gradient flow and gradient descent when step size is smaller than 1 L . More situations where this assumption holds is discussed in Appendix B. The starting point of optimization is initialized from a prior distribution p ( θ ) , and we define a corresponding energy V ( θ ) = -log p ( θ ) + C where C is a parameter independent constant. The output distribution q could be characterized as follows:

The KL divergence between output distribution and prior can be derived as follows:

Based on a physics analog of Helmholtz free energy change in isothermal process, we define energy term and entropy term separately:

The energy term only depends on the prior. If Gaussian prior is used, ∆ f 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 lead to decrease in KL divergence, because the decrease in energy term is offset by the change in entropy term.

The entropy term determines the gain or loss of entropy caused by an optimization procedure. The pure gain means that the optimization flow maps a region of parameter space into another region with larger area. For example, consider the simple function f ( x ) = 2 x in a one-dimensional setting. The Jacobian of this function, which describes how much the function stretches or compresses space, is constant at 2 . Imagine we start with an input x that is uniformly distributed within the interval [0 , 1] . After applying the function, the output f ( x ) 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 is inevitable. In Section 7, we show that for gradient 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 F.2.

Finally, we have

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 are based on the optimization flow, energy change, and loss of entropy discussed in Section 3.

The most simple example of TransBL 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

The above expectation could be regarded as attaching an additional weight to the stochastic predictor obtained from ordinary optimization, and we should use a weighted average of the results obtained by different initialization. This process doesn't involve any more training of parameter θ than the optimization procedure that TransBL is based on, though additional computation might be required for the value of ∆ f V ( θ ) and ∆ f S ( θ ) .

The unnormalized weight of above importance sampling is w λ ( θ ) = e -λr ( f ( θ )) e -∆ f V ( θ ) e ∆ f S ( θ ) . That means not all results obtained by training are 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:

We define the above value as sampling efficiency, and a non-diminishing efficiency at certain limit indicates that Bayesian learning is at most a constant factor slower than optimization. The numerator of the above expression is determined by the partition function because E θ ∼ p [ w λ ( θ )] = Z λ . Therefore, only the denominator depends on the optimization procedure, and we expect E θ ∼ p [ w 2 λ ( θ )] to be as small as possible for efficient Bayesian learning. Notice that this quantity again relies on energy change ∆ f V ( θ ) and entropy change ∆ f S ( θ ) . We will give theoretical analysis of these two values for infinitely wide neural network in Section 7.

4.1 An Illustrative Example

For illustration purpose, we consider a univariate loss function that presents two distinct global minima with zero loss. One of these is characterized as a 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 a small weight to the solution found within the sharp minimum. Consequently, TransBL can adeptly recreate the posterior, as shown in the Figure 6b. Due to space limitations, all figures for the illustration are moved to Appendix F.1.

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 a good estimation of posterior. A formal argument relies on applying Donsker and Varadhan's variational formula [17] to obtain the following equation:

The first term is independent of training procedure and only relies on the definition of prior. The second term measures the KL divergence between output distribution and Gibbs measure. This is also the gap between expected loss bound obtained by some training methods and the optimal one. Notice that the above KL divergence could also be expressed in terms of weight for TransBL, energy change, and entropy loss:

5.2 Efficiency of Optimization Flows for TransBL

In this section, we show that, in order to 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 a different divergence.

In order to see this, we define χ 2 divergence as D χ 2 ( p ∥ q ) = E x ∼ q [( p ( x ) q ( x ) ) 2 ] -1 .

Theorem 2. For sampling efficiency of TransBL whose proposal distribution is q , we have:

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 it minimizes the expected loss upper bound if and only if the sampling efficiency of TransBL is 1. A more general correspondence between generalization and sampling efficiency is desired. However, the complication lies in different divergences used in two quantities. By comparing Eq. (4) and Eq. (6),

we see that χ 2 divergence is used for measuring sampling efficiency. Although an upper bound of KL divergence D KL ( p λ ∥ q ) could be obtained, this is not directly comparable to the D KL ( q ∥ p λ ) used in generalization bound because of asymmetry of KL divergence. That difference also justifies the interpolation between optimization and Bayesian learning shown in the next section.

6 Interpolation of Optimization and Bayesian Learning

We recall that the output distribution q for the optimization and the posterior p λ for Bayesian learning are connected by an additional weight w λ : p λ ( θ ) = 1 Z λ w λ ( f -1 ( θ )) q ( θ ) . These two distributions have distinct properties. The distribution q bears worse generalization, but is easy to sample from given an established optimization oracle. The posterior p λ , on the other hand, is better shaped, but hard to sample from. Therefore, a natural question is whether we can find a distribution p β λ 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 a function v β : R + → R + where β is the parameter used for interpolation, we can define a modified weight as v β ( w λ ( θ )) . The interpolation distribution is

Note that the superscript doesn't mean power but means the interpolation is controlled by β . The expectation of a function F ( θ ) on p β λ could be computed similarly as Eq. (2):

We note Eq. (4) still holds when we change q into p β λ , and the gap between the expected loss bound obtained by p β λ and the posterior is:

Notice that the above identity degenerates into Eq. (5) if we use v β ( · ) = 1 , which corresponds to the original optimization.

The sampling efficiency of p β λ is again defined as the ratio of effective sample size and sample size: eff β λ = ( E θ ∼ p [ v β ( w λ ( θ )])) 2 / E θ ∼ p [ v β ( w λ ( θ )) 2 ] and enjoys a relation similar to Theorem 2:

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 w λ ( θ ) . Asmall 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 to large value of w λ ( θ ) , according to Eq. (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 a distribution with heavier tail than posterior. Due to 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 deteriorate too much.

Searching for the Pareto optimality among various trade-off methods is not a concern in this paper, although obviously this is an important issue and is a subject for future work. In this paper, we only consider a simple weight clipping method: v β ( w ) = { w w ≤ β β otherwise .

At the limit of β → 0 , p β λ degenerates to output distribution q of optimization. At the limit of β →∞ , p β λ converges to posterior p λ . A lower bound of sampling efficiency eff β λ ≥ ( Z β λ /β ) 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 β ∈ (0 , ∞ ) , D KL ( p β λ ∥ p λ ) and eff β λ are both monotonically decreasing functions of β .

7 Overparameterized Neural Network

The energy change ∆ f V ( θ ) and entropy change ∆ f S ( θ ) play a pivot role in previous sections. All quantities of interest, including KL divergence in the PAC-Bayesian bounds, weights in TransBL and sampling efficiency, depend on these two terms. Therefore, it is essential to understand the dynamics of these quantities. In this section, we will explore some simple infinitely wide neural network. In particular, this analysis will provide us insight for further developments in the deep Bayesian learning and PAC-Bayesian bounds.

7.1 Network Definition

We consider a feedforward network with d fully connected layers. The parameter θ = vec ( W 1 , . . . , W d ) is a vector of flattened weights. For a specific input s a , the forward propagation of the network is defined as follows:

The output y d ( θ, s a ) is a scalar. The point-wise function σ i ( y ) is an activation function with bounded first to fourth order derivative. The hidden layer width at i -th layer is n i such that y i ( θ, s a ) , x i ( θ, s a ) ∈ R n i . W i ∈ R n i × n i -1 are trainable weights. The coefficient c i = 1 / √ n i -1 is 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, 59]. The network is initialized with Gaussian prior, such that ( W i ) j,k ∼ N (0 , 1) .

We introduce the following notation to simplify the result. For a multi-variable scalar function f ( X ) , ∇ X f ( X ) is a row vector when X is a vector. If X is a matrix, then ( ∇ X f ( X )) i,j = ∂f ( X ) ∂X j,i and ∇ X f ( X ) share the same shape as X T . For a multi-variable vector-value function F ( X ) , we define ( ∇ X F ( X )) i,j = ∂F i ( X ) ∂X j .

The network is trained by gradient flow or infinitesimal step size gradient descent. Let output vector be ( Y ( θ )) a = y d ( θ, s a ) and total loss function as L ( Y ) = ∑ m a =1 l s ( Y a , z a ) , the gradient flow is defined as:

7.2 Energy Term

The energy for the prior is defined as V ( θ ) = ∑ d i =1 1 2 ∥ W i ∥ 2 Fro . We define y ( i ) d ( θ, s a ) = Tr( ∇ W i y d ( θ, s a ) W i ) and the dynamics of V ( θ ( t )) is:

We find that y ( i ) d ( θ, s a ) can be computed by a 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 i from i -th layer.

For deriving the dynamics of ( Y ( i ) ( θ )) a = y ( i ) d ( θ, s a ) , we define Θ ( i ) ( θ ) = ( ∇ θ Y ( i ) ( θ ))( ∇ θ Y ( θ )) T and it follows

We note that the definition of Y ( d ) follows Y = Y ( d ) , therefore Θ ( d ) , as a gradient Gram matrix, is just neural tangent kernel (NTK) which has been studied in [39, 3, 18].

7.3 Entropy Term

We write down the gradient flow in an abstract form d d t θ ( t ) = -g ( θ ( t )) and denote the solution as θ ( t ) = f ( t, θ (0)) = θ (0) -∫ t 0 g ( θ ( t )) dt . This gradient flow is a special example of optimization flow in Section 3. Here the entropy change is ∆ t S ( θ ) = log J f ( t, · ) ( θ ) . According to Liouville formula, we have d d t (∆ t S ( θ (0))) = -div g ( θ ( t )) = -Tr( ∇ g ( θ ( t ))) . Given that the gradient has form g ( θ ) = -( ∇ Y L ( Y ( θ )) ∇ θ Y ( θ )) T , the entropy change could be formulated as follows:

The content in the first trace term in Eq. (9) is called Gauss-Newton matrix, which only represents curvature information of loss l s . The matrix H a = ∇ θ ∇ θ Y a is the Hessian of y d ( θ, s a ) . 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 on initialization and training. For all 1 ≤ i ≤ j < k ≤ d , α = 1 , . . . , n j , we introduce auxiliary vectors ( ξ ( i,j ) ) α = c 2 i ∥ x i -1 ∥ 2 ∥∇ y i ( y j ) α ∥ 2 and γ ( j ) k = ( ∇ x j y k ) σ ′′ j ( y j ) . Notice that we left out parameters ( θ, s a ) for concision.

Theorem 4. Let ⊙ be element-wise product, we have

The above formula is a multi-variable version of ( f d -1 ◦ · · · ◦ f 1 ) ′′ = ∑ d -1 j =1 ( f d -1 ◦ · · · ◦ f j +1 ) ′ f ′′ j (( f j -1 ◦ · · · ◦ f 1 ) ′ ) 2 . ξ ( i,j ) plays the role of squared gradient and is a kind of diagonal NTK in middle layers.

We define (Γ ( i ) ( θ )) a = γ ( i ) d ( θ, s a ) and matrices Φ ( i ) ( θ ) = ( ∇ θ Γ ( i ) ( θ ))( ∇ θ Y ( θ )) T . Then we have for all 1 ≤ i < d ,

7.4 Infinite Width Limit

The dynamics in previous two sections can be dramatically simplified in infinite width limit.

Theorem 5. When hidden layers width n = n 1 , . . . , n d -1 approaches infinity, the random vectors Y ( i ) ( θ (0)) and Γ ( i ) ( θ (0)) at initialization converges in law to a multivariate Gaussian distribution with zero mean and covariance matrix Σ . Moreover, the following quantities converge in probability to constant and don't change during finite training time:

The specific formulas of Σ , Θ ( i ) , Φ ( i ) , and Ξ ( i,j ) are shown in Appendix D.

The last line of above limits is for all α = 1 , . . . , n j , therefore theorem 4 could be simplified to:

Theorem 5 demonstrates a function space picture, where the dynamics of Y ( i ) ( θ ( t )) and Γ ( i ) ( θ ( t )) only depend on these values themselves but not on internal dynamics of the network. The auxiliary vectors Y ( i ) ( θ ( t )) and Γ ( i ) ( θ ( t )) could be obtained by solving ODEs (8) and (10), d d t V and d d t S can be integrated according to Eqs. (7) and (9). Therefore, the KL divergence is accessible as an expectation in Eq. (1). Based on this KL divergence, we 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, we show the asymptotic behaviour under mild assumption on loss function.

10

0

10

1

Training Time

0

.

00

0

.

05

0

.

10

0

.

15

Training Loss

Function Space

Parameter Space

10

0

10

1

Training Time

-

0

.

2

-

0

.

1

0

.

0

V

(

θ

(

t

))

-

V

(

θ

(0))

Function Space

Parameter Space

10

0

10

1

Training Time

-

6

-

4

-

2

0

Function Space

Parameter Space

Corollary 6. If there exist C 1 , C 2 such that ∣ ∣ ∣ d d y l s ( y, t ) ∣ ∣ ∣ ≤ C 1 , ∣ ∣ ∣ d 2 d y 2 l s ( y, t ) ∣ ∣ ∣ ≤ C 2 , then D KL ( q t ∥ p ) ≤ O (1) t + O (1) t 2 , eff λ ≥ O (1) exp ( -O (1) t -O (1) t 2 ) for any finite t , where O (1) represents positive constant irrelevant to t , q t is the output distribution of an infinitely wide network after training for t length of time, and eff λ is the sampling efficiency of importance sampling with q t as proposal distribution.

8 Experiments

We first illustrate that Hessian trace doesn't vanish for overparameterized network and our analysis induces an efficient estimation of this value. Next, we verify our 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.

The synthetic dataset is constructed on a unit circle. We choose 20 values τ uniformly distributed between 0 to π and we let 2dimensional input to the network be s ( τ ) = [cos( τ ) sin( τ )] .

For each input s ( τ ) , we are concerned about Hessian trace Tr( ∇ θ ∇ θ y d ( θ, s ( τ ))) . Fast Hessian trace estimation by itself is a challenging problem, due to the high dimensional parameters and complex structure of the network. We

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

τ

-

4

-

2

0

2

4

Our Approximation

Hutchinson

-

10

0

10

follow previous works [36, 6, 78] and adopt Hutchinson'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 a fixed value to reduce computation. More details are shown in Appendix F.3. We apply the Hutchinson's method with 1000 independent random Rademacher vectors. The 3 σ interval is plotted in the Figure 1. It is shown that our approximation aligns with ground truth well. Based on this approximation, the distribution of Hessian is also calculated in Figure 1.

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 ( τ ) = 1 2 sin(3 τ ) and loss to be mean squared error (MSE). For dynamics in parameter space, we run SGD with finite step size 0 . 01 and mini-batch size 1 . For dynamics in function space, we solve Eqs. (7) to (10) with fixed matrices Θ ( i ) , Φ ( i ) and Ξ ( i,j ) . The result is plotted in Figure 2, where for SGD, the training time is defined as step size times iteration number. The ground truth of Hessian trace on one input is again estimated by Hutchinson's method. We can see from all three kinds of outputs that the function space

dynamics produces similar output as 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 [75]. 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 l ( y, t ) = 1 / (1 + exp( yt )) and surrogate loss l s ( y, t ) = log(1 + exp( -yt )) for gradient descent. For Gibbs measure, we fix λ = 180 . The entropy change is approximately evaluated by integrating Eq. (9) with finite step size and fixed Θ ( d ) .

Wetrain 10 5 independent network and the results are shown in Figure 3. The shadow region represents 3 σ 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 Z λ ≈ -43 . 6 .

2

4

6

8

10

Training Time

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

1

.

0

GD Train

GD Test

TransBL Test

-

120

-

100

-

80

-

60

-

40

log(

w

)

No Training

Trained

-

80

-

70

-

60

-

50

-

40

log

β

10

-

2

10

-

1

10

0

Sampling Efficiency

0

.

750

0

.

775

0

.

800

0

.

825

0

.

850

Test Accuracy

We also show a clear trade-off between sampling efficiency and test accuracy for TransBL with weight clipping, where the parameter β plays a pivotal role in balancing these aspects. When β is set high, the model leans towards a Bayesian posterior framework. Although this setting noticeably reduces the number of effective samples obtained, we observed an enhancement in the model's generalization ability, as depicted in the right figure of Figure 3. Conversely, with a smaller β , the model's behavior tends to resemble that of a typical deep ensemble, indicating a high sampling efficiency but worse generalization.

9 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 χ 2 divergence, which determines sampling efficiency for importance sampling, are related to the change of energy and entropy loss during training. The dynamics of these two quantities in DNNs are studied, leading to a 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-based optimization as proposal distribution, we bridge the gap between optimization and Bayesian learning problems, and provide a flexible interpolation for accuracy-computation trade-off.

References

Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring. In ICML , pages 1591-1598. PMLR, 2012.

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.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via overparameterization. In ICML , pages 242-252. PMLR, 2019.

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.

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.

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):1-34, 2011.

Peter Bartlett, Dylan J Foster, and Matus Telgarsky. Spectrally-normalized margin bounds for neural networks. arXiv preprint arXiv:1706.08498 , 2017.

Jeremy Bernstein and Yisong Yue. Computing the information content of trained neural networks. arXiv preprint arXiv:2103.01045 , 2021.

Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In ICML , pages 1613-1622. PMLR, 2015.

Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems , volume 33, pages 1877-1901, 2020.

Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems , 32:10836-10846, 2019.

Yuan Cao and Quanquan Gu. Generalization error bounds of gradient descent for learning overparameterized deep relu networks. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 3349-3356, 2020.

O Catoni. Pac-bayesian supervised classification: The thermodynamics of statistical learning. institute of mathematical statistics lecture notes-monograph series 56. IMS, Beachwood, OH. MR2483528 , 5544465, 2007.

Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In ICML , pages 1683-1691. PMLR, 2014.

Eugenio Clerico, George Deligiannidis, and Arnaud Doucet. Conditional gaussian pac-bayes. arXiv preprint arXiv:2110.11886 , 2021.

Arnak S Dalalyan. Theoretical guarantees for approximate sampling from smooth and log-concave densities. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 79(3):651-676, 2017.

Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time-iii. Communications on pure and applied Mathematics , 29(4):389-461, 1976.

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.

David Duvenaud, Dougal Maclaurin, and Ryan Adams. Early stopping as nonparametric variational inference. In Artificial Intelligence and Statistics , pages 1070-1077. PMLR, 2016.

Ethan Dyer and Guy Gur-Ari. Asymptotics of wide networks from feynman diagrams. In International Conference on Learning Representations , 2020.

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.

Cong Fang, Hanze Dong, and Tong Zhang. Mathematical models of overparameterized neural networks. Proceedings of the IEEE , 109(5):683-703, 2021.

MA Ganaie, Minghui Hu, et al. Ensemble deep learning: A review. arXiv preprint arXiv:2104.02395 , 2021.

Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julien. Pac-bayesian theory meets bayesian inference. In Advances in Neural Information Processing Systems , volume 29, 2016.

Alison L Gibbs and Francis Edward Su. On choosing and bounding probability metrics. International statistical review , 70(3):419-435, 2002.

Peter W Glynn and Chang-han Rhee. Exact estimation for markov chain equilibrium expectations. Journal of Applied Probability , 51(A):377-389, 2014.

Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size-independent sample complexity of neural networks. In Conference On Learning Theory , pages 297-299. PMLR, 2018.

Alex Graves. Practical variational inference for neural networks. Advances in neural information processing systems , 24, 2011.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations , 2020.

Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In ICML , pages 1225-1234. PMLR, 2016.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Jeremy Heng and Pierre E Jacob. Unbiased hamiltonian monte carlo with couplings. Biometrika , 106(2): 287-302, 2019.

Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. Stochastic variational inference. Journal of Machine Learning Research , 14(5), 2013.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks , 2(5):359-366, 1989.

Zhengmian Hu and Heng Huang. On the random conjugate kernel and neural tangent kernel. In International Conference on Machine Learning , pages 4359-4368. PMLR, 2021.

Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation , 18(3):1059-1076, 1989.

Pavel Izmailov, Sharad Vikram, Matthew D Hoffman, and Andrew Gordon Gordon Wilson. What are bayesian neural network posteriors really like? In ICML , volume 139, pages 4629-4640. PMLR, 2021.

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 (Statistical Methodology) , 82(3):543-600, 2020.

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems , volume 31, 2018.

Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.10196 , 2017.

Pitas Konstantinos, Mike Davies, and Pierre Vandergheynst. Pac-bayesian margin bounds for convolutional neural networks-technical report. arXiv preprint arXiv:1801.00171 , 2017.

Jaehoon Lee, Jascha Sohl-dickstein, Jeffrey Pennington, Roman Novak, Sam Schoenholz, and Yasaman Bahri. Deep neural networks as gaussian processes. In ICLR , 2018.

Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl-Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems , 32:8572-8583, 2019.

Jason D Lee, Max Simchowitz, Michael I Jordan, and Benjamin Recht. Gradient descent only converges to minimizers. In Conference on learning theory , pages 1246-1257. PMLR, 2016.

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:1806.05159 , 2018.

Etai Littwin, Tomer Galanti, and Lior Wolf. On random kernels of residual architectures. In Uncertainty in Artificial Intelligence , pages 897-907. PMLR, 2021.

Chaoyue Liu, Libin Zhu, and Mikhail Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems , 33, 2020.

Stephan M, t, Matthew D. Hoffman, and David M. Blei. Stochastic gradient descent as approximate bayesian inference. Journal of Machine Learning Research , 18:1-35, 2017.

David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computation , 4 (3):448-472, 1992.

Dougal Maclaurin, David Duvenaud, and Ryan P Adams. Early stopping is nonparametric variational inference. arXiv preprint arXiv:1504.01344 , 2015.

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 arXiv:1804.11271 , 2018.

Andreas Maurer. A note on the pac bayesian theorem. arXiv preprint cs/0411099 , 2004.

David A McAllester. Some pac-bayesian theorems. Machine Learning , 37(3):355-363, 1999.

Chris Mingard, Guillermo Valle-Pérez, Joar Skalse, and Ard A Louis. Is sgd a bayesian sampler? well, almost. Journal of Machine Learning Research , 22, 2021.

Radford M Neal. Bayesian learning for neural networks , volume 118. Springer Science & Business Media, 2012.

Kirill Neklyudov and Max Welling. Orbital mcmc. In International Conference on Artificial Intelligence and Statistics , pages 5790-5814. PMLR, 2022.

Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of over-parametrization in generalization of neural networks. In ICLR , 2018.

Theodore Papamarkou, Jacob Hinkle, M Todd Young, and David Womble. Challenges in markov chain monte carlo for bayesian neural networks. arXiv preprint arXiv:1910.06539 , 2019.

Daniel Park, Jascha Sohl-Dickstein, Quoc 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.

Konstantinos Pitas. Dissecting non-vacuous generalization bounds based on the mean-field approximation. In International Conference on Machine Learning , pages 7739-7749. PMLR, 2020.

Matthias Seeger. Pac-bayesian generalisation error bounds for gaussian process classification. Journal of machine learning research , 3(Oct):233-269, 2002.

Daniel Seita, Xinlei Pan, Haoyu Chen, and John Canny. An efficient minibatch acceptance test for metropolis-hastings. In Proceedings of the Twenty-Seventh International Joint Conference oncArtificial Intelligence , pages 5359-5363, 7 2018.

Mariia Seleznova and Gitta Kutyniok. Neural tangent kernel beyond the infinite-width limit: Effects of depth and initialization. In International Conference on Machine Learning , pages 19522-19560. PMLR, 2022.

Ruoqi Shen and Yin Tat Lee. The randomized midpoint method for log-concave sampling. Advances in Neural Information Processing Systems , 32, 2019.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

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-19237, 2020.

Gerald Teschl. Ordinary differential equations and dynamical systems , volume 140. American Mathematical Soc., 2012.

Achille Thin, Yazid Janati El Idrissi, Sylvain Le Corff, Charles Ollion, Eric Moulines, Arnaud Doucet, Alain Durmus, and Christian X Robert. Neo: Non equilibrium sampling on the orbits of a deterministic transform. Advances in Neural Information Processing Systems , 34, 2021.

Leslie G Valiant. A theory of the learnable. Communications of the ACM , 27(11):1134-1142, 1984.

Twan Van Laarhoven. L2 regularization versus batch and weight normalization. arXiv preprint arXiv:1706.05350 , 2017.

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge University Press, 2019.

Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In ICML , pages 681-688. Citeseer, 2011.

Christopher KI Williams. Computing with infinite networks. Advances in neural information processing systems , pages 295-301, 1997.

Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Advances in Neural Information Processing Systems , volume 33, pages 4697-4708, 2020.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint arXiv:1902.04760 , 2019.

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems , 32, 2019.

Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE International Conference on Big Data (Big Data) , pages 581-590. IEEE, 2020.

Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM , 64(3):107-115, 2021.

Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wilson. Cyclical stochastic gradient mcmc for bayesian deep learning. In ICLR , 2020.

Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P. Adams, and Peter Orbanz. Non-vacuous generalization bounds at the imagenet scale: a PAC-bayesian compression approach. In ICLR , 2019.

A Related Works

Generalization of Overparameterized Network Generalization for neural network has been studied based on stability [11, 12] norm [7, 57, 41, 27, 45], compression [4] , and special function class [5, 2]. [21, 81] achieves non-vacuous PAC-Bayesian generalization bounds based compression and perturbation robustness. Clerico et al. [15] derives PAC-Bayes bound and learning method for training a Gaussian process. A realizable PAC-Bayesian bound is also derived in [8] with sampling realizable networks instead of training by gradient descent. Mean field theory [22, 60, 66] could also be applied to obtain generalization bound.

Early Stopping Restricting training time is a popular trick in deep learning. Related theoretical explanation includes a 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, our 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 Acomprehensive 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 to this fields in two aspect. First, we show that Hessian trace doesn't vanish in infinite width limit. Second, we propose a novel Hessian trace estimation method in Appendix F.3, which is much more efficient than commonly used Hutchinson'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 to run MCMC till converge. Several orbit MCMC [68, 56] has been explored, seeking to incorporate non-equilibrium map into the transition kernel to 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 χ 2 divergence as discussed in Appendix F.2. Telescopic sum argument [26] with coupling technique [38, 32] have also been proposed to tackle this problem. Another source of bias is inaccurate implementation of transfer kernel. The discretization error of various Langevin dynamics method has been analyzed [1, 16, 64]. The bias from noisy loss function [62] and noisy gradient [72, 14] could be reduced but could not vanish. Papamarkou et al. [58] 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 f exist everywhere. This assumption is not as restrictive as it seems.

For gradient descent, f ( θ ) = θ -h ∇ l . If l ∈ C 2 , ∇ l is L-Lipschitz and h ≤ 1 /L , then f is 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.

For gradient flow, d d t θ = -∇ l . Our assumption immediately follows existence and uniqueness of ODE solution. For example, if l ∈ C 2 and ∇ l is L-Lipschitz, then f t ( θ ) is defined for all t and θ (Corollary 2.6 [67]) and is a diffeomorphism (Theorem 2.10 [67]).

Momentum SGD and Adagrad always satisfy our assumption on enriched state space (Appendix F.2).

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 7, we only analyze the infinitesimal step-size. The main reason for this choice is that it allows the training dynamics of 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, it 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, 73, 42, 51, 18, 39, 43, 76, 3]. The infinite-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, 46, 63] have tried to understand the impact of finite width, known as "finite-width corrections". Delving into this adds complexity to the theoretical analysis, and hence it's left as a 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.

C Additional Proofs

Proof of Theorem 2. First we show that sampling efficiency is determined by the χ 2 divergence.

The inequality in Eq. (6) is from Theorem 5 in [25].

Proof of Theorem 3. Based on some algebra and calculus, we have

For sampling efficiency, we have

Proof of Eq. (27). We discard the parameter θ, s a without loss of generality.

We finish the proof by combining above results.

Proof of Theorem 4.

Only σ ′ j ( y j ) for i ≤ j < d depends on W i in the above formula.

D Analysis of Infinitely Wide Network

We first give the analytic form of limits in Theorem 5.

The covariant matrix Σ is defined recursively. We first define series of Gaussian random variables y ( i ) k ( s a ) , γ ( i ) k ( s a ) , where variables with different subscript k are independent. We also let y k ( s a ) = y ( d ) k ( s a ) . The covariance of y ( i ) k ( s a ) and y ( j ) k ( s b ) are denoted as Σ k [ y ( i ) ( s a ) , y ( j ) ( s b )] . Similarly,

we can define ( i ) ( j ) ( i ) ( j )

In the infinite width limit, the covariance at initialization of ( Y ( i ) ( θ (0))) a and ( Y ( j ) ( θ (0))) b is Σ d [ y ( i ) ( s a ) , y ( j ) ( s b )] . Similarly, the covariance at initialization of ( Y ( i ) ( θ (0))) a and (Γ ( j ) ( θ (0))) b is Σ d [ y ( i ) ( s a ) , γ ( j ) ( s b )] , the covariance at initialization of (Γ ( i ) ( θ (0))) a and (Γ ( j ) ( θ (0))) b is Σ d [ γ ( i ) ( s a ) , γ ( j ) ( s b )] .

Other matrices limits are

̸

Before embarking on the proof, we first give an intuition of above results. We first notice that y ( i ) d ( θ, s a ) can be defined recursively:

̸

̸

We can give a similar definition by forward propagation for γ ( i ) j ( θ, s a ) = ( ∇ x i y j ( θ, s a )) σ ′′ i ( y i ( θ, s a )) , which covers the case when j ≤ i :

Therefore, both y ( i ) d and γ ( i ) d can be regarded as output from some other "neural network". We compares these forward pass with original neural network as follows:

Both y ( i ) d and γ ( i ) d propagate in a same way as original network until i -th layer. After that, y ( i ) d propagate by linear transformation, with matrix diag( σ ′ j ( y j ( θ, s a ))) being controlled by original network. γ ( i ) d uses a different activation function σ ′′ i which is different than original network with σ i in i -th layer, and propagate by linear transformation in a similar way as y ( i ) d .

When we consider the y d , y ( i ) d and γ ( i ) d as output from a multiple branch neural network, the Σ in Eq. (21) is just conjugate kernel, and Θ ( i ) , Φ ( i ) , Ξ ( i,j ) in Eq. (22) are fragment of NTK of the larger network.

Proof of Theorem 5. We first prove the convergence to Gaussian distribution. It follows the proof of convergence of Θ ( i ) , Φ ( i ) , Ξ ( i,j ) . Finally we prove that these values don't change during training.

We note that in the forward propagation as shown in Eqs. (27) to (29), we only need to 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 is a non-linear element-wise vector operation, e.g. diag( σ ′ j ( y j ( θ, s a ))) y ( i ) j ( θ, s a ) = σ ′ j ( y j ( θ, s a )) ⊙ y ( i ) j ( θ, s a ) .

Therefore, the calculation of y d , y ( i ) d and γ ( i ) d 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 σ , the general result in [76] 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 is left is 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 that has been studied in [39, 3, 18]. Moreover, (Ξ ( i,j ) ) α is just the diagonal NTK if we consider a sub neural network which coincide with original network in previous j -1 layers, but use one single hidden unit ( y j ) α at j -th layer as output. Therefore, we only need to take care about Θ ( i ) and Φ ( i ) for i < d .

The first term is easy to calculate and shares similar form with classical NTK after computing the trace Tr ( ∇ W j y ( i ) d ( s a ) T ∇ W j y d ( s b ) ) .

Combining above results, we have

The limit of first term is a.s. lim ∇ y j y d ( s a ) ∇ y j y d ( s b ) T = ∏ d -1 k = j E [ σ ′ k ( y k ( s a )) σ ′ k ( y k ( s b ))] . In order to calculate the second term c 2 j ( x ( i ) j -1 ( s a ) x j -1 ( s b )) , we note that E w ∼N (0 ,I ) [ ww T ] = I and ( W j ) α is a standard normal random vector, then we have

so the limit is Σ j [ y ( i ) ( s a ) , y ( s b )] .

All other terms comes from second order derivative of non-linear unit σ k and doesn't show up in classical NTK result. We note in Tr ( ( c j x j -1 ( s a )( ∇ x k y d ( s a ) T ⊙ y ( i ) k ( s a ) ⊙ σ ′′ k ( y k ( s a ))) T ∇ y j y k ( s a )) T ∇ W j y d ( s b ) ) = ( c 2 j x j -1 ( s a ) T x j -1 ( s b ))( ∇ y j y d ( s b ) ∇ y j y k ( s a ) T ( ∇ x k y d ( s a ) T ⊙ y ( i ) k ( s a ) ⊙ σ ′′ k ( y k ( s a )))) , c 2 j x j -1 ( s a ) T x j -1 ( s b ) = O (1) converge to a fixed value, in ∇ y j y d ( s b ) ∇ y j y k ( s a ) T ⊙∇ x k y d ( s a ) = ∇ x k y d ( s b ) diag σ ′ k ( y k ( s a )) ∇ y j y k ( s b ) ∇ y j y k ( s a ) T ⊙∇ x k y d ( s a ) , ∇ y j y k ( s b ) ∇ y j y k ( s a ) T converges to fixed values times identity matrix, and ∇ x k y d ( s b ) T ∇ x k y d ( s b ) also converges to fixed values times identity matrix. The rest terms converge to its expectation which contains y ( i ) k ( s a ) , σ ′′ k ( y k ( s a )) and σ ′ k ( y k ( s b )) .

Note that γ ( i ) d ( θ, s a ) = ( ∇ x i y d ( θ, s a )) σ ′′ i ( y i ( θ, s a )) and when j ≤ i the value y i ( θ, s a ) depends on W i . Therefore we conduct a separate discussion. When j > i , the limit value could be derived in a way similar to (Θ ( i ) ) a,b . For j ≤ i ,

The summations are again from the second order derivative of non-linear unit and the limit could be derived in a similar way to Φ ( i ) . We focus on first term is somewhat different since it depends on third order derivative of non-linear unit.

Finally we prove that Θ ( i ) , Φ ( i ) , Ξ ( i,j ) are constant during training. Since these values are all inner product of gradients, we just need to show gradients ∇ W j y d ( θ, s a ) , ∇ W j y ( i ) d ( θ, s a ) and ∇ W j γ ( i ) d ( θ, s a ) doesn't change too much. We write down the formula of these gradient without the

summation terms which could be controlled similarly.

We rely on following Local Lipschitz property to prove the stability of gradients during training.

Lemma 7 (Local Lipschitz Condition) . Let v j -1 ( θ ) to denote any of following values:

there is a constant K > 0 such that for every R > 0 , with probability increasing to 1 at limit of n →∞ , or equivalently, with high probability over random initialization (w.h.p.o.r.i.) of θ , for any θ 1 , θ 2 ∈ B ( θ, R ) , where B ( θ, R ) = { θ ′ | ∥ θ ′ -θ ∥ < R } , the following holds

Moreover, let v j ( θ ) to denote any of following values:

there is a constant K > 0 such that for every R > 0 , w.h.p.o.r.i. of θ , for any θ 1 , θ 2 ∈ B ( θ, R ) ,

By using the norm upper bound in Lemma 7, we have that for any R > 0 and θ ∈ B ( θ (0) , R ) , the gradient norm Tr ( ( ∇ W j y d ( θ, s a )) T ∇ W j y d ( θ, s a ) ) = c 2 j ∥ σ j -1 ( y j -1 ( θ, s a )) ∥ 2 ∥ ∥ ∇ y j y d ( θ, s a ) ∥ ∥ 2 ≤ K is bounded w.h.p.o.r.i. . For any finite training time, R could be selected large enough to ensure θ ( t ) ∈ B ( θ (0) , R ) during training. The norm bound of difference in Lemma 7 also ensures that the change of gradient norm is bounded by 1 √ n K ∥ θ ( t ) -θ (0) ∥ in B ( θ (0) , R ) w.h.p.o.r.i. which is diminishing at the limit n → 0 .

Proof of Lemma 7. All vectors v j in Lemma 7 could be computed sequentially by three kinds of operations: matrix multiplication ζ = Wε , element-wise non-linearity ε = ϕ ( ζ ) with ϕ having bounded derivative and element-wise multiplication ε = ϕ ( ζ ) ⊙ ζ ′ , where the element-wise non-linear function ϕ is bounded and Lipschitz.

For the first kind of operation, given that W follows standard normal distribution, we have the ζ is also normally distributed and enjoy following high probability bound.

Eq. (2.19) of [71] For n dimensional standard Gaussian random vector ζ ∼ N (0 , I n ) , for any t

Therefore we have that for ζ = Wε , ∥ ζ ∥ ≤ K √ n ∥ ε ∥ for some K with high probability.

For element-wise non-linearity ε = ϕ ( ζ ) where ζ is a standard Gaussian random vector with i.i.d. coordinates, we derive the Chernoff bound as follows.

For element-wise multiplication ε = ϕ ( ζ ) ⊙ ζ ′ , since ϕ is bounded, ∥ ε ∥ ≤ O (1) ∥ ζ ′ ∥ all the time.

Therefore, we have ∥ ε ∥ = ∥ ϕ ( ζ ) ∥ ≤ K √ n w.h.p.o.r.i. for large enough K .

Recursively apply above three operations gives the bound of norm ∥ v j ∥ .

In order to control the norm of difference, we note that for matrix multiplication:

The first term could be controlled with high probability as ∥ W 1 ( ε 1 -ε 2 ) ∥ ≤ ∥ W 0 ( ε 1 -ε 2 ) ∥ + ∥ ( W 1 -W 0 )( ε 1 -ε 2 ) ∥ ≤ K √ n ∥ ε 1 -ε 2 ∥ + R ∥ ε 1 -ε 2 ∥ ≤ O (1) √ n ∥ ε 1 -ε 2 ∥ for large enough n , and W 0 is the parameter at initialization, R is the maximum distance between W 1 and W 0 .

For element-wise non-linearity ε = ϕ ( ζ ) , since ϕ is Lipschitz, we have the difference as

For element-wise multiplication ε = ϕ ( ζ ) ⊙ ζ ′ , since ϕ is bounded and Lipschitz, we have the difference as

The first term is reduced by noticing ζ ′ 1 has O (1) elements.

Recursively apply above three operations gives the bound of norm of difference.

Proof of Corollary 6. According to Theorem 5, in the infinite width limit, we have

We use O (1) to denote constants that is irrelevant to t. Given that ∇ Y L are bounded, we have

It follows that the Hessian is also bounded.

The time derivative and finite time change of energy and entropy is

The KL divergence is

The last inequality is because ∑ d i =1 Y ( i ) ( θ (0)) follows Gaussian distribution and the expectation of norm is finite.

The sampling efficiency could also be controlled.

The first inequality comes from r ( θ ) ≥ 0 . The last inequality follows from the fact that, for random vectors with tail similar to log-normal distribution, the expectation exists. More specifically, let σ i (Σ) is eigenvalues of covariance matrix Σ , we have E x ∼N (0 , Σ) [exp( c ∥ x ∥ )] ≤ ∏ i E y i ∼N (0 ,σ i (Σ)) [exp( c ∥ y i ∥ )] = ∏ i exp ( c 2 σ 2 i (Σ) / 2 ) (1 + erf ( cσ i (Σ) / 2))) .

E Additional Experiment Results

In Figure 4, we show the correlation between KL divergence and generalization gap during training. We also show how weight distribution evolves.

2

4

6

8

10

Training Time

0

1

2

3

4

5

D

KL

(

q

t

‖

p

)

KL Divergence

0

.

02

0

.

04

0

.

06

0

.

08

2

4

6

8

10

Training Time

-

120

-

100

-

80

-

60

-

40

log(

w

)

F Additional Discussion on Algorithms

F.1 An illustrative example of TransBL

(a) Prior

(b) Loss

(c) Posterior

(a) Gradient flow

(b) TransBL

transformed by adding a weight

In the presented motivating example, we explore a univariate loss function that presents two distinct global minima with zero loss. One of these is characterized as a sharp minimum, while the other represents a flat minimum. If the function were to be randomly initialized and then optimized, it might converge to either of the two localities. However, insights from PAC Bayesian indicate that the flat minimum is surrounded by a higher posterior probability density. A direct initialization from the prior, followed by training using gradient flow, often results in a significant deviation of the ensemble from the posterior. This is primarily because the optimization process fails to recognize the presence of a 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.1. In Figure 7a, with β = 0 . 2 , the curve leans more towards the optimization result, revealing a broader spread and

less-defined peaks. This shows a scenario where optimization has a stronger influence than Bayesian learning. Yet with β = 5 , the distribution is almost akin to what one would expect from a Bayesian posterior. In essence, the parameter β serves as a tuning knob, allowing us to interpolate and traverse the spectrum from pure optimization-driven results to outcomes heavily influenced by Bayesian learning. This interpolation mechanism offers a flexible approach to merge the strengths of both methodologies.

(a)

β

= 0

.

2

(b)

β

= 1

(c)

β

= 5

F.2 Optimization Algorithms

In the main paper, we analyze the change of energy ∆ V and entropy ∆ S 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 Eq. (9), the entropy change is related to the curvature of loss l s 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 a price of lower sampling efficiency.

Momentum SGD For parameter x and momentum v , the Momentum SGD with fraction γ , step size h and gradient oracle g can be formulated into following two steps:

By defining y = [ x v ] , Momentum SGD turns into a transform T ( y ) composed by two sub transforms T ( y ) = L 2 ◦ L 1 ( y ) that correspond to two lines separately.

The determinant of Jacobian for L 1 is det( γI ) = γ d , where d is dimension of parameters. The determinant of Jacobian for L 2 is 1 , therefore, the determinant of Jacobian for Momentum SGD is γ d .

Adagrad For parameter x , first-order momentum v and second-order momentum m , the Adagrad is:

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-preserving, and second step compress momentum by γ in each dimension. Therefore the determinant of Jacobian for Adagrad is γ d .

Both above two algorithms enjoy simple form of Jacobian determinant. However, the sampling efficiency degenerates since χ 2 divergence increases when we consider joint probability distribution of parameter and auxiliary variables, like momentum.

F.3 Approximate DNN Hessian Trace by Forward Propagation

We first recall that due to Theorem 4, we have

Our motivation is based on the stability of ξ ( i,j ) . We have shown in Theorem 5 that in the infinite width limit, ξ ( i,j ) converge to a vector with same elements and is stable during training. For the wide but finite width network, this property is largely preserved.

If we replace ξ ( i,j ) with its limit value, we obtain following formula:

In order to calculate the above value efficiently, we notice that above summation can be regarded as propagation of a tangent vector.

We first change the network definiton into

Notice that compared to original network, we add vectors b i , and when b i is set as 0 , the output is same as original network.

It can be easily seen that ∇ x j y d = ∇ b j y d , therefore we could just let σ ′′ j ( y j ( θ ( t ) , s a )) ( ∑ j i =1 (Ξ ( i,j ) ) a ) be the tangent vector for b i and let it propagates along the forward pass, the result tangent vector in output space is just an estimation of Tr( H a ( θ ( t ))) .

F.4 Final PAC Bayes bound

The 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 Figure 4) by solving the above ODE.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n**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\nQUESTIONS:\n- 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\nLIMITATIONS:\n-    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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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$.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nOnly some minor ones.\n\nMinor weakness:\n1. line 275, equation 9 is too long\n2. line 265, Equation (27) is not properly referenced.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nAlthough 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\nQUESTIONS:\nQ.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\nLIMITATIONS:\nThe structure of the paper is difficult to understand.\nThe theoriy that are ultimately obtained is unclear compared to the exisitng PAC Bayes bound.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\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\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\na) 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\nLIMITATIONS:\n- 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'.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}