Transformer Blocks

Attention and the feed-forward network do different work. Attention lets each token row gather information from other positions. The feed-forward network then transforms the coordinates within each row, using the same weights at every position.

Why both? Attention compares a row with the rows it is allowed to read and combines their value vectors in proportion to the shares it computes, which sum to one across the positions read. In the live block the row for cat reads <start>, the and itself, and the first head gives those three positions shares of 0.3988, 0.3307 and 0.2705. What the row receives is a mixture in those proportions, and a mixture can only contain what went into it. Making something new out of the gathered content is the second job. The feed-forward network widens the row into a larger intermediate one, replaces every negative intermediate value with zero, and narrows the result back to the row’s width. The zeroing gives the network a changing slope, so what it answers is not simply another weighted sum of the row it read.

The same weights serve every position because the network’s question is the same wherever a row came from. Where a token sits is already written into its row by the position step, so the network reads the position out of the row rather than needing a separate set of weights for each place. That is also what lets one block read a longer prefix with the same parameters. The live example reads three tokens, and a fourth token would meet the same expansion and contraction matrices.

Separate heads give attention several independently learned comparison spaces. They are not assigned grammatical roles in advance. A head might learn a useful relationship, but its role must be investigated rather than inferred from its name.

Each branch produces a contribution that is added to its input. These residual additions preserve a direct route for information and gradients. The branch can learn an adjustment without needing to recreate the entire incoming representation.

-0.7264 + (-0.1005) = -0.8269
-0.8269 + (-0.1214) = -0.9483

Layer normalization standardizes coordinates within each token row. Our block normalizes before attention and before the feed-forward network. This is a pre-normalization arrangement; the original Transformer paper placed normalization after its residual additions.

Each addition can change the scale of a row, so a branch reads a standardized copy and meets values on the same scale whatever the earlier additions did. The copy is only what the branch reads. Its contribution is added to the un-normalized row, so the direct route is never rescaled inside the block, and the row is normalized once more only before the vocabulary projection in Part 3.

Let X contain the positioned token rows. LN means layer normalization, applied independently to each row. Attention is causal self-attention: each query can read its own position and earlier positions. The first addition produces H.

H = X + Attention(LN(X))

F = ReLU(LN(H) W₁ + b₁) W₂ + b₂
Y = H + F

The first feed-forward matrix expands the row into a wider intermediate representation. ReLU replaces negative intermediate values with zero. The second matrix returns to the model width so its contribution can be added to H. This implementation uses zero biases.

The SDK's attention projections store output units as matrix rows. The two feed-forward matrices shown here use input coordinates as rows, so they multiply on the right without a transpose. Their shapes, rather than their letters, determine the multiplication.

The live block uses four coordinates per token, two attention heads, and eight intermediate feed-forward coordinates. For one token row, the transformations have these shapes:

Input row:           (1, 4)
Expansion W₁:        (4, 8)
Intermediate row:   (1, 8)
Contraction W₂:     (8, 4)
Branch output:      (1, 4)
Residual output:    (1, 4)

The final row can enter another block with the same model width. A vocabulary projection after the final block turns these hidden coordinates into next-token scores.

The reference decoder combines embeddings, sinusoidal positions, one pre-normalized block, final row normalization, and a vocabulary projection. It exposes its intermediate arrays so the causal relationship can be checked.

import numpy as np
from oop_ml.numpy.modern import CausalLanguageModel

model = CausalLanguageModel(n_vocab=7, n_model=4, n_heads=2)
trace = model.respond_to(np.array([0, 1, 2]))
print(trace.attention)
print(trace.logits)

This object implements inference with fixed seeded parameters. It does not implement whole-decoder training, cross-attention, padding masks, or a key/value cache. The attention primitive itself already supports differentiation; composing a complete training pipeline is a further step.

Stacking blocks creates repeated opportunities to gather and transform information. Training determines which representations and comparisons are useful. Randomly initialized blocks have the structure of the computation but have not learned grammar, facts, or instruction following.

A decoder-only language model is one transformer arrangement. Encoders can use unmasked attention, and encoder-decoder models add cross-attention to another sequence. The causal decoder here is the arrangement we need for the next lessons.