Anthropic's recent mechanistic interpretability paper, Toy Models of Superposition, helps to demonstrate the conceptual richness of very small feedforward neural networks. Even when being trained on synthetic, hand-coded data to reconstruct a very straightforward function (the identity map), there appears to be non-trivial mathematics at play and the analysis of these small networks seems to providing an interesting playground for mechanistic interpretability. 

While trying to understand their work and train my own toy models, I ended up making various notes on the underlying mathematics. This post is a slightly neatened-up version of those notes, but is still quite rough and un-edited and is a far-from-optimal presentation of the material. In particular, these notes may contain errors, which are my responsibility.

1. Directly Analyzing the Critical Points of a Linear Toy Model

Throughout we will be considering feedforward neural networks with one hidden layer. The input and output layers will be of the same size and the hidden layer is smaller. We will only be considering the autoencoding problem, which means that our networks are being trained to reconstruct the data. The first couple of subsections here are largely taken from the Appendix to the paper "Neural networks and principal component analysis: Learning from examples without local minima." by Pierre Baldi and Kurt Hornik. (Neural networks 2.1 (1989): 53-58).

Consider to begin with a completely liner model., i.e. one without any activation functions or biases. Suppose the input and output layers have $D$ neurons and that the middle layer has $d<D$ neurons.  This means that the function that the model is implementing is of the form $x\mapsto ABx$, where $x\in R^D$, $B$ is a $d\times D$ matrix, and $A$ is a $D\times d$ matrix. That is, the matrix $B$ contains the weights of the connections between the input layer and the hidden layer, and the matrix $A$ is the weights of the connections between the hidden layer and the output layer. It is important to realise that even though - for a given set of weights - the function that is being implemented here is linear, the mathematics of this model and the dynamics of the training are not completely linear.

The error on a given input $x$ will be measured by $\Vert x-ABx{\Vert }^2$ and on the data set $\{x^t{\}}_{t=1}^T$, the total loss is 

Define $\Sigma$ to be the matrix whose $(i,j{)}^{th}$ entry ${\sigma }_{ij}$ is given by 

Clearly this matrix is symmetric.  

Assumption. We will assume that the data is such that a) $\Sigma$ is invertible and b) $\Sigma$ has distinct eigenvalues. 

Let ${\lambda }_1>\cdots >{\lambda }_D$ be the eigenvalues of $\Sigma$.

1.1 The Global Minimum

Proposition 1. (Characterization of Critical Points) Fix the dataset and consider $L$ to be a function of the two matrix variables $A$ and $B$. For any critical point $(A,B)$ of $L$, there is a subset $I\subset \{1,\dots ,D\}$ of size $d$for which 

1. $AB$ is an orthogonal projection onto a $d$-dimensional subspace spanned by orthonormal eigenvectors of $\Sigma$ corresponding to the eigenvalues $\{{\lambda }_i{\}}_{i\in I}$; and
2. $L(A,B)=tr\Sigma -{\sum }_{i\in I}{\lambda }_i={\sum }_{i\notin I}{\lambda }_i$.

Corollary 2. (Characterization of the Minimum) The loss has a unique minimum value that is attained when $I=\{1,\dots ,d\}$, which corresponds to the situation when $AB$  is an orthogonal projection onto the $d$-dimensional subspace spanned by the eigendirections of $\Sigma$ that have the largest eigenvalues. 

Remarks. We won't try to spell out all of the various connections to other closely related things, but for those who want some more keywords to go away and investigate further, we just remark that the minimization problem being studied here is about finding a low-rank approximation to identity and is closely related to Principal Component Analysis. See also the Eckart–Young–Mirsky Theorem.

We begin by directly differentiating $L$ with respect to the entries of $A$ and $B$. Using summation convention on repeated indices, we first take the derivative with respect to $b_{j^{\prime }{k}^{\prime }}:$

Setting this equal to zero and interpreting this equation for all $j^{\prime }=1,\dots ,d$ and $k^{\prime }=1,\dots ,D$ gives us that 

Then, separately, we differentiate $L$ with respect to $a_{i^{\prime }{j}^{\prime }}$ : 

Setting this equation equal to zero for every $i^{\prime }=1,\dots ,D$ and $j^{\prime }=1,\dots ,d$ we have that: 

Thus 

Since we have assumed that $\Sigma$ is invertible, the first equation immediately implies that $A^T=A^{T}AB$. If we assume in addition that $A$ has full rank (a reasonable assumption in any case of practical interest), then $A^{T}A$ is invertible and we have that 

which in turn implies that 

where we have written $P_A$ to denote the orthogonal projection on to the column space of $A$.  

