It has been empirically observed

that different local optima, obtained from training deep neural networks

don’t generalize in the same way for the unseen data sets, even if they

achieve the same training loss. Recently, this problem has received

renewed interest from both the empirical and theoretical deep learning

communities. From the theoretical standpoint, most of the generalization

bounds developed for explaining this phenomenon only consider a worst-case

scenario, thus ignoring the generalization capabilities among different

solutions. In this blog we are interested in the following question:

We discuss several related proposals and introduce our recent work that

builds a connection between the model generalization and the Hessian of

the loss function.

img.animated-gif {

width: 450px;

height: auto;

}

.column {

float: left;

width: 33.33%;

padding: 5px;

}

.row::after {

content: “”;

clear: both;

display: table;

}

## Empirical Observations on Hessian and Generalization

Hochreiter

and Schmidhuber, and more recently, Chaudhari et. al. and Keskar et al. argue that the

local curvature, or “sharpness”, of the converged solutions for deep

networks is closely related to the generalization property of the

resulting classifier. The sharp minimizers, which led to lack of

generalization ability, are characterized by a significant number of large

positive eigenvalues in (nabla^2 hat{L}(x)), the loss function being

minimized.

Neyshabur et al.

point out the sharpness itself may not be enough to determine the

generalization capability. They argue that it is necessary to include

the scales of the solution with sharpness to be able to explain

generalization. To this end, they suggest an “expected sharpness” based

on the PAC-Bayes bound: $$E_{usim N(0,sigma^2)^m}[hat{L}(w+u)] –

hat{L}(w)$$

As shown in the figure below, intuitively if the local minima is

“smooth” then perturbing the model won’t cause the objective to change

much. Thus it provides a way of measuring the smoothness of the local

minima.

## How to Perturb the Model?

It is well known that adding noise

to the model during training can help improving the model generalization.

However, how to set proper noise levels remains unknown. Most

state-of-the-art methods assume same level of perturbation along all

directions, but intuition suggests that this may not be appropriate.

Let’s look at a toy example. We construct a small 2-dimensional sample set

from a mixture of (3) Gaussians, and binarize the labels by thresholding

them from the median value. Then we use a (5)-layer MLP model with only

two parameters (w_1) and (w_2) for prediction, and the cross entropy

as the loss. The variables from different layers are shared. Sigmoid is

used as the activation function. The loss landscape of a model trained

using (100) samples is plotted below:

If we look at the optima indicated by the orange vertical bar, the

smoothness of the loss surface along different directions is quite

different. Should we impose the same level of perturbations along all

directions? Perhaps not.

We argue that more noise should be put along the “flat” directions.

Specifically, we suggest adding uniform or (truncated) Gaussian noise

whose level, in each coordinate direction, is roughly proportional to

$$sigma_i approx frac{1}{sqrt{nabla^2_{i,i} hat{L}+ rho N_{gamma,

epsilon}(w_i)} },$$ where (rho) is the local Lipschitz constant of the

Hessian (nabla^2hat{L}), and (N_{gamma, epsilon}(w_i)=gamma |w_i|

+ epsilon).

## A Metric on Model Generalization

As discussed by Dinh et al. and Neyshabur et al., the spectrum of

Hessian itself may not be enough to determine the generalization power of a

model. In particular, for a multi-layer perceptron with RELU as the activation

function, one may re-parameterize the model and scale the Hessian spectrum

arbitrarily without affecting the model prediction and generalization.

Employing some approximations, we propose a metric on model generalization,

called PACGen, that depends on the scales of the parameters, the Hessian, and

the high-order smoothness terms that are characterized by the Lipschitz contant

of Hessian.

$$Psi_{gamma, epsilon}(hat{L}, w^ast) = sum_i log

left((|w_i^ast| + N_{gamma,epsilon}(w_i^ast))maxleft(sqrt{nabla^2_i

hat{L}(w^ast) + rho(w^ast)sqrt{m}N_{gamma,epsilon}(w_i^ast)},

frac{1}{N_{gamma,epsilon}(w_i^ast)}right)right),$$ where we assume

(hat{L}(w)) is locally convex around (w^ast).

Even though we

assume the local convexity in our metric, in application we may calculate the

metric on any point on the loss surface. When (nabla^2_i hat{L}(w^ast) +

rho(w^ast)sqrt{m}N_{gamma,epsilon}(w_i^ast) < 0), we simply treat it as

(0 ). The metric is plotted as the color wrapped on the surface in the figure

below.

As displayed in the figure, the metric score around the global optimum,

indicated by the green vertical bar, is high, suggesting possible poor

generalization capability as compared to the local optimum indicated by the red

bar.

On the other hand, the overall expected loss is determined by

both the loss and the metric. For that we plotted a color plane on the bottom

of the figure. The color on the projected plane indicates an approximated

generalization bound, which considers both the loss and the generalization

metric. The local optimum indicated by the red bar, though has a slightly

