Chapter 6 of 0–21

Simplifying the model

Which parts matter: ablations, the functional form and the effective weights.

Before asking what the trained network computes, it helps to find out which of its parts matter at all. The ARENA course does this first, deriving a simpler “functional form” of the network [3]; the paper reports the same facts in its appendix [1]. All numbers below are measured on the paper's trained network at step 40,001.

Ablations: switch a part off and see what breaks

An ablation removes or replaces one part of the computation and measures how much the output changes. A zero ablation replaces the part's output with zeros; a mean ablation replaces it with its average over all inputs, which removes only the information that depends on the input (zeroing can also remove a useful constant offset).

A skip connection is a path that carries a layer's input around the layer and adds it back — the “x +” in x^{(1)} = x^{(0)}_2 + \ldots and x^{(2)} = x^{(1)} + \ldots of chapter 5.

“our networks do not significantly use the skip connection around the MLP”— Nanda et al. 2023, Sec. 3, p. 3 · source

Consequence 1: logits come from the neurons through one matrix

If the skip connection around the MLP is ignored, x^{(2)} \approx W_{out}\,\mathrm{MLP}(a,b) (plus constants), so

The paper calls W_L = W_U W_{out} (a 113 × 512 matrix) the neuron–logit map W_L: column m says how much one unit of neuron m's activation adds to each answer's logit. ARENA calls the same matrix (transposed) W_{logit}. A product of matrices that maps one meaningful quantity directly to another is an effective weight matrix; we will study W_L in the Fourier basis in chapter 15.

Consequence 2: each head is a two-way choice between a and b

Attention to “=” is negligible, so each head's softmax is effectively over two scores, which (chapter 4) is a sigmoid of their difference:

Moreover the query comes from the “=” position, whose vector x^{(0)}_2 is the same for every input. So the score of position 0 depends only on the token there: it is a fixed lookup table C_j over the 113 numbers, s_0 = C_j[a], and likewise s_1 \approx C_j[b] (up to the small difference between the positional embeddings of positions 0 and 1, see below). Hence

This is the paper's App. C.1.3 derivation [1]; ARENA calls C_j the effective attention weights W_{attn} [3].

<p class="fig-expect"><strong>What to expect:</strong> About half to a and half to b; almost nothing to '='.</p>
<p class="fig-meta">
  <a href="../chapters/06-simplifying-the-model.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. A.1; <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> section 3, “Development of commutativity”
  · Computed by <code>checkpoint_metrics.py</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Here is that two-way choice in the real network. Each dot is one input: its attention to a against the difference of the two scores, with the sigmoid drawn on top (pick another head in the dashboard's control bar).

<p class="fig-expect"><strong>What to expect:</strong> The points follow σ closely, because attention to '=' is negligible, and stay in its nearly straight middle part.</p>
<p class="fig-meta">
  <a href="../chapters/06-simplifying-the-model.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> App. C.1.3
  · Computed by <code>web/js/model/grokmodel.js</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

And the choice on every input at once: for each head, the attention from “=” to a, to b and to itself over all 12,769 questions.

<p class="fig-expect"><strong>What to expect:</strong> In the trained network the attention to a and to b form periodic stripes and the attention to '=' is close to 0 everywhere.</p>
<p class="fig-meta">
  <a href="../chapters/06-simplifying-the-model.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#nanda2023" title="Nanda et al. 2023">[1]</a> Fig. 4
  · Computed by <code>web/js/model/grokmodel.js</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

Consequence 3: the input is symmetric in a and b

The problem has commutativity: a + b = b + a. A network that has learned the task should treat a and b alike. In the trained network the positional embeddings of positions 0 and 1 are both small — length 0.12 and 0.12, against 1.18 for a typical token embedding — and point in similar directions (cosine similarity 0.62), whereas the “=” position's embedding is large (2.76). So whether a number sits at position 0 or 1 barely changes its vector, and a and b are processed (almost) identically:

<p class="fig-expect"><strong>What to expect:</strong> Rows 0 and 1 stay small (paper run: norms <span class="num" data-num="paper.pos_norm0">0.12</span> and <span class="num" data-num="paper.pos_norm1">0.12</span>, against <span class="num" data-num="paper.we_norm_mean">1.18</span> for a typical token embedding) and partly aligned (cosine <span class="num" data-num="paper.pos_cos01">0.62</span>), so a and b are treated almost alike; row 2 (the “=” position) is large.</p>
<p class="fig-meta">
  <a href="../chapters/06-simplifying-the-model.html">Learn step by step →</a>
   · Source: <a class="cite" href="../references.html#arena152" title="ARENA 1.5.2">[3]</a> “Answer some initial questions”
  · Computed by <code>browser (GrokModel)</code> (<a href="https://github.com/BurnyCoder/visualizing-grokking-modular-addition">code</a>)
</p>

The simplified network

Putting it together, the trained network behaves like

where E_t is token t's embedding plus its positional vector. Everything that depends on the input enters through the embeddings E_a, E_b, the lookup tables C_j, a ReLU, and the neuron–logit map W_L. The rest of the site explains what each of these does.