Claim. We next claim that $\Sigma$ commutes with $P_A$.

Proof of claim. Plugging (5) into (3), we have:

Then, right-multiply by $A^T$ and use the fact that $P_A^T=P_A$ to get:

The right-hand side is manifestly a symmetric matrix, so we deduce that $\Sigma P_A$ is symmetric. If the product of two symmetric matrices is symmetric then they commute, so this indeed shows that $\Sigma$ commutes with $P_A$ and completes the proof of the claim.

Now let $U$ be the orthogonal matrix which diagonalizes $\Sigma$,  i.e. the matrix for which 

where $\Lambda$ is a diagonal matrix with entries ${\lambda }_1>{\lambda }_2>\cdots >{\lambda }_D>0$. 

Claim. We next claim that $P_A=U{P}_{U^{T}A}{U}^T$ and that $P_{U^{T}A}$ is diagonal.  

Proof of Claim. Firstly, using the standard formula for orthogonal projections, we have 

 which implies that

To show that $P_{U^{T}A}$ is diagonal, we show that it commutes with the diagonal matrix $\Lambda$ (any matrix that commutes with a diagonal matrix must itself be diagonal). Starting from $P_{U^{T}A}\Lambda$, we first insert the identity matrix in the form $U^{T}U$, and then use (8) and (9) thus: 

Then recall that we have already established that $P_A$ commutes with $\Sigma$. So we can swap them and then performing the same trick in reverse: 

This shows that $P_{U^{T}A}$ commutes with $\Lambda$ and completes the proof of the claim.

So, given that $P_{U^{T}A}$ is an orthogonal projection of rank $d$ and is diagonal, there exists a set of indices $I=\{i_1,\dots ,i_d\}$ with $1\le i_1<i_2<\cdots <i_d\le D$ such that the $(i,j{)}^{th}$ entry of $P_{U^{T}A}$ is zero if $i\ne j$ and 1 if $i=j\text{and}i\in I$. And since $P_A=U{P}_{U^{T}A}{U}^T$, we see that 

where $U_I$ is formed from $U$ by simply setting to zero the $j^{th}$ column if $j\notin I$. This is manifestly an orthogonal projection onto the span of $\{u_{i_1},\dots ,u_{i_d}\}$, where $u_1,u_2,\dots ,u_D$ is an orthonormal basis of eigenvectors of $\Sigma$ (and indeed the columns of $U$).  Combining these observations with (5), we have that

This proves the first claim of the proposition.

To prove the second part, write $AB=\left[p_{ij}\right]$ and compute thus: 

But we know from (7) and (11) that $P_{U_I}\Sigma P_{U_I}=\Sigma P_{U_I}$ and so this last line is actually just equal to

Focussing on the second term and using (11), then (9) and (8), then cancelling $U^{T}U=I$, and then - to reach the last line - cyclicly permuting the matrices inside the trace operator to produce another $U^{T}U$ cancellation, we have:

The diagonal form of $P_{U^{T}A}$ means that this final expression is equal to ${\sum }_{i\in I}{\lambda }_i$, meaning that 

Since $tr\Sigma ={\sum }_{i=1}^D{\lambda }_i$ (the trace is always equal to the sum of the eigenvalues), this completes the proof of the proposition. $\square$

Remarks. Equation (10) above tells us that $col\left(A\right)=span⟨u_{i_1},\dots ,u_{i_d}⟩$, which means that there exists an invertible matrix $C$ with $A=U_{I}C$. Then, using (4), we compute that 

So we have:

1.2 Characterizing Other Critical Points

This subsection is something of an aside, but it is included for completeness.

Proposition 3. (Other Critical Points are Saddle Points.) Fix the dataset and consider $L$ to be a function of the two matrix variables $A$ and $B$. Every other critical point is a saddle point, i.e. if $(A,B)$ is a critical point but not equal to the unique minimum, then exist $\tilde{A}$ and $\tilde{B}$ which are arbitrarily close to $A$ and $B$ respectively and at which a lower loss is achieved. 

Proof. Since $(A,B)$ is not the unique global minimum, we know from Corollary 2 that $I\ne \{1,\dots ,d\}$. This means that there are distinct indices $j$ and $k$ for which $j\in I$, $k\notin I$ and $k<j$. In particular, bear in mind that ${\lambda }_k>{\lambda }_j$. 

Now, given any $ϵ>0$, put 

And let us form the new matrix ${\tilde{U}}_I$ by starting with $U_I$ and replacing the column $u_j$ with ${\tilde{u}}_j$. Write 

