The Block

Passo 5 de 7. Entra: \(X\) da página de embeddings e a saída da atenção \(a\) da página de máscaras e cabeças, as duas \(5 \times 4\). Sai: \(y\), \(5 \times 4\) — a mesma forma de \(X\), então outro bloco poderia vir depois. ← Máscaras e cabeças · próximo passo: A saída →

Dentro de um bloco: conversar e depois pensar

Um bloco tem duas subcamadas e cada uma tem um papel bem definido:

• A atenção é a única parte em que os tokens trocam informação. É uma reunião: cada token pergunta, os outros respondem e ele sai com um resumo do que ouviu.
• O MLP (ou feed-forward) processa cada token sozinho, com os mesmos pesos para todas as posições. É o trabalho individual depois da reunião: pegar o que chegou e transformar em algo útil. É também onde fica a maior parte do "conhecimento" armazenado do modelo.

Nenhuma das duas substitui o vetor do token. Cada uma calcula uma correção que é somada a ele pela conexão residual:

LayerNorm: a mesma escala para toda linha

O \(\text{Norm}\) de um bloco do GPT-2 é a LayerNorm, aplicada a cada linha (cada token) separadamente:

em que \(\mu\) e \(\sigma^2\) são a média e a variância dos \(d\) números da linha, \(\epsilon = 10^{-5}\) só protege contra a divisão por zero e \(\gamma\), \(\beta\) são vetores aprendidos que começam em 1 e 0 (o brinquedo os mantém assim). Por que ela é necessária: cada bloco soma a sua correção ao fluxo; o tamanho de uma linha deriva de bloco para bloco e de token para token; uma subcamada cuja entrada muda de escala precisa reaprender o tempo todo o que é um número "grande". Cada subcamada lê uma cópia posta numa escala comum; o fluxo em si mantém o seu tamanho. No GPT-2 small cada linha tem 768 números; no brinquedo, 4.

O exemplo atravessando o bloco

O bloco parte de \(x = X\) (as cinco linhas da página de embeddings) e roda as duas equações acima — com os pesos do brinquedo, cada número abaixo pode ser conferido à mão.

1. \(u = \text{LN}(x)\) e \(a = \text{Attn}(u)\) — o capítulo de atenção fez esses dois: \(u\) na página de Q, K, V e \(a\) na página de máscaras e cabeças. Para o último token, ␣no: \(a_5 = [-0{,}739; 1{,}411; -0{,}521; -0{,}520]\).

2. O primeiro residual, \(h = x + a\). Para ␣no:

   \[ \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} \]

   É o único lugar do bloco em que um token recebe números dos outros.

3. A segunda LayerNorm, em \(h\). Média \(\mu = (-1{,}239 + 0{,}911 - 1{,}021 + 0{,}980)/4 = -0{,}092\), desvio-padrão \(\sigma = 1{,}041\), então \(u' = (h - \mu)/\sigma = [-1{,}102; 0{,}964; -0{,}892; 1{,}030]\).

4. O MLP, cada token sozinho. \(W_1\) é \(4 \times 8\) e \(W_2\) é \(8 \times 4\) (no GPT-2: 768 → 3072 → 768). O brinquedo usa \(W_1 = [\,I \mid -I\,]\): as 8 unidades ocultas recebem \(u'\) e \(-u'\). A GELU, \(\text{gelu}(z) = \tfrac{1}{2} z \left(1 + \tanh\!\big(0{,}798\,(z + 0{,}0447\,z^3)\big)\right)\), deixa passar os positivos e espreme os negativos para perto de 0. Para o primeiro número de ␣no: \(\text{gelu}(-1{,}102) = -0{,}149\) e \(\text{gelu}(1{,}102) = 0{,}952\). Depois \(W_2 = [\,I \,;\, I\,]\) soma as duas metades, \(m = \text{gelu}(u') + \text{gelu}(-u')\) — um \(|u'|\) suave: \(m_5 = [0{,}803; 0{,}640; 0{,}560; 0{,}718]\).

5. O segundo residual, \(y = h + m\). Para ␣no:

   \[ \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} \]

As cinco linhas:

\(y\) é \(5 \times 4\), como \(x\). A página da saída o lê. Com o ramo de atenção desligado, \(h = x\) e a última linha seria \(y_5 = [-0{,}248; -0{,}248; -0{,}248; 3{,}088]\): feita de ␣no na posição 5 e de mais nada — a página da saída mostra quanto isso custa.

O script abaixo passa uma frase de 5 tokens por um bloco completo do tipo moderno (RMSNorm, GQA e SwiGLU, que o que mudou desde 2017 explica), pequeno o bastante para ler (\(d = 8\)), imprime a forma em cada etapa e depois conta os parâmetros do mesmo bloco na largura de um modelo real:

"""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

O fluxo residual entra \(5 \times 8\) e sai \(5 \times 8\); só as etapas internas mudam de forma. A matriz de atenção é \(5 \times 5\) por cabeça (cada token contra cada token) e a camada oculta do MLP é mais larga que \(d\). Na largura real, 32 blocos somam 5,7B de parâmetros; com a tabela de embedding e a camada de saída (\(2 \times 128\text{k} \times 4096 \approx 1\text{B}\)) chega-se perto de 7B. O Llama 3 8B usa um MLP ainda mais largo (14336 em vez de 11008), o que fecha os 8B.

As formas dizem para onde vai cada número; o painel abaixo mostra os números. É a bloco_a_mao do laboratório (as mesmas linhas, com a LayerNorm e a GELU do GPT-2) reduzida a uma head, 3 tokens e \(d = 4\), com pesos simples o bastante para conferir cada entrada à mão. Percorra linha a linha e depois teste os casos: uma linha multiplicada por 10, uma mudança só no último token e a atenção removida (a pergunta logo abaixo).

← Máscaras e cabeças · próximo passo: A saída →