💬 NLP

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

本文介绍了局部分支路由（Local Branch Routing, LBR），这是一种基于标记级的测试时扩展框架，通过扩展局部前瞻树并使用轻量级路由器选择最优分支，从而高效地增强语言模型的推理能力，进而实现端到端强化学习，并在数学推理任务上超越了现有的离散思维链和软标记基准。

原作者： 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

发布于 2026-06-25

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： 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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是对论文《局部分支路由》（Local Branch Routing）的解释，使用了简单的语言和日常类比。

核心问题：思考太慢或视野太窄

想象你正在尝试解决一个非常难的数学题，或者在规划一次复杂的旅行。你有一个聪明的助手（AI）在协助你。

目前的 AI 助手通常以两种方式之一工作：

“单路径”漫步者： 他们循序渐后地思考，一旦脑中蹦出第一个念头就立刻认定。如果他们在早期走错了路，可能会陷入困境，因为他们从未观察过其他的选项。
“全图”探索者： 他们试图同时写出所有可能的解题路径，检查所有路径，然后选出最好的一个。这非常准确，但就像为了找一句话而去读遍图书馆里所有的书一样——这会消耗过多的时间和计算能力。

这篇论文的作者想要寻找一个“金发姑娘原则”（意指恰到好处）的解决方案：一种既能通过观察几种不同的可能性来做出更好选择，又不会因检查所有内容而陷入泥潭的方法。

解决方案：局部分支路由 (LBR)

作者提出了一种名为局部分支路由 (Local Branch Routing) 的新方法。你可以把它理解为一种**“先看一眼，再做决定”**的策略。

以下是其工作原理，分步骤使用徒步旅行者选择小径的类比进行说明：

1. “向前看”（生长树结构）

AI 不会立即选择下一个词（或路径），而是会停顿一下。它会将接下来的几个词（或路径标记）想象成真实的。

论文术语： 扩展一个小的局部前瞻树 (Expands a small local lookahead tree)。
类比： 想象你站在一个分叉路口。你不是直接选一条路，而是快速在路径 A 上走 3 步，在路径 B 上走 3 步，然后在路径 C 上也走 3 步。你还没有决定走哪条路，你只是在脑海中“行走”，以观察地形如何。

2. “路由器”（决策者）

一旦 AI “走过”了这些短路径，它就会观察结果。它会问：“这些短路径中，哪一条看起来最有希望？”

论文术语： 使用轻量级路由器选择深度为 1 的子树 (Uses a lightweight router to select the depth-1 subtree)。
类比： 一位聪明的向导（路由器）观察了你探索的地形。也许路径 A 是一处悬崖，路径 B 是个沼泽，但路径 C 通往一片美丽的草地。向导指向路径 C 并说：“好，那我们就正式选择这一条吧。”

3. “剪枝与移动”（向前推进）

AI 正式写下路径 C 的第一步。它丢弃了路径 A 和路径 B 的想法（剪枝）。然后，它将起点移动到该第一步的终点，并重复这个过程：再次向前看，挑选最好的下一步，然后继续前进。

论文术语： 剪枝–移动–生长解码过程 (Prune–shift–grow decoding process)。
类比： 你在草地路径上迈出了第一步。现在你到了一个新的位置。你再次向前看，挑选下一个最好的步骤，然后继续行走。

为什么这种方法比其他方法更好？

论文将此与另外两种常见的 AI 思考方式进行了对比：

对比“离散思维链” (Discrete Chain-of-Thought) —— “单路径”漫步者：
- 问题： “单路径”漫步者在看到路径长什么样之前，就必须决定走哪条路。这就像是在没开门之前就选择一扇门。
- LBR 的优势： LBR 在决定之前先“打开门”（走过路径）。论文表明，在走过几步之后，这些“隐藏状态”（路径的心理图像）包含了有助于做出更好决定的宝贵线索。
对比“软标记分支” (Soft-Token Branching) —— “模糊的混合体”：
- 问题： 一些方法试图通过将所有路径融合在一起，变成一种“模糊”的平均值来同时观察所有路径。这就像看一张照片，其中三条路径重叠在一起。你无法清晰地看到任何单一路径的细节。
- LBR 的优势： LBR 保持路径的离散性（独立且清晰）。它分别走路径 A、路径 B 和路径 C，并将它们进行明确的比较。论文发现，保持路径分离可以让 AI 看到特定的细节（如悬崖或草地），而这些细节在“模糊”的混合中会丢失。

