Features, Directions, and Superposition

import numpy as np

rng = np.random.default_rng(0)
d, n_features = 64, 2000
directions = rng.normal(size=(n_features, d))
directions /= np.linalg.norm(directions, axis=1, keepdims=True)

gram = directions @ directions.T
np.fill_diagonal(gram, 0.0)
print(f"max |cosine| between any pair: {np.abs(gram).max():.3f}")
# 2000 directions in 64 dimensions, and no pair is more than weakly aligned.
# That headroom is the resource superposition spends.

Hold on to

• Features are directions; neurons are arbitrary basis vectors
• Superposition packs many features into fewer dimensions
• Polysemantic neurons are the expected consequence, not an anomaly

Try each one before opening the solution. Getting it wrong first is most of where the learning happens.

1. Run the near-orthogonality experiment across dimensions 16, 64, 256 and plot maximum pairwise alignment against dimension.
   Hint

   Generate random unit vectors, take the Gram matrix, and read off the largest off-diagonal magnitude.

   Solution

   Maximum pairwise alignment falls sharply as dimension grows: cramped and highly interfering at 16 dimensions, comfortable at 64, and near-orthogonal at 256 even with thousands of directions. That headroom is the resource superposition spends — a model can store many more features than it has dimensions provided they are sparse enough that interference rarely matters. It also means the number of features is not bounded by the number of neurons, which is why counting neurons tells you nothing about capacity.

   import numpy as np
   
   rng = np.random.default_rng(0)
   for d in (16, 64, 256):
       dirs = rng.normal(size=(2000, d))
       dirs /= np.linalg.norm(dirs, axis=1, keepdims=True)
       gram = dirs @ dirs.T
       np.fill_diagonal(gram, 0.0)
       print(f"d={d:4d}  max|cos|={np.abs(gram).max():.3f}  "
             f"mean|cos|={np.abs(gram).mean():.4f}")

   Check your work

   Paste this after your own code. If it runs without raising, you have it.

   import numpy as np
   rng = np.random.default_rng(0)
   def max_align(d, n=2000):
       v = rng.normal(size=(n, d)); v /= np.linalg.norm(v, axis=1, keepdims=True)
       g = v @ v.T; np.fill_diagonal(g, 0.0)
       return np.abs(g).max()
   assert max_align(256) < max_align(16)
   print("ok")

2. Find a polysemantic neuron in a small open model by collecting its top-activating inputs. Name at least two unrelated triggers.
   Hint

   Run a corpus through the model, record one neuron's activation per token, and take the top 50.

   Solution

   You will typically find a neuron whose top activations span genuinely unrelated contexts — a programming construct, a phrase in another language, and a formatting pattern, say. The temptation is to invent a story that unifies them; resist it. The straightforward reading is that this neuron is one coordinate shared by several superposed feature directions, and no single label describes it because no single feature owns it.

   import torch
   
   activations = []
   handle = model.blocks[LAYER].mlp.register_forward_hook(
       lambda m, i, o: activations.append(o[0, :, NEURON].detach()))
   
   with torch.no_grad():
       for batch in corpus_batches:
           model(batch)
   
   handle.remove()
   flat = torch.cat(activations)
   top = flat.topk(50).indices
   # Map those positions back to their tokens and read the contexts.

Sign in to track your progress through the lab.