We want to calculate the loss of the model at $(\tilde{A},\tilde{B})$. We ought to bear in mind that it is not a critical point, so we cannot assume the intermediate results in the proof of Proposition 2, but it turns out that the bits that are most useful for this computation rely only on algebra and (13), (14). We start from the equivalent of line (11) which is that $\tilde{A}\tilde{B}={\tilde{U}}_I{\tilde{U}}_I^T=P_{{\tilde{U}}_I},$  which implies that $\tilde{A}\tilde{B}=P_{\tilde{A}}$. And so just as in (12) above, we have 

Now, looking at the final term on the right-hand side, we have $P_{{\tilde{U}}_I}\Sigma P_{{\tilde{U}}_I}^T=P_{\tilde{A}}\Sigma P_{\tilde{A}}$ and (by cycling permutation) $tr\left(P_{\tilde{A}}\Sigma P_{\tilde{A}}\right)=tr\left(P_{\tilde{A}}\Sigma \right)$. And since 

we have:

We also use $tr\left(P_{{\tilde{U}}_I}\Sigma \right)=tr\left(P_{\tilde{A}}\Sigma \right)$ and (16) on the second term on the right-hand side of (15) to ultimately arrive at:

So we are interested in computing the diagonal elements of $P_{U^T\tilde{A}}$. Fix $i\in \{1,\dots ,D\}$. The $i^{th}$ diagonal entry is given by: 

This can be computed directly from the definition of ${\tilde{U}}_I$ to give that the $i^{th}$ entry on the diagonal is equal to

Therefore 

Since ${\lambda }_k>{\lambda }_j$ this shows that in an arbitrarily small neighbourhood of the critical point $(A,B)$ we can find a point $(\tilde{A},\tilde{B})$ where smaller loss is achieved. We will not bother doing so here, but one can also check that $(A,B)$ is not a local maxima by using the fact that for fixed (full rank) $A$, the function $z\mapsto \parallel x-Az{\parallel }^2$ is convex. $\square$

2. Sparse Data, Weight Tying, and Gradients

Abstractly analyzing critical points is not at all the same as training real models. In this section we start to think about data and the optimization process.

2.1 Sparse Synthetic Data

Here we describe the kind of training data used in Anthropic's toy experiments

Fix a number $S\in [0,1]$. This parameter is the sparsity of the data. We will typically be most interested in the case where $S$ is close to 1. 

Let $\{B_i^t{\}}_{t=1}^{\infty }{}_{i=1}^D$ be an independent and identically distributed family of Bernoulli random variables with parameter $(1-S)$. And let $\{U_i^t{\}}_{t=1}^{\infty }{}_{i=1}^D$ be an IID family of $\text{Uniform}([0,1])$ random variables. Write $X_i^t=B_i^t{U}_i^t$ and $X^t:=(X_1^t,\dots ,X_D^t)$. Our datasets $\{x^t{\}}_{t=1}^T\in R^D$ will be drawn from the IID family $\{X^t{\}}_{t=1}^{\infty }$. Notice that

• Independently, for each data point $x^t\in R^D$ and for each $i\in \{1,\dots ,D\}$, we will have $P(x_i^t=0)=S$. 
• So, the expected number of non-zero entries for each data point is $(1-S)D$. To bring this in line with the way people say things like "$k$-sparse", we can say that the data is, on average, $(1-S)D$-sparse.
• $E\left(x_i^t\right)=\frac{1}{2}(1-S)$.

Remark. Judging from some of the existing literature on the linear model that we analyze in Section 1 (e.g. "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks." by Andrew M. Saxe, James L. McClelland and Surya Ganguli), it seems like it's tempting to make an assumption/simplification/approximation that $\Sigma =I$. I still don't feel like I understand how justifiable that is - for me this question is a potential 'jumping-off' point for further analysis of the whole problem. Recall that the matrix $\Sigma$ is equal to  ${\sum }_{t=1}^T{x}^t\otimes x^t$.  Certainly the probability that an off-diagonal entry of $x^t\otimes x^t$ is equal to zero is $1-(1-S{)}^2$ whereas for the diagonal entries it is just $S$. And note that $E\left(x_i^t{x}_j^t\right)=\frac{1}{4}(1-S{)}^2$ if $i\ne j$ and $E((x_i^t{)}^2)=\frac{1}{3}(1-S)$.   But the diagonal entries are still independent and I'm not sure why thinking of them as equal makes sense. 

