I Trained a Tiny Network to Compress Data. It Drew a Pentagon.

I wasn’t expecting geometry to show up. I was reproducing a small piece of Anthropic’s 2022 interpretability paper, “Toy Models of Superposition” [1], mostly because the central claim sounded implausible enough that I wanted to check it myself rather than take it on faith. The claim: a neural network can represent more features than it has dimensions to work with, by packing them in at angles to each other and tolerating a bit of interference. Ask it to compress five things into two dimensions under the right conditions, and it doesn’t pick two winners and give up on the rest. It arranges all five into a perfect pentagon.

I didn’t have PyTorch or any autograd library available, and no internet access to install one either, so everything below is plain NumPy, and I derived the backward pass by hand. That turned out to be the right kind of annoying. Deriving the gradients yourself forces you to actually understand what the model is doing to the data, rather than trusting a .backward() call you’ve never had to think about. If you have a math background and you’ve never hand-derived backprop through even a tiny network, I’d genuinely recommend it as an exercise. It’s ten minutes of chain rule that makes everything downstream click.

All the data in this post is synthetic. I generate it myself in code, there’s no external dataset involved, which is also how the original paper does it. All images, unless otherwise noted, are by the author.

The problem this is trying to explain

Here’s the motivating puzzle, and it’s a real one in interpretability research. If you look inside a trained neural network hoping to find individual neurons that cleanly represent individual concepts, one neuron for “is this a dog,” one for “is this red,” you mostly don’t find that. Instead you find neurons that seem to respond to several unrelated things at once, a neuron that fires for both cat faces and the front ends of cars, say. This is called polysemanticity, and it makes interpretability much harder, because you can’t just read off what a network “believes” by inspecting individual units.

The paper’s proposal is that polysemanticity isn’t noise or failure. It’s a real strategy the network uses on purpose, because it has more concepts to represent than it has neurons to represent them with, and most of those concepts are rarely active at the same time. If two features are almost never “on” simultaneously, the network can afford to let them share a direction in activation space, since the interference only costs something on the rare occasions both happen to fire together. This packing strategy is what the paper calls superposition, and its toy model is designed to be the simplest possible setting where you can watch it happen and actually measure it.

The model, and the math I had to work out to train it

The setup is small on purpose. You have n synthetic features, each one a number between 0 and 1 that’s zero most of the time (that’s the sparsity) and nonzero the rest of the time. You compress them down through a bottleneck of m hidden dimensions, where m is smaller than n, and then try to reconstruct the original features on the way back out through a ReLU.

Concretely, with a single weight matrix W of shape (m, n) used for both the compression and the reconstruction:

Training minimizes a weighted squared error between $x$and $hat{x}$

Since I didn’t have autograd, I needed $partial mathcal{L} / partial W$ and $partial mathcal{L} / partial b$ by hand. This isn’t bad once you set it up as two matrix multiplications sharing the same weights. Let $z = W^top h + b$

where $odot$ is elementwise multiplication and $[z > 0]$

That second equation is the only part that took me a minute to get right, and it’s a nice small reminder of why tied weights are a little more interesting to differentiate than they look.

Here’s the full implementation, batched over samples, with Adam written out explicitly since I didn’t have that for free either:

import numpy as npdef generate_batch(batch_size, n_features, sparsity, importance, rng):    """Sparse synthetic features: each is independently 'on' with probability    (1 - sparsity); when on, its value is Uniform(0, 1)."""    values = rng.uniform(0, 1, size=(batch_size, n_features))    mask = rng.uniform(0, 1, size=(batch_size, n_features)) > sparsity    return values * maskclass ToyModel:    def __init__(self, n_features, n_hidden, importance, rng, lr=1e-3):        self.n_features = n_features        self.n_hidden = n_hidden        self.importance = importance        self.W = rng.normal(0, 1 / np.sqrt(n_features), size=(n_hidden, n_features))        self.b = np.zeros(n_features)        self.lr = lr        self.mW = np.zeros_like(self.W); self.vW = np.zeros_like(self.W)        self.mb = np.zeros_like(self.b); self.vb = np.zeros_like(self.b)        self.t = 0        self.beta1, self.beta2, self.eps = 0.9, 0.999, 1e-8    def forward(self, X):        H = X @ self.W.T          # (batch, n_hidden)        Z = H @ self.W + self.b   # (batch, n_features)        Xhat = np.maximum(Z, 0)        return H, Z, Xhat    def loss(self, X, Xhat):        diff = X - Xhat        return (self.importance[None, :] * diff**2).sum(axis=1).mean()    def step(self, X):        batch = X.shape[0]        H, Z, Xhat = self.forward(X)        diff = X - Xhat        dXhat = -2 * self.importance[None, :] * diff / batch        dZ = dXhat * (Z > 0)        db = dZ.sum(axis=0)        dH = dZ @ self.W.T        dW = H.T @ dZ + dH.T @ X        self._adam_update(self.W, dW, 'mW', 'vW')        self._adam_update(self.b, db, 'mb', 'vb')        return self.loss(X, Xhat)    def _adam_update(self, param, grad, mname, vname):        self.t += 1 if mname == 'mW' else 0        m = getattr(self, mname); v = getattr(self, vname)        m[:] = self.beta1 * m + (1 - self.beta1) * grad        v[:] = self.beta2 * v + (1 - self.beta2) * (grad ** 2)        mhat = m / (1 - self.beta1 ** self.t)        vhat = v / (1 - self.beta2 ** self.t)        param -= self.lr * mhat / (np.sqrt(vhat) + self.eps)

