Compose residual attention and normalization

Learn it

Words to know

transformer block

One layer of a transformer: attention, then a small feed-forward network, each with a residual connection and normalization.

residual connection

Adding a step's input to its output, so the step only has to learn a change.

normalization

Rescaling a group of numbers to a standard size, which keeps training stable.

causal mask

The rule that a token may attend only to itself and earlier tokens, never to later ones.

How it works

1. Attention: every token attends to itself and earlier tokens. Here the input X is used as queries, keys and values at once.
2. Residual: add the attention output to X.
3. Normalize each token's row with layer normalization. Its epsilon is passed as eps=1e-5. Call the result h.
4. Feed-forward with a residual: return h + ReLU(h). A real block has a learned MLP here; this exercise uses a bare ReLU so the numbers can be followed by hand.

Worked example

Input
X = [[1, 0, -1]]

Output
[[2.44949, 0, -1.22474]]

The two additions keep what came before instead of replacing it. Even where ReLU wiped a value out, the residual keeps the original -1.2247.

Why it matters

A GPT is token and position embeddings, a stack of these blocks, and one linear layer on top. You assemble exactly that in the next chapter.

Your turn

Using PyTorch, run one simplified transformer block with no learned weights.

X

a table of shape [time, width]: one row per token

Returns

The nested list h + ReLU(h), where h = layer_norm(X + causal_attention(X, X, X)) with eps=1e-5.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
X = [[1, 0, -1]]

Output
[[2.44949, 0, -1.22474]]

Input
X = [[1, 1]]

Output
[[0, 0]]

Input
X = [[1, -1]]

Output
[[2, -0.999999]]

Input
X = [[1, 0, -1], [0, 2, 0]]

Output
[[2.44949, 0, -1.22474], [-0.656374, 2.82605, -0.756652]]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Three library calls do the work: attention, layer normalization and ReLU. Your job is to connect them with two additions.

2. Hint 2: The key step

   Call scaled_dot_product_attention(x, x, x, is_causal=True), add x to the result, then normalize over the last dimension with layer_norm.

3. Hint 3: The outline

   Make x from X as a float64 tensor. With F = torch.nn.functional: a = F.scaled_dot_product_attention(x, x, x, is_causal=True); h = F.layer_norm(x + a, (x.shape[-1],), eps=1e-5); return (h + torch.relu(h)).tolist().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. Each token attends to itself and the tokens before it. is_causal=True applies the causal mask for you. Passing x three times uses each row as its own query, key and value.
2. Residual: add the input back, then normalize each row. The second argument of layer_norm is the shape to normalize over: the width of one row, written as a one-item tuple.
3. Feed-forward step with its own residual. h passes through unchanged and the ReLU output is added on top, so nothing learned earlier is lost.

The same numbered steps are the comments in the solution.

import torch


def solve(X):
    F = torch.nn.functional
    x = torch.tensor(X, dtype=torch.float64)
    # 1. Each token attends to itself and the tokens before it
    attended = F.scaled_dot_product_attention(x, x, x, is_causal=True)
    # 2. Residual: add the input back, then normalize each row
    h = F.layer_norm(x + attended, (x.shape[-1],), eps=1e-5)
    # 3. Feed-forward step with its own residual
    return (h + torch.relu(h)).tolist()

Watch out

• (x.shape[-1],) needs the trailing comma. Without it Python sees a plain number in brackets, and layer_norm raises TypeError because it wants a tuple.
• Add the input before normalizing, not after. The order here is attention, add, normalize.

Check yourself