The Block

Step 5 of 7. In: \(X\) from the embeddings page and the attention output \(a\) from the masks and heads page, both \(5 \times 4\). Out: \(y\), \(5 \times 4\) — the same shape as \(X\), so another block could follow. ← Masks and heads · next step: Output →

Inside a block: talk, then think

A block has two sublayers, and each has a well-defined job:

• Attention is the only part where tokens exchange information. It is a meeting: each token asks, the others answer, and it leaves with a summary of what it heard.
• The MLP (or feed-forward) processes each token alone, with the same weights at every position. It is the individual work after the meeting: take what arrived and turn it into something useful. It is also where most of the model's stored "knowledge" lives.

Neither of them replaces the token's vector. Each computes a correction that the residual connection adds to it:

LayerNorm: the same scale for every row

\(\text{Norm}\) in a GPT-2 block is LayerNorm, applied to each row (each token) separately:

where \(\mu\) and \(\sigma^2\) are the mean and the variance of the row's \(d\) numbers, \(\epsilon = 10^{-5}\) only guards against dividing by zero, and \(\gamma\), \(\beta\) are learned vectors that start at 1 and 0 (the toy keeps them there). Why it is needed: every block adds its correction to the stream, so the size of a row drifts from block to block and from token to token, and a sublayer whose input changes scale has to keep re-learning what a "large" number is. Each sublayer reads a copy put on a common scale; the stream itself keeps its size. In GPT-2 small each row has 768 numbers; in the toy, 4.

The running example through the block

The block starts from \(x = X\), the five rows of the embeddings page, and runs the two equations above — with the toy's weights, every number below can be checked by hand.

1. \(u = \text{LN}(x)\) and \(a = \text{Attn}(u)\) — the attention chapter did these: \(u\) on the Q, K, V page, \(a\) on the masks and heads page. For the last token, ␣the: \(a_5 = [-0.739, 1.411, -0.521, -0.520]\).

2. The first residual, \(h = x + a\). For ␣the:

   \[ \begin{aligned} h_5 &= [-0.5, -0.5, -0.5, 1.5] + [-0.739, 1.411, -0.521, -0.520] \\ &= [-1.239, 0.911, -1.021, 0.980] \end{aligned} \]

   This is the only place in the block where a token receives numbers from the others.

