跳转至

Dynamic Latent Routing

会议: NeurIPS 2026
arXiv: 2605.14323
领域: LLM 其他
关键词: 动态潜变量路由, 离散表征, 通用狄克斯特拉搜索, 残差流操纵, 强化学习策略拼接

一句话总结

针对语言模型后训练中离散潜变量破坏预训练序列结构且依赖多阶段预热的瓶颈,提出通用狄克斯特拉搜索(GDS)理论并导出动态潜变量路由(DLR),在单阶段训练中通过策略头搜索在残差流施加块级连续转向向量,在低数据量微调下全面超越有监督微调。

研究背景与动机

当前神经语言模型虽然输入是离散文本,但其内部计算几乎完全依赖连续的隐层表征展开。若能为模型引入离散潜编码,就如同在时序维度将高阶宏观子策略显式拼接,使模型能够在更抽象的动作空间中自主规划计算链路。然而,分层强化学习与现有离散编码研究长期面临理论与实用双重鸿沟:一方面,学术界一直缺乏关于「在时变环境中沿时序维度拼接子策略何以保证全局最优」的严格理论刻画;另一方面,现有的离散潜变量注入方法均存在显著缺陷。

现有的离散编码注入方法(如 Pause Tokens、Meta-Tokens、TokenAssorted 等)普遍将编码作为额外的离散 Token 插入到自然语言文本序列中。这种方式严重破坏了模型在海量文本预训练中建立的自回归序列结构与上下文相对位置模式,导致模型必须依赖昂贵的大规模预热微调阶段来适应这种人为扰动,在样本受限的低数据微调场景中性能甚至严重劣化,远落后于标准有监督微调(SFT)。此外,现有方法割裂了编码字典构建、路由策略学习与底座模型更新,需要复杂的两阶段乃至多阶段流水线。

为打破上述理论空白与架构限制,本文从时变 MDP 出发证明了通过时序拼接最优子策略能够逼近全局最优目标达成策略,并将优先队列搜索过程松弛为由可学习策略头引导的残差流干预机制。核心 idea:将离散潜变量作为块级连续转向向量直接注入残差流以保护预训练序列完整性,并通过单阶段「搜索-选择-更新」闭环实现路由策略、编码字典与模型参数的端到端联合优化。

方法详解

整体框架

动态潜变量路由(DLR)的宏观流程分为「候选序列动态搜索」、「最优潜编码选择」以及「多目标联合更新」三个核心步骤。输入自然语言序列首先按粗粒度时序划分为长度固定的块(chunk),由挂载在指定隐层的线性路由头在前向传播时根据首个 Token 表征以设定采样温度抽样候选离散编码序列;随后,底座语言模型并行评估候选序列在残差流施加转向向量后的条件似然并挑选最优序列;最后,计算包含通用语言模型损失、互信息增益、策略蒸馏及双词元齐夫正则化的联合损失函数,实现策略网络、离散字典与底座参数的同步梯度更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列与粗粒度分块<br/>划分长度为 K 的连续 Token 块"] --> B["块级残差流转向向量注入<br/>首 Token 表征过线性路由头生成候选分布"]
    B --> C["候选路由序列并行采样搜索<br/>以温度 τ 抽取 N 条候选编码路径"]
    C --> D["条件对数似然评估与最优选择<br/>挑选最大化 p(x|a) 的黄金编码序列 a*"]
    D --> E["单阶段联合损失反向传播<br/>通用能力+信息增益+策略优化+Zipf正则"]

关键设计

1. 理论根基与通用狄克斯特拉搜索(GDS):证明时序子策略拼接的全局最优性

针对传统强化学习贝尔曼迭代依赖单步状态精细更新而缺乏时序宏观子策略有效性理论证明的痛点,本文形式化了时变奖励函数下的动态马尔可夫决策过程(DMDP)。在时变环境下静态策略次优性成立的前提下,证明了策略时序拼接后的价值函数分解定理。提出了通用狄克斯特拉搜索(GDS)算法:利用优先队列管理达到目标状态集合的候选子策略拼接链,以误差容限自适应剪枝。理论证明在有界时变奖励假设下,GDS 保证首次从队列弹出的可行策略即为全局最优目标达成策略,奠定了通过离散宏观编码搜索组装复杂长程行为的数学基石。

