← início
Nível 3 · model/TransformerModel.java

Transformer

O degrau final. Duas mudanças o separam do nível 2. O contexto deixa de ser uma janela rígida e passa a ser ponderado pelo próprio modelo; e os gradientes deixam de ser escritos um a um e passam a ser derivados por um grafo.

Contexto: 64 caracteresTreino: ~25 min (2000 passos)Perda alvo: < 1,7
$ java -jar target/mini-gpt-java.jar train --model transformer

O nível por dentro

A limitação que este nível resolve

O MLP concatena a janela. Isso tem duas consequências ruins. Primeira: cada posição ganha um bloco próprio de pesos, então o que a rede aprende na posição 3 não transfere para a posição 4. Segunda: a janela é rígida, com oito caracteres sempre, sejam eles relevantes ou não.

A pergunta que a atenção respondeEm vez de "quais são os oito caracteres anteriores?", a atenção pergunta: dos caracteres anteriores, de quais eu preciso agora, e quanto? Os pesos dessa mistura não são fixos: são calculados a cada posição, a partir do próprio conteúdo.

Consultas, chaves e valores

Uma cabeça de atenção: as três projeções, os escores, a máscara e a misturax(T × E)Qo que procuroQ = x·W_QKo que ofereçoK = x·W_KVo que entregoV = x·W_VQ·Kᵀ / √d + máscara(T × T)softmax por linhapesos · Va saída de cada posição é a média ponderada dos valores que ela escolheu olhar
Uma cabeça de atenção. As três projeções saem do mesmo x, daí o nome self-attention.
Q = x·W_Q     K = x·W_K     V = x·W_V
atenção(Q, K, V)  =  softmax( Q·Kᵀ / √d  +  máscara ) · V
Q (query)
o que a posição atual procura: uma pergunta, em forma de vetor
K (key)
o que cada posição anterior oferece: um rótulo com que a pergunta é comparada
V (value)
o que cada posição entrega, se for escolhida
Q·Kᵀ
o produto escalar de cada pergunta com cada rótulo: alto quando combinam
√d
o fator de escala, com d = E / número de cabeças

O softmax por linha transforma essas afinidades em pesos que somam 1. A saída de uma posição é a média ponderada dos valores das posições que ela decidiu olhar. Mude o texto e os pesos mudam.

Multi-head é o mesmo mecanismo repetido em subespaços independentes. Com E = 128 e 4 cabeças, cada cabeça trabalha em d = 32 dimensões e pode se especializar (uma acompanha a palavra em curso, outra o começo da frase). As saídas são concatenadas de volta a 128.

Por que dividir por √d

O produto escalar de dois vetores aleatórios de dimensão d tem desvio padrão proporcional a √d. Sem correção, com d = 32 os escores chegam ao softmax grandes demais; ele satura, um peso vai a quase 1 e os outros a quase 0, e a derivada do softmax saturado é quase nula. O modelo para de aprender antes de começar.

Um padrão que se repeteEscala e inicialização existem, quase sempre, pelo mesmo motivo: manter os valores na faixa em que a derivada ainda é grande. Foi por isso que o nível 2 inicializou os pesos com 1/√fan_in, e é por isso que aqui se divide por √d. Duas aparições da mesma preocupação.

A máscara causal

Nada no mecanismo acima impede a posição 5 de olhar a posição 9. E a posição 9 é justamente o caractere que a posição 5 deveria prever. Sem impedimento, o modelo aprende a copiar a resposta: a perda de treino despenca e o modelo generaliza zero.

escores[i][j]  =  −∞      sempre que  j > i
i
a posição que está prevendo
j
a posição sendo olhada
−∞
porque e^{−∞} = 0: depois do softmax, o peso é exatamente zero

O Tensor.causalMask aplica isso antes do softmax, e não depois. Zerar depois exigiria renormalizar em seguida e ainda deixaria o gradiente fluir pelo caminho proibido: o mascaramento tem que acontecer onde o softmax possa vê-lo.

Posição: a atenção não tem ordem

Embaralhe as posições de entrada e a atenção devolve as mesmas saídas, embaralhadas junto. Ela é uma operação sobre um conjunto, não sobre uma sequência, e uma sequência de caracteres embaralhada não é português. A correção é somar, a cada posição, um vetor que depende só do índice:

Tensor x = Tensor.add(
Tensor.rows(tokEmb, tokIds), // quem é o caractere
Tensor.rows(posEmb, posIds)); // onde ele está
model/TransformerModel.java — token mais posição

