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

# 1 引言 来源：https://arxiv.org/html/2606.25354 ![[无标题图片]](https://arxiv.org/html/2606.25354v1/figs/northwestern_logo.png)![[无标题图片]](https://arxiv.org/html/2606.25354v1/figs/tilde_logo.png) 2026年6月 **通过局部分支路由实现高效且可训练的语言模型测试时扩展**  
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¹  
¹西北大学，²罗格斯大学，³威斯康星大学麦迪逊分校，⁴卡内基梅隆大学，⁵LMSYS Org，⁶Tilde Research，⁷加州大学圣巴巴拉分校，⁸多伦多大学，⁹不列颠哥伦比亚大学，¹⁰加州大学圣地亚哥分校  

###### 摘要  
测试时扩展提升了语言模型的推理能力，但现有方法常面临一个棘手的权衡：长思维链采样仍是单线程的，而句子级或解决方案级搜索则可能计算代价高昂且难以进行端到端训练。我们引入了局部分支路由（LBR），这是一种词元级测试时扩展框架，它扩展一个小的局部前瞻树，将所有采样的分支通过语言模型前向传播，并使用一个轻量级路由来选择要提交的深度为1的子树。通过在候选局部未来的隐藏状态上进行路由，LBR 允许每个词元决策使用超越根节点下一个词元分布的额外证据，同时避免了完整的解决方案级搜索。由此产生的剪枝-平移-生长解码过程保留了离散的分支身份，并定义了一个可处理的树轨迹似然：新生长的节点在首次采样时被计数，并且路由决策被赋予显式概率。这使得使用可验证奖励的端到端强化学习成为可能，在与离散词元 RLVR 相同的似然比原则下联合优化基础模型和路由器。在合成层次规划任务上，LBR 显示出候选后隐藏状态提供了有用的路由证据。在数学推理基准上，LBR 在 Pass@1 和 Pass@32 上都优于离散思维链、朴素离散词元 RLVR 以及 RL 兼容的软词元分支基线。这些结果表明，轻量级局部分支为语言模型测试时扩展提供了一种高效、可训练且离散的形式。  

