Demystify Optimization and Generalization of Over-parameterized PAC-Bayesian Learning

We use bold-faced letters for vectors and matrices otherwise representing scalar. We use $\|\cdot\|_{2}$ to denote the Euclidean norm of a vector or the spectral norm of a matrix, while denoting $\|\cdot\|_{F}$ as the Frobenius norm of a matrix. For a neural network, we denote $\phi(x)$ as the activation function and we adopt ReLU activation in this work. Given a logical variable $a$, then $\mathbb{I}(a)=1$ if $a$ is true otherwise $\mathbb{I}(a)=0$. Let ${\bf I}_{d}$ be the identity matrix with dimension of $\mathbb{R}^{d\times d}$. We denote $[n]=\{1,2,\ldots,n\}$. The least eigenvalue of matrix $\bf A$ is denoted as $\lambda_{0}({\bf A})=\lambda_{\min}(\bf A)$.

In PAC-Bayesian learning we use probabilistic neural networks instead of deterministic networks, where the weights always follow a certain distribution. In this work, we adopt the Gaussian distribution for the weights, and define a two-layer probabilistic neural network governed by the following expression,

where ${\bf x}\in\mathbb{R}^{d}$ is the input, ${\bf W}\in\mathbb{R}^{m\times d}$ are weight matrix in the first layer, and ${\bf v}\in\mathbb{R}^{m}$ is the weight vector in the second layer. For simplicity, we represent the weight matrix as a set of column vectors, i.e., ${\bf W}=({\bf w}_{1},{\bf w}_{2},\dots,{\bf w}_{m})^{\top}$. In particular, we use the re-parametrization trick to treat the weight in the first layer and leave the weights in the second to be initialized by uniform distribution,

where $\odot$ denotes the element-wide product operation, $\boldsymbol{\mu}$ and $\boldsymbol{\sigma}$ are mean and variance respectively. We then fix the second layer ${\bf v}$ and only optimize the first layer. This is a trick that has been widely used in the NTK paper (Arora et al., 2019a; Du et al., 2018) for simplifying analysis, the same results can be obtained on deep probabilistic neural network.

Suppose data $S=\{({\bf x}_{i},y_{i})\}_{i=1}^{n}$ are i.i.d. samples from a non-degenerate distribution $\mathcal{D}$. Defining that the $\mathcal{H}$ represents the hypothesis space, $h(\bf{x})$ is the prediction of hypothesis $h\in\mathcal{H}$ over sample $\bf{x}$. $R_{\mathcal{D}}(h)=\mathbb{E}_{({\bf x},y)\sim\mathcal{D}}[\ell(y,h(\bf{x}))]$ represents the generalization error of classifier $h$ and $\widehat{R}_{S}(h)=\frac{1}{n}\sum_{i=1}^{n}\ell\left(y_{i},h\left({\bf x}_{i}\right)\right)$ represents the empirical error of classifier $h$, where $\ell(\cdot)$ is the loss function.

The PAC-Bayes theorem (Langford & Seeger, 2001; Seeger, 2002; Maurer, 2004) concludes that,

Here $R_{\mathcal{D}}(Q)=\mathbb{E}_{({\bf x},y)\sim\mathcal{D},h\sim Q}[\ell(y,h({\bf x}))]=\mathbb{E}_{h\sim Q}[R_{\mathcal{D}}(h)]$ and $\widehat{R}_{S}(Q)=\mathbb{E}_{h\sim Q}[\widehat{R}_{S}(h)]$. $\mathrm{KL}(Q\|Q_{0})=\mathbb{E}_{Q}\left[\ln\frac{Q}{Q_{0}}\right]$ is the Kullback-Leibler (KL) divergence and ${\rm kl}(q\|q^{\prime})=q\log(\frac{q}{q^{\prime}})+(1-q)\log(\frac{1-q}{1-q^{\prime}})$ is the binary KL divergence. Furthermore, combining with pinsker’s inequality for binary KL divergence, ${\rm kl}(\hat{p}\|p)\geq(p-\hat{p})^{2}/(2p)$, when $\hat{p}<p$, yields,

The above bound is the classical bound. On the other hand, combining with the inequality $\sqrt{ab}\leq\frac{1}{2}(\bar{\lambda}a+\frac{b}{\bar{\lambda}})$, for all $\bar{\lambda}>0$, lead to a PAC-Bayes-$\lambda$ bound, which is proposed by (Thiemann et al., 2017),

