Training Large Reasoning Models Efficiently via Progressive Thought Encoding

Probleemstelling

Grote Redeneringsmodellen (Large Reasoning Models - LRMs) presteren uitstekend op complexe taken zoals wiskunde en wetenschap, maar ze kampen met een fundamenteel efficiëntieprobleem tijdens het trainen met Versterkend Leren (RL).

• De bottleneck: RL-training vereist lange "rollouts" (generatie van lange denkketens of Chain-of-Thought) om op basis van uitkomsten beloningen te krijgen. Deze autoregressieve decodatie domineert de tijd en het geheugengebruik.
• Het dilemma: Om het geheugen te beperken, worden vaak strategieën zoals "sliding-window caches" (schuifvensters) gebruikt, waarbij oude tokens worden verwijderd. Echter, het verwijderen van tussenliggende denkstappen (intermediate thoughts) verstoort het vermogen van het model om lange contexten te integreren, wat leidt tot een significante daling in redeneerkwaliteit en trainingseffectiviteit.
• De vraag: Kan men LRMs efficiënt trainen onder strikte geheugenbudgetten zonder in te leveren op de redeneeraccuraatheid?

Methodologie: Progressive Thought Encoding

De auteurs introduceren Progressive Thought Encoding, een methode voor parameter-efficiënt fine-tuning die het probleem van het verliezen van context oplost zonder het cache-geheugen te vergroten.

Kernprincipes:

1. Niet wegwerpen, maar coderen: In plaats van verplaatste (evicted) tokens simpelweg te verwijderen, encodeert de methode hun informatie in vaste vectorrepresentaties.
2. Integratie in LoRA: Deze contextuele informatie wordt dynamisch ingebed in lichte LoRA-adapters (Low-Rank Adaptation). Hierdoor onthoudt het model de essentiële redeneersignalen zonder de cache-grootte te hoeven vergroten.
3. Online Adaptatie tijdens RL: Tijdens de RL-rollout (bijvoorbeeld met GRPO - Grouped Reinforcement Policy Optimization):
   • Het model genereert tokens totdat de cache vol is.
   • Wanneer tokens worden verwijderd (geëjecteerd), worden ze gebruikt om een compacte latente representatie bij te werken.
   • Een globale query-vector ($q_g$) fungeert als een samenvatting van alle tot nu toe verplaatste context.
   • De LoRA-weights worden bijgewerkt met een berekende delta ($\Delta W$) die de informatie van de verplaatste tokens vastlegt.
   • Dit proces verloopt continu: elke keer dat de cache vol raakt, wordt de contextstaat ($S_e$) genormaliseerd en opgeteld bij de nieuwe verplaatste context, waardoor het model een "streaming" geheugen behoudt.
4. Cache-beleid: Tijdens training worden vraag-tokens permanent bewaard (als anker), terwijl alleen de "denk-tokens" (thinking tokens) onderhevig zijn aan een schuifvenster-strategie.

Belangrijkste Bijdragen

1. Formalisatie van het probleem: De auteurs identificeren de inefficiëntie van RL-training voor LRMs onder lange rollouts en formuleren dit als een optimalisatieprobleem met cache-beperkingen.
2. Nieuwe Methode: Ze stellen "Progressive Thought Encoding" voor, een techniek die leert van verplaatste tokens om redeneervermogen te behouden binnen een beperkt geheugen.
3. Empirisch Bewijs: Uitgebreide experimenten tonen aan dat de methode zowel de trainings-efficiëntie als de robustheid van inferentie aanzienlijk verbetert, en een nieuwe standaard zet voor schaalbaar redeneren.

Resultaten

De methode werd getest op drie modellen (Qwen2.5-3B, Qwen2.5-7B, en DeepSeek-R1-Distill-Llama-8B) over zes uitdagende wiskundige benchmarks (waaronder MATH500, Olympiade, AMC en AIME 2024/2025).

Kernresultaten:

• Accuraatheid: De methode behaalde gemiddeld +19,3% verbetering ten opzichte van LoRA-finetuning en +29,9% ten opzichte van niet-gefine-tuned LRMs.
• AIME Prestaties: Op de zeer moeilijke AIME2024/2025 benchmarks werd een verbetering van +23,4% in nauwkeurigheid bereikt onder dezelfde strakke cache-budgetten.
• Geheugenefficiëntie:
  • De piekgeheugengebruik (Peak GPU Memory) werd met bijna 50% verlaagd ten opzichte van volledige cache-training.
  • Het berekende vermogen (TFLOPs) voor attention werd aanzienlijk gereduceerd (bijv. van 7,4 TFLOPs naar 4,6 TFLOPs voor het 8B-model).
• Robuustheid: In tegenstelling tot baselines die snel degradëren bij kleiner wordende cache-vensters, behoudt de voorgestelde methode hoge nauwkeurigheid, zelfs bij zeer kleine vensters (bijv. 768 tokens).
• Schaalbaarheid: De methode stelt het model in staat om effectief te redeneren met generatielengten tot 64K tokens, zelfs met een vast cache-venster van 1K tokens, waarbij de prestaties lineair blijven verbeteren.

Betekenis en Impact

Dit werk is significant omdat het de fundamentele trade-off tussen geheugenbeperkingen en redeneerkwaliteit doorbreekt.

• Efficiëntie: Het maakt RL-training van grote redeneringsmodellen haalbaar op hardware met beperkt geheugen (bijv. consumer GPU's of kleinere clusters), wat de drempel voor onderzoek en toepassing verlaagt.
• Kwaliteit: Het toont aan dat het "vergeten" van tokens niet noodzakelijk is; door de informatie slim te comprimeren in de modelweights, kan het model een "full-context" redeneervermogen nabootsen.
• Toekomst: Het opent de weg voor schaalbare RL-training en adaptieve verwijderingsstrategieën, en suggereert dat toekomstige systemen nog langer en complexere denkketens kunnen hanteren zonder exponentiële kostenstijgingen.

Kortom, Progressive Thought Encoding biedt een elegante oplossing om grote redeneringsmodellen te trainen en in te zetten binnen realistische hardware-beperkingen, zonder in te leveren op de complexiteit van de taken die ze moeten oplossen.