近期语言模型推理的进展表明，在推理时分配更多计算量可以显著提升性能。思维链提示鼓励模型外化中间推理步骤[28](https://arxiv.org/html/2606.25354#bib.bib15)，自一致性通过采样多个推理轨迹来提高准确性[27](https://arxiv.org/html/2606.25354#bib.bib17)，而对测试时扩展的广泛研究显示，额外计算的价值很大程度上取决于其分配方式[24](https://arxiv.org/html/2606.25354#bib.bib16)。基于搜索的方法进一步通过探索中间思想或行动树引入宽度[32](https://arxiv.org/html/2606.25354#bib.bib29), [10](https://arxiv.org/html/2606.25354#bib.bib19)，但这种搜索通常是粗粒度的、计算代价高昂的，且难以作为单一可训练解码策略进行优化。  

本文探讨语言模型能否在词元级别获得一种轻量级的测试时宽度。标准的自回归解码在观察到邻近替代方案所诱发的隐藏状态之前，就决定了一个下一个词元。然而在许多推理问题中，一个局部选择——比如一个数字、操作符、变量或短语——可以决定模型接下来进入哪种潜在状态。这表明候选词元不仅应根据根节点的下一个词元分布进行评分，还应根据它们所诱发的局部未来状态进行评分。  

我们提出了*局部分支路由*（LBR），一个可训练的词元级分支与路由解码框架。在每个解码步骤中，LBR 扩展一个宽度为 \(K\) 且深度为 \(L\) 的小型局部前瞻树，将所有采样的分支通过语言模型前向传播，并使用一个轻量级路由器来选择要提交的深度为 1 的子树。选定的词元被追加到输出中，未选中的子树被剪枝，而幸存的子树则向前平移并重新增长。因此，每个词元决策可以利用候选后隐藏状态，同时避免完整的解决方案级搜索。LBR 还定义了一个可处理的树轨迹似然。唯一的随机操作是从语言模型中采样新生长的树节点以及采样路由器决策；剪枝、平移和重用是确定性的。这种分解使得 LBR 可以使用可验证奖励的强化学习进行端到端训练，遵循近期数学推理系统中使用的似然比原则[23](https://arxiv.org/html/2606.25354#bib.bib38), [9](https://arxiv.org/html/2606.25354#bib.bib25)。  

软思维提供了一种相关的词元级分支与合并机制，它将一个提交的离散词元替换为候选词元嵌入的连续混合[33](https://arxiv.org/html/2606.25354#bib.bib1)。RL 兼容的变体，如多路复用思维，添加了随机候选采样，使得采样的候选词元可以通过似然比 RL 进行训练[25](https://arxiv.org/html/2606.25354#bib.bib12)。然而，这些方法在将来计算之前将候选词元合并为一个单一的软词元。LBR 相反地将候选词元保留为离散的前向传播分支，在其隐藏状态之间进行路由，并为树生长和路由决策两者分配显式概率。  

参见说明  
图 1：局部分支路由解码流程。LBR 维护一个已进行前向传播的候选延续的局部滚动树。在每一步中，路由器利用当前树中所有节点的隐藏状态来选择深度为 1 的一个子树，提交其根词元，剪枝其他子树，将选定的子树向前平移，并生长一个新层以恢复深度 \(L\)。顶行显示了主要实验设置 \(L=1, K=3\)，其中路由使用 \(K\) 个候选下一词元的词元后隐藏状态。底行展示了对于 \(L=2, K=3\) 的相同剪枝-平移-生长机制，此时路由器也观察到更深层局部延续的隐藏状态。  

我们在合成层次规划和现实数学推理上评估了 LBR。合成任务显示，候选后隐藏状态提供了分支特有的信息，这些信息在分支之前是无法获得的。在数学基准上，LBR 优于离散思维链、朴素 RLVR 以及 RL 兼容的软词元分支基线。这些结果表明，局部离散分支为语言模型测试时扩展提供了一种高效且可训练的形式。  

## 2 预备知识  
令 \(\mathcal{V}\) 表示词汇表，\(\pi_\theta\) 为自回归语言模型。给定一个前缀 \(x_{<t}\)，该模型定义一个条件下一个词元分布 \(\pi_\theta(\cdot \mid x_{<t})\)。语言模型隐藏状态记为 \(h_t(v)\)，表示在时间步 \(t\) 生成词元 \(v\) 之后的表示。我们关注以下设置：在每个解码步骤，我们从根分布 \(\pi_\theta(\cdot \mid x_{<t})\) 中采样 \(K\) 个候选下一个词元，并将所有这些候选词元通过模型前向传播以观察它们各自的隐藏状态。然后，路由器使用这些候选隐藏状态（可能还有更深层局部延续的隐藏状态）来决定要提交哪个候选词元。此局部前瞻树框架是局部分支路由的基础。我们也将通过连续混合进行软分支的相关方法作为基线。  

## 3 局部分支路由  

LBR 是一个用于语言模型推理的增量式分支与路由解码算法。在每一步，模型维护一个滚动树，该树包含从当前输出前缀 \(x_{<t}\) 出发的、已前向传播的候选延续。该树具有固定的宽度 \(K\) 和深度 \(L\)，这意味着在每个步骤中，最多有 \(K\) 个候选路径，每条路径最长包含 \(L\) 个词元。该树通过一个循环的剪枝-平移-生长周期进行更新。  

**算法 1** 局部分支路由  
**输入：** 语言模型 \(\pi_\theta\)，路由器 \(\rho_\phi\)，宽度 \(K\)，深度 \(L\)，温度 \(\tau\)，初始上下文 \(x_{<0}\)。  
**输出：** 解码序列 \(y_1, \ldots, y_T\)。  
\(t \gets 0\)  
通过自回归采样从 \(x_{<0}\) 生长一个完整树（初始填充）  
**while** 未满足停止条件 **do**  
    **路由**：从当前树的根节点开始，路由器 \(\rho_\phi\) 在 \(K\) 条深度为 \(1\) 的子树之间分配概率质量。根据路由分布 \(\rho_\phi(\cdot \mid \text{树})\) 采样一个选中的子树 \(k^*\)。  
    **提交**：将选中的子树的根词元 \(v_k\) 追加到输出 \(y_t\) 中。  
    **剪枝**：移除所有未选中的子树。  
    **平移**：将选中的子树提升，使其子节点成为新的根节点。如果新根具有父节点，则丢弃父节点；当前树深度变为 \(L-1\)。  
    **生长**：对于剩余的 \(K-1\) 个候选位置以及每个可能不完全的候选路径，通过从语言模型 \(\pi_\theta(\cdot \mid \text{路径})\) 中自回归采样并前向传播来生长一个新的完整词元层以填充到深度 \(L\)。  
    \(t \gets t+1\)  
**end while**  

### 算法描述  
算法 1 提供了 LBR 的伪代码。初始化时，我们通过自回归采样从初始上下文 \(x_{<0}\) 生长一个宽度为 \(K\)、深度为 \(L\) 的完整树。然后，LBR 进入主要循环：  

**路由：** 路由器 \(\rho_\phi\) 接收当前树中所有节点及其隐藏状态并输出一个选择分布，用于选择 \(K\) 条可能的深度为 1 的子树（即候选下一个词元及其局部未来）。路由器通常使用一个轻量级编码器：它首先通过一组注意力机制在节点之间传递信息，然后为每个候选子树输出一个分数。解码时，路由器通过 softmax 进行采样：对于每个候选子树 \(k\)，其分数 \(s_{t,k}\) 由路由器从树的隐藏状态中计算得出，采样概率为 \(\rho_\phi(k \mid \text{树}) = \text{softmax}(s_{t,1:K}/\tau)\)。  

**提交：** 选中的子树的根词元 \(v_{k^*}\) 被追加到输出序列中，成为 \(y_t\)。  

**剪枝：** 所有其他 \(K-1\) 条深度为 1 的子树已不再是候选；它们被剪枝，相关的隐藏状态可以从计算图中分离。  

**平移：** 选中的子树被向上平移一个层级：其子节点成为新的根节点。如果被提升的节点有父节点（在初始深度 \(L=1\) 时，根节点是原始根），则丢弃父节点。经过平移后，树中幸存的部分深度为 \(L-1\)。  

**生长：** 由于选中的子树的子节点在平移后成为新的根节点，局部树此时缺少完整深度 \(L\) 的最外层。为了恢复完整深度，我们为每个候选路径生长一个新词元层。具体来说，对于 \(K\) 个候选位置中的每一个（注意平移后一些候选路径可能是不完全的），我们从语言模型 \(\pi_\theta\) 中自回归采样一个词元，将该词元前向传播通过模型以获取其隐藏状态，并将其作为子节点附加到路径的末端。生长是完全确定性的，除了依赖于 \(\pi_\theta\) 的采样。生长后，树恢复其完整形状：一个深度为 \(L\)、宽度为 \(K\) 的完整树，其中根节点对应提交词元 \(y_t\) 之后的模型状态。  

循环然后进入步骤 \(t+1\)。图 1 提供了 \(L=1\) 和 \(L=2\) 的图示。整个解码过程持续进行，直到满足停止条件（例如，遇到结束词元或达到最大长度）。  

### 似然分解  
为了在强化学习中使用 LBR，我们需要计算在路由器策略和基础语言模型策略下生成完整序列的轨迹似然。LBR 的关键属性是只有两种随机操作：  
1. **树生长：** 在每一步生长词元层时，新节点是从语言模型 \(\pi_\theta\) 中采样的。由于生长是完全确定性地在每个候选位置采样一个词元，整个生长步骤生成 \(K\) 个新节点（深度 \(L\)），每个节点对应一条特定的候选路径。在轨迹似然中，我们将这些生长步骤的贡献视为相关路径上下文下语言模型词元概率的乘积。  
2. **路由：** 在每个解码步骤 \(t\)，路由器 \(\rho_\phi\) 分配选择分布并采样一个选中的子树。  
所有其他操作——提交已选根词元、剪枝、平移和生长模板——都是确定性的，不对似然做出贡献。  
因此，在长度为 \(T\) 的解码序列上，轨迹似然为：  
\[
P_{\theta,\phi}(y_{1:T} \mid x_{<0}) = \prod_{t=1}^{T} \underbrace{\rho_\phi(k_t \mid \text{tree}_t)}_{\text{routing}} \prod_{\text{生长节点 } v \in \text{growth}_t} \pi_\theta(v \mid \text{context}(v)),
\]
其中 \(\text{tree}_t\) 是步骤 \(t\) 路由决策时的树状态，\(\text{growth}_t\) 是步骤 \(t\) 生长操作期间采样的节点集，\(\text{context}(v)\) 是节点 \(v\) 在树中的路径前缀（以 \(x_{<0}\) 开始，通过所有先前已提交词元和当前树中的祖先节点）。值得注意的是，树节点在首次采样时被计入似然，但随后在树结构中移动、复制或重用这些节点（通过剪枝、平移和生长）不会引入额外的概率因子。这种分解允许我们在主模型和路由器之间定义一个联合轨迹似然，从而使用强化学习进行端到端训练。

## 4 路由器架构  

路由器 \(\rho_\phi\) 是一个轻量级神经网络，它接收局部树的隐藏状态，并输出选择 \(K\) 条深度为 1 的子树之一的分布。我们为所有实验使用一个统一的路由器架构，如图 2 所示。该架构由三个模块组成：路径编码器、子树聚合器和候选比较器。  

**路径编码器。** 对于每个候选根词元 \(k\)，我们将与该候选关联的路径集合表示为 \(\{p_{k,m}\}_{m=1}^{M_k}\)，其中每条路径 \(p_{k,m}\) 是一个词元序列 \([u_{t,k}, u_{t+1,k}^{(m)}, \ldots, u_{t+L-1,k}^{(m)}]\)，代表候选根词元及其可能的局部延续。每条路径 \(p_{k,m}\) 被编码为一个单独向量 \(g_{t,k,m} = \text{PathEnc}_\phi([h_t(u_{t,k}), h_{t+1}(u_{t+1,k}^{(m)}), \ldots, h_{t+L-1}(u_{t+L-1,k}^{(m)})])\)，其中 \(h_t(\cdot)\) 表示当前时间步的语言模型隐藏状态，PathEnc 是一个小型 Transformer 编码器（单层，隐藏维度 \(p\)）将可变的路径长度映射到固定维度。  

**子树聚合器。** 对于候选根词元 \(k\)，其相关的路径集合同一子树的多个可能的局部延续。为了获得该候选子树的单向量表示，我们使用基于注意力的池化：通过一个简单的多头软注意块 \(\text{AttnPool}_\phi\)，将多条路径编码聚合成单个候选向量 \(g_{t,k} = \text{AttnPool}_\phi(\{g_{t,k,m}\}_{m=1}^{M_k})\)。该池化机制允许路由器关注与路由决策最相关的局部未来。在我们的主要设置 \(L=1\) 中，每个候选子树只有一条路径（单层所有前向传播的候选），因此子聚合器是恒等映射。  

**候选比较器。** 现在我们有 \(K\) 个候选向量 \(g_{t,1}, \ldots, g_{t,K}\)。为了在候选之间进行路由，我们通过另一个集注意力块 \(\text{SetAttn}_\phi^{\text{cand}}\) 来促进它们之间的信息交换。该块允许每个候选向量通过多头自注意机制在候选集上收集信息。之后，一个共享的线性层（参数 \(w_\phi\)）将更新后的候选向量映射到路由分数：  
\[
\tilde{g}_{t,1}, \ldots, \tilde{g}_{t,K} = \text{SetAttn}_\phi^{\text{cand}}(g_{t,1}, \ldots, g_{t,K}), \quad s_{t,k} = w_\phi^\top \tilde{g}_{t,k},
\]  
路由概率通过 softmax 获得：\(\rho_\phi(k \mid \text{tree}_t) = \text{softmax}(s_{t,1:K}/\tau)\)。  

**训练路由器。** 在 RL 训练期间，路由器与基础语言模型通过使用我们在第 3 节中描述的轨迹似然进行联合优化。路由器权重通过强化学习目标（例如，REINFORCE 或其变体，在 RLVR 设置中采用基于奖励的似然比）的梯度更新进行更新。在推理时，我们可以通过温度 \(\tau\) 控制采样的随机性。  

参见说明  
图 2：路由器架构。在每个解码步骤，路由器（i）通过路径编码器 \(\text{PathEnc}_\phi\) 独立编码每个候选子树的路径隐藏状态；（ii）通过基于注意力的池化 \(\text{AttnPool}_\phi\) 聚合每个子树的多条路径；（iii）通过集注意力块 \(\text{SetAttn}_\phi^{\text{cand}}\) 在候选之间交换信息；及（iv）使用共享线性头输出路由分数。路由器被训练为在给定局部树的情况下预测最有用的子树的离散选择。  

### 4.1 与软分支的比较  

一个相关的思想线是将软令牌与连续混合机制相结合。软思维[33](https://arxiv.org/html/2606.25354#bib.bib1)将候选词元的嵌入平均为单个连续表示，然后由语言模型在分支点进行处理。多路复用思维[25](https://arxiv.org/html/2606.25354#bib.bib12)通过添加随机候选采样使其可 RL 训练，并使用似然比目标优化语言模型。然而，这些方法在进一步推理开始之前将候选词元压缩为单一共享表示。作为回应，LBR 通过保持候选隐藏状态分离并在其中进行路由来保持候选特定的分支信息，而不是合并它们。图 3 通过比较 LBR 与基于软思维的管道，说明了这种差异。在我们的层次规划任务中，软合并丢失了分支特定信息（参见第 B 节），而 LBR 通过路由保留了离散分支身份。  

参见说明  
图 3：LBR 与软思维的概念比较。在层次规划任务中，LBR 将候选词元保持为离散分支，保留其身份，并通过比较决策状态进行路由。软思维将候选词元嵌入合并为单一连续表示，失去离散身份并使后续层难以区分分支。  

## 5 实验设置  

### 5.1 合成层次规划任务  

我们设计一个合成任务来调查 LBR 在存在结构概念选择（即，每当模型只需一次全局选择即可确定整个规划结构时）时的行为。该任务，*基数翻译规划*，由 Radix[29](https://arxiv.org/html/2606.25354#bib.bib29) 引入。在该任务中，模型生成一个平面图来调度操作，图结构由两个高级决策决定：节点布局的顺序（广度优先与深度优先）和数字基础（二进制与十进制）。一旦确定了这些选择，接下来的图结构就完全确定了。因此，一个离散模型可以通过最初正确分支来显著降低后代的困惑度。然而，在标准自回归解码中，这些选择必须在看到其局部结果之前基于有限的上下文做出。我们假设 LBR 可以通过检查候选局部未来状态并在它们之间进行路由来改进这些关键分支上的决策。  

我们将模型微调到 0 训练损失（在合成的 Radix 数据上），然后评估四个解码策略：贪婪离散 CoT、具有 32 个样本的自一致性（SC32）、软思维（使用均匀混合）和 LBR（\(K=3, L=1\)）。我们报告 Pass@1（通过贪婪解码或路由）和 Pass@32（通过 32 个自一致性样本或 32 个 LBR 路由样本）。我们使用一个小的联合微调步骤进行 LBR 路由，同时使用验证奖励：当模型选择导致图结构完全正确的分支时，给予 +1 的奖励，否则为 0。我们报告通过/失败准确率，其中“通过”意味着完整的图生成满足约束条件。  

### 5.2 数学推理任务  

我们使用三种主流方法或与 LBR 密切相关的基线评估 LBR：离散思维链 (CoT)、具有 32 个样本的离散自一致性 (SC32) 以及通过软思维 (Soft Thinking) 的多路复用思维风格的软分支。对于数学推理，我们使用 DeepSeek-R1-Distill-Qwen-1.5B[12](https://arxiv.org/html/2606.25354#bib.bib5) 作为基础模型。我们按照标准的 RLVR 程序对基础模型和路由器进行端到端联合训练：模型被指示在思维链中生成推理步骤，LBR 路由在生成的每个步骤中使用。我们报告 GSM8K[11](https://arxiv.org/html/2606.25354#bib.bib2)、MATH[14](https://arxiv.org/html/2606.25354#bib.bib10)、AMC 2023、AIME 2024 和 OmniMath[16](https://arxiv.org/html/2606.25354#bib.bib11) 上的 Pass@1 和 Pass@32 准确率。  

### 5.3 模型和训练细节  

基础模型是 DeepSeek-R1-Distill-Qwen-1.5B。路由器是 12.4M 参数的轻量级网络。对于数学推理，我们使用将 LBR 集成到 RLVR 管道中的联合训练。我们遵循标准 RLVR 训练，但显著的变化是 LBR 在每个解码步骤生成一个树结构，路由器选择一条路径。我们使用 GRPO[26](https://arxiv.org/html/2606.25354#bib.bib26)（组相对策略优化）和 REINFORCE[31](https://arxiv.org/html/2606.25354#bib.bib31) 的 RL 训练，并使用 KL 惩罚来保持接近参考模型。奖励由可验证格式和答案正确性决定。对于数学推理，我们将基础模型和路由器一起训练。我们使用 8 次 A100 GPU 运行。关于超参数的更多细节见附录 A。  

## 6 结果与分析  

### 6.1 合成层次规划任务  

表 1 显示了合成基数翻译规划任务的结果。LBR 显著优于所有基线，在 Pass@1 和 Pass@32 上都实现了 96.6% 的准确率。离散贪婪解码仅达到 65.2%，自一致性借助 32 个样本将其提高到 89.8%。软思维（均匀混合）优于贪婪解码（79.5%），但明显低于 LBR（96.6%）。这些结果表明，路由中使用的候选后隐藏状态对于解决需要高难度结构决策的任务中的关键分支特别有价值。  

**表 1：合成层次规划（基数翻译）上的解码准确率（%）。**  

| 方法 | Pass@1 | Pass@32 |  
| :--- | :--- | :--- |  
| 贪婪离散 CoT | 65.2 | – |  
| 自一致性（32 样本） | – | 89.8 |  
| 软思维（均匀混合） | 79.5 | – |  
| LBR \(K=3, L=1\) | **96.6** | **96.6** |  

参见说明  
图 4：合成层次规划任务中的路由决策质量。在第一个概念选择（节点顺序和基数）上，LBR 正确路由的准确率远高于随机猜测（50% vs. 25% vs. 11.8%）。  

### 6.2 数学推理基准  

表 2 显示了多个数学基准的结果。LBR 一致地改进了离散 CoT 和软思维基线。在 Pass@1 上，LBR 在 GSM8K 上实现了 +2.2 的改进（89.1% → 91.3%），在 MATH 上实现了 +4.1 的改进（69.0% → 73.1%）。LBR 在更具挑战的竞赛数学基准（AMC、AIME、OmniMath）上也优于所有基线。  

在 Pass@32 上，LBR 同样优于自一致性基线。例如，在 MATH 上，LBR 实现了 82.7% 的 Pass@32，而离散 SC32 为 79.2%，软思维为 80.5%。这些改进表明，在词元级的分支和路由与在推理时进行采样计算是互补的，并且当路由策略经过训练以选择有希望的局部路径时，LBR 实现了更高效的采样分布。  

**表 2：数学推理基准上的准确率（%）。粗体表示每个指标的最佳结果；下划线表示第二佳。**  

| 方法 | GSM8K PM1 | GSM8K PM32 | MATH PM1 | MATH PM32 | AMC 2023 PM1 | AMC 2023 PM32 | AIME 2024 PM1 | AIME 2024 PM32 | OmniMath PM1 | OmniMath PM32 |  
| :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- |  
| 离散 CoT (贪婪) | 89.1 | – | 69.0 | – | 51.2 | – | 26.5 | – | 33.0 | – |  
| 离散 CoT (SC32) | – | 91.5 | – | 79.2 | – | 64.5 | – | 42.8 | – | 53.8 |  
| 软思维 + RLVR | 88.0 | 90.5 | 68.1 | 80.5 | 49.5 | 64.0 | 24.2 | 42.4 | 32.1 | 52.9 |  
| LBR + RLVR | **91.3** | **92.5** | **73.1** | **82.7** | **54.6** | **67.4** | **28.7** | **44.5** | **36.2** | **56.1** |  

### 6.3 消融研究  

我们进行消融研究来分析 LBR 的各个组成部分。  

**树宽度 \(K\) 的影响。** 图 5（左）显示了在 \(K=2,3,5,10\) 时 MATH 上的性能。性能随 \(K\) 增加，在 \(K=5\) 附近达到饱和。较大宽度提供更多多样性，但也增加了计算开销；我们发现在大多数实验中 \(K=3\) 是一个很好的平衡。  

**树深度 \(L\) 的影响。** 图 5（右）将 \(L=1\)（仅一个前瞻词元）与 \(L=2\)（两个前瞻词元）进行比较。有趣的是，对于 \(L=2\)，性能几乎相同。这表明在许多情况下，一个词元的局部前瞻就足够了；路由后可能存在足够的信息。更深的树可能带来有限的额外收益，同时增加计算成本。  

参见说明  
图 5：宽度 \(K\) 和深度 \(L\) 的消融研究。左：在 MATH 上，性能随 \(K\) 增加，在 \(K=5\) 附近饱和。右：在合成和 MATH 上，\(L=1\) 和 \(L=2\) 的效果相似，表明单层前瞻通常就足够。  

**路由器的贡献。** 表 3 显示了在合成任务上不同路由策略的性能。正确路由（基于未来信息选择）与随机路由（在所有候选之间均匀采样）进行比较。使用候选后隐藏状态的路由决策显著优于随机基线。  

**表 3：合成任务上的路由消融研究。**  

| 路由策略 | Pass@1 |  
| :--- | :--- |  
| LBR（正确路由） | **96.6** |  
| 随机路由（均匀） | 82.1 |  

### 6.4 与 RL 训练的效率比较  

表 4 比较了训练效率。LBR 在合成和数学任务上的训练时间与标准的 RLVR 管道大致相当，同时提供了更好的性能。路由器仅增加 12.4M 参数，并且使用轻量级的集注意力操作，这使得 LBR 在计算上对大模型来说是可行的。  

**表 4：训练效率比较。**  

| 方法 | 端到端训练时间 (A100 GPU 小时) | 每步额外计算量 | 参数增长率 |  
| :--- | :--- | :--- | :--- |  
| 离散 CoT + RLVR | 128 | 0% | 0% |  
| 软思维 + RLVR | 140 | +9% | +0.8% |  
| LBR + RLVR | 145 | +13% | +0.8% |  

## 7 相关工作  

**测试时扩展。** 测试时扩展在大型语言模型（如 OpenAI o1、DeepSeek-R1[1](https://arxiv.org/html/2606.25354#bib.bib5)、Claude 3.5 Sonnet、Gemini 2.0 Flash Thinking 等）中引起了广泛关注。早期的研究工作，如思维链[28](https://arxiv.org/html/2606.25354#bib.bib15)和自一致性[27](https://arxiv.org/html/2606.25354#bib.bib17)，探索了在推理时增加计算的方法。布朗等人[3](https://arxiv.org/html/2606.25354#bib.bib3) 首次通过逻辑推理展示了测试时扩展的潜力。威利等人[30](https://arxiv.org/html/2606.25354#bib.bib30) 和斯内尔等人[24](https://arxiv.org/html/2606.25354#bib.bib16) 对测试时计算分配进行了系统研究。LBR 在这些方向的基础上，提供了一种结构化、可训练的 token 级分支和路由新机制。  

**树搜索解码。** 基于搜索的解码方法，如束搜索和蒙特卡洛树搜索[29](https://arxiv.org/html/2606.25354#bib.bib29)；[32](https://arxiv.org/html/2606.25354#bib.bib29), [10](https://arxiv.org/html/2606.25354#bib.bib19), [7](https://arxiv.org/html/2606.25354#bib.bib8), [8](https://arxiv.org/html/2606.25354#bib.bib9), [18](https://arxiv.org/html/2606.25354#bib.bib24), [5](https://arxiv.org/html/2606.25354#bib.bib4), [4](https://arxiv.org/html/2606.25354#bib.bib7)，探索了多条路径以改进生成。思维树[32](https://arxiv.org/html/2606.25354#bib.bib29) 在句子级别进行搜索，而 LBR 在 token 级别运行，允许在每个解码步骤进行更细粒度的决策。  

**可训练解码策略。** 最近的努力使解码策略通过强化学习变得可训练[21](https://arxiv.org/html/2606.25354#bib.bib23), [20](https://arxiv.org/html/2606.25354#bib.bib22), [9](https://arxiv.org/html/2606.25354#bib.bib25), [6](https://arxiv.org/html/2606.25354#bib.bib6), [19](https://arxiv.org/html/2606.25354#bib.bib26)。特别是 RLVR[23](https://arxiv.org/html/2606.25354#bib.bib38), [9](https://arxiv.org/html/2606.25354#bib.bib25) 使用可验证奖励训练语言模型进行推理。LBR 通过引入一个与基础语言模型一起训练的显式路由器，将 RLVR 扩展到树搜索空间。这允许在 token 级分支决策上进行端到端训练。  

**软分支与连续解码。** 软思维[33](https://arxiv.org/html/2606.25354#bib.bib1)和多路复用思维[25](https://arxiv.org/html/2606.25354#bib.bib12) 通过连续混合探索了软分支的概念。这些方法在进一步处理之前将候选 token 合并为一个连续表示。其框架在推理质量和训练稳定性方面存在缺陷。LBR 通过保持候选 token 离散并在它们之间进行路由提供了替代方案，在本文评估的任务中显示出更强的性能。  

## 8 结论与展望  

我们提出了局部分支路由（LBR），一种用于语言模型的高效可训练 token 级测试时扩展框架。LBR 扩展了一个小的局部前瞻树，使用路由器选择有希望的路径，并定义了一个可处理的轨迹似然，允许通过 RL 进行端到端训练。在合成层次规划任务上，LBR 准确率显著提升（96.6% vs. 65.2% 贪婪 CoT）。在数学推理基准上，LBR 一致地改进了离散 CoT 和软分支基线。这些结果为语言模型推理中可训练 token 级决策的有效性提供了有力证据。  

**限制和未来工作。** LBR 增加了每步计算量，但额外的开销是适度的（+13%）。我们计划探索动态宽度分配和分层路由。将 LBR 集成到现成的基础模型（例如使用 RLVR）中需要进行联合训练，这可能需要高端资源。未来的工作包括扩展到更大模型、改进路由器架构以及研究集成树搜索的变体。  

## 参考文献  
...  
（注：因篇幅限制，参考文献列表在此省略；只翻译了部分引用内容）  

## 附录 A 实验细节  

### A.1 合成层次规划任务  

*基数翻译规划*任务[29](https://arxiv.org/html/2606.25354#bib.bib29) 要求模型生成一个由四个操作组成的平面图，节点布局由两个高级概念决定：节点顺序（广度优先 BFS 与深度优先 DFS）和数字基础（二进制与十进制）。有四个可能的图结构，每个对应一种选择组合。我们使用 BFS+二进制、BFS+十进制、DFS+二进制、DFS+十进制的等比例混合，生成 10k 个训练样本和 1k 个测试样本。任务定义为自回归序列生成，序列表示图的谱。例如，对于二进制基础，操作数为二进制整数；对于十进制为十进制整数。  

**模型。** 我们使用 GPT-2 基础架构（约 85M