💬 NLP

Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing

Dieses Paper führt Local Branch Routing (LBR) ein, ein Framework zur skalenbaren Testzeit-Optimierung auf Token-Ebene, das das logische Schlussfolgern von Sprachmodellen effizient verbessert, indem es lokale Lookahead-Bäume erweitert und einen leichtgewichtigen Router verwendet, um optimale Zweige auszuwählen, wodurch ein End-to-End-Reinforcement-Learning ermöglicht wird und bestehende Discrete Chain-of-Thought- sowie Soft-Token-Baselines bei mathematischen Denkaufgaben übertrifft.

Ursprüngliche Autoren: Yutong Yin, Mingyu Jin, Jin Pan, Changyi Yang, Zijie Xia, Dhruv Pai, Shuming Hu, Zhen Zhang, Chenyang Zhao, Jinman Zhao, Wujiang Xu, Raymond Li, Xin Eric Wang, Julian McAuley, Zhaoran Wang

Veröffentlicht 2026-06-25

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Yutong Yin, Mingyu Jin, Jin Pan, Changyi Yang, Zijie Xia, Dhruv Pai, Shuming Hu, Zhen Zhang, Chenyang Zhao, Jinman Zhao, Wujiang Xu, Raymond Li, Xin Eric Wang, Julian McAuley, Zhaoran Wang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das große Problem: Zu langsam oder zu engstirnig denken

Stellen Sie sich vor, Sie versuchen, ein sehr schwieriges mathematisches Problem zu lösen oder eine komplexe Reise zu planen. Sie haben einen smarten Assistenten (die KI), der Ihnen hilft.

Derzeit arbeiten KI-Assistenten meist auf eine von zwei Arten:

Der „Ein-Pfad“-Wanderer: Er denkt Schritt für Schritt und legt sich sofort auf die erste Idee fest, die ihm in den Sinn kommt. Wenn er früh eine falsche Abbiegung nimmt, kann er stecken bleiben, weil er nie die anderen Optionen betrachtet hat.
Der „Vollständige-Landkarte“-Entdecker: Er versucht, jeden mögliche Lösungspfad gleichzeitig aufzuschreiben, prüft sie alle und wählt dann den besten aus. Das ist sehr genau, aber es ist, als würde man versuchen, jedes Buch in einer Bibliothek zu lesen, um einen einzigen Satz zu finden – das kostet zu viel Zeit und Rechenleistung.

Die Autoren dieser Arbeit wollten eine „Goldlöckchen“-Lösung finden: einen Weg, um sich ein paar verschiedene Möglichkeiten anzusehen, um eine bessere Entscheidung zu treffen, ohne sich durch das Überprüfen von allem zu verzetteln.

Die Lösung: Local Branch Routing (LBR)

Die Autoren schlagen eine neue Methode namens Local Branch Routing vor. Betrachten Sie dies als eine „Vorausschau-und-dann-Entscheiden“-Strategie.

So funktioniert es, Schritt für Schritt, anhand der Analogie eines Wanderers, der einen Pfad wählt:

1. Das „Vorausschauen“ (Den Baum wachsen lassen)

Anstatt sofort das nächste Wort (oder den nächsten Wegweiser) zu wählen, hält die KI inne. Sie stellt sich die nächsten Wörter (oder Wegweiser) so vor, als wären sie real.

Der Fachbegriff der Arbeit: Expands a small local lookahead tree. (Erweitert einen kleinen lokalen Lookahead-Baum.)
Die Analogie: Stellen Sie sich vor, Sie stehen an einer Weggabelung. Anstatt einfach einen Pfad zu wählen, gehen Sie schnell 3 Schritte auf Pfad A, 3 Schritte auf Pfad B und 3 Schritte auf Pfad C. Sie legen sich noch nicht fest; Sie „gehen“ diese Wege nur in Ihrem Geist ab, um zu sehen, wie das Gelände aussieht.

2. Der „Router“ (Der Entscheidungsträger)

Sobalt die KI diese kurzen Pfade „abgelaufen“ ist, sieht sie sich die Ergebnisse an. Sie fragt: „Welcher dieser kurzen Pfade sieht am vielversprechendsten aus?“

Der Fachbegriff der Arbeit: Uses a lightweight router to select the depth-1 subtree. (Verwendet einen leichtgewichtigen Router, um den Tiefen-1-Teilstammbaum auszuwählen.)
Die Analogie: Ein kluger Führer (der Router) betrachtet das Gelände, das Sie erkundet haben. Vielleicht führt Pfad A zu einer Klippe, Pfad B zu einem Sumpf, aber Pfad C führt zu einer wunderschönen Wiese. Der Führer zeigt auf Pfad C und sagt: „Okay, entscheiden wir uns offiziell für diesen.“

3. Das „Beschneiden und Verschieben“ (Vorwärtsbewegen)

