G-TRACER: Expected Sharpness Optimization

The connection between generalization performance and the loss-surface geometry of deep-learning architectures in the neighborhood of local minima has long been the subject of interest and speculation, dating back to the MDL-based arguments of (Hinton & van Camp, 1993) and (Hochreiter & Schmidhuber, 1997). The connection is an intuitively appealing one, in that the sharp local minima of the highly nonlinear, non-convex optimization problems associated with modern overparameterized deep learning architectures are more likely to be brittle and sensitive to perturbations in the parameters and training data, and thus lead to worse performance on unseen data. We can build some intuition for this from a probabilistic modelling perspective, given a dataset $\mathcal{D}=\{(x_{i},y_{i})_{i=1}^{n}\}$ consisting of $n$ independent input random variables $x_{i}$ with distribution $p(x)$ and corresponding targets (or labels) $y_{i}$ with distribution $p(y|x)$ and treating the parameters $w\in\Theta\subseteq\mathbb{R}^{p}$ of a deep neural network (DNN) $f(\cdot,w):\mathbb{R}^{d_{x}}\rightarrow\mathbb{R}^{d_{y}}$ as a random variable. Given a loss function $l(y_{i},f(x_{i},w))$ our goal is to find a $w^{*}$ that minimizes the expected loss: $\mathbb{E}_{p(x,y)}[l(y,f(x,w))]$. Writing the finite-sample version of this expected loss as $L(w)=\sum_{i=1}^{n}l(y_{i},f(x_{i},w))$, we can form a generalized posterior distribution (Bissiri et al., 2016) $p(w|\mathcal{D})=p(w)\frac{1}{Z}\exp\{-L(w))\}$ (with normalizer $Z$) over the weights, which coincides with the Bayesian posterior in the special case that the loss is the negative log-likelihood $L(w)=-\frac{1}{n}\sum_{i=1}^{n}\log p(y_{i}|x_{i},w)$ and then, together with an output (conditional predictive) probability distribution $p(y|x,w)$, we can form the predictive distribution by marginalization:

At a local maximum (or mode) $w_{k}$ of $p(w|\mathcal{D}$) we can can form the Laplace approximation (valid asymptotically, for large $n$):

where (assuming, for simplicity, a flat prior) $H=\nabla_{w}^{2}L|_{w=w_{k}}$ and the normalizer (which for the negative log-likelihood loss is the evidence, or marginal likelihood, and which we will also refer to as the pseudo-marginal likelihood) is given by $Z_{k}=p(w_{k}|\mathcal{D})(2\pi)^{\frac{d}{2}}\det(H)^{-\frac{1}{2}}$. Thus, in the neighborhood of each local maximum of $p(w|\mathcal{D})$, we approximate the posterior by a multivariate Gaussian with covariance given by the inverse Hessian of the negative loss: $p(w_{k}|\mathcal{D})\sim\mathcal{N}(w_{k},H^{-1})$. Modern DNNs are characterized by multimodal losses (Wilson & Izmailov, 2020), and so, informally, we can decompose the posterior predictive distribution into disjoint contributions from each of the local maxima, the sum of which dominate the overall integral:

where $Z=\sum_{k^{\prime}\in\{k\}}Z_{k^{\prime}}$, which is an expectation with respect to a probability measure with density given by: $\frac{1}{Z}\sum_{k^{\prime}\in\{k\}}Z_{k^{\prime}}p_{k^{\prime}}(w|\mathcal{D})$, and which, by writing:

can be viewed as a Gaussian mixture model (GMM) with mixing coefficients:

Thus the relative contribution of each component is given by the relative size of the pseudo-marginal likelihoods $Z_{k}$. For very high-dimensional $w\in\mathbb{R}^{p}$ (modern DNN architectures often have billions or even trillions of parameters) even small differences in the width of the Gaussian approximation will have exponentially large effects on the magnitude of $Z_{k}$ (which can be thought of as the the volume associated with the local maximum). How does all this relate to flatness? The Gaussian curvature $K$, providing an intrinsic (and thus coordinate-free) measure of curvature is given by:

and contributions to the mixture thus scale inversely with $\sqrt{K}$. In other words, the flatter the solution, the more it contributes to the mixture model against which the output probability distribution is integrated, in order to form the posterior predictive distribution. In a typical high-dimensional setting, the effect of small differences in curvature (or flatness) is exponentially magnified. To see this, we can consider two local minima i) with Hessian $H$, and ii) an $\epsilon$-flattened minimum ($0<\epsilon\ll 1$) with Hessian $H^{\prime}=(1-\epsilon)H$, (so that each eigenvalue of $H$ is simply shrunk by a constant factor $(1-\epsilon)$). We have $K^{\prime}=\det(H^{\prime})=(1-\epsilon)^{p}\det(H)$, so that the $\epsilon$-flattened minimum with curvature $K^{\prime}$ has exponentially lower Gaussian curvature. The corresponding Gaussian approximations have covariances $\Sigma^{\prime}\approx(1+\epsilon)\Sigma$ and $\Sigma$, and the ratio of the corresponding pseudo-marginal likelihoods scales as:

Thus we can see that the contribution from the flatter minimum dominates in the high-dimensional limit. In an empirical study, Huang et al. (2019) train a ResNet18 architecture on the Street View House Number (SVHN) dataset and estimate the volume around local minima using Monte-Carlo integration, finding that the volumes of basins surrounding minima that generalize well are at least 10,000 orders of magnitude larger than those of minima that generalize poorly.

A complementary approach is to characterize the loss-surface Hessian, since, at such local minimum of the loss, for a perturbation $\Delta w$, we have:

There has therefore been a large literature attempting to characterize the loss-surface Hessian $\nabla^{2}L(w)$ and to relate these characteristics to generalization. In many practically relevant cases, multiple minima are associated with zero (or close to zero) training error, and explicit or implicit regularization is needed to find solutions with the best generalization error. Overparameterization is associated with the bulk of the Hessian spectrum lying close to zero and thus to highly degenerate minima (Sagun et al., 2017). Wei & Schwab (2020) further show that given a degenerate valley in the loss surface, SGD on average decreases the trace of the Hessian, which is strongly suggestive of a connection between locally flat minima, overaparameterization and generalization.

Despite the intuitive appeal and plausible justifications for flat solutions to be a goal of DNN optimization algorithms, there have been few practical unqualified successes in exploiting this connection to improve generalization performance. A notable exception is a recent algorithm, Sharpness Aware Minimization (SAM) (Foret et al., 2020), which seeks to improve generalization by optimizing a saddle-point problem of the form:

An approximate solution to this problem is obtained by differentiating through the inner maximization, so that, given an approximate solution $\Delta w^{*}:=\rho\frac{\nabla L(w^{k})}{\|\nabla L(w^{k})\|_{2}}$ to the inner maximization (dual norm) problem:

the gradient of the SAM objective is approximated as follows:

While the method has gained widespread attention, and state-of-the-art performance has been demonstrated on several benchmark datasets, it remains relatively poorly understood, and the motivation and connection to sharpness is questionable given that the Euclidian norm-ball isn’t invariant to changes in coordinates. Given a 1-1 mapping $g:\Theta^{\prime}\rightarrow\Theta$ we can reparameterize our DNN $f(\cdot,w)$ using the "pullback" $g^{*}(f)(\cdot,\nu):=f(\cdot,g(\nu))$ under which, crucially, the underlying prediction function $f(\cdot,w):\mathbb{R}^{d_{x}}\rightarrow\mathbb{R}^{d_{y}}$ (and therefore the loss) itself is invariant, since, for $\nu=g^{-1}(w)$, we have $f(\cdot,w)=f(\cdot,g(\nu))$. Under this coordinate transformation, however, the Hessian at a critical point transforms as (Dinh et al., 2017):

In particular, Dinh et al. (2017) explicitly show, using layer-wise transformations $T_{\alpha}:(w_{1},w_{2})\rightarrow(\alpha w_{1},\alpha^{-1}w_{2})$, that deep rectifier feedforward networks possess large numbers of symmetries which can be exploited to control sharpness without changing the network output. The existence of these symmetries in the loss function, under which the geometry of the local loss can be substantially modified (and in particular, the spectral norm and trace of the Hessian) means that the relationship between the local flatness of the loss landscape and generalization is a subtle one.

This bound justifies and motivates the SAM objective:

