Demystifying the Optimization and Generalization of Deep PAC-Bayesian Learning
Authors: Wei Huang
Chunrui Liu
Yilan Chen
Richard Yi Da Xu
Miao Zhang
Tsui-Wei Weng
Venue: Submitted to ICLR 2023 (2023)
Keywords: PAC-Bayes
Probabilistic Neural Netowrks
Neural Tangent Kernel

Abstract: In addition to being a successful generalization bound analysis tool, the PAC-Bayesian bound can also be incorporated into an objective function to train a probabilistic neural network, which we refer to simply as {\it PAC-Bayesian Learning}. PAC-Bayesian learning has been proven to be able to achieve a competitive expected test set error numerically, while providing a tight generalization bound in practice, through gradient descent training. Despite its empirical success, the theoretical analysis of deep PAC-Bayesian learning for neural networks is rarely explored. To this end, this paper proposes a theoretical convergence and generalization analysis for PAC-Bayesian learning. For a deep and wide probabilistic neural network, we show that when PAC-Bayesian learning is applied, the convergence result corresponds to solving a kernel ridge regression when the probabilistic neural tangent kernel (PNTK) is used as its kernel. Based on this finding, we further obtain an analytic and guaranteed PAC-Bayesian generalization bound for the first time, which is an improvement over the Rademacher complexity-based bound for deterministic neural networks. Finally, drawing insight from our theoretical results, we propose a proxy measure for efficient hyperparameter selection, which is proven to be time-saving on various benchmarks.

Under review as a conference paper at ICLR 2023
DEMYSTIFYING THE OPTIMIZATION AND GENERALIZA- TION OF DEEP PAC-BAYESIAN LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
PAC-Bayes has long been generalization analysis framework where the expected population error can be bounded by the sum of training eтror and the divergence between posterior and prior distribution. In addition to being a successful gener- alization bound analysis tool, the PAC-Bayesian bound can also be incorporated into an objective function train a probabilistic neural network, which we refer simply as PAC-Bayesian Learning. able to achieve competitive expected test set error numerically, while providing a tight generalization bound in practice. through gradient descent training. Despite its empirical success, the thcoretical analysis of deep PAC-Bayesian learning for neural networks is rarely explored. To this end, this paper proposes theoretical convergence and generalization analysis for PAC-Bayesian learning. For a deep and wide probabilistic neural network, we show that when PAC-Bayesian learning ap- plied, the convergence result corresponds to solving a kernel ridge regression when the probabilistic neural tangent kernel (PNTK) is used as its kernel. Based on this finding, we further obtain an analytic and guaranteed PAC- PAC-Bayesian generalization bound for the first time, which is an improvement over the Rademacher complexity- based bound for deterministic neural networks. Finally, drawing insight from our theoretical results, we propose proxy measure for efficient hyperparameter selection, which is proven to be time-saving on various benchmarks.
learning has been proven to be
INTRODUCTION
Deep learning has demonstrated powerful learning capability due to its over- preeparranttiiinaioon struc- ture, in which various network architectures have been responsible for its significant leap in per- formance (LeCun et al.] 2015). Over-fitting and complex hyperparameters are two of the major challenges in deep learning, hence designing generalization guarantees for deep networks is an important research goal (Zhang et al. (02T). Recently, a learning framework that trains a proba- bilistic neural network with PAC- Bayesian bound objective function has been proposed (Bégin ctt al. 2016 Dziugaite & Roy, 2017, Neyshabur et al., [017b] Raginsky et al.) [017]] Neyshabur Smith & Le| 2 Pérez-Ortiz et al. [2200] Guan & Lu (022), which is known as PAC-Bayesian learning. While providing a tight generalization bound, AAC-Bayesian learning has been proven to be able to achieve a competitive expected test set error (Ding et al. 2022).
al. 70l London 20t7|Smmthh
Furthermore, this generalization bound computed from the training data can obviate the need for splitting data into training, testing, and validation set, which is highly applicable for training deep network with scarce data (Pérez-Ortee Ortiz et al., 2020; Grünwald & Mehta 2020). Meanwhile, these advancements on PAC-Bayesian bounds have been widely adapted with different deep neural network structures including convolutional neural network (Zhou et al. 2018 Pérez-Ortiz et al. 20200) binary
activated multilayer networks (Letarte et al., 2019), partially aggregated neural networks (Biggs Guedj] 2020). and graph neural network (Liao et al. 2020).
Due to the impressive empirical success of PAC-Bayesian learning, there is increasing interest in understanding its theoretical properties. However, it is either restricted to specific technique variant such as Entropy-SGi which minimizes an objective indirectly by approximating stochastic gradient ascent the so-called local entropy (Dziugaite & Roy| 2018a) and differential privacy (Dziugaite & Roy| (2188), or relies heavily on empirical exploration (Neyshabur et al.) [017a] Dziugaite et al. 2020). To our best knowledge, there has been investigation far into why the training of

---