2. 块级残差流转向向量注入:无损保持预训练自回归文本结构

针对离散 Token 插入破坏预训练位置关系与注意力模式的痛点,DLR 放弃向输入文本注入额外 Token,转而在隐层残差流上执行连续向量偏置。对于长度为 \(T\) 的自然语言序列 \(x_{1:T}\),引入抽象比率 \(K\),将其划分为 \(M = \lceil T/K \rceil\) 个连续块,每个 Token \(t\) 所属块索引为 \(m(t) = \lfloor (t-1)/K \rfloor\)。在线性路由头 \(W_{\text{rt}} \in \mathbb{R}^{C \times H}\)(\(C\) 为字典大小)的作用下,仅提取每个块首个 Token 处的注入层隐藏状态 \(h_{mK+1}^{(l^*)}\),生成路由 logits:

\[z_m = W_{\text{rt}} \, \text{sg}\big(h_{mK+1}^{(l^*)}\big) \in \mathbb{R}^C\]

通过阻断梯度(\(\text{sg}\))保证路由决策与隐层主干梯度解耦。随后,该块内所有 Token 共享对应的离散编码 \(a_{m(t)}\),并将字典中可学习的转向向量 \(e_{a_{m(t)}} \in \mathbb{R}^H\) 乘以固定缩放因子 \(\alpha\) 原位加到残差流上:

\[h_t^{(l^*)} \leftarrow h_t^{(l^*)} + \alpha \, e_{a_{m(t)}}, \quad t \in \{1,\dots,T\}\]

由于输入文本序列与自注意力拓扑完全保持原样,模型在初始化阶段便拥有标准的零样本先验能力,消除了低数据训练中巨大的语法重适应成本。

3. 单阶段动态搜索、选择与联合优化(Search-Select-Update):闭环神经松弛

为解决传统离散潜变量需要两阶段训练、冷启动预热或多模型级联的问题,DLR 构建了 GDS 的端到端神经松弛机制。在每个训练步,首先通过路由头概率分布 \(p_\theta(a|x)\) 以采样温度 \(\tau\) 并行抽取 \(N\) 条候选编码序列 \(\{a^{(1)}, \dots, a^{(N)}\}\)(搜索阶段);随后在残差流中应用转向向量,挑选出使得序列条件生成似然最大的候选编码序列 \(a^* = \arg\max_{a^{(i)}} p_\theta(x \mid a^{(i)})\)(选择阶段);最后利用构建的四项联合损失 \(\mathcal{L}_{\text{DLR}}\) 进行反向传播(更新阶段):

\[\mathcal{L}_{\text{DLR}}(x, a^*) = -\log p_\theta(x) - \log \frac{p_\theta(x \mid a^*)}{\text{sg}(p_\theta(x))} - \alpha_{\text{policy}} \log p_\theta(a^* \mid x) + \alpha_{\text{reg}} D_{\text{KL}}\big(p_\theta(a) \parallel p_{\text{bi-zipf}}(a)\big)\]

损失函数巧妙融合了四个物理含义清晰的目标:\(-\log p_\theta(x)\) 充当通用基线,约束无转向下的底座表征能力;信息增益项 \(-\log [p_\theta(x \mid a^*) / \text{sg}(p_\theta(x))]\) 作为内在奖励,强力驱动被选中的潜变量为文本生成提供实质性帮助;\(-\alpha_{\text{policy}} \log p_\theta(a^* \mid x)\) 则是策略蒸馏项,将前向搜索筛选出的胜出编码反向监督给路由头;双词元齐夫先验正则项 \(\alpha_{\text{reg}} D_{\text{KL}}(p_\theta(a) \parallel p_{\text{bi-zipf}}(a))\) 防止所有块路由坍缩到单一优势编码,维持字典使用率与时序转移多样性。

一个完整示例