and resulting algorithm. While the bound in Theorem 2 suggests that the ridge-penalty should vary with the radius of the perturbation, in practice (Foret et al., 2020) the penalty term is fixed (or simply set to zero) even when different perturbation radii are searched over. Subsequent refinements of SAM (Kim et al., 2022) ignore the ridge penalty term altogether, and the choice of an optimal perturbation radius is what drives the success of the method. It is not clear, however, why this adversarial parameter-space perturbation should help generalization more than evaluating (and approximating) the expectation in the very bound which motivates the SAM procedure in the first place, which would lead instead to an objective (ignoring, for now, the ridge penalty term) of the following form:

Moreover, the worst-case adversarial perturbation used by SAM is likely to be noisier and is also naturally a significantly looser bound than the expectation-based bound.

Ignoring for simplicity the contribution from the prior term (which would correspond to a ridge-regularization term under the assumption $p(w)\sim N(0,\sigma_{p}I)$), leads to following objective, which we seek to minimize over $w$:

where $\mathcal{H}(q)=-\mathbb{E}_{q(w)}[q(w)]$ is the entropy of $q$. For the choice $q(w)\sim N(w,\sigma^{2}I)$, the optimization problem associated with the variational objective becomes (absorbing some constants into $\rho$):

so that we can see that $\rho$ determines the variance of Gaussian perturbation over which the loss is averaged. More generally, choosing $q\sim N(w,\Sigma)$ leads to the following variational objective:

so that large values of $\rho$ will correspond to distributions with larger volume, since for $x\sim N(0,\Sigma)$, $x$ lies within the ellipsoid $x^{T}\Sigma^{-1}x=\chi^{2}(\alpha)$ with probability $1-\alpha$, with the volume of the ellipsoid proportional to $\det(\Sigma)^{\frac{1}{2}}$ (Anderson, 2003). We show in section 2.3 that expanding the expectation under $q$ to second order, we have:

so that, in the neighborhood of a local minimum, where the Hessian is positive-definite, the curvature of the loss-surface is penalized over a region whose volume is determined by $\rho$. While intuitively appealing, this flatness inducing penalty is not invariant to coordinate transformations, so that scale changes (such as occur, for example, when applying batch-normalization or weight-normalization) which have no effect on the output of the learned probability distribution, can nevertheless still result in arbitrary changes to the penalty. More generally, as discussed above, any geometric notion of loss surface flatness must be independent of arbitrary rescaling of the network parameters. Motivated by these considerations, we apply steepest descent in the KL-metric (also known as natural gradient descent (Amari, 1998)) to our variational objective, under the assumption $q(w)\sim N(w,\Sigma)$ in order to find solutions to:

where $\rho$ is a positive real-valued regularization parameter.

We show in the sequel that, assuming an isotropic Gaussian prior, $p(w)\sim N(0,\eta I)$, performing gradient descent w.r.t. the natural gradient then leads to the following iterative update equations:

where $\alpha_{t}$ and $\beta$ are the learning rates for the mean and precision updates, respectively, and $\Lambda:=\Sigma^{-1}$ is the precision matrix. Approximating the expectations to second order and further simplifying leads to the following update equations (see below for a detailed derivation):

where $H=\nabla^{2}_{w}L(w)$ is the Hessian. The update rule for $\overline{H}$ is an exponential smoothing and the update for the mean consists of a preconditioned (by the inverse smoothed Hessian) gradient, together with, crucially, a penalty term proportional to the (affine-invariant) ratio of the Hessian and the smoothed Hessian. Finally, via an empirical Fisher (diagonal) Hessian approximation (see below for details and a discussion of alternatives) and dropping the preconditioner, we arrive at a modified SGD-type update which we call SGD-TRACER.

Following Khan & Rue and Zhang et al., we make the assumption $q(w)\sim N(\mu,\Sigma)$ and seek to optimize the variational objective in Equation (22) w.r.t. the variational parameters $\phi=(\mu,\Sigma)$ using natural gradient descent, which allows us to derive an algorithm that respects the intrinsic geometry of the parameter space, and thus derive an algorithm that seeks sharp minima in an approximately coordinate-independent way.

