Demonstrações interativas da aula de self-attention e transformers
Oito visualizações que acompanham os slides, em duas metades. Primeiro, o mecanismo de self-attention: o pipeline de queries, keys e values passo a passo numa frase curta, a matriz de atenção softmax(QK⊤/√d) como heatmap, por que escalar os scores por 1/√d evita a saturação da softmax, e a saída da atenção como média ponderada dos values. Depois, como o bloco Transformer envolve esse mecanismo: a atenção scaled dividindo os scores por √d, a self-attention multi-head com a matriz de cada cabeça em heatmap, o positional encoding sinusoidal de Vaswani et al. (2017) e o bloco Add & Norm (residual + LayerNorm) que mantém o sinal estável em profundidade.
Cada token xt ∈ ℝ³ da frase "o gato caça o rato" é projetado em três vetores por matrizes aprendidas: a query qt = xtWq, a key kt = xtWk e o value vt = xtWv (aqui d = 2, então cada projeção tem 2 componentes, mostradas como barras). Escolha o token de consulta t nas abas e use ⏭ para avançar pelas etapas vistas em aula: embeddings, projeções, scores qt⊤ki/√d, softmax e soma ponderada dos values. Note que o token repetido "o" tem o mesmo embedding nas duas posições, logo as mesmas projeções.
Embeddings xi e projeções qi, ki, vi de cada token (a coluna destacada é o token de consulta t)
Scores, pesos ai(t) e saída c(t) para a query escolhida
💡 Avance até a etapa da softmax e troque o token de consulta nas abas: os pesos ai(t) mudam, mas as keys e os values continuam os mesmos, cada query faz uma consulta diferente sobre a mesma memória. Compare as colunas dos dois "o": embeddings idênticos geram q, k e v idênticos, então as duas posições recebem o mesmo peso de atenção. Clique em 🎲 e observe que outros Wq, Wk, Wv produzem outro padrão de atenção: é exatamente isso que o treinamento ajusta.
Na forma matricial vista em aula, attention(Q, K, V) = softmax(QK⊤/√d)·V, a matriz softmax(QK⊤/√d) tem shape seq_size × seq_size: a linha t contém os pesos da query do token t sobre todos os tokens, e cada linha soma 1. Abaixo, a frase dos slides em português: "o animal não cruzou a rua porque ele estava cansado". Com pesos "treinados" (escolhidos à mão para imitar uma cabeça de atenção útil), o pronome "ele" atende ao substantivo animado "animal". Passe o mouse sobre as células para ver cada par e clique numa linha para detalhá-la à direita.
Matriz de atenção: linhas são queries, colunas são keys (hover mostra o par, clique seleciona a linha)
Linha selecionada: pesos ai(t) da query sobre cada token
💡 Com os pesos treinados, olhe a linha de "ele": o maior peso vai para "animal", a atenção resolve a correferência do exemplo dos slides. Veja também as linhas dos artigos "o" e "a" (atendem substantivos) e a de "não" (atende verbos). Troque para pesos aleatórios e sorteie algumas vezes com 🎲: antes do treinamento a matriz não tem estrutura útil, são os Wq e Wk aprendidos que criam o padrão.
Como visto em aula, se as componentes de q e k são i.i.d. com média 0 e variância 1, então 𝔼[q·k] = 0 e var[q·k] = d: a magnitude típica dos scores cresce com √d. Scores grandes saturam a softmax (pesos quase 0 ou 1) e os gradientes somem. Dividir por √d traz a variância de volta para 1, independente da dimensão. Mova o slider de d e compare a distribuição dos scores e a softmax resultante com e sem a escala.
Distribuição de q·k (ciano, var = d) e de q·k/√d (laranja, var = 1), alturas normalizadas
Softmax sobre m = 10 scores: sem escala (ciano) vs com escala (laranja)
💡 Em d = 2 as duas distribuições quase coincidem e as duas softmax são parecidas. Suba para d = 128 ou 512: sem escala os scores se espalham (desvio padrão √d ≈ 22.6 em d = 512) e a softmax vira quase one-hot, com o maior peso perto de 1 e os demais perto de 0. Com a escala, a distribuição laranja não muda com d e a softmax permanece suave. Use 🎲 para ver que o efeito não depende da amostra.
A saída de uma query é c(t) = Σ ai(t)·vi com ai(t) ≥ 0 e Σ ai(t) = 1: uma combinação convexa, então c(t) sempre cai dentro do fecho convexo dos values. Abaixo, seis values vi ∈ ℝ² arrastáveis e scores fixos si. O slider γ multiplica os scores antes da softmax, a = softmax(γ·s), fazendo o papel da magnitude dos scores da seção anterior (sem escala, γ grande; com escala, γ moderado).
Values vi (arraste os pontos; o raio acompanha o peso), saída c em roxo e média simples em verde
Pesos ai = softmax(γ·si); a linha tracejada marca o caso uniforme 1/m
💡 Com γ = 0 a softmax é uniforme e c coincide com a média simples (× verde). Em torno de γ ≈ 1 a saída é uma mistura suave. Com γ = 6 a softmax fica quase one-hot e c gruda no value de maior score: a atenção vira uma seleção. Arraste um value de peso alto e veja c acompanhá-lo; arraste um de peso baixo e quase nada muda. Note que c nunca sai do polígono tracejado.
Como visto em aula, se as componentes de q, k ∈ ℝd são i.i.d. com 𝒩(0, 1), então 𝔼[q·k] = 0 e var[q·k] = d: a magnitude típica dos scores cresce com √d. Scores grandes saturam a softmax, quase todos os pesos da atenção viram 0 ou 1 e o gradiente some. A correção do transformer é usar score(q, k) = q⊤k/√d, que mantém a variância em 𝒪(1). Mova o slider de d e compare o histograma dos scores e os pesos da softmax com e sem a escala.
Histograma de q·k (laranja) e de q·k/√d (ciano) sobre 1500 pares aleatórios
Pesos da softmax de uma query sobre m = 10 keys: sem escala (laranja) vs com escala (ciano)
💡 Com d = 2 as duas distribuições quase coincidem e as softmax são parecidas. Suba para d = 256 ou 512: o histograma laranja se espalha (desvio ≈ √d) enquanto o ciano fica parado, e a softmax sem escala concentra quase todo o peso em uma única key (saturação). Clique em 🎲 algumas vezes: qual key "ganha" muda, mas o padrão de saturação não.
Cada token xt da frase dos slides, "The animal didn't cross the street because it was tired", é projetado em qt = xtWq, kt = xtWk e vt = xtWv. Com D = 16 e H cabeças, cada cabeça trabalha em dimensão D/H com suas próprias matrizes, calcula softmax(QK⊤/√d) (uma matriz m × m de pesos, aqui em heatmap) e as saídas são concatenadas e projetadas por WO. Na aba decoder, a máscara causal soma −∞ aos scores com j > i, zerando a atenção ao futuro. Clique em uma linha de qualquer heatmap para escolher a cabeça e o token de consulta.
Matriz de atenção softmax(QK⊤/√d) de cada cabeça (linha = query, coluna = key; clique para selecionar)
Pesos ai(t) da cabeça selecionada e, abaixo, concat(head1, …, headH) → WO
💡 Com H = 1 há um único padrão de atenção; com H = 8 cada cabeça olha a frase de um jeito diferente (algumas quase diagonais, outras espalhadas). Selecione o token "it" e compare os pesos entre cabeças: em algumas, "animal" recebe peso alto, ecoando a correferência. Troque para a máscara causal e veja o triângulo superior zerar: a linha do primeiro token vira 1 na própria posição. Clique em 🎲 e note que os padrões mudam, pois Wq, Wk, Wv aqui são aleatórios, não treinados.
Self-attention pura é invariante à ordem dos tokens, então somamos ao embedding uma codificação de posição. A definição dos slides é Pi,2k = sin(i/n2k/d) e Pi,2k+1 = cos(i/n2k/d), onde i é a posição, d a dimensão do embedding e n o hiperparâmetro de escala (10000 no paper). Dimensões à esquerda oscilam rápido, à direita devagar, como os bits de um contador binário contínuo. O gráfico ao lado mostra o produto escalar normalizado Pi·Pj/‖Pi‖²: pico em j = i e decaimento suave, o que dá ao modelo a noção de proximidade relativa. Arraste verticalmente sobre o heatmap ou use o slider para escolher i.
Matriz P: posição i (vertical) × dimensão j (horizontal); ciano = +1, laranja = −1 (arraste para mover i)
Pi·Pj/‖Pi‖² em função de j (pico em j = i)
💡 Com n = 10000 cada posição tem assinatura única e o produto escalar decai suavemente. Reduza para n = 10: todas as colunas viram ondas rápidas, o padrão se repete e aparecem picos secundários altos no gráfico da direita, posições distantes ficam parecidas (viola o requisito de codificação única). Aumente d com n grande e veja as colunas extras quase constantes, frequências lentas que distinguem posições muito distantes.
Cada sub-bloco do transformer é embrulhado em output = LayerNorm(x + Subblock(x)). A LayerNorm normaliza por token (por linha): LN(xi) = γ⊙(xi − μi)/σi + β, com μi e σi calculados sobre as d dimensões do próprio token. À esquerda, um batch de 6 tokens com médias e escalas diferentes entra na LN e sai padronizado linha a linha. À direita, empilhamos N blocos Feed Forward aleatórios e medimos ‖x(ℓ)‖ em cada profundidade: sem o caminho identidade o sinal colapsa ou explode conforme o ganho g dos pesos; com residual ele sobrevive, e com residual + LN a norma fica estável, exatamente o que permite treinar N = 6 (ou muito mais) blocos.
Batch de 6 tokens × d = 16: antes (esquerda) e depois da LayerNorm (direita); ciano = +, laranja = −
log₁₀‖x(ℓ)‖ por profundidade ℓ: sem residual (laranja), x + F(x) (verde), LN(x + F(x)) (ciano)
💡 Na LN, suba μin para 3 e σin para 3: a entrada vira faixas desbalanceadas, mas a saída continua igual (a LN apaga média e escala de cada token); γ e β reintroduzem escala e deslocamento aprendidos. Na profundidade, com g = 0.7 a curva laranja despenca (sinal some em poucos blocos) e com g = 1.4 ela explode; a verde cresce devagar e a ciano fica plana para qualquer g. Veja no readout que cos(x(0), x(N)) fica alto só com residual: o caminho identidade preserva o sinal original.