Grokking

Grokking é uma generalização tardia: uma rede neural alcança acurácia perfeita nos dados de treinamento logo no início, permanece no nível do acaso nos dados de validação por um longo período e, então, muito depois de parecer ter se sobreajustado, generaliza de repente. Power et al. (2022), da OpenAI, deram nome ao efeito após treinar pequenos transformers em tarefas algorítmicas, como aritmética modular, nas quais a acurácia na validação passou do nível do acaso a 100% muito depois de a acurácia no treinamento ter chegado ao platô. A palavra vem do romance Um Estranho numa Terra Estranha, de Robert Heinlein, em que grok significa compreender algo completamente.

Como o grokking se desenrola#

A adição modular é a tarefa clássica de grokking: a rede recebe pares (a, b) e deve produzir (a + b) mod p, com apenas uma fração de todos os pares disponível para treinamento. A execução abaixo treina um pequeno perceptron multicamadas com o otimizador AdamW do PyTorch usando 30% dos 9.409 pares para p = 97 e termina em menos de um minuto numa CPU:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
p = 97  # task: predict (a + b) mod 97 from one-hot encodings of a and b
a, b = torch.cartesian_prod(torch.arange(p), torch.arange(p)).T
x, y = torch.cat([F.one_hot(a, p), F.one_hot(b, p)], dim=1).float(), (a + b) % p
idx = torch.randperm(p * p)
train, val = idx[: p * p * 3 // 10], idx[p * p * 3 // 10 :]  # train on 30% of all pairs

model = torch.nn.Sequential(torch.nn.Linear(2 * p, 256), torch.nn.ReLU(), torch.nn.Linear(256, p))
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1.0)  # strong weight decay
for step in range(1, 20001):
    loss = F.cross_entropy(model(x[train]), y[train])
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 4000 == 0:
        train_acc, val_acc = [(model(x[i]).argmax(1) == y[i]).float().mean().item() for i in (train, val)]
        print(f"step {step:5d}: train accuracy {train_acc:.2f}, validation accuracy {val_acc:.2f}")

Na etapa 4.000, a precisão no treinamento já é 1,00, enquanto a precisão na validação é 0,00. A precisão na validação sobe então para 0,16 na etapa 8.000, 0,54 na etapa 12.000, 0,84 na etapa 16.000 e 0,97 na etapa 20.000. Com weight_decay=0.0, a mesma execução mantém a precisão na validação em 0,00 durante todas as 20.000 etapas.

Nanda et al. (2023) explicaram o platô por meio da engenharia reversa de um transformer de uma camada treinado em adição modular. A rede aprende a transformar a adição em rotação ao redor de um círculo, usando características de Fourier discretas e identidades trigonométricas. O treinamento se divide em três fases contínuas: memorização, formação do circuito e limpeza. O circuito de generalização cresce gradualmente enquanto a precisão na validação permanece estável, e o salto acontece quando a perda de pesos remove a solução memorizada que ainda dominava as saídas.

O que faz o grokking acontecer#

• Regularização: a perda de pesos foi a intervenção mais eficaz nas experiências originais, e a execução acima não generaliza sem ela.
• Tamanho do conjunto de dados: frações menores de dados de treinamento exigem muito mais otimização antes de generalizar; abaixo de um determinado limite, a rede nunca generaliza dentro de um orçamento prático.
• Norma dos pesos: Omnigrok relaciona o grokking ao tamanho dos pesos, pois as soluções que generalizam têm uma norma dos pesos menor do que as soluções memorizadas. Partindo de pesos iniciais grandes, os autores induziram o grokking em tarefas de imagem, linguagem e moléculas; ao restringir a norma dos pesos, eliminaram-no nas tarefas algorítmicas.
• Estabilidade numérica: Prieto et al. (2025) descobriram que, sem regularização, o treinamento com entropia cruzada continua aumentando a escala dos logits até que erros de ponto flutuante na softmax impeçam a aprendizagem, fenômeno que chamam de colapso da Softmax. A ativação StableMax e o otimizador ⊥Grad produzem grokking sem perda de pesos.

Grokking vs. conceitos relacionados#

• Sobreajuste: uma diferença entre o desempenho no treinamento e na validação que persiste ou aumenta. No grokking, a diferença diminui por conta própria com a continuação do treinamento.
• Dupla descida: o erro de teste diminui, aumenta e volta a diminuir à medida que o tamanho do modelo ou o tempo de treinamento aumenta. A dupla descida por época é outra forma de melhoria tardia, na qual o erro de teste aumenta por algum tempo durante o treinamento antes de voltar a diminuir.
• Interpretabilidade mecanística: estudo dos algoritmos dentro das redes treinadas. O grokking em aritmética modular tornou-se um dos principais estudos de caso da área, porque foi possível identificar por completo o algoritmo aprendido.

Grokking na prática#

O grokking fica mais evidente em pequenos conjuntos de dados algorítmicos regidos por uma única regra exata. Ele também é relevante para o desenho de experimentos, pois a interrupção antecipada encerra uma execução quando a métrica de validação deixa de melhorar, o que interromperia uma execução com grokking durante o platô. O treinamento Ultralytics YOLO para após patience épocas sem melhoria na aptidão de validação, 100 por padrão, e executa todas as épocas quando patience=0.