# Why I stopped being into basin broadness

There was a period where everyone was really into basin broadness for measuring neural network generalization. This mostly stopped being fashionable, but I'm not sure if there's enough written up on why it didn't do much, so I thought I should give my take for why I stopped finding it attractive. This is probably a repetition of what others have found, but I thought I might as well repeat it.

Let's say we have a neural network $f_w(x) : \mathbb{R}^n$. We evaluate it on a dataset $(x, y) \sim D$ using a loss function $L(\hat y, y) : \mathbb{R}$, to find an optimum $w^* = \arg \min_w E_{(x, y)\sim D}[L(f_w(x), y)]$. Then there was an idea going around that the Hessian matrix (i.e. the second derivative of $E_{(x, y)\sim D}[L(f_w(x), y)]$ at $w^*$) would tell us something about $w^*$ (especially about how well it generalizes).

If we number the dataset $(x_i, y_i)$, we can stack all the network outputs $\hat y_i(w) = f_w(x_i)$ which fits into an empirical loss $\hat L(\hat y) = \frac{1}{n} \sum_{i=1}^{n} L(\hat y_i, y_i)$. The Hessian that we talked about before is now just the Hessian of $\hat L(\hat y(w))$. Expanding this out is kind of clunky since it involves some convoluted tensors that I don't know any syntax for, but clearly it consists of two terms:

*   The Hessian of $\hat L$ with a pair of the Jacobian of $\hat y$ on each end (this can just barely be written without crazy tensors: $( J_w \hat y(w))^T (H_{\hat y} \hat L(\hat y)) \mid_{\hat y(w)} J_w \hat y(w)$)
*   The gradient of $\hat L$ with a crazy second derivative of $\hat y$.

Now, the derivatives of $\hat L$ are "obviously boring" because they don't really refer to the neural network weights, which is confirmed if you think about it in concrete cases, e.g. if $L(\hat y, y) = -y \log(\hat y)-(1-y) \log(1-\hat y)$ with $y=1$ or $y=0$, the derivatives just quantify how far $\hat y$ is from $y$. This obviously isn't relevant for neural network generalization, except in the sense that it tells you which direction you *want* to generalize in.

Meanwhile, $J_w \hat y(w)$ is incredibly strongly related to neural network generalization, because it's literally a matrix which specifies how the neural network outputs change in response the weights. In fact, it forms the core of the neural tangent kernel (a standard tool for modelling neural network generalization), because the NTK can be expressed as $J_w \hat y(w) (J_w \hat y(w))^T$.

The "crazy second derivative of $\hat y$" can I guess be understood separately for each $\hat y_i$, as then it's just the Hessian $H_w \hat y_i(w)$, i.e. it reflects how changes in the weights interact with each other when influencing $\hat y_i$. I don't have any strong opinions on how important this matrix is, though because $J_w \hat y(w)$ is so obviously important, I haven't felt like granting $H_w \hat y_i(w)$ much attention.

The NTK as the network activations?
-----------------------------------

*Epistemic status: speculative, I really should get around to verifying it. Really the prior part is speculative too, but I think those speculations are more theoretically well-grounded. But if I'm wrong with either, please call me a dummy in the comments so I can correct.*

Let's take the simplest case of a linear network, $f_w(x) = w^Tx$. In this case, $J_w \hat y(w) = x^T$, i.e. the Jacobian is literally just the inputs to the network. If you work out a bunch of other toy examples, the takeaway is qualitatively similar (the Jacobian is closely related to the neuron activations), though not exactly the same.

There are of course some exceptions, e.g. $f_{a, b}(x) = abx$ at $a=b=0$ just has a zero Jacobian. Exceptions this extreme are probably rare, but more commonly you could have some softmax in the network (e.g. in an attention layer) which saturates such that no gradient goes through. In that case for e.g. interpretability, it seems like you'd often still really want to "count" this, so arguably the activations would be *better* than the NTK for this case. (I've been working on a modification to the NTK to better handle this case.)

The NTK and the network activations have somewhat different properties and so it switches which one I consider most relevant. However, my choice tends to be more driven by analytical convenience (e.g. the NTK and the network activations lie in different vector spaces) than by anything else.