Demonstrações interativas da aula de modelos de linguagem decoder-only
Quatro visualizações que acompanham os slides: a máscara causal Mij = −∞ se j > i comparada com a atenção bidirecional e a sliding window, o laço autoregressivo P(xt | x<t) com prefill e decode usando KV-cache, o tamanho do cache nas variantes MHA, MQA e GQA (2·G·dk·L floats por token), e as leis de escala de Chinchilla com C ≈ 6·N·D e a regra dos 20N tokens.
No encoder (BERT) cada posição vê toda a sequência. No decoder a self-attention é mascarada: soma-se Mij = 0 se j ≤ i, −∞ se j > i aos scores antes da softmax, então a posição i só atende a j ≤ i e o token nunca "vê o futuro" durante o treino. A matriz abaixo mostra os pesos softmax(QK⊤/√dk + M) de uma cabeça com scores sorteados: linha i é quem consulta, coluna j é quem é consultado. A aba sliding window restringe ainda mais, cada posição atende só às W posições anteriores (Mistral), com custo O(n·W) em vez de O(n²).
Matriz de atenção: roxo = peso após a softmax, × = mascarado (passe o mouse e clique numa linha)
Pesos da linha selecionada: a softmax redistribui a massa só entre as posições visíveis
💡 Selecione a linha i = 1 no modo causal: o primeiro token só pode atender a si mesmo, então o peso é 1 por definição. Troque entre bidirecional e causal observando a mesma linha: os scores são idênticos, mas a softmax renormaliza só sobre j ≤ i. Na sliding window com n = 12 e W = 2, conte as células ativas por linha: no máximo W + 1, e é daí que vem o custo O(n·W).
Um decoder-only gera um token por vez amostrando de P(xt | x1, …, xt−1). A inferência tem duas fases vistas em aula: o prefill processa o prompt inteiro em paralelo e armazena K e V de cada token; cada passo de decode computa K e V apenas do token novo e reutiliza o cache. Sem cache, todo K e V é recomputado a cada passo e o custo total cresce como O(n²·d); com cache, cai para O(n·d) por token. O modelo aqui é um trigrama por contagem treinado num pequeno corpus (o do "rato que roeu a rolha" dos slides), o suficiente para ver o laço autoregressivo funcionando.
Sequência e células K, V por posição: verde = no cache (reutilizado), laranja = computado neste passo
Distribuição P(xt | x<t) do próximo token (8 mais prováveis)
K, V computados (acumulado) por passo de decode: com cache vs sem cache
💡 Gere passo a passo e observe que com cache só uma coluna fica laranja por passo, enquanto no modo sem cache todas ficam: é exatamente a diferença entre as curvas linear e quadrática do gráfico de custo. Repare também que a distribuição muda a cada token gerado, pois o contexto x<t cresceu: isso é o laço autoregressivo. Use 🎲 para amostrar outra continuação do mesmo prompt.
O KV-cache custa 2·G·dk·L floats por token, onde G é o número de pares K,V por camada. Na MHA cada uma das H cabeças tem seu par (G = H); a MQA compartilha um único par entre todas (G = 1); a GQA divide as cabeças em G grupos (1 < G < H), o compromisso usado em Llama 2/3 e Mistral. Ajuste H, G, as camadas L e o contexto n e veja a memória do cache em fp16 (dk = 128 fixo).
Cabeças de query (ciano) e pares K,V compartilhados (verde): cada grupo de H/G cabeças lê o mesmo par
Memória do KV-cache (fp16) para o contexto n: MHA vs GQA atual vs MQA
💡 Carregue o preset do GPT-3 e veja o cache de MHA passar de 9 GB já com n = 2048: por isso servir LLMs com MHA pura é caro. No preset do Llama 2 70B (H = 64, G = 8) a redução é de 8× com qualidade quase idêntica. Arraste G até 1 para virar MQA e até H para virar MHA: GQA interpola entre os dois extremos.
Treinar um modelo de N parâmetros com D tokens custa C ≈ 6·N·D FLOPs (2 no forward, 4 no backward, por parâmetro e por token). Para um orçamento C fixo, a curva abaixo (forma paramétrica de Hoffmann et al., L(N, D) = E + A/Nα + B/Dβ com D = C/6N) mostra a perda de cada divisão entre modelo e dados: modelos pequenos demais ou grandes demais (subtreinados) perdem para o ponto compute-ótimo, em que N* ∝ C0.5 e D* ∝ C0.5, a razão D*/N* ≈ 20 tokens por parâmetro. Arraste o ponto branco sobre a curva e mova o orçamento C.
Curva isoFLOP: perda L(N, C/6N) em função de N (arraste o ponto branco; verde = ótimo, laranja = regra 20N)
N* (ciano) e D* (laranja) compute-ótimos em função de C, com modelos reais marcados
💡 Com log₁₀ C = 23.8 (o orçamento do Chinchilla) o ótimo fica perto de 70B parâmetros e 1.4T tokens, e o ponto do Chinchilla cai em cima das retas do gráfico da direita. Arraste o ponto branco para N = 280B nesse mesmo orçamento e veja a perda subir: é o Gopher, maior porém subtreinado. Repare que o Llama 3 8B fica bem abaixo da reta de N*: ele é deliberadamente sobretreinado (D/N ≈ 1875), porque compute-ótimo não é inferência-ótimo.