Structural Inference: Interpreting Small Language Models with Susceptibilities

Deze paper introduceert een lineair responsframework dat een klein taalmodel als een Bayesiaans statistisch mechanisch systeem behandelt, waarbij susceptibiliteiten worden gebruikt om functionele modules zoals multigram- en inductiehoofden te identificeren en te attribueren.

De kern

Stel je voor dat een kunstmatige intelligentie (zoals een chatbot) een enorme, donkere kamer is vol met duizenden kleine, knipperende lampjes. We weten dat deze lampjes samenwerken om zinnen te maken, maar we weten niet precies wat elk lampje doet. Soms denken we dat ze "leren", maar hoe weten we wat ze echt begrijpen en wat ze alleen maar nabootsen?

Dit paper introduceert een nieuwe manier om in die donkere kamer te kijken, genaamd "Structural Inference" (Structurele Afleiding). Het gebruikt een slimme truc uit de natuurkunde om te zien hoe de lampjes reageren op veranderingen.

Hier is de uitleg in simpele taal:

1. De Analogie: De Magnetische Kamer

Stel je voor dat de AI een heel groot blok ijzer is. Normaal gesproken is het ijzer niet magnetisch; de atomen (de lampjes) staan willekeurig.

• De Proef: De onderzoekers sturen een heel klein, zacht magnetisch veld naar het ijzer. Dit is in hun geval geen echt magnetisme, maar een verandering in de data die de AI leest.
  • Voorbeeld: Ze laten de AI eerst gewone teksten lezen (zoals nieuwsberichten) en daarna plotseling veel code (zoals GitHub) of juridische teksten.
• De Reactie (De "Susceptibiliteit"): Als je een klein magnetisch veld op ijzer richt, reageren sommige atomen direct en draaien ze mee. Andere atomen blijven stil of draaien juist de andere kant op.
  • In de AI noemen ze deze reactie "susceptibiliteit" (gevoeligheid).
  • Negatieve reactie: Het lampje "duwt" de nieuwe informatie weg. Het zegt: "Nee, dit past niet bij wat ik nu doe."
  • Positieve reactie: Het lampje "trekt" de nieuwe informatie aan. Het zegt: "Ja! Dit is precies wat ik moet doen!"

2. Wat ontdekten ze? (De "Inductiekringen")

Door te kijken naar welke lampjes (in de AI heet dit attention heads) hoe reageren op welke teksten, konden de onderzoekers de interne structuur van de AI in kaart brengen. Ze vonden twee belangrijke groepen:

• De "Woord-Indelingers": Sommige lampjes reageerden heel sterk op het begin en einde van woorden. Ze helpen de AI om te begrijpen waar een woord begint en eindigt, net als een kind dat leert om woorden uit een stroom van letters te halen.
• De "Inductiekringen" (De Pattern-herkenners): Dit is het coolste deel. Ze vonden een speciaal team van lampjes dat werkt als een detective.
  • Als de AI een zin leest als: "De kat zat op de mat. De kat...", dan weten deze lampjes: "Ah, dit patroon komt vaker voor! De volgende woord is waarschijnlijk 'zat'."
  • Ze herkennen patronen in de tekst en voorspellen wat er als volgende komt.

3. Het Gevecht: "Uitdrukken" vs. "Onderdrukken"

De onderzoekers ontdekten dat de AI niet alleen dingen "voorspelt", maar ook dingen "onderdrukt".

• Uitdrukken: Een lampje zegt: "Ik denk dat het woord 'appel' hier past!"
• Onderdrukken: Een ander lampje zegt: "Wacht, 'appel' past hier niet, dat is gek. Ik ga die kans verkleinen."

Dit is als een vergadering. Sommige mensen schreeuwen een idee naar voren (positieve reactie), terwijl anderen zeggen: "Nee, dat is een slecht idee" (negatieve reactie). De AI gebruikt dit gevecht om de beste antwoord te kiezen.

4. Waarom is dit belangrijk?

Vroeger moesten onderzoekers de AI "openbreken" (delen uitschakelen) om te zien wat er gebeurde. Dat was als een auto demonteren om te zien hoe de motor werkt.

