Post-Grokking Collapse at the Representation-Readout Interface in Muon-Trained Transformers

Resumo Técnico: Colapso Pós-Grokking na Interface de Representação–Leitura em Transformers Treinados com Muon

Declaração do Problema
O artigo investiga uma instabilidade específica que ocorre em modelos Transformer treinados em aritmética modular (especificamente adição e subtração modular) após atingirem o "grokking" — o fenômeno onde a generalização emerge muito tempo após a acurácia de treinamento ter saturado. Embora trabalhos anteriores tenham estabelecido que o Muon (um otimizador ortogonalizante) acelera o grokking em comparação ao AdamW, este estudo relata que modelos treinados com Muon frequentemente sofrem um colapso recorrente e abrupto na performance de generalização pós-grokking. Crucialmente, essa falha não se limita ao Muon; modelos AdamW também exibem instabilidade. O artigo demonstra que nem o otimizador nem o AdamW são estritamente estáveis em cinco sementes (seeds): a linha de base AdamW selecionada registra avaliações abaixo do limiar em quatro de cinco sementes (atingindo um mínimo de 27,59%), enquanto o Muon registra tais avaliações em todas as cinco. O problema central é que esses modelos perdem a capacidade de generalizar, apesar de o conjunto de treinamento permanecer resolvido, e as medidas padrão de progresso mecanístico (suporte espectral e distribuição de potência) falham em detectar a falha, reportando um circuito intacto mesmo quando a acurácia cai para níveis próximos ao acaso.

Metodologia
Os autores empregam um regime de treinamento de otimizador dividido (split-optimizer), onde as matrizes de pesos ocultos são atualizadas pelo Muon, enquanto os embeddings de token/posição e o unembedding de saída (leitura/readout) são atualizados pelo AdamW. O estudo utiliza tarefas de aritmética modular ($a+b \pmod p$ e $a-b \pmod p$) com $p=113$ (e variantes com $p=97$) em Transformers do tipo decoder-only de várias profundidades (1, 2 e 4).

Componentes metodológicos principais incluem:

1. Ramificação Correspondente (Matched Branching): As trajetórias são ramificadas a partir de estados de parâmetros bit-idênticos para isolar fatores causais. Um ramo continua com a otimização padrão, enquanto outros congelam grupos específicos de parâmetros (matrizes ocultas, embeddings ou readout) para testar a necessidade.
2. Decomposição de Fourier: Os autores decompõem os estados residuais internos do modelo em modos de Fourier 2D sobre a grade dos operandos. Eles definem "famílias" específicas de modos (ex: a família da adição $(k, k)$, família da subtração $(k, -k)$) e usam intervenções de filtragem para medir a suficiência e a necessidade dessas famílias para a performance da tarefa.
3. Suíte de Intervenção: O estudo aplica intervenções como dispersão de fase (phase scrambling), realocação de frequência e substituição de cross-readout para determinar a natureza exata do algoritmo aprendido e a fonte do desalinhamento.
4. Ablação: O caminho de ortogonalização normalizada do Muon é ablado para testar seu papel na dispersão espectral e instabilidade observadas.

Principais Descobertas e Resultados

