Phase 07 - Lesson 02

Self-Attention do Zero

This lesson includes a graded coding exercise that runs in your browser, unlocked with lifetime access.

Atenção é uma tabela de consulta onde cada palavra pergunta "quem importa para mim?" - e aprende a resposta.

Tipo: Build Linguagens: Python Pré-requisitos: Fase 3 (Núcleo de Deep Learning), Fase 5 Lição 10 (Sequência-para-Sequência) Tempo: ~90 minutos

Objetivos de Aprendizado

• Implementar self-attention com produto escalar escalonado do zero usando apenas NumPy, incluindo projeções de query/key/value e a soma ponderada por softmax
• Construir uma camada de multi-head attention que divide as cabeças, computa atenção em paralelo e concatena os resultados
• Rastrear como a matriz de atenção captura relações entre tokens e explicar por que escalonar por sqrt(d_k) evita a saturação do softmax
• Aplicar mascaramento causal para converter atenção bidirecional em atenção autorregressiva (estilo decoder)

O Problema

RNNs processam sequências um token por vez. Quando você chega ao token 50, a informação do token 1 já foi espremida por 50 etapas de compressão. Dependências de longo alcance acabam esmagadas em um estado oculto de tamanho fixo - um gargalo que nenhuma quantidade de gating de LSTM resolve por completo.

O artigo de atenção de Bahdanau de 2014 mostrou a solução: deixar o decoder olhar de volta para cada posição do encoder e decidir quais delas importam para o passo atual. Mas ainda estava acoplado a uma RNN. O artigo de 2017 "Attention Is All You Need" fez uma pergunta mais incisiva: e se a atenção for o único mecanismo? Sem recorrência. Sem convolução. Só atenção.

Self-attention permite que cada posição de uma sequência atenda a todas as outras posições em um único passo paralelo. É isso que torna os transformers rápidos, escaláveis e dominantes.

O Conceito

A Analogia da Consulta a Banco de Dados

Pense na atenção como uma consulta suave a um banco de dados:

Traditional database:
  Query: "capital of France"  -->  exact match  -->  "Paris"

Attention:
  Query: "capital of France"  -->  similarity to ALL keys  -->  weighted blend of ALL values

Cada token gera três vetores:

• Query (Q): "O que estou procurando?"
• Key (K): "O que eu contenho?"
• Value (V): "Que informação eu forneço se for selecionado?"

O produto escalar entre uma query e todas as keys produz os scores de atenção. Um score alto significa "esta key combina com minha query". Esses scores ponderam os values. A saída é uma soma ponderada dos values.

Computação de Q, K, V

Cada embedding de token é projetado através de três matrizes de pesos aprendidas:

Input embeddings (sequence of n tokens, each d-dimensional):

  X = [x1, x2, x3, ..., xn]       shape: (n, d)

Three weight matrices:

  Wq  shape: (d, dk)
  Wk  shape: (d, dk)
  Wv  shape: (d, dv)

Projections:

  Q = X @ Wq    shape: (n, dk)      each token's query
  K = X @ Wk    shape: (n, dk)      each token's key
  V = X @ Wv    shape: (n, dv)      each token's value

Visualmente, para um token:

             Wq
  x_i ------[*]------> q_i    "What am I looking for?"
       |
       |     Wk
       +----[*]------> k_i    "What do I contain?"
       |
       |     Wv
       +----[*]------> v_i    "What do I offer?"

A Matriz de Atenção

Uma vez que você tem Q, K, V para todos os tokens, os scores de atenção formam uma matriz:

