Transformer Blocks: Preserve Width While Mixing Context and Features

Single-head attention produces one distribution over visible sources at each position. Multiple heads produce several distributions using different learned feature projections. A transformer block combines this context mixing with a position-wise nonlinear network, while residual paths preserve the original representation.

You will implement MultiHeadAttention, LayerNorm, FeedForward, and TransformerBlock. The essential contracts are which features each head can see, which axis is normalized, and which tensor a residual adds back.

Project the full representation before splitting heads

Let embedding width be D and head count H. Require positive D and H with D % H == 0; each head has width d_h = D/H.

Construct full-width Q, K, and V with three Linear(D,D) projections. Split each projected row into H slices of width d_h. Equivalently, each head could use its own Linear(D,d_h) projections. In both forms, every head's projections can inspect every input feature. Splitting raw input features first would impose a different, restricted architecture.

For T=3, D=4, H=2:

x                     (3,4)
Q, K, V               (3,4) each
head 0 slices         columns [0:2], shape (3,2)
head 1 slices         columns [2:4], shape (3,2)
each head's scores    (3,3), scaled by sqrt(2)
each head's output    (3,2)
concatenated output   (3,4)
final W_o projection  (3,4)

Concatenate by position, then by head. Averaging the two head vectors would leave width 2 and discard their separate feature slots. Use one final Linear(D,D) projection; do not add an extra output projection inside every head. This projection/split/concatenate structure follows multi-head attention.

See the projection difference with one vector

For input [1,2,3,4], suppose the first projected query row uses weight [0,0,1,0]. The first feature of head 0's query is 3, even though it depends on an input feature in the second half of the vector.

A raw-input split would give head 0 only [1,2], making that projection impossible. An output projection after attention cannot generally repair the missing ability to choose attention weights using all input features. This is why matching the final shape alone is too weak a test.

With biases enabled on all four full-width projections, attention owns 4*(D*D+D) parameters. At D=4 this is 80, regardless of the valid number of heads. More heads change how projected features form score matrices, not the overall projection width.

Normalize features separately at each position

For one length-D vector x:

mu = sum(x)/D
variance = sum((x[k]-mu)**2)/D
normalized[k] = (x[k]-mu)/sqrt(variance+eps)
output[k] = gamma[k]*normalized[k] + beta[k]

Initialize gamma to ones and beta to zeros. Both are trainable vectors, adding 2D parameters. Use population variance, dividing by D rather than D-1; this also works for D=1. Epsilon is inside the square root. These choices match the LayerNorm definition.

For [1,3], mean is 2 and variance is 1. With eps=1e-5, normalized output is approximately [-0.999995,0.999995]. For a constant vector [3,3], the centered values are zero and output is beta, with no division by zero.

Keep mean and variance in the autograd graph. Computing them from .data gives the right forward values but wrong derivatives. Do not normalize across the T positions: that would allow a future position to affect an earlier one even if attention masking were correct.

Compose two residual updates

The block uses pre-normalization:

n1 = LayerNorm1(x), independently for each position
a  = causal MultiHeadAttention(n1)
y  = x + a
n2 = LayerNorm2(y), independently for each position
f  = FeedForward(n2), independently for each position
z  = y + f

Addition here is elementwise, not Python list concatenation. Save the unnormalized input for each residual. In particular, the second branch receives y, the representation after the attention residual.

The feed-forward network is Linear(D,4D) -> ReLU -> Linear(4D,D) and shares weights across all positions. It mixes features within one position. Attention is the part that moves information between positions.

Pre-norm leaves an identity path around each sublayer. Research on normalization placement studies its effect on gradient behavior and training; it is a design choice here, not a guarantee that arbitrary learning rates or depths will be stable.

Derive checks stronger than shape

For D=4, the default FFN has 4*16+16 + 16*4+4 = 148 parameters. Two normalizations add 16. Including the 80 attention parameters, the block has 244 unique parameters. This count catches missing biases, extra head projections, and duplicate registration.

If every parameter in both attention and feed-forward branches is zero, both branch outputs are zero and the entire block must return its input exactly. The normalization parameters do not remove that identity path. If the output changes, inspect the residual source or the placement of normalization.

Differentiate one LayerNorm output with respect to the input and compare with a finite difference; a forward-only check will not catch detached statistics. Change a future input row and confirm the whole block preserves its earlier outputs.

Checkpoint. For D=12 and H=3, what is the scaling denominator and concatenated width?

Check your reasoning

sqrt(4)=2 and width 12. Scaling by sqrt(12) would change every head's distribution.

Now the block can be stacked without changing width or sequence length. The capstone adds the token interfaces, output vocabulary projection, and next-token objective.