• A Natureza do Colapso: A falha é um "desalinhamento oculto–leitura" (hidden–readout misalignment). À medida que a perda de treinamento se aproxima de zero, o gradiente desaparece, mas os dois otimizadores respondem de forma diferente. O passo de atualização do Muon torna-se amplamente independente da magnitude do gradiente (elasticidade $\approx -0,03$), enquanto o do AdamW permanece proporcional (elasticidade $\approx +1,5$). Isso faz com que a representação oculta e a leitura se afastem a uma taxa de 8,0 vezes por parâmetro. O conjunto de treinamento permanece resolvido, mas a leitura não consegue mais decodificar a representação com a qual foi treinada.
• Mascaramento vs. Falha de Circuito: O artigo distingue dois modos de falha. Em uma falha de circuito, a própria família de Fourier deixa de computar a tarefa. Em um mascaramento de circuito, a família alinhada à tarefa (ex: a família da adição) continua a resolver a tarefa perfeitamente de forma isolada (100% de acurácia), mas sua contribuição é sobrecarregada por um remanescente adversário no modelo completo. Os autores mostram que reescalar apenas o componente alinhado à tarefa restaura a acurácia total do modelo, provando que o circuito estava intacto, mas foi "vencido em votação".
• Diferenças entre Otimizadores: O Muon atinge o grokking significativamente mais rápido (mediana de 13.011 passos vs. 30.486 para AdamW), mas é universalmente instável em todas as configurações testadas. O AdamW também é instável em cinco sementes, embora a severidade do colapso (medida pela contagem de avaliações abaixo do limiar) seja dois ordens de magnitude menor do que a do Muon. Em um varrimento de semente única, as configurações AdamW mostraram instabilidade condicional onde a severidade escalava com a taxa de aprendizado, mas a replicação em cinco sementes confirmou que a linha de base "estável" selecionada não era estritamente estável entre as sementes. A instabilidade no Muon é atribuída à sua regra de atualização, que espalha a representação por um amplo espectro (326 pares conjugados efetivos vs. ~5 para AdamW), deixando a família alinhada à tarefa com uma parcela menor do poder total.
• Seleção de Família pela Tarefa: A tarefa, não o otimizador, seleciona a família de Fourier. Modelos treinados em adição usam os modos $(k, k)$; modelos treinados em subtração usam $(k, -k)$. A base dentro da família é exata e não compartilhada; dois modelos treinados a partir da mesma semente divergem para usar frequências e fases diferentes, e suas leituras tornam-se quase ortogonais, falhando em decodificar as representações um do outro.
• Profundidade e Alinhamento: Em modelos mais profundos, o código de Fourier é escrito por um MLP no penúltimo bloco, mas só se torna "legível" para o unembedding no bloco final. O desalinhamento ocorre especificamente na interface entre o resíduo do bloco final e o unembedding.
• Prevenção: Congelar ou as matrizes ocultas ou os parâmetros não ocultos (embeddings e readout) após a formação do circuito previne completamente o colapso. O "Muon Estável", onde os embeddings e a leitura são congelados pós-grokking, alcança estabilidade estrita através de 451.400 passos pós-grokking e cinco sementes, enquanto o Muon padrão falha repetidamente.

Significância e Alegações
O artigo afirma identificar uma limitação fundamental no treinamento com otimizadores divididos e na dinâmica pós-grokking: o fluxo residual não possui uma base privilegiada e, quando a superfície de perda se achata, a perda não fornece pressão para manter o alinhamento específico entre a representação e a leitura. Os autores argumentam que:

1. Medidas padrão são insuficientes: Medidas espectrais (suporte e distribuição de potência) podem permanecer invariantes durante um colapso porque a falha é uma mudança de base (desalinhamento) e não uma perda do próprio algoritmo.
2. Grokking é mascaramento reversível: A transição da memorização para a generalização é o inverso da transição da generalização para o colapso; ambos envolvem o aumento ou a queda da participação da família alinhada à tarefa em relação ao remanescente adversário.
3. A interface é o locus da falha: A instabilidade não é uma propriedade da representação oculta ou da leitura isoladamente, mas de sua identificação conjunta. Fixar um lado da interface estabiliza o sistema.
4. A velocidade do Muon tem um custo: Embora o Muon acelere a descoberta do circuito, sua regra de atualização cria uma representação que é mais suscetível ao desalinhamento com a leitura sob condições normais de treinamento, necessitando de intervenção (congelamento) para estabilidade de longo prazo.

O estudo não propõe novas aplicações, mas oferece um relato mecanístico de um modo de falha específico em interpretabilidade mecanística, destacando que um modelo pode possuir um circuito correto, mas falhar em utilizá-lo devido ao desvio representacional.