Link to original article

Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: Hessian and Basin volume, published by Vivek Hebbar on July 10, 2022 on The AI Alignment Forum. Thanks to Thomas Kwa for the question which prompted this post. Note: This is mostly a primer / introductory reference, not a research post. However, the details should be interesting even to those familiar with the area. When discussing “broad basins” in the loss landscape of a DNN, the Hessian of loss is often referenced. This post explains a simple theoretical approximation of basin volume which uses the Hessian of loss. Suppose our minimum has loss=0. Define the basin as the region of parameter space draining to our minimum where loss < threshold T. Simplest model If all eigenvalues of the Hessian are positive and non trivial, we can approximate the loss as a paraboloid centered on our minimum: The vertical axis is loss, and the horizontal plane is parameter space. The shape of the basin in parameter space is the shadow of this paraboloid, which is an ellipsoid. The principal directions of curvature of the paraboloid are given by the eigenvectors of the Hessian. The curvatures (second derivative) in each of those directions is given by the corresponding eigenvalue. Radii of the ellipsoid: If we start at our minimum and walk away in a principal direction, the loss as a function of distance traveled is L(x)=12λix2, where λi is the Hessian eigenvalue for that direction. So given our loss threshold T, we will hit that threshold at a distance of x=√2Tλi. This is the radius of the loss-basin ellipsoid in that direction. The volume of the ellipsoid is Vbasin=Vn∏i√2Tλi, where the constant Vn is the volume of the unit n-ball in n dimensions. Since the product of the eigenvalues is the determinant of the Hessian, we can write this as: So the basin volume is inversely proportional to the square root of the determinant of the Hessian. Everything in the numerator is a constant, so only the determinant of the Hessian matters in this model. The problem with this model is that the determinant of the Hessian is usually zero, due to zero eigenvalues. Fixing the model If we don't include a regularization term in the loss, the basin as we defined it earlier can actually be infinitely big (it's not just a problem with the paraboloid model). However, we don't really care about volume that is so far from the origin that it is never reached. A somewhat principled way to fix the model is to look at volume weighted by the initialization distribution. This is easiest to work with if the initialization is Gaussian. To make the math tractable, we can replace our ellipsoid with a "fuzzy ellipsoid" -- i.e. a multivariate Gaussian. Now we just have to integrate the product of two Gaussians, which should be easy. There are also somewhat principled reasons for using a "fuzzy ellipsoid", which I won't explain here. However, this is only somewhat principled; if you think about it further, it starts to become unclear: Should we use the initialization Gaussian, or one based on the expected final L2 norm? What about cases where the norm peaks in the middle of training, and is smaller at the start and finish? If we have an L2 regularization term in the loss, then the infinite volume problem usually goes away; the L2 term makes all the eigenvalues positive, so the formula is fine. If we have weight decay, we can interpret this as L2 regularization and add it to the loss. For a relatively simple approximation, I recommend the formula: Where: Loss is the unregularized loss λ is the amount of weight decay (or L2 regularization 12λθ2) c=k/σ2, where σ is the standard deviation of the initialization Gaussian, and k is a constant on the order of unity. I have not calculated the theoretically most appropriate value of k. For a crude model, k=1 is probably good enough. T is the loss threshold. If you really care about...