Grokking

Il grokking è una generalizzazione ritardata: una rete neurale raggiunge presto un'accuratezza perfetta sui dati di addestramento, resta a livello casuale sui dati di validazione per un lungo periodo e poi, molto tempo dopo essere apparentemente andata in overfitting, generalizza all'improvviso. Power et al. (2022) di OpenAI hanno dato questo nome all'effetto dopo aver addestrato piccoli transformer su attività algoritmiche come l'aritmetica modulare, in cui l'accuratezza di validazione è passata dal livello casuale a quella perfetta molto tempo dopo che l'accuratezza di addestramento si era stabilizzata. Il termine proviene dal romanzo di Robert Heinlein Straniero in terra straniera, in cui grok significa comprendere qualcosa completamente.

Come si manifesta il grokking#

L'addizione modulare è l'attività classica per studiare il grokking: la rete riceve coppie (a, b) e deve restituire (a + b) mod p, avendo a disposizione solo una frazione di tutte le coppie per l'addestramento. Il codice qui sotto addestra un piccolo perceptron multistrato con l'ottimizzatore AdamW di PyTorch sul 30% delle 9.409 coppie per p = 97 e termina in meno di un minuto su una 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}")

Al passaggio 4.000, l'accuratezza di addestramento è già 1,00 mentre quella di validazione è 0,00. L'accuratezza di validazione sale poi a 0,16 al passaggio 8.000, a 0,54 al passaggio 12.000, a 0,84 al passaggio 16.000 e a 0,97 al passaggio 20.000. Con weight_decay=0.0, la stessa esecuzione mantiene un'accuratezza di validazione pari a 0,00 per tutti i 20.000 passaggi.

Nanda et al. (2023) hanno spiegato il plateau con il reverse engineering di un transformer a uno strato addestrato sull'addizione modulare. La rete impara a trasformare l'addizione in una rotazione attorno a un cerchio, usando caratteristiche discrete di Fourier e identità trigonometriche. L'addestramento si articola in tre fasi continue: memorizzazione, formazione del circuito e rifinitura. Il circuito che generalizza cresce gradualmente mentre l'accuratezza di validazione rimane costante; il salto arriva quando il decadimento dei pesi rimuove la soluzione memorizzata che ancora dominava gli output.

Cosa provoca il grokking#

• Regolarizzazione: Il decadimento dei pesi è stato l'intervento più efficace negli esperimenti originali e, senza di esso, l'esecuzione qui sopra non generalizza affatto.
• Dimensione del dataset: Frazioni di addestramento più piccole richiedono molta più ottimizzazione prima di generalizzare; sotto una certa soglia, la rete non generalizza mai entro un budget pratico.
• Norma dei pesi: Omnigrok collega il grokking alla dimensione dei pesi, poiché le soluzioni che generalizzano hanno una norma dei pesi inferiore a quelle basate sulla memorizzazione. Partendo da pesi iniziali grandi, gli autori hanno indotto il grokking in attività di visione, linguaggio e molecole; vincolando la norma dei pesi, lo hanno eliminato dalle attività algoritmiche.
• Stabilità numerica: Prieto et al. (2025) hanno rilevato che, senza regolarizzazione, l'addestramento con entropia incrociata continua ad aumentare la scala dei logit finché gli errori in virgola mobile nella softmax non arrestano l'apprendimento; chiamano questo fenomeno Softmax Collapse. La loro attivazione StableMax e l'ottimizzatore ⊥Grad producono grokking senza decadimento dei pesi.

Grokking a confronto con concetti correlati#

• Overfitting: Una differenza tra le prestazioni di addestramento e di validazione che persiste o aumenta. Nel grokking, la differenza si riduce da sola continuando l'addestramento.
• Doppia discesa: L'errore di test diminuisce, aumenta e poi diminuisce di nuovo mentre crescono le dimensioni del modello o la durata dell'addestramento. La doppia discesa per epoca è un'altra forma di miglioramento ritardato: durante l'addestramento l'errore di test aumenta per un po', prima di diminuire di nuovo.
• Interpretabilità meccanicistica: Lo studio degli algoritmi interni alle reti addestrate. Il grokking nell'aritmetica modulare è diventato uno dei principali casi di studio perché è stato possibile ricostruire per intero l'algoritmo appreso.

Il grokking nella pratica#

Il grokking è più evidente nei piccoli dataset algoritmici governati da un'unica regola precisa. È comunque importante per la progettazione degli esperimenti, perché l'arresto anticipato interrompe un'esecuzione quando la metrica di validazione smette di migliorare, bloccando un'esecuzione soggetta a grokking durante il plateau. L'addestramento di Ultralytics YOLO si arresta dopo patience epoche senza miglioramenti nella fitness di validazione, 100 per impostazione predefinita, ed esegue ogni epoca quando patience=0.