Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

Resumo Técnico: Autoatenção Exclusiva Além do AdamW

Declaração do Problema
Os mecanismos de autoatenção padrão em Transformers permitem que a saída para um token na posição $i$ inclua uma mistura ponderada de vetores de valor, incluindo o próprio vetor de valor do token ($v_i$). Embora expressiva, essa configuração pode criar um "atalho representacional" onde a saída da atenção permanece fortemente alinhada com $v_i$, mesmo que os caminhos residuais já preservem informações locais do token. Este fenômeno, denominado "viés de similaridade de atenção", foi abordado por Zhai [2] através da Autoatenção Exclusiva (XSA), que remove explicitamente o componente da saída da atenção paralelo ao próprio vetor de valor do token.

No entanto, o estudo original da XSA avaliou o método exclusivamente dentro de uma pilha de treinamento NanoGPT utilizando o otimizador AdamW. Os autores do trabalho original identificaram explicitamente a compatibilidade da XSA com outros otimizadores, como o Muon, como uma questão em aberto. Esta é uma lacuna crítica porque a pesquisa moderna de modelos de linguagem de pequena escala (ex: a pilha nanochat) utiliza cada vez mais estratégias de otimização híbridas (Muon para parâmetros de matrizes, AdamW para embeddings) e características arquiteturais como embeddings de valor e codificação de posição rotativa (RoPE). Não está claro se as restrições geométricas da XSA permanecem estáveis ou benéficas quando transferidas para estes regimes de treinamento modernos baseados em Muon.

Metodologia
O estudo conduz um experimento de reprodução e transferência controlado usando a pilha de treinamento nanochat. A metodologia envolve:

1. Implementação: Uma forma parametrizada de XSA é integrada ao nanochat. A operação é definida como $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$, onde $y_i$ é a saída de atenção padrão, $v_i$ é o vetor de valor do token, e $\alpha \in [0, 1]$ controla a força da projeção. Em $\alpha=1$, o componente autoalinhado é totalmente removido.
2. Validação de Correção: A implementação é validada numericamente para garantir que, em $\alpha=0$, o caminho XSA seja equivalente à atenção padrão. Execuções curtas de treinamento verificam se o baseline e os caminhos $\alpha=0$ exibem um comportamento de perda inicial correspondente.
3. Configuração Experimental: Dois níveis de profundidade de modelo são avaliados sob condições de treinamento equivalentes (mesma ordenação de dados, família de otimizadores, cronograma de taxa de aprendizado e comprimento de contexto):
   • d12: 12 camadas, dimensão oculta 768, 12 cabeças.
   • d24: 24 camadas, dimensão oculta 1024, 16 cabeças.
   • Otimizador: Parâmetros de matriz são otimizados com Muon; os parâmetros restantes usam AdamW.
   • Dados: ClimbMix-400B.
4. Métricas: O estudo monitora a perda de treinamento, bits-por-byte (bpb) de validação, pontuações agregadas de downstream (Base CORE e ChatCORE) e diagnósticos de atenção (alinhamento de cosseno entre a saída e o auto-valor, comportamento diagonal).

Principais Contribuições

• Reprodução da Pilha Muon: A primeira integração da XSA em uma pilha de treinamento baseada em Muon, abordando diretamente a questão de compatibilidade de otimizador deixada pelo estudo original.
• Implementação Parametrizada: Uma implementação flexível que permite atenção padrão, XSA total e projeção parcial via um parâmetro de força $\alpha$, juntamente com controles de seleção de camada.
• Correção e Instrumentação: Validação do limite $\alpha=0$ e correção de artefatos de registro (convenções de índice de passo e suavização) que anteriormente distorciam as comparações iniciais de treinamento.
• Avaliação de Duas Profundidades: Uma análise comparativa nas profundidades d12 e d24 usando um conjunto abrangente de métricas.
• Análise Mecanística e de Sistemas: Diagnósticos de autoalinhamento de atenção e uma avaliação do overhead de tempo de execução.

Resultados
Os resultados empíricos indicam que a XSA é estável sob a pilha baseada em Muon, mas produz benefícios downstream mistos dependendo da profundidade do modelo e da métrica específica:

• Estabilidade de Otimização: As trajetórias de perda de treinamento para XSA e modelos baseline permanecem próximas em ambas as profundidades. Não há evidência de divergência, NaNs ou instabilidade, confirmando que a XSA não interrompe a geometria de atualização do Muon.
• Desempenho de Validação (Bits-per-Byte): A XSA melhora consistentemente o bpb de validação retido:
  • d12: Reduzido de 0.8484 para 0.8457 (melhoria relativa de 0,32%).
  • d24: Reduzido de 0.7193 para 0.7171 (melhoria relativa de 0,31%).
• Métricas Downstream (CORE):
  • d24: A XSA melhora tanto o Base CORE (0.2593 para 0.2606, +0,50%) quanto o ChatCORE (0.3638 para 0.3682, +1,21%).
  • d12: Embora o bpb melhore, o Base CORE diminui de 0.1602 para 0.1508 (-5,87%).
• Diagnósticos Mecanísticos: Diagnósticos de atenção mostram redução no autoalinhamento e redução no comportamento diagonal sob XSA, consistente com o mecanismo pretendido de promover a mistura contextual.
• Overhead de Tempo de Execução: O overhead é caracterizado como modesto. Embora uma estimativa inicial tenha sugerido ~5%, medições repetidas de throughput colocaram a diferença próxima à variabilidade do sistema em execuções curtas, sugerindo que não há penalidade proibitiva.

Significância e Alegações
O artigo conclui que a Autoatenção Exclusiva é um viés indutivo arquitetural compatível para pilhas de treinamento baseadas em Muon, mas sua utilidade não é universal. A principal significância reside em demonstrar que a restrição representacional da XSA se transfere limpa para um regime de otimização moderno sem desestabilizar o treinamento.

No entanto, o autor enfatiza que os ganhos de downstream são dependentes de escala e métrica. A melhoria no bpb de validação em ambas as profundidades sugere um benefício representacional, mas a degradação do Base CORE em d12 indica que os benefícios de remover o autoalinhamento podem depender da profundidade do modelo, da capacidade dos caminhos residuais/de valor para preservar informações locais, ou da tarefa de avaliação específica. O estudo não afirma que a XSA é uma substituição universalmente superior para a atenção padrão, mas sim uma restrição estável e de baixa interrupção que oferece benefícios em contextos específicos (notadamente em d24 neste estudo).

O trabalho serve como uma verificação de robustez em vez de uma reprodução byte a byte, destacando as complexas interações entre modificações arquiteturais, geometria de otimizador e métricas de avaliação na pesquisa moderna de modelos de linguagem.