Resumo Técnico: Como Domar o Grokking: A Geometria da Representação como um Sinal de Controle

Declaração do Problema

O grokking é um fenômeno em redes neurais sobreparametrizadas onde os modelos inicialmente memorizam os dados de treinamento e apenas exibem uma generalização forte após uma otimização prolongada. Embora estudos recentes tenham vinculado o grokking ao weight decay, compressão de características, formação de circuitos e dinâmica de aprendizado de representação, os fatores específicos que governam o tempo e a emergência dessa generalização tardia permanecem incompletamente compreendidos. Um tema comum na literatura existente é que as representações aprendidas se comprimem progressivamente em estruturas de menor dimensão. No entanto, trabalhos anteriores tratam majoritariamente a geometria da representação como uma propriedade emergente, e não como um mecanismo de controle direto. Este artigo investiga se a manipulação explícita da geometria das representações ocultas pode servir como um sinal de controle para influenciar o início e a velocidade do grokking.

Metodologia

Os autores introduzem a Regularização de Dimensionalidade Geométrica (GeomDR), um regularizador espectral ao nível da representação projetado para controlar diretamente a dimensionalidade efetiva das representações ocultas durante o treinamento.

Caracterização Geométrica

Para uma matriz de representação de uma camada oculta $Z \in \mathbb{R}^{N \times d}$, o método primeiro centraliza os dados para obter $\tilde{Z}$. Em seguida, calcula a matriz de covariância empírica $C = \frac{1}{N-1}\tilde{Z}^\top\tilde{Z}$. A dimensionalidade efetiva ($D_{eff}$) é calculada usando a razão de participação dos autovalores ($\mu_i$) de $C$:
$$D_{eff} = \frac{(\sum_i \mu_i)^2}{\sum_i \mu_i^2}$$
Adicionalmente, uma distância de vizinhança local ($\rho$) é definida como a distância euclidiana média para os $k$ vizinhos mais próximos das representações normalizadas.

O Objetivo GeomDR

O GeomDR impõe uma penalidade na variância contida em direções além de um subespaço alvo de dimensão $d^*$. Dados os autovalores ordenados em ordem decrescente ($\mu_1 \ge \mu_2 \ge \dots \ge \mu_d$), a perda por camada é:
$$L_{GeomDR}^{(\ell)} = \sum_{j=d^*+1}^{d} \mu_j^{(\ell)}$$
O objetivo total de treinamento combina a perda da tarefa ($L_{task}$) e o regularizador geométrico:
$$L = L_{task} + \lambda(t)L_{GeomDR}$$
Crucialmente, o regularizador não é ativo desde o início. Ele é introduzido após uma fase inicial de treinamento sem restrições (etapa de ativação $t_s$) e sua força $\lambda(t)$ é aumentada gradualmente usando uma rampa de cosseno. Este design mimetiza a trajetória natural do grokking, onde a memorização precede a generalização.

Protocolo Experimental

O framework foi avaliado em três tarefas algorítmicas conhecidas por exibir grokking:

1. Adição Modular: Predizer $(a + b) \mod 97$.
2. Divisão Modular: Predizer $(a \cdot b^{-1}) \mod 97$.
3. Composição de Permutação: Predizer a composição de duas permutações em $S_5$.

Os experimentos utilizaram duas arquiteturas: um Perceptron Multicamadas (MLP) e um Transformer. O estudo variou sistematicamente os cronogramas de intervenção (força $\lambda_{max}$ e duração da rampa $T_{ramp}$), dimensionalidades alvo ($d^*$) e tempos de início de intervenção ($t_s$).

Principais Contribuições