以多步进位算术推理 959,271 + 040,756 = 1,000,027 为例:在生成过程中,每个答案字符位置映射为一个抽象编码(\(K=1\))。当计算遭遇深层进位级联时,模型路由头直接为连续四级进位位置激活特定离散编码 \(t_2\)(对应级联进位子任务),而在跨越 9 的进位边界激活编码 \(t_6\)(对应和为 9 边界检测),并在无进位的简单求和位指派独立基准编码。无需额外探测分类器,仅通过监测残差流所选择的离散代号序列,人类即可直观读取神经网络内部执行的进位运算状态;若人为将进位错误样本的关键位编码置换为正确的 \(t_2\),模型能以近 31% 的修复率直接修正最终结果。

实验关键数据

主实验

在单 epoch 低数据受限场景下,评估涵盖 Qwen3(0.6B、1.7B、4B、8B)及 Llama3.2(1B、3B)六种基座模型在四大问答与推理数据集(GSM8K、ScienceQA、StrategyQA、CSQA)上的表现,主结果对比(以 4 种典型规模模型为例)如下:

模型架构 微调方法 GSM8K (Acc %) ScienceQA (Acc %) StrategyQA (Acc %) CSQA (Acc %) 全局平均增益
Qwen3-0.6B SFT 46.0 \(\pm\) 2.7 48.0 \(\pm\) 2.1 47.0 \(\pm\) 3.7 64.1 \(\pm\) 2.7 基线
Qwen3-0.6B PauseToken 46.2 \(\pm\) 2.7 47.9 \(\pm\) 2.1 24.6 \(\pm\) 3.2 64.7 \(\pm\) 2.7 -4.4 pp
Qwen3-0.6B TokenAssorted 15.7 \(\pm\) 2.0 13.1 \(\pm\) 1.4 29.4 \(\pm\) 3.4 64.2 \(\pm\) 2.7 -20.0 pp
Qwen3-0.6B DLR (\(C=1\)) 45.8 \(\pm\) 2.7 48.9 \(\pm\) 2.1 50.4 \(\pm\) 3.7 64.0 \(\pm\) 2.7 +1.7 pp
Qwen3-0.6B DLR (本文) 49.4 \(\pm\) 2.7 55.3 \(\pm\) 2.1 54.6 \(\pm\) 3.7 66.4 \(\pm\) 2.6 +5.8 pp
Qwen3-1.7B SFT 60.2 \(\pm\) 2.6 56.4 \(\pm\) 2.1 51.7 \(\pm\) 3.7 74.3 \(\pm\) 2.5 基线
Qwen3-1.7B PauseToken 60.9 \(\pm\) 2.6 56.2 \(\pm\) 2.1 32.0 \(\pm\) 3.5 76.4 \(\pm\) 2.4 -4.3 pp
Qwen3-1.7B TokenAssorted 28.5 \(\pm\) 2.4 17.3 \(\pm\) 1.6 33.0 \(\pm\) 3.5 76.3 \(\pm\) 2.4 -21.8 pp
Qwen3-1.7B DLR (\(C=1\)) 63.5 \(\pm\) 2.6 60.0 \(\pm\) 2.0 55.3 \(\pm\) 3.7 77.4 \(\pm\) 2.3 +3.4 pp
Qwen3-1.7B DLR (本文) 65.7 \(\pm\) 2.6 64.1 \(\pm\) 2.0 60.3 \(\pm\) 3.7 78.4 \(\pm\) 2.3 +6.5 pp
Qwen3-8B SFT 78.9 \(\pm\) 2.1 61.8 \(\pm\) 2.0 68.2 \(\pm\) 3.5 85.6 \(\pm\) 2.1 基线
Qwen3-8B PauseToken 79.4 \(\pm\) 2.1 65.2 \(\pm\) 2.0 38.8 \(\pm\) 3.6 82.9 \(\pm\) 2.2 -7.1 pp
Qwen3-8B TokenAssorted 63.0 \(\pm\) 2.5 69.7 \(\pm\) 1.9 41.2 \(\pm\) 3.7 83.1 \(\pm\) 2.2 -9.4 pp
Qwen3-8B DLR (\(C=1\)) 78.9 \(\pm\) 2.2 66.4 \(\pm\) 2.0 70.3 \(\pm\) 3.5 83.8 \(\pm\) 2.1 +1.2 pp
Qwen3-8B DLR (本文) 84.3 \(\pm\) 2.0 80.6 \(\pm\) 1.7 72.9 \(\pm\) 3.4 87.2 \(\pm\) 1.9 +7.6 pp
Llama3.2-1B SFT 30.9 \(\pm\) 2.5 38.0 \(\pm\) 2.0 48.0 \(\pm\) 3.7 65.2 \(\pm\) 2.7 基线
Llama3.2-1B PauseToken 31.3 \(\pm\) 2.5 34.0 \(\pm\) 2.0 25.5 \(\pm\) 3.3 64.9 \(\pm\) 2.7 -6.6 pp
Llama3.2-1B TokenAssorted 14.2 \(\pm\) 1.9 33.6 \(\pm\) 2.0 29.3 \(\pm\) 3.4 65.6 \(\pm\) 2.7 -9.9 pp
Llama3.2-1B DLR (\(C=1\)) 38.4 \(\pm\) 2.6 43.7 \(\pm\) 2.1 48.6 \(\pm\) 3.7 68.3 \(\pm\) 2.6 +4.2 pp
Llama3.2-1B DLR (本文) 41.1 \(\pm\) 2.7 49.1 \(\pm\) 2.1 51.4 \(\pm\) 3.7 71.4 \(\pm\) 2.5 +7.7 pp

