Resumo Técnico: Resíduos de Atenção Multi-Cabeça (MHAR)

1. Declaração do Problema

Os Transformers padrão propagam informações através da profundidade via um único fluxo residual aditivo, onde cada subcamada lê apenas o estado imediatamente anterior ($h_\ell = h_{\ell-1} + f_\ell(h_{\ell-1})$). Embora características úteis computadas em camadas anteriores permaneçam na soma contínua, elas tornam-se individualmente inacessíveis.

Resíduos de Atenção (Kimi, 2025) relaxam isso tratando todas as saídas de subcamadas anteriores como uma memória, permitindo que cada subcamada recupere uma combinação aprendida e ponderada por softmax de todo o histórico de profundidade. No entanto, este mecanismo sofre de um gargalo estrutural: utiliza uma única consulta de roteamento compartilhada ($q \in \mathbb{R}^d$) para pontuar o histórico de profundidade. Consequentemente, todas as $d$ coordenadas de características são forçadas a ler das mesmas camadas nas mesmas proporções através de uma única distribuição $\alpha$.

O artigo argumenta que isso cria um "compromisso forçado". À medida que a largura do modelo aumenta, diferentes subespaços de características provavelmente desenvolvem preferências divergentes sobre quais camadas históricas são mais relevantes. Uma única consulta não pode satisfazer simultaneamente essas preferências distintas, colapsando preferências de profundidade distintas em uma única leitura. O custo desse compromisso cresce com a largura do modelo, podendo tornar o roteamento de consulta única prejudicial em grandes escalas.

2. Metodologia: Resíduos de Atenção Multi-Cabeça (MHAR)

Os autores propõem os Resíduos de Atenção Multi-Cabeça (MHAR), uma generalização estrita dos resíduos de atenção que elimina o compromisso forçado sem adicionar parâmetros ou computação significativa.

Mecanismo Central

Em vez de uma única consulta $q$, o MHAR remodela a consulta de roteamento em $H$ cabeças por subespaço, cada uma com seu próprio softmax independente sobre o histórico de profundidade.

• Remodelação (Reshaping): O vetor de consulta único $q \in \mathbb{R}^d$ é remodelado em um tensor de forma $(H, d/H)$.
• Roteamento Independente: Para cada cabeça $h$, o modelo computa uma distribuição de profundidade distinta $\alpha^{(h)}$ baseada no corte do vetor de origem correspondente àquela cabeça.
• Leitura Diagonal de Bloco: A mistura recuperada é diagonal de bloco; cada subespaço lê o histórico de profundidade de forma independente.
• Recuperação: Definir $H=1$ recupera exatamente os resíduos de atenção de cabeça única originais.

Detalhes de Implementação

• Eficiência de Parâmetros: A transformação é uma remodelação, adicionando zero parâmetros sobre a linha de base de cabeça única.
• Eficiência de Computação: O overhead é insignificante, envolvendo apenas $H$ pequenos softmaxes sobre o eixo de profundidade (comprimento $\le 2L+1$) em vez de um único grande.
• Otimização de Sistema: Reconhecendo que o roteamento de profundidade é limitado por memória (exigindo tráfego de $O(L^2 B T d)$), os autores fornecem kernels Triton fundidos. Esses kernels realizam o softmax online com RMSNorm fundido, eliminando a materialização intermediária de chaves e fontes empilhadas. Isso reduz o tempo de roteamento em $2\times$ em relação ao torch.compile e $6.4\times$ em relação às implementações de referência em modo eager. Isso eleva o rendimento de treinamento para $0.55\text{--}0.88\times$ da linha de base.

Configuração de Hiperparâmetros

O número de cabeças de roteamento $H$ é definido como igual ao número de cabeças KV no modelo. Os autores consideram isso um padrão próximo do ideal e livre de hiperparâmetros, que alinha a granularidade do roteamento com a granularidade de consumo do modelo.

3. Principais Contribuições