In this work, we study the PAC-Bayesian learning by training the probabilistic network described by Equation (1) with an objective PAC-Bayes bound (5). Specifically, the objective function can be expressed as,

where $\lambda$ is a hyperparameter introduced in a heuristic manner to make the method more flexible, and $\bar{\lambda}=2$. Since the term regarding $\delta$ is a constant, we omit it in the objective function. Instead of optimizing ${\bf W}$ itself directly, we introduce reparameterization trick (Kingma et al., 2015) to perform gradient descent,

where $\eta$ is the learning rate.

The neural tangent kernel was originated from (Jacot et al., 2018), whose definition is given by,

where $\boldsymbol{\theta}\in\mathbb{R}^{p\times 1}$ with $p$ being the number of parameters in the model. We demonstrate how it is used to describe the training dynamics of deep neural networks.

The dynamics of gradient flow for parameters are given by,

Then the dynamics of output functions $f({\bf X})={\rm vec}(f({\bf x})_{{\bf x}\in{\bf X}})\in\mathbb{R}^{nm\times 1}$ follow,

It is proven that when the width of neural networks is infinite, $\boldsymbol{\Theta}({\bf X},{\bf X},t)$ converge in probability to $\boldsymbol{\Theta}({\bf X},{\bf X})$ (Jacot et al., 2018), i.e., $\boldsymbol{\Theta}({\bf X},{\bf X},t)=\boldsymbol{\Theta}({\bf X},{\bf X})$. Assuming that the loss is the mean squared error loss, $\widehat{R}_{S}=\frac{1}{2}\|f({\bf X})-{\bf Y})\|^{2}_{2}$, then we have,

At the end of the training, namely, $t\rightarrow\infty$, the solution is equivalent to the well-known kernel regression, of which the kernel is NTK.

To simplify the analysis, we first consider the optimization of probabilistic neural networks of the form (1) with objective $\widehat{R}_{S}(Q)$ and adopting the mean squared error loss. In other words, we neglect the KL divergence term at this stage and will show that the corresponding results can be extended to the target function with KL divergence. To perform the optimization analysis, we define the neural tangent kernel for probabilistic neural network, and name them as probabilistic neural tangent kernel (PNTK).

Different from standard (deterministic) neural network, the probabilistic network consist of two sets of parameters, i.e., $\boldsymbol{\mu}_{r}$ and $\boldsymbol{\sigma}_{r}$, with $r=[m]$. The gradient descent based on reparameteriztion trick (Equation 7) is performed on each of the two parameter set separately. As a result, we have two corresponding tangent kernels. It is shown that the NTK will converge in probability to a deterministic kernel in the infinite-width limit (Jacot et al., 2018), as we will show in later derivation, the PNTK will converge to a limiting kernel at initialization and during training with a large width condition. The detailed expression for the limiting kernel is given as follows

The definition above gives the exact formulation of the limiting NTK for random networks. In addition, its positive definiteness is a critical factor to guarantee the convergence of gradient descent, thus we state the assumption as follows,

With all the preliminaries above, we state our main theorem,

Our theorem establishes that if $m$ is large enough, the expected training error converges to zero at a linear rate. In particular, the least eigenvalue of PNTK governs the convergence rate. To prove the theorem, we split the process into three parts, namely, PNTK at initialization, PNTK during training and changes of weights during training, and all the detailed proofs for our theorems and lemmas can be found in the Appendix.

We first study the behavior of tangent kernels with ultra-wide condition, namely $m={\rm ploy}(n,1/\lambda_{0},1/\delta)$, at initialization. The following lemma demonstrates that if $m$ is large, then $\boldsymbol{\Theta}^{\mu}_{0}$ and $\boldsymbol{\Theta}^{\sigma}_{0}$ have a lower bound on smallest eigenvalue with high probability. The proof is by the standard concentration bound.

