Behaviour Manifolds and the Hessian of the Total Loss—Notes and Criticism

[Spencer Becker-Kahn]

In this note I will discuss some computations and observations that I have seen in other posts about “basin broadness/​flatness”. I am mostly working off the content of the posts Information Loss --> Basin flatness and Basin broadness depends on the size and number of orthogonal features. I will attempt to give one rigorous and unified narrative for core mathematical parts of these posts and I will also attempt to explain my reservations about some aspects of these approaches. This post started out as a series of comments that I had already made on the posts, but I felt it may be worthwhile for me to spell out my position and give my own explanations.

Work completed while author was a SERI MATS scholar under the mentorship of Evan Hubinger.

Basic Notation and Terminology

We will imagine fixing some model architecture and thinking about the loss landscape from a purely mathematical perspective. We will not concern ourselves with the realities of training.

Let $\Theta$ denote the parameter space of a deep neural network model $f$. This means that each element $\theta \in \Theta$ is a complete set of weights and biases for the model. And suppose that when a set of parameters $\theta \in \Theta$ is fixed, the network maps from an input space $R^n$ to an output space $O$. When it matters below, we will take $O=R^k$, but for now let us leave it abstract. So we have a function

$$f:\Theta \times R^n\to O,$$

such that for any $\theta \in \Theta$, the function $f(\cdot ,\theta ):R^n\to O$ is a fixed input-output function implemented by the network.

Let $D=\{(x^d,y^d){\}}_{d=1}^D\subset R^n\times O$ be a dataset of $D$ training examples. We can then define a function $F:\Theta \to O^D$, by

$$F\left(\theta \right)=(f(x^1,\theta ),\dots ,f(x^D,\theta )).$$

This takes as input a set of parameters $\theta$ and returns the behaviour of $f(\cdot ,\theta )$ on the training data.

We will think of the loss function as $l:O^D\to R$.

Example. We could have $D=\{(x^d,y^d){\}}_{d=1}^D$ , $O=R^k$ , and

$$\begin{array}{}l(o^1,\dots ,o^D)=\frac{1}{2}\sum _{d=1}^D\Vert o^d-y^d{\Vert }^2.& \text{(*)}\end{array}$$

We also then define what we will call the total loss

$$L:\Theta \to R$$

by

$$L\left(\theta \right)=l(F\left(\theta \right))=l(f(x^1,\theta ),\dots ,f(x^D,\theta )).$$

This is just the usual thing: The total loss over the training data set for a given set of weights and biases. So the graph of $L$ is what one might call the ‘loss landscape’.

Behaviour Manifolds

By a behaviour manifold (see [Hebbar]), we mean a set $\Sigma \subset \Theta$ of the form

$$\Sigma =F^{-1}((o_1,\dots ,o_D))=\{\theta \in \Theta :F\left(\theta \right)=(o_1,\dots ,o_D)\}$$

where $(o_1,\dots ,o_D)\in O^D$ is a tuple of possible outputs. The idea here is that for a fixed behaviour manifold $\Sigma$, all of the models given by parameter sets $\theta \in \Sigma$ have identical behaviour on the training data.

Assume that $\Theta$ is an appropriately smooth $N$-dimensional space and let us now assume that $O=R^k$.

Suppose that $N>kD$. In this case, at a point $\theta \in \Theta$ at which the Jacobian matrix $JF\left(\theta \right)$ has full rank, the map $F$ is a submersion. The submersion theorem (which—in this context—is little more than the implicit function theorem) tells us that given $o\in O^D$, if $F$ is a submersion in a neighbourhood of a point $\theta \in F^{-1}\left(o\right)$, then $F^{-1}\left(o\right)$ is a smooth $(N-kD)$-dimensional submanifold in a neighbourhood of $\theta$ . So we conclude that in a neighbourhood of a point in parameter space at which the Jacobian of $F$ has full rank, the behaviour manifold is an $(N-kD)$-dimensional smooth submanifold.

Reservations

Firstly, I want to emphasize that when the Jacobian of $F$ does not have full rank, it is generally difficult to make conclusions about the geometry of the level set, i.e. about the set that is called the behaviour manifold in this setting.