结果：他们发现了什么？

作者在两类任务上进行了测试：

合成规划（一个虚构的游戏）： 他们创建了一个 AI 需要在图中导航的谜题。他们发现 LBR 表现得好得多，因为它可以通过行走短路径找到“线索”，从而做出正确的转向。
数学推理（真实的数学问题）： 他们在困难的数学基准测试（如竞赛中使用的测试）上测试了 LBR。
- 结果： LBR 比标准的“单路径”方法和“模糊混合”方法解决了更多的正确问题。
- 效率： 它实现这一点时，并不需要检查宇宙中所有的可能解。它只是检查了少数局部选项，做出了聪明选择，然后继续前进。

总结

局部分支路由 (Local Branching Routing) 就像是给了 AI 一个“手电筒”，让它在做决定之前可以窥探几步之后的景象。它并不试图预见整个未来（那太昂贵了），但它也不会盲目猜测。通过观察一些短期的可能性、清晰地进行比较并挑选出最好的路径，AI 在解决困难推理问题时变得更加聪明和准确。

技术摘要：通过局部分支路由实现高效且可训练的语言模型测试时扩展

问题陈述

语言模型（LLM）推理领域的最新进展表明，在推理阶段分配额外的计算量（测试时扩展）可以显著提高性能。然而，现有方法面临着一个关键的权衡：

思维链（CoT）采样： 虽然有效，但标准的 CoT 仍是单线程的，在观察到替代方案的后果之前，就必须提交当前的 token。
基于搜索的方法： 如思维树（Tree of Thoughts）或蒙特卡洛树搜索（MCTS）等方法通过探索中间思维树来引入宽度。然而，这些方法通常粒度较粗、计算成本高，且难以作为单一的可端到端训练的解码策略进行优化。
软 Token 方法： 如 Soft Thinking 或 Multiplex Thinking 等技术通过将候选嵌入合并为连续混合体来提供 token 级的宽度。虽然这些方法是可训练的，但它们模糊了候选分支的离散身份，可能掩盖特定概念的信息。

本文认为，语言模型可以受益于一种轻量级的、token 级的测试时扩展机制，该机制允许决策过程受到局部未来候选隐藏状态的启发，而不会产生全解搜索的高昂成本或软混合带来的歧义。

方法论：局部分支路由 (Local Branch Routing, LBR)

作者提出了 局部分支路由 (LBR)，这是一种可训练的、token 级的分支与路由解码框架。LBR 通过一个滚动式的局部前瞻树来增强标准的自回归生成。

解码流水线

在每个解码步骤 $t$ ，给定已提交的前缀 $x_{<t}$ ，LBR 执行一个四阶段的“剪枝–移动–生长”（prune–shift–grow）过程：

生长 (Grow)： LBR 不仅仅采样单个下一个 token，而是扩展一个以 $x_{<t}$ 为根、宽度为 $K$ 、深度为 $L$ 的局部树 $\mathcal{A}_t$ 。每一个采样的节点都会被前向传播通过语言模型以计算其隐藏状态 $h_t(v)$ 。这创建了一个局部隐藏状态树，而非仅仅是一组候选字符串。
路由 (Route)： 一个轻量级的路由器 $\rho_\phi$ 观察前向传播后的树，并选择一个深度为 1 的子树进行提交。路由器使用**集合注意力机制（set-attention mechanism）**将每个候选子树（例如，当 $L=1$ 时为单个 token，当 $L>1$ 时为一个 token 及其后续内容）编码为一个向量。这些向量通过跨子树注意力进行比较，从而产生相对评分，使路由器能够在兄弟未来的竞争中做出对比性决策。
剪枝 (Prune)： 路由器选择最佳子树（索引为 $k^*_t$ ）。所有其他子树均被丢弃。所选子树的根 token 成为提交的下一个 token $x_t$ 。
移动与重构 (Shift and Regrow)： 所选子树被向前移动（移除其根节点，其后代成为下一个步骤的新前缀）。随后会生长出一个新层以恢复树的深度 $L$ ，循环往复。

