SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhengmian Hu, Heng Huang
College Park, MD 20740
huzhengmian@gmail.com,henghuanghh@gmail.com
This paper proposes a novel algorithm, Transformative Bayesian Learning (TransBL), which bridges the gap between empirical risk minimization (ERM) and Bayesian learning for neural networks. We compare ERM, which uses gradient descent to optimize, and Bayesian learning with importance sampling for their generalization and computational complexity. We derive the first algorithm-dependent PAC-Bayesian generalization bound for infinitely wide networks based on an exact KL divergence between the trained posterior distribution obtained by infinitesimal step size gradient descent and a Gaussian prior. Moreover, we show how to transform gradient-based optimization into importance sampling by incorporating a weight. While Bayesian learning has better generalization, it suffers from low sampling efficiency. Optimization methods, on the other hand, have good sampling efficiency but poor generalization. Our proposed algorithm TransBL enables a trade-off between generalization and sampling efficiency.
Deep Neural Networks (DNNs) have achieved remarkable success in machine learning and related applications. It repeatedly outperformed conventional machine learning approaches, resulting in ground-breaking research such as human-level performance in computer vision [31], substantial progress in natural language processing [10], and mastering the game of Go [65].
The success of DNNs is the result of a critical combination of the complexity and generalization. On the one hand, universal approximation theorem [34] guarantees that any continuous function can be approximated arbitrarily well by using a deep network. On the other hand, deep architectures, together with large-scale training data and back-propagation algorithm, present a good generalization ability towards unseen data.
Although expressive power and generalization ability are both desirable on their own, they are mutually incompatible. Being one of the main contributions of statistical learning theory, the probably approximately correct (PAC) learning [69] allows us to establish an upper bound of generalization gap by capacity of a hypothesis space. Generally speaking, if a model is capable of fitting random labels, then it must generalize poorly.
Recent result [79] shows that overparameterized DNNs do fit random label perfectly. However, it is observed that more steps of stochastic gradient descent (SGD) are needed to train a neural network to fit random labels. This phenomenon suggests that the learning capability of DNNs increases with the number of training steps. In light of this observation, the algorithm-dependent generalization bounds which control the complexity by maximum training steps are preferred over uniform convergence bounds.
In this paper, we follow the PAC-Bayesian approach to derive our algorithm-dependent generalization bounds. PAC-Bayesian bounds were first introduced by McAllester [53] and further developed in [61, 52, 13]. These bounds apply for stochastic learning algorithms and focus on expected generalization error over a probability on parameter space. The generalization is related to the KL divergence between the output distribution of a learning algorithm and a prior distribution.
Typically, deep neural networks contain a large number of parameters and are trained by numerous training epochs through some gradient descent updates. The training dynamics of deep neural networks is complicated, and hence direct theoretical analysis on the KL divergence is intractable. Fortunately, such situation can be largely simplified in the infinite width limit [55, 73, 42, 51, 18, 39, 43, 76, 3]. Under this limit, the output distribution of stochastic neural networks drawn from the prior is approximated by a Gaussian distribution. Moreover, the training dynamics of gradient-based optimization is governed by the kernel gradient descent, which guarantees that the evolution of network output only depends on the function values themselves. Thus, our first problem to study in this paper is: Q1 - In the infinitely wide neural network limit, can we theoretically derive the formula of this KL divergence?
Optimizing the PAC-Bayesian bound of expected error gives rise to the Gibbs measure. It is also widely termed as posterior, as it shares the same form as posterior in Bayes' rule, given that error is interpreted as likelihood. Drawing samples from this posterior is very hard in practice. Markov Chain Monte Carlo (MCMC) methods have been explored for deep Bayesian learning. However, it suffers from slow convergence and high computational cost on high-dimensional parameter spaces. Non-exact minimization of PAC-Bayesian bound gives rise to variational approximation (VA) which is more computationally efficient, but biased due to the difference between variational distribution and true posterior.
Given that MCMC and VA have their own disadvantages in Bayesian learning, our paper investigates the second problem: Q2 - Does there exist a Bayesian learning method with non-diminishing sampling efficiency even for infinite wide neural network?
Finally, we notice that the gradient-based optimization is efficient in training DNNs, but comes with larger generalization error. Bayesian learning, on the other hand, optimizes the expected loss but has lower efficiency. Our third question to study is: Q3 - Does there exist an interpolation between optimization and Bayesian learning for trade-off between computation efficiency and generalization error?
In this paper, we give positive answers to above three questions and the main contributions of this paper are summarized as follows:
1) We analyze the infinitely wide neural network trained by gradient flow, or equivalently infinitesimal step size gradient descent. We show that the infinite width limit largely simplifies the training dynamics, and the KL divergence between the output distribution and Gaussian prior can be formulated as a function of training time and training data.
2) As a byproduct of our analysis on the generalization and sampling efficiency, we prove that the trace of Hessian for DNN is not diminishing under infinite width limit and depends on the initialization and training. To the best of our knowledge, this dynamics of Hessian trace is new and maybe of independent interest.
3) We show that if the determinant of Jacobian of optimization flow is available, we can compute a weight for each optimized predictor, such that the weighted output distribution is just posterior. The sampling efficiency in infinite width limit is also derived. We call this type of algorithm as Transformative Bayesian Learning (TransBL) because it is transformed from an optimization procedure.
4) We show that modifying the additional weight in TransBL gives rise to an interpolation between optimization and Bayesian learning. The behaviour of TransBL is increasingly similar to optimization when the weight is changed toward being uniform. This interpolation doesn't alter training dynamics, thus enables flexible trade-off between sampling efficiency and generalization.
We first explain the setting and briefly review PAC-Bayesian and Bayesian learning as background. Other related works is discussed in Appendix A.
Problem Setup Given input space  $\mathcal{X}$  and label space  $\mathcal{Y}$ , we assume that labeled training data are drawn independently from an unknown data distribution  $D$  over  $\mathcal{X} \times \mathcal{Y}$ . The training set is denoted as  $S^{\mathrm{train}} = \{(s_a, z_a) | 1, \ldots, m\}$ , where  $m$  is the size of the training sample. A predictor is a function  $h: \mathcal{X} \to \mathbb{R}$  and the set of predictors is denoted as  $\mathcal{H}$ . We introduce two losses  $l: \mathbb{R} \times \mathcal{Y} \to [0,1]$  and  $l^s: \mathbb{R} \times \mathcal{Y} \to \mathbb{R}$ , where the first one is used to measure the error rate and the second smooth surrogate loss with polynomially bounded second order derivative is used for providing learning signal in gradient-based training. We define expected loss as  $R(h) = \mathbb{E}_{(s,z) \sim D}[l(h(s), z)]$  and empirical loss as  $r(h) = \frac{1}{m} \sum_{a=1}^{m} l(h(s_a), z_a)$ . The expected loss is the central quantity which we are interested in but cannot be unobserved in general. Therefore, the currently popular approach to statistical learning is to derive an upper bound of expected loss.
PAC-Bayesian Bounds We will follow the PAC-Bayesian approach to the generalization problem. Theorem 1 (Theorem 1.2.6. in [13]). For any positive  $\lambda$  and  $\delta \in (0,1)$ , with at least  $1 - \delta$  in the probability of training samples, for all distribution  $q$  on space  $\mathcal{H}$ , we have
$$
\mathbb {E} _ {h \sim q} [ R (h) ] \leq \Phi_ {\frac {\lambda}{m}} ^ {- 1} \left(\mathbb {E} _ {h \sim q} [ r (h) ] + \frac {D _ {\mathrm {K L}} (q \| p) + \log \frac {1}{\delta}}{\lambda}\right).
$$
The KL divergence is  $D_{\mathrm{KL}}(q \| p) = \mathbb{E}_{x \sim q} \left[ \log \frac{q(x)}{p(x)} \right]$  and the auxiliary function is defined as
$$
\Phi_ {a} ^ {- 1} (v) = \frac {1 - \exp (- a v)}{1 - \exp (- a)} \leq \frac {a v}{1 - \exp (- a)}.
$$
Although there are huge amounts of research on PAC-Bayesian approach, all of them rely on the KL divergence between output distribution  $q$  of stochastic predictor learned from certain algorithm and a prior  $p$ , therefore the analysis of this KL divergence is a natural question.
Bayesian Learning We define Gibbs measure as  $p_{\lambda}(\theta) = \frac{1}{Z_{\lambda}} e^{-\lambda r(\theta)} p(\theta)$  and partition function as  $Z_{\lambda} = \mathbb{E}_{\theta \sim p}[e^{-\lambda r(\theta)}]$ . The Gibbs measure is related to posterior in Bayesian inference [24], and minimize right-hand side in theorem 1, inducing a much tighter bound:
$$
\mathbb {E} _ {h \sim q} [ R (h) ] \leq \Phi_ {\frac {\lambda}{m}} ^ {- 1} (- \frac {1}{\lambda} \log (Z _ {\lambda} \delta)) \leq \frac {1}{m (1 - \exp (- \frac {\lambda}{m}))} \log \frac {1}{Z _ {\lambda} \delta}.
$$
Despite the fact that expected loss is optimized, Bayesian learning bears significantly more difficulty than popular optimization-based algorithms. It is noted that the sampling from a high dimensional distribution with irregular shape may lead to high fluctuations and low efficiency, yet in typical neural network, the number of parameters to be trained is very large, and loss landscape is complicated. Many sampling techniques have been developed and explored, but the state-of-the-art is still far from satisfactory. In this paper, we concentrate on theoretical understanding of the Bayesian approach and consider a special class of importance sampling in Section 4 where we can compare the computation efficiency of Bayesian method and optimization approach.
3 Training by Optimization
We now derive the KL divergence under deterministic optimization setting. We assume all predictors are parameterized by some parameters  $\theta \in \Theta$ . The expected loss  $R(h)$  and empirical loss  $r(h)$  could therefore be regarded as functions of parameter  $\theta$ . An optimization flow is a function  $f: \Theta \to \Theta$ , which typically relies on the training data to update the model's parameters  $\theta$  and minimize the empirical loss  $r(h)$ . The updating can be achieved through various methods such as single or multiple steps of gradient descent, gradient flow with infinitesimally small step sizes, or more sophisticated methods detailed in Appendix B. We require the Jacobian determinant  $J_{f}(\theta) = \operatorname*{det}(\nabla f(\theta))$  exists and the optimization flow is a bijective, i.e. the inverse  $f^{-1}$  exists everywhere. This is satisfied for gradient flow and gradient descent when step size is smaller than  $\frac{1}{L}$ . More situations where this assumption holds is discussed in Appendix B. The starting point of optimization is initialized from a prior distribution  $p(\theta)$ , and we define a corresponding energy  $V(\theta) = -\log p(\theta) + C$  where  $C$  is a parameter independent constant. The output distribution  $q$  could be characterized as follows:
$$
\Pr_ {q} (\theta \in A) = \Pr_ {p} (f (\theta) \in A), q (\theta) = p \left(f ^ {- 1} (\theta)\right) \mid J _ {f ^ {- 1}} (\theta) \mid .
$$
The KL divergence between output distribution and prior can be derived as follows:
$$
D _ {\mathrm {K L}} (q \| p) = \mathbb {E} _ {\theta \sim q} [ \log \frac {q (\theta)}{p (\theta)} ] = \mathbb {E} _ {\theta \sim p} [ \log \frac {q (f (\theta))}{p (f (\theta))} ] = \mathbb {E} _ {\theta \sim p} [ \log \frac {p (\theta)}{p (f (\theta))} - \log | J _ {f} (\theta) | ].
$$
Based on a physics analog of Helmholtz free energy change in isothermal process, we define energy term and entropy term separately:
$$
\Delta_ {f} V (\theta) = V (f (\theta)) - V (\theta), \quad \Delta_ {f} S (\theta) = \log | J _ {f} (\theta) |.
$$
The energy term only depends on the prior. If Gaussian prior is used,  $\Delta_f V(\theta)$  corresponds to the change of squared norm of parameters. Although the KL divergence is positively related to the energy term, and energy term is monotonically increasing with norm of trained parameters, naively compressing parameters to reduce parameter norm doesn't lead to decrease in KL divergence, because the decrease in energy term is offset by the change in entropy term.
The entropy term determines the gain or loss of entropy caused by an optimization procedure. The pure gain means that the optimization flow maps a region of parameter space into another region with larger area. For example, consider the simple function  $f(x) = 2x$  in a one-dimensional setting. The Jacobian of this function, which describes how much the function stretches or compresses space, is constant at 2. Imagine we start with an input  $x$  that is uniformly distributed within the interval [0, 1]. After applying the function, the output  $f(x)$  becomes uniformly distributed over this larger interval [0, 2]. Therefore the entropy increase by  $\log(2)$ . However, most optimization procedures pursue minimization of certain objective functions which is only small in an extremely small region of parameter space, and thus loss of entropy is inevitable. In Section 7, we show that for gradient descent algorithm, the loss of entropy is related to the local curvature. We also extend the discussion to algorithms with enriched state space, such as Momentum SGD and Adagrad in Appendix F.2.
$$
D _ {\mathrm {K L}} (q \| p) = \mathbb {E} _ {\theta \sim p} [ \Delta_ {f} V (\theta) - \Delta_ {f} S (\theta) ] \tag {1}
$$
The above KL divergence represents an increase in free energy, or the information gain from the training, depending on the point of view. Both energy term and entropy term can be evaluated empirically, though the computational cost depends on the choice of optimization method, network architecture and prior.
4 Transformative Bayesian Learning
In this section, we introduce a class of importance sampling algorithms solving the Bayesian learning problem. These algorithms are based on the optimization flow, energy change, and loss of entropy discussed in Section 3.
The most simple example of TransBL algorithm comes from using the output distribution  $q$  of an optimization flow  $f$  as the proposal distribution. For any function  $F$  on parameter space  $\Theta$ , we have
$$
\begin{array}{l} \mathbb {E} _ {\theta \sim p _ {\lambda}} [ F (\theta) ] = \mathbb {E} _ {\theta \sim q} \left[ F (\theta) \frac {p _ {\lambda} (\theta)}{q (\theta)} \right] = \frac {1}{Z _ {\lambda}} \mathbb {E} _ {\theta \sim p} \left[ F (f (\theta)) e ^ {- \lambda r (f (\theta))} \frac {p (f (\theta))}{q (f (\theta))} \right] \tag {2} \\ = \frac {1}{Z _ {\lambda}} \mathbb {E} _ {\theta \sim p} \left[ F (f (\theta)) e ^ {- \lambda r (f (\theta))} e ^ {- \Delta_ {f} V (\theta)} e ^ {\Delta_ {f} S (\theta)} \right] \\ \end{array}
$$
The above expectation could be regarded as attaching an additional weight to the stochastic predictor obtained from ordinary optimization, and we should use a weighted average of the results obtained by different initialization. This process doesn't involve any more training of parameter  $\theta$  than the optimization procedure that TransBL is based on, though additional computation might be required for the value of  $\Delta_f V(\theta)$  and  $\Delta_f S(\theta)$ .
The unnormalized weight of above importance sampling is  $w_{\lambda}(\theta) = e^{-\lambda r(f(\theta))}e^{-\Delta_fV(\theta)}e^{\Delta_fS(\theta)}$ . That means not all results obtained by training are treated equally. In particular, the solution with lower empirical loss, lower energy increase, and producing higher entropy is more important than others.
One benefit of TransBL is that we can measure the computation efficiency by comparing it to the training by optimization. A simple way to do this is to calculate the ratio of effective sample size and
sample size:
$$
\operatorname {e f f} _ {\lambda} = \frac {\left(\mathbb {E} _ {\theta \sim p} [ w _ {\lambda} (\theta) ]\right) ^ {2}}{\mathbb {E} _ {\theta \sim p} [ w _ {\lambda} ^ {2} (\theta) ]} \tag {3}
$$
We define the above value as sampling efficiency, and a non-diminishing efficiency at certain limit indicates that Bayesian learning is at most a constant factor slower than optimization. The numerator of the above expression is determined by the partition function because  $\mathbb{E}_{\theta \sim p}[w_{\lambda}(\theta)] = Z_{\lambda}$ . Therefore, only the denominator depends on the optimization procedure, and we expect  $\mathbb{E}_{\theta \sim p}[w_{\lambda}^2 (\theta)]$  to be as small as possible for efficient Bayesian learning. Notice that this quantity again relies on energy change  $\Delta_fV(\theta)$  and entropy change  $\Delta_fS(\theta)$ . We will give theoretical analysis of these two values for infinitely wide neural network in Section 7.
For illustration purpose, we consider a univariate loss function that presents two distinct global minima with zero loss. One of these is characterized as a sharp minimum, while the other represents a flat minimum. A direct initialization from the prior, followed by training using gradient flow, often results in an ensemble with significant deviation from the posterior. This is because the optimization process fails to recognize the presence of the sharp minimum, while insights from PAC Bayesian indicate that the flat minimum is surrounded by a higher posterior probability density. TransBL method applies a small weight to the solution found within the sharp minimum. Consequently, TransBL can adeptly recreate the posterior, as shown in the Figure 6b. Due to space limitations, all figures for the illustration are moved to Appendix F.1.
5 Connections between Bayesian Learning and Optimization
5.1 A Bayesian Perspective for Optimization
We show that, in order to achieve low expected loss, the optimization flow in deep learning should reshape the initial distribution to a good estimation of posterior. A formal argument relies on applying Donsker and Varadhan's variational formula [17] to obtain the following equation:
$$
\mathbb {E} _ {\theta \sim q} [ r (\theta) ] + \frac {1}{\lambda} D _ {\mathrm {K L}} (q \| p) = - \frac {1}{\lambda} \log \mathbb {E} _ {\theta \sim p} [ \exp (- \lambda r (\theta)) ] + \frac {1}{\lambda} D _ {\mathrm {K L}} (q \| p _ {\lambda}) \tag {4}
$$
The first term is independent of training procedure and only relies on the definition of prior. The second term measures the KL divergence between output distribution and Gibbs measure. This is also the gap between expected loss bound obtained by some training methods and the optimal one. Notice that the above KL divergence could also be expressed in terms of weight for TransBL, energy change, and entropy loss:
$$
D _ {\mathrm {K L}} (q \| p _ {\lambda}) = \log Z _ {\lambda} + \mathbb {E} _ {\theta \sim p} [ - \log w _ {\lambda} (\theta) ] = \log Z _ {\lambda} + \mathbb {E} _ {\theta \sim p} [ \lambda r (f (\theta)) + \Delta_ {f} V (\theta) - \Delta_ {f} S (\theta) ]. \tag {5}
$$
5.2 Efficiency of Optimization Flows for TransBL
In this section, we show that, in order to achieve high sampling efficiency, the optimization flow in deep learning is again expected to reshape the initial distribution to a good estimation of posterior. The only difference to the previous section is that the deviation between output distribution and posterior is measured under a different divergence.
In order to see this, we define  $\chi^2$  divergence as  $D_{\chi^2}(p\| q) = \mathbb{E}_{x\sim q}[(\frac{p(x)}{q(x)})^2 ] - 1.$
Theorem 2. For sampling efficiency of TransBL whose proposal distribution is  $q$ , we have:
$$
D _ {\mathrm {K L}} \left(p _ {\lambda} \| q\right) \leq \log \left(1 + D _ {\chi^ {2}} \left(p _ {\lambda} \| q\right)\right) = - \log \operatorname {e f f} _ {\lambda} \tag {6}
$$
With the above result, we can establish the equivalence between optimal generalization bound and optimal sampling efficiency, i.e. an optimization flow is optimal in the sense it minimizes the expected loss upper bound if and only if the sampling efficiency of TransBL is 1. A more general correspondence between generalization and sampling efficiency is desired. However, the complication lies in different divergences used in two quantities. By comparing Eq. (4) and Eq. (6),
we see that  $\chi^2$  divergence is used for measuring sampling efficiency. Although an upper bound of KL divergence  $D_{\mathrm{KL}}(p_\lambda \| q)$  could be obtained, this is not directly comparable to the  $D_{\mathrm{KL}}(q \| p_\lambda)$  used in generalization bound because of asymmetry of KL divergence. That difference also justifies the interpolation between optimization and Bayesian learning shown in the next section.
6 Interpolation of Optimization and Bayesian Learning
We recall that the output distribution  $q$  for the optimization and the posterior  $p_{\lambda}$  for Bayesian learning are connected by an additional weight  $w_{\lambda}$ :  $p_{\lambda}(\theta) = \frac{1}{Z_{\lambda}} w_{\lambda}(f^{-1}(\theta))q(\theta)$ . These two distributions have distinct properties. The distribution  $q$  bears worse generalization, but is easy to sample from given an established optimization oracle. The posterior  $p_{\lambda}$ , on the other hand, is better shaped, but hard to sample from. Therefore, a natural question is whether we can find a distribution  $p_{\lambda}^{\beta}$  as intermediate and interpolation between them.
Given our formulation of TransBL, it turned out to be quite easy to construct such an interpolation by modifying the weight. For a function  $v_{\beta} : \mathbb{R}^{+} \to \mathbb{R}^{+}$  where  $\beta$  is the parameter used for interpolation, we can define a modified weight as  $v_{\beta}(w_{\lambda}(\theta))$ . The interpolation distribution is
$$
p _ {\lambda} ^ {\beta} (\theta) = \frac {1}{Z _ {\lambda} ^ {\beta}} v _ {\beta} (w _ {\lambda} (f ^ {- 1} (\theta))) q (\theta), Z _ {\lambda} ^ {\beta} = \mathbb {E} _ {\theta \sim p} [ v _ {\beta} (w _ {\lambda} (\theta)) ].
$$
Note that the superscript doesn't mean power but means the interpolation is controlled by  $\beta$ . The expectation of a function  $F(\theta)$  on  $p_{\lambda}^{\beta}$  could be computed similarly as Eq. (2):
$$
\mathbb {E} _ {\theta \sim p _ {\lambda} ^ {\beta}} [ F (\theta) ] = \frac {1}{Z _ {\lambda} ^ {\beta}} \mathbb {E} _ {\theta \sim p} [ F (f (\theta)) v _ {\beta} (w _ {\lambda} (\theta)) ].
$$
We note Eq. (4) still holds when we change  $q$  into  $p_{\lambda}^{\beta}$ , and the gap between the expected loss bound obtained by  $p_{\lambda}^{\beta}$  and the posterior is:
$$
D _ {\mathrm {K L}} (p _ {\lambda} ^ {\beta} \| p _ {\lambda}) = \log \frac {Z _ {\lambda}}{Z _ {\lambda} ^ {\beta}} + \frac {1}{Z _ {\lambda} ^ {\beta}} \mathbb {E} _ {\theta \sim p} \left[ \log \left(\frac {v _ {\beta} (w _ {\lambda} (\theta))}{w _ {\lambda} (\theta)}\right) v _ {\beta} (w _ {\lambda} (\theta)) \right].
$$
Notice that the above identity degenerates into Eq. (5) if we use  $v_{\beta}(\cdot) = 1$ , which corresponds to the original optimization.