Technische Samenvatting: SuperThoughts: Reasoning Tokens in Superpositie

Probleemstelling

Large Language Models (LLM's) lossen complexe problemen op door expliciete Chain-of-Thought (CoT) sequenties te genereren, waarbij elk token een eenheid van computationele verwerking vertegenwoordigt (één forward pass). Hoewel recente modellen zoals OpenAI o1 en DeepSeek-R1 aantonen dat verhoogde compute tijdens de testtijd de prestaties aanzienlijk verbetert, is deze sequentiële token-generatie computationeel duur.

Bestaande benaderingen die proberen discrete token-generatie te omzeilen door te redeneren in continue latente ruimtes (bijv. COCONUT, latent reasoning), worden geconfronteerd met een kritieke uitdaging: het gebrek aan intermediaire supervisie. Standaard CoT-training vertrouwt op token-niveau cross-entropy loss bij elke stap, wat dichte gradiëntsignalen levert. Wanneer redeneren plaatsvindt in een ongestructureerde latente ruimte, verdwijnt deze supervisie, wat leidt tot instabiliteit tijdens de training, representatieve drift en slechte prestaties op taken met een lange horizon. Daartegenover staan methoden die simpelweg de context comprimeren of het aantal discrete tokens verminderen via fine-tuning; deze verminderen echter niet het aantal vereiste forward passes, wat de impact op de inference-throughput beperkt.

De centrale vraag die dit werk adresseert is: Kan een model worden getraind om te redeneren in een rijkere, superponeerde latente ruimte, terwijl het de intermediare supervisiesignalen van discrete token-generatie behoudt?

Methodologie

De auteurs stellen SuperThoughts voor, een framework dat paren van opeenvolgende CoT-tokens comprimeert tot enkele latente representaties, waardoor het model twee tokens per stap kan decoderen. Deze aanpak halveert het aantal vereiste forward passes voor een redeneerketen van lengte $L$ (reducerend naar $L/2$ stappen) terwijl de discrete token-supervisie behouden blijft.

Architectuur

Het model bestaat uit drie primaire componenten:

1. Compressor: Encodeert een paar token-embeddings $(c_{2i-1}, c_{2i})$ naar een enkele latente vector $x_i$. De auteurs verkennen twee implementaties:
   • Lineaire Projectie: Concateneert embeddings en projecteert ze via een leerbare matrix.
   • Transformer Block: Verwerkt het paar door een kleine Transformer-laag en extraheert de hidden state van het tweede token.
2. Main Module: De primaire redeneer-backbone (geïnitialiseerd vanuit een pre-trained LLM). Het verwerkt de gecomprimeerde vector $x_i$ om een hidden state $h_i$ te produceren en voorspelt het volgende oneven-geïndexeerde token ($c_{2i+1}$).
3. Multi-Token Prediction (MTP) Module: Een lichtgewicht module die het daaropvolgende even-geïndexeerde token voorspelt ($c_{2i+2}$). Deze neemt als input het vorige even token, het nieuw voorspelde oneven token en de hidden state van de Main module, projecteert deze naar een vector $x'_i$ om de volgende voorspelling te genereren. Beide modules delen dezelfde output language modeling head.

Trainingsstrategie

Om de instabiliteit van latent reasoning aan te pakken, wordt een tweestaps trainingsprotocol gehanteerd:

1. Latent Distillatie (Fase 1): De Compressor wordt getraind om de gecomprimeerde latente ruimte af te stemmen op het pre-trained taal-manifold. Een "teacher"-model verwerkt de discrete token-sequentie, terwijl een "student"-model de gecomprimeerde paren verwerkt. De Compressor wordt geoptimaliseerd om de Smooth-L1 afstand tussen de student's hidden states (na verwerking van $x_i$) en de teacher's hidden states (na verwerking van het bijbehorende tokenpaar) te minimaliseren. Dit zorgt ervoor dat de gecomprimeerde representaties de redeneertoestand van de oorspronkelijke discrete sequentie behouden.
2. Joint End-to-End Training (Fase 2): De Main module, MTP module en de Compressor worden gezamenlijk getraind met behulp van standaard cross-entropy loss voor alle voorspelde tokens (zowel Main als MTP outputs).

Adaptieve Inference

In het besef dat niet alle redeneerstappen even moeilijk zijn, introduceren de auteurs een confidence-based adaptief mechanisme. Op elke stap, indien de voorspellingszekerheid (maximale waarschijnlijkheid) van de MTP module onder een drempelwaarde $\tau$ valt, verwerpt het model de MTP-voorspelling. Het valt dan terug op standaard decoding voor die stap, waarbij het enkele token direct naar de Main module voert. Dit stelt het model in staat om meer rekenkracht toe te wijzen aan "moeilijke" tokens, terwijl de efficiëntie op "gemakkelijke" tokens behouden blijft.

Belangrijkste Bijdragen

1. SuperThoughts Architectuur: Een nieuw framework dat tokenparen fuseert in superponeerde embeddings en twee tokens per stap decodeert met behulp van een Main module en een MTP module, wat effectief het aantal vereiste forward passes halveert.
2. Tweestaps Trainingsprotocol: Een methode om latent reasoning te stabiliseren door eerst de gecomprimeerde ruimte af te stemmen via distillatie alvorens gezamenlijke end-to-end training met discrete supervisie uit te voeren.
3. Adaptief Inference Mechanisme: Een dynamisch fallback-systeem dat throughput ruilt voor nauwkeurigheid bij onzekere redeneerstappen, wat prestatieverlies op complexe taken voorkomt.
4. Empirische Evaluatie: Uitgebreide evaluatie op wiskundige en wetenschappelijke redeneer-benchmarks (MATH500, AMC23, OlympiadBench, GPQA-Diamond) die een significante reductie in CoT-lengte aantoont met minimale nauwkeurigheidsverlies.

Resultaten

De auteurs evalueerden SuperThoughts op Qwen2.5-Math modellen (1.5B, 7B, en 14B varianten) over meerdere benchmarks:

• Lengte Reductie: Met het adaptieve mechanisme bereikt SuperThoughts een reductie van 20–35% in CoT-lengte.
• Nauwkeurigheid: De methode behoudt een nauwkeurigheid binnen 1–2 punten van de standaard discrete CoT baseline op de meeste taken. Zo bereikte de 1.5B model met $\tau=0.999$ een MATH500 nauwkeurigheid van 73.0% (tegenover 72.4% baseline) met een 36% reductie in CoT-lengte.
• Schaal-effecten: Grotere modellen (7B) tolereren agressieve compressie beter dan kleinere modellen (1.5B). Zonder adaptieve inference leed het 1.5B model aan aanzienlijke nauwkeurigheidsdalingen (14–21 punten), terwijl het 7B model slechts 5–12 punten daalde.
• Vergelijking met HAMburger: SuperThoughts presteerde beter dan de HAMburger baseline (Liu & Zhang, 2025) op elk compressieniveau, waarbij het een hogere nauwkeurigheid behaalde voor dezelfde CoT-lengte of kortere CoT-lengtes voor dezelfde nauwkeurigheid.
• Inference Versnelling: Hoewel de theoretische versnelling proportioneel is aan de lengtereductie, is de werkelijke wall-clock tijd versnelling iets lager vanwege de overhead van de Compressor en MTP modules. Deze overhead neemt echter af naarmate het model groter wordt; voor het 14B model vertaalde een ~33% reductie in lengte zich in een ~28% reductie in generatietijd.

Betekenis

Het artikel beweert dat SuperThoughts succesvol aantoont dat modellen kunnen redeneren in een superponeerde latente ruimte zonder de stabiliteit die door discrete token-supervisie wordt geboden op te offeren. Door tokenparen te comprimeren en adaptieve inference te gebruiken, biedt het framework een praktisch pad om de computationele kosten van langdurige redeneertaken te verlagen. Het werk suggereert dat grotere modellen voldoende per-stap capaciteit bezitten om superponeerde tokens betrouwbaar te verwerken, wat erop wijst dat toekomstige schaling nog grotere compressie (bijv. superpositie van drie of meer tokens) mogelijk kan maken. Bovendien merken de auteurs op dat naarmate native MTP modules de standaard worden in toekomstige LLM's (bijv. Qwen3-Next, MiMo), de toepassing van SuperThoughts nog eenvoudiger en effectiever zal worden.