3. The second LayerNorm, on \(h\). Mean \(\mu = (-1.239 + 0.911 - 1.021 + 0.980)/4 = -0.092\), standard deviation \(\sigma = 1.041\), so \(u' = (h - \mu)/\sigma = [-1.102, 0.964, -0.892, 1.030]\).

4. The MLP, each token alone. \(W_1\) is \(4 \times 8\) and \(W_2\) is \(8 \times 4\) (GPT-2: 768 → 3072 → 768). The toy uses \(W_1 = [\,I \mid -I\,]\), so the 8 hidden units receive \(u'\) and \(-u'\), and GELU, \(\text{gelu}(z) = \tfrac{1}{2} z \left(1 + \tanh\!\big(0.798\,(z + 0.0447\,z^3)\big)\right)\), lets positive numbers through and squashes negative ones to near 0. For the first number of ␣the: \(\text{gelu}(-1.102) = -0.149\) and \(\text{gelu}(1.102) = 0.952\). Then \(W_2 = [\,I \,;\, I\,]\) adds the two halves, \(m = \text{gelu}(u') + \text{gelu}(-u')\) — a smooth \(|u'|\): \(m_5 = [0.803, 0.640, 0.560, 0.718]\).

5. The second residual, \(y = h + m\). For ␣the:

   \[ \begin{aligned} y_5 &= [-1.239, 0.911, -1.021, 0.980] + [0.803, 0.640, 0.560, 0.718] \\ &= [-0.436, 1.551, -0.461, 1.698] \end{aligned} \]

All five rows:

\(y\) is \(5 \times 4\), like \(x\). The output page reads it. With the attention branch switched off, \(h = x\) and the last row would be \(y_5 = [-0.248, -0.248, -0.248, 3.088]\): built from ␣the at position 5 and nothing else — the output page shows what that costs.

The script below pushes a 5-token sentence through a complete block of the modern kind — RMSNorm, GQA and SwiGLU, which what changed since 2017 explains — small enough to read (\(d = 8\)), prints the shape at every step, and then counts the parameters of the same block at the width of a real model (its labels are in Portuguese; the shapes and numbers are what matter):

"""Um bloco pré-LN inteiro em numpy: as formas em cada etapa e onde os parâmetros moram.

Primeiro uma frase de 5 tokens atravessa um bloco pequeno (d = 8) e cada etapa imprime sua forma —
o fluxo residual entra (n, d) e sai (n, d). Depois a mesma conta de parâmetros para d = 4096, a
largura de um modelo classe 8B, com GQA e SwiGLU.
"""
import numpy as np

rng = np.random.default_rng(0)

def rmsnorm(x, eps=1e-6):
    return x / np.sqrt((x ** 2).mean(-1, keepdims=True) + eps)

def softmax(s):
    e = np.exp(s - s.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

def silu(x):
    return x / (1 + np.exp(-x))

def swiglu_width(d):
    return round(8 * d / 3 / 256) * 256 if d >= 256 else 8 * d // 3

def params(d, h, h_kv):
    dh = d // h
    attn = d * d + 2 * d * h_kv * dh + d * d            # W_Q, W_K e W_V (GQA), W_O
    ffn = 3 * d * swiglu_width(d)                       # W_1, W_3, W_2
    return attn, ffn, 2 * d                             # duas RMSNorm, um ganho cada

n, d, h, h_kv = 5, 8, 2, 1
dh, f = d // h, swiglu_width(d)
Wq, Wk, Wv, Wo = (rng.normal(size=s) / np.sqrt(d) for s in [(d, d), (d, h_kv * dh), (d, h_kv * dh), (d, d)])
W1, W3, W2 = rng.normal(size=(d, f)) / np.sqrt(d), rng.normal(size=(d, f)) / np.sqrt(d), rng.normal(size=(f, d)) / np.sqrt(f)

def show(nome, t):
    print(f"  {nome:<34}{str(t.shape):>10}")

x = rng.normal(size=(n, d))
print(f"n = {n} tokens, d = {d}, {h} heads de query, {h_kv} head de K/V")
show("x (fluxo residual, entrada)", x)

u = rmsnorm(x)
q = (u @ Wq).reshape(n, h, dh).transpose(1, 0, 2)
k = (u @ Wk).reshape(n, h_kv, dh).transpose(1, 0, 2).repeat(h // h_kv, axis=0)   # GQA: K/V repetidos
v = (u @ Wv).reshape(n, h_kv, dh).transpose(1, 0, 2).repeat(h // h_kv, axis=0)
show("Q, por head", q)
show("K, V (1 head, repetido)", k)
causal = np.triu(np.full((n, n), -np.inf), 1)
A = softmax(q @ k.transpose(0, 2, 1) / np.sqrt(dh) + causal)
show("A = softmax(QKᵀ/√d_h + M)", A)
o = (A @ v).transpose(1, 0, 2).reshape(n, d) @ Wo
x = x + o
show("x + atenção", x)

u = rmsnorm(x)
g = silu(u @ W1) * (u @ W3)
show("SwiGLU, camada oculta", g)
x = x + g @ W2
show("x + MLP (saída do bloco)", x)

print("\nd = 4096, 32 heads de query, 8 de K/V:")
attn, ffn, norm = params(4096, 32, 8)
total = attn + ffn + norm
for nome, p in [("atenção", attn), ("MLP (SwiGLU)", ffn), ("normalização", norm)]:
    print(f"  {nome:<14}{p / 1e6:8.1f} M  {100 * p / total:4.0f}%")
print(f"  {'bloco':<14}{total / 1e6:8.1f} M")
print(f"  × 32 blocos  {32 * total / 1e9:8.2f} B")

n = 5 tokens, d = 8, 2 heads de query, 1 head de K/V
  x (fluxo residual, entrada)           (5, 8)
  Q, por head                        (2, 5, 4)
  K, V (1 head, repetido)            (2, 5, 4)
  A = softmax(QKᵀ/√d_h + M)          (2, 5, 5)
  x + atenção                           (5, 8)
  SwiGLU, camada oculta                (5, 21)
  x + MLP (saída do bloco)              (5, 8)

d = 4096, 32 heads de query, 8 de K/V:
  atenção           41.9 M    24%
  MLP (SwiGLU)     135.3 M    76%
  normalização       0.0 M     0%
  bloco            177.2 M
  × 32 blocos      5.67 B

The residual stream enters \(5 \times 8\) and leaves \(5 \times 8\); only the internal steps change shape. The attention matrix is \(5 \times 5\) per head (every token against every token), and the MLP's hidden layer is wider than \(d\). At real width, 32 blocks come to 5.7B parameters; with the embedding table and the output layer (\(2 \times 128\text{k} \times 4096 \approx 1\text{B}\)) that gets close to 7B. Llama 3 8B uses an even wider MLP (14336 instead of 11008), which makes up the 8B.

Shapes say where each number goes; the panel below shows the numbers. It is the lab's block_by_hand — the same lines, GPT-2's LayerNorm and GELU — shrunk to one head, 3 tokens and \(d = 4\), with weights simple enough to check every entry by hand. Step through it, then try the cases: a row scaled by 10, a change to the last token only, and attention removed (the question just below).

← Masks and heads · next step: Output →