Met deze nieuwe methode hoeven ze de AI niet te beschadigen. Ze sturen gewoon een klein "stootje" (een verandering in de tekst) en kijken hoe de AI reageert.

• Het is alsof je een piano bespeelt: je hoort welke toetsen (lampjes) klinken als je op een bepaalde toets (data) drukt.
• Hiermee kunnen ze zien dat de AI niet zomaar een "zwarte doos" is, maar een complex machine met specifieke onderdelen die specifieke taken hebben, zoals het herkennen van patronen of het structureren van zinnen.

Samenvatting

Dit paper zegt: "Laten we niet raden wat de AI doet, maar laten we haar een klein beetje veranderen en kijken hoe ze reageert."

Door te kijken naar deze reacties, hebben ze bewezen dat kleine AI-modellen al heel slimme, gespecialiseerde onderdelen hebben die werken als een goed georganiseerd team: sommigen zoeken patronen, sommigen houden de structuur in de gaten, en sommigen houden de verkeerde ideeën tegen. Het is een nieuwe manier om de "ziel" van de machine te begrijpen zonder haar te beschadigen.

Technische samenvatting

Probleemstelling

De interne microscopische organisatie van neurale netwerken, die hun complexe gedrag mogelijk maakt, blijft slecht begrepen. Bestaande methoden voor mechanische interpretatie, zoals ablaties (het uitschakelen van onderdelen) of invloedfuncties, hebben beperkingen. Ablaties kunnen bijvoorbeeld worden verstoord door "zelfherstel" (self-repair) in het netwerk, waarbij andere lagen de effecten van een ablatie compenseren. Er is behoefte aan een theoretisch onderbouwde, schaalbare methode om de interne structuur van modellen te ontrafelen en te begrijpen hoe deze reageert op veranderingen in de data-distributie.

Methodologie: Structurele Inferentie en Susceptibiliteit

Het paper introduceert een nieuw interpretatiekader genaamd Structural Inference, gebaseerd op statistische mechanica en Bayesiaanse leertheorie. De kern van de methode is het concept van susceptibiliteit (gevoeligheid).

1. Theoretisch Kader:

   • Het neurale netwerk wordt behandeld als een Bayesiaans statistisch mechanisch systeem.
   • Een kleine, gecontroleerde verstoring (perturbatie) in de data-distributie (bijvoorbeeld het verschuiven van de Pile-dataset naar meer GitHub-code of juridische tekst) induceert een lineaire respons in de verwachte gedragingen van specifieke netwerkcomponenten (zoals attention heads).
   • Deze respons wordt de susceptibiliteit ($\chi$) genoemd. Wiskundig wordt deze gedefinieerd als de afgeleide van de posterior-verwachting van een observable (bijv. het verlies van een component) ten opzichte van de verstoring in de data, geschaald met de inverse temperatuur ($\beta$) en steengrootte ($n$).
   • Formeel: $\chi = -\text{Cov}_\beta[\phi, \Delta L]$, waarbij $\phi$ de observable is en $\Delta L$ de verandering in het verlies door de data-verschuiving.

2. Praktische Schatting (Local Susceptibility):

   • Omdat het aftrekken van de volledige Bayesiaanse posterior computatietechnisch onhaalbaar is, gebruiken de auteurs lokale susceptibiliteiten.
   • Ze localiseren de posterior rond een lokaal minimum $w^*$ (de getrainde modelgewichten) door een Gaussische prior te gebruiken.
   • Steekproeven worden getrokken met Stochastic Gradient Langevin Dynamics (SGLD) om de verwachtingen te schatten.
   • De methode berekent per-token susceptibiliteiten, wat betekent dat voor elk token $(x, y)$ wordt bepaald hoe gevoelig een specifieke component is voor het voorspellen van $y$ in context $x$ onder een verschuiving in de data.

