Technische Samenvatting: Compressed Latent Reasoning (CoLaR)

Probleemstelling

Large Language Models (LLMs) bereiken superieure prestaties in wiskundig redeneren door middel van Chain-of-Thought (CoT) prompting, waarbij expliciete, op tokens gebaseerde redeneerketens worden gegenereerd. Deze langdurige sequenties zijn echter computationeel duur en inefficiënt, wat zorgt voor een hoge serverbelasting en de schaalbaarheid belemmert, met name onder omstandigheden met een hoge concurrency. Hoewel bestaande efficiëntiemethoden proberen redundante tokens over te slaan of compacte notaties te gebruiken, blijven zij beperkt door ijle, op tokens gebaseerde representaties. Bovaleert, recente latent-gebaseerde benaderingen voor redeneren (bijv. Coconut, CODI) vertrouwen vaak op vaste lengtes van redeneerketens en deterministische latente generatie, wat hun aanpassingsvermogen en exploratiecapaciteiten beperkt.

Methodologie: Compressed Latent Reasoning (CoLaR)

CoLaR is een nieuw framework dat is ontworpen om redeneerprocessen dynamisch te comprimeren naar een dichte latente ruimte, terwijl het vermogen behouden blijft om diverse redeneerpaden te verkennen. De methode werkt via een tweestaps trainingsaanpak:

1. Supervised Fine-Tuning (SFT) met Auxiliaire Doelstellingen

CoLaR gaat verder dan standaard next-token voorspelling door een aanvullende next compressed embedding prediction doelstelling te introduceren.

• Dynamische Compressie: Tijdens de training wordt een compressiefactor $c$ willekeurig gesampled uit een vooraf gedefinieerd bereik $[1, c_{max}]$. De embeddings van $c$ opeenvolgende redeneertokens worden samengevoegd tot één gecomprimeerde embedding met behulp van een Embedding Compress module. Deze module schaalt de som van de embeddings door $1/\sqrt{c}$ om de oorspronkelijke distributie-eigenschappen te behouden, waardoor vervorming door eenvoudige mean pooling wordt voorkomen.
• Dual-Head Architectuur: Het model maakt gebruik van twee koppen:
  • Latent Head: Een probabilistische kop (een two-headed MLP) die het gemiddelde ($\mu$) en de standaarddeviatie ($\sigma$) van de volgende gecomprimeerde embeddingdistributie voorspelt. Dit maakt autoregressieve voorspelling van dichte latente variabelen mogelijk in plaats van deterministische waarden.
  • Language Head: Voorspelt het uiteindelijke antwoord en biedt cruciaal dichte supervisie door de oorspronkelijke redeneertokens te voorspellen die overeenkomen met de gecomprimeerde embeddings. Dit benadert een multi-label classificatietaak om ervoor te zorgen dat het model de semantiek van de gecomprimeerde groepen leert.
• Loss Functie: De trainingsverliesfunctie combineert de taalmodellering-loss (voor tokenvoorspelling) met een latente loss. De latente loss maakt gebruik van een soft-MSE formulering die Mean Squared Error combineert met een entropie-regularisatieterm. Dit moedigt het model aan om de exploratiecapaciteit (diverse latents) te behouden terwijl het de data fit, wat de onderbezetting (under-fitting) problemen aanpakt die worden waargenomen bij standaard Negative Log-Likelihood (NLL) loss op eenvoudigere datasets.

2. Reinforcement Learning (RL) voor Exploratie en Exploitatie

Om voorbij het nabootsen van de CoT-leraar te gaan en het potentieel van latent redeneren te ontsluiten, gebruikt CoLaR Group Relative Policy Optimization (GRPO).

• Probabilistische Sampling: Omdat de Latent Head een waarschijnlijkheidsverdeling uitvoert, kan het model diverse latente redeneerketens genereren voor dezelfde input.
• Reward Mechanisme: Voor elke vraag wordt een groep outputs (latente ketens + antwoorden) gegenereerd. Beloningen (rewards) worden toegekend op basis van de correctheid van het antwoord (1 voor correct, 0 voor incorrect).
• Balansstrategie: De GRPO-doelstelling berekent groep-genormaliseerde beloningen. Cruciaal is dat de beloning wordt gemiddeld over het aantal tokens/latents in de keten. Dit ontwerp moedigt het model inherent aan om kortere, correcte redeneerpaden te exploiteren (omdat de beloning per token hoger is), terwijl het nog steeds diverse paden exploreert om correcte oplossingen te vinden.