Examples. The following simple examples are to emphasize that there is not a straightforward intuitive relationship that says “when the Jacobian has less than full rank, there are fewer directions in parameter space along which the behaviour changes and therefore the behaviour manifold is bigger than $(N-kD)$-dimensional”:

1. Consider $g:R^2\to R$ given by $g(x,y)=x^2+y^2$. We have $\nabla g=(2x,2y)$. This has rank 1 everywhere except the origin: At the point $(0,0)$ it has less than full rank. And at that point, the level set is just a single point, i.e. it is 0-dimensional.

2. Consider $h:R^2\to R$ given by $h(x,y)=x^2.$ We have $\nabla h=(2x,0)$. Again, this has less than full rank at the point $(0,0).$ And at that point, the level set is the entire $y$-axis, i.e. it is 1-dimensional.

3. Consider $j:R^2\to R$ given by $j(x,y)=1.$ We of course have $\nabla j=(0,0)$. This has less than full rank everywhere, and the only non-empty level set is the entire of $R^2$, i.e. 2-dimensional.

Remark. We note further, just for the sake of intuition about these kinds of issues, that the geometry of the level set of a smooth function can in general be very bad: Every closed subset is the zero set of some smooth function, i.e. given any closed set $C\subset R^n$ , there exists a smooth function $g:R^n\to R$ with $C=\{x\in R^n:g(x)=0\}.$ Knowing that a level set is closed is an extremely basic fact and yet without using specific information about the function you are looking at, you cannot conclude anything else.

Secondly, the use of the submersion theorem here only makes sense when $N>kD$. But this is not even commonly the case. It is common to have many more data points (the $D$) than parameters (the $N$), ultimately meaning that the dimension of $O^D$ is much, much larger than the dimension of the domain of $F$. This suggests a slightly different perspective, which I briefly outline next.

Behavioural Space

When the codomain is a higher-dimensional space than the domain, we more commonly picture the image of a function, as opposed to the graph, e.g. if I say to consider a smooth function $g:R\to R^2$, one more naturally pictures the curve $g\left(R\right)$ in the plane, as a kind-of ‘copy’ of the line $R$, as opposed to the graph of $g$. So if one were to try to continue along these lines, one might instead imagine the image $F\left(\Theta \right)$ of parameter space in the behaviour space $O^D.$ We think of each point of $O^D$ as a complete specification of possible outputs on the dataset. Then the image $F\left(\Theta \right)\subset O^D$ is (loosely speaking) an $N-$dimensional submanifold of this space which we should think of as having large codimension. And each point $F\left(\theta \right)$ on this submanifold is the outputs of an actual model with parameters $\theta$. In this setting, the points $\theta \in \Theta$ at which the Jacobian $JF\left(\theta \right)$ has full rank map to points $F\left(\theta \right)\in F\left(\Theta \right)$ which have neighbourhoods in which $F\left(\Theta \right)$ is smooth and embedded.

The Hessian of the Total Loss

A computation of the Hessian of $L$ appears in both Information Loss --> Basin flatness and Basin broadness depends on the size and number of orthogonal features, under slightly different assumptions. Let us carefully go over that computation here, in a slightly greater level of generality. We continue with $O=R^k$, in which case $O^D=R^{k\times D}$. The function we are going to differentiate is:

$$L\left(\theta \right)=l(F\left(\theta \right))=l(f(x^1,\theta ),\dots ,f(x^D,\theta )).$$

And since each $f(x^d,\theta )\in R^k$ for $d=1,\dots ,D$, we should think of $F\left(\theta \right)$ as a $k\times D$ matrix, the general $(p,d{)}^{th}$entry of which is $f^p(x^d,\theta )$.

We want to differentiate twice with respect to $\theta$. Firstly, we have

$$\frac{\partial }{\partial {\theta }_i}L\left(\theta \right)=\sum _{p=1}^k\sum _{d=1}^D{\nabla }^{(p,d)}l(F\left(\theta \right))\cdot \frac{\partial f^p(x^d,\theta )}{\partial {\theta }_i}$$