3. Interpretatie van Tekens:

   • Negatieve susceptibiliteit: Wordt geïnterpreteerd als expressie. De component helpt het voorspellen van het token; veranderingen die het verlies verlagen, verlagen ook de kans op het token (of versterken de voorspelling).
   • Positieve susceptibiliteit: Wordt geïnterpreteerd als suppressie. De component werkt tegen het voorspellen van het token; veranderingen die het verlies verlagen, verhogen juist de kans op het token (de component "onderdrukt" de voorspelling).

4. Structurale Inferentie via PCA:

   • De auteurs bouwen een responsmatrix op waarin de rijen verschillende data-distributies (probes) zijn en de kolommen verschillende modelcomponenten (attention heads).
   • Door Principal Component Analysis (PCA) toe te passen op deze matrix, kunnen ze patronen in de data identificeren die corresponderen met functionele modules in het model. De hoofdcomponenten (PCs) vertegenwoordigen de "modes" van de data-distributie, en de ladingen (loadings) tonen welke heads bijdragen aan deze modes.

Belangrijkste Bijdragen

• Nieuw Paradigma: De ontwikkeling van een interpretatiekader dat voortkomt uit Bayesiaanse leertheorie en statistische fysica, wat een principieel verband legt tussen data-structuur en interne modelstructuur.
• Methodologie van Structurele Inferentie: Een methode om interne netwerkmogelijkheden te ontdekken en te attribueren aan patronen in de data, zonder afhankelijk te zijn van ablaties die lastig te interpreteren zijn door self-repair.
• Empirische Validatie: Het succesvol toepassen van deze methode op een klein taalmodel (3M parameters) om bekende circuits te isoleren en te karakteriseren.

Resultaten

De methode werd getest op een 2-laags Transformer (alleen attention, geen MLP) getraind op een subset van The Pile. De resultaten tonen aan dat attention heads betekenisvol verschillende susceptibiliteiten vertonen voor verschillende data-verschuivingen:

1. PC1: Woordsegmentatie: De eerste hoofdcomponent is uniform over alle heads en correleert sterk met de taak van het segmenteren van tokenstromen in woorden (onderscheid tussen woordstart, woorddeel en woordend).
2. PC2: Het Inductie Circuit: De tweede component scheidt duidelijk het bekende inductie circuit (heads die patronen zoals $A \dots B \dots A \dots B$ herkennen en voorspellen) van andere heads.
   • De inductie-heads (bijv. 1:6, 1:7) tonen een sterke respons op inductiepatronen (negatieve susceptibiliteit = expressie).
   • De resterende heads (waaronder de "multigram heads" in laag 1) tonen een positieve respons (suppressie) op deze patronen. Dit bevestigt eerdere bevindingen dat sommige heads inductiepatronen actief onderdrukken.
3. PC3: Haakjes Matching: De derde component correleert met het voorspellen van sluitende haakjes en delimiter-patronen, wat overeenkomt met eerdere bevindingen over "Dyck heads".
4. Robuustheid: De analyse werd herhaald over meerdere training seeds en toonde aan dat de geïdentificeerde structuren (zoals het inductie circuit) consistent terugkomen, wat de betrouwbaarheid van de methode onderstreept.

Betekenis en Impact

Dit werk biedt een fundamenteel nieuwe manier om neurale netwerken te analyseren:

• Theoretische Diepgang: Het koppelt interpretatie direct aan de wiskundige theorie van generalisatie en singular learning theory (Watanabe), in plaats van puur empirische observaties.
• Overcoming Self-Repair: In tegenstelling tot ablaties, die kunnen worden geneutraliseerd door het netwerk, meet susceptibiliteit de inherente gevoeligheid van de componenten, waardoor het effecten van suppressie en expressie direct zichtbaar maakt.
• Schaalbaarheid: Hoewel de huidige studie op een klein model is uitgevoerd, is de methode (gebaseerd op SGLD) theoretisch schaalbaar naar grotere modellen. Het biedt een blauwdruk voor het ontrafelen van complexe circuits in grote taalmodellen (LLMs) door te kijken naar hoe ze reageren op veranderingen in hun trainingsdata.

Kortom, het paper bewijst dat het analyseren van de "lineaire respons" van een model op data-verschuivingen een krachtig hulpmiddel is om de functionele organisatie van neurale netwerken te ontrafelen.