训练目标

LBR 定义了一个可处理的树轨迹似然（tree-trajectory likelihood），能够实现针对可验证奖励（RLVR）的端到端强化学习。

随机性： 只有两个操作具有随机性：从基础模型中采样新生成的节点 $G_t$ ，以及采样路由器决策 $k^*_t$ 。剪枝、移动和重用节点的过程是确定性的。
似然度： 轨迹 $F$ 的对数概率是新生成的 token 以及路由器选择概率的对数之和：
$\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)$
优化： 该框架使用似然比目标（例如 GRPO）进行训练，其中来自验证器的优势 $A(F)$ 既奖励基础模型的 token 生成，也奖励路由器的选择决策。

核心贡献

具有离散身份的 Token 级分支： 与通过连续混合合并候选者的软 Token 方法不同，LBR 将候选者保留为离散的前向传播分支。这使得路由器能够利用后候选隐藏状态（post-candidate hidden states）——即在候选 token 被生成并由模型处理后才显现出的信息。
可处理的端到端训练： 通过定义特定的树轨迹似然，LBR 允许基础模型和路由器在与离散 token 推理系统相同的 RLVR 原则下进行联合优化。
高效的测试时扩展： 该框架在 token 级别引入了宽度，且具有浅层前瞻（ $L$ 很小，例如 1 或 2），避免了全深度搜索的计算开销，同时仍能利用超越根节点下一个 token 分布的证据。
集合注意力路由器： 作者引入了一种使用跨子树注意力来比较兄弟候选者的路由器架构，证明了路由是一个对比性决策，而非独立的重排序。

实验结果

合成层次规划

在一项涉及基数平移图可达性的合成任务上：

性能： LBR 达到了最高的靶点准确率 (0.967)，优于离散 CoT (0.914) 和 Soft Thinking (0.845)。
机制分析： 探测实验表明，后候选隐藏状态在 LBR 中比分支前的状态更具预测性，能够更好地预测最终目标。这证实了前向传播候选者允许模型访问标准自回归解码无法获取的下游信息。
表示： 与通过连续混合模糊概念身份的 Soft Thinking 不同，LBR 保留了离散的分支身份，使得正确的路径更容易被恢复。

数学推理基准测试

在 DeepSeek-R1-Distill-Qwen 模型（1.5B 和 7B）的六个基准测试（Minerva, AIME, MATH500 等）上进行了评估：

性能提升： LBR 在以下方面相比之下提升了 Pass@1 和 Pass@32：
- 离散思维链 (Discrete CoT)。
- 原始离散 token RLVR。
- 与 RL 兼容的软 token 分支基线（如 Multiplex Thinking）。
深度消融： 将前瞻深度从 $L=1$ 增加到 $L=2$ 进一步提升了性能，尤其是在 7B 主干模型上，这表明更深的局部延续提供了更多的路由证据。
路由器消融： 去除跨子树注意力（使用独立路由器）会导致性能下降，验证了兄弟分支的对比性比较对于性能至关重要。
训练动态： 在 RLVR 训练期间，随着路由器变得更加果断，其熵降低；同时，基础模型在路由位置的熵也随之降低，这表明了两者之间的协同进化（co-adaptation），即模型学会了生成更具辨识度的局部分支。

意义与主张

本文声称，局部分支路由 (Local Branch Routing) 提供了一种高效、可训练且离散的语言模型测试时扩展方式。其主要意义在于证明了：

后候选证据是有价值的： 路由决策受益于由候选 token 诱导的隐藏状态，而非仅仅是分支前的分布。
离散分支优于软合并： 保留离散分支身份比软 token 混合能实现更清晰的决策，因为后者可能会模糊概念身份。
可扩展性： 该框架在单线程 CoT 和高成本的全树搜索之间提供了一个折中方案，通过轻量级的局部分支实现了有效的扩展。

作者承认了**局部性（locality）**的局限性：LBR 受限于浅层前瞻窗口，并未进行完整解的全局搜索。建议未来的工作可以将 LBR 与自适应深度或选择性全局搜索相结合。