The Hessian rank bounds the learning coefficient

post by Lucius Bushnaq (Lblack) · 2024-08-08T20:55:36.960Z · LW · GW · 9 comments

Contents

  Introduction
    Some words on notation
  Theorem:
  Proof:
  Future work
    Acknowledgments
None
9 comments

TL;DR: In a neural network with $d$ parameters, the (local) learning coefficient $\lambda$ can be upper and lower bounded by the rank of the network's Hessian $d_1$:

$\frac{d_1}{2}\le \lambda \le \frac{d_1}{2}+\frac{d-d_1}{4}$.

The lower bound is a known result. The upper bound is a claim by me, and this post contains the proof for it.^[1] If you find any problems, do point them out. 

Edit 16.08.2024: The original version of this post had a three in the denominator of the upper bound. Dmitry Vaintrob [LW · GW] spotted an improvement to make it a four.

Introduction

The learning coefficient $\lambda$ is a measure of loss basin volume and model complexity. You can think of it sort of like an effective parameter count of the neural network. Simpler models that do less stuff have smaller $\lambda$.

Calculating $\lambda$ for real networks people actually use is a pain. My hope is that these bounds help make estimating it a bit easier.

In a network with $d$ parameters, the learning coefficient is always a number

$0\le \lambda \le \frac{d}{2}$.

An existing result in the literature says that if you’ve calculated the rank of the network’s Hessian $d_1$,^[2] you get a tighter lower bound

$\frac{d_1}{2}\le \lambda .$

I claim that we can also get a tighter upper bound

$\lambda \le \frac{d_1}{2}+\frac{d-d_1}{4}$,

where $d-d_1$ will be the dimension of the Hessian kernel, meaning the number of zero eigenvalues it has.^[3]

This is neat because it means we can get some idea of how large $\lambda$ is using only standard linear algebra methods. All we need to know is how many zero eigenvalues the Hessian has.^[4] Singular Learning Theory introductions often stress that just looking at the Hessian isn’t enough to measure basin volume correctly. But here we see that if you do it right, the Hessian eigenspectrum can give you a pretty good idea of how large $\lambda$ is. Especially if there aren't that many zero eigenvalues.

Intuitively, the lower bound works because a direction  in the parameters $w$ that isn't free to vary to second order in the Taylor expansion won't become any more free to vary if you pile on a bunch of higher order terms. The Second order term strictly dominates the higher order ones, they can't cancel it out.

Qualitatively speaking, the upper bound works for the same reason. The higher order terms in the Taylor expansion of the loss can only matter so much. The Hessian is the leading term, so it can impact $\lambda$ the most, adding $\frac{1}{2}$ per Hessian rank to it. The remaining terms of order $O\left(w^3\right)$ and above can only add a maximum of $\frac{1}{4}$ for the remaining directions. 

The proof for the upper bound will just be a small modification of the proof for theorem 7.2 on pages 220 and 221 of Algebraic Geometry and Statistical Learning Theory. Maybe read that first if you want more technical context.

Some words on notation

In the following, I’ll mostly stick to the notation and conventions of the book Algebraic Geometry and Statistical Learning Theory. You can read about all the definitions there. I'm too lazy to reproduce them all. 

To give some very rough context, $K\left(w\right)$ is sort of like the 'loss' at parameter configuration $w$, $\phi \left(w\right)$ is our prior over parameters, and $Z\left(n\right)$ is the partition function after updating on $n$ data points.^[5]

Theorem:

Let $W\subset R^d$ be the set of parameters of the model. If there exists an open set $U\subset W$ such that

$\{w\in U:K(w)=0,\phi (w)>0\}$

is not an empty set, and we define $d_1=$ rank$\left(H\right)$ as the rank of the Hessian $H$ at a $w_0\in U$

$H_{i,j}=\frac{{\partial }^{2}K\left(w\right)}{\partial w_i\partial w_j}{|}_{w=w_0}$

with $w_i,w_j$ elements of some orthonormal basis $\{w_1,\dots w_d\}$ of $R^d$, then

$\lambda \le \frac{d_1}{2}+\frac{d-d_1}{4}$.

Proof:

We can assume $w_0=0$ without loss of generality. If ${ϵ}_1,{ϵ}_2$ are sufficiently small constants,

$Z\left(n\right)=\int \exp (-nK(w\left)\right)\phi \left(w\right)dw\ge \underset{\left|w^{\left(1\right)}\right|\le {ϵ}_1,\left|w^{\left(2\right)}\right|\le {ϵ}_2}{\int }\exp \{-nK(w\left)\right\}\phi \left(w\right)dw.$

Here,  $w^{\left(1\right)}\in W/\text{ker}\left(H\right),w^{\left(2\right)}\in \text{ker}\left(H\right)$.

If we pick $\{w_1,\dots w_d\}$ to be the Hessian eigenbasis, then for sufficiently small $\left|w\right|>0$

$K\left(w\right)=\frac{1}{2}\sum _{i,i=1}^{d_1}{H}_{i,i}{w}_i^{\left(1\right)}{w}_i^{\left(1\right)}+O\left(\right|w^{\left(1\right)}{|}^2\left|w^{\left(2\right)}\right|)+O(\left|w^{\left(2\right)}{|}^2\right|w^{\left(1\right)}\left|\right)+O\left(\right|w{|}^4)$.

There is no $O\left(\right|w^{\left(2\right)}{|}^3)$ term because we are at a local optimum, so no direction can have a leading term that is an odd power.

