Toy Models of Superposition

Setting. Features $x\in\mathbb R^n$ have independent coordinates, $x_i=0$ with probability $S$ (the sparsity; $1-S$ is the density) and $x_i\sim U[0,1]$ otherwise, with importances $I_i$. The hidden layer $h=Wx\in\mathbb R^m$, $m\lt n$, is decoded with the same matrix, and the two basic models differ only in the readout: $$x'=W^\top Wx+b\ \ \text{(linear)},\qquad x'=\mathrm{ReLU}(W^\top Wx+b)\ \ \text{(ReLU output)},\qquad L=\mathbb E\sum_iI_i(x_i-x'_i)^2.$$ Column $W_i$ is feature $i$'s direction, of norm near 1 if the feature is represented and near 0 if not, and the feature dimensionality $$D_i=\frac{\|W_i\|^2}{\sum_j(\hat W_i\cdot W_j)^2}$$ is 1 for a dedicated dimension, $1/2$ for an antipodal pair ($W_j=-W_i$) and 0 for a dropped feature, while $D^*=m/\|W\|_F^2$ counts dimensions per represented feature overall. Superposition means more represented features than dimensions. All findings are empirical, apart from closed-form losses of a few candidate configurations.

A ReLU readout superposes sparse features and a linear readout does not. With $n=20$, $m=5$ and $I_i=0.7^i$, the linear model always learns the top five features, as PCA would. The ReLU output model matches it at density 1, but already at density 0.3, and more so below, it represents extra features through non-orthogonal directions, least important first and in antipodal pairs. Since the two models share architecture, data and loss, superposition needs the output nonlinearity together with sparsity; sparsity alone does nothing for a linear readout.

The loss is feature benefit against interference, and sparsity makes interference rare. For the linear model with Gaussian inputs (as in Saxe et al. on deep linear networks) the loss is, up to constants, $\sum_iI_i(1-\|W_i\|^2)^2+\sum_{i\ne j}I_j(W_j\cdot W_i)^2$: a benefit for each represented feature and a penalty for every overlap, so representing more than $m$ features never lowers the loss. For the ReLU model, split the expectation by which features are active; as $S\to1$ inputs with one active feature dominate (the all-zero input only penalizes positive biases), with loss $$L_1=\sum_i\int_0^1I_i\big(x_i-\mathrm{ReLU}(\|W_i\|^2x_i+b_i)\big)^2dx_i+\sum_{i\ne j}\int_0^1I_j\,\mathrm{ReLU}\big((W_j\cdot W_i)\,x_i+b_j\big)^2dx_i.$$ The ReLU makes negative interference cost nothing, a negative bias turns small positive interference negative, and with unit norms, zero bias and equal importances the interference term is a Thomson-type energy for points on a sphere. The smallest case shows the orders of magnitude (my computation, not in the paper; Monte Carlo agrees). With $n=2$, $m=1$, $b=0$, density $p$ and importance $r$ for the second feature, compare dropping that feature with the antipodal pair $W=[1,-1]$, which is exact when one feature is active and loses the smaller one from both outputs when both are: $$L_{\rm drop}=\frac{rp}3,\qquad L_{\rm antipodal}=(1+r)\,p^2\,\mathbb E\min(U,V)^2=\frac{(1+r)p^2}6,$$ so the pair wins exactly when $p\lt2r/(1+r)$. Through a linear readout it costs $(1+r)p/3$, or $(1+r)(p/3-p^2/4)$ with the best bias, against $r(p/3-p^2/4)$ for dropping, so it loses at every density.

The choice between dropping, superposing and dedicating a dimension changes abruptly. The optimal loss is the minimum of a few smooth curves, one per candidate configuration, so where two cross its derivative jumps, a first-order transition. For two features in one dimension, across importance and density, the trained diagram has three regions with sharp borders (extra feature dropped, antipodal pair, extra feature given the dimension) and matches the one built from the closed-form losses of $W=[1,0]$, $[0,1]$, $[1,-1]$. "Phase change" here means a jump in the optimum of a finite system, with no large-system limit. The candidate list is not complete: an external replication solved the bias-free two-feature case exactly and found a further phase with similar weights, and regions where the optimum moves continuously.

Identical features settle into uniform polytopes, and unequal ones deform them. With $n=400$ identical features in $m=30$ dimensions, $D^*$ stays at 1 and at $1/2$ over wide ranges of sparsity, and the individual $D_i$ cluster at $3/4$ (tetrahedron), $2/3$ (triangle), $1/2$ (antipodal pair), $2/5$ (pentagon), $3/8$ (square antiprism) and 0. Each value is $m/n$ for $n$ points in $m$ dimensions, and the reason is a tight frame (my gloss; the paper reads the shapes as Thomson-problem solutions): the denominator of $D_i$ is $\hat W_i^\top WW^\top\hat W_i$, so for points of equal norm $D_i=m/n$ at every vertex exactly when $WW^\top=(n/m)I_m$. Polytopes whose vertices differ appear split into smaller polytopes in orthogonal subspaces, with no interference across factors. Unequal features deform the shapes. Five features in two dimensions form a regular pentagon; making one denser pushes the others away from it, and making it sparser lets them move closer to it until the pentagon changes abruptly into two antipodal pairs with that feature dropped. Correlated features prefer orthogonal directions, and with too little room collapse onto their sum $(a+b)/\sqrt2$, as PCA would; anticorrelated features prefer to be antipodal.

With a ReLU on the hidden layer, neurons become polysemantic and can still compute. In $h=\mathrm{ReLU}(Wx)$, $x'=\mathrm{ReLU}(W^\top h+b)$, neurons are monosemantic on dense data and turn polysemantic as sparsity grows, and both kinds coexist in one model. This is a weak test, since the model gains nothing from the hidden ReLU and avoids using it when it can, say with positive biases that make every neuron linear. Computing $|x|$ requires the hidden ReLU: in $h=\mathrm{ReLU}(W_1x)$, $y'=\mathrm{ReLU}(W_2h+b)$ with active inputs uniform on $[-1,1]$, the target $|x|=\mathrm{ReLU}(x)+\mathrm{ReLU}(-x)$ needs two neurons per feature, and with $n=3$, $m=6$ the model learns exactly that. With $n=100$, $m=40$ it is monosemantic on dense data and on sparse data computes $|x|$ for more features than it has neurons, so superposition supports computation as well as storage. Many neurons carry one primary feature with a large weight and a few secondary ones with small weights, so their strongest activations look clean and their weaker ones mixed, as in language-model neurons.

Superposition predicts vulnerability to adversarial inputs, and feature dimensionalities change in jumps during training. In superposition a feature's end-to-end weights look like $(W^\top W)_i=(1,\epsilon,-\epsilon,\ldots)$, and the off-diagonal entries give an attacker a direction that moves feature $i$'s output without activating it. With the optimal $L_2$ attack in closed form, vulnerability rises more than threefold as superposition forms and tracks features per dimension, $1/D$. During ordinary training, feature dimensionalities jump between the polytope values, each jump with a small drop in the loss. Both observations are correlational and confined to the toy model.