Toy models of Superposition was the first mechanistic interpretability paper I chose to read and it is the first AI paper I have chosen to implement. One of the key advantages of this paper is how easy it actually is to implement (Far easier than it is to reason about in my opinion). The core premise of the paper is an analysis of results the Anthropic team found when they trained small toy models in order to display superposition, when an AI model represents multiple features in the same dimension in order to save neurons.
The paper centers around an extremely simple model. A simple encoder and decoder structure where high feature data is multiplied by a weight matrix with lower dimensionality. This model is trained on synthetic data where we generate a vector of values we classify as “features”. We create as many features as we have columns in our weight matrix. The goal is to observe how the model chooses to represent the features when there are fewer neurons than dimensions.
When implementing the paper, my initial approach was to do it using as little as Anthropic’s public implementation code and work from there. I did heavily use their code for the Matplotlib graphs so our representations would look similar, as is disclosed in the Jupyter notebook. Initially, I also did not include batching or models with different sparsity. As implementing these became necessary to follow Anthropic’s code I gradually added them.
One unique addition I made was adding a LinearToyModel class. This allowed me to compare how the linear model they described in the paper compares to the ReLU one they discuss. As discussed in the paper, the linear model does not display superposition as the ReLU allows the model to overcome slight noise in the loss.
I first went about implementing their intro diagram which displays a scatter plot of the features. In line with the paper, I created a weight matrix of 5 features and 2 neurons. This way you can plot each neuron as a dimension and see if superposition occurs. In the original paper, Anthropic only implemented this graph for a ReLU model but out of pure interest I did it on a linear model too. As you can see from the diagrams below, the vectors stayed orthogonal on the linear model (Did not display superposition), this means that the model dropped some features. On the ReLU model superposition is shown.
One feature that I wished was available on the original paper was the ability to see how these vectors move throughout the training process and with different levels of sparsity. To achieve this, I copied my code from the Jupyter notebook into a python file and wrote the matrix weights at each loss step into a .bin file (More performant than saving a number of matplotlib images). I then used Claude Opus 5 to make this interactive diagram which looks similar to the one I created from the paper.
While I had originally only planned to train 1 model for each figure (in contrast to Anthropic’s implementation), while working on the second diagram I realised that to get the best value out of the graphs I would need to train multiple models of different sparsity. In response, I added another dimension to my weight matrix. I also added batching to improve the performance of my models. I leaned heavily on Anthropic’s implementation for this code because it was my first time with either of these concepts. I hope that in the future I will be able to do these from scratch myself however.
The next figure I recreated was included in the original paper as a way to show interference between features. This is displayed by representing the features in the weight matrix on a bar chart. The length of the bar corresponds to the length of each feature vector (This means how intensely the feature is represented). The colour of the bar corresponds to how polysemantic it is. To find polysemanticity, we normalise each feature in the weight matrix and then get the dot product of these features and the features in the original weight matrix. We then get the length of this vector. The idea is that we then get a number to represent how close together these vectors point. The diagrams show that as sparsity increases, we generally get more features that display more interference.
This figure is also accompanied by another diagram which plots how each feature interferes with each other by having a square closer to red if they do interfere and a square closer to blue if they do not. Each diagram would have a strong diagonal (showing that each feature obviously interfered significantly with itself) and a small amount of interference with other features. As sparsity increased in the models trained, the amount of interference among the features did too.
My third diagram was a simple line graph of the feature sparsity and the dimensions allocated per feature. The dimensions per feature was found by dividing the number of neurons by the sum of the squared entries of the matrix. In the paper this is done by squaring the Frobenius norm of the matrix but I chose to simply square the matrix and sum each row instead as it was easier for me to reason about. The feature sparsity was simply 1/feature probability . My only other change to this graph was changing the axes. They originally corresponded to the formula from the paper but I changed them to a word description of what they represented.
The last graph is a scatter plot of the sparsity versus the number of dimensions per feature. The difference between this and the last plot is that it graphs it for each feature. A small amount of noise is added to the placement to make the points easier to see. A key finding from the paper is that the number of features per dimension bunches at different fractions. The most obvious example of this is at 0.5 where the points almost form a straight line. This makes it one of my favourite graphs in the paper.
Overall, I found this to be a really good learning experience. One of the things I found hardest was wrapping my head around some of the matrix operations. Once I got the hang of this, however, it made the paper much easier. My goals for the future are to get better at implementing these papers without having to rely on the original implementations. I would also like to expand my work on this paper in future to include more interactive graphs like the intro diagrams and implement some of the diagrams that anthropic did not include in their colab.
If you have any suggestions of papers I should read/implement or you are just getting into these kind of papers (like me) and have any questions on my implementation, feel free to reach out.