Resumo Técnico: Duas Velocidades de Aprendizado: Uma Decomposição Representação-Leitura de Grokking e Double Descent

Declaração do Problema

Sinais padrão para monitorar o treinamento de redes neurais profundas, como perda de treinamento e precisão, frequentemente falham em capturar a dinâmica complexa da generalização. Dois fenômenos proeminentes complicam esse cenário:

1. Grokking: Um modelo alcança desempenho de treinamento quase perfeito enquanto a precisão de teste permanece em níveis de acaso por um período prolongado, seguida por um surgimento abrupto da generalização.
2. Double Descent por Época: A perda de treinamento diminui monotonicamente, enquanto a perda ou erro de teste exibe uma subida e descida não monotônicas.

Explicações existentes são frequentemente específicas da tarefa (por exemplo, dependendo de estruturas algébricas em aritmética modular) ou dependem de transições de regime binárias (por exemplo, de aprendizado "preguiçoso" para "rico"). Há uma falta de um framework unificado, agnóstico à tarefa, capaz de diagnosticar esses fenômenos em arquiteturas e tarefas diversas, distinguindo generalização genuína de artefatos de aprendizado espúrios induzidos por receitas de treinamento não padrão.

Metodologia: Decomposição Representação-Leitura

Os autores propõem um framework que decompõe uma rede neural $f(x)$ em dois componentes distintos:

1. Codificador ($\phi_\theta$): Um mapa não linear do espaço de entrada para um espaço de representação de $N$ dimensões.
2. Leitura ($W$): Um mapa linear do espaço de representação para os logits de saída.

Essa decomposição permite o rastreamento separado do aprendizado de representação (mudanças em $\phi_\theta$) e da calibração da leitura (mudanças em $W$). Para analisar a interação entre esses componentes, os autores empregam quatro medidas diagnósticas agnósticas à tarefa derivadas da geometria representacional e da teoria do Kernel Tangente Neural (NTK):

1. Dimensão Crítica ($N_{crit}$): Quantifica a robustez da separabilidade linear no espaço de representação. É definida como a menor dimensão de projeção esperada necessária para preservar a separabilidade com probabilidade $\geq 0,5$. Uma $N_{crit}$ mais baixa indica "manifolds" menos emaranhados e maior generalizabilidade.
2. Medidas Geométricas Relevantes para a Tarefa (GLUE): Um conjunto de quatro métricas ($D, R, \rho_c, \rho_a$) derivadas da teoria da Geometria Ligada à Eficiência de Desemaranhamento (GLUE). Estas caracterizam a dimensão efetiva do manifold, o raio e o alinhamento de centros e eixos entre manifolds condicionados por rótulos.
3. Precisão de Sonda Linear: Um classificador linear regularizado treinado em representações congeladas. A lacuna entre a precisão da sonda e a precisão do modelo diagnostica o desalinhamento da leitura (sobreajuste ou subajuste em relação às capacidades do codificador).
4. Alinhamento NTK-Rótulo da Última Camada: Mede o alinhamento entre a matriz Gram de características (última camada NTK) e o kernel de rótulos. Isso quantifica o quão bem a geometria de características atual suporta o fluxo de gradiente em direção a soluções discriminativas de rótulos, avaliado separadamente nos conjuntos de treinamento e teste.

Principais Contribuições e Resultados

1. Duas Velocidades de Aprendizado no Grokking

Aplicando a decomposição a tarefas canônicas de grokking (adição modular, composição de permutação, paridade esparsa) em arquiteturas MLP e Transformer, os autores refutam a conta estrita "de preguiçoso para rico", que postula que o aprendizado de representação está dormente até o início do grokking.