Belangrijkste Bijdragen

1. Dynamic-Speed Latent Reasoning: CoLaR introduceert een framework dat autoregressief latente variabelen voorspelt die de semantiek van meerdere woordtokens inkapselen. Dit maakt een dynamische aanpassing van de redeneersnelheid tijdens de inferentie mogelijk door simpelweg een gewenste compressiefactor te prompten.
2. Probabilistic Latent Head & RL Integratie: In tegen tegenover eerdere deterministische latente methoden, gebruikt CoLaR een probabilistische Latent Head. Gecombineerd met RL, stelt dit het model in staat om diverse redeneerpaden te verkennen en meer compacte oplossingen te exploiteren, wat de efficiëntie aanzienlijk verbetert zonder de nauwkeurigheid op te offeren.
3. Nieuwe Trainingsdoelstellingen: De introductie van de next compressed embedding prediction taak en de soft-MSE latente loss stelt het model in staat om effectief dichte representaties te leren en de balans tussen exploratie en exploitatie te beheren.

Experimentele Resultaten

Uitgebreide evaluaties werden uitgevoerd op vier basisschool wiskunde datasets (GSM8k, GSM8k-hard, SVAMP, MultiArith) en de meer uitdagende MATH dataset.

• Prestaties vs. Latent Baselines: Op basisschool datasets behaalde CoLaR een 14,1% hogere nauwkeurigheid dan state-of-the-art latent-gebaseerde baselines (bijv. Coconut) bij vergelijkbare compressieverhoudingen.
• Efficiëntie vs. Expliciete CoT: Vergeleken met expliciete CoT reduceerde CoLaR de lengte van de redeneerketen met 53,3% met slechts een 4,8% prestatievermindering.
• RL Verbeteringen op Uitdagende Taken: Op de MATH dataset resulteerde het toepassen van RL op CoLaR in een winst van 5,36% in nauwkeurigheid ten opzichte van de niet-RL versie, terwijl de lengte van de redeneerketen met 82,8% werd verminderd.
• Scaling: Experimenten die schalen van 1B naar 8B parameters testten, bevestigden dat de efficiëntie- en nauwkeurigheidswinsten van Colar standhouden bij grotere modelcapaciteiten. CoLaR-8B-RL overtrof zijn CoT-leraar op de GPQA benchmark (37,5% vs 35,7%) terwijl de redeneerketen met 69% werd ingekort.
• Generalisatie: Het model demonstreerde een robuuste generalisatie naar onbekende compressiefactoren en out-of-domain datasets (bijv. MultiArith, GPQA).

Betekenis

Het artikel stelt dat CoLaR de kritieke flessenhals van computationele kosten in LLM-redeneren aanpakt door de verschuiving van ijle token-niveau verwerking naar dichte latente-ruimte redeneren. De betekenis ligt in:

• Dynamische Aanpasbaarheid: Het vermogen om de redeneersnelheid (compressiefactor) tijdens de inferentie aan te passen, biedt een flexibele afweging tussen latentie en nauwkeurigheid, geschikt voor variërende real-world applicatiebeperkingen.
• Efficiëntie zonder Opoffering: Door gebruik te maken van probabilistisch latent redeneren en RL, bereikt CoLaR substantiële reducties in de lengte van de redeneerketen (tot 82,8%) terwijl het de nauwkeurigheid op complexe taken handhaaft of zelfs verbetert, wat het idee uitdaagt dat efficiëntie ten koste moet gaan van prestaties.
• Exploratie-Exploitatie in de Latente Ruimte: Het werk demonstreert dat het introduceren van stochasticiteit en reinforcement learning in latente redeneerprocessen modellen in staat stelt om efficiëntere redeneerpaden te ontdekken die deterministische, vaste-lengte benaderingen missen.

De auteurs erkennen beperkingen, waaronder het feit dat de huidige prestaties op sommige benchmarks de expliciete CoT benaderen (in plaats van consistent te overtreffen) en uitdagingen in het generaliseren naar niet-gehele compressiefactoren vanwege discrete tokenisatiebeperkingen. Echter, het framework vestigt een veelbelovende richting voor efficiënt, schaalbaar en adaptief LLM-redeneren.