Split into heads, gate, merge back
Medium15 min10 hidden tests
An interpretability team switches attention heads off to see which ones a model needs. They work on the attention output just before its final projection, a tensor of shape (B, T, D): B sequences, T tokens and D features, where head h owns the block of D / H columns starting at h * D / H. Write gate_heads(x, n_heads, gates) that splits x into heads of shape (B, H, T, D / H), multiplies head h by gates[h] (0 switches it off, 1 keeps it), and merges the heads back into (B, T, D). Return both the split heads and the merged result.
Gate the heads
head h owns columns h·Dh to (h + 1)·Dh − 1D = 6 splits into H = 3 heads of Dh = 2 columns, so heads has shape (1, 3, 2, 2). heads[0, 1] = [[2, 3], [8, 9]]; times its gate 0.5 it goes back into columns 2 to 3 of the merged output. Tap a gate to cycle 1, 0, 0.5, 2.
3 heads of width 2.
Examples
Two tokens, two heads; the second head is switched off
gate_heads(np.arange(8).reshape(1, 2, 4), 2, [1.0, 0.0])
returns ([[[[0, 1], [4, 5]], [[2, 3], [6, 7]]]], [[[0, 1, 0, 0], [4, 5, 0, 0]]])
Three heads of width 2
gate_heads(np.arange(12).reshape(1, 2, 6), 3, [1.0, 0.5, 2.0])
returns ([[[[0, 1], [6, 7]], [[2, 3], [8, 9]], [[4, 5], [10, 11]]]], [[[0, 1, 1, 1.5, 8, 10], [6, 7, 4, 4.5, 20, 22]]])
Requirements
xhas shape (B, T, D) and D is divisible byn_heads(H).gatesis a list of H floats.heads[b, h, t]isx[b, t, h * Dh:(h + 1) * Dh], withDh = D / H.- The merged array puts head h back in columns
h * Dhto(h + 1) * Dh - 1, after multiplying it bygates[h]. - Return a tuple
(heads, merged): float arrays of shape (B, H, T, Dh) and (B, T, D).
The idea behind it
Multi-head attention runs H small attentions side by side. A layer projects each token to a D-wide vector, then treats it as H heads of width Dh = D / H, each attending on its own. In code, that split is a reshape plus a transpose: (B, T, D) to (B, T, H, Dh) to (B, H, T, Dh), so each head sees its own (T, Dh) matrix and a batched matrix product handles all heads at once. After attention, the heads are merged by the inverse steps and passed through an output projection. A bare reshape from (B, T, D) to (B, H, T, Dh) runs without error but mixes tokens across heads, and because the matching bare reshape undoes it, a round trip test cannot catch it. Gating or zeroing single heads, as here, is a standard interpretability tool called head ablation.