From the above lemma, we implicitly show that both tangent kernels regarding $\boldsymbol{\mu}_{r}$ and $\boldsymbol{\sigma}_{r}$ are equivalent in the infinite-width limit. It is also true in general as long as the reparametrization trick is adopted for training probabilistic neural networks. The next problem is that PNTKs are time-dependent matrices, thus varying during gradient descent training. To account for this problem, we build a lemma stating that if the weights ${\bf w}(t)$ is close to ${\bf w}(0)$ during gradient descent training, then the corresponding PNTKs $\boldsymbol{\Theta}^{\mu}(t)$ and $\boldsymbol{\Theta}^{\sigma}(t)$ are close to $\boldsymbol{\Theta}^{\infty}$ and have a lower bound over smallest eigenvalue $\lambda_{0}$. Here we introduce ${\bf w}_{r}(t)\equiv\boldsymbol{\mu}_{r}(t)+\boldsymbol{\sigma}_{r}(t)\odot\boldsymbol{\epsilon}_{r}(0)$, where $\boldsymbol{\epsilon}_{r}(0)$ copies sample value from that in ${\bf w}_{r}(0)$. Notice ${\bf w}_{r}(t)$ in our definition is not exact weights at training step $t$, but we introduce it to facilitate our proof.

By the above lemma, we demonstrate that if the change of weight is bounded, then the tangent kernel matrix is closed to its expectation. The next lemma will show that the change of weight is bounded during gradient descent training when the PNTK is close to the limiting kernel Equation (13).

The main gap between PNTK and final results for output function can be filled by writing down the gradient flow dynamics,

Combined with fact that PNTKs are bounded during training, therefore the behavior of the loss is traceable. To finalize the proof for Theorem 4.4, we construct a final lemma.

According to Equation (6), there is a KL divergence term in the objective function. We expand the KL-divergence for two Gaussian distribution, $p({w}({t}))=\mathcal{N}({\mu}_{t},{\sigma}_{t})$, and $p({w}({0}))=\mathcal{N}({\mu}_{0},{\sigma}_{0})$,

We compare the gradient of mean weights $\boldsymbol{\mu}$ and variance weights $\boldsymbol{\sigma}$. With a direct calculation, we have $\frac{\partial f({\bf x}_{i})}{\partial\boldsymbol{\mu}_{r}}=\frac{1}{\sqrt{m}}{v}_{r}\mathbb{I}({\bf w}^{\top}_{r}{\bf x}_{i}\geq 0){\bf x}_{i}$ and $\frac{\partial f({\bf x}_{i})}{\partial\boldsymbol{\sigma}_{r}}=\frac{1}{\sqrt{m}}{v}_{r}\mathbb{I}({\bf w}^{\top}_{r}{\bf x}_{i}\geq 0){\bf x}_{i}\odot\boldsymbol{\epsilon}_{r}$. It is shown that there is one more random variable $\boldsymbol{\epsilon}_{r}$ associated with the gradient regarding variance weights, which leads to the expected gradient norm to be zero. To verify this result, we conduct an experiment and leave the results in the Appendix E.2. Therefore, it is approximately equivalent to fix $\boldsymbol{\sigma}$ during gradient descent training. With this assumption we arrive at the conclusion:

The theorem above reveals that the regularization effect of the KL term in PAC-Bayesian learning, and presents an explicit expression for the convergence result of the output function.

Recall that in Theorem 3.2, the PAC-Bayesian bound with respect to the distribution at initialization and after optimization is given. Therefore, combined with results from Theorem 4.10, we can directly provide a generalization bound for PAC-Bayesian learning with ultra-wide networks.

In this theorem, we establish a reasonable generalization bound for PAC-Bayesian learning framework, thus providing a theoretical guarantee. We further demonstrate the advantage of PAC-Bayesian learning by comparing it with the Rademacher complexity-based generalization bound for deterministic neural network with kernel ridge solution.

Theorem 4.12 is obtained by following Theaorem 5.1 in (Hu et al., 2019), who presents a Rademacher complexity-based generalization bound for ultra-wide neural networks with kernel ridge regression solution. Similar analysis for kernel regression with NTK can be found in (Arora et al., 2019a; Cao & Gu, 2019).

The main difference between two generalization bound is $\frac{{\bf Y}^{\top}(\boldsymbol{\Theta}+{\lambda}/{\sigma^{2}_{0}}{\bf I})^{-1}{\bf Y}}{n}$ versus $\sqrt{\frac{{\bf Y}^{\top}(\boldsymbol{\Theta}+{\lambda}/{\sigma^{2}_{0}}{\bf I})^{-1}{\bf Y}}{n}}$, which is due to that PAC-Bayesian bound count the KL divergence while Rademacher bound calculate the reproducing kernel Hilbert space (RKHS) norm. We find that the rates of convergence are different. One is $O(1/n)$ and another is $O(1/\sqrt{n})$. Therefore, we conclude that the PAC-Bayesian bound has an improvement over Rademacher complexity-based bound.