消融实验

针对超参数搜索预算、采样多样性、块压缩比及因果干预的系统性消融实验如下:

消融配置/变量 ScienceQA Acc (%) GSM8K Acc (%) 性能变化幅度与行为特征
DLR 默认基线 (\(N=4, \tau=1, K=4, C=32\)) 55.3 \(\pm\) 2.1 49.4 \(\pm\) 2.7 维持高利用率与均衡多样性
无搜索阶段 (\(N=1\)) 45.5 \(\pm\) 2.1 38.7 \(\pm\) 2.6 全局大幅下降 -9.8 pp 至 -10.7 pp,证实前向搜索必要性
扩展搜索预算 (\(N=8\)) 55.1 \(\pm\) 2.1 49.2 \(\pm\) 2.7 差异在 \(\pm 0.5\) pp 以内,说明 \(N=4\) 已饱和搜索边际收益
贪心确定性采样 (\(\tau=0\)) 39.8 \(\pm\) 2.0 42.5 \(\pm\) 2.7 坍缩至 argmax 死锁点,平均掉点 -6.9 pp
过度随机探索 (\(\tau=2\)) 51.2 \(\pm\) 2.1 39.4 \(\pm\) 2.6 偏离策略网络引导接近均匀噪声,GSM8K 掉点达 -10.0 pp
单 token 逐位操纵 (\(K=1\)) 42.3 \(\pm\) 2.1 45.7 \(\pm\) 2.7 失去时间维度宏观抽象能力,SciQA 掉点达 -13.0 pp
静态单转向向量 (\(C=1\)) 48.9 \(\pm\) 2.1 45.8 \(\pm\) 2.7 无法按样本动态指派策略,整体平均落后完整版 -3.6 pp
全局剔除转向向量 (\(\alpha \to 0\)) 49.1 \(\pm\) 2.1 43.2 \(\pm\) 2.7 损失潜编码引导,Qwen 系列平均跌落 6.2 至 17.4 pp
随机打乱编码指派 (Random Replace) 50.5 \(\pm\) 2.1 44.6 \(\pm\) 2.7 破坏因果匹配,Qwen 系列平均跌落 4.8 至 11.7 pp
去除双词元齐夫正则 (\(\alpha_{\text{reg}}=0\)) 48.2 \(\pm\) 2.1 45.7 \(\pm\) 2.7 离散编码字典完全坍缩至少数向量,利用率急剧衰减