Assuming an isotropic Gaussian prior, $p(w)\sim N(0,\eta I)$, performing gradient descent w.r.t. this natural gradient then leads to the following iterative update equations:

where $\alpha_{t}$ and $\beta$ are the learning rates for the mean and precision updates, respectively. We work with each of these update equations in turn. Starting with the update equation for the mean $\mu$, the key observation is that the expectation $\mathbb{E}_{q}[\nabla_{w}L(w)]$ is taken with respect to the distribution $q(w)$ which is an exponential moving average of the expected Hessian $\mathbb{E}_{q}[\nabla_{w}^{2}L(w)]$. This updating happens naturally as a consequence taking natural gradient steps, and is what leads to an approximately coordinate free algorithm in the sequel. Applying Bonnet’s theorem Khan & Rue (2021) and forming the second-order approximation to the loss:

Writing $w-\mu=\Sigma^{\frac{1}{2}},\nu$ where $\nu\sim\mathcal{N}(0,I)$ we have:

where we used the fact that $\mathbb{E}_{\nu\sim\mathcal{N}(0,I)}[{\nu^{T}\Omega\nu]}=\sum_{i,j}\Omega_{i,j}\mathbb{E}[\nu_{i}\nu_{j}]=\text{Tr}(\Omega)$, and where $H$ is the Hessian $\nabla^{2}_{w}L(w)$. We therefore have have that:

Choosing the prior variance $\eta$ to be infinite and thus ignoring terms involving $\eta$, corresponding to an improper prior (and consistent with the discussion above), leads to the following update for the mean:

Thus, in order to blur the loss with multivariate Gaussian noise in a way that aligns with the intrinsic geometry of the parameter space, we can (to second order) augment the loss with a term involving the Trace of the Hessian. Taking now the update equation for the covariance, we can simplify by Price’s theorem (Khan & Rue, 2021) together with a Taylor expansion, to get, to second order: (see Appendix for details) $\mathbb{E}_{q}[\nabla_{w}^{2}L(w)]\approx\nabla_{w}^{2}L(w)|_{w=\mu}$:

We next substitute, as is common in the literature on approximate second order approximation Martens (2020), the Generalized Gauss-Newton matrix (GGN) for the Hessian, given by:

where $J_{f}$ is the Jacobian of the output function $f$ and $H_{L}$ is the Hessian of the loss w.r.t. the output distribution. The GGN is a positive definite approximation to the Hessian which converges to the Hessian as the fitted residuals go to zero, (Kunstner et al., 2019). The most practically relevant losses, cross-entropy (classification), and squared error (regression) correspond to exponential family output distributions with natural parameters given by $f(x,w)$, together with for the log-loss $l(y,f(x,w))=-\log(p(y|x,w)$, and for these choices, the GGN is equivalent to the Fisher Information Matrix. While the evaluation of the GGN matrix, in particular the matrix multiplies involving the Jacobians $J_{f}$, can be relatively costly, the FIM can be expressed as an expectation of outer products of gradients w.r.t. the output distribution $p(y|x,w)$:

which, following Martens (2020), can be estimated using a single Monte Carlo sample from the output distribution: $\tilde{y}\sim p(y|x_{i},w)$. Using this biased Fisher approximation in our setting thus requires gradients to be calculated through an expectation $\nabla_{w}\mathbb{E}_{p(y|x,w)}[L(w;y)]$ approximated using a Monte Carlo sample from the model’s output distribution. Since the expectation is taken w.r.t. a distribution which depends on $w$, it is necessary to reparameterize so that the discrete Monte Carlo sample is expressed as the deterministic transformation of a $g_{w}(z)$ (depending on $w$) of a sample $z\sim h_{\theta}(z)$ from a distribution not depending on $w$, so that $\mathbb{E}_{p(y|x,w)}[L(w;y)]=\mathbb{E}_{z\sim h_{\theta}(z)}[L(w;g_{w}(z)]$. In the discrete case (corresponding to classification), since the argmax function is non-differentiable, the standard approach is the Gumbel-Softmax reparameterization Jang et al. (2016), which uses the softmax function as a continuous relaxation of the argmax function together with i.i.d. samples distributed as Gumbel(0,1).

It’s important to note that is different from simply evaluating $\log p(y|x,w)$ on the training labels, a widely-used approximation known as the empirical Fisher $F_{\text{emp}}$:

which, despite lacking the same convergence guarantees, performs competitively in many settings (Kunstner et al., 2019). We find empirically in our experiments that the empirical Fisher performs competitively with the MC approximation to the GGN (Khan et al., 2018; Kingma & Ba, 2014) and has the advantage of being straightforward and cheap to compute from the already computed gradients (in the case of Adam-TRACER, the smooth squared gradients are already computed and maintained for use as a preconditioner). Given the conceptual and computational simplicity of this approach, and despite its known suboptimality, in the following, we substitute the empirical Fisher for the Hessian. Recent advances in approximate second-order methods in optimization, notably Yao et al. (2020), suggest avenues for improvement, and we leave investigations of alternatives, such as the smoothed (Hessian-free) Hessian diagonal sketch used in AdaHessian, for future work.

Substituting the empirical Fisher approximation for the Hessian in the update equation for the precision, we have the following update:

Rewriting the update equations in terms of this exponentially smoothed FIM $\overline{F}$, absorbing a factor $\rho$ in to $\alpha_{t}$, and writing the iteration in terms of the parameter $w$, we obtain:

Crucially, the penalty term $\rho\text{Tr}(F\overline{F}^{-1})]$ can be seen to be invariant to affine coordinate transformations, since it is the trace of the ratio of two (0,2) tensors which transform in the same way. Indeed under an affine coordinate transformation with Jacobian $J$ we have $F\rightarrow J^{T}F^{{}^{\prime}}J$ and $\overline{F}\rightarrow J^{T}\overline{F}^{{}^{\prime}}J$ so that:

By penalizing the ratio of the (squared) gradients and the exponentially smoothed gradients, the trace ratio penalty in effect is penalizing the change in (squared) gradient, in a coordinate-free way. More generally, given a coordinate change given by a diffemorphism $\Phi:\mathbb{R}^{p}\rightarrow\mathbb{R}^{p}$ and with Jacobian $J(w)$, then given the exponential decay in the update equation for the Fisher, subject to $\Phi$ having sufficient regularity, and for sufficiently small $\beta$, the penalty term is approximately coordinate free under general smooth diffeomorphisms $\Phi$ (see SM for details).

We now make two simplifications. First, we use a mean-field approximation of the FIM by its diagonal, as is done in Adam (Kingma & Ba, 2014) and Adagrad (Duchi et al., 2011), thus:

Secondly, it is standard practice to add Tikhonov regularization or damping by a small positive real constant $\delta$ when using 2nd-order optimization methods, giving in this case the preconditioner: $(\overline{F}+\delta I)^{-1}$. This is justified by recognizing that the local quadratic model from which the second-order update is ultimately derived is a second-order approximation to the KL divergence and is thus only valid locally. For directions corresponding to small eigenvalues, parameter updates can lie outside the region where the approximation is reasonable (Martens, 2020). This is true, a fortiori, when diagonal approximations are used, as is the case here. As our emphasis here is on geometric regularization, we drop the preconditioner entirely by choosing $\delta$ to be sufficiently large that the preconditioner is equal to the identity (up to a constant, which is absorbed into the learning rate).

Finally, as most current deep learning frameworks don’t straightforwardly support access to per-example gradients, which can in principle be achieved with negilible additional cost (see, for example, BackPACK Dangel et al. (2020) second-order Pytorch extensions), for simplicity and efficiency, we use the gradient magnitude (GM) approximation (Bottou et al., 2016), as used in standard optimizers Adam and RMSprop, replacing the sum of squared gradients with the square of summed gradients:

Writing the resulting FIM diagonal as $(\nabla_{w}L(w))^{2}$, we finally end up with the following simple update equations:

which are summarized in Algorithm 1. We show in appendix A.3 that the algorithm converges to a neighborhood of a local minimum of $L(w)$ of size $\mathcal{O}(\rho^{2})$. We note in passing that, in this simplest form (after applying the gradient magnitude approximation), the update equations amount to regularizing with a (scale-adjusted) gradient norm. In principle (particularly for the large batch case) we would expect to see significant improvements by moving to per-gradient calculations (which are in principle no more expensive but require additional work in most current ML frameworks).