Technische Samenvatting: Hoe leren Transformers tokens te associëren: leidende termen van gradiënten brengen mechanistische interpreteerbaarheid

Probleemstelling

Grote Taalmodellen (LLMs) gebaseerd op zelf-attention hebben opmerkelijke capaciteiten getoond in het vastleggen van feitelijke kennis en het genereren van samenhangende tekst, grotendeels gedreven door hun vermogen om semantische associaties te leren (bijvoorbeeld het verband tussen "vogel" en "vlog" of de uitwisselbaarheid van "auto" en "vrachtwagen"). Hoewel recente studies specifieke mechanismen hebben geïdentificeerd, zoals inductiehoofden of lineaire semantische relaties, ontbreekt er een principiële, theorie-gebaseerde uiteenzetting van hoe deze associaties ontstaan tijdens de training van op attention gebaseerde transformers op realistische natuurlijke taaldata.

Vorige theoretische werken maken vaak gebruik van onrealistische aannames, zoals synthetische gestructureerde taal, vereenvoudigde architecturen (zonder positie-encoding of residuale verbindingen), of niet-standaard trainingsregimes (bijvoorbeeld sequentiële component-voor-component training). Deze afwijkingen van de praktijk beperken de generaliseerbaarheid van hun inzichten tot moderne LLM's. Dit artikel vult de lacune door de volgende vraag te stellen: Hoe ontstaan semantische associaties tijdens de training van op attention gebaseerde taalmodellen op natuurlijke taaldata?

Methodologie

De auteurs stellen een theoretisch kader voor dat is gebaseerd op trainingsdynamiek, met name door analyse van de leidende-termenbenadering van gradiënten tijdens de vroege fasen van training.

1. Modelopzet: De studie richt zich op een standaard $L$-laagse attention-gebaseerde transformer met positie-encoding (relatief), causale masking en residuale streams, getraind via standaard full-batch gradiëntafstijging op natuurlijke taalcorpora.
2. Leidende-termenbenadering: Met de erkenning dat transformers kerngedrag (inclusief semantische relaties) vroeg in de training verwerven, leiden de auteurs gesloten-formule expressies af voor de gewichtsmatrices (Output $W_O$, Value $V^{(l)}$, Query-Key $W^{(l)}$, en Positie-encoding $P^{(l)}$) door de gradiëntupdates te benaderen. Zij tonen aan dat voor een voldoende aantal stappen $s$ (waarbij $s \leq \eta^{-1} \min(\dots)$), de gewichten dicht bij hun leidende gradiënttermen blijven.
3. Ontbinding in Basisfuncties: De theoretische analyse toont aan dat de geleerde gewichten kunnen worden uitgedrukt als eenvoudige composities van drie fundamentele basisfuncties, afgeleid van corpusstatistieken:
   • Bigrammapping ($\bar{B}$): Vangt afhankelijke volgorde van volgende tokens op ($P(e_j|e_i)$).
   • Uitwisselbaarheidsmapping ($\Sigma \bar{B} = \bar{B}^\top \bar{B}$): Vangt functionele gelijkenis tussen tokens op, gebaseerd op de gelijkenis van hun voorafgaande tokenverdelingen (bijvoorbeeld synoniemen of gedeelde grammaticale rollen).
   • Contextmapping ($\bar{\Phi}$): Encodeert langere prefix-suffix co-occurrence, en vertegenwoordigt de gemiddelde context waarin een token voorkomt.

Belangrijkste Bijdragen

Het artikel levert drie primaire bijdragen:

1. Expliciete Karakterisering van Gewichten: Het presenteert de eerste expliciete, gesloten-formule karakterisering van gewichten in op attention gebaseerde transformers getraind op realistische tekstcorpora onder de loss voor voorspelling van het volgende token. De gewichten blijken composities te zijn van de bovengenoemde drie basisfuncties.
2. Mechanistische Interpretatie van Semantische Associaties: De auteurs interpreteren de in de gewichten geleerde kenmerken als specifieke composities van bigram-, uitwisselbaarheids- en contextmappings. Zij tonen aan hoe deze composities het model in staat stellen semantische associaties te vangen (bijvoorbeeld hoe "vis" gerelateerd is aan "vijver" via contextmapping, of hoe "rood" gerelateerd is aan "vrachtwagen" via bigrammapping).
3. Empirische Validatie: De theoretische voorspellingen worden gevalideerd op zowel een toy 3-laags attention-model getraind op het TinyStories-dataset als een praktisch, grootschalig LLM (Pythia-1.4B) getraind op OpenWebText.

Resultaten

• Theoretische Alignering: Experimenten op een 3-laagse transformer getraind op TinyStories tonen aan dat de cosinus-similariteit tussen de theoretisch voorspelde leidende termen en de daadwerkelijk geleerde gewichten extreem hoog blijft (bijvoorbeeld >0,99 voor Attention- en Value-matrices), zelfs na 100 epochs training.
• Semantische Structuur: Kwalitatieve analyse bevestigt dat de basisfuncties betekenisvolle linguïstische structuren vangen. Zo correleert de bigrammapping $\bar{B}$ bijvoeglijke naamwoorden met de zelfstandige naamwoorden die zij modificeren, terwijl de contextmapping $\bar{\Phi}$ entiteiten correleert met hun typische omgevingen (bijvoorbeeld "vis" met "vijver").
• Generalisatie naar Realistische LLM's: Wanneer toegepast op Pythia-1.4B, tonen de theoretische leidende-termenkenmerken sterke overeenstemming met de geleerde embedding en attention-correlaties, met name in de vroege fasen van training. De analyse toont aan dat, hoewel de gewichten naarmate de training vordert afwijken naar rijkere kennisrepresentaties, zij voor een aanzienlijke duur de fundamentele associatieve kenmerken behouden die door de theorie worden voorspeld.
• Samenwerking van Componenten: De analyse van de volledige modelberekening (Vergelijking 12) illustreert hoe de residuale stream een gemiddelde voorspelling van het volgende token biedt, terwijl het zelf-attention blok (gedreven door $\bar{Q}$ en $\bar{\Phi}^\top \Sigma \bar{B}$) deze voorspelling verfijnt door selectief te focussen op tokens die het meest indicatief zijn voor het volgende token, gebaseerd op corpusstatistieken.

Betekenis

Het artikel claimt een mechanistische basis te bieden voor representatieleren in grote taalmodellen door deep learning trainingsdynamiek te verbinden met linguïstische theorie. Door de analyse te verankeren in natuurlijke tekstverdelingen en standaard trainingsprocedures, verkleint het werk de kloof tussen formele theoretische analyse en praktisch LLM-gebruik aanzienlijk.

De auteurs betogen dat hun bevindingen een rigoureuze manier bieden om de functie van elke transformer-component te contextualiseren. In plaats van het model te zien als een black box, suggereert de theorie dat het gedrag van het model zich ontwikkelt door eerst deze fundamentele associatieve kenmerken (bigrammen, uitwisselbaarheid, context) te vestigen, die dienen als ankers voor het leren van complexere semantische relaties. Dit werk opent paden voor interpreteerbaarheidsonderzoek door te suggereren dat gewichtsmatrices over verschillende componenten kunnen worden ontbonden in eenvoudige functies van gedeelde corpusstatistieken, wat bredere hypothesen mogelijk maakt over hoe concepten ontstaan in modellen, verder dan specifieke mechanismen.