关键发现

  • 复杂推理任务增益显著:DLR 在需要多步推理的 ScienceQA 与 GSM8K 上提升尤为剧烈(在 Qwen3-8B 上 SciQA 暴涨 +18.8 pp,在 Llama-1B 上 GSM8K 提升 +10.2 pp),而在常识问答 CSQA 上提升相对平缓,印证了离散宏观路由在长程结构化分解中的不可替代性。
  • 离散编码自发形成因果专门化:在没有人类子任务标签监督的情况下,DLR 编码序列 \(n\)-gram 随长度增加对科学领域的纯度(Topic Purity)显著单调递增;对单个编码实施消融时呈现鲜明的主题正负翻转(如消融 code 0 使生物题准确率上升 +3.6 pp 但使物理题掉点 -4.0 pp),说明编码已演化为具备因果负载的领域调节旋钮。
  • 低数据微调下的抗崩塌特性:在仅有单个训练 epoch 的极低数据约束下,传统 Token 插入方法(如 TokenAssorted)在小模型上完全崩塌(在 Qwen3-0.6B 上 GSM8K 跌至 15.7%),而 DLR 凭借残差流原位干预机制展现出坚固的稳定性,全面刷新 24 个模型-数据集评估单元的领先纪录。

亮点与洞察

  • 离散概念与连续流的有机统一:传统离散潜变量通常试图直接修改离散 Token 符号流,本文创新性地将其降维映射为在连续残差流中原位叠加的转向偏置,兼顾了离散路由的可解释性与连续表征的无损兼容性。
  • GDS 的神经平滑松弛极其巧妙:摆脱了强化学习传统的贝尔曼逐状态值迭代,采用优先队列的“搜索-选择-更新”架构,在神经网络中以轻量策略头和生成似然评估代替显式排队,单阶段打通了离散探索与参数学习。
  • 白盒级机制干预与错误外科修复:六位算术分析表明编码直接对应进位级联与边界判断电路,通过直接替换特定位置的离散代号即可外科手术式地纠正模型的深层推理错误,无需复杂的后验激活探针。

局限与展望

  • 注入层选择需要模型特定搜索:最优转向层位 \(l^*\) 依赖于基座模型参数量(小模型如 0.6B/1.7B 偏好浅层 \(l^*=1\),中大型如 4B/8B 偏好深层 \(l^*=22/24\)),跨架构泛化时仍需保留离散层位超参探索。
  • 前向采样带来的额外算力开销:在训练阶段每个步骤需要抽取 \(N=4\) 条候选编码序列并并行前向计算以获取似然排序,虽然推理时仅需单次前向,但微调计算量较朴素 SFT 略有上升。
  • 长文本自适应块长的拓展空间:当前抽象比率 \(K\) 采用固定窗口划分,未来可探索内容驱动的动态分块边界,根据语义密度自适应调整宏观策略的有效步长。

相关工作与启发

  • vs Pause Tokens / Meta-Tokens: 后两者均以额外 Token 形式插入序列,需要大量预热微调才能缓解对预训练注意力的冲击;DLR 采用残差流转向向量注入,低数据微调零破坏、零冷启动。
  • vs TokenAssorted: TokenAssorted 依赖外部 VQ-VAE 编码器进行两阶段离散标签预打标;DLR 在单阶段内利用生成似然实现端到端闭环联合优化,编码完全由任务自身信息增益驱动。
  • vs 静态表征工程 (Representation Engineering): 传统转向向量通常静态固定或全局施加;DLR 学习了动态、块级依赖且输入自适应的离散路由策略,具备更细粒度的时序上下文承载能力。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从时变 MDP 理论导出 GDS 并将其松弛为单阶段端到端潜变量残差路由机制。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 种主流基座、4 大权威基准、全量消融与细粒度因果截断分析,论证极其扎实。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论证明完整,从宏观动机到机制公式推导清晰紧凑,图表表现力强。
  • 价值: ⭐⭐⭐⭐⭐ 为轻量级大模型后训练、内在可解释性与因果控制提供了高度可复现且效果卓越的新范式。