Die KI schreibt nun offiziell den ersten Schritt von Pfad C auf. Sie verwirft die Ideen von Pfad A und Pfad B (Pruning/Beschneiden). Dann verschiebt sie ihren Ausgangspunkt an das Ende dieses ersten Schrittes und wiederholt den Prozess: erneut vorausschauen, den besten nächsten Schritt wählen und weitergehen.

Der Fachbegriff der Arbeit: Prune–shift–grow decoding process. (Prune–Shift–Grow-Dekodierungsprozess.)
Die Analogie: Sie machen den ersten Schritt auf dem Wiesenpfad. Jetzt sind Sie an einem neuen Ort. Sie schauen wieder voraus, wählen den nächsten besten Schritt und gehen weiter.

Warum ist dies besser als andere Methoden?

Die Arbeit vergleicht dies mit zwei anderen gängigen Arten, wie eine KI denkt:

Vs. „Discrete Chain-of-Thought“ (Der „Ein-Pfad“-Wanderer):
- Das Problem: Der „Ein-Pfad“-Wanderer muss entscheiden, welchen Weg er einschlägt, bevor er sieht, wie der Pfad aussieht. Es ist, als würde man eine Tür wählen, ohne sie zu öffnen.
- Der LBR-Vorteil: LBR öffnet die Tür (geht den Pfad ab), bevor es entscheidet. Die Arbeit zeigt, dass die „Hidden States“ (das mentale Bild des Pfades) nach dem Gehen einiger Schritte wertvolle Hinweise enthalten, die helfen, eine bessere Entscheidung zu treffen.
Vs. „Soft-Token Branching“ (Die „verschwommene Mischung“):
- Das Problem: Einige Methoden versuchen, alle Pfade gleichzeitig zu betrachten, indem sie sie zu einem „verschwommenen“ Durchschnitt vermischen. Es ist, als würde man ein Foto betrachten, auf dem alle drei Pfade übereinander projiziert sind. Man kann die Details keines einzelnen Pfades klar erkennen.
- Der LBR-Vorteil: LBR hält die Pfade diskret (getrennt und klar). Es geht Pfad A ab, dann Pfad B, dann Pfad C und vergleicht sie deutlich. Die Arbeit fand heraus, dass das Trennen der Pfade es der KI ermöglicht, spezifische Details (wie eine Klippe oder eine Wiese) zu sehen, die in einer „verschwommenen“ Mischung verloren gehen würden.

Die Ergebnisse: Was haben sie herausgefunden?

Die Autoren testeten dies bei zwei Arten von Aufgaben:

Synthetische Planung (Ein erfundenes Spiel): Sie erstellten ein Puzzle, bei dem die KI durch einen Graphen navigieren musste. Sie fanden heraus, dass LBR viel besser darin war, es zu lösen, da es die „Hinweise“, die durch das Abgehen der kurzen Pfade gefunden wurden, nutzen konnte, um die richtige Abbiegung zu machen.
Mathematisches Denken (Echte Matheaufgaben): Sie testeten LBR an schwierigen mathematischen Benchmarks (wie sie bei Wettbewerben verwendet werden).
- Das Ergebnis: LBR löste mehr Probleme korrekt als die standardmäßige „Ein-Pfad“-Methode und die „verschwommene Mischung“.
- Effizienz: Dies gelang ohne die Notwendigkeit, jeden möglichen Lösungsweg im Universum zu überprüfen. Es hat lediglich ein paar lokale Optionen geprüft, eine kluge Wahl getroffen und ist dann weitergegangen.

Das Fazit

Local Branch Routing ist wie das Geben einer „Taschenlampe“ an eine KI, die es ihr ermöglicht, ein paar Schritte voraus zu spähen, bevor sie eine Entscheidung trifft. Sie versucht nicht, die ganze Zukunft zu sehen (was zu teuer wäre), aber sie rät auch nicht einfach blindlings. Indem sie ein paar kurze Möglichkeiten betrachtet, sie klar miteinander vergleicht und die beste auswählt, wird die KI klüger und präziser beim Lösen schwieriger Denkaufgaben.

Technisches Resümee: Effiziente und trainierbare Skalierung der Sprachmodell-Testzeit via Local Branch Routing

Problemstellung

Jüngste Fortschritte im Bereich des Reasoning von Sprachmodellen (LLMs) zeigen, dass die Zuweisung zusätzlicher Rechenleistung während der Inferenz (Test-Time Scaling) die Leistung signifikant steigern kann. Bestehende Ansätze stehen jedoch vor einem kritischen Kompromiss:

Chain-of-Thought (CoT) Sampling: Während CoT effektiv ist, bleibt das Standard-CoT Single-Threaded und verpflichtet sich auf ein Token, bevor die Konsequenzen von Alternativen beobachtet werden können.
Suchbasierte Methoden: Ansätze wie Tree of Thoughts oder Monte Carlo Tree Search (MCTS) führen Breite durch das Erkunden von Bäumen intermediärer Gedanken ein. Diese sind jedoch oft grobkörnig, rechenintensiv und schwer als ein einziger, end-to-end trainierbarer Decoding-Policy zu optimieren.
Soft-Token-Methoden: Techniken wie Soft Thinking oder Multiplex Thinking bieten eine Token-Ebene der Breite durch das Verschmelzen von Kandidaten-Embeddings in kontinuierliche Mischungen. Obwohl diese trainierbar sind, verwischen sie die diskrete Identität der einzelnen Zweige, was konzeptspezifische Informationen verschleiern kann.

Das Paper postuliert, dass Sprachmodelle von einem leichtgewichtigen, Token-basierten Test-Time Scaling-Mechanismus profitieren könnten, der es ermöglicht, Entscheidungen durch die Hidden States lokaler zukünftiger Kandidaten zu informieren, ohne die Kosten einer vollständigen Lösungssuche oder die Ambiguität von Soft-Mixturen zu verursachen.

Methodik: Local Branch Routing (LBR)

Die Autoren schlagen Local Branch Routing (LBR) vor, einen trainierbaren, Token-basierten Branch-and-Route-Decoding-Framework. LBR erweitert die standardmäßige autoregressive Generierung um einen rollierenden lokalen Lookahead-Baum.

Decoding-Pipeline

In jedem Decoding-Schritt $t$ , gegeben ein committed Prefix $x_{<t}$ , führt LBR einen vierstufigen „Prune–Shift–Grow“-Prozess aus:

Grow: Anstatt ein einzelnes nächstes Token zu sampeln, expandiert LBR einen lokalen Baum $\mathcal{A}_t$ der Breite $K$ und Tiefe $L$ , der in $x_{<t}$ wurzelt. Jeder gesampelte Knoten wird durch das Sprachmodell weitergeleitet (forwarded), um seinen Hidden State $h_t(v)$ zu berechnen. Dies erzeugt einen Baum aus lokalen Hidden States anstelle einer bloßen Menge von Kandidaten-Strings.
Route: Ein leichtgewichtiger Router $\rho_\phi$ beobachtet den weitergeleiteten Baum und wählt einen Tiefe-1-Subtree aus, um diesen zu committen. Der Router verwendet einen Set-Attention-Mechanismus, um jeden Kandidaten-Subtree (z. B. ein einzelnes Token für $L=1$ , oder ein Token plus seine Fortsetzungen für $L>1$ ) in einen Vektor zu kodieren. Diese Vektoren werden mittels Cross-Subtree-Attention verglichen, um ein relatives Scoring zu erzeugen, was dem Router ermöglicht, eine kontrastive Entscheidung unter Geschwister-Zukunftspfaden zu treffen.
Prune: Der Router wählt den besten Subtree (Index $k^*_t$ ). Alle anderen Subtrees werden verworfen. Das Root-Token des ausgewählten Subtrees wird das nächste committed Token $x_t$ .
Shift and Regrow: Der ausgewählte Subtree wird nach vorne verschoben (sein Root wird entfernt, und seine Nachkommen werden das neue Prefix für den nächsten Schritt). Eine neue Ebene wird gewachsen, um die Baumtiefe $L$ wiederherzustellen, und der Zyklus wiederholt sich.

Trainingsziel

LBR definiert eine handhabbare Tree-Trajectory-Likelihood, die ein End-to-End Reinforcement Learning mit verifizierbaren Belohnungen (RLVR) ermöglicht.

Stochastizität: Nur zwei Operationen sind stochastisch: das Sampling der neu gewachsenen Knoten $G_t$ aus dem Basismodell und das Sampling der Router-Entscheidung $k^*_t$ . Pruning, Shifting und das Wiederverwenden von Knoten sind deterministisch.
Likelihood: Die Log-Wahrscheinlichkeit einer Trajektorie $F$ ist die Summe der Log-Wahrscheinlichkeiten der neu gewachsenen Tokens und der Router-Selektionswahrscheinlichkeiten:
$\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)$
Optimierung: Das Framework wird unter Verwendung eines Likelihood-Ratio-Objectives (z. B. GRPO) trainiert, bei dem der Advantage $A(F)$ von einem Verifier sowohl das Token-Generation des Basismodells als auch die Selektionsentscheidungen des Routers belohnt.

Zentrale Beiträge

