Multi-head attention·10 min

Concat + project

0 / 1

# grade this stage (from your course directory)
npx theinference test --stage concat-project --watch

Run all 12 heads, concatenate their outputs back to [T, 768], then mix them with the output projection. That final step is what makes them a team — and it completes GPT-2's attention sublayer.

Concat 12 heads, then one projection

head0[T,64]  head1[T,64]  ...  head11[T,64]
      |__________ concat along columns __________|
                            |
                            v
                    concat [T,768]      (head h -> cols h*64 : (h+1)*64)
                            |
                            @ c_proj[768,768] + b
                            v
                       out [T,768]

What to build

Your task

Implement runMHA

Implement runMHA(x, layer): run each of the 12 heads (reuse runHead), concatenate to [T, 768] (head h → columns h·64 : (h+1)·64), then apply c_proj: concat @ c_proj.weight + c_proj.bias. Load h.{layer}.attn.c_proj.{weight,bias}.

tests/adapters.ts

export function runMHA(x: number[][], layer: number): number[][] {
  // x: [T, 768], layer: which block's weights  ->  [T, 768]
}

Concat then ONE projection — c_proj mixes the heads

The 12 heads run independently — until c_proj. That final [768, 768] matmul is what lets head 3's finding flow into head 7's channels; without it the heads never talk.

Returning the bare concatenation is a real bug, not a shortcut: the concat is not the attention sublayer's output. The sublayer output is concat @ c_proj.weight + c_proj.bias. (Same Conv1D [in, out] layout — x @ W, no transpose.)

How it's graded

Run on block 0; your [T, 768] output is compared with allclose, atol 1e-4.

References

Why multiple heads Pre-LN + residual stream