Technische Samenvatting: Leren van State-Tracking uit Code met Lineaire RNN's

Probleemstelling

State-tracking (toestandsbewaking) is een fundamentele capaciteit voor het begrijpen van programma-executie, speltoestanden en wereldmodellen. Hoewel theoretisch werk heeft vastgesteld dat er een kloof bestaat waarbij Transformers uitblinken in associatief geheugen en Recurrent Neural Networks (RNN's) goed presteren bij state-tracking, bestaat er een aanzienlijke kloof tussen hoe deze modellen worden geëvalueerd versus hoe ze worden getraind.

De meeste bestaande benchmarks voor state-tracking (bijv. permutatiecompositie) maken gebruik van een sequence-to-sequence opstelling: het direct mappen van een sequentie van acties naar een sequentie van toestanden. Dit verschilt fundamenteel van het next-token prediction paradigma dat wordt gebruikt om moderne Large Language Models (LLMs) te trainen, waarbij het model het volgende token in een stroom moet voorspellen die acties afwisselt met gedeeltelijke onthullingen van de toestand (bijv. print-statements). Bovendien gaat echte code-executie vaak gepaard met partiële observeerbaarheid en probabilistische transities (bijv. vertakkingen op basis van verborgen toestanden of willekeurige keuzes), wat een setting creëert waar de ware toestand niet volledig observeerbaar is, wat het onderhoud van een belief-distributie vereist in plaats van een deterministische toestand.

Methodologie

1. Herformulering van State-Tracking via Code (REPL-traces)

Om de kloof tussen abstracte automaten-benchmarks en LLM-trainingsdoelstellingen te overbruggen, zetten de auteurs permutatiecompositie-taken om in Python REPL-traces.

• Mechanisme: In plaats van een sequentie van permutaties $\sigma_1, \dots, \sigma_m$ direct te mappen naar cumulatieve toestanden, verwerkt het model een sequentie van Python-commando's (variabele toewijzingen, swaps) afgewisseld met print-statements die gedeeltelijke toestandinformatie onthullen.
• Doelstelling: Het model wordt getraind via next-token prediction om het volgende token in de trace te voorspellen, wat een commando of een onthulde toestandswaarde kan zijn. Dit dwingt het model om state-tracking impliciet te leren door middel van de voorspellingsopgave.

2. Architecturale Vergelijking

De studie vergelijkt drie architecturale klassen op deze taken:

• Transformers: Standaard attention-gebaseerde modellen.
• Lineaire RNN's (DeltaNet [0, 1]): Lineaire RNN's beperkt tot eigenwaarden in $[0, 1]$, die bekend staan als ongeschikt voor complexe state-tracking zoals parity of permutatiecompositie.
• Lineaire RNN's (DeltaNet [-1, 1]): Lineaire RNN's met uitgebreide eigenwaarde-bereiken $[-1, 1]$, die theoretisch in staat zijn om deterministische permutatiecompositie op te lossen.

3. Theoretisch Kader: PFSA-SR

Om de moeilijkheid van state-tracking in echte code te analyseren, introduceren de auteurs de Probabilistic Finite-State Automaton with State Reveals (PFSA-SR).

• Definitie: Een tupel $(Q, \Sigma, \delta, \rho, q_0)$ waarbij $\delta$ een probabilistische transitiekernel is en $\rho$ een reveal-functie die inputs naar subsets van toestanden mapt.
• Dynamiek: Het systeem evolueert via probabilistische transities, maar de observeerder ontvangt "harde" restricties (reveals) die de verzameling mogelijke toestanden inkrimpen (support pruning), in plaats van zachte waarschijnlijkheden te bieden.
• Belief Update: Tracking vereist het onderhouden van een belief-distributie $b_t$ over toestanden. De update omvat een reveal-stap (het wegstrepen van inconsistente toestanden en hernormaliseren) gevolgd door een transitie-stap. Cruciaal is dat de hernormalisatiestap $f(x) = x/\|x\|_1$ niet-lineair is.

4. Representatie-analyse

Het onderzoek onderzoekt twee representaties voor belief-toestanden in lineaire RNN's:

• Joint Representatie: Expliciteert alle $n!$ permutaties. De auteurs stellen dat dit numeriek instabiel is voor lineaire RNN's omdat de input-injectieterm $B(x_t)$ niet adaptief de waarschijnlijkheidsmassa van overlevende configuraties kan aanvullen na een gedeeltelijke reveal, wat leidt tot exponentiële norm-afname.
• Marginale Representatie: Houdt de marginale waarschijnlijkheid van elke positie onafhankelijk bij (een $n \times n$ dubbel stokastische matrix). Hoewel dit het mogelijk maakt dat $B(x_t)$ direct zekerheid injecteert voor onthulde entries, tonen de auteurs aan dat het nog steeds vatbaar is voor exponentiële afname onder adversariële reveal-sequenties, waarbij specifieke niet-onthulde entries nooit direct worden aangevuld.

Belangrijkste Resultaten

1. Supervisie-dichtheid en Architectuur

• Transformers: Vereisen dichte toestandssupervisie om state-tracking te leren. Naarmate de afstand tussen print-statements (reveals) groter wordt, daalt de prestatie van de Transformer snel. Ze falen in het generaliseren naar volledige permutaties, zelfs met dichte supervisie.
• Lineaire RNN's (DeltaNet [-1, 1]): Presteren uitstekend in deze setting. Ze leren toestanden betrouwbaar te tracken, zelfs met ijle supervisie, en vertonen robuuste lengte-extrapolatie (generaliseren naar sequenties langer dan de training).
• Lineaire RNN's (DeltaNet [0, 1]): Falen in het oplossen van de taak, wat bevestigt dat het uitgebreide eigenwaarde-bereik $[-1, 1]$ noodzakelijk is voor complexe state-tracking.

2. Interpreteerbaarheid van State-Tracking

Analyse van de getrainde DeltaNet [-1, 1] laat zien dat state-tracking gelokaliseerd is naar specifieke componenten:

• Eigenwaarde-extensie: Ten minste één attention head in een specifieke laag (bijv. Laag 12, Head 3) leert zijn decay-parameter $\beta$ in te stellen op 2, wat resulteert in een eigenwaarde van $-1$.
• Subruimte-structuur: Deze specifieke head opereert in een aanzienlijk lager-dimensionale lineaire subruimte vergeleken met andere heads, wat consistent is met de theoretische vereisten voor het tracken van permutaties.
• Interventie: Het terugschalen van de $\beta$-range van deze specifieke head naar $[0, 1]$ veroorzaakt een onmiddellijke prestatie-instorting, terwijl het schalen van andere heads geen effect heeft, wat bevestigt dat deze head de toegewezen state-tracking mechanisme is.

3. Limieten van Lineaire RNN's in Probabilistische Settings

Het paper demonstreert dat hoewel lineaire RNN's deterministische taken oplossen, ze inherente barrières ervaren in probabilistische settings (PFSA-SR):

• Exponentiële Afname: Onder de joint representatie veroorzaken gedeeltelijke reveals dat de ongenormaliseerde toestand-norm exponentieel afneemt ($1/2 \to 1/4 \to 1/8 \dots$), omdat de lineaire recursie niet in staat is de noodzakelijke niet-lineaire hernormalisatie uit te voeren.
• Adversariële Kwetsbaarheid: Zelfs met de marginale representatie kan een tegenstander een sequentie van probabilistische swaps en selectieve reveals construeren die de waarschijnlijkheidsmassa van niet-onthulde entries exponentieel doet afnemen totdat deze verdwijnt in de eindige precisie.
• Conclusie: Stabiele belief-handhaving in probabilistische settings met partiële observeerbaarheid lijkt onhaalbaar voor lineaire RNN's zonder ofwel niet-lineaire hernormalisatie in de recursie, ofwel frequente volledige toestand-resets.

Betekenis en Claims

Het paper claimt een realistischer testbed te bieden voor state-tracking dat aansluit bij het next-token prediction doel van taalmodellen, waarbij zij verder gaan dan synthetische sequence-to-sequence benchmarks.

1. De Kloof Overbruggen: Het toont aan dat lineaire RNN's die in staat zijn tot state-tracking (specifiek DeltaNet met uitgebreide eigenwaarden) deze capaciteiten kunnen leren via next-token prediction op code-achtige traces, terwijl Transformers dat niet kunnen.
2. Theoretische Barrière: Het identificeert een fundamentele beperking van lineaire RNN's in real-world code scenario's. Echte code bevat vaak probabilistische transities en partiële observeerbaarheid (bijv. verborgen branching, willekeurige keuzes). De auteurs tonen aan dat het tracken van toestanden in dergelijke omgevingen niet-lineaire hernormalisatie vereist, wat lineaire RNN's intern niet kunnen uitvoeren. Dit suggereert dat de "complexiteitskloof" die in real-world taken wordt waargenomen, voort kan komen uit de probabilistische aard van code-executie in plaats van alleen sequentielengte of complexiteit.
3. Implicaties voor Benchmarking: De auteurs suggereren dat het benchmarken van architecturen uitsluitend op Deterministische Eindige-Toestandsautomaten (DFA) taken, de potentie van niet-lineaire RNN's en de beperkingen van lineaire RNN's maskeert. Het voorgestelde PFSA-SR framework en de bijbehorende faalmodi (exponentiële norm-afname) bieden een theoretische verklaring voor waarom standaard lineaire recurrente modellen moeite hebben wanneer redeneren over onzekerheid vereist is.

Het paper sluit bescheiden af met de opmerking dat de experimenten zich richten op synthetische permutatiegroepen en dat toekomstig werk nodig is om deze bevindingen uit te breiden naar real-world executie-traces en om hybride architecturen te onderzoeken die lineaire recursie afwisselen met periodieke niet-lineaire normalisatie.