for $i=1,\dots ,N$.

Then for $j=1,\dots ,N$ we differentiate again:

$$\frac{{\partial }^2}{\partial {\theta }_j\partial {\theta }_i}L\left(\theta \right)=\sum _{p,q=1}^k\sum _{d,d^{\prime }=1}^D{\nabla }^{(q,d^{\prime })}{\nabla }^{(p,d)}l(F\left(\theta \right))\frac{\partial f^q(x^{d^{\prime }},\theta )}{\partial {\theta }_j}\frac{\partial f^p(x^d,\theta )}{\partial {\theta }_i}$$$$\begin{array}{}+\sum _{p=1}^k\sum _{d=1}^D{\nabla }^{(p,d)}l(F\left(\theta \right))\frac{{\partial }^2{f}^p(x^d,\theta )}{\partial {\theta }_j\partial {\theta }_i}.& \text{(1)}\end{array}$$

This is now an equation of $(N\times N)$ matrices.

At A Local Minimum of The Loss Function

If $\theta$ is such that $F\left(\theta \right)$ is a local minimum for $l$ (which means that the parameters are such that the output of the network on the training data is a local minimum for the loss function), then the second term on the right-hand side of (1) vanishes (because the term includes the first derivatives of $l$, which are zero at a minimum). Therefore: If $F\left({\theta }^{\ast }\right)$ is a local minimum for $l$ we have:

$$\frac{{\partial }^2}{\partial {\theta }_j\partial {\theta }_i}L\left({\theta }^{\ast }\right)=\sum _{p,q=1}^k\sum _{d,d^{\prime }=1}^D{\nabla }^{(q,d^{\prime })}{\nabla }^{(p,d)}l(F\left({\theta }^{\ast }\right))\frac{\partial f^q(x^{d^{\prime }},{\theta }^{\ast })}{\partial {\theta }_j}\frac{\partial f^p(x^d,{\theta }^{\ast })}{\partial {\theta }_i}.$$

If, in addition, the Hessian of $l$ is equal to the identity matrix (by which we mean ${\nabla }^{(q,d^{\prime })}{\nabla }^{(p,d)}l={\delta }_{pq}{\delta }_{d{d}^{\prime }}$ - as is the case for the example loss function given above in (*)), then we would have:

$$\frac{{\partial }^2}{\partial {\theta }_j\partial {\theta }_i}L\left({\theta }^{\ast }\right)=\sum _{p=1}^k\sum _{d=1}^D\frac{\partial f^p(x^d,{\theta }^{\ast })}{\partial {\theta }_j}\frac{\partial f^p(x^d,{\theta }^{\ast })}{\partial {\theta }_i}$$$$\begin{array}{}=\sum _{d=1}^D\frac{\partial f(x^d,{\theta }^{\ast })}{\partial {\theta }_i}\cdot \frac{\partial f(x^d,{\theta }^{\ast })}{\partial {\theta }_j}.& \text{(2)}\end{array}$$

Reservations

In Basin broadness depends on the size and number of orthogonal features, the expression on the right-hand side of equation (2) above is referred to as an inner product of “the features over the training data set”. I do not understand the use of the word ‘features’ here and in the remainder of their post. The phrase seems to imply that a function of the form

$$x^d⟼\frac{\partial f(x^d,{\theta }^{\ast })}{\partial {\theta }_j},$$

defined on the inputs of the training dataset, is what constitutes a feature. No further explanation is really given. It’s completely plausible that I have missed something (and perhaps other readers do not or will not share my confusion) but I would like to see an attempt at a clear and detailed explanation of exactly how this notion is supposed to be the same notion of feature that (say) Anthropic use in their interpretability work (as was claimed to me).

Criticism

I’d like to tentatively try to give some higher-level criticism of these kinds of approaches. This is a tricky thing to do, I admit; it’s generally very hard to say that a certain approach is unlikely to yield results, but I will at least try to explain where my skepticism is coming from.

