Technische Samenvatting: Transformers Representeren Hoog Gestructureerde Wereldmodellen Lineair

Probleemstelling

Het artikel onderzoekt of transformers die zijn getraind op sequentiële redeneersporen, interne "wereldmodellen" ontwikkelen die de onderliggende algebraïsche structuur van het taakdomein weerspiegelen, in plaats van slechts de oppervlakkige presentatie. Hoewel eerdere werken (bijv. Li et al. [2023], Nanda et al. [2023]) hebben aangetoond dat modellen die zijn getraind op Othello spontaan lineaire representaties van bezetting van het bord leren (één feature per cel), is het onduidelijk of dit ook geldt voor taken met verschillende structurele beperkingen. Specifiek vragen de auteurs zich af: in een combinatorische constraint-satisfaction-taak zoals Sudoku, waar geldigheid wordt bepaald door overlappende substructuren (rijen, kolommen en 3×3 vakken) in plaats van geïsoleerde cellen, representeert het model de toestand cel-voor-cel of substructuur-voor-substructuur?

Methodologie

De auteurs trainden een causal decoder-only transformer van 8 lagen (8 attention-heads, $d_{model}=576$) op 2,7 miljoen Sudoku-oplossingsporen gegenereerd door een Norvig-stijl solver met backtracking. Het model behaalde 98,4% nauwkeurigheid per cel en 97,5% per raster op een vastgehouden testset.

Om de interne berekening te analyseren, hanteerden de auteurs een "probe-and-patch"-methodologie:

1. Lineair Probing: Logistische regressie-classifiers werden getraind op residual stream-activaties bij het [clues_end]-token (dat het einde van de initiële aanwijzingen markeert) om specifieke concepten te decoderen. Ze testten drie families van probes:
   • Cel-toestand: Het voorspellen van het cijfer in een specifieke cel (81 probes).
   • Cel-kandidaten: Het voorspellen van geldige kandidaat-cijfers voor een specifieke cel (729 probes).
   • Substructuur-toestand: Het voorspellen of een cijfer aanwezig is in een specifieke substructuur (rij, kolom of vak) (243 probes).
2. Causale Interventie (Activation Patching): Om te verifiëren of gedecodeerde representaties causaal worden gebruikt, trokken de auteurs specifieke probe-richtingen af van de residual stream en maten ze de resulterende daling in logits van het doel-token.
3. Attention-analyse: De auteurs decomposeerden attention-heads in Query-Key (QK) en Output-Value (OV) circuits om te bepalen welke informatie heads aandacht geven en hoe ze de residual stream modificeren.
4. Logit Lens & Direct Logit Attribution (DLA): Deze hulpmiddelen werden gebruikt om te volgen hoe voorspellingen evolueren over de lagen heen en om veranderingen in logits toe te schrijven aan specifieke componenten (heads of neuronen).
5. Neuron Ablatie: Specifieke neuronen in de laatste MLP-laag werden geïdentificeerd en geablateerd om hun causale rol in besluitvorming te testen.

Belangrijkste Bijdragen en Resultaten

1. Emergentie van een Substructuur-Wereldmodel

In tegenstelling tot de "cel-voor-cel"-representatie die in Othello-modellen werd gevonden, bouwt de Sudoku-transformer een wereldmodel op dat is georganiseerd rond de constraint-algebra van de puzzel.

• Lineaire Scheidbaarheid: Lineaire probes voor "is cijfer $d$ aanwezig in substructuur $S$?" bereiken perfecte exact-match-nauwkeurigheid (1,0) in middenlagen (Lagen 4–7) voor alle 243 (substructuur, cijfer)-paren.
• Beperkingen op Celniveau: Daarentegen bereiken probes die de toestand van individuele cellen voorspellen een plateau bij ongeveer 80% nauwkeurigheid. Probes voor celkandidaten bereiken 100% nauwkeurigheid, maar vertonen structurele uitlijning: kandidaat-vector voor cellen die een substructuur delen, zijn sterk gecorreleerd, wat suggereert dat ze lineaire combinaties zijn van substructuurniveau-features.
• Stabiliteit: Deze substructuurrepresentatie ontstaat in Laag 4 en blijft geometrisch stabiel gedurende het hele oplossingsproces, zoals bevestigd door cross-layer en cross-position transfer-experimenten.
• Causaal Gebruik: Activation patching bevestigt dat deze richtingen causaal noodzakelijk zijn; het overschrijven ervan om aan te geven dat een cijfer aanwezig is in een substructuur, zorgt ervoor dat de voorspelling van het model voor die plaatsing instort.

