Transformers with RL or SFT Provably Learn Sparse Boolean Functions, But Differently

Technische Samenvatting: Transformers met RL of SFT leren bewijsbaar schaarse Booleaanse functies, maar op verschillende manieren

Probleemstelling
Grote Taalmodellen (LLMs) gebaseerd op de Transformer-architectuur hebben opmerkelijke redeneercapaciteiten aangetoond via Chain-of-Thought (CoT)-prompting. Hoewel Supervised Fine-Tuning (SFT) en Reinforcement Learning (RL) de belangrijkste methoden zijn om deze capaciteiten te verbeteren, blijven de theoretische mechanismen die ten grondslag liggen aan hun leerdynamiek slecht begrepen. Specifiek ontbreekt het aan een rigoureuze analyse die vergelijkt hoe RL (met name met procesbeloningen) en SFT (specifiek zonder teacher forcing) Transformers in staat stellen om complexe, schaarse Booleaanse functies te leren die anders moeilijk op een end-to-end manier te leren zijn. Dit werk adresseert deze lacune door de leervaardigheid van $k$-schaarse Booleaanse functies te onderzoeken—functies die worden bepaald door een kleine subset van $k$ invoerbits—met behulp van een Transformer met één laag die is uitgerust met tussenliggende redeneerstappen.

Methodologie
De auteurs stellen een unificerend theoretisch kader voor om de leerdynamiek van zowel RL als SFT te analyseren op een specifieke klasse van $k$-schaarse Booleaanse functies. Deze functies zijn recursief decomposeerbaar in vaste 2-schaarse Booleaanse functies (bijvoorbeeld pariteit, AND, OR), waardoor een redeneerketen ontstaat die vergelijkbaar is met CoT.

1. Modelarchitectuur: Het onderzoek maakt gebruik van een vereenvoudigde Transformer met één laag bestaande uit positionele encoding, een self-attention-laag en een feedforward-laag. Het model is ontworpen om iteratief tussenliggende redeneerstappen ($y^{(t)}$) te genereren voordat een definitief antwoord wordt geproduceerd.
2. "Voorgeladen" masker: Om theoretische hanteerbaarheid te waarborgen en de structurele priors die voortvloeien uit pretraining te simuleren, leggen de auteurs een specifiek "voorgeladen" masker op de attention-weights. Dit masker beperkt de attention van redeneertokens tot hun relevante kindknopen (zoals gedefinieerd door de recursieve decompositie) en de invoersequentie, waardoor foutpropagatie effectief wordt gecontroleerd.
3. Leerparadigma's:
   • RL met procesbeloning: De Transformer wordt geoptimaliseerd via policy gradient om een beloningsfunctie te maximaliseren waarbij beloningen worden toegewezen aan elke tussenliggende redeneerstap ($r_t$), en niet alleen aan het eindresultaat.
   • SFT zonder teacher forcing: De Transformer wordt getraind om een verlies (hinge loss) te minimaliseren over gegenereerde redeneerketens in vergelijking met ground-truth labels. Cruciaal genereert het model zijn eigen tussenliggende stappen tijdens de training (autoregressieve generatie) in plaats van dat het wordt gevoerd met ground-truth-voorvoegsels (teacher forcing), waarmee het mismatch tussen training en inferentie wordt aangepakt.
4. Kernanalytisch hulpmiddel: De auteurs introduceren een "kritiek gradiëntcomponent" ($\gamma$). Zij tonen aan dat de scheiding van dit component—waarbij het gradiëntsignaal voor relevante posities (kindknopen) strikt groter is dan voor irrelevante posities—dient als een voldoende voorwaarde voor bewijsbare leervaardigheid.

Belangrijkste bijdragen en resultaten
Het artikel stelt voldoende voorwaarden vast waaronder een Transformer met één laag bewijsbaar $k$-schaarse Booleaanse functies kan leren via zowel RL als SFT. De belangrijkste bevindingen zijn:

1. Unificerende leervaardigheidsvoorwaarde: Zowel RL als SFT slagen als het kritieke gradiëntcomponent relevante posities scheidt van irrelevante posities. Deze voorwaarde wordt geverifieerd voor drie fundamentele voorbeelden: $k$-PARITY, $k$-AND en $k$-OR.
2. Verschillende leerdynamiek:
   • RL (Procesbeloning): Onder de scheidingsvoorwaarde leert RL de volledige CoT-redeneerketen gelijktijdig in één gradiëntupdate. De procesbeloning biedt informatieve signalen voor alle stappen, ongeacht de huidige staat van andere stappen, waardoor het model direct de volledige redeneercapaciteit kan verwerven.
   • SFT (Geen teacher forcing): Daarentegen vertoont SFT een stap-voor-stap leergedrag. Vanwege de autoregressieve aard van generatie kan het model stap $t$ niet betrouwbaar leren totdat stap $t-1$ correct is gegenereerd. Bijgevolg vereist SFT $T$ gradiëntupdates (waarbij $T$ de ketenlengte is) om de volledige keten te leren, waarbij per update één stap wordt opgelost.
3. Moeilijkheid van eindbeloning: Het artikel bewijst dat RL dat alleen met een finale (outcome-based) beloning wordt geoptimaliseerd, deze functies niet efficiënt leert. De variantie van de policy gradient is exponentieel klein, waardoor het signaal wordt overspoeld door ruis, wat de noodzaak van procesbeloningen voor complexe redeneertaken verklaart.
4. Uitbreiding naar seriële CoT: De analyse wordt uitgebreid buiten de binaire-boomdecompositie naar seriële CoT-structuren (veelvoorkomend in praktische LLM's), waarmee wordt bevestigd dat dezelfde leerdynamiek (gelijktijdig voor RL, stapsgewijs voor SFT) geldt onder passende masking.

Betekenis en claims
De auteurs claimen dat dit werk de eerste theoretische karakterisering biedt van de leerdynamiek van RL met procesbeloningen voor Transformers, waarmee een kritieke lacune in de literatuur wordt opgevuld. Door teacher forcing en data-augmentatiestrategieën te verwijderen die vaak worden gebruikt in eerdere theoretische werken (bijvoorbeeld Kim & Suzuki, 2025), biedt het onderzoek een realistischere vergelijking tussen RL en SFT.

De betekenis ligt in het onthullen dat hoewel beide methoden leervaardigheid kunnen bereiken, ze dit doen via fundamenteel verschillende mechanismen:

• RL met procesbeloning nabootst de efficiëntie van SFT met teacher forcing door de volledige keten gelijktijdig te leren.
• SFT zonder teacher forcing induceert natuurlijk een stapsgewijs curriculum, een intrinsieke eigenschap van autoregressieve training die geen extern curriculumlearning of datafiltering vereist.

Het artikel concludeert dat empirische vergelijkingen tussen RL en SFT zorgvuldig factoren moeten controleren zoals beloningsontwerp (proces versus outcome) en trainingsregimes (teacher forcing versus autoregressieve generatie), aangezien deze factoren de leerdynamiek aanzienlijk veranderen en de vergelijking van de optimalisatiemethoden zelf kunnen verstoren.