Os dois embeddings são aprendidos. A soma parece ingênua, e funciona porque o espaço tem 128 dimensões, com direções de sobra para carregar as duas informações sem que uma apague a outra.

LayerNorm e conexões residuais

Um bloco: LayerNorm, atenção, residual, LayerNorm, feed-forward, residualx(B·T × E)LayerNormatenção causal multi-headresidual+LayerNormfeed-forwardE → 4E → E, com tanh no meio+para o bloco seguintex'(B·T × E)a derivada de x + f(x) é 1 + f′(x): sempre existe um caminho por onde o gradiente passa inteiro
Um bloco. O nível 3 empilha dois deles por padrão (--blocks).
x  ←  x  +  atenção( LayerNorm(x) )
x  ←  x  +  feedForward( LayerNorm(x) )
LayerNorm
normaliza cada posição para média 0 e variância 1, depois reescala por γ e β aprendidos
x + …
a conexão residual: a sub-camada aprende uma correção do que já estava lá

O residual é o que torna a profundidade viável. Como a derivada de x + f(x) em relação a x é 1 + f'(x), existe sempre um caminho por onde o gradiente chega intacto às camadas de baixo. Sem ele, o gradiente atravessa um produto de fatores e desaparece.

Pesos amarrados

A cabeça de saída poderia ser mais uma matriz E × V. Em vez disso, o projeto reaproveita a tabela de embeddings transposta: logits = x · tokEmbᵀ. A mesma matriz que diz "este caractere é este vetor" responde "qual caractere se parece com este vetor?", e a economia é de V × E parâmetros, o que também regulariza.

O autodiff, em três frases

Um Tensor guarda três coisas: o valor, o gradiente acumulado e uma referência às entradas que o produziram, junto com a função que propaga o gradiente para elas. Cada operação (matmul, layerNorm, softmaxRows, …) constrói um nó novo e registra esse caminho de volta. Ao final, loss.backwardAll() percorre o grafo em ordem topológica reversa, aplicando cada propagação uma vez.

É por isso que o nível 2 vem antesQuem já derivou dW2 = Aᵀ·dLogits passo a passo reconhece a mesma expressão dentro do backward do matmul. O autodiff é o nível 2 escrito uma vez por operação, em vez de uma vez por modelo.

O lote achatado

Quase tudo no Transformer age em cada posição independentemente: embeddings, LayerNorm, as projeções Q/K/V, o feed-forward e a cabeça de saída. Empilhando as B sequências de comprimento T numa única matriz (B·T, E), essas camadas viram uma multiplicação grande em vez de B pequenas, o que importa muito em CPU. Só a atenção mistura posições, e por isso só ela é feita sequência por sequência.

O que esperar do texto

  • A menor perda dos três níveis. A meta declarada do projeto é validação abaixo de 1,7 num corpus real de ~1 MB.
  • Palavras reais na maior parte do tempo, e frases que começam a fechar: a concordância de gênero e número aparece sozinha.
  • Com os padrões (2000 passos, lote 16), cerca de 25 minutos numa CPU comum, sem GPU. Mais --steps leva mais fundo, ao custo de tempo.
  • A distância entre a perda de treino e a de validação é o que você deve vigiar: se a primeira cai e a segunda para, o modelo passou a decorar.

Para fazer com as mãos

Nenhum destes exercícios pede uma biblioteca nova, e todos cabem numa sessão. O terceiro e o quarto de cada nível são os que mais ensinam, porque quebram alguma coisa de propósito.

E1Uma cabeça ou quatro

Compare --heads 1 com --heads 4, mantendo --embed 128. O número de parâmetros é praticamente o mesmo, então qualquer diferença de perda vem da estrutura, não do tamanho.

E2Encurte o contexto

Rode com --context 16 e com 64. Quanto o contexto longo vale, em nats? E quanto ele custa, em segundos por passo? (A atenção é quadrática em T.)

E3Fure a máscara

Remova a chamada a Tensor.causalMask e treine por 200 passos. A perda de treino despenca; a de validação, não. Explique o que o modelo aprendeu a fazer, e por que é inútil.

E4Empilhe mais um bloco

Rode com --blocks 3. Custa mais tempo por passo e melhora a perda? Depois pense: se as conexões residuais não existissem, o que aconteceria com o gradiente ao atravessar três blocos?

O arquivo deste nívelTudo o que esta página explica está em model/TransformerModel.java. O Javadoc da classe traz a mesma matemática, ao lado da linha que a implementa.