Technische Samenvatting: Inductie in Beide Richtingen

Probleemstelling

Hoewel de interne mechanismen van autoregressieve (AR) transformers uitgebreid zijn bestudeerd vanuit het perspectief van mechanistische interpreteerbaarheid, blijven de operationele mechanismen van Diffusion Language Models (DLMs) grotendeels onverkend. DLMs genereren tekst via iteratieve denoising in plaats van sequentiële next-token voorspelling, waarbij ze bidirectionele aandacht en verschillende trainingsdoelen gebruiken. Er bestaat een kritische kloof in het begrip of DLMs inductie implementeren—een canoniek mechanisme voor in-context leren waarbij een model een herhaalde context identificeert en het daaropvolgende token kopieert—op dezelfde wijze als AR-modellen. Specifiek is het onduidelijk hoe DLMs omgaan met inductie wanneer het doel-token verborgen is achter een mask-token [M] en wordt omringd door bidirectionele context, in tegenstelling tot de strikt causale context die beschikbaar is voor AR-modellen.

Methodologie

De auteurs voerden een gecontroleerde mechanistische studie uit waarbij gematchte attention-only AR-transformers en absorbing-mask DLMs werden vergeleken.

• Modelarchitectuur: De studie maakte gebruik van kleinschalige modellen (dieptes $L \in \{1, 2, 3\}$, verborgen dimensie $d_{model}=512$, 16 attention heads) getraind op FineWeb-Edu sequenties. Beide modelfamilies deelden identieke architecturen en trainingsrecepten, waarbij ze alleen verschilden in hun attention masks (causaal voor AR, bidirectioneel voor DLM) en trainingsdoelen (next-token voorspelling versus absorbing-mask diffusie).
• Evaluatietaak: De inductiecapaciteiten werden gemeten met behulp van synthetische sequenties van lengte 512, waarbij de eerste helft werd herhaald in de tweede helft. De taak hield in dat gemaskeerde tokens in de tweede helft werden voorspeld op basis van de niet-gemaskeerde bron in de eerste helft (forward inductie) en vice versa (reverse inductie).
• Analysetechnieken:
  • Causale Ablaties: Gemiddelde ablaties werden uitgevoerd op individuele attention heads om de causale noodzakelijkheid van deze heads voor inductie te bepalen.
  • Gewichtsdecompositie: De studie maakte gebruik van QK (Query-Key) en OV (Output-Value) decomposities om te analyseren hoe attention logits worden gevormd en hoe informatie van de bron naar de doelpositie stroomt.
  • Lineaire Probing: Een lineaire probe werd getraind op de residual stream om de aanwezigheid van globale mask-fractie informatie te detecteren.
  • Patch Experimenten: "Patching"-technieken werden gebruikt om causaal te interveniëren op de representatie van de globale mask-rate om veranderingen in prediction entropy te observeren.
• Modelvarianten: Voor de mechanistische analyse gebruikten de auteurs "folded no-LayerNorm" DLMs, waarbij getrainde LayerNorm-modellen werden omgezet naar lineaire kaarten om additieve decompositie van de residual stream te faciliteren, terwijl werd geverifieerd dat de prestaties vergelijkbaar bleven met de originele LayerNorm-modellen.

Belangrijkste Bijdragen en Resultaten

1. Bidirectionele Inductie Circuit

De studie identificeert dat DLMs een bidirectioneel inductie circuit leren dat richting-symmetrisch is. In tegen tegenstelling tot AR-modellen, die vertrouwen op een enkel causaal pad, maken DLMs gebruik van twee afzonderlijke paden:

• Previous-Token en Next-Token Heads: In de eerste laag (Laag 0) letten specifieke heads op de directe linker- en rechterburen van het gemaskeerde token, waarbij deze lokale contextkenmerken in de residual stream schrijven.
• Inductie Heads: In daaropvolgende lagen (Laag 1+) gebruiken inductie heads QK-gewichten om de kenmerken die door de Laag 0 heads bij de gemaskeerde positie zijn geschreven, te matchen met de overeenkomstige kenmerken bij de bron-antwoord positie.
• Mechanisme: De QK-decompositie onthult dat inductie heads Q-compositie (query gebouwd uit de lokale context van de mask) en K-compositie (key gebouwd uit de lokale context van de bron) gebruiken om de overeenkomstige bron te lokaliseren. Het OV-pad kopieert vervolgens het bron-token naar de gemaskeerde positie.