In all experiments, the NTK parameterization is chosen for initializing parameters, which follows Equation (1). Specifically, the initial mean for weights, $\mu_{0}$, is sampled from a truncated Gaussian distribution with a mean of 0 and a standard variance of $1$, truncating at 2 standard deviations. To make sure that the variance is positive, the initial variance for weight is transformed from the given value of $\rho_{0}$ through formula $\sigma_{0}=\log(1+\exp(\rho_{0}))$.

For the experiment in section 5.2, we consider a three hidden layer ReLU fully-connected network of the training objective derived from the PAC-Bayesian lambda bound in Equation (5), used an ordinary MSE function as loss, trained with full-batch gradient descent using learning rates equal to 1 on a fixed subset of MNIST ($|D|=128$) of ten-classification. An random initialized prior with no connection to data is used since it is inline with our theoretical setting and we only intend to observe the change in parameters rather than the performance of the actual bound.

In section 5.3, we use both fully connected and convoluted neural network structures to perform experiments on MINIST and CIFAR-10 datasets for demonstrating the effectiveness of our training-free PAC-Bayesian network bound for searching hyperparameters under different datasets and network structures. We build a 3 layers fully-connected neural network with 600 neurons on each layer. Another convolutional architecture with a total of 13 layers with around 10 million learnable parameters is constructed. We adopt data-dependent prior since it is a practical and popular method (Dziugaite et al., 2020; Perez-Ortiz et al., 2021). Specifically, this data-dependent prior is pre-trained on a subset of total training data with empirical risk minimization. The networks for posterior training is then initialized by the weights learnt from prior. At last, the generalization bound is computed using Equation (5). Bound computing-related settings are referred to the work done by Pérez-Ortiz et al. (2020), such as, confidence parameters for the risk certificate and chernoff bound, and the number of monte carlo samples for estimating the risk certificate.

After $T=2^{17}$ steps of gradient descent updates from different random initialization, we plot the changes of (input/output/hidden layer) $\boldsymbol{\mu}$ and $\boldsymbol{\rho}$ for weights corresponded to the variation in width $m$ for each layer on Figure 1. We observe that the relative Frobenius norm change in input/output layer’s weights scales as $1/\sqrt{m}$ while hidden layers’ weights scales as $1/m$ during the training, which verifies our lemma 4.8.

PAC-Bayesian learning framework provides competitive performance with non-vacuous generalization 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 $\rho_{0}$, and the KL penalty weight ($\lambda$). Since these three values do not change during the training, we refer them as hyperparameters. Choosing the right hyperparameters via grid search is obviously prohibitive, as each attempt to compute the generalization bound can involve significant computational resources.

Another plausible approach is to design some kind of predictive, so-called “training-free” metric such that we can approximate the error bound without going through an expensive training process. In light of this goal, we have already developed a generalization bound in the theorem D.1 via NTK. Since NTK changes are held constant during training, we can predict the generalization bound by this proxy metric, which can be formulated as follows:

where $\boldsymbol{\Theta}$ denotes the gram matrix. ${\bf Y}{\bf Y}^{\top}$ is a $n\times n$ label similarity matrix (if two data have the same label, their joint entry in the matrix is 1 and 0 otherwise), and $n$ is the number of data used. To demonstrate the computational practicality of this training-free metric, we compute $\mathcal{PA}$ using only a 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, they can be found in NAS (Chen et al., 2021; Deshpande et al., 2021), MAE Random Sampling (Camero et al., 2021), pruning at initialization (Abdelfattah et al., 2021). To the best of our knowledge, there is currently no training-free method for selecting hyperparameters in the PAC-Bayesian framework, which we consider to be one of the novelties of this paper.

The Figure 2 demonstrates a strong correlation between $\mathcal{PA}$ and the actual generalization bound. Finally, we demonstrate that by searching through all possible combinations of hyperparameters using $\mathcal{PA}$, it is possible to select a hyperparameter leading towards a result that is comparable to the best generalization bound, but without the excessive computation. To put things in perspective, in Table 1, we show that grid search can be much more expensive than $\mathcal{PA}$. For instance, under the setting of CNN and CIFAR10 dataset, we save 42.5 times computational time to find the bound that is 1% lower than the best generalization bound.