2. Substructuur-uitgelijnde Attention-mechanismen

Attention-heads in de middenlagen (Lagen 4–6) zijn gespecialiseerd om te opereren op specifieke substructuren.

• Routing: Attention-heads concentreren hun focus op specifieke rijen, kolommen, vakken, bands of stacks.
• Constraint-handhaving: De OV-circuits van deze heads implementeren een symmetrisch constraint-checking-mechanisme. Als een cijfer aanwezig is in de geattendeerde substructuur, onderdrukt de head de logits voor dat cijfer in de relevante cellen; als het afwezig is, promoot het het cijfer.
• Lokalisatie: Causale ablatie van specifieke heads (bijv. een head die aandacht geeft aan kolommen 4–6) verhoogt de logits van illegale plaatsingen alleen binnen die specifieke substructuur, terwijl andere substructuren onaangetast blijven.

3. De Naked-Single Circuit

De auteurs identificeerden een spaarzaam, monosemantisch circuit in de laatste MLP-laag dat verantwoordelijk is voor het vastleggen van "naked single"-plaatsingen (cellen met precies één geldige kandidaat).

• Pre-MLP Toestand: Voor de laatste MLP rangschikt de residual stream al het juiste cijfer als eerste in 98,77% van de unieke naked-single-cases.
• Gedwongen Neuronen: De laatste MLP bevat 91 toegewijde neuronen die uitsluitend afvuren wanneer hun geassocieerde cel precies één kandidaat heeft. Deze neuronen zijn monosemantisch en activeren niet voor andere aantallen kandidaten.
• Versterkingsmechanisme: Wanneer een naked-single-neuron afvuurt, past het een uniforme, massieve boost toe (+3,00 logits) op alle plaatsing-tokens voor die specifieke cel. Aangezien het juiste cijfer al door de middenlagen als hoogst gerangschikt was, behoudt deze uniforme versterking de voorsprong en concentreert de softmax-kans op het juiste cijfer.
• Causale Bevestiging: Het ablaten van deze neuronen verlaagt de kans op de doel-naked single met ongeveer 0,59 zonder andere gelijktijdige naked singles op hetzelfde bord te beïnvloeden, wat hun rol als de uiteindelijke toewijzingspoort bevestigt.

Betekenis en Claims

Het artikel claimt aan te tonen dat de geometrie van emergente wereldmodellen wordt gevormd door de constraint-algebra van het domein in plaats van de oppervlakkige decompositie. Bij Sudoku, waar geldigheid wordt gedefinieerd door de doorsnede van rijen, kolommen en vakken, leert het model de toestand te representeren in termen van deze substructuren.

Bovendien biedt de studie een volledig mechanistisch, end-to-end algoritmisch verslag van hoe een transformer een combinatorische redeneertaak oplost:

1. Representatie: Middenlagen construeren een lineair, substructuur-uitgelijnd wereldmodel.
2. Constraint-propagatie: Gespecialiseerde attention-heads handhaven constraints door ongeldige kandidaten binnen die substructuren te onderdrukken.
3. Besluitvorming: Een spaarzaam aantal monosemantische neuronen in de laatste laag detecteert specifieke logische toestanden (naked singles) en versterkt het vooraf berekende juiste antwoord.

De auteurs concluderen dat mechanische interpretabiliteitstools deze complexe redeneerprocessen succesvol kunnen reverse-engineeren, waarbij blijkt dat de resulterende besluitvormingscircuits spaarzaam, monosemantisch en volledig interpreteerbaar zijn.

Beperkingen

De auteurs merken op dat deze bevindingen specifiek zijn voor een enkele modelarchitectuur (8 lagen, 8 heads) die is getraind op een enkele puzzeltype (Sudoku). Ze claimen niet dat substructuur-wereldmodellen of naked-single-circuits noodzakelijkerwijs zullen verschijnen in modellen van verschillende maten of die zijn getraind op niet-backtracking-paden. Bovendien is de analyse van hoe attentionpatronen worden geselecteerd (QK-circuits) niet volledig gedetailleerd, en rust het causale bewijs op een specifiek patching-protocol.