Induction Heads Interpolate N-Grams

Technische Samenvatting: Inductiekoppen Interpoleren N-grammen

Probleemstelling
Inductiekoppen (induction heads), geïdentificeerd als het primaire circuit voor in-context leren (ICL) in transformers, worden traditioneel gekenmerkt in de "hard-attention" limiet. In dit regime voert het circuit exacte patroonherkenning en maximum-likelihood (ML) tellingen uit. Hoewel dit mechanistisch aantrekkelijk is, is deze visie statistisch incompleet voor hogere-orde afhankelijkheden: exacte $k$-gram matches zijn vaak te zeldzaam in eindige steekproeven om betrouwbare voorspellingen te ondersteunen, en de ML-schatter wijst een nulwaarschijnlijkheid toe aan ongeziene voortzettingen. Klassieke taalmodellen pakken dit aan via smoothing- en back-off mechanismen (bijv. Jelinek–Mercer, Katz back-off), die waarschijnlijkheidsmassa herverdelen naar ongeziene gebeurtenissen en interpoleren over contextlengtes. De centrale vraag die dit werk adresseert, is of inductiekoppen een rijkere klasse van schatters implementeren dan strikte tellingen, specifiek in staat tot het uitvoeren van dergelijke smoothing.

Methodologie
De auteurs analyseren transformers die getraind zijn op orde-$k$ Markov-ketens, een setting die een balans biedt tussen analytische traceerbaarheid en de trade-offs van eindige steekproeven die inherent zijn aan context-gebaseerde voorspelling.

1. Theoretisch Kader: De studie maakt gebruik van een "ontkoppelde" (disentangled) transformer-architectuur (waarbij MLP's worden verwijderd en concatenatie wordt gebruikt voor residual streams) om een exacte analyse te faciliteren. De auteurs construeren een specifieke twee-laagse transformer met Relative Positional Encoding (RPE) en bewijzen dat deze een soft context-matching estimator implementeert.
2. Mechanisme Constructie:
   • Laag 1 (Copying Heads): Gebruikt $k$ attention heads om tokens te kopiëren van specifieke relatieve lags ($t-1, \dots, t-k$), waardoor een representatie van de huidige context wordt geconstrueerd.
   • Laag 2 (Context Matching): Gebruikt een enkele attention head om de query-context te vergelijken met alle kandidaat-contexten in de historie. De attention scores zijn ontworpen om een verschuivings-gealigneerd inproduct te berekenen, waarbij matches worden gewogen op basis van de grootte van de overlap (Hamming overlap) in plaats van een exacte match te vereisen.
   • BOS Token: De auteurs voegen een Beginning-of-Sequence (BOS) token toe als een sequentie-onafhankelijk attention-doel om additieve pseudo-counts te introduceren.
3. Statistische Analyse: De auteurs leiden de closed-form estimator af die door dit circuit wordt geïmplementeerd en vergelijken deze met Bayes-optimale voorspellers onder twee taak-priors:
   • Independent Dirichlet Prior: Waarbij contexten ongecorreleerd zijn.
   • Hierarchical Dirichlet Prior: Waarbij overgangen van hogere orde gecentreerd zijn op hun lagere-orde suffixes, wat correlaties creëert tussen contexten die een gedeelde suffix hebben.

Belangrijkste Bijdragen
Het artikel stelt vast dat inductiekoppen twee complementaire smoothing-mechanismen implementeren wanneer ze opereren bij eindige attention-gewichtsschalen:

1. Soft Context Matching (Interpolatie):

   • Het circuit aggregeert bijdragen van zowel exacte als partiële context-matches.
   • Matches worden exponentieel gewogen op basis van hun overlapgrootte ($|M|$).
   • Dit induceert een data-afhankelijke interpolatie over context-ordes, analoog aan Jelinek–Mercer smoothing. In tegen tegenstelling tot klassieke JM smoothing, die vaste globale gewichten gebruikt, zijn de gewichten van de transformer afhankelijk van de gerealiseerde tellingen in de specifieke sequentie.
   • De estimator verenigt mixture modellen (zoals JM) en back-off strategieën (zoals Katz): wanneer exacte matches zeldzaam zijn, "back-off" het model automatisch naar lagere-orde partiële matches zonder expliciete regels.

2. BOS als Additieve Pseudo-Counts:

   • De BOS token fungeert als een vaste, sequentie-onafhankelijke sink voor attention.
   • Dit mechanisme stelt het model in staat om constante pseudo-counts toe te voegen aan context-specifieke transitie-tellingen.
   • Theoretisch herstelt dit add-$\alpha$ smoothing (bijv. Laplace smoothing) in de limiet van grote attention-gewichten, wat een verklaring op circuit-niveau biedt voor hoe architecturale conventies prior-achtige regularisatie implementeren.

3. Empirische Validatie:

   • De auteurs trainen zowel ontkoppelde als standaard (volledig geparameteriseerde) transformers op Markov-keten-taken.
   • Resultaten tonen aan dat getrainde modellen de voorspelde attention-patronen herstellen:
     • Zonder BOS vertrouwen ze op soft context matching om te interpoleren over orden.
     • Met BOS gebruiken ze pseudo-counts om bijna Bayes-optimale prestaties te behalen onder onafhankelijke priors.
     • Onder hiërarchische priors, waar lagere-orde contexten gestructureerd bewijs bieden, presteren de modellen beter dan vaste add-$\alpha$ baselines door het interpolatiemechanisme te benutten.
   • Mechanistische analyse bevestigt dat getrainde standaard transformers de specifieke gewichtsstructuren (bijv. block-shifted matrices) leren die de theorie voorspelt.

Resultaten

• KL Divergentie: In settings waar add-$\alpha$ smoothing Bayes-optimaal is (Independent Dirichlet), matchen getrainde transformers met BOS de Bayes-optimale voorspeller. In hiërarchische settings, waar gestructureerd bewijs bestaat, stelt het interpolatiemechanisme transformers in staat om vaste smoothing-baselines te overtreffen.
• Attention Patronen: Visualisaties van getrainde modellen bevestigen het twee-fasen circuit: Laag 1 voert lag-specifieke kopieeracties uit, en Laag 2 voert soft matching uit, waarbij de attention-massa over volledige en partiële matches wordt verdeeld op basis van overlap.
• Asymmetrische Lag Gewichten: In hiërarchische taken leren getrainde modellen hogere gewichten toe te kennen aan matches op de meest recente token (lag 1) vergeleken met oudere tokens (lag 2), wat consistent is met de structuur van de prior waarbij recente tokens betere voorspellers zijn van de volgende staat.

Significantie
Dit artikel overbrugt de kloof tussen de mechanistische interpreteerbaarheid van inductiekoppen en klassieke statistische smoothing. Het demonstreert dat transformers niet louter ongeregulariseerde tellingen uitvoeren; in plaats daarvan leren ze te regulariseren in in-context estimatie. Door dynamisch te interpoleren over context-ordes en architecturale tokens (BOS) te gebruiken voor additieve smoothing, implementeren inductiekoppen geavanceerde statistische procedures die zich aanpassen aan de datadistributie, wat hun robuustheid in in-context leer-taken verklaart.