Induction Heads Interpolate N-Grams

Sintesi Tecnica: Le teste di induzione interpolano gli n-grammi

Definizione del Problema
Le teste di induzione, identificate come il circuito primario per l'apprendimento in-context (ICL) nei transformer, sono tradizionalmente caratterizzate nel limite dell' "attenzione hard". In questo regime, il circuito esegue un esatto pattern matching e un conteggio di massima verosimiglianza (ML). Sebbene questa visione sia meccanicisticamente accattivante, è statisticamente incompleta per le dipendenze di ordine superiore: i match esatti di $k$-grammi sono spesso troppo rari in campioni finiti per supportare una predizione affidabile, e lo stimatore ML assegna probabilità zero alle continuazioni non viste. I modelli linguistici classici affrontano questo problema tramite meccanismi di smoothing e back-off (ad esempio, Jelinek–Mercer, Katz back-off), che ridistribuiscono la massa di probabilità verso eventi non visti e interpolano attraverso diverse lunghezze di contesto. La domanda centrale affrontata da questo lavoro è se le teste di induzione implementino una classe più ricca di stimatori oltre il semplice conteggio, essendo capaci di eseguire tale smoothing.

Metodologia
Gli autori analizzano transformer addestrati su catene di Markov di ordine-$k$, un setting che bilancia la trattabilità analitica con i compromessi dei campioni finiti inerenti alla predizione basata sul contesto.

1. Framework Teorico: Lo studio utilizza un'architettura transformer "disentangled" (rimuovendo i MLP e utilizzando la concatenazione per gli stream residui) per facilitare l'analisi esatta. Gli autori costruiscono un particolare transformer a due strati con Codifica Posizionale Relativa (RPE) e dimostrano che implementa uno stimatore di soft context-matching.
2. Costruzione del Meccanismo:
   • Strato 1 (Teste di Copia): Utilizza $k$ teste di attenzione per copiare i token da specifici lag relativi ($t-1, \dots, t-k$), costruendo una rappresentazione dell'attuale contesto.
   • Strato 2 (Context Matching): Utilizza una singola testa di attenzione per confrontare il contesto di query con tutti i contesti candidati nella storia. I punteggi di attenzione sono progettati per calcolare un prodotto interno con shift allineato, pesando i match in base al numero di token sovrapposti (overlap di Hamming) piuttosto che richiedere match esatti.
   • Token BOS: Gli autori incorporano un token di Inizio Sequenza (BOS) come target di attenzione indipendente dalla sequenza per introdurre pseudo-conteggi additivi.
3. Analisi Statistica: Gli autori derivano lo stimatore a forma chiusa implementato da questo circuito e lo confrontano con predittori Bayes-ottimali sotto due prior di task:
   • Prior di Dirichlet Indipendente: Dove i contesti sono non correlati.
   • Prior di Dirichlet Gerarchico: Dove le transizioni di ordine superiore sono centrate sui loro suffissi di ordine inferiore, creando correlazioni tra contesti che condividono un suffisso.

Contributi Chiave
Il paper stabilisce che le teste di induzione implementano due meccanismi di smoothing complementari quando operano a scale di peso di attenzione finite:

1. Soft Context Matching (Interpolazione):

   • Il circuito aggrega i contributi sia dai match di contesto esatti che da quelli parziali.
   • I match sono pesati esponenzialmente in base alla dimensione del loro overlap ($|M|$).
   • Ciò induce un'interpolazione dipendente dai dati attraverso gli ordini di contesto, analogamente allo smoothing di Jelinek–Mercer. A differenza dello smoothing JM classico, che utilizza pesi globali fissi, i pesi del transformer dipendono dai conteggi realizzati nella specifica sequenza.
   • Lo stimatore unifica i modelli di miscela (come JM) e le strategie di back-off (come Katz): quando i match esatti sono rari, il modello effettua automaticamente il "back-off" verso match parziali di ordine inferiore senza regole esplicite.

2. BOS come Pseudo-conteggi Additivi:

   • Il token BOS agisce come un "sink" fisso e indipendente dalla sequenza per l'attenzione.
   • Questo meccanismo permette al modello di aggiungere conteggi pseudo costanti ai conteggi di transizione specifici del contesto.
   • Teoricamente, questo recupera lo smoothing add-$\alpha$ (ad esempio, lo smoothing di Laplace) nel limite di grandi pesi di attenzione, fornendo una spiegazione a livello di circuito di come le convenzioni architettoniche implementino la regolarizzazione di tipo prior.

3. Validazione Empirica:

   • Gli autori addestrano sia transformer disentangled che standard (completamente parametrizzati) su task di catene di Markov.
   • I risultati mostrano che i modelli addestrati recuperano i pattern di attenzione previsti:
     • Senza BOS, si affidano al soft context matching per interpolare tra gli ordini.
     • Con BOS, utilizzano i pseudo-conteggi per raggiungere prestazioni vicine all'ottimo Bayes sotto prior indipendenti.
     • Sotto prior gerarchici, dove i contesti di ordine inferiore forniscono evidenza strutturata, i modelli superano le baseline a smoothing add-$\alpha$ fisso sfruttando il meccanismo di interpolazione.

Risultati

• Divergenza KL: In contesti dove lo smoothing add-$\alpha$ è Bayes-ottimale (Dirichlet Indipendente), i transformer addestrati con BOS corrispondono al predittore Bayes-ottimale. Nei contesti gerarchici, dove esiste evidenza strutturata, il meccanismo di interpolazione permette ai transformer di superare le baseline di smoothing fisso.
• Pattern di Attenzione: Le visualizzazioni dei modelli addestrati confermano il circuito a due stadi: lo Strato 1 esegue la copia specifica per lag, e lo Strato 2 esegue il soft matching, distribuendo la massa di attenzione tra match completi e parziali in base all'overlap.
• Pesi di Lag Asimmetrici: Nei task gerarchici, i modelli addestrati imparano ad assegnare pesi maggiori ai match sul token più recente (lag 1) rispetto ai token precedenti (lag 2), in modo coerente con la struttura del prior dove i token recenti sono predittori migliori dello stato successivo.

Significatività
Il paper colma il divario tra l'interpretabilità meccanicistica delle teste di induzione e lo smoothing statistico classico. Dimostra che i transformer non eseguono semplicemente un conteggio non regolarizzato; piuttosto, imparano a regolarizzare la stima in-context. Interpolando dinamicamente attraverso gli ordini di contesto e utilizzando token architettonici (BOS) per lo smoothing additivo, le teste di induzione implementano procedure statistiche sofisticate che si adattano alla distribuzione dei dati, spiegando la loro robustezza nei task di apprendimento in-context.