Resumo Técnico: Deriva de Atenção em Decodificação Especulativa Autoregressiva

Declaração do Problema

A decodificação especulativa acelera a inferência de Modelos de Linguagem de Grande Escala (LLM) utilizando um modelo "rascunhador" (drafter) leve para prever tokens futuros, os quais são então verificados por um modelo "alvo" maior. Embora eficaz, os atuais modelos rascunhadores (especificamente EAGLE-3 e cabeças de Previsão de Múltiplos Tokens) exibem fragilidade significativa sob variações de implantação, como perturbações de templates e entradas de contexto longo. Praticantes frequentemente precisam re-treinar rascunhadores para motores de inferência ou prompts específicos para manter o desempenho, sugerindo uma instabilidade arquitetônica subjacente.

O artigo identifica um fenômeno anteriormente não reportado chamado Deriva de Atenção. À medida que um rascunhador gera uma sequência de tokens especulativos, seu mecanismo de atenção desloca-se progressivamente do "sumidouro de atenção" (um token âncora estável, tipicamente próximo ao início do prompt) e concentra-se cada vez mais em seus próprios tokens gerados recentemente. Essa deriva correlaciona-se com uma degradação nas taxas de aceitação, particularmente em cenários de contexto longo e quando os templates de entrada se desviam das distribuições de treinamento.

Metodologia e Análise

Os autores investigam as causas raiz da deriva de atenção através de uma combinação de observação empírica e análise arquitetônica em múltiplas famílias de modelos (Llama 3.1, Qwen 3/3.5 e GPT-oss).

1. Observação Fenomenológica

Utilizando mapas de calor de atenção e métricas quantitativas, os autores demonstram que:

• Migração do Sumidouro: Em modelos com sumidouros de atenção fortes (por exemplo, Llama, Qwen), a atenção do rascunhador ao token sumidouro diminui à medida que a cadeia de especulação aprofunda, enquanto a atenção ao token rascunhado mais recente aumenta.
• Universalidade: Essa deriva é observada tanto em rascunhadores EAGLE-3 (treinados post-hoc sobre alvos congelados) quanto em cabeças MTP (treinadas conjuntamente com o alvo), sugerindo que é uma propriedade de designs autoregressivos que consomem suas próprias saídas.
• Ausência de Sumidouro: Mesmo em modelos projetados para suprimir sumidouros (por exemplo, GPT-oss com atenção em porta ou vieses aprendíveis), ocorre uma forma de deriva onde a atenção se concentra em tokens recorrentes do template em vez do início do prompt.

2. Análise da Causa Raiz: Magnitude do Estado Oculto

O artigo traça a deriva de atenção ao caminho residual não normalizado entre os passos de especulação em arquiteturas padrão pré-norm:

• Crescimento Monotônico: Em transformers pré-norm, as conexões residuais acumulam magnitude sem normalização. Os autores mostram que a magnitude do estado oculto do rascunhador ($\|h^*_k\|$) cresce monotonicamente com a profundidade da especulação ($k$).
• Empilhamento de Camadas vs. Autoregressão: Esse crescimento faz com que o rascunhador se comporte menos como um preditor autoregressivo autônomo e mais como uma sequência de camadas adicionais de transformer empilhadas sobre o alvo. No passo $k$, o rascunhador efetivamente aproxima o modelo alvo na profundidade $N+k$, em vez de prever o próximo token baseado em uma distribuição estável.
• Entradas Desequilibradas: Os estados ocultos fundidos do modelo alvo ($h_{FC}$) são frequentemente dominados por camadas mais profundas com magnitudes maiores, exacerbando ainda mais o desequilíbrio de escala quando alimentados no rascunhador.

3. Intervenções Propostas

Para combater a acumulação de magnitude e estabilizar o processo de rascunho, os autores propõem duas modificações arquitetônicas:

