Fast Generalization after Interpolation via Critically Damped Momentum Optimization

Resumo Técnico: Generalização Rápida após Interpolação via Otimização de Momento Criticamente Amortecido

Enunciado do Problema
Em regimes de alta dimensionalidade e baixa amostragem (comuns em genômica, medicina e finanças), os modelos de aprendizado de máquina frequentemente alcançam um desempenho de treinamento quase perfeito (interpolação), enquanto falham em generalizar para dados não vistos. Esse hiato surge porque o cenário de otimização contém muitas configurações de parâmetros que produzem perdas de treinamento idênticas, mas diferem significativamente em complexidade e desempenho de teste. Embora métodos de regularização (ex: decaimento de peso) sejam tradicionalmente usados para enviesar a otimização em direção a soluções mais simples, o fenômeno do "grokking" revela um descompasso temporal: os modelos frequentemente memorizam os dados primeiro e só então transitam para a generalização muito mais tarde. Essa transição tardia é computacionalmente ineficiente, exigindo ordens de magnitude mais épocas para a generalização do que para o ajuste (fitting). O problema central é como selecionar eficientemente uma solução interpolante de baixa complexidade imediatamente após o modelo atingir o limiar de interpolação, em vez de depender de dinâmicas implícitas lentas.

Metodologia
Os autores propõem o GROKtimizer, uma estratégia de otimização bifásica projetada para separar explicitamente a fase de ajuste da fase de generalização.

1. Estrutura Teórica:

   • Dinâmica Pós-Interpolação: O artigo modela o regime pós-interpolação como um sistema dinâmico quadrático local. Sob a suposição de que o cenário de perda próximo a uma solução interpolante pode ser aproximado por uma forma quadrática envolvendo a matriz Hessiana ($H$) e um termo de regularização $L_2$ ($\lambda$), a dinâmica ao longo das direções planas (espaço nulo de $H$) assemelha-se a um oscilador amortecido.
   • Amortecimento Crítico: Os autores demonstram que o Momento Criticamente Amortecido (CDM) fornece a trajetória ótima para minimizar a norma dos parâmetros dentro deste poço plano. Ao ajustar o parâmetro de momento $\beta$ para $\beta = 1 - 2\sqrt{\lambda\eta}$ (onde $\eta$ é a taxa de aprendizado), o otimizador alcança o amortecimento crítico.
   • Aceleração de Convergência: Teoricamente, o CDM oferece uma aceleração quadrática sobre o Gradiente Descendente (GD) clássico neste regime. Enquanto o GD converge com uma escala de tempo proporcional a $\lambda^{-1}$, o CDM converge com uma escala de tempo proporcional a $\lambda^{-1/2}$. Isso resulta em uma redução provável de ordem quadrática no número de épocas necessárias para atingir a solução de norma mínima.
   • Estimativa da Taxa de Aprendizado: Para implementar o CDM na prática sem calcular a Hessiana completa, os autores propõem um método de Iteração de Potência corrigido por Aitken para estimar o maior autovalor ($\lambda_{max}$) da Hessiana. Isso permite o cálculo online da taxa de aprendizado ótima $\eta_{max}$ no limite da estabilidade.

2. Implementação Algorítmica (GROKtimizer):

   • Fase 1 (Interpolação): O otimizador executa com zero decaimento de peso para priorizar a convergência rápida ao limiar de interpolação. Esta fase utiliza o momento padrão para ajustar os dados de treinamento rapidamente.
   • Fase 2 (Generalização): Uma vez alcançada a interpolação, o otimizador muda para o Momento Criticamente Amortecido com decaimento de peso. O parâmetro de momento é ajustado para a condição de amortecimento crítico derivada na teoria, e o decaimento de peso é aplicado para conduzir os parâmetros em direção à solução de norma mínima dentro do manifold interpolante.

Principais Contribuições

1. Caracterização do Regime Pós-Interpolação: Os autores caracterizam formalmente a fase de generalização tardia como um problema de oscilador amortecido ao longo das direções planas do cenário de perda, vinculando o fenômeno do grokking à geometria do espaço nulo da Hessiana.
2. Aceleração Teórica: Eles provam que o momento criticamente amortecido acelera a convergência para soluções interpolantes de baixa norma de forma quadrática mais rápida que o gradiente descendente padrão, estabelecendo a otimalidade entre métodos de primeira ordem sob o modelo quadrático local.
3. GROKtimizer: Introduzem um cronograma de otimização bifásico que alterna dinamicamente do ajuste rápido não regularizado para a minimização de norma criticamente amortecida.
4. Validação Empírica: O método é validado em benchmarks sintéticos de grokking e em conjuntos de dados reais de alta dimensionalidade, demonstrando melhorias consistentes sobre otimizadores padrão (Adam, AdamW, SGD, Muon).

Resultados Experimentais
O artigo avalia o GROKtimizer em vários benchmarks:

• Benchmarks Sintéticos: Em tarefas como Adição Modular, Adição Binária e Paridade Esparsa, o GROKtimizer supera significamente os baselines. Por exemplo, na tarefa Gaussiana, ele reduziu a perda de validação em várias ordens de magnitude em comparação com o baseline mais forte.
• Dados Reais de Alta Dimensionalidade: Em tarefas de classificação biomédica (Leukemia, TCGA) e previsão de propriedades moleculares (QM9), o GROKtimizer alcançou a maior acurácia e a menor perda de validação. Os ganhos foram mais pronunciados em cenários de escassez de dados, onde o número de características excede amplamente o número de amostras.
• Eficiência de Amostra: Em experimentos de MNIST com poucos dados (10–50 amostras por classe), o GROKtimizer mostrou a maior vantagem, confirmando que a minimização explícita da norma melhora a eficiência de amostragem quando o espaço de solução é subdeterminado.
• Pré-treinamento de Modelos de Linguagem: Experimentos em modelos WikiText2 e BabyLM mostraram que, embora a correlação direta entre a norma dos parâmetros e a perda de validação seja menos clara em transformers (devido a simetrias de escala e atualizações adaptativas), o padrão bifásico ainda se mantém. A mudança para a fase de minimização de norma após a interpolação produziu reduções adicionais na perda de validação.

Significância e Alegações
O artigo alega que o GROKtimizer oferece uma solução fundamentada para o problema do "grokking", tratando o treinamento pós-interpolação como um problema de otimização explícito em vez de um período de espera passivo. Ao alavancar a geometria do cenário de perda, o método acelera a seleção de soluções de baixa complexidade.

Os autores reconciliam suas descobertas com a hipótese de mínimos planos (flat-minima), argumentando que a dimensionalidade do espaço nulo da Hessiana serve como uma medida robusta de achatamento e complexidade. Eles sugerem que minimizar a norma dentro deste espaço nulo é equivalente a navegar em direção a mínimos mais planos e generalizáveis.

Limitações
Os autores reconhecem algumas limitações:

• A abordagem é projetada primariamente para regimes de restrição de dados e pode ser menos crítica em pré-treinamentos de larga escala de visão ou linguagem onde a abundância de dados domina o desempenho.
• O método assume que o modelo possui capacidade suficiente para atingir a interpolação; não é uma solução para modelos subparametrizados.
• A análise teórica depende de uma aproximação quadrática local do cenário de perda, embora os resultados empíricos sugiram que o método é robusto a desvios desta suposição.
• A norma dos parâmetros é usada como um proxy para complexidade, o que pode não capturar todos os aspectos da generalização do modelo em cada arquitetura.