The data (and the loss) are model two main ideas: Firstly, that the coordinate directions of the input space act as a natural set of features for the data. And secondly, when $S$ is close to 1, the sparsity of the data is supposed to capture the fact that features really do often tend to be sparse in real-world data, i.e. we see that for any given object or any given word/idea that appears in a language, it is the case that most images don't contain that object and most sentences don't contain that word or idea. 

2.2 Weight Tying and The Gradient Flow

In practice, when we train an autoencoder like this, we do so with weight tying. Roughly speaking, this means that we only consider the case where $A=B^T$. Proposition 1 does indeed allow for a global minimum in which $A=B^T$: This is achieved by essentially taking $C=I$ in the equations ($\star$) at the end of Section 1.1, i.e. we have:

But note that we don't actually want to try to repeat the analysis of Section 1 on a loss of the form ${\sum }_t\parallel x^t-W^{T}W{x}^t{\parallel }^2$. This would be a higher-order polynomial function of the entries of $W$ and so it's genuinely a different and potentially more complicated functional. The way that weight-tying is done in practice is more similar to saying that we insist during training that updates are made that preserve the equality $A=B^T$.

Equations (1) and (2) in subsection 1.1 are obtained as a direct result of differentiating the loss with respect to individual entries of the matrices (or individual 'weights' if we interpret this model as a feedforward neural network without activations). Our computations show that:

In an appropriate continuous time limit, if we set the learning rate to 1, the weights during training evolve according to the differential equations:

Remarks. Notice that there is a certain deliberate sloppiness here: One doesn't really have a fixed matrix $\Sigma$ and then run this gradient flow for all time; the matrix $\Sigma$ is a function of (a batch of) training data. So we need to be careful about any further manipulations or interpretations of these equations. 

Those caveats having been noted, if we additionally add in the weight-tying constraint $A=B^T$ , we get:

We can even make the substitution $W=\overline{W}{U}^T$ to introduce the form:

 In components (and without summation convention) the equation reads 

Let $\{{\overline{w}}^i{\}}_{i=1}^D$ denote the set of columns of $\overline{W}$ so that (24) can becomes:

Expanding the brackets and executing the sum over $l$ gives:

Then the sum over $k$ further simplifies the first term to give: 

Finally, just peel off the $k=j$ term from the remaining summation to arrive at the equation 

Remarks. (cf. the previous two Remarks)  If we assume that $\Sigma =I$, then $W=\overline{W}$ and the equation above arises as gradient descent on the energy functional 

It's plausible that a reasonable line of argument to justify this is that since no particular directions in the data are special, it means that over time, on average, the effects of different eigenvalues of $\Sigma$ just somehow 'average out'. But I don't endorse or understand how that argument would actually go. Regardless, if we just assume this for now, as is explained in the Anthropic paper, we can think of the two terms in (28) as being in competition. The first term suggests that model 'wants' to learn the $k^{th}$ feature by arranging $|w^k|=1$. However, as it tries to does so, it incurs a penalty - given by ${\sum }_{i\ne k}\lambda |⟨w^i,w^k⟩{|}^2$  - that can reasonably be interpreted as the extent to which the hidden representation $w^k\in R^d$ of that feature interferes with its attempts to represent and reconstruct the other features. 

3. The ReLU Output Model 

3.1 The Distribution of the Data and the Integral Loss

Perhaps a better way to try to incorporate information about the distribution of the data into the analysis here is to directly let $\mu$ be the distribution (i.e. in the proper measure-theoretic since) of $X_t$ on $R^D$ and to consider

In the Anthropic paper and in my own work, we are ultimately more interested in a model with biases and ReLUs at the output layer. 

Performing an analysis anything like that done in Section 1 seems much harder for this model, but perhaps more progress can be made studying the integral above. 

The synthetic data we described in the previous section is all contained in the cube $Q:=[0,1{]}^D\subset R^D$ . In the sparse regime i.e. with $S$ close to 1, the vast majority of the data is concentrated around the lower-dimensional skeletons of the cube. For $l=1,\dots ,D$, if we write $Q_l$ for the set of points in the cube with only $l$ non-zero entries, i.e.

then $C$ is the disjoint union  

Without a closer analysis of binomial tail bounds I can't immediately tell how well-justified it is to say, ignore ${\cup }_{l\ge 2}{Q}_l$ and focus the analysis just on 1-sparse vectors in the dataset. i.e. You might want to say that $\mu ({\bigcup }_{l\ge 2}{Q}_l)$ is sufficiently small such that that region contributes only negligibly to the integral. Then you can start to work with more manageable expressions To my mind this is another concrete potential 'jumping-off' point if one were to do more investigation. In particular, it is in the direction of the observations made in Toy Models of Superposition to suggest a link between this problem and 'Thomson Problem'.