The perspective and the computations that are presented here (which in my opinion are representative of the mathematical parts of the linked posts and of various other unnamed posts) do not use any significant facts about neural networks or their architecture. In particular, in the mathematical framework that is set up, the function $f$ is more or less just any smooth function. And the methods used are just a few lines of calculus and linear algebra applied to abstract smooth functions. If these are the principal ingredients, then I am naturally led to expect that the conclusions will be relatively straightforward facts that will hold for more or less any smooth function $f$.

Such facts may be useful as part of bigger arguments—of course many arguments in mathematics do yield truly significant results using only ‘low-level’ methods—but in my experience one is extremely unlikely to end up with significant results in this way without it ultimately being clear after the fact where the hard work has happened or what the significant original insight was.

So, naively, my expectation at the moment is that in order to arrive at better results about this sort of thing, arguments that start like these ones do must quickly bring to bear substantial mathematical facts about the network, e.g. random initialization, gradient descent, the structure of the network’s layers, activations etc. One has to actually use something. I feel (again, speaking naively) that after achieving more success with a mathematical argument along these lines, one’s hands would look dirtier. In particular, for what it’s worth, I do not expect my suggestion to look at the image of the parameter space in ‘behaviour space’ to lead (by itself) to any further non-trivial progress. (And I say ‘naively’ in the preceding sentences here because I do not claim myself to have produced any significant results of the form I am discussing).

[Vladimir_Nesov]

I worry that using $O^D$ as the space of behaviors misses something important about the intuitive idea of robustness, making any conclusions about $F$ or $O^D$ or behavior manifolds harder to apply. A more natural space (to illustrate my point, not as something helpful for this post) would be $O^{R^n}$, with a metric that cares about how outputs differ on inputs that fall within a particular base distribution $\gamma$, something like $$d(g,h)=E_{X\sim \gamma }|g\left(x\right)-h\left(x\right)|.$$

The issue with $O^D$ is that models in a behavior manifold only need to agree on the training inputs, and always include all models with arbitrarily crazy behaviors at all inputs outside the dataset, even if we are talking about inputs very close to those in the dataset (which is what $\gamma$ above is supposed to prevent). So the behavior manifolds are more like cylinders than balls, ignoring crucial dimensions. Since generalization does work (so learning tends to find very unusual points of them), it’s generally unclear how a behavior manifold as a whole is going to be relevant to what’s actually going on.

[Spencer Becker-Kahn]

I agree that the space $O^D$ may well miss important concepts and perspectives. As I say, it is not my suggestion to look at it, but rather just something that was implicitly being done in another post. The space $O^{R^n}$ may well be a more natural one. (It’s of course the space of functions $R^n\to O$, and so a space in which ‘model space’ naturally sits in some sense. )

[Vivek Hebbar]

The perspective and the computations that are presented here (which in my opinion are representative of the mathematical parts of the linked posts and of various other unnamed posts) do not use any significant facts about neural networks or their architecture.

You’re correct that the written portion of the Information Loss --> Basin flatness post doesn’t use any non-trivial facts about NNs. The purpose of the written portion was to explain some mathematical groundwork, which is then used for the non-trivial claim. (I did not know at the time that there was a standard name “Submersion theorem”. I had also made formal mistakes, which I am glad you pointed out in your comments. The essence was mostly valid though.) The non-trivial claim occurs in the video section of the post, where a sort of degeneracy occuring in ReLU MLPs is examined. I now no longer believe that the precise form of my claim is relevant to practical networks. An approximate form (where low rank is replaced with something similar to low determinant) seems salvageable, though still of dubious value, since I think I have better framings now.

Secondly, the use of the submersion theorem here only makes sense when $N>kD$.

Agreed. I was addressing the overparameterized case, not the underparameterized one. In hindsight, I should have mentioned this at the very beginning of the post—my bad.

(Sorry for the very late response)

[Vivek Hebbar]

All in all, I don’t think my original post held up well. I guess I was excited to pump out the concept quickly, before the dust settled. Maybe this was a mistake? Usually I make the ~opposite error of never getting around to posting things.

[Johannes Treutlein]

I think there should be a space both for in-progress research dumps and for more worked out final research reports on the forum. Maybe it would make sense to have separate categories for them or so.