Transformers Joguinho

Construindo o transformers joguinho do zero

Achei útil fazer um MiniTransformer simples em PyTorch para entender o mecanismo de atenção antes de brincar com os modelos prontos do Hugging Face. Não precisa complicar. Eu usei uma rede com uma head de atenção multi-head, encodeur posicional sinusoidal e duas camadas de feed-forward. O objetivo era rodar em CPU por uns dias e perceber onde o processo travava, não rodar em GPU. O primeiro erro que eu cometi foi tentar treinar com batch size alto logo de cara. Eu estava usando sequências de 512 tokens e batch de 64, e o custo de memória disparava pra nada. A atenção tem complexidade quadrática no tamanho da sequência, então o produto atenção x query vira uma matriz enorme muito rápido. O workaround foi simplesmente reduzir a janela de atenção para 256 tokens e batch para 8. O treino ficou mais lento em épocas, mas cada step passava sem estourar memória. Depois eu descobri que existe flash-attention (versão 2), mas isso exige CUDA e já é outro nível de complexidade.

O que é transformers joguinho

Transformers joguinho é basicamente um exercício didático de implementação de um modelo transformer pequeno, geralmente orientado a texto. A proposta é escrever o mecanismo self-attention, a parte de feed-forward position-wise, e a camada de saída em poucas linhas de código. Isso ajuda a entender o que acontece por baixo dos panos antes de usar bibliotecas como Transformers do Hugging Face. Uma coisa que muita gente não percebe é que você não precisa de mil camadas para ter um resultado decente num dataset pequeno. Dois blocks de transformer com 4 heads de atenção e dimensão de embeding de 128 já resolvem classificações simples de texto, como sentiment analysis com IMDB ou categorização de datasets pequenos. O modelo não vai generalizar para tarefas complexas, mas para fins de aprendizado isso é suficiente.

Componentes essenciais

Vou listar só o que realmente importa montar. Qualquer coisa além disso é detail de implementação que você encontra na documentação do PyTorch sem problema.

Self-attention

O self-attention calcula scores de atenção entre todos os pares de posições da sequência. Você projeta queries, keys e values usando pesos aprendíveis, faz o produto matricial QK, aplica scale por raiz quadrada da dimensão das keys, e passa por softmax. O resultado é ponderado sobre V. Se você está treinando multi-head attention, divide o embeding em N heads, processa cada head separadamente, e depois concatena tudo de volta.

Encoding posicional

Como o transformer não tem natureza recursiva como RNNs, ele não sabe a ordem dos tokens por si só. O encoding posicional sinusoidal adiciona informação de posição ao embedding de cada token. Eu usei seno para dimensões pares e cosseno para ímpares, com frequências decrescentes de acordo com a dimensão. Isso funciona bem porque permite ao modelo generalizar para comprimentos de sequência diferentes durante o inferência.

Feed-forward network

Cada posição passa por uma rede feed-forward independente, geralmente com dois camadas fully connected e uma ativação ReLU no meio. A dimensão interna costuma ser 4x a dimensão do embeding. Isso é o que dá capacidade de expressividade ao modelo além do próprio mecanismo de atenção.

Treino e problemas comuns

Um dos problemas mais chatos que eu encontrei foi o modelo simplesmente não convergir com learning rate alto. Eu comecei com 1e-3 e o loss oscilava sem baixar. Reduzi para 1e-4 com weight decay de 0.01 e AdamW, aí estabilizou. Learning rate scheduler do tipo warmup linear por 10% dos steps e depois decay cosine também ajudou bastante. Sem warmup, os primeiros steps tendem a dar gradientes explosivos porque os pesos começam aleatórios. Outro ponto: normalização de camadas. Eu estava esquecendo o layer norm residual em alguns lugares. A arquitetura original do paper usa pre-norm (layer norm antes do sub-module) em vez de post-norm. Com post-norm, o modelo ficou instável nas primeiras épocas. Troquei para pre-norm e o treino estabilizou numa época.

Se você vai usar dropout, não coloque em todo lugar. Dropout de 0.1 nos outputs de attention e 0.1 na FFN funciona. Dropout mais alto que isso começa a prejudicar mais do que ajudar em modelos pequenos como esse.

Limitações reais

MiniTransformer não serve para nada sério. A capacidade dele é limitada pela quantidade de parâmetros. Um modelo com 2 blocks, 4 heads e dimensão 128 tem aproximadamente 1.5 milhão de parâmetros. Isso não é suficiente para tarefas que exigem compreensão profunda de linguagem. Para comparação, o GPT-2 small tem 124 milhões. Se você precisa de algo que funcione de verdade em produção, use um modelo pré-treinado do Hugging Face. BERT base, RoBERTa, ou TinyLlama dependendo da tarefa. O valor do mini transformer é estritamente educacional. Ele mostra como a peça se encaixa, mas não substitui modelos treinados em larga escala.

Outra limitação prática é o custo computacional da atenção quadrática. Se você aumentar a sequência para 1024 tokens, o tempo de treino dobra ou triplica mesmo com batch menor. A solução seria implementar sparse attention ou usar técnicas como Longformer, mas aí você sai do escopo de um projeto didático.

Exemplo de código

Aqui vai um esqueleto mínimo. A intenção é mostrar a estrutura, não ser um código completo e production-ready. Importante: normalize o embedding de entrada com layer norm antes de passar pelo primeiro bloco de transformer. Muita gente pula isso e se pergunta por que o treino demora pra decolar. O código abaixo presume um vocabulário já tokenizado e convertido em IDs numéricos. Se você está começando do zero, use o tokenizer do BERT como base, ele é simples de importar e já vem treinado.

Uma dica final: salve o modelo a cada 500 steps e faça validação periódica. Eu perdi horas num treino porque o modelo tinha overfitted nos primeiros 2000 steps e eu não tinha salvo checkpoint. Sem validação, você não sabe quando parar.

Recursos práticos

O repositório do pytorch.org tem um exemplo oficial de transformer chamado "AnnotatedTransformer" que é um bom ponto de partida. Também vale a pena olhar a implementação do mini-gpt da karpathy no repositório dele, que explica o GPT em poucas linhas. Para quem quer ir além e adicionar causal masking, a função generate_causal_mask do PyTorch resolve rapidinho. Se o objetivo é apenas brincar e entender, baixar um checkpoint pequeno do Hugging Face e rodar inference com ele pode ser mais produtivo do que treinar do zero. Mas se a meta é entender o mecanismo de atenção na prática, o transformers joguinho é o caminho mais direto.