Technische Samenvatting: CoLT – Redeneren met Chain of Latent Tool Calls

Probleemstelling

Chain-of-Thought (CoT) prompting is essentieel gebleken voor het verbeteren van de redeneercapaciteiten van Large Language Models (LLM's) door complexe problemen op te splitsen in sequentiële stappen. Echter, expliciete CoT vereist het genereren van redeneerketens token voor token, wat leidt tot hoge computationele kosten tijdens de inferentie. Hoewel methoden voor impliciete CoT (bijv. COCONUT, CODI, COLAR) zijn voorgesteld om te redeneren in een continue latente ruimte om de tokenlengte te verminderen, lijden deze over het algemeen aan twee kritieke beperkingen:

1. Structurele Rigiditeit: Ze vereisen vaak significante wijzigingen aan de modelarchitectuur en uitgebreide training.
2. Toepasbaarheid: Deze beperkingen beperken hun bredere adoptie over verschillende taken en modelstructuren.

Bovendien werken bestaande latente redeneermethoden vaak volledig in een "black-box" latente ruimte, wat potentieel ten koste gaat van de interpreteerbaarheid en de vooraf getrainde redeneerkracht van het basismodel.

Methodologie: Het CoLT-framework

De auteurs stellen CoLT (Chain-of-Latent-Tools) voor, een nieuw framework dat latent redeneren implementeert als "parametrische tool calls". In plaats van volledig binnen de latente ruimte te redeneren, staat CoLT het hoofd-LLM toe om in de expliciete tokenruimte te opereren, terwijl het uitbesteden van de expansie van redeneringsstappen naar externe, differentieerbare decoders.

Kernmechanisme

1. Seed Tokens en Latent Tool Calls:
   Tijdens het redeneerproces genereert het hoofd-LLM speciale "seed tokens" om een gecondenseerde redeneringsstap te vertegenwoordigen. Deze bestaan uit:

   • Body Tokens (<BDY>): Dragers van latente embeddings die de gecomprimeerde informatie van de redeneringsstap bevatten.
   • Trigger Tokens (<TRG>): Indicatoren die het einde van een latente tool call signaleren en specificeren welke externe decoder moet worden aangeroepen.

   Een typische latente tool call sequentie verschijnt als: [Vraag] [Vorige Stappen] <BDY> ... <TRG>.

2. Differentieerbare Decoding:
   Wanneer een trigger token wordt geactiveerd, worden de verborgen toestanden (hidden states) van de seed tokens geëxtraheerd en doorgegeven aan een externe differentieerbare decoder (bijv. een lichtgewicht Transformer, RNN of Multi-hot decoder).

   • De decoder pakt de verborgen toestanden weer uit naar volledige, expliciete teksttokens (bijv. wiskundige vergelijkingen).
   • Deze gedecodeerde tokens worden aan de context toegevoegd, waarna het hoofd-LLM zijn autoregressieve generatie voortzet met de bijgewerkte, leesbare context.

3. Trainingsdoelstellingen:
   Het framework wordt end-to-end getraind met behulp van twee gesuperviseerde verliescomponenten:

   • Main Loss ($L_{main}$): Zorgt ervoor dat het hoofdmodel de seed tokens en trigger-signalen correct genereert.
   • Latent Loss ($L_{lat}$): Zorgt ervoor dat de decoder de teksttokens correct reconstrueert uit de seed embeddings.
   • Totale Loss: $L_{sup} = L_{main} + L_{lat}$.

4. Compatibiliteit met Reinforcement Learning (RL):
   Omdat de decoders differentieerbaar zijn en sampling ondersteunen, kan het gehele redeneerproces worden gemodelleerd als een gesprek met meerdere rondes. Dit maakt de toepassing van Group Relative Policy Optimization (GRPO) mogelijk om diverse redeneerpaden te verkennen en correcte ketens te versterken zonder handmatig ruis te injecteren.

Belangrijkste Bijdragen

1. Nieuw Framework: CoLT introduceert een methode om latent redeneren te implementeren als tool calls, waardoor het hoofdmodel in de vooraf getrainde tekstruimte kan redeneren terwijl het tussenliggende stappen comprimeert in seed tokens.
2. Decoder Flexibiliteit: De auteurs ontwerpen en evalueren diverse decoderstructuren (Transformer, RNN, Multi-hot), waarmee zij aantonen dat hoewel Transformers het best presteren, het framework aanpasbaar is aan verschillende vormen.
3. Efficiëntie en Prestaties: CoLT bereikt een hogere nauwkeurigheid en kortere redeneerlengtes vergeleken met baseline latente modellen (inclusief COCONUT, CODI, COLAR en SIM-CoT).
4. RL Integratie: Het framework is compatibel met reinforcement learning algoritmen, waardoor prestatiewinsten op moeilijke datasets mogelijk zijn door middel van exploratie.

Experimentele Resultaten

De auteurs evalueerden CoLT op vier wiskundige redeneerbenchmarks: GSM8k, GSM8k-hard, SVAMP, en MultiArith.

• Nauwkeurigheid vs. Lengte: CoLT presteerde beter dan bestaande latente redeneer-baselines (inclusief COCONUT, CODI, COLAR, en SIM-CoT) in zowel nauwkeurigheid als in de lengte van de redeneerketen. Zo behaalde CoLT op GSM8k-Aug een nauwkeurigheidsverbetering van 5% ten opzichte van COLAR (2x), terwijl de lengte van de redeneertokens werd verminderd van 12,7 naar 7,73.
• Generalisatie: Het model vertoonde een sterke out-of-domain generalisatie, met name op de MultiArith dataset waar het de supervised CoT baseline overtrof.
• Ablatie-studies:
  • Decoderlagen: Het verhogen van het aantal decoderlagen en de lengte van de seed tokens verbeterde de nauwkeurigheid, zij het met afnemende meeropbrengsten in verhouding tot de computationele kosten.
  • Decoderstructuur: Transformer-decoders leverden de beste prestaties op, gevolgd door RNN's. Multi-hot decoders presteerden slecht, wat suggereert dat recurrente structuren beter geschikt zijn voor het uitpakken van gepakte seeds.
• Reinforcement Learning: Het toepassen van GRPO op CoLT-modellen op GSM8k-Aug en de MATH-dataset (met gebruik van DeepSeek-R1-Distill-Qwen-1.5B) verbeterde de prestaties verder, wat de bekwaamheid van het framework bewijst om te leren van diverse redeneerpaden.

Betekenis en Claims

Het artikel beweert dat CoLT een praktische oplossing biedt voor de afweging tussen redeneerefficiëntie en modelcapaciteit. Door het hoofdmodel te verankeren in de vooraf getrainde tekstruimte en externe decoders te gebruiken om de "latente" expansie af te handelen, behoudt CoLT de interpreteerbaarheid van redeneerketens (de uiteindelijke output is volledig tekstueel) terwijl de inferentiekosten aanzienlijk worden verlaagd.

De auteurs benadrukken dat CoRT minimale wijzigingen vereist aan de oorspronkelijke modelstructuur, wat het breder toepasbaar maakt dan eerdere methoden voor latent redeneren die zware architecturale wijzigingen vereisen. Daarnaast suggereert de compatibiliteit van het framework met reinforcement learning een pad om latent redeneren op te schalen naar complexere taken, mits de granulariteit van de tool calls en de decoderstructuren geschikt worden afgestemd op het specifieke domein. Het artikel concludeert dat hoewel de focus momenteel ligt op wiskundig redeneren, het concept van latente tool calls potentieel heeft voor bredere toepassingen, inclusclusief entiteitsretrieval en multimodale redenering.