Token-Ebene Branching mit diskreter Identität: Im Gegensatz zu Soft-Token-Methoden, die Kandidaten in kontinuierliche Mischungen verschmelzen, bewahrt LBR die Kandidaten als diskrete weitergeleitete Zweige. Dies erlaubt dem Router, Post-Candidate Hidden States zu nutzen – Informationen, die erst nach der Verarbeitung eines Kandidaten-Tokens durch das Modell offengelegt werden.
Handhabbare End-to-End-Trainierbarkeit: Durch die Definition einer spezifischen Likelihood für die Tree-Trajektorie ermöglicht LBR die gemeinsame Optimierung des Basismodells und des Routers unter demselben RLVR-Prinzip, das in diskreten Token-Reasoning-Systemen verwendet wird.
Effiziente Test-Time Scaling: Das Framework führt Breite auf der Token-Ebene mit einem flachen Lookahead ( $L$ ist klein, z. B. 1 oder 2) ein, wodurch der Rechenaufwand einer vollständigen Tiefensuche vermieden wird, während gleichzeitig Evidenz über die Root-Next-Token-Verteilung hinaus genutzt wird.
Set-Attention Router: Die Autoren führen eine Router-Architektur ein, die Cross-Subtree-Attention nutzt, um Geschwister-Kandidaten zu vergleichen, und demonstrieren, dass Routing eine kontrastive Entscheidung und kein unabhängiges Reranking ist.

Experimentelle Ergebnisse

Synthetische hierarchische Planung

Auf einer synthetischen Aufgabe bezüglich Radix-translatierter Graph-Erreichbarkeit:

Performance: LBR erreichte die höchste Zielgenauigkeit (0.967) und übertraf damit Discrete CoT (0.914) und Soft Thinking (0.845).
Mechanismus-Analyse: Probing-Experimente zeigten, dass Post-Candidate Hidden States in LBR signifikant prädiktiver für das finale Ziel sind als Pre-Branching-States. Dies bestätigt, dass das Forwarding von Kandidaten es dem Modell ermöglicht, auf Downstream-Informationen zuzugreifen, die bei standardmäßiger autoregressiver Dekodierung nicht verfügbar sind.
Repräsentation: Im Gegensatz zu Soft Thinking, welches Konzept-Identitäten durch kontinuierliche Mischungen verwischt, bewahrte LBR die diskreten Zweig-Identitäten, was die Rekuperation des korrekten Pfades erleichterte.

Mathematische Reasoning-Benchmarks

Evaluierung auf DeepSeek-R1-Distill-Qwen Modellen (1.5B und 7B) über sechs Benchmarks (Minerva, AIME, MATH500, etc.):

Leistungssteigerungen: LBR verbesserte sowohl Pass@1 als auch Pass@32 gegenüber:
- Discrete Chain-of-Thought.
- Vanilla diskretem-Token RLVR.
- RL-kompatiblen Soft-Token-Branching-Baselines (z. B. Multiplex Thinking).
Depth Ablation: Eine Erhöhung der Lookahead-Tiefe von $L=1$ auf $L=2$ verbesserte die Performance weiter, insbesondere beim 7B-Backbone, was darauf hindeutet, dass tiefere lokale Fortsetzungen mehr Routing-Evidenz liefern.
Router Ablation: Das Entfernen der Cross-Subtree-Attention (unter Verwendung eines unabhängigen Routers) führte zu einer geringeren Performance, was validiert, dass der kontrastive Vergleich von Geschwister-Zweigen entscheidend ist.
Trainingsdynamik: Während des RLVR-Trainings sank die Router-Entropie, als der Router entscheidungsfreudiger wurde, während auch die Entropie des Basismodells an den gerouteten Positionen sank, was auf eine Ko-Adaption hindeutet: Das Modell lernt, unterscheidbarere lokale Zweige zu generieren.

Bedeutung und Ansprüche

Das Paper behauptet, dass Local Branch Routing eine effiziente, trainierbare und diskrete Form des Test-Time Scalings für Sprachmodelle bietet. Die primäre Bedeutung liegt in der Demonstration, dass:

Post-Candidate Evidence wertvoll ist: Routing-Entscheidungen profitieren von den Hidden States, die durch Kandidaten-Tokens induziert werden, nicht nur von der Pre-Branching-Verteilung.
Diskretes Branching dem Soft-Merging überlegen ist: Die Bewahrung diskreter Zweig-Identitäten ermöglicht eine klarere Entscheidungsfindung im Vergleich zu Soft-Token-Mischungen, welche die Konzept-Identität verschleiern können.
Skalierbarkeit: Das Framework bietet einen Mittelweg zwischen Single-Threaded CoT und der teuren vollständigen Baumsuche und ermöglicht effektives Scaling durch leichtgewichtiges lokales Branching.

Die Autoren räumen die Einschränkung der Lokalität ein: LBR ist auf ein flaches Lookahead-Fenster beschränkt und führt keine globale Suche über vollständige Lösungen durch. Zukünftige Arbeit wird vorgeschlagen, LBR mit adaptiver Tiefe oder selektiver globaler Suche zu kombinieren.