1. Arquitetura MHAR: Uma modificação livre de parâmetros que divide a consulta de roteamento em cabeças independentes, removendo o compromisso forçado do roteamento de profundidade de consulta única.
2. Kernels Eficientes: Kernels Triton fundidos que tornam o roteamento de profundidade prático em escala, reduzindo significativamente o overhead de memória e a latência.
3. Conversão de Preservação de Identidade: Um método para implantar MHAR em modelos pré-treinados usando "resíduos de atenção delta" (adicionando o termo roteado ao fluxo existente em vez de substituí-lo), permitindo o treinamento intermediário sem picos de perda.
4. Insight Mecanístico: A sondagem direta de consultas treinadas confirma que o desacordo de subespaço cresce com a largura e é o principal driver da degradação do roteamento de cabeça única.

4. Resultados Experimentais

Pré-treinamento (Do Zero)

Treinado no FineWeb-Edu em três escalas (100M, 350M, 1B):

• Perda de Validação: O MHAR melhora consistentemente a perda de validação em relação aos Transformers padrão e aos resíduos de atenção de cabeça única.
  • 100M: $-0.049$ nats em relação à linha de base.
  • 350M: $-0.080$ nats em relação à linha de base.
  • 1B: $-0.063$ nats em relação à linha de base.
• Comportamento de Escala: A vantagem do MHAR sobre o roteamento de cabeça única se alarga com a escala. Enquanto o roteamento de cabeça única ajuda em 100M ($-0.039$), ele degrada para ser prejudicial em 350M ($+0.055$) e 1B ($+0.140$).
• Equivalência de Computação: O MHAR alcança um ganho de equivalência de computação de $1.3\text{--}1.5\times$ sobre a linha de base.

Treinamento Intermediário (8B)

Usando um controle de correspondência de cronograma em um modelo de 8B (Marin-8B) com annealing em um corpus orientado a capacidades:

• Preservação de Identidade: A conversão não introduz choque de otimização; as curvas de perda coincidem com a linha de base durante os primeiros 100 passos.
• Ganhos em Downstream: O MHAR produz melhorias estatisticamente significativas em tarefas de raciocínio:
  • GSM8K: $+3.2$ pontos (McNemar pareado $p=0.004$).
  • GPQA: $+3.1$ pontos (McNemar pareado $p=0.038$).
  • Ganhos em tarefas de MMLU, MATH e código não foram estatisticamente significativos nesta configuração específica.

Estudos de Ablação

• Contagem de Cabeças: O número ideal de cabeças $H$ acompanha o número de cabeças KV. Em regimes subtreinados, mais cabeças sempre ajudam; em regimes convergidos, $H=KV$ é o "joelho" da curva.
• Largura vs. Profundidade: A sondagem de consultas treinadas revela que o "desacordo de largura" (preferências de roteamento divergentes entre subespaços) cresce monotonicamente com a largura do modelo, validando a hipótesia do compromisso forçado.
• Alinhamento de Cabeças: Alinhar as cabeças de roteamento aos grupos KV de atenção não produz benefício mensurável sobre o roteamento de subespaço arbitrário, sugerindo que o benefício vem da multiplicidade de cabeças e não de seu alinhamento específico com as estruturas de atenção.

5. Significância e Alegações

O artigo afirma que os Resíduos de Atenção Multi-Cabeça resolvem uma limitação arquitetônica fundamental na forma como os Transformers acessam o histórico de profundidade. Ao permitir que diferentes subespaços de características atendam a diferentes camadas de forma independente, o MHAR escala efetivamente onde o roteamento de consulta única falha.

Os autores enfatizam que:

• A melhoria não se deve a parâmetros ou computação adicionados, mas sim à remoção de um gargalo estrutural.
• O método é prático em escala devido aos kernels fundidos que mitigam a natureza limitada por memória do roteamento de profundidade.
• A abordagem é versátil, funcionando tanto do zero quanto via treinamento intermediário de preservação de identidade.
• A necessidade de roteamento multi-cabeça é uma consequência direta da largura do modelo; conforme os modelos crescem, o "compromisso forçado" de uma única consulta torna-se um passivo, tornando o MHAR um componente crítico para arquiteturas conscientes de profundidade.