2. Context Toegang vs. Mechanisme Sterkte

De auteurs demonstreren dat de superieure inductie-prestaties van DLMs vergeleken met AR-modellen niet komen door een fundamenteel sterker inductiemechanisme, maar door toegang tot bidirectionele context.

• Wanneer beperkt tot alleen linkercontext (om de AR-setting na te bootsen), presteren DLMs op hetzelfde niveau als of zelfs onder het niveau van AR-modellen.
• DLMs presteren alleen beter dan AR-modellen wanneer zowel de linker- als de rechterburen van het gemaskeerde token zichtbaar zijn.
• De inductie-score in de volledige bidirectionele setting benadert de som van de scores uit de settings met alleen links en alleen rechts, wat bevestigt dat het voordeel voortkomt uit het benutten van informatie uit beide richtingen.

3. Impliciete Timestep via Globale Mask Rate

Het paper levert causaal bewijs dat DLMs de globale fractie van gemaskeerde tokens berekenen en gebruiken als een impliciete timestep, ondanks het ontbreken van expliciete timestep embeddings.

• Representatie: Een lineaire probe kan de globale mask-fractie met hoge nauwkeurigheid ($R^2 \approx 0.82-0.91$) herstellen uit de residual stream na de eerste attention laag.
• Causale Nuttigheid: Het patchen van de "mask-rate richting" van een staat met hoge corruptie naar een staat met lage corruptie verhoogt causaal de prediction entropy op de gemaskeerde positie. Dit bevestigt dat het model het globale corruptieniveau gebruikt om zijn onzekerheid te moduleren, waarbij het de mask-rate effectief behandelt als een timestep.

4. Circuit Robuustheid en Emergentie

• Emergentie: Inductie ontstaat abrupt tijdens de training (rond 8–24k stappen) en vereist ten minste twee lagen om te functioneren, wat consistent is met de noodzaak voor één laag om kenmerken te schrijven en een daaropvolgende laag om ze op te halen.
• Lokale Window: Het inductiemechanisme leunt zwaar op een lokaal venster van ongeveer vijf tokens (het gemaskeerde token plus twee buren aan elke zijde). De prestaties dalen aanzienlijk als deze directe buren gemaskeerd zijn, zelfs als de globale corruptierate laag is.

Betekenis en Claims

Dit werk vormt een fundamentele stap in het begrijpen van de interne mechanismen van DLMs, waarbij de aanname dat ze simpelweg AR-gedrag kopiëren, wordt overstegen. De auteurs beweren dat:

1. Mechanistische Divergentie: DLMs inductie implementeren via een symmetrisch, bidirectioneel circuit dat structureel verschilt van de causale inductie circuits in AR-modellen, specifelijk in de manier waarop query- en key-kenmerken worden geconstrueerd vanuit gemaskeerde posities.
2. Contextueel Voordeel: Het prestatieverschil tussen DLMs en AR-modellen in in-context leren toe te schrijven is aan de beschikbaarheid van bidirectionele context, en niet aan een inherent superieur inductie-algoritme zelf.
3. Impliciet State Management: DLMs afleiden intern temporele informatie (timestep) uit de globale staat van de sequentie (mask-fractie) zonder expliciete embeddings, een mechanisme dat causaal verbonden is met prediction entropy.

Het paper concludeert dat deze bevindingen een basis bieden voor het begrijpen van in-context leren in niet-autoregressieve modellen, hoewel het beperkingen erkent met betrekking tot het gebruik van kleine, attention-only modellen en synthetische taken, waarbij wordt opgemerkt dat verder onderzoek nodig is om te bepalen of deze mechanismen schalen naar grotere, MLP-bevattende DLMs en natuurlijke taaltaken.