Every result below came from actually running this, not from the paper’s numbers.

Experiment one: the network quietly gives up on the features that don’t matter

First check: does the compression even work, and what happens to the features that matter less? I trained with 20 features going into 5 hidden dimensions, sparsity 0.9 (each feature is zero 90 percent of the time), and importance decaying geometrically across the 20 features, so feature 0 matters most and feature 19 matters least.

The right panel is the one I found genuinely satisfying to see appear from real numbers instead of a description in a paper. The network isn’t reconstructing all 20 features with mediocre accuracy. It’s making a clear decision: represent the important features well, and past a fairly sharp threshold, just stop bothering with the rest. Nobody told it to do that. It fell out of gradient descent on a plain weighted MSE loss.

Experiment two: five features, two dimensions, and the pentagon

This is the result that got me to write this up. Squeeze 5 equally important features into just 2 hidden dimensions, and watch what the 5 learned feature vectors actually look like as arrows in that 2D space, at three different sparsity levels.

I checked that last claim numerically rather than eyeballing it: the five angles in the high-sparsity run came out at 14.2, 85.8, 158.3, 229.8, and 301.6 degrees, which are 71.5 to 72.5 degrees apart, essentially a perfect regular pentagon, accurate to about half a degree. There is no term in the loss function that rewards symmetry. A pentagon is just the most efficient way to place 5 points on a circle so that every pair is as far apart as every other pair, which minimizes the worst-case interference between any two features. Gradient descent found that arrangement on its own because it’s genuinely the optimal packing, not because anyone told it what a pentagon was.

Experiment three: what interference actually looks like

The pentagon picture is nice for 5 features in 2 dimensions because you can actually see it. With more features you can’t draw the picture anymore, but you can look directly at how much any two features are stepping on each other’s toes, by computing $W^top W$, whose diagonal tells you how well each feature reconstructs itself and whose off-diagonal entries tell you how much reconstructing one feature corrupts another.

This lines up with experiment one almost exactly. The same cutoff around feature 12 or 13 shows up independently in both, which is a good sanity check that this is measuring a real effect and not an artifact of one particular plot.

Experiment four: turning the sparsity dial

The last experiment is the one that actually earns the word “phase transition,” a term the paper uses and that I was skeptical of until I saw it myself. Fix 30 features and 5 hidden dimensions, and sweep sparsity from 0 (features are almost always active) up to 0.99 (features are active only 1 percent of the time), retraining from scratch at each level.

The flat stretch at low sparsity and the sharp climb after roughly 0.8 is what makes this a phase transition rather than a smooth tradeoff. There’s a real regime change in the strategy the network adopts, not just a gradual dial turning. Below some threshold, superposition costs more in interference than it’s worth. Above it, it’s clearly worth it, and the network commits to using it.

What this actually tells you

None of this required a large model, a GPU, or a real-world dataset. Thirty synthetic numbers and an afternoon were enough to watch a real instance of a phenomenon that’s currently central to how people think about interpreting much bigger models. The headline claim survives contact with actual code: networks really do represent more concepts than they have neurons, they do it by exploiting sparsity, and the geometry they land on to do it, antipodal pairs, regular polygons, isn’t decorative. It’s the mathematically efficient packing for the amount of interference the loss function is willing to tolerate.

The practical stakes are bigger than a pentagon. If a large language model is representing thousands of concepts inside a few thousand neurons, and it almost certainly is, then polysemanticity isn’t a bug you can fix by staring harder at individual neurons. It’s the predictable consequence of compression under sparsity, and it’s a big part of why mechanistic interpretability has had to develop tools, like sparse autoencoders, specifically designed to undo this packing and pull individual features back out. Having now watched superposition happen in a system small enough to fully see, I understand why that line of research exists in a way I don’t think I would have from reading about it alone.

···

References:

[1] N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg and C. Olah, Toy Models of Superposition (2022), Transformer Circuits Thread