思维链追踪无法揭示全部:MoE模型中推理内部状态的双层读取

29 八月 20269 视图

研究人员将推理内部状态压缩为紧凑的64维语义框架J64,可通过专家路由统计进行恢复。这有助于更精确地评估解题过程、改进答案选择并控制生成停止:AUC提升达0.135,准确率提升最高达5.9个百分点。

思维链追踪无法揭示全部:MoE模型中推理内部状态的双层读取

基于轨迹(traces)进行解释存在风险

在混合专家(MoE)架构的模型中,每个 token 仅经过少数“专家”处理。处理后会留下一个形式化的痕迹——记录了哪些专家被调用以及对应的权重。这类轨迹虽便于用作解释,但它更像一份没有通话内容的电话记录:能看到谁给谁打了电话,却看不到他们商定了什么。

最近的一篇预印本(Chen 等人,arXiv:2608.17638)表明,基于轨迹的解释遗漏了一个重要的信息层。作者提出不仅要读取可见的路由路径,还要读取模型内部推理状态。核心思路是双层读取:先构建一个紧凑的语义空间,再通过常规路由统计信息恢复一个廉价代理。

双层状态读取器

J64:用 64 个轴替代完整词表

第一层基于空间 J 构建,其规模与模型词表相当。该空间被蒸馏为一个紧凑的语义框架 J64——仅含 64 个轴。这并非任意的主题投影,而是一组基于模型自身推理状态训练得到的不变量。

J64 的主要价值在于,它能捕捉到可见轨迹中不存在的状态。两条相同的路由痕迹可能隐藏着不同的内部过程,J64 有助于区分它们。此外,它还能分离两种不同的负荷:模型在推理中投入的努力程度,以及该样本对模型而言本身的难度。

在保留数据上的验证显示,与基线方法(将同一 rollout 视为 token 填充度并使用相同聚合方式)相比,ROC 曲线下面积的提升为 0.096–0.135。也就是说,改进完全来自状态读取方式本身,而非模型或数据格式的变化。

R64:基于廉价统计的同等视角

J64 有一个实际缺点:它需要访问模型的内部表示。为使该方法适用于生产环境,作者从专家路由的原生统计信息中重建了 J64。该代理称为 R64,在推理期间不会产生显著额外开销。

在三个模型和两个架构家族上,R64 与 J64 之间的中位逐轴相关性达到 0.69–0.86。在 gpt-oss-20b 模型上,R64 保留了原始 J64 预测增益的 95–100%。换言之,对许多任务而言,无需计算完整的语义空间,仅从现有路由机制中提取统计信息即可。

测试时决策的两种时间尺度

状态读取不仅可作为事后解释,还可作为决策工具。作者考虑了两种场景:分析已完成候选以及实时控制生成过程。

从现成候选中选择

当模型已生成多个答案候选时,J64 和 R64 可改善对其中一条分支的选择。加权投票也经过了单独验证:若用 R64 的权重替代简单多数投票,在八种设置中有七种结果优于普通多数投票。这意味着隐藏状态可作为更精确的答案聚合信号。

生成过程中的停止与重采样

第二种场景是生成过程中需要随 token 出现实时决策。滑动读取窗口被输入到累积式“停止并重采样”策略中:在特定时刻,模型停止当前分支并重新开始生成。该策略的工作点仅基于训练问题确定,因此在测试时不会针对已知答案进行调整。

J64 相比对照组(使用打乱的同源样本)在准确率上提升了 1.1–5.9 个百分点。仅基于路由信息的 R64 保留了其中 0.9–3.2 个百分点的增益。即使是廉价代理,对生成控制也足够有效。

能否通过路由器改变推理

工作的最后一部分超越了被动解释。研究者对路由器进行编辑,以影响与 J64 对应的机制。此后,模型行为以可预测的方式发生变化——这种变化符合 J64 解释所预示的方向。模型可诊断的“卡顿”从数值猜测转向精确的符号执行。

这是一个重要信号:J64 的轴不仅捕捉相关性,还体现了推理中的因果作用。如果路由器编辑产生预期变化,则说明被读取的状态确实参与决策,而非仅仅是附带产物。

MoE 轨迹仍然是一个便捷但过于粗糙的模型接口。双层读取——先进行经济的语义投影,再基于路由构建代理——使我们能够看到可见专家调用之前及之外发生的过程。这类内部“仪表盘”有望成为调试和管理大型模型的下一项标准工具。

常问问题

思维链追踪无法揭示全部:MoE模型中推理内部状态的双层读取