For sufficiently large $n$, transforming $Z\left(n\right)$ by ${w^{\prime }}^{\left(1\right)}=n^{\frac{1}{2}}{w}^{\left(1\right)},{w^{\prime }}^{\left(2\right)}=n^{\frac{1}{4}}{w}^{\left(2\right)}$ yields

$Z\left(n\right)\ge n^{-\frac{d_1}{2}}{n}^{-\frac{d-d_1}{4}}\underset{\left|{w^{\prime }}^{\left(1\right)}\right|\le 1,\left|{w^{\prime }}^{\left(2\right)}\right|\le 1}{\int }\exp \{-nK({w^{\prime }}^{\left(1\right)},{w^{\prime }}^{\left(2\right)}\}\phi (n^{-\frac{1}{2}}{w^{\prime }}^{\left(1\right)},n^{-\frac{1}{4}}{w^{\prime }}^{\left(2\right)})d{w}^{\prime }.$

Rearranging and substituting gives

$Z\left(n\right)n^{\frac{d_1}{2}+\frac{d-d_1}{3}}\ge \underset{\left|w^{\prime }\right|\le 1}{\int }\exp \{-\frac{1}{2}{\sum }_{i=1}^{d_1}{H}_{i,i}{w^{\prime }}_i^{\left(1\right)}{w^{\prime }}_i^{\left(1\right)}+n^{-\frac{1}{4}}O(\left|{w^{\prime }}^{\left(1\right)}{|}^2\right|{w^{\prime }}^{\left(2\right)}\left|\right)+O\left(\right|{w^{\prime }}^{\left(2\right)}{|}^2\left|{w^{\prime }}^{\left(1\right)}\right|)+O(\left|w^{\prime }{|}^4\right)\}\phi (n^{-\frac{1}{2}}{w}^{\prime \left(1\right)},n^{-\frac{1}{4}}{w}^{\prime \left(2\right)})d{w}^{\prime }.$

The right-hand side will converge to a positive constant $c>0$ as $n\to \infty$:

$\underset{\left|w^{\prime }\right|\le 1}{\int }\exp \{-\frac{1}{2}{\sum }_i^{d_1}{H}_{i,i}{w}_i^{\prime \left(1\right)}{w^{\prime }}_i^{\left(1\right)}+O(\left|{w^{\prime }}^{\left(2\right)}{|}^2\right|{w^{\prime }}^{\left(1\right)}\left|\right)+O\left(\right|{w^{\prime }}^{\left(2\right)}{|}^4\left)\right\}\phi \left(0\right)d{w}^{\prime }=c$ .

Thus, we have 

$\underset{n\to \infty }{lim}Z\left(n\right)n^{\frac{d_1}{2}+\frac{d-d_1}{4}}\ge c$   .                   (1)

 

By Theorem 7.1(1), page 218 of Algebraic Geometry and Statistical Learning Theory,^[6] we also know 

$\underset{n\to \infty }{lim}Z\left(n\right)\frac{n^{\lambda }}{{\left(\log n\right)}^{m-1}}=c^{\ast }$  ,                  (2)

where $c^{\ast }>0$ is some positive constant and $m$ is the order of the largest pole of 

$\zeta \left(z\right)=\int K\left(w{)}^z\phi \right(w)dw$.

Putting (1) and (2) together, we get

$\lambda \le \frac{d_1}{2}+\frac{d-d_1}{4}$.

Future work

I think it is possible to prove progressively tighter upper and lower bounds for the learning coefficient $\lambda$ using information from progressively higher-order terms in the Taylor expansion of the loss.

For example, if the third and fourth-order terms involving $w^{\left(2\right)}$  happened to all be zero as well for a particular model, we could've transformed $w^{\prime \left(2\right)}=n^{\frac{1}{6}}{w}^{\left(2\right)}$ in the proof instead, and obtained the even tighter upper bound $\lambda \le \frac{d_1}{2}+\frac{d-d_1}{6}$.

I have some early results for an upper bound formula like this, though right now it only works under specific conditions^[7].

Since the higher-dimensional tensors in these $O\left(\right|w{|}^n)$ terms can be very computationally intensive to deal with, the real challenge in making these progressively stricter bounds for $\lambda$ practically useful would be finding tricks to make handling those tensors cheaper. A SPAR team I'm mentoring is currently trying out some ideas on that front on real networks.

Acknowledgments

Credit to Dmitry Vaintrob for spotting the improvement that tightened the upper bound from $\frac{d_1}{2}+\frac{d-d_1}{3}$ in the original version of this post to $\frac{d_1}{2}+\frac{d-d_1}{4}$. Thanks to Daniel Murfet for discussion and feedback on the proof idea.

 

1. ^{^}

   I couldn’t find the upper bound in the literature and Daniel Murfet [LW · GW] hadn’t heard of it either, which is why I wrote this. But it could be around somewhere and I just missed it.

2. ^{^}

   For the learning coefficient, this will be the smallest Hessian rank out of all points in the set of minimum loss. For the local learning coefficient, it'll be the smallest Hessian rank out of all minimum loss points in the local neighbourhood we specify.

3. ^{^}

   This will hold for the Hessian rank at any point with minimum loss, though the upper bound from the point with the smallest Hessian rank will of course be the tightest.

4. ^{^}

   Or, in practice, how many Hessian eigenvalues are numerically small enough to not matter at the noise scale we choose.

5. ^{^}

   I.e. the integral of our Bayesian posterior over the parameters after updating on $n$ data points.

6. ^{^}

   Careful, at least my version of the book has a typo here. I think the fraction should be the other way around from the way it's written in the book.

7. ^{^}

   Only for MSE loss, and only if we are at a global optimum where we match the output labels exactly. 

9 comments

Comments sorted by top scores.