Superposition in a toy model

Runs in your browser

Survey: §5.2 The Superposition Hypothesis

A network can represent more features than it has dimensions if the features are rarely active at the same time. This model has 5 input features and only 2 hidden dimensions. Drag the sparsity slider and retrain to see when it starts sharing dimensions.

Sparsity S (probability a feature is zero) 0.90Importance decay (I_i = rⁱ) 0.90

Presets:

0 / 5features represented

2hidden dimensions

–loss

0steps

With 2 dimensions the model keeps only the 0 most important features.

Columns W_i of the 2 × 5 weight matrix: how each feature is embedded in the hidden space. Darker means more important. The unit circle is shown for reference.

W^TW. Off-diagonal entries are interference between features.

‖W_i‖. Above the dashed line (0.5) a feature counts as represented.

Σ_j≠i(Ŵ_i·W_j)²: how much feature i shares its direction with the others.

What is being computed

This is the ReLU-output toy model of Elhage et al. (2022): h = Wx, x′ = ReLU(Wᵀh + b). Training minimises importance-weighted reconstruction error, Σ_i I_i(x_i − x′_i)². Each feature is 0 with probability S and otherwise drawn from Uniform[0, 1]. Training is Adam (learning rate 0.01, batch 512) for 1,500 steps with a seeded random number generator, so the same settings give the same result every time.

Checked results (seed 1, importance decay 0.9, tested in tests/sim.test.ts): at S = 0 the model represents only 2 of 5 features, at S = 0.7 it represents 4, and at S = 0.9 it represents all 5. A feature counts as represented when ‖W_i‖ > 0.5. This reproduces the qualitative finding of the original paper: sparsity drives superposition.

Go further

Notebook: Superposition Hypothesis · Open in Colab

References

Nelson Elhage, Tristan Hume, Catherine Olsson, Neel Nanda, Tom Henighan, Yuntao Bai, Ben Mann, Amanda Askell, and Chris Olah. Toy models of superposition. Transformer Circuits Thread, 2022a. URL https://transformer-circuits.pub/2022/toy_model/.