Under review as a conference paper at ICLR 2023
FAC-Bayesian learning is successful and why the PAC- MAC-Baaesi bound is tight unseen data after training. For example, it is still unclear when applying gradient descent to BAC-Baaesian learning:
Q1: How effective is gradient descent training on training set?
Q2: How tight is the generalization bound compared to those learning frameworks using non- probabilistic neural networks?
The answers to these questions can be highly non-trivial due to the inherent non-convex problem of over-paramettrrizatonn (Jain & Kar 2017) and additional randomness introduced by probabilistic neural networks (Specht 1990) as well as additional challenges brought by the divergence between posterior/prior distribution pairs known as Kullback-Leibler (KL) divergence. Nevertheless, this paper shows that it is possible answer the above questions by leveraging the recent advances in deep learning theory with over- wver-parameterized setting. It has been shown that wide networks optimized with gradient descent can achieve near-zero training error, and the critical factor that governs the training process is the neural tangent kernel (NTK), which can be proven to be unchanged during gradient descent training (Jacot al.] 2018). thus providing a guarantee for achieving global minimum (Du et al. 2019; Allen-Zhu et al. 2019). Under the PAC-Bayesian framework, NTK is no longer calculated from the derivative of the weights directly, but instead is calculated based on the gradient of the distribution parameters of the weights. We call this Probabilistic NTK (PNTK), based on which we build a convergence analysis to characterize the optimization process of PAC-Bayes learning. Due to the explicit solution obtained by optimization analysis, we further formulate the generalization bound of Bay-Baynn learning for the first time, and demonstrate its advantage by comparing with the theoretical generalization bound of learning framework with non-stochastic neural networks
(Arora et al. 2019a]
Cao
Gu 2019 Hu et al. 2019).
We summarize our contributions follows:
With a detailed characterization of gradient descent training of the PAC-Bayes objective function, we derive that the final solution is kernel ridge regression with its kernel being the PNTK.
Based on the optimization solution, we derive an analytical and guaranteed PAC- AAC-Bayesian bound for deep networks for the first time. Moreover, bound differs from other PAC- Bayes bounds. Recent papers require distribution of posterior, while our bound is completely independent of computing the distribution of posterior.
The performance of PAC- Bayesian learning depends on the selection of a large number of hyperparameters. We design a training- free proxy based on our thcoretical bound and show it is effective and time-saving.
Our technique of analyzing optimization and generalization of probabilistic neural networks through wvef-parametcrization has wide range of applications such as the Variational Auto-encoder (Kingma & Welling 2013 Rezende et al,, 2014) and deep Bayesian net- works (MacKay| 1992] Neal 2012), we believe our technique can provide the basis for the analysis of wverpparameterized probabilistic neural networks.
2 RELATED WORK
PAC- MAC-Bayeiaan analysis. A Probably Approximately Correct (PAC) Bayes framework (McAllester 1999a b) can incorporate knowledge about the learning algorithm and probability distribution over set of hypotheses, thus providing a test performance (generalization) guarantee. Subsequently, the PAC-Bayesian method is adopted to analyze the generalization bound of the probabilistic neural networks (Langford & Caruanal (2002b). While the original PAC-Bayes theory only works with a bounded loss function. [Haddouche et al. (202T) expanded the PAC-Bayesian theory to learning problems with unbounded loss functions. Furthermore, several improved PAC-Bayesian bounds suitable for different scenarios are introduced by Bégin et al. [2014 2016). As result of the flexibility and generalization properties of PAC-Bayes, it is widely used to analyze complex, non-convex, and overparameterized optimization problems, especially over- wee-paaameterizee neural networks (Guedij 2019). Neyshabur et al. (2017b) presented generalization bound for feedforward neural networks with ReLÚ activations in terms of the product the spectral norm the layers and the Frobenius norm of the weights.
2

---

Under review as a conference paper at ICLR 2023
PAC-Bayesian learning. In addition obtaining the theoretical analysis for the generalization properties of deep learning, it is important to achieve a numerical bound on generalization for practical deep learning algorithms. [angford & Caruana (2002a) introduced method to train Bayesian neural network and used a refined PAC-Bayesian bound for computing the error upper bound. Later, Neyshabur et al. (2017a) extended Langford al. [Langford & Caruanal (2002a)'s 's work by developing a training objective function derived from relaxed PAC- BBy-ssaa bound. In the standard application of PAC-Bayes, the prior is typically chosen be a spherical Gaussian centered at the origin. However, without incorporating the information of data, the KL divergence might be unreasonably large, limiting the performance of the PAy-Bayes method. To address this gap, large volume of literature proposes to obtain localized CAC-Bayes bounds via distribution-ddeeeddee priors through data (Ambroladze et al. 2007; Negrea et al.] 2019 Dziugaite et al. 2020 Perez- et al.) 202T). Furthermore, Darrharmore,(Dxiuuaate & Roy 2018b);; Tinsi & Dalalyan (2022) showed how differentially private data-dependent prior yiclds a valid PAC-Bayes bound for a situation where the data distribution presumed to be unknown. More recently, research has focused on providing a
bound for more realistic architectures, such as convolutional neural network (Zhou et al.) 2018) binary activated multilayer networks (Letarte et al., 2019), partially aggregated neural networks (Biggs & Guedj] 2020), and graph neural networks (Liao et al. 2020). We denote the practical use of the PAC-Bayesian algorithm train weeetearameterized neural networks as PAC-Bayesian learning and the target of this work is to demystify the success behind deep learning trained via the PAC-Bayesian bound through PNTK.
3 PRELIMINARY
Notation. We use bold-faced letters for vectors and matrices and nan-bold-faced letters for scalars. We use || llz to denote the Euclidean norm of a vector or the spectral norm of a matrix, while denoting || || as the Frobenius norm of a matrix. For neural network, we denote (xx as the activation function. We denote [n] = (1.2...., The least eigenvalue of matrix A is denoted as AA Amin (A).
3.1 DEEP PROBABILISTIC NEURAL NETWORK
In PAC-Bayesian learning we use probabilistic neural networks (PNNs) instead of deterministic networks. where the weights always follow certain distribution. In this work, we adopt the Gaussian distribution for the weights, and define a L-layer probabilistic neural network governed by the following recursive expression,
x() 1 (W (-1), 11,,11, f=vix(1) Vm
(1)
where (0) = x E Rd is the input, w E gm is the weight matrix at the first layer, ( E Rm is the weight at the I-th layer for 2 < 1 L. and v E RT is the weight vector at the output layer. To keep weights follow Gaussian distribution during gradient descent training, we introduce the re-parameterization trick (Kingma & Welling 2013 Kingma et al. 2015):
( -
w(()
((l), ((I) I), <ISL: v V= + Vo g()), g(() NNO, I).
(2)
where denotes the dlement-wide product operation, thus gd)) for 1 S S L and g(9) share the same size as their corresponding weight matrix or vector. The key insight of re- re-parameterization is to sample& (1) for I L and g(,) from normal distribution A/ I). and leave w w!! VP.
and v. to be deterministic.
We adopt random initialization for mean weights, where ( V. ~ NO ct II for l E [1, L]. With an abuse of notation, we omit the size of mean 0 and variance I which in coordinate with their weight matrix or vector. On the other hand, we use an absolute constant to initialize variance weights, namely (!()) Va &g 1, where is matrix or vector with all elements to be 1.
3.2 PAC-BAYESIAN LEARNING
Suppose data S (xx;,)))); are i.i.d. samples from a non- distribution D. Define H to be the hypothesis space, h(x) to be the prediction of hypothesis h. H over for Let
3

---

Under review as a conference paper at ICLR 2023
(h) ((x,y))>D(((e {xx)]] represent the generalization enror of classifier h and Rs(h) Di=1 (y: (x:) represent the empirical error of classifier h, where E(-) the loss function.
In PAC-Bayes, the prior Q(0) E H is the prior distribution in H at initialization or before training, and the posterior E H is the distribution of parameters after training. To make the evaluation of prediction based on weight parameters W feasible for 1 [L]. we adopt the Gaussian distribution
for the parameters, and the expected value for population risk and empirical error are (Q =
((x.)=D.h-- [y, [x))] E (h)). (s(Q) = En-o Rs)h). The PAC-Bayes theory
(Langford Seeger 2001 Seeger 2002 Maurer 2004 gives the following theorem:
Theorem 3.1. Then for any å E (0. 11, the following inequality holds uniformly for all posteriors
distributions Q E H with probability of at least 6,
kl (((Rs())IBB((Q) < (I(QIIQ(O)) + log 2V
(3)
where (L(QIQQ()) = leQ In Q(0) is the Kullback-Leibler (KL) divergence and kl(gll4') = q glog(;) + (1-11000) 9) log( 1 is the binary KL divergence. Furthermore, combined with Pinsker's incquality for binary KL divergence, kl(pllp) z (p p)2 when ò < p. yields.
(L(Q)Q(0)) log 2 2
RD( R2) 1 (Q)(Q)
(4)
n
Equation is a classical result. This result can be further combined with the incquality vab s (Aa + {) for all À > O, which leads to PAC-Bayes-2 bound in Theorem 3.2, as proposed by
Thiemann et al. (2017):
Theorem 3.2. Let Qo E H be some prior distribution over H. Then for any d E (0, 11, the following inequality holds uniformly for all posteriors distributions E H with probability of least 1
is(Q) (2( < 1 1-2/2
KL(QIIQ(0) + log Zyn
+ nd(1 n121-\/2)
(5)
In this work, inspired by Catoni (2007 y: Rivasplata et al. (2019), we aim to promote the training objective as PAC-Bayes bound and choose Equation (5) as the training objective. We highlight that the original interest of Theorem 3.2 (Thiemann et al.) 2017) to allow the optimization of quasiconvex objective both À and Q However, since our main goal is to study the optimization and generalization properties of PNNs, we set directly = 1 and we omit the factor two and express the objective function follows:
L(Q) Rs(Q) + KL QIQ(())) En~Q
n Σ (yn (xx)) KL(QIQQ())
n
n
where À is a hyperparameter introduced in a heuristic manner to make the method more flexible. Since the term regarding is a constant, we omit in the objective function. We set to be the squared loss in the training objective function because it has a nice property such that the final solution of the output function explicit in the infinite-width limit. The global convergence can be extended to cross-entropy loss like existing works (ji & Telgarsky| 2019 Chen et al. 2019). Instead of optimizing w(!) and v directly, the gradient descent with repurameterization trick leads to
aL(Q) (t)
O (t)
(1
(t =)
(7)
where 17 is the learning rate. For simplicity, we omit the gradient descent expression for V and V.. and will omit the corresponding terms regarding V,. Vo in the following text unless otherwise specified. To simplify theoretical analysis, this work considers gradient flow instead, and the same results can be extended gradient descent case with careful analysis.
4

---

Under review as a conference paper at ICLR 2023
4 MAIN THEORETICAL RESULTS
In this section, Theorem 4.2 gives precise characterization of how the objective function without KL divergence decreases zero. We then extend the convergence characterization the full objective, and find the final solution is kernel ridge regression, as demonstrated by Theorem 4.3 As consequence, we are able to establish an analytic generalization bound through Theorem
4.1 OPTIMIZATION ANALYSIS
To simplify the analysis, first consider the optimization of probabiliistic neural networks of the form (1) with objective (Q). In other words, we neglect the KL divergence term at this stage and show that the corresponding results can be extended to the target function with KL divergence in the next section. Given this premise, we show that for L--ayer probabilistic neural network, the gradient flow of output function admits the following dynamics
df(::t) af(X:t) dt
00, Ōt
af(X:t) 00.
+
= (y -
(8)
002
d0. δt
U where 0₂ = ((W ME1 vn) and 0, ((W a M1=1 v,)) are collection of mean weights and variance weights. Besides, (X, lungent kernels (PNTKs) defined as follows,
(X,X;;)) X;1) E Raxn and O,
E Ruxa are probabilistic neural
Definition 4.1 (Probabilisticc Neural Tangent Kemel). The tangent kernels associated with the output function f(X; t) at parameters Bp and G, are defined as,
, (X, 0f(X:t)
(ff(x;t)
with ((::))V wey (X;1( Vx..((X
0
2. XX x; 0f(X:()
[0((:1))
wf V (X::)) T f(:
00.
80.
(9)
Different from standard (deterministic) neural networks, the probabilistic network consist of two sets of parameters Oμ and O,, thus the PNTK has two corresponding tangent kernels.
One of the key findings of this work that the PNTKs (X, X) and O, (X, X) will both converge (X,X)
to a limiting deterministic kernel denoted as (X, X) at initialization and during training if m is sufficiently large, namely limm-+oo θ (X, X) = (X, and limm- O, (X, @ (X.X).
As result, in the fiinite-width limit, dynamics of output function with gradient flow is linear:
df (f( t) - f(x) (X,))
(10)
dt
By leveraging this insight, we arrive at our main convergence theory for deep probabilistic neural networks, which is stated formally as follows
Theorem 4.2 (Convergence of probabilistic networks with large width). Suppose (-) is H -LLipschhizz n² (o n n² log(2/8))
An > O, and the network's 's width is ofm m = 2 20LL) max
x3(k())
8 xg(K(
with the initialization. Then, with a probability of at least 1 over the random initialization, we have
Rs(Q(t)) I exp (- (o) Rs(((()))
(11)
where we define K00 (x,,x,) (x x and K !!)) (x,, x,) = limm
(1)
Our theorem establishes that if m large enough, the expected training emror converges to zero at a linear rate. In particular, the least eigenvalue of PNTK governs the convergence rate. Besides, find the change of weight is bounded during training, which is consistent with the requirement of PAC-Bayes theory that the loss function is bounded.
5

---

Under review as a conference paper at ICLR 2023
4.2 TRAINING WITH KL DIVERGENCE
According to Equation (6). there is a KL divergence term in the objective function. We expand the KL-divergence for two Gaussian distributions, P(I) NN o²), and P(0) NNoo of),
(L(((t)(P(O)) =
1 ( log dg 2
(ee p
d0 -1)
0t 44-1
(12)
2
We compare the gradient of mean and variance weights w Wo With a direct calculation, we have 8f(x.) of
0f(x))
a'(w()) x(==1)xx(-1) and (f(x,) Bw vmm 8x(1)
af(x.)
@ g()), It
dw !h mm 0x
shown that there is one more random variable (()) associated with the gradient regarding variance weights, which results in the expected gradient norm being zero. Therefore, it is equivalent to fix w. during gradient descent training and we arrive at the conclusion that probabilistic neural network is performing kernel ridge regression in the infinite-width limit:
Theorem 4.3. Consider gradient descent
objective function 16).
Suppose 2
poly(n, 1/,,, 1/6, 1/E). Then, with a probability of at least 5 over the random initialization, we have
fxx. Q(t))((()))))))))))))) (t=x = 0 (X) +N
(13)
where Q()
aligns with the definition of the empirical loss function.
Theorem reveals the regularization effect of the KL term in PAC-Bayesian learning, and presents an explicit expression for the convergence result of the output function.
4.3 GENERALIZATION ANALYSIS
We use squared loss to train the probabilistic neural network but adopt a general and suitable loss é E [0, 1] to evaluate the PNN's generalization. Recall that in Theorem 3.2 the PAC-Bayesian bound concerning the distribution at initialization and after optimization is given. Therefore, combined with the results from Theorem 4.3 ultra-wide condition.
we provide generalization bound for PAC-Bayesian learning with
Theorem 4.4 (PAC-Bayesian bound with NTK). Suppose data S (xxiyy:)] are i.i.d. samples from non-degenerate distribution DD and m 2 poly(n, >o .8-1). Consider any loss function É IR x E + [0, 1] that is 1-Lipschitz in the first argument such that fly,y) 0. Then with probability at least 1 - 5 over the random initialization and the training samples, the probabilistic neural network (PNN) trained by gradient descent for T 2 22 npo risk R& (Q) that is bounded as follows:
log z) iterations has population
(Q) < y (x,
y O
X + n
log 2v4
+
+
8
no?
n
(14)
The proof can found in the Appendix this theorem, we establish reasonable generalization bound for the PAC-Bayesian learning framework, thus providing theoretical guarantee. Com- pared to the PAC-bayes bound (5), our bound is analytic and computable. We further demonstrate the advantage of PAC-Bayesian learning by comparing it with the Rademacher complexity-base generalization bound for deterministic neural networks with a kernel ridge solution.
Theorem 4.5 (Rademacher bound with NTK). Suppose data =x),, are i.i.d. samples from non-degenerate distribution D. and m 2 poly(n, > 1..).. Consider any loss function
IR x RR + [0, 1] that is 1-Lipschitz in the first argument such that E(y, = 0. Then with probability of at least 1 - 8 over the random initialization and truining sumples, the deterministic neural network trained by gradient descent for T 2 l noo
log z iterations has population risk Rop
that is bounded as follows:
T X. X.
A/1
X. X) +/ o
log " An5
Rp S
(15)

---

Under review as a conference paper at ICLR 2023
Theorem 4.5 is obtained by following Theorem 5.1 in Hu et al. (2019), which presents Rademacher complexity- based generalization bound for ultra-wide neural networks with a kernel ridge regression solution. Similar analysis for kernel regression without regularization based on NTK can be found in Arora et al. (019a)) Cao & Gu (019).
The main difference between two generalization bounds is y (X
versus (eg (X.))+//))) by which is due to the fact that the PAC-Bayesian bound count the KL divergence while Rademacher bound calculate the reproducing kernel Hilbert space (RKHS) norm. We find that the convergence rate of the focused term are different. One is O(1/n) and the other is (11/ Vn). Therefore, we conclude that the PAC-Bayesian bound has a numerical improvement over the Rademacher complexity-basee bound when is large.
PROOF SKETCH
To prove Theorem 4.2 first show that PNTKs at initialization are close to the limiting kernel given the width is large enough. Then we prove the distance between PNTKs and limiting kernel during training is also bounded, meanwhile loss has a linear convergence rate by induction. Our proof framework is similar to Du et al. (2019); Arora et al. (2019a) s. However, The main difference is that our network architecture is much more complex (e.g. probabilistic network contains two of parameters) and each set involves its own randomness which requires bounding many terms more elaborately. The detailed proof can be found Appendix
The proof of Theorem 4.3 utilizes an argument of linearization of the network model in the infinite- width limit. This allows us obtain an ordinary differential equation for output function with the solution of kernel ridge regression. The details are given in Appendix B
For generalization analysis, we defer the proofs of Theorem 4.4 Appendix C Our proof is based on a characterization of the empirical error and KI. divergence term via the explicit solution found in
Theorem
EXPERIMENTS
As an extension our finding of the PAC-Bayesian bound in Theorem 4.4 we provide a training-free free metric to approximate the PAC-Bayesian bound via PNTK, which can be used to select the best hyperparameters without involving any training and eliminate excessive computation time. Besides, provide empirical verification of theory in Appendix D.I and comparison of theoretical bounds with empirical bounds in Appendix D.2
6.1 EXPERIMENTAL. SETUP
In all experiments, the NTK parameterization is chosen to initialize the parameters, which follows Equation (1). Specifically, the initial mean weights O,., are sampled from truncated Gaussian distribution with a mean of zero and one standard variance of 1, truncating at two standard deviations. To ensure that the variance is positive, the initial variance for weight is transformed from the given value of Ca through the formula
log(1 + exp(po)).
describe the of both fully connected and convoluted neural network structures to perform experiments on MINIST and CIFAR 10 datasets to demonstrate the effectiveness of our training-fir free PAC-Bayesian network bound for searching hyperparameters under different datasets and network structures. In particular, we build a 3-layer fully-connected neural network with 600 neurons on each layer. On the other hand, the convolutional architecture equipped with a total of 13 layers with around 10 million learnable parameters. We adopt data-dependent prior since practical and popular method (Perez-Ortiz et al.) 2021 Fortuin 2022). Specifically, this data-dependent prior is pre-trained on subset of total training data with empirical risk minimization. The networks for posterior training are then initialized by the weights learned from the prior. Finally, the generalization bound is computed using Equation (5). The relevant settings are referred to in the work by Peree-errreeeeee Ortiz et al. (2020), such as confidence parameters for the risk certificate and Chernoff bound, and the 150,000 times of Monte Carlo samples estimate the risk certificate.
In section 6.2

---

Under review as a conference paper at ICLR 2023
Correlation between y. and bourd under FCN with different priors
Correlation between and bound under FCN with different A
Correlation between PA and bound
under PCN with different
Erron Bound
800
30%
10h
0175
Error Bound
ort Aex
012
B0t Ara
0100
10.6 11.3
10.0
1.00
Lh0
364
7..
Correlatton berween 154 and bound
Correlation between PA and bound under CNN with different A
and bound Correlation with different
under CNN with differrnn pro
20 Ara
Enror Bound
6th
Error Bound
praeeera 9ra
02
scvuror Ara
27.4
P.
Figure 1: The first shows correlation results of FCN structure the MNIST dataset. Kendall-tau correlations between generalization bound with respect to the proportion of prior data, coefficient of KL penalty, and /n 0.89, 0.89, and 0.93 1 % level of significance. Similar results found in the CNN structure with the CIFARIO dataset where Kendall-tar -tu correlations are 0.89, 0.83, and 0.57, as shown in the second row.
Setup
Risk cert.
Computation time (hours)
Data
Method
Network
ge
Single
Total
Exhaustive Scarch
324.00
CNN
.0059
16.92
10964.16
Rayesian Search
FCN
.0010
.0212
0.50
18.00
MNIST
CNN
.0059
.0110
16.92
609.12
FCN
.0010
.0264
0.03
19.44
PA
CNN
.0085
.0160
0.03
19.44
Exhaustive Search
706.32
CNN
.0142
.1969
45.00
29.160.00
Bayesian
FCN
.174
0.5377
1.09
39.24
CIFARIO
Search
CNN
.0142
.1969
45.00
1,620.00
FCN
.178
0.5490
0.03
19.44
PA
CNN
.0142
.1970
0.03
19.44
Table The performance i.e., risk certificates (cross-entropy (a-e and accuracy (01) and computation time against three hyperparameters searching methods (exhaustive search, Bayesian search, and P.A, the training free method). For the lowest risk certificate and computational time are highlighted in boldface, and second best are highlighted by underlining.
6.2 SELECTING HYPERPARAMETERS VIA TRAINING- FREE METRIC
The PAC-Bayesian learning framework provides competitive performance with non-vacuous general- ization bounds. However, the tightness of this generalization bound depends on the hyperparameters used, such as the proportionality of data used for the prior, the initialization of Pu. and the KL penalty weight (A). Since these three values do not change during the training, refer to them as hyperparameters. Choosing the right hyperparameters via a grid search obviously prohibitive, as each attempt compute the generalization bound can involve significant computational resources.
8

---

Under review as a conference paper at ICLR 2023
Another plausible approach is to design some kind of predictive, "training-free metric that can approximate the error bound without going through an expensive training process. In light of this goal, we have already developed generalization bound in theorem 4.4 via NTK. Since NTK changes are held constant during training, we can predict the generalization bound by this proxy metric, which can be formulated follows:
Pr
(//21-11 yy 2
+
θ + 1/c yy
(16)
where θ is an empirical NTK associated with mean weights, measured on a finite-width neural network at initialization. yy is n x n label similarity matrix (if two data have the same label, their joint entry in the matrix is one and zero otherwise), and n is the number of data used. Note that the proposed proxy metrie in Eq. (16) share the same spirit of kernel alignment, label similarity metric, which has been widely used in the application of deep active learning (Wang et al. (202I), model selection for fine-tuning (Deshpande al. 202T), and neural architecture search (NAS) (Mok et al. 2022). To demonstrate the computational practicality of this training-free free metric, we compute PAA using only subset of the data for each class (325 per class for FCN and 75 per class for CNN). We should also mention that training- free methods for searching neural architectures are not new, and can be found in NAS (Chen et al. 2021 Deshpande et al. [202T) MAE Random Sampling (Camero) al., [202I), pruning at initialization (Abdelfattah et al. (021). To the best of our knowledge, there is currently no training-free fareing-free method for selecting hyperparameters in the PAC-Bayesian framework, which we consider be one of the novelties of this paper.
Figure D demonstrates a strong correlation between P.A and the actual generalization bound. Finally, demonstrate that by searching through all possible combinations of hyperparameters using P.A. is possible to select a hyperparameter leading towards a result that is comparable to the best generalization bound, but without excessive computation. To put things in perspective, in Table [ we compare the risk certificates and computation time for three hyperparameters finding methods (exhaustive search, Bayesian search and PAA) on the two architectures (FCN and CNN) and two datasets (MNIST and CIFARIO). Unlike exhaustively searching where the best set of hyperparameters are selected from 648 different hyperparameter combinations (9 data-dependent prior with different subsets data for prior training, 9 different values of KL penalty, and 8 different values of Pb). Bayesian search takes only 36 iterations to find the lowest bound since evaluates the information in past iterations of searching and efficiently selecting the next set of hyperparameters based on the prior knowledge, Yet, reducing the number of search iterations cannot sufficiently reduce the overall computation time when training a large and complex model. For instance, under the CIFARIO 10 dataset, it takes 45 hours to train CNN with the bound. In contrast, using the training- -fre method of P.A save 83.33 times the computational time find the bound that is close the lowest risk certificate in accuracy.
DISCUSSION
In this work, we theoretically prove that the learning dynamics of deep probabilistic neural networks using training objectives derived from PAC-Bayes bounds are exactly described by the NTK in an wwer-parameterized setting. Empirical investigation reveals that this agrees well with the actual training process. Furthermore, the expected output function trained with a PAC-Bayesian bound converges the kernel ridge regression under mild assumption. Based on this finding, we obtain an explicit generalization bound with respect to NTK for PAC-Bayesian learning, which improves over the generalization bound obtained through NTK on a probabibaii neural network. Finally, we show that the PAC-Bayesian bound score, the training-free method, can effectively select the hyperparameters which leads to lower generalization bound without cost excessive computation time cost which the brute-force grid search would incur. In summary, we establish our theoretical analysis on PAC-Bayes with a random initialized prior. Notice that neural tangent kernel cannot characterize the feature learning process in deep learning (Damian et al. 2022 Ba et al. 2022). This paper does not try to capture the feature learning for probabilistic neural networks given the NTK techniques used, but does provide sufficient new important convergence and generalization analysis for PAC- -ayyesaaa learning. One promising direction would be to study PAC-Bayesie learning with data-dependent priors by NTK.

---

Under review as a conference paper at ICLR 2023
8 REPRODUCIBILITY STATEMENT
To ensure the results and conclusions of our paper are reproducible, we make the following efforts:
Theoretically, state the full of assumptions and include complete proofs of our theoretical
results in Section 4 Appendix A and C
Experimentally, we provide our code, and instructions needed to reproduce the main experimental results, And we specify all the training and implementation details in Section 6 and Appendix D
REFERENCES
Mohamed S Abdelfattah, Abhinav Mehrotra, Lukasz Dudziak, and Nicholas Lane. Zero-cost proxies for lightweight nas. arXiv preprint arXiv:2101.08/34 2021.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252. PMLR, 2019.
Amiran Ambroladze, Emilio Herrado-Hemánder. and John Shawe-Taylor Tighter pac-bayes bounds. Advances in neural information processing systems, 19:9, 2007.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grainee analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332. PMLR, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv: 17955, 2019b.
Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High- dimensional asymptotics of feature learning: How one gradient step improves the representation. arXiv preprint ar arXiv:2205.014: 2022.
Luc Bégin, Pascal Germain, François Laviolette, and Jean-Francis Roy. Pac-bayesian theory for transductive learning. In Artificial Intelligence and Statistics, pp. 105-113, PMLR, 2014.
Lue Bégin, Pascal Germain, François Laviolette, and Jean-Francis Roy. Pac-bayesian bounds based the rényi divergence. In Artificial Intelligence and Statistics, pp. 435-444, PMLR, 2016.
Felix Biggs and Benjamin Guedj. Differentiable pac-bayes objectives with partially aggregated neural networks. arXiv preprint Stti::2006.122228 2020.
Andrés Camero, Hao Wang, Enrique Alba, and Thomas Bäck. Bayesian neural architecture search using a training-free performance metric. Applied Soft Computing. 106:107356, 2021.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in Neural Information Processing Systems, 25:1836-10846 2019.
Olivier Catoni. Pac-bayesian supervised classification: the thermodynamics of statistical leamning arXiv preprint urXiv:07/2.02245 2007.
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in four gpu hours: theoretically inspired perspective. arXiv preprint arXiv:2102.1/535 2021.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much over-parameterizztion is sufficient to learn deep relu networks? arXiv preprint arXiv:1911. 2019.
Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can leam representations with gradient descent. In Conference on Learning Theory, pp. 5413-5452, PMLR, 2022.
Aditya Deshpande, Alessandro Achille, Avinash Ravichandran. Hao Li, Luca Zancato, Charless Fowlkes, Rahul Bhotika, Stefamo Soatto, and Pietro Perona. A linearized framework and a new benchmark for model selection for fine-tuning. arXiv preprint arXiv:2102. 00084, 2021.
10

---

Under review as a conference paper at ICLR 2023
Nan Ding, Xi Chen, Tomer Levinboim, Beer Changpinyo, and Radu Soricut. Pactran: Pac-bayesian metrics for estimating the transferability pretrained models to classification tasks. arXiv preprint arXiv:2203.05112, 2022.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference Machine Learning, pp. 1675-1685. PMLR, 2019.
Gintare Karolina Dziugaite and Daniel Roy. Entropy-sgd optimizes the prior of a pac-bayes bound: Generalization properties of entropy-sgd and data-dependent priors. In International Conference on Machine Learning, pp. 1377-1386. PMLR, 201 8a.
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. CoRR, Abs/1703..1008, 2017. URL https:
Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent pac-bayes priors via differential privacy. arXiv preprint arXiv:1802.09555 2018b.
Gintare Karolina Dziugaite, Kyle Hsu, Waseem Gharbieh, and Daniel M Roy. On the role of data in pac-bayes bounds. arXiv preprint (rXXi:2006.1092) 2020.
Vincent Fortuin. Priors in bayesian deep learning: A review. International Statistical Review, 2022.
Peter D Grünwald and Nishant A Mehta. Fast rates for general unbounded loss functions: From emm to generalized bayes. J. Mach. Learn. Res., 21:56-1, 2020.
Jiechao Guan and Zhiwu Lu. Fast-rate pac-bayesian generalization bounds for meta-learning. In International Conference on Machine Learning, pp. 7930-7948. PMIR, 2022.
Benjamin Guedj. A primer on pac- mec-bbaesian learning. arXiv preprint arXiv: 2019.
Maxime Haddouche, Benjamin Guedj, Omar Rivasplata, and John Shawe- Taylor. Pac-bayes un- leashed: generalisation bounds with unbounded losses. Entropy, 23(10):1330, 2021.
Wei Hu, Zhiyuan Li, and Dingli Yu. Simple and effective regularization methods for training on noisily labeled data with generalization guarantee. arXiv preprint /xX5v:/905//1365 2019.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint 07X752180 2018.
Prateek Jain and Purushottam kar. Non-convex optimization for machine learning. arXiv preprint arXiv:1712.07897. 2017.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint DXXX::1909.12292 2019.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.61 14, 2013.
Durk Kingma, Tim Salimans, and Max Welling. Variational dropout and the local reparameterization trick. Advances in neural information processing systems,
2015.
John Langford and Rich Caruana. (not) bounding the true error. In T. Dietterich, S. Becker, and Z. Ghahramani (eds.), Advances in Neural Information Processing Systems, volume 14. MIT Press, 2002a.
John Langford and Rich Caruana. (not) bounding the true error. Advances in Neural Information Processing Systems, 2:809-818 816, 2002b.
John Langford and Matthias Seeger. Bounds for averaging classifiers. 2001.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521((553):436 444, 2015.
11

---

Under review as a conference paper at ICLR 2023
Jachoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks as gaussian processes. arXiv preprint 2017.
Jaehoon Lee, Lechao Xiao, Samuel S Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. arXiv preprint 1etp:::92200000000 2019.
Gaël Letarte, Pascal Germain, Benjamin Guedj, and François Laviolette. Dichotomize and generalize: Pac-bayesian binary activated deep neural networks. arXiv preprint arXiv: /905. 10259, 2019.
Renjie Liao, Raquel Urtasun, and Richard Zemel. A pac-bayesian approach to generalization bounds for graph neural networks. arXiv preprint arXiv:2012.07669, 2020.
Ben London. A pac- pac-bayesian analysis of randomized learning with application to stochastic gradient descent. arXiv preprint (xXiv:170.006/1 2017.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computa- fion, 4(3):488 472, 1992.
Andreas Maurer. note on the pac bayesian theorem. arXiv preprint cs/041 1099, 2004.
David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learning theory, pp. 164-170, 1999a.
David A McAllester. Some -ay-saaei theorems. Machine Learning, 37(3):355-366 1999b.
Jisoo Mok, Byunggook Na, Ji-Hoon Kim, Dongyoon Han, and Sungroh Yoon. Demystifying the neural tangent kernel from a practical perspective: Can it be trusted for neural architecture search without training? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11861-11870, 2022.
Radford M Neal. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
Jeffrey Negrea, Mahdi Haghifam, Gintare Karolina Dziugaite, Ashish Khisti, and Daniel M Roy. Information-theond generalization bounds for sgld via data-dependent estimates. arXiv preprint arXiv: 2019.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring general- ization in deep learning. arXiv preprint arXiv:1706 170.:007678477 2017a.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally- normalized margin bounds for neural networks. arXiv preprint arXiv:1707.09560 2017b.
María Pérez-Ortiz, Omar Rivasplata, John Taylor,aylor, and Csaba Szepesvári. Tighter risk certifi- cates for neural networks. CoRR, abs/2007. 12911, 2020. URL http:ttts///aii..rrrrt wwg./bbsorg//bb.....
2007 12911
Maria Perez-Ortiz, Omar Rivasplata, Benjamin Guedj, Matthew Gleeson, Jingyu Zhang, John Shawe- Taylor, Miroslaw Bober, and Josef Kittler. Learning pac-bayes priors for probabilistic neural networks. arXiv preprint arXiv:2 109. 10304, 2021.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: a nonasymptotic analysis. Conference on Learning Theory, pp. 1674-1703. PMLR, 2017.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning. pp. 1278-1286. PMLR, 2014.
Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvari. Pac- bac-bbyes with backprop. arXiv preprint 178.07090..0. 2019.
12

---

Under review as a conference paper at ICLR 2023
Matthias Seeger. Pac-bayesian generalisation error bounds for gaussian process classification. Journal of machine learning research, 3(Oct):233-25) 269, 2002.
Samuel L Smith and Quoc V Le. A bayesian perspective on generalization and stochastic gradient descent. arXiv preprint 1rXiv:1710.06451, 2017..
Donald Specht. Probabilistic neural networks. Neural networks, 3(1):109-118, 1990.
Niklas Thiemann, Christian Igel, Olivier Wintenberger, and Yevgeny Seldin. strongly quasiconvex pac-bayesian bound. In International Conference on Algorithmic Learning Theory, pp. 466-492. PMLR, 2017.
Laura Tinsi and Arnak Dalalyan. Risk bounds for aggregated shallow neural networks using gaussian priors. In Conference Learning Theory, pp. 227-253. PMLR, 2022.
Haonan Wang, Wei Huang, Andrew Margenot, Hanghang Tong, and Jingrui He. Deep active learning by leveraging training dynamics. arXiv preprint arXiv:21 2021.
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.
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. arXiv preprint arXiv:1804.05880 05862, 2018.
13

---

Under review as a conference paper at ICLR 2023
A
PROOF OF THEOREM 4.2
Theorem A.1 (Restatement Theorem
Suppose al.) is H -ipischittz and the network's width
with the initialization. Then, with probability of at least 1 5 over the random initialization, we have,
of m Q 20(L) max
log(Ln// LL//8) ))(
n² log(2/5)
A((K!5))
Rs(Q - exp (o Rs(((())
where we define K (x,, = (x: TTxx and K (x,,x, limm (x. (x (1)
x3
Proof Sketch of Theorem
A.7
study the behavior of output function under gradient flow, we first
write down its dynamics df(X;1) 1](X;1)
,0, Os(X;1)00 00,
00 ōt 00, Ōt = (X (xx
dt
where @" and O the PNTKs of the whole network, composed by the NTK of each layer. We observe that if @* and @" converge to deterministic kernel, then the dynamics of output function admit linear system, which is tractable during evolution.
Before demonstrating the main steps, we introduce Neural Network Gaussian Process (NNGP) of our studied neural network in the infinite- width limit (Lee et al. 2017), which defined as follows:
K (x,, xj) = )
K (x,,,)) = lim (x (1)
m--00
where subscript i.j denote the index input samples. Instead showing the O" and @ are close @ in infinite-width limit, we use (L) as an anchor kernel. With the relation of NKT and NNGP, we can simplify our proof. Therefore, to prove Theorem A. three core steps are:
Step Show at initialization Amin (0), Amin (O. (0)) 2 Amin (K(o)) and the required condition
on m.
Step 2 Show during training Amin ( (0), Amiu (+)(0) 2 Amin (K (k()) and the required condition on m.
Step Show during training the empirical loss has linear convergence rate.
In proof, we mainly focus deriving the condition
by analyzing Amin (0) and
Amin (0, (0)) at initialization through Lemma A.2 and Lemma A.3 For step 2, we construct Lemma
A.5 and Lemma A.6 demonstrate that Amin O (0)), Amin ( Z Amin((((") This This leads to the conclusion that the required condition on during train. Finally, we summarize all the previous lemmas and conclude that the training error converges at a linear rate through Lemma A.7
A.1 STEP 1. BOUNDING LEAST EIGENVALUE OF PNTK AT INITIALIZATION
We first study the behavior of tangent kernels with an ultra-wide condition, namely m poly(n, 1/Ao, 1/8) at initialization. Lemmas A.2 and A.3 demonstrate that if m is large, then the feature of each layer is upproximately normalized, @*(0) and (0) have a lower bound on the smallest eigenvalue with high probability.
Lemma A.2 (Initial norm at initialization). Suppose a()) is H-Lipschitz. If m 22 (Tgge()) (
where == (cr + 2 (22)))))))) +H), while Wp and W. are initialized by the form described
in Section 3.7 then with probability at least - over random initialization, for each E [L] and
i [n), have
< |x (1 (0) 2
where the geometric series function (c( Σーー C².
14

---

Under review as a conference paper at ICLR 2023
( L0 Rinin (k(t))
n²
Lemma A.3 (PNTK at initialization). Suppose o(.) is H H-Lisschitz. If S2
while Wn y and W initialized by the form described in Section 3.1
then with probability least
1 - 6, we have
Amin (((((()) 2 min (K(1.) 14
Amin (0)) 2 Amin (K(4)
Proof of Lemma A.2 The proof is by induction method. The induction hypothesis is that with
probability at least (I - 1) nz over (O)..)...... (O) for every 1 S s l 1, we have
gc(1")
(1) (0)ll2 S 1+ (c
S
2gc(L)
2ge (L)
where the geometric series function gc(1) -ーー C².
Note that there are two randomness each ( for E [L], which can be seen from the expression:
W w w (())
The first randomness comes from the initialization of ( and the second randomness is from random variable &. We then unify the two randomness into one, namely W() ~N (ct, + cg) -)) through the following argument:
(w())
P(W.))-
P(W
203
V2TcC
Plugging the density function P(W pnj)) of variable μr into the above expression, we can obtain,
(we))
w 205
(wit)
o)²
P(W!))
mill
kit
26
dw())
e
V2Tca
C2
(1 .) ato
cg)
2m (c,, + cg)
With the result of (P(W vny) at hand, we continue to bound ||x] the randomness from (0)) According to the feed-forward expression, we know that
( I|3. We calculate the expectation over
X ! (0)
Σ (w xx(-((0)))
Then we have
x (0)1 = E [o (o) Tx (1 xx(- (0))²
E
(c
2 (/
15

---

Under review as a conference paper at ICLR 2023
Because a()) is H H-Lisschitz, for 1 1 2, we have
E 1 (o (az)²
E -Ex
< H|a - 11 B2- / (((0((22)0((21)
(a $ H|a - 1| Bz.,woo.) { 1)2)]
$ H|a
(2l000)) Z + H (122N(0.1)
2
H|a - 11 11- (2)0(0))
H| + 1))
C a - 11
$ 24 c
where we define C = (c + cg) (20 +HH).
For the variance we have
(c) + e2)
| ( (0)1
Var ow (ww (0)
Var
(o) + 22
(0) xx)=
m
(gg + 22
E
!
m
C2
where C2 a00) + 80) 3 / 2A + (4) ( 512H and the last inequality we used the formula for the first four absolute moments of Gaussian.
Applying Chebyshev's inequality and plugging in our assumption on m, we have with probability 1 nl over (
8
||x 9 (0)122
El|x 2 (0))
S
2ge (L)
Thus with probability 1 - d, over W(1) ..... w
||x (0)l - ||x ( (0) (0)11
Cgoll 11)
+ LL(
2g0(L) (L)
(I(l)
(go(L)
Using union bounds over [n], we prove the lemma.
Proof of Lemma A.3
For a weight matrix, we decompose it into m weight vectors, namely W())
[w/). w/ .... expressed as
Then the derivative of output the parameters (t) and way can be
(1)
0f(x))
0f(x;) (1) Vm Ox(1)
o' (1-1)x
(w(
(17)
0f(x))
0x(x;)
d' d'(w!! (1-1)) (1 o 5!!)
(1)
mm Ox(1)
16

---

Under review as a conference paper at ICLR 2023
According to the definition of PNTK for each layer:
(1) = V w/o t) t)V t(
V
W!O W (xx:f( t(
Through a standard calculation, we show that PNTKs can be expressed as |1-
n.i (x)
(E)
x
11-
1
u (ww
x (1- 1 (o ((w!l)) x -))
n.ij (x(4-11) 11-1)
(L)
v ((w! (1)) (1-1) )o' ((w 11) E
m
Note that the difference between O (L)
! and (L)) can be decomposed as follows:
(i)
9 K(H) (0 (0) k() + (K!) - -K(4))
(L)) K(n) ( (0) k() + (K(L) K(U))
We split the proof process into two phases:
First we use concentration inequality to show that if m S
n² login? (0) X (k:z))
we have
Amin (K (K(1))
Amin (K(L))
o -
K!6)
+ ! - (((L)
<
4
Second, we show that if m = (0(1)))
(01ooottn//00200 Xmin (Kili))
then we have,
Amin (K(2)
K k(l)
Phase 1, bounding θ a Plugging the derivative result regarding mean weights in Equation 17 into the definition of PNTK (Eqution) 9) yields:
(1) B.g (0) (x} (L- 1 (L-1) 1) Σ (w x((--)())))) 1/ ((w
((w x) x)1-1)
m
By an analysis, we find that for all pairs of i, (.)))) (0) is the average of m i..d. random variables, with the expectation
K kC b0.1 = (ch + (2) -Ew-N(o.)) (x T 11-1) o d'(w" x1 (1 o''w'' [x]
x (-11))
Then by Hoeffding's inequality, know that the following inequality holds with probability at least 1 8'
log(2/8")
n.d (0) ((.)
(1) (1) 20.15
S
2m
Because NTK matrix is of size n X n, we then apply a union bound over all i, E [r]] (by setting 8' = 5/nn). and obtain that
log (og(2n2/0)
C (1)
(0)
((() $
2m
17

---

Under review as a conference paper at ICLR 2023
Thus have,
+ (P) K(L)
+ (()) K(L)
F
+ (L)
(0)
K (LL)
M.ij
o0.13
w
n log (2ng 15)
m
Finally, if n22 log(2n²/8))
Amin ( (K(4)
which implies
12
n² log n 18) A (K12))
then with probability
at least 1 Ó, 6,
Amiu (K(1))
0 (0) -00-- K(()
4
Phase 1, bounding O ! (
Plugging the derivative result regarding mean weights in Equation 17 into the definition of PNTK (Eqution 9 yields:
(0).
θ O (0) (x: (1-1))
1
((w o ((w!") (1-1)) (o'(w( xj (1-)))))+ .
Note that the tangent kernel θ (L) (0) differs from O (()) (0) with an additional term <²- It is known 2.ij
o.ij X1 independently with o' ((w T (1-1) 2 0) T o Tx( 2 20).
the c?
( Because [[x]] -1, the expectation of θ (0) equals the expectation of θ 2.ú (0). Thus for all 0.1y
(0).
pairs of i, (0) the average m i.i.d. random variables with the expectation
E o D (0)
= K
Now calculate the concentration bound. It is known that ş is independent and sub-exponential. Then, by sub-exponential tail bound, we know that the following holds with probability at least -
1og(8/8))
θ (L) (0)
(L) oc.ij
2m
This bound of the same order to concentration bound for θ (1)) 14,3 (l,1(()) (0).
Thus we can take all the
arguments for @ (L) (0) above to finalize the proof.
log(8n2 Amnn (L)
which implies m ? log(nn then with probability at least
X (K (K(b)
1 6,
Amin (K(L)
e (L() (0) ((L))
4
Phase 2, bounding
k (1)
K(L)
We show with probability 1 over the W(), for any
1 - <.,,, n,
log(Ln/0)
11
x
(1)
K
00.1
(
< <
m
The error constant & depends on the choice of activation function, and satisfies
E S C 20())
18

---

Under review as a conference paper at ICLR 2023
with C being positive constant. The 20(1) term comes form perturbation propagation through the neural network. The proof is by induction, and detailed proof can be found in the proof of Theorem E.1 in Du et al. (2019). Applying the union bound to the number of paths concludes the theorem, and the condition of m follows:
("("
n² log (og
m S
X (K(2)
Remark A.1. The concentration bound is over two randomness, one is initialization of W μ. and the other is Gaussian variable E.
A.2 STEP 2. BOUNDING LEAST BIGENVALUE OF PNTK DURING TRAINING.
The next problem is that PNTKs are time-dependent matrices, thus varying during training. To account for this problem, establish following lemmas stating that if the weight W (t) close to
(0) during gradient descent training, then the corresponding PNTKs O (T)) (()) and @ (1) are
close to their initialization O (L) (t). O (L)(t) respectively.
Importantly, we introduce an auxiliary weight matrix ( = W auxiliary weight vector v(t) vt + (t) O E" (0), where {(O) and (0( are the exact value of random variables at initialization. Then we demonstrate lemmas in step 2 as follows:
P (t) + W (t) O E(O), and an
Lemma A.4. W (0) and (0) are initialized by the form described in Section 3.1 and suppose
for every I E [4]. ( I Cw,ov [x(]()]]l,
| I C2,0 and (/ w w S /mR
for some constant Cw,0s Ca.0 > 0 and R S Cu Ifa(.) is H-Lipschitz, then with probability at least 1 - 8, we have
x (1) (0)12 S Ro,, (1) + log(2/5))
where C. = =2, C-
Lemma A.5. If (,) and W. (0) are initialized by the form described in Section 3.1 and uppose (.) is H- Lipschitz and smoooh.h Suppose for I E [L]. (/) I Cw.oym,
[v(0)]22 S m. |<<0)] S a4.0m² 1 S $||x(000)1 S Cr,0- If |
c2.0
(V)(0)) vitt) v(0) I /mR where R S cgez (L) (L)-1 ( (K(4)) + 1og(2/6))
R < cge, 11-inn (K(4)m + log(2/6)) and R S cge, (L) for some small constant c and c2 2/ Howo then with probability at least we have
Amin (K()))
θ (1)) 2
Amin (K(L)
( (0)
Lemma A.6. If (s((,t)) s exp(-Amin (K()))(') Rs 0) holds for S 1' S t, we have for any 0 058st
(s) - w0 (5)())))(())
{(s) ((s))v(0)
< R' mm
1 44 - for some small constant with
where R'
Apym
max (2 T.m,o:
Proof of Lemma A.4 The proof sketch is by induction method.
For = 0, where the target is input which is fixed, thus satisfying the hypothesis. Now suppose the consider 1' l.
induction hypothesis holds for I' 0,...,
19

---

Under review as a conference paper at ICLR 2023
(t) - (0)
o ( (
m
OW (t)
- (1x (0))
(0)) -000 2000
(1100///1 +
( - + (0000--0000)
2
x(- (1)) - x" - 1100 2
H 1 m
(w)(0)- - w0 2
+ 1000---0000
x² xa=1(0
H (cw,oVm+ + RVm(1 + log(2/5))) HR (1 1) m
HVMR(1 + 1og(2/5)cz. m
<
H H (ca9e, (1 - 1) + 11 + log(2/6))
<e (1) (1 (0/2))))))
For simplicity we define (t) = w1 x; (t)
Proof of Lemma A.5.
Now we bound the distance between e ( (t) and @ (l.)) (0) through the following inequality:
θ (1)) (1) (1)
e (z) n.ig (0)
(t) T (6- 1) (t)
Σ
(t)() ot) (z)) of (t)
m
m
(0) (1-11(0) x5
(0)
(0) (2.,(()) o' (0)
m
x3 (t))(()) (t) x
(0) 1) (0) - u (0) 0' (s)r(()) (E)((())
+
(0)
Σ 2, (o' (z,, (t)) o' (t), o' (e,,
+
(t)
Σ (c-(t) - (0)3) o (z), of (z))(()))
< <H² 12 x
(1-1)
(1) x (0) T (0)
x
+ cg m
2
1
er (0) (o' (t), of (z),,(t) - o' (0)) of (z))
+ +4H22
1
(t)² - (0)
2.0
=T/ 1
1
20

---

Under review as a conference paper at ICLR 2023
For 1 by Lemma A.4 we have
F- -H= 20
(L-1)
11 (t) x) (t) - x (0)
x/ (0)
x
<H2 2
(x:
(1) (1))) x (1-1) (0)) x
(0)) (1) (1) + HH -22 x; (x) - xf-" -0)
S 2 11 09c, (L) + log(2/6)) (ca.0 + Hc,.0g.. (L) + log (2))))
+ +2 3 0e, (L) RC + log(2/6))
<303 2 H gc (L) (ogg/))))
For 12 we have
13 =2,00
1
Σ e+ (0)2 (t) -116)))-r007))) (o00 (0)) o') (0),
<2
1
m
er (0) (z)) (t),
000 m
+ (0)²
(o' (1)) - o' (0)) (z..- (0))
H m.0
m
(0) (1) - (0)
+ (0)² (1, - 200
m
3 3 2
m
Σ
(1) - (0)
+
Σ
()r(1) - (0)
vm
Using the same proof for Lemma A.4 it is easy to see that
m
2
Σ
(1) - =0)
S 2 200 (1 1og(2/)))
+
Thus
1 < 20,0, Lg. (L) (L) 12g(2/0))
For 123.
13.3 = 2 200
Σ
(1)2 v. (0)2
(t(() 0, (0) (+( + (+) e,c00/
< 441226 22,0 m
< 12H 2 2 R(1 + R(+lo0(21)))))) (2)())
Therefore we can bound the perturbation
2
Σ
00 (t) - (L)
n.yy (0)
F
(2 + 3H22 H 22 (og(/2))()/)))
2 2200 + log((//))]nnR
21

---

Under review as a conference paper at ICLR 2023
Recall the bound on R, which is R I cge, (L) *minn (k(1))) (11 + log(2/6)) -2, we have the desired result for (
Amin (K(4))
4
Then we bound the distance between O 002 (1) and 62) (0) through the following inequality:
θ (1) (0)
1) (t) T (1- (1- 1) (t)
1
(t) (2), o' (z)r(t)) (t))
xj
m
x (1-1)
(0) Txx xj
(6-1)
(0)
Ur (0) (z, o' (0) c())
(1-1)())
(t) x
x((-1) (0) ト m
v (0) o' (z), of (e))(()) .2
+
v0) (o' (t), of (z), d' (z), 0' (0))
+
(1) x3 (1) (t)
(v.(1)) - (c(0))-.(00))))))))) o' (z)r (z)) (z))))) -
+
x x (0) x (0)
Σ
(0) (0' (0)) (3 (ss, (1))) c()))
< SH 12,00 + 1og(2/6))
x
11-1) (1) x!" 41- -0) x 11 (0) (1 log(2/6))
m Σ
+ 1, + 12g()/
(0)² (o' (z (t)) 0' (t) o' (2 (z,, of (2) (0))
+ +4H222 (20 + 1og(2//))
1
Σ
(t) - e2)
+ 4H22222 24122220 220000000 m
Σ
(t) (0))
-(1 + log(()))))
r + ( + p.J.
For 1235, by the tail bound fora chi-square variable, we have
log(2 (2/))
1 s 2 (1
m
Therefore we can bound the perturbation
2 ( (t) - (1)) (0)
Σ
1,j=1
(2 + +3H? H 10 "0 (D)(1 + log(2/5))²
+ 2 2 21 + l2g(2/0))2 + 41122 2000000 + log(2/8)
nnn
22

---

Under review as a conference paper at ICLR 2023
Recall the bound on R, which is R I c9c, (L) *miin (k()))n ((11))nn1111 + log(2/6)) -3, desired result for :
have the
Amin (K(4))
o (1) (0)
(0)
2
We first consider the derivative of w( and have:
Proof of Lemma A.6
W
(1)
(s)
ds
L-1+
II , ()) (,0(3) |1
f.
Σ (ui sx (s) ((8)T
m
k=1+1
2-141
Σ |u f(x:: 8) +>> |||| 11 (s)1l2
II ww (5) II | (3))l2
S1
1=1
k=1+1
d
Σ (y f(xxis))x (5)
V (s(s)
(s)
ds
where
J(n) diag (o' ((w(())xx((
o' ((w!))) [x(())((((-1))
R
are the derivative matrices induced by the activation function.
To bound
(>))ll we can just apply Lemma A.4 and get
1) (s)llz 09c, (h) log log(2/0)) + Ca,0 2c, + log(2/8)).
To bound
[w())s) (s)
2- we use our assumption
II
w (1 (s)l2
II (/ + || ( k=/+1
k=1+1
II (Cw,oV'm+ R' m1) + log(2/8)) kmi+1
(Cw,0 + R²) l--t 1-1 x + log(2/0))
< (2cw,,) 1-1 m L (1 + (2))).))
Note that | (s)]l2 < H. Plugging in these two bounds back, obtain
d ds
(1) (6)
<4n1c,,002,00 log (2))) bi [(x;; 8)
S S4nps (og log(2/0))) v /uly +|;
< + log((//))) Aoo no R'vm. Vm.
Similarly, we have
d vu(s)
n 22nc.0 Σ
1y f(x;is)
dt
Se Aos 1
7 Tndoo'' vm.
23

---

Under review as a conference paper at ICLR 2023
Next, we consider the derivative of w and have:
1-1+1
d
(s)
=1
(
(v= [(x; Q())))x a ">( o (() BQ(a)
ds
m
v(()T
II k=141
1 (o)
where have used we the definition of loss Rs(Q) = @ Rf( and interchanged integration and differentiation. Similarly, for Vo we have,
d ds
a (s)
n
F Baco) -wwe-fxxa [(x;)))
{(4) =0
Integrating the derivative of weights, we obtain
(s) w/
+ -w))())))!()))) {0(0
- W (s)
d
(5))
S Rt' v
dp K
F
((5)- v(0
d
<
v((())
$ SR'
ds
6'=0
A.3 SETP 3. TOWARDS LINEAR CONVERGENCE RATE OF EMPIRICAL LOSS
Now we process to analyze the convergence rate of empirical error. Combined with fact that least eigenvalue of PNTKs and change of weights are bounded during training, the behavior of the loss is traceable. To finalize the proof for Theorem 4.2 we show:
Lemma A.7. If R' < R, we have (Q()())) I Rsp(-AAt)Rs((((1)))
Proof of Lemma
A.7
According to the gradient flow of output function, we have
tx
Σ
Of(xxtt (w! (t) ow!! dt (0((x)t) dv
) dt ) (0f(x)) dvo(!))
0f(x:t) dw (t)
)
dt
aw!! dt
+
OV
Ov. dt
!w!! ow!!! ) (or(x)) f(x,)
0f(x)) 0f(x,)
(y: - f(x)))
Σ
?w!! ow!!
3=1
1m1
0f(x)) f(x))
+ Ox(x)) as(x;)
} Ov.
+
Ovp ov,
dv.
2 Σ (y, f(x::t)) (L)
Then the dynamics of loss can be caleulated,
24

---

Under review as a conference paper at ICLR 2023
d
1
E ((xQ -x /
Rs(((()) = dd
di
+ (0) (x- (yx - y - f(x, ((1)
S 5-20 [
where we have used the condition Rs(Q, t) S exp(-Aot)R=sQ..))
R. Therefore, we have the desired result:
Finally, we provide a bound for Rs(Q())):
| y S(X 2(0))
Σ
y + yo yi Q(0)) + /(x))))((x)0))))
=
2
Σ (1+ (1)))1)
(())
Recall that Lemma A.5 and Lemma A.6
R < c9c, (L) Amin (k()) (og(2((/)))
Rt' 16(1 + (2/ cn 0220 (cx)² Vnily f(x, Q(0))l2
dovm
Thus R' R yields
Ω nd log(2/22/)) (2/8)10
2
B PROOF OF THEOREM 4.3
Theorem B.1 (Restatement of Theorem 4.3 Consider gradient descent on objective function (o)
Suppose 2 poly(n, 1/Ao, 1/8, 1/E). Then, with a probability of least 8 over the random initialization, we have
(() le=) 1= = (x +X/C) '
where fxx Q())) E aligns with the definition of the empirical loss function.
To proceed the proof, we first establish the result of kernel ridge regression in the infinite- -aitt limit, and then bound the perturbation on the predict. According the linearization rules for infinitely-wer wide networks (Lee et al. 2019), the output function can be expressed
Proof of Theorem B.
tx, t) (x) (t) (0) + de (x) T (0, (t) (,)))).
where (x) = f(x, Q(0)), and &e (x) = Ve, Q(0).. Recall that 0,, does not change during training, calculated in the proof of Lemma A.6 then the KL divergence reduces
1 (1) (0))
KL. 2
22
25

