Weight Decay Regimes in Grokking Transformers: Cheap Online Diagnostics

Resumo Técnico: Regimes de Decaimento de Pesos em Transformers de Grokking

Declaração do Problema
Grokking refere-se ao fenômeno em que transformers treinados em tarefas algorítmicas (especificamente aritmética modular) exibem uma transição abrupta de memorização para generalização muito depois que a perda de treinamento satura. Embora estudos mecanísticos tenham identificado a formação de circuitos internos (por exemplo, representações de Fourier) e cabeças de indução como impulsionadores dessa transição, a literatura existente carece de um "diagnóstico online barato" para rastrear essas dinâmicas em tempo real. Além disso, o papel do decaimento de pesos ($\lambda$) como parâmetro de controle para esses regimes não foi mapeado sistematicamente através das escalas de modelos, nem a contribuição específica da coordenação das cabeças de atenção foi isolada dos efeitos da magnitude dos pesos.

Metodologia
Os autores treinam transformers apenas decodificadores em tarefas de aritmética modular (mod+, mod-, mod$\times$, mod$\div$) usando o otimizador AdamW. O estudo emprega uma varredura densa de valores de decaimento de pesos e varia as escalas de modelos (de 0,82M a 85M parâmetros) e arquiteturas (Transformers, MLPs, LSTMs, Mamba).

Componentes metodológicos-chave incluem:

1. Parâmetros de Ordem Online: Dois diagnósticos escalares são computados a cada passo de treinamento usando apenas ativações de atenção:
   • Média de Similaridade Cosseno entre Cabeças de Atenção ($\bar{s}(t)$): Mede a coordenação (sincronização) entre as cabeças de atenção.
   • Desvio Padrão da Entropia ($\sigma_H(t)$): Mede a dispersão da entropia de Shannon entre as cabeças, rastreamento da diferenciação.
2. Escala de Tamanho Finito: Os autores ajustam uma função logística à probabilidade de grokking ($P(\text{grok})$) versus decaimento de pesos para estimar um limiar crítico $\lambda_c$. Eles também ajustam uma lei de potência ao "tempo para grokking" ($t_{\text{grok}}$) para determinar o expoente crítico $\nu$.
3. Intervenções Causais: Para distinguir entre estrutura de padrão de cabeças e magnitude de pesos, os autores realizam experimentos pareados onde as cabeças de atenção são re-inicializadas versus onde as normas de pesos são cortadas no pico da fase de diferenciação.
4. Verificação Formal: Identidades diagnósticas (limites e propriedades algébricas) são verificadas por máquina em Lean 4 para garantir a boa formação matemática.

Principais Contribuições

1. Limiar Crítico Quantitativo: O estudo identifica um limiar crítico agudo de decaimento de pesos para o transformer canônico de 4 camadas e 8 cabeças em tarefas mod+: $\lambda_c = 0,0158$ (IC 95% [0,0109, 0,0200]).
2. Diagrama de Regimes: Um mapa de dois eixos ($\lambda$, $N$) revela três regimes distintos:
   • Memorização: $\lambda < \lambda_c$ (grokking próximo de zero).
   • Grokking Desenvolvimental: $\lambda \ge \lambda_c$ (atingindo ~100% de taxa de grok, com tempo para grokking diminuindo à medida que $\lambda$ aumenta).
   • Colapso: $\lambda = 10$ (todas as cabeças convergem para padrões idênticos, $\bar{s}=1,0$).
3. Dinâmicas de Duas Fases: A trajetória de treinamento é caracterizada por:
   • Fase 1 (Sincronização): As cabeças convergem ($\bar{s}$ sobe de ~0,93 para 0,995) coincidindo com o início do grokking.
   • Fase 2 (Diferenciação): As cabeças divergem ($\bar{s}$ cai para ~0,88, $\sigma_H$ aumenta) enquanto a precisão de teste permanece alta.
4. Dimensão por Cabeça como Modulador de Amplitude: A amplitude da diferenciação ($\sigma_H$) é empiricamente dominada pela dimensão por cabeça ($d/H$), saturando monotonicamente à medida que $d/H$ aumenta.
5. Escopo Trans-Arquitetural: A transição controlada pelo decaimento de pesos é observada em arquiteturas sem atenção (MLP, LSTM, Mamba), embora o limiar crítico $\lambda_c$ se desloque significativamente (abrangendo aproximadamente uma ordem de magnitude) dependendo da arquitetura.

Resultados

• Expoente de Escala: O tempo para grokking segue uma lei de potência $t_{\text{grok}} \propto (\lambda - \lambda_c)^{-\nu}$ com um expoente empírico $\nu = 0,757$ (IC 95% [0,725, 0,799]). Os autores observam explicitamente que este valor não corresponde aos expoentes de referência padrão (por exemplo, $\nu=1/2$ ou Ising 3D $\nu \approx 0,63$) sob sua densidade de grade atual, adiando a identificação da classe de universalidade para trabalhos futuros.
• Fragilidade de Horizonte Longo: Em horizontes estendidos (20.000 épocas), um colapso "anti-grokking" dependente da semente é observado em algumas trajetórias, caracterizado por um padrão de cinco estágios terminando em decaimento de precisão. Isso é identificado como uma fragilidade e não um ciclo universal.
• Isolamento Causal: Re-inicializar as cabeças de atenção reduz significativamente a amplitude da Fase 2 ($\sigma_H$) em comparação com controles, enquanto o corte de norma de pesos não o faz. Isso isola o efeito à estrutura do padrão das cabeças e não à magnitude dos pesos.
• Replicação Multi-Tarefa: O padrão de controle por decaimento de pesos mantém-se através de quatro operações modulares e duas escalas de modelo, confirmando a robustez do mapa de regimes além da adição.

Significado e Alegações
O artigo alega fornecer uma estrutura de diagnóstico barata e computável online para rastrear dinâmicas de grokking usando apenas ativações de atenção, complementando diagnósticos mais caros de paisagem de perda, como o coeficiente de aprendizado local refinado (rLLC). Estabelece o decaimento de pesos como um parâmetro de controle empírico primário que separa os regimes de memorização, grokking desenvolvimental e colapso.

Os autores são modestos em seu escopo:

• Eles alegam nenhuma classe de universalidade para o expoente observado $\nu \approx 0,76$, notando que os expoentes de referência testados caem fora de seus intervalos de confiança.
• Eles não alegam que o decaimento de pesos é necessário para generalização em todas as arquiteturas, apenas que atua como um parâmetro de controle dentro de arquiteturas de atenção padrão.
• Eles explicitamente adiam alegações de mecanismo causal sobre por que a diferenciação ocorre para trabalhos de acompanhamento, focando em vez disso em quando e como ela se manifesta.
• As descobertas são limitadas a aritmética modular em pequenos modelos de atenção de transformer; alegações sobre grandes modelos de linguagem ou universalidade mais ampla estão fora do escopo.

O trabalho se posiciona como uma contribuição fenomenológica e diagnóstica, fornecendo um "mapa de regimes" e "parâmetros de ordem" que complementam trabalhos teóricos concurrentes sobre escala de tamanho finito, contração de Lyapunov e relaxamento de vidro.