DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

I Trained a Transformer in the Browser, Then Found the Algorithm Inside It

Every weight of a language model is sitting on a disk you can read, and nobody can tell you why it produced the token it produced.

Saliency and attribution methods answer where did it look. Mechanistic interpretability asks the harder question: what algorithm is it running, and can I predict its behaviour on an input I have not shown you yet?

This page tests that claim end to end on a model it trains live in your browser: a real attention-only transformer — token and positional embeddings, 2 layers, 2 heads, d_model = 16, hand-written forward and backward pass, Adam, no libraries.

No pre-trained weights, no stored results. It starts from random numbers when the page loads.

Live: https://dev48.infy.uk/ai/days/day62-mechanistic-interpretability.html

Pick a task whose correct algorithm you already know

The task: a sequence of distinct random tokens, with one block of 3 repeated later at a random offset. After the second copy, the model must continue it.

The only rule that works is find the earlier copy of my own token, look at the token after it, copy that — the induction circuit Olsson et al. found in real models in 2022. It provably needs two layers, so a 2-layer model is the smallest thing that can express it.

Crucially, the controls are computed rather than claimed: nothing about a token alone predicts what follows it, and because the earlier copy moves, no fixed positional rule works either. The page brute-forces those ceilings so you know what "the model learned something" actually means here.

Keep the residual stream a sum

Every component reads a normalised copy of the residual stream and adds its output back. Nothing overwrites anything. So the stream at the last layer is literally:

embed + pos + sum of head outputs
Enter fullscreen mode Exit fullscreen mode

Because the final LayerNorm scale is a single number fixed by the forward pass, the unembedding is linear in that sum. Therefore each component contributes an exactly-defined number of logits to every token.

That is what makes direct logit attribution an identity rather than a heuristic — and the page checks that the attributions sum to the actual logits.

Check the backward pass at initialisation, not at convergence

Hand-written gradients get checked against finite differences. Do it at initialisation, where the gradients are large and varied. At convergence everything is near zero and a broken gradient looks fine.

Ablation, three ways, with a control that is bit-identical or wrong

Zero-ablation, mean-ablation and resample-ablation disagree, and each answers a slightly different question. The essential discipline is the control: ablating a component that should not matter must leave the output bit-identical. If it does not, your ablation is touching something you did not intend and every number after it is suspect.

Activation patching, kept honest by two identities

Patching a clean activation into a corrupted run should recover the clean logits when you patch everything, and change nothing when you patch nothing. Those two endpoints are cheap to check and they catch indexing bugs that otherwise masquerade as findings.

The scores that make it a circuit claim

Two measurements over the test set:

  • Previous-token score — how much of position i's attention lands on i-1.
  • Induction score — how much of a repeated token's attention lands on the position immediately after the earlier copy of itself.

That second one is the signature. Layer 0 learns previous-token heads; layer 1 uses them via K-composition to build the induction head. The page shows the QK and OV circuits separately, because they answer different questions: where to look, and what to move once you are looking.

Superposition, and a sparse autoencoder on the model's own activations

d_model = 16 is smaller than the number of things the model needs to represent, so features share directions. A sparse autoencoder trained on the model's activations pulls some of them apart — on a model small enough that you can check whether the features it finds are real.

Repo: https://github.com/dev48v/ai-from-zero

Top comments (0)