Technische Samenvatting: Efficiënte en Trainbare Taalmodel Test-Time Scaling via Local Branch Routing

Probleemstelling

Recente vooruitgang in de redeneerkracht van taalmodellen (LLM's) laat zien dat het toewijzen van extra rekenkracht tijdens de inferentie (test-time scaling) de prestaties aanzienlijk kan verbeteren. Echter, bestaande benaderingen kampen met een kritieke afruil:

• Chain-of-Thought (CoT) Sampling: Hoewel effectief, blijft standaard CoT enkeldraadig, waarbij men zich vastlegt op één token voordat de gevolgen van alternatieven worden waargenomen.
• Zoekgebaseerde Methoden: Benaderingen zoals Tree of Thoughts of Monte Carlo Tree Search (MCTS) introduceren breedte door bomen van tussenliggende gedachten te verkennen. Deze zijn echter vaak grofmazig, rekenintensief en moeilijk end-to-end te optimaliseren als een enkel, trainbaar decoderingsbeleid.
• Soft-Token Methoden: Technieken zoals Soft Thinking of Multiplex Thinking bieden token-niveau breedte door kandidaat-embeddings te versmelten tot continue mengsels. Hoewel trainbaar, vervagen deze methoden de discrete identiteit van kandidaat-takken, wat potentieel concept-specifieke informatie kan vertroebelen.

Het artikel stelt dat taalmodellen kunnen profiteren van een lichtgewicht, token-niveau test-time scaling mechanisme dat beslissingen kan laten informeren door de verborgen toestanden (hidden states) van lokale toekomstige kandidaten, zonder de kosten van volledige oplossing-niveau zoektocht of de ambiguïteit van zachte mengsels te dragen.

Methodologie: Local Branch Routing (LBR)

De auteurs stellen Local Branch Routing (LBR) voor, een trainbaar, token-niveau branch-and-route decoderingsframework. LBR breidt standaard autoregressieve generatie uit met een rollende lokale lookahead-boom.

Decoderingspipeline

Bij elke decoderingsstap $t$, gegeven een gecommitteerde prefix $x_{<t}$, voert LBR een vierfasig "prune–shift–grow" proces uit:

1. Grow: In plaats van één enkel volgend token te samplen, breidt LBR een lokale boom $\mathcal{A}_t$ uit met breedte $K$ en diepte $L$, geworteld bij $x_{<t}$. Elke gesamplede node wordt doorgestuurd (forwarded) door het taalmodel om de verborgen toestand $h_t(v)$ te berekenen. Dit creëert een boom van lokale verborgen toestanden in plaats van enkel een verzameling kandidaat-strings.
2. Route: Een lichtgewicht router $\rho_\phi$ observeert de doorgestuurde boom en selecteert één diepte-1 subtak om te committeren. De router gebruikt een set-attention mechanisme om elke kandidaat-subtak (bijv. een enkel token voor $L=1$, of een token plus zijn continuaties voor $L>1$) te coderen naar een vector. Deze vectoren worden vergeleken via cross-subtree attention om een relatieve score te produceren, waardoor de router een contrastieve beslissing kan nemen tussen evenknieke toekomsten (sibling futures).
3. Prune: De router selecteert de beste subtak (index $k^*_t$). Alle andere subtakken worden weggegooid. De wortel-token van de geselecteerde subtak wordt het gecommitteerde volgende token $x_t$.
4. Shift and Regrow: De geselecteerde subtak wordt naar voren verschoven (de wortel wordt verwijderd en de nakomelingen worden de nieuwe prefix voor de volgende stap). Een nieuwe laag wordt gegroeid om de boomdiepte $L$ te herstellen, en de cyclus herhaalt zich.

Trainingsdoelstelling

LBR definieert een behapbare tree-trajectory likelihood die end-to-end reinforcement learning met verifieerbare beloningen (RLVR) mogelijk maakt.

• Stochastiek: Slechts twee operaties zijn stochastisch: het samplen van de nieuw gegroeide nodes $G_t$ vanuit het basismodel en het samplen van de router-beslissing $k^*_t$. Pruning, shifting en het hergebruiken van nodes zijn deterministisch.
• Likelihood: De log-waarschijnlijkheid van een traject $F$ is de som van de log-waarschijnlijkheden van de nieuw gegroeide tokens en de router-selectiekansen:
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• Optimalisatie: Het framework wordt getraind met een likelihood-ratio doelstelling (bijv. GRPO) waarbij de advantage $A(F)$ van een verifier zowel de token-generatie van het basismodel als de selectiebeslissingen van de router beloont.

Belangrijkste Bijdragen

1. Token-niveau Vertakking met Discrete Identiteit: In tegen tegenoverkomende soft-token methoden die kandidaten versmelten tot continue mengsels, behoudt LBR kandidaten als discrete doorgestuurde takken. Dit stelt de router in staat om gebruik te maken van post-candidate hidden states—informatie die pas wordt onthuld nadat een kandidaat-token is gegenereerd en verwerkt door het model.
2. Behapbare End-to-End Training: Door een specifieke likelihood voor de tree-traject te definiëren, maakt LBR het mogelijk om het basismodel en de router gezamenlijk te optimaliseren onder hetzelfde RLVR-principe dat wordt gebruikt in discrete-token redeneersystemen.
3. Efficiënte Test-Time Scaling: Het framework introduceert breedte op token-niveau met een ondiepe lookahead ($L$ is klein, bijv. 1 of 2), waardoor de computationele overhead van volledige diepte-zoektocht wordt vermeden terwijl nog steeds gebruik wordt gemaakt van bewijs dat verder gaat dan de root next-token distributie.
4. Set-Attention Router: De auteurs introduceren een router-architectuur die cross-subtree attention gebruikt om evenknieke kandidaten te vergelijken, waarmee zij aantonen dat routing een contrastieve beslissing is in plaats van onafhankelijke her-rangschikking (reranking).

Experimentele Resultaten

Synthetische Hiërarchische Planning

Op een synthetische taak betreffende radix-getranslateerde graaf-bereikbaarheid:

• Prestaties: LBR bereikte de hoogste doel-accuratesse (0.967), waarmee het Discrete CoT (0.914) en Soft Thinking (0.845) overtrof.
• Mechanisme Analyse: Probing-experimenten toonden aan dat post-candidate hidden states in LBR significant voorspellender zijn van het uiteindelijke doel dan pre-branching states. Dit bevestigt dat het doorsturen van kandidaten het model in staat stelt om downstream informatie te benutten die onbeschikbaar is voor standaard autoregressieve decodering.
• Representatie: In tegenstelling tot Soft Thinking, die concept-identiteiten vertroebelde via continue mengsels, behield LBR de discrete tak-identiteiten, waardoor het juiste pad beter traceerbaar bleef.

Wiskundige Redeneer Benchmarks

Geëvalueerd op DeepSeek-R1-Distill-Qwen modellen (1.5B en 7B) over zes benchmarks (Minerva, AIME, MATH500, etc.):

• Prestatiewinsten: LBR verbeterde zowel Pass@1 als Pass@32 ten opzichte van:
  • Discrete Chain-of-Thought.
  • Vanilla discrete-token RLVR.
  • RL-compatibele soft-token branching baselines (bijv. Multiplex Thinking).
• Diepte Ablatie: Het verhogen van de lookahead diepte van $L=1$ naar $L=2$ verbeterde de prestaties verder, met name op de 7B backbone, wat suggereert dat diepere lokale continuaties meer routing-bewijs bieden.
• Router Ablatie: Het verwijderen van cross-subtree attention (het gebruik van een onafhankelijke router) resulteerde in lagere prestaties, wat valideert dat de contrastieve vergelijking van evenknieke takken cruciaal is.
• Trainingsdynamiek: Tijdens de RLVR-training nam de router-entropie af naarmate de router besluitvaardiger werd, terwijl de entropie van het basismodel op gerouteerde posities ook afnam, wat wijst op co-adaptatie waarbij het model leert om meer onderscheidbare lokale takken te genereren.

Betekenis en Claims

Het artikel claimt dat Local Branch Routing een efficiënte, trainbare en discrete vorm van taalmodel test-time scaling biedt. De primaire betekenis ligt in het aantonen dat:

1. Post-candidate bewijs waardevol is: Beslissingen over routing profiteren van de verborgen toestanden die worden geïnduceerd door kandidaat-tokens, niet alleen van de pre-branching distributie.
2. Discrete vertakking superieur is aan soft merging: Het behouden van discrete tak-identiteiten maakt duidelijkere besluitvorming mogelijk vergeleken met soft-token mengsels, die de concept-identiteit kunnen vertroebelen.
3. Schaalbaarheid: Het framework biedt een middenweg tussen enkeldraadige CoT en dure volledige boom-zoektochten, waardoor effectieve scaling mogelijk wordt via lichtgewicht lokale vertakking.

De auteurs erkennen de beperking van lokaliteit: LBR is beperkt door het ondiepe lookahead-venster en voert geen globale zoektocht uit over volledige oplossingen. Toekomstig werk suggereert het combineren van LBR met adaptieve diepte of selectieve globale zoektocht.