1. Pré-Norm em Estados Ocultos do Rascunhador: Aplicar RMSNorm aos estados ocultos do rascunhador após a conexão residual (entre passos de especulação) impede o crescimento monotônico da magnitude. Isso força o rascunhador a aprender uma regra autoregressiva invariante à profundidade em vez de refinamentos específicos à profundidade.
2. RMSNorm por Estado Oculto: Adicionar camadas RMSNorm antes de fundir os estados ocultos do alvo ($h_{low}, h_{mid}, h_{high}$) na entrada do rascunhador para equilibrar a magnitude do sinal de entrada.

Os autores também avaliam a Atenção em Porta (suprimindo a dependência de um sumidouro fixo) como uma intervenção separada. Eles descobrem que, embora a porta elimine o sumidouro, ela não impede o crescimento de magnitude, e combinar porta com pós-norm pode levar à super-regularização (colapso de entropia).

Resultados Chave

A arquitetura pós-norm proposta foi avaliada contra o atual estado da arte pré-norm EAGLE-3 em vários benchmarks e condições:

• Perturbação de Template: Sob perturbações de template (por exemplo, removendo templates de chat ou tokens de início de sequência), o rascunhador pós-norm melhorou o comprimento de aceitação em até 2× comparado ao pré-norm. Modelos pré-norm sofreram quedas catastróficas (perda de até 52%), enquanto modelos pós-norm permaneceram resilientes (perda máxima de 5%).
• Tarefas de Contexto Longo: Em tarefas de contexto longo (LongBench, chat multi-turno com atenção de janela deslizante), o pós-norm melhorou o comprimento de aceitação em 1,18×. Em configurações de contexto longo com atenção total, modelos pré-norm colapsaram para aceitação próxima de zero, enquanto o pós-norm manteve desempenho utilizável.
• Benchmarks Padrão: Em sete benchmarks padrão (chat multi-turno, matemática, codificação), o pós-norm mostrou uma melhoria consistente de 1,10× no comprimento de aceitação.
• Eficiência de Treinamento: Rascunhadores pós-norm generalizam para profundidades de inferência além do seu horizonte de treinamento. Isso permite treinar com profundidades "tempo de treino-teste" (TTT) mais curtas (por exemplo, reduzindo de 8 para 4 passos), cortando o tempo de treinamento em aproximadamente um terço sem sacrificar o desempenho.
• Robustez a Ruído: Arquiteturas pós-norm toleram significativamente mais ruído nos estados ocultos (retenção de 58% de aceitação em ruído $\alpha=0,5$ vs. 5% para pré-norm), sugerindo melhor robustez à quantização e deslocamentos de distribuição.

Significado e Alegações

O artigo alega que a Deriva de Atenção é um modo de falha crítico, anteriormente não reportado, na decodificação especulativa que decorre da interação entre caminhos residuais não normalizados e realimentação autoregressiva.

• Diagnóstico: Os autores argumentam que a deriva de atenção é um sintoma de uma questão mais profunda: a acumulação da magnitude do estado oculto, o que faz com que o rascunhador aprenda transformações dependentes da profundidade em vez de previsão de token estável.
• Solução: Eles afirmam que simplesmente corrigir o sumidouro de atenção (por exemplo, via porta) é insuficiente; as dinâmicas de magnitude subjacentes devem ser abordadas. A arquitetura pós-norm proposta é apresentada como uma correção simples e eficaz que estabiliza o processo de rascunho.
• Generalização: As descobertas sugerem que o pós-norm deve ser considerado uma escolha arquitetônica padrão para rascunhadores de decodificação especulativa, particularmente para sistemas de produção que exigem robustez contra prompts variados e contextos longos.

O artigo conclui que nomear e caracterizar a "deriva de atenção" fornece um diagnóstico útil para a robustez do rascunhador, e que intervenções arquitetônicas visando a acumulação de magnitude são essenciais para uma decodificação especulativa confiável.