My aim for this project was to find connections between singular learning theory and the edge-of-stability phenomenon in gradient descent; this was intended to be ambitious, and I picked this question partly because I wanted to learn about singular learning theory. I succeeded in doing this; I now know some basics singular learning theory, and think it’s very cool, and may use it in my research. While I didn’t obtain substantial theoretical results, this project yielded some ideas for potentially promising directions, as well as some maybe-interesting experimental results. This post will:

Singular learning theory

An important feature of deep neural networks is that they are typically non-identifiable. That is, if a network is capable of learning a function $f$, it is typically capable of learning the function in many different ways.

A simple and common example of this is the ReLU activation. When a ReLU activation sits between two linear layers, it’s possible to change the input and output weights without affecting the function expressed:

$$ \begin{align*} y &= w_2 \;\text{ReLU}(w_1 x + b_1)& \\ y &= \frac{w_2}{a} \;\text{ReLU}(aw_1 x + ab_1) \quad &(a>0) \\ y &= \pm \,\text{ReLU}(w'x + b') &\text{(putting $a = |w_2|$)} \end{align*} $$

The ReLU-between-linear-layers has a free parameter—it uses three parameters, but the set of functions it can express corresponds to a two-dimensional manifold. This is true in bigger layers, too—you can scale all the input and output weights for a single ReLU activation by (opposite) constant factors without affecting what the layer does. Similarly, many practical neural networks have back-to-back linear layers, which let you rotate and scale the intermediate space without affecting the function the model expresses. These are non-dataset-dependent examples, but degeneracy for a model is typically dependent on the loss function, and in fact can correspond meaningfully to the model’s ability to learn the dataset.

A solution with free parameters is interesting partly because it looks unlike a quadratic basin. Intuitively, the basin is bigger; in fact, along the directions corresponding to the free parameters, the second derivatives of the loss are zero. There are other ways that a basin can be “bigger” than quadratic— a quartic, for example—and these also correspond to a degenerate Hessian. In some deep learning contexts, the “size” of a basin is measured by the inverse determinant of the Hessian, but this is only a good measure for basins that look locally quadratic.

Instead, we can measure the size of a basin using the volume of the “sub-level set”, $V[\left\{x \mid L(x) < \varepsilon \right\}]$. For a regular quadratic basin, $V \sim \varepsilon^{d/2}$, where $d$ is the number of parameters. It can be shown that $V \sim \varepsilon^{\ell}$ for any singularity of an analytic model, for some positive $\ell$, but $\ell$ can be (in fact, typically is) less than the number of parameters. $\ell$ is called the learning coefficient and loosely can be thought of as being (half of) the “effective dimensionality” of the model.

Knowledge to be discovered [in a statistical model] corresponds to a singularity.

Sumio Watanabe [watanabe09]

Singular learning theory, a line of research pioneered by Sumio Watanabe and his students, posits that singularities are important for understanding learning. Results include the free energy formula, which describes a tradeoff between loss and effective dimensionality. SLT grew up in the context of Bayesian learning; there is currently work ongoing to port it to the context of gradient descent, where it may be useful for interpretability.

This sequence of blog posts [carroll23] and this set of slides [pepin24] are good introductions.

DSLT 0. Distilling Singular Learning Theory — AI Alignment Forum

simon-pepin.github.io

The edge of stability