• Aprendizado Contínuo de Representação: $N_{crit}$ diminui gradual e continuamente ao longo do treinamento, incluindo a fase pré-grokking. O aprendizado de representação é lento, mas não ausente; procede continuamente em vez de como uma transição de fase súbita.
• Leitura Viciada para Treinamento: Antes do início do grokking, a direção de aprendizado da leitura é sistematicamente viciada em direção aos dados de treinamento. A lacuna de alinhamento NTK (treino menos teste) é grande e positiva, indicando que a leitura sobreajusta o conjunto de treinamento enquanto as representações ainda não são suficientemente robustas para suportar a generalização.
• Mecanismo: O grokking surge não de uma mudança súbita nos regimes de aprendizado, mas da equalização gradual da direção da leitura à medida que a qualidade da representação ($N_{crit}$) eventualmente alcança a velocidade da calibração da leitura.

2. Diferenças Arquiteturais em Mecanismos Geométricos

O framework revela que diferentes arquiteturas alcançam a generalização através de caminhos geométricos distintos.

• MLPs vs. Transformers: Em tarefas de adição modular e composição de permutação, MLPs e Transformers exibem tendências opostas no alinhamento de centros ($\rho_c$) e no alinhamento de eixos ($\rho_a$). Enquanto ambas as arquiteturas reduzem a dimensão e o raio do manifold, elas manipulam a geometria interclasse de maneira diferente para alcançar o mesmo comportamento superficial de grokking.

3. Diagnóstico de Aprendizado Espúrio

O framework distingue com sucesso a generalização genuína de artefatos causados por receitas de treinamento não padrão.

• Artefato de Grokking no MNIST: Um exemplo de grokking anteriormente relatado no MNIST (induzido por normas de peso iniciais altas e conjuntos de treinamento pequenos) é diagnosticado como espúrio. O framework revela degradação de representação (aumento de $N_{crit}$ e diminuição da precisão da sonda linear) e sobreajuste da leitura durante a fase "pré-grokking". A aparente generalização atrasada é um artefato da leitura eventualmente corrigindo uma representação degradada, em vez de uma melhoria genuína no aprendizado de características.
• Double Descent por Época: Em cenários de double descent induzidos por ruído de rótulos, o framework identifica um processo de dois estágios: uma fase inicial de deterioração da representação seguida por uma fase onde as atualizações da leitura são dominadas por sinais espúrios (rótulos ruidosos). Remover os ingredientes não padrão (por exemplo, ruído de rótulos) elimina tanto o padrão de double descent quanto as assinaturas diagnósticas.

4. Framework de Análise Top-Down

Os autores posicionam essa decomposição como um complemento "top-down" à interpretabilidade mecânica "bottom-up" (por exemplo, descoberta de circuitos). Ao caracterizar a estrutura global das representações (geometria do manifold, separabilidade, alinhamento), o framework fornece observáveis consistentes através das divisões treino/teste e tarefas diversas, oferecendo um sinal confiável para monitorar a dinâmica de aprendizado onde curvas de perda padrão falham.

Significância

O artigo estabelece a decomposição representação-leitura como um framework unificado e agnóstico à tarefa para entender a dinâmica de aprendizado. Sua principal significância reside em:

• Refinamento de Contas Teóricas: Vai além da dicotomia binária preguiçoso-rico, demonstrando que o aprendizado de representação é contínuo e que o grokking é impulsionado pelas velocidades relativas da melhoria da representação versus a calibração da leitura.
• Utilidade Diagnóstica: Fornece assinaturas específicas para identificar dinâmicas de aprendizado espúrias (por exemplo, degradação de representação ou desalinhamento da leitura) que podem ser confundidas com fenômenos de aprendizado genuínos em configurações de treinamento não padrão.
• Generalizabilidade: O framework aplica-se consistentemente em tarefas algorítmicas, conjuntos de dados do mundo real (MNIST) e arquiteturas (MLP, Transformer, ResNet), oferecendo uma ferramenta principiante para pesquisa de interpretabilidade e monitoramento de modelos críticos para segurança.

Os autores concluem que, embora o framework seja atualmente aplicado na última camada linear, ele oferece uma base robusta para futuras investigações em camadas intermediárias, estruturas específicas da tarefa e cenários de aprendizado contínuo.