---

-y Under review as a conference paper at ICLR 2023
Then the gradient flow equation for f) becomes,
d0,(() (L(Q(I)
dt
(Jiik(X, - [ 02(
p (O,,tt)
(X) y ( - -x)
which is an ordinary differrntiall equation regarding (t), and the solution is,
(t) Φμ (x) (op X. X) + A/c "
y(I e ( (xX)+)/621)))))
Plug this result into the linearization of expected output function, we have,
fuck (x, t) + (x, ( [X X A/( 'y(I (0 (xx)x
Then we take the time to be infinity and have
fakk (x) t- + (x, (X, + +
The next step show that
f(x, Q(1) - fntx(x) - O(E).
where & = Einit + log( Es Eg with [xx, Q(0)) - Eiit and limm-+oo θ - (2)(() S Eg.
The proof relies a careful analysis on the trajectories induced by gradient flows for optimizing the neural network and the NTK predictor. The detailed proof can be found in the proof of Theorem 3.2 in Arora et al. (2019b), and we can replace kernel ridge regression here by kernel regression.
C PROOFS OF SECTION 4.3
C.I PROOF OF THEOREM
Theorem C.1 (Restatement of Theorem 4.4). Suppose data S = {(xi, Bi), are i..d samples from non-degenerate distribution DD and m 2 poly(n, (,11)11. Consider any loss function L. x 1) that I-Lipschitz the first argument such that f(y, = 0. Then with probability of at least 1 - 5 over the random initialization and the training samples, the probabilistic neural network (PNN) trained by gradient descent for T S2( L log 4) iterations has populatioon risk (D( that is bounded as follows:
1po0
(X
y (X, + A/22T) 2y
2ym log 7
Rp(Q)<
nc2
Proof of Theorem C.I The generalization bound consists two terms, one is the empirical error, and another is KL divergence.
(1) We first bound the empirical error Σ (Sock Q(t oo)) oo)) y:)2 with following inequality,
7
Σ (fni Q(t = xc)) - -
O (X, (
(x,
16 X 1/21)
= =V/ (
x +
26

---

Under review as a conference paper at ICLR 2023
Then further bound the error term as follows:
(fmuk(Xi), ,y1)
1
{(fmu((x.)) E(yer vi)
72
n
1
fnik(Xe)) 3b
71
1=1
2
fuck fuek((;)) 1/i
nn
X, X) + 1/221) Py (
2
n
(2) The next step to calculate the KL divergence. According the solution of differential equation in Theorem B.I we have,
(t) 0₂(00 =
then = oo yields:
- 1,
"y
Therefore, the KL divergence is,
KL 1/e y (0 (X
(x XX "y
1 < 2 x X) + +//2T)
Finally, by Equation 5 we achieve the PAC- PAC-Bayesian generalization bound,
(X
y
( X,
x)
1) y
log 2
(x(Q)= S
+ +0
nc
C.2 PROOF OF THEOREM 4.5
Theorem C.2 (Restatement of Theorem 4.5). Suppose data S y:)) are i..d. samples from non-degenerate distribution DD and m 2 poly(n, .8-).. Consider any loss function TR x R + [0, that is 1- I-Lipschitz in the first argument such that { = 0. Then with probubility of at least - over the random initialization and training samples, the deterministic neural network trained by gradient descent for T 2 O 1/A0 that is bounded follows:
log z iterations has population risk Rp
y
y T (X +Xx21 2y + +0 #
log And
Rp
n
Proof of Theorem C.2 In this proof, Rademaaher-complexitg analysis. Let H. be the re- producing kernel Hilbert space (RKHS) corresponding to the kernel k(:,)) -). It is known that the RKHS norm of function [x)((x) e (x, X) (X.
X/e2T) -ty y a TTxx, X) is (0 X X) + x/cg1) -ly Then
[[] VaTk k(X, X)o, where O [X, X and a
can bound the [/mtillz
27

---

Under review as a conference paper at ICLR 2023
/ x 1t +1/22))
s S (
For function class FB {f(x) T k(x X
If <
B}. it shown that its empirical
Rademacher complexity can be bounded as Arora et al. 2019a).
LB sup TE
B [T X)
Rs(Fn)
Ixihns LSEF
n
Assume that Tr|k X)] 2 n. Recall the standard generalization bound from Rademacher complex- ity, with probability at least 6, we have,
log(2/8)
sup SEF
E Ep(((f(x))
E (f(())) s $2Rs(F) + +3
2n
There have,
y (X,,) y
-2y ( (X 1/ 1
+0
log Aad
RD S
D ADDITIONAL EXPERIMENTS
This section contains additional experimental results. Training is performed with a server with a CPU with 5,120 cores, and 32 GB Nvidia Quadro V100.
D.1 VALIDATION OF THEORETICAL, RESULTS
We first provide empirical support showing that the training dynamics of wide probabilistic neural networks using the training objective derived from PAC-Bayes bound captured by PNTK, which validates Lemma A.6
Consider three hidden layer ReLU fully-connected network of the training objective derived from the Acy-Bayesian lambda bound in Equation (5), using an ordinary MSE function as loss. The neural network is trained with a full-batch gradient descent using learning rates equal subset of MNIST (ID) = 128) of ten classifications. A random initialized prior with no connection to data used since it is in line with our theoretical setting and we only intend to observe the change in parameters rather than the performance of the actual bound.
on fixed
After T 217 steps of gradient descent updates from different random initialization, we plot the changes of w:() and ( of input/output/hiddden layer with respect to width m for each layer on Figure 2 We observe that the relative Frobenius norm change in the input/output layer's weights
scales as 1 /m while the hidden layers' weight scales is 1/m during the training, which verifies
Lemma A.6 6
D.2 COMPARISON BETWEEN THFORETICAL BOUNDS AND EMPIRICAL. BOUNDS
We make a comparison between theoretical bounds (Equations 14 15 and empirical bounds. The experiments are performed on two different network structures, fully connected neural network and convoluted neural network on MNIST and CIFAR 10 datasets. In particular, we build a 3-layer fully-connected neural network with 600 neurons on each layer. The convolutional architecture is equipped with total of 13 layers with around 10 million learnable parameters. We adopt the same reper-parameter for both theoretical bounds and empirical bounds. The result is shown in Figure First, for theoretical bounds, find that the MAC-s bound is smaller than the Rademacher bound. Secondly, we find that both theoretical bounds are larger than empirical bounds, which meets expectations.
28

---

Under review as a conference paper at ICLR 2023
(t)) //O.I1//I///000
[[W - (o
-
2 2-9
2
2 2
2
2 13
2-11
2 1
2-13
(input)
2
Layer (input)
Layer 2
Layer
2
2
Layer
Layer
2 2
Layer (output)
2 21
2 2
1VM
1
24
25
26
27
28
29
2
221
24
2
26
27
28
29
2
2
m
Figure 2: Relative Frobenius norm change in μ and respectively during training with MSE loss which is derived from the classic PAC-Bayesian bound, where m is the width of the network.
Compar ison of Gereral ization Bounds
62
01
3.66
49
3.12
0.02
0.01
FCN, MNIST
CNN, MNIST
CIFARI
CIFARI0
Figure 3: Comparing NTK Rademacher bound, NTK PAC-Bayesian bound and Empirical Bound with different datasets and network structures.
V
fl
IV,.1
35
30
14,1
28. 54
29. 2
25
21. 93
20
15
10
5
0.
65
0.83
0.87
0
t=0
t=10
t=100
Figure 4: Comparison between the gradient of mean μ and standard deviations
29

---

Under review as a conference paper at ICLR 2023
Kendall taumo. 70
0.09
Error Bound
0.08
0.07
0.06
8
0.05
0.04
500
1000
1500
PA
Figure 5: Correlation between aggregated proxy P.A and generalization bound.
D.3 COMPARISON OF GRADIENT NORM WITH RESPECT TO MEAN WEIGHT AND VARIANCE WRIGHT
We then conduct an experiment compare the gradient of norm with respect to 0, and 0.. The result is shown in Figure 4 We can see that the gradient norm of f(x) is much larger than that of Vo,f(xx fx), which implies that 0 effectively fixed during gradient descent training.
D.4 CORRELATION BETWEEN GENERALIZATION BOUND PROXY METRIC AND GENERALIZATION BOUND
In Figure we observe a positive and significant correlation between PA and generalization bound held among different values of selected hyperparameter while fixing other hyperparameters. Furthermore, we provide a Figure presenting the correlation for aggregated values of Po and A, under the cireumstance where 50% data is used for prior training. We clearly see that lower PA corresponds to the lower bound, with strong positive Kendall-tau correlation of 0.7.
D.5
GRID SEARCH
For selecting hyperparameters, we conduct grid search over Po. percent of prior data, and KL penalty X. Notably, we do grid sweep over the data for prior training with different proportion in [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9] since 0.2 is the minimum proportion required for obtaining a reasonably lower value generalization bound (Dziugaite et al., 2020). For the rest, we run over Po value [0.03, 0.05, 0.07, 0.09, 1. 0.3, 0.5, 0.7] for FCN ([0.05, 0.07, 0.09, 0.1, 0.3, 0.5, 0.7, 0.9) for CNN) and KI. penalty at [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1] for both structures.
30


Reviews:
[{"review_id": "xT76b0HE63g", "reviewer": "Reviewer_Fsnz", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper authors study wide stochastic networks, where the weights are sampled from a Gaussian with diagonal covariance. Authors aim to give a convergence analysis when the means and variances are obtained via gradient descent trained on the PAC-Bayesian bound objective.  \n\nAuthors obtain these results by introducing probabilistic neural tangent kernels that explain the dynamics of the training of mean weights and variances.", "summary_of_the_review": "I recommend rejecting this paper. Main reason for my recommendation is that the results are not clearly written and likely have mistakes in them. ", "strength_and_weaknesses": "Strengths\n\n- Results obtained are quite involved due to the nature of the convergence analysis\n\nWeaknesses\n\n- Paper lacks clarity; the results obtained are difficult to follow and authors don’t give a palatable proof sketch and defer to the appendix for the detail which is not feasible to review for the conference.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "wIAwdtKXKpt", "reviewer": "Reviewer_ag6P", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work theoretically proves the training dynamics of deep and wide probabilistic neural networks and generalization guarantees described by the minimum eigenvalue of the probabilistic NTK. The PAC-Bayesian framework follows the previous classical NTK based, lazy-training results on DNNs, e.g., using the squared loss and gradient descent to prove the kernel equivalence. Besides, a training-free proxy based the derived theoretical results is used for hyper-parameter selection. \n", "summary_of_the_review": "In summary, this paper presents the NTK-based analysis for deep and wide probabilistic neural networks in training dynamics and generalization guarantees. Nevertheless, there is not enough enthusiasm to vote for acceptance due to the not significant results and the not new technical proof.", "strength_and_weaknesses": "**Pros:** \n\n1. Deriving the training dynamics and kernel equivalence for deep and wide probabilistic neural networks under the NTK regime\n2. Building a PAC-Bayesian framework for generalization guarantees of probabilistic neural networks\n\n**Cons:**  \n\n1. This paper focuses on a certain style of DNNs, where the weights always follow a certain distribution. The authors follow the re-parameterization trick to ensure the weight distribution to be Gaussian during training. I think this is a significant assumption/setting by Eq. (2), which is different from classical deep learning theory literature. In this case, the fourth point in the contribution requires well-polished to avoid some over-claim.\n\n2. In theorem 4.2, the authors forgot a significant assumption: lambda_0(K) > 0. Another drawback is the exponential order the depth L due to the perturbation propagation. \n\n3. More importantly, the matrix in Definition 4.1 is a Jacobian matrix instead of NTK. The authors put the probabilistic NTK into an important position but no formulation of the final PNTK is given. Only \\Theta^{\\inf} and K^{(l)}_{\\inf} in an inner product formulation is not enough until I find it in the appendix at page 17, 18. The paper requires a better organization.\n\n4. Besides, I don’t agree with the authors claiming that PAC-Bayesian bound has an improvement over the Rademacher complexity-based bound. Comparing Eq. (14) and (15), even though the first term in both equations is different, the final convergence rate is the same at the $O(1/\\sqrt{n})$ order due to the second term. In fact, Rademacher complexity bound is more general than NTK based bound. This is because, if the neural networks don’t work in a NTK regime, Eq. (14) can not be obtained but Rademacher complexity bound independent of NTK still holds.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "GJuXkHdQQH5", "reviewer": "Reviewer_7Cjm", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers the problem of learning a probabilistic neural network using the regularizer that comes from the PAC-Bayes generalization bounds. The main assumption is that the weights are sampled from Gaussian distribution with given mean and variance. The training algorithm is based on applying the gradient descent on the mean and variance of each weight. The training loss function is based on the PAC-Bayes bound which consists of an empirical risk and KL term. Then, in the limit that the width goes to infinity the authors provide a closed-form expression on the output of the neural network similar to NTK results.  The author also provides some numerical results which show the effectiveness of their approach.", "summary_of_the_review": "The paper considers an important problem: using PAC-Bayes objective to train a probabilistic neural network. The contribution of the results is not clear in terms of the differences between NTK literature. ", "strength_and_weaknesses": "The problem setup in the paper is very interesting. However, I found the paper quite difficult to follow. The connection between theorems is not very clear. More importantly, the author did not discuss how their proof technique is differ from the NTK results.\n\nQuestions:\n\n1- I think equation 3 only works if the loss function is 0-1. Later in the paper the squared loss is used. Could you clarify this point?\n\n2- What does $+- \\xi$ in Eq. (13) mean?\n\n3- What is the intuition behind Equation (16)?\n\n4- In the numerical results section, why is there no plot on the actual generalization versus the theoretical bounds?\n\n5- In Thm 4.2 the KL term is dropped and the GD is only on the empirical risk. Later, in Thm 4.3 the authors also consider the KL term. I don't quite understand the intuition behind the proof technique here. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-3QsZDTwGHT", "reviewer": "Reviewer_DZ2U", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "\nThis paper aims to theoretically investigate the optimization and generalization of deep PAC-Bayesian learning in the setting of wide neural networks. Technically, based on \nthe analysis ideas of the neural tangent kernel (NTK) for the deterministic models under the overparameterized settings, this paper characterizes the convergence of the probabilistic neural \nnetworks, which builds the connection with the kernel ridge regression via the probabilistic neural tangent kernel (PNTK). Then, the authors give the PAC-Bayesian generalization bound, \nwhich is an improvement over the Rademacher complexity-based bound. Further, inspired by the obtained bound, they propose a proxy measure for efficient hyperparameter selection. Finally, some experimental results are also \nprovided.", "summary_of_the_review": "In summary, this paper theoretically characterizes the optimization and generalization of deep probabilistic neural networks via the analytical framework of the neural tangent kernel. It is valuable if it is the first time to provide these results.\nHowever, this paper needs some necessary revisions as in the 'Weaknesses' part. ", "strength_and_weaknesses": "\n## Strength\n1. Overall, this paper is well-organized and written clearly.\n2. The proofs of the theoretical results seem correct although I have not checked them line-by-line.\n3. The obtained theory results are valuable if it is the first time for deep probabilistic neural networks. \n4. Based on the theory results, the proposed proxy measure for efficient hyperparameter selection is effective with experimental support.\n\n\n\n## Weaknesses\n1. As highlighted for deterministic models, the analysis framework based on the neural tangent kernel has limitations, such as it cannot characterize the feature learning process in deep learning. This paper also has these limitations, which are for probabilistic models.\nThis paper does not discuss these limitations and more discussions should be added.\n2. The technical novelty may be limited since the NTK-based analyses have been widely used in deep learning theory although this paper considers the probabilistic neural networks. \n3. In my view, the comparisons between obtained the PAC-Bayesian generalization bound and the Rademacher complexity-based one might not be fair because of the following reasons. First, the definitions of the generalization errors for these two are different. Second, from these two bounds, although we \ncan find the key differences are the first term, as the authors claimed, one is for $O(\\frac{1}{n})$, and another is $0(\\frac{1}{\\sqrt{n}})$, but these bounds totally depending on $o(\\frac{1}{\\sqrt{n}})$. Further, the authors claimed that the obtained PAC-Bayesian bound is non-vacuous in the experiments.\nThus, I recommend the authors do additional experiments to compare these two bounds and related discussions should be added.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject