The block around attention.

An attention layer computes useful mixtures, but a transformer diagram also contains boxes named Add, Norm, and Feed Forward. These operations are not decoration. They determine how information and gradients pass through the stack and how each position transforms its newly gathered context.

We will build the block used in our final model. Its input and output both have shape (B, T, D), which makes it possible to stack blocks without changing the representation width.

1. Add means a residual update

Let F be a learned transformation that returns the same shape as its input. A residual update is:

Every coordinate of F’s output is added to the matching coordinate of x. There is no concatenation and no new feature axis. The block can learn an update while the original representation has a direct path forward.

If F returns zero, the block returns x. During differentiation, a change in x has a direct path to the output through the addition, as well as a path through F. This direct contribution is one reason residual connections can help optimization. It does not guarantee that every deep network will train well; normalization, initialization, and the rest of the computation still matter. The foundations learning lesson explains how local derivatives carry a loss signal backward through a computation.

Residual addition is different from adding positional encodings. Both are elementwise sums, but one introduces an input signal and the other combines a representation with a learned update.

2. Layer normalization acts within a token

For one D-dimensional token vector x, compute its mean and population variance across features:

Layer normalization then computes:

Epsilon is a small positive constant that keeps the denominator well-defined. Gamma and beta are learned feature-wise scale and shift parameters. Every token uses its own mean and variance, while sharing the same learned gamma and beta for that normalization layer.

With (B, T, D) input and LayerNorm(D), the statistics do not mix batch items or token positions. Layer normalization therefore does not secretly let a causal position read later tokens through shared sequence-wide statistics. Its statistics are also computed from the current input during both training and inference; this is different from batch normalization’s usual running-statistics behavior. See Layer Normalization ↗ and PyTorch’s LayerNorm documentation ↗ .

For [1, 2, 3, 4], the mean is 2.5 and the variance is 1.25. Ignoring epsilon for a moment and taking gamma = 1, beta = 0 gives approximately [-1.342, -0.447, 0.447, 1.342]. With epsilon included, the normalized variance is slightly below one; learned gamma and beta can change the final mean and variance further.

3. The feed-forward network mixes features

The position-wise feed-forward network applies the same small neural network independently to every token:

The first projection expands D features to an inner width, such as 64 from a model width of 32. ReLU replaces negative intermediate values with zero. The second projection returns to D features so the result can be added to the residual stream.

Without the nonlinearity, the two affine transformations could be combined into one affine transformation. The activation allows a richer input-dependent transformation of each token’s feature vector.

Attention mixes information between positions. The FFN transforms features within a position, using the context that has already arrived there. Its weights are shared across positions in the same block, but different blocks generally have their own FFN parameters. The original transformer’s position-wise FFN is described in §3.3 of Attention Is All You Need ↗ .

def layer_norm(x: FloatArray, gamma: FloatArray | None = None, beta: FloatArray | None = None, eps: float = 1e-5) -> FloatArray:
    mean = x.mean(axis=-1, keepdims=True)
    variance = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
    normalized = (x - mean) / np.sqrt(variance + eps)
    gamma = np.ones(x.shape[-1]) if gamma is None else gamma
    beta = np.zeros(x.shape[-1]) if beta is None else beta
    return normalized * gamma + beta


def feed_forward(x: FloatArray, w1: FloatArray, b1: FloatArray, w2: FloatArray, b2: FloatArray) -> FloatArray:
    hidden = np.maximum(x @ w1 + b1, 0.0)  # ReLU, separately per token.
    return hidden @ w2 + b2               # Return to the residual width.

def layer_norm(x: Tensor, gamma: Tensor | None = None, beta: Tensor | None = None, eps: float = 1e-5) -> Tensor:
    # normalized_shape=(D,) means statistics over features of EACH token.
    return F.layer_norm(x, (x.shape[-1],), weight=gamma, bias=beta, eps=eps)


def feed_forward(x: Tensor, w1: Tensor, b1: Tensor, w2: Tensor, b2: Tensor) -> Tensor:
    hidden = torch.relu(x @ w1 + b1)
    return hidden @ w2 + b2

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

The NumPy function exposes the normalization arithmetic. The PyTorch function delegates to the differentiable library operation. In our trainable blocks below, nn.LayerNorm owns gamma and beta; they are initialized to one and zero, respectively, and can then change during training. The NumPy block is a forward demonstration using those default affine parameters.

4. The order of Add and Norm matters

Two common arrangements are:

The 2017 transformer diagram places normalization after each residual addition. Our teaching model uses pre-norm: normalize the input to a sublayer, compute its update, and add the update to the unnormalized residual stream. We also apply a final normalization at the end of each stack.

These formulas are not interchangeable. The location of normalization changes the forward function and the gradient path. Xiong et al. ↗ analyze this distinction and its implications for optimization. We choose one ordering consistently so our equations, diagrams, and executable model agree.

5. Assemble the two-sublayer block

Our encoder block computes:

The two normalization modules have separate learned parameters. MHA here means self-attention using the normalized states for all three projected inputs, subject to the source-padding mask.

class EncoderBlock:
    """A pre-norm forward pass. Normalization gain=1 and bias=0 here."""
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
        self.attn = MultiHead(width, heads, seed)
        self.ffn = ffn_parameters(width, inner, seed + 1)

    def __call__(self, x: FloatArray, source_visible: BoolArray) -> FloatArray:
        normalized = layer_norm(x)
        x = x + self.attn(normalized, normalized, source_visible)[0]
        return x + feed_forward(layer_norm(x), *self.ffn)

class EncoderBlock(nn.Module):
    def __init__(self, width: int = 32, heads: int = 4, inner: int = 64, seed: int = 10) -> None:
        super().__init__()
        self.attn = MultiHead(width, heads, seed)
        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.ffn = nn.Sequential(nn.Linear(width, inner), nn.ReLU(), nn.Linear(inner, width))

    def forward(self, x: Tensor, source_visible: Tensor) -> Tensor:
        normalized = self.norm1(x)
        x = x + self.attn(normalized, normalized, source_visible)[0]
        return x + self.ffn(self.norm2(x))

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

The complete files supply MultiHead, layer_norm, and the NumPy FFN parameters from earlier sections. PyTorch registers its layers as submodules, so model.parameters() will include all their learned weights.

6. Two experiments that isolate the roles

First, set the attention and FFN updates to zero. A pre-norm block’s residual route becomes an identity map. A final stack normalization, if applied after that block, can still change the representation. Keep those two boundaries separate when you test.

Second, hold all token vectors fixed except one, and apply only the FFN to each row. Only the changed row’s output can change. Perform the same intervention before unmasked self-attention, and other rows can change too. That is the difference between feature-wise transformation and communication across positions.

Dropout is another common component: during training it randomly removes selected contributions, usually with rescaling. Our tiny model omits dropout so repeated forward passes are deterministic and cross-library comparisons stay simple. Omitting it here is a teaching choice, not a claim that regularization is unnecessary for larger tasks.

Carry this forward

We now have a reusable contextual transformation. An encoder applies it with source visibility. A decoder needs causal target visibility and, when there is an encoder, another attention sublayer for source memory. The next lesson puts those pieces into complete architectures and explains when each is useful.