A Pre-Training Analogue of Grokking in Language Models: Tracing Delayed Grammatical Generalization

Resumo Técnico: Um Análogo de Grokking em Pré-treinamento para Modelos de Linguagem

Declaração do Problema

O fenômeno do grokking — onde redes neurais transitam da memorização de dados de treinamento para a generalização para exemplos não vistos após um longo período de treinamento — tem sido extensamente estudado em configurações supervisionadas com divisões explícitas de treino/validação e muitos épocas (Power et al., 2022). No entanto, essa formulação clássica não se traduz diretamente para o pré-treinamento de Grandes Modelos de Linguagem (LLMs). No pré-treinamento, os modelos são otimizados para a predição do próximo token sobre um grande corpus não rotulado (ex: C4) com repetição limitada de dados e sem uma divisão natural de treino/validação supervisionada sobre tarefas de avaliação downstream. Consequentemente, não há um mecanismo padrão para medir se o modelo primeiro ajusta exemplos específicos e apenas mais tarde generaliza para variações não vistas da mesma tarefa durante a fase de pré-treinamento.

Metodologia

Para abordar essa lacuna, os autores propõem um framework baseado em exposição para construir uma divisão proxy de treino/validação para estudar a generalização tardia durante o pré-treinamento.

1. Divisão Proxy Baseada em Exposição:

   • O estudo utiliza o BLiMP (Warstadt et al., 2020), um conjunto de dados de pares mínimos contendo uma frase aceitável e uma inaceitável que visam fenômenos gramaticais específicos.
   • Para cada par mínimo, os autores extraem uma frase crítica: o menor intervalo contínuo que contém o contraste gramatical e o contexto necessário.
   • Divisão Proxy-Treino (Exposta): Exemplos onde a frase crítica aparece literalmente nos tokens do C4 vistos durante a janela de pré-treinamento analisada.
   • Divisão Proxy-Validação (Não Exposta): Exemplos onde a frase crítica não aparece literalmente na janela de pré-treinamento.
   • Esta divisão distingue com base na exposição literal, não se o modelo viu a regra gramatical subjacente em outros contextos.

2. Configuração Experimental:

   • Modelos: Modelos decodificadores apenas do tipo LLaMA com 35M, 60M e 130M de parâmetros.
   • Treinamento: Treinados no C4 em inglês por 2.000 passos (aprox. 200M de tokens não-padding) com checkpoints a cada 100 passos.
   • Avaliação: A acurácia é medida nos pares mínimos do BLiMP. O modelo está correto se atribuir uma probabilidade maior à frase aceitável.
   • Métrica: O Atraso de Generalização Tardia ($\Delta t$) é definido como a diferença nos passos de treinamento entre quando a divisão proxy-treino (exposta) e a divisão proxy-validação alcançam um limiar de 80% de acurácia.

3. Validação e Análise:

   • Teste de Divisão Aleatória Pareada: Para garantir que o atraso observado não seja um artefato do tamanho da divisão ou diferenças arbitrárias entre os exemplos, os autores geram 1.000 divisões aleatórias preservando os tamanhos originais dos grupos. Eles comparam o atraso baseado em exposição contra esta distribuição nula.
   • Análise de Representação: Eles constroem vetores de conceitos gramaticais ($v_i = h_{gram} - h_{ungram}$) para medir a separabilidade de sentenças gramaticais vs. gramaticais incorretas e computam o rank efetivo desses vetores para analisar a dimensionalidade do espaço de representação.
   • Análise de Atenção: Eles rastreiam os pesos de atenção do token crítico para o token de contexto necessário para resolver a dependência gramatical, medindo mudanças na força e entropia da atenção.

Principais Resultados

1. Observação de Generalização Tardia: Através de cinco fenômenos gramaticais (ex: Concordância Sujeito-Verbo, Existential There Raising), a divisão proxy-treino (exposta) consistentemente alcança alta acurácia (80%) antes da divisão proxy-validação (não exposta). Isso confirma um análogo de grokking no pré-treinamento, onde os modelos performam bem em exemplos diretamente expostos antes de generalizar para exemplos estruturalmente similares, mas sem exposição literal.
2. Validação do Atraso: Testes de permutação de divisão aleatória pareada confirmam que este atraso é estatisticamente significativo ($p < 0,05$) e não pode ser explicado apenas pelo tamanho da divisão. O atraso está especificamente ligado à exposição literal da frase crítica.
3. Mudanças nos Vetores de Conceitos:
   • Separabilidade: O AUROC dos vetores de conceitos gramaticais geralmente aumenta durante o pré-treinamento, indicando melhor separabilidade entre sentenças gramaticais e gramaticais incorretas.
   • Dimensionalidade: Após a transição de generalização tardia, o rank efetivo dos vetores de conceitos gramaticais aumenta. Isso sugere que os conceitos gramaticais tornam-se distribuídos em um subespaço de maior dimensão, em vez de colapsarem em poucas direções dominantes.
4. Reorganização da Atenção: A transição para a generalização é acompanhada por mudanças localizadas na atenção. Especificamente, a atenção do token crítico para o token de contexto relevante aumenta em um pequeno número de cabeças de atenção específicas. Essas mudanças são específicas ao fenômeno (ex: diferentes camadas/cabeças para Concordância Sujeito-Verbo vs. Concordância Determiner-Noun) em vez de um deslocamento uniforme em todo o modelo.

Principais Contribuições

1. Framework: Introdução de um framework baseado em exposição que permite o estudo de dinâmicas semelhantes ao grokking durante o pré-treinamento de LLMs sem supervisão explícita downstream.
2. Evidência Empírica: Demonstração de que, através de cinco fenômenos gramaticais, exemplos com exposição literal da frase crítica alcançam alta acurácia significativamente antes daqueles sem tal exposição.
3. Robustez: Validação de que o atraso observado não é um artefato de divisão de dataset via testes de divisão aleatória pareada.
4. Insight Mecanístico: Evidência de que a transição comportamental para a generalização é acompanhada por mudanças específicas na geometria do vetor de conceito gramatical (aumento de dimensionalidade) e reorganização localizada da atenção.

Significância e Alegações

O artigo afirma fornecer um "análogo de pré-treinamento da configuração clássica de grokking". Ao utilizar a exposição do corpus como um proxy para os dados de treinamento, os autores demonstram que a transição tardia do ajuste para a generalização é observável mesmo no cenário não supervisionado de predição do próximo token do pré-treinamento.

Os autores enfatizam que suas descobertas sugerem que a exposição do corpus pode ser usada para estudar a generalização tardia do tipo grokking. As análises mecanísticas revelam que essa mudança comportamental não é meramente uma mudança na métrica de performance, mas está enraizada nas representações internas do modelo: os conceitos gramaticais tornam-se mais separáveis e ocupam subespaços de maior dimensão, enquanto os mecanismos de atenção se reorganizam para focar mais fortemente em tokens de contexto relevantes para a dependência.

O estudo permanece modesto em seu escopo, observando que a "divisão proxy-validação" representa exemplos sem exposição literal, não necessariamente exemplos com gramática inteiramente inédita. O framework é atualmente limitado a fenômenos gramaticais onde frases críticas podem ser definidas como intervalos sub-sentenciais.