Mini-GPT
Um GPT decoder-only escrito do zero em PyTorch, um componente por vez, cada um com seus testes, um experimento e um documento que explica a matemática.
- Python
- PyTorch
Desafio
Entender exatamente o que acontece dentro de um modelo de linguagem quando ele recebe uma frase, sem os atalhos do nn.Transformer, e justificar cada escolha de projeto com um número medido no próprio modelo.
Abordagem
Construir o modelo em dezoito componentes, do tokenizer ao tokenizer BPE que o substituiu, cada um com seu arquivo de teste, um experimento e um documento. As peças escritas à mão, como a LayerNorm e a cross-entropy, têm testes que garantem igualdade com as do PyTorch. Cada escolha de projeto fica registrada com a principal alternativa descartada e o motivo, quase sempre uma medição. O treino usa AdamW com warmup e gradient clipping, e os checkpoints guardam os pesos, o estado do otimizador, a posição no data loader e os estados aleatórios.
Resultado
Um treino de 15 épocas leva a melhor loss de validação a 2,551, perplexidade de 12,82, e o modelo ainda melhorava quando o treino parou. Retomar de um checkpoint reproduz bit a bit o treino sem pausa, o que um teste confere, e um checkpoint treinado vem no repositório, então dá para gerar texto sem treinar antes.
Arquitetura
Decisões e trade-offs
Um tokenizer BPE substituiu o tokenizer por caracteres original
A geração parava no meio da palavra, e o BPE cobre cerca de 2,1x mais caracteres no mesmo context length.
Um token embedding bem maior: o vocabulário foi de 112 para 4.096 tokens.
Pre-norm: LayerNorm antes da attention e do feed-forward
O caminho residual fica livre de normalizações, e o gradiente atravessa toda a profundidade do modelo.
O residual stream nunca é normalizado e sua escala cresce durante o treino, então uma LayerNorm final fica antes do LM head.
Weight tying entre o token embedding e o LM head
Como no GPT-2, economiza V×d parâmetros e dá gradiente a todas as linhas do embedding em todo passo.
LayerNorm e cross-entropy escritas à mão
O objetivo é didático; testes garantem igualdade com o nn.LayerNorm e o F.cross_entropy do PyTorch.
A cross-entropy própria é cerca de 0,14 ms mais lenta por passo.
Os checkpoints guardam todo o estado do treino, não só os pesos
Retomar reproduz bit a bit o treino sem pausa. Sem o estado do AdamW, os primeiros passos depois de retomar andam mais que o normal e o resultado deixa de ser reproduzível.
A geração sorteia com temperature 0,8 e top-p 0,95
A decodificação greedy entra em loop, e o top-p adapta o número de candidatos à confiança do modelo.
Evidência principal
Perplexidade de validação de 12,82 após um treino de 15 épocas