higher loss, has a similar overall bound compared to the “sharp” global

optimum.

## Reconstruction from the Model

Assuming the model is generating

(y=f(x)=p(y|x)), since we know the true distribution (p(x)), we can

reconstruct the join distribution (p(x,y)=p(x)p(y|x)). To be specific, we

first sample (x) from (p(x)), then use the model to predict

(hat{y}=f(x)). The samples sampled from the true distribution, the

distribution induced by the “sharp” minima and the “flat” minima, are shown

below.

The prediction function from the “flat” minima seems to have a simpler dicision

boundary as compared that of the “sharp” minima, even though the “sharp” minima

fits the labels better.

## When PAC-Bayes meets Hessian Lipschitz

In order to get

a quantative evaluation of the local optimum and the model perturbation, we

need to look at the PAC-Bayes bound:

**[PAC-Bayes-Hoeffding
Bound]** Let (l(f, x,y)in [0,1]), and (pi) be any fixed distribution

over the parameters (mathcal{W}). For any (delta>0) and (eta>0),

with probability at least (1-delta) over the draw of (n) samples, for

any (w) and any random perturbation (u), $$mathbb{E}_u [L(w+u)]leq

mathbb{E}_u [hat{L}(w+u)] + frac{KL(w+u||pi) + log

frac{1}{delta}}{eta} + frac{eta}{2n}$$

Suppose (hat{L}(w^ast)) is the loss function at a local optima

(w^ast), when the perturbation level of (u) is small,

(mathbb{E}_u[hat{L}(w^ast+u)]) tends to be small, but

(KL(w^ast+u|pi)) may be large since the posterior is too focused on a

small neighboring area around (w^ast), and vice versa. As a consequence,

we may need to search for an optimal perturbation level for (u) so that

the bound is minimized.

For that we need some assumptions on the smoothness of the Hessian:

**[Hessian Lipschitz]** A twice differentiable function

(f(cdot)) is (rho)-Hessian Lipschitz if:

$$forall w_1, w_2, |nabla^2 f(w_1) – nabla^2 f(w_2)|leq rho

|w_1-w_2|,$$

where (|cdot|) is the operator norm.

Suppose the empirical loss function (hat{L}(w)) satisfies the

Hessian Lipschitz condition in some local neighborhood (Neigh_{gamma,

epsilon}(w)), then the perturbation of the function around a fixed

point can be bounded by terms up to the third-order,

(Nesterov) $$hat{L}(w+u) leq hat{L}(w) + nabla hat{L}(w)^T u +

frac{1}{2}u^T nabla^2 hat{L}(w) u + frac{1}{6}rho|u|^3~~~~forall

u~~s.t.~~ w+uin Neigh_{gamma,epsilon}(w)$$

Combine the PAC-Bayes bound and we get $$mathbb{E}_u [L(w+u)]leq hat{L}(w)

+frac{1}{2} sum_i nabla_{i,i}^2 hat{L}(w)mathbb{E} [u_i^2] +

frac{rho}{6}mathbb{E}[|u|^3] + frac{KL(w+u||pi) + log

frac{1}{delta}}{eta} + frac{eta}{2n}$$

Pluging in different perturbation distributions we can explicitly calculate the

KL divergence and solve for the best level of perturbation for each parameter.

## PACGen on Real Data Sets

The PACGen metric is calculated on a PyTorch

model with various batch size and learning rate. We had similar

observations as Keskar et al. that as the batch size grows, the gap between

the test loss and the training loss tends to get larger. Our proposed metric

(Psi_{gamma, epsilon}(hat{L}, w^ast)) also shows the exact same

trend.

Similarly if we fix the training batch size as (256), as the learning rate

decreases, the gap between the test loss and the training loss increases, which

is also consistent with the trend calculated from (Psi_{gamma,

epsilon}(hat{L}, w^ast)).

## Pertubed Optimization

PAC-Bayes bound suggests we should optimize the perturbed loss instead of the true loss for a better generalization. In particular, the level of perturbation on each parameter is set according to the local smoothness property. We observe improved performance for the perturbed model on CIFAR-10, CIFAR-100, as well as the tiny Imagenet data sets.

## Conclusion and Citation Credit

We connect the smoothness of the solution with the model generalization in the PAC-Bayes framework. We theoretically show that the generalization power of a model is related to the Hessian and the smoothness of the solution, the scales of the parameters, as well as the number of training samples. Based on our generalization bound, we propose a new metric to test the model generalization and a new perturbation algorithm that adjusts the perturbation levels according to the Hessian. Finally, we empirically demonstrate the effect of our algorithm is similar to a regularizer in its ability to attain better performance on unseen data. For more details, including details regarding the proof and its assumptions, please refer to our pre-print.

*Identifying Generalization Properties in Neural Networks *Huan Wang, Nitish Shirish Keskar, Caiming Xiong and Richard Socher.