Transformer Blocks: Preserve Width While Mixing Context and Features

Project before splitting

For D=4, H=2, each head has width 2.

Q/K/V projections inspect all four input features, then split their output columns.

Splitting raw inputs first restricts the attention mechanism.

1 / 8

Prove the distinction

Input [1,2,3,4]; first query row weight [0,0,1,0].

Head 0 can receive query feature 3.

If it only saw raw input [1,2], this projection would be impossible.

2 / 8

Track head shapes

T=3: Q/K/V (3,4) -> two sets (3,2).

Each head produces (3,3) scores scaled by sqrt(2), then (3,2) output.

Concatenate to (3,4), apply one output projection.

3 / 8

Normalize one position

For [1,3]: mean 2, population variance 1.

With eps=1e-5, output before affine scaling is about [-0.999995,0.999995].

Statistics must stay in the graph; normalization across positions can leak the future.

4 / 8

Reconstruct the residuals

y = x + MHA(LN1(x))
z = y + FFN(LN2(y))

Normalize each position. Add elementwise. Preserve the unnormalized residual input.

5 / 8

Separate two kinds of mixing

Attention mixes information across visible positions.

FFN D -> 4D -> D with ReLU mixes features independently per position.

Both preserve the width needed for residual addition.

6 / 8

Count a whole block

At D=4: attention 80, FFN 148, two norms 16. Total 244 unique scalars.

Predict what a different head count changes: score partitions, not total projection width.

7 / 8

Use behavioral checks

Zero both branches: the block must return x exactly.

Change a future row: earlier outputs stay equal.

Finite-difference LayerNorm inputs: detached statistics must fail. Then stack blocks.

8 / 8
←/→ or click edges to navigate · ? help · N notes · F fullscreen