1. Observação do Colapso Precedente: Os autores demonstram que o colapso da dimensionalidade precede consistentemente a transição da memorização para a generalização (grokking) em configurações de linha de base. Isso sugere que a compressão geométrica não é meramente um correlato, mas um motor mecanístico da transição.
2. Introdução do GeomDR: Eles propõem o GeomDR, um regularizador espectral que suprime diretamente o espectro de covariância fora de um subespaço alvo, controlando assim a dimensionalidade efetiva das representações ocultas.
3. Controle Sistemático das Dinâmicas: O artigo fornece um estudo empírico abrangente mostrando que intervenções geométricas podem alterar substancialmente as dinâmicas de generalização tardia. Dependendo do cronograma de intervenção e da dimensionalidade alvo, o GeomDR pode acelerar o grokking, retardá-lo ou suprimi-lo inteiramente.

Resultados

Aceleração do Grokking

Em todas as configurações avaliadas, o GeomDR alterou consistentemente as dinâmicas de grokking:

• Desempenho do MLP: Na adição modular, o GeomDR acelerou o grokking em até 52 vezes (reduzindo de 362K para ~7K passos). Melhorias semelhantes foram observadas para divisão modular (39x) e composição de permutação (~23x). Em todos os casos de MLP, o GeomDR alcançou uma taxa de sucesso de 100%, onde a linha de base era instável ou falhava.
• Desempenho do Transformer: Embora a dinâmica de grokking da linha de base para Transformers já fosse mais rápida que a dos MLPs, o GeomDR ainda proporcionou aceleração consistente (ex: ~2x na adição modular), mantendo uma taxa de sucesso de 100%.

Influência dos Hiperparâmetros

• Dimensionalidade Alvo ($d^*$): Dimensões alvo moderadas a grandes ($16 \le d^* \le 64$) produziram a maior aceleração. A compressão agressiva ($d^* \le 8$) frequentemente retardou ou desestabilizou o aprendizado.
• Tempo de Intervenção ($t_s$): Uma "janela temporal crítica" foi identificada. Intervenções aplicadas muito cedo (antes que representações úteis se formem) ou muito tarde (após a memorização estar rígida) foram menos eficazes. Tempos de ativação intermediários (ex: $t_s = 2000$ para MLPs) produziram o grokking mais rápido e estável.
• Cronograma: Forças de regularização e durações de rampa intermediárias geralmente performaram melhor; intervenções excessivamente fortes poderiam desestabilizar o treinamento.

Correlações Geométricas

O estudo confirmou que as intervenções que aceleraram o grokking também aceleraram o início do colapso da dimensionalidade. Nos modelos GeomDR, o intervalo temporal entre o colapso da dimensionalidade e a generalização bem-sucedida foi significamente reduzido em comparação com as linhas de base, sugerindo um acoplamento mais estreito entre a reorganização geométrica e a emergência da generalização.

Escalabilidade e Eficiência

Os autores demonstraram que o GeomDR permanece eficaz mesmo quando o espectro de covariância é estimado a partir de um subconjunto aleatório de representações (ex: 5%–25% dos dados), oferecendo uma aceleração de tempo de execução de 2.4x–2.7x sem afetar materialmente a aceleração do grokking.

Significância e Alegações

O artigo afirma que a geometria da representação serve como um efeito de controle eficaz para o grokking. Ao modificar diretamente o espectro de covariância das ativações ocultas, pesquisadores podem influenciar sistematicamente quando e quão rapidamente uma rede transita da memorização para a generalização.

Os autores postulam que este trabalho fornece evidências de que intervenções geométricas oferecem uma abordagem prática para estudar e influenciar a generalização tardia. Em vez de observar passivamente as mudanças representacionais, o método permite a manipulação direta das propriedades geométricas para testar hipóteses sobre dinâmicas de aprendizado. Os resultados sugerem que a relação entre a geometria da representação e o grokking é robusta através de diferentes tarefas e arquiteturas, embora a magnitude da melhoria varie (sendo mais pronunciada em MLPs do que em Transformers, provavelmente devido aos vieses indutivos mais fortes destes últimos).

O artigo conclui que, embora não estabeleça a dimensionalidade da representação como o único mecanismo causal, demonstra que intervenções geométricas são suficientes para alterar o tempo do grokking, oferecendo uma nova metodologia experimental para pesquisa em otimização e generalização.