Scores = Q @ K^T    shape: (n, n)

              k1    k2    k3    k4    k5
        +-----+-----+-----+-----+-----+
   q1   | 2.1 | 0.3 | 0.1 | 0.8 | 0.2 |   <- how much q1 attends to each key
        +-----+-----+-----+-----+-----+
   q2   | 0.4 | 1.9 | 0.7 | 0.1 | 0.3 |
        +-----+-----+-----+-----+-----+
   q3   | 0.2 | 0.6 | 2.3 | 0.5 | 0.1 |
        +-----+-----+-----+-----+-----+
   q4   | 0.9 | 0.1 | 0.4 | 1.7 | 0.6 |
        +-----+-----+-----+-----+-----+
   q5   | 0.1 | 0.3 | 0.2 | 0.5 | 2.0 |
        +-----+-----+-----+-----+-----+

Each row: one token's attention over the entire sequence

Por Que Escalonar?

Os produtos escalares crescem com a dimensão dk. Se dk = 64, os produtos escalares podem estar na faixa das dezenas, empurrando o softmax para regiões onde os gradientes desaparecem. A solução: dividir por sqrt(dk).

Scaled scores = (Q @ K^T) / sqrt(dk)

Isso mantém os valores em uma faixa onde o softmax produz gradientes úteis.

Softmax Transforma Scores em Pesos

O softmax converte scores brutos em uma distribuição de probabilidade ao longo de cada linha:

Raw scores for q1:   [2.1, 0.3, 0.1, 0.8, 0.2]
                            |
                         softmax
                            |
Attention weights:   [0.52, 0.09, 0.07, 0.14, 0.08]   (sums to ~1.0)

Agora cada token tem um conjunto de pesos dizendo o quanto atender a cada outro token.

Soma Ponderada dos Values

A saída final para cada token é uma soma ponderada de todos os vetores de value:

output_i = sum( attention_weight[i][j] * v_j  for all j )

For token 1:
  output_1 = 0.52 * v1 + 0.09 * v2 + 0.07 * v3 + 0.14 * v4 + 0.08 * v5

Pipeline Completo

                    +-------+
  X (input)  ----->|  @ Wq  |-----> Q
                    +-------+
                    +-------+
  X (input)  ----->|  @ Wk  |-----> K
                    +-------+                     +----------+
                    +-------+                     |          |
  X (input)  ----->|  @ Wv  |-----> V ---------->| weighted |----> output
                    +-------+          ^          |   sum    |
                                       |          +----------+
                              +--------+--------+
                              |    softmax      |
                              +---------+-------+
                                        ^
                              +---------+-------+
                              | Q @ K^T / sqrt  |
                              +-----------------+

Fórmula em uma linha:

Attention(Q, K, V) = softmax( Q @ K^T / sqrt(dk) ) @ V

Construa

Passo 1: Softmax do zero

O softmax converte logits brutos em probabilidades. Subtraia o máximo para estabilidade numérica.

import numpy as np

def softmax(x):
    shifted = x - np.max(x, axis=-1, keepdims=True)
    exp_x = np.exp(shifted)
    return exp_x / np.sum(exp_x, axis=-1, keepdims=True)

logits = np.array([2.0, 1.0, 0.1])
print(f"logits:  {logits}")
print(f"softmax: {softmax(logits)}")
print(f"sum:     {softmax(logits).sum():.4f}")

Passo 2: Atenção com produto escalar escalonado

A função central. Recebe as matrizes Q, K, V e retorna a saída da atenção mais a matriz de pesos.

def scaled_dot_product_attention(Q, K, V):
    dk = Q.shape[-1]
    scores = Q @ K.T / np.sqrt(dk)
    weights = softmax(scores)
    output = weights @ V
    return output, weights

Passo 3: Classe de self-attention com projeções aprendidas

Um módulo completo de self-attention com matrizes de pesos Wq, Wk, Wv inicializadas com escalonamento estilo Xavier.

class SelfAttention:
    def __init__(self, d_model, dk, dv, seed=42):
        rng = np.random.default_rng(seed)
        scale = np.sqrt(2.0 / (d_model + dk))
        self.Wq = rng.normal(0, scale, (d_model, dk))
        self.Wk = rng.normal(0, scale, (d_model, dk))
        scale_v = np.sqrt(2.0 / (d_model + dv))
        self.Wv = rng.normal(0, scale_v, (d_model, dv))
        self.dk = dk

    def forward(self, X):
        Q = X @ self.Wq
        K = X @ self.Wk
        V = X @ self.Wv
        output, weights = scaled_dot_product_attention(Q, K, V)
        return output, weights

Passo 4: Rode em uma frase

Crie embeddings fictícios para uma frase e observe os pesos de atenção.

sentence = ["The", "cat", "sat", "on", "the", "mat"]
n_tokens = len(sentence)
d_model = 8
dk = 4
dv = 4

rng = np.random.default_rng(42)
X = rng.normal(0, 1, (n_tokens, d_model))

attn = SelfAttention(d_model, dk, dv, seed=42)
output, weights = attn.forward(X)

print("Attention weights (each row: where that token looks):\n")
print(f"{'':>6}", end="")
for token in sentence:
    print(f"{token:>6}", end="")
print()

for i, token in enumerate(sentence):
    print(f"{token:>6}", end="")
    for j in range(n_tokens):
        w = weights[i][j]
        print(f"{w:6.3f}", end="")
    print()

Passo 5: Visualize a atenção com heatmap ASCII

Mapeie os pesos de atenção para caracteres para uma visualização rápida.

def ascii_heatmap(weights, tokens, chars=" ░▒▓█"):
    n = len(tokens)
    print(f"\n{'':>6}", end="")
    for t in tokens:
        print(f"{t:>6}", end="")
    print()

    for i in range(n):
        print(f"{tokens[i]:>6}", end="")
        for j in range(n):
            level = int(weights[i][j] * (len(chars) - 1) / weights.max())
            level = min(level, len(chars) - 1)
            print(f"{'  ' + chars[level] + '   '}", end="")
        print()

ascii_heatmap(weights, sentence)

Use

O nn.MultiheadAttention do PyTorch faz exatamente o que construímos, mais a divisão em múltiplas cabeças e a projeção de saída:

import torch
import torch.nn as nn

d_model = 8
n_heads = 2
seq_len = 6

mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=n_heads, batch_first=True)

X_torch = torch.randn(1, seq_len, d_model)

output, attn_weights = mha(X_torch, X_torch, X_torch)

print(f"Input shape:            {X_torch.shape}")
print(f"Output shape:           {output.shape}")
print(f"Attention weight shape: {attn_weights.shape}")
print(f"\nAttn weights (averaged over heads):")
print(attn_weights[0].detach().numpy().round(3))

A diferença principal: a multi-head attention roda múltiplas funções de atenção em paralelo, cada uma com suas próprias projeções Q, K, V de tamanho dk = d_model / n_heads, e depois concatena os resultados. Isso permite que o modelo atenda a diferentes tipos de relação simultaneamente.

Entregue

Esta lição produz:

• outputs/prompt-attention-explainer.md - um prompt para explicar a atenção através da analogia da consulta a banco de dados

Exercícios

1. Modifique scaled_dot_product_attention para aceitar uma matriz de máscara opcional que define certas posições como infinito negativo antes do softmax (é assim que funciona o mascaramento causal/de decoder)
2. Implemente a multi-head attention do zero: divida Q, K, V em n_heads blocos, rode atenção em cada um, concatene e projete através de uma matriz de pesos final Wo
3. Pegue duas frases diferentes de mesmo comprimento, passe-as pela mesma instância de SelfAttention e compare seus padrões de atenção. O que muda? O que permanece igual?

Termos-Chave

Leitura Adicional

• Attention Is All You Need (Vaswani et al., 2017) - o artigo original do transformer
• The Illustrated Transformer (Jay Alammar) - o melhor passo a passo visual da arquitetura completa
• The Annotated Transformer (Harvard NLP) - implementação linha a linha em PyTorch com explicações