GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

1. 研究背景与问题 (Problem)

背景：
推理时扩展（Inference-Time Scaling, ITS）是提升大语言模型（LLM）推理性能的关键机制。在传统的离散 Token 模型中，ITS 通常通过采样多个推理轨迹（如 Self-Consistency 或 Best-of-N）来实现，这依赖于明确的 Token 级条件概率分布。然而，潜在推理模型（Latent Reasoning Models）（如 COCONUT, CODI）直接在隐藏状态空间进行多步推理，不生成中间文本 Token。

核心问题：
现有的潜在推理 ITS 方法主要依赖启发式扰动（Heuristic Perturbations），例如 Dropout 或固定高斯噪声，来生成多样化的候选轨迹。作者指出这些方法存在以下严重局限性：

1. 缺乏显式分布定义： 这些扰动没有定义明确的潜在思维条件采样分布，导致探索过程难以控制或优化。
2. 扰动强度与质量不匹配： 更强的扰动（更大的噪声）并不一定带来更好的采样质量。实验表明，过强的扰动往往导致分布偏移（Distribution Shift）过大，却未能生成更有用的推理路径，甚至降低最终决策的准确性。
3. 盲目性： 现有的启发式方法（如固定 Dropout 率）是手工设定的，对骨干网络、任务和潜在几何结构高度敏感，容易陷入“探索不足”（Under-exploration）或“过度探索”（Over-exploration）的极端。

核心疑问： 在连续推理空间中，什么样的探索才是有效的？仅仅是增加随机性，还是需要受控的、可优化的采样策略？

---

2. 方法论：高斯思维采样器 (GTS)

为了解决上述问题，作者提出了 高斯思维采样器（Gaussian Thought Sampler, GTS）。

2.1 核心思想

GTS 将潜在空间的探索重新定义为从学习到的条件分布中对连续推理状态进行采样。它不再注入固定的启发式噪声，而是预测一个**上下文相关（Context-dependent）**的高斯扰动分布。

2.2 模型架构

• 冻结骨干网络： 保持原有的潜在推理骨干网络（Backbone）参数冻结。
• 可学习采样器： 添加一个轻量级的模块（由 MLP 组成），用于预测每个推理步骤 $k$ 的扰动分布参数。
• 条件高斯策略：
  • 输入：确定性隐藏状态 $h_k^{det}$ 作为上下文 $c_k$。
  • 输出：预测均值 $\mu_\phi(c_k)$ 和对角线协方差的对数标准差 $\log \sigma_\phi(c_k)$。
  • 采样：$z_k \sim \mathcal{N}(\mu_\phi(c_k), \text{diag}(\sigma_\phi^2(c_k)))$。
  • 扰动状态：$\tilde{h}_k = h_k^{det} + z_k$，该扰动状态被反馈回骨干网络进行后续推理。

2.3 训练策略：GRPO 风格优化

由于推理时的最终决策（正确/错误）相对于采样器参数是不可微的，GTS 采用**强化学习（RL）**进行优化，具体使用 GRPO (Group Relative Policy Optimization) 风格的策略：

• 奖励设计 (Reward Design)：
  • 正确性信号： 如果采样轨迹解码出的答案与真值一致，给予正奖励；否则给予负奖励。
  • 置信度塑形 (Confidence Shaping)： 引入基于生成答案归一化对数概率的置信度项，鼓励高置信度的正确答案，抑制高置信度的错误答案。
• 优化目标：
  • 计算采样轨迹的密度比 $\rho$（当前策略与参考策略之比）。
  • 使用截断的 PPO/GRPO 损失函数最大化期望奖励。
  • 加入 KL 散度正则化项，防止策略偏离参考分布过远，确保训练稳定性。

---

3. 关键贡献 (Key Contributions)

1. 诊断分析： 揭示了启发式潜在 ITS 的关键缺陷——有效的扩展需要控制采样质量，而不仅仅是注入随机性。强扰动往往导致分布偏移过大但采样增益（Sampling Gain）下降。
2. 提出 GTS： 将潜在扰动重新形式化为条件采样问题，将启发式扰动转化为显式的、可优化的概率采样策略。
3. 实证验证： 在多个基准测试（GSM8K, MultiArith, SVAMP 等）和两种不同的潜在推理架构（COCONUT, CODI）上，GTS consistently 优于基于 Dropout 和标准高斯噪声的基线方法。

---

4. 实验结果 (Results)

4.1 推理时扩展性能 (Scaling Behavior)

• Pass@N 指标： 随着采样预算 $N$（采样轨迹数量）的增加，GTS 展现出比启发式基线（Dropout, Standard Gaussian）更强的扩展能力。
• 对比基线：
  • Dropout (p=0.1)： 在小到中等预算下表现尚可，但在大预算下提升有限。
  • Dropout (p=0.5)： 往往导致性能下降，特别是在 COCONUT 模型上，表现为严重的过度探索。
  • GTS： 在中等到大预算下 consistently 取得最佳性能，表明学习到的采样策略能生成更多“有用”的轨迹多样性。

4.2 采样质量分析 (Sampling Quality)

作者引入了 采样增益 (Sampling Gain, SG) 和 分布偏移 (Distribution Shift, JS 散度) 指标：

• SG (Sampling Gain)： 衡量采样轨迹是否提高了模型对正确答案的预测置信度。
• 结果：
  • 启发式方法（特别是强 Dropout）虽然增加了分布偏移（JS 散度高），但 SG 往往为负或很低（即扰动让模型偏离了正确答案区域）。
  • GTS 在保持适度分布偏移的同时，实现了最高的 SG 和 SG 率（SG > 0.5 的比例），证明其扰动是有方向性且有利于决策的。

4.3 泛化能力 (Out-of-Distribution)

在未见过的算术推理基准（MultiArith, SVAMP, GSM8K-Hard）上，GTS 依然表现出比固定扰动更稳定的扩展性能，证明了学习到的扰动策略具有更好的泛化性。

4.4 行为分析

• 信噪比 (SNR) 分析： GTS 能够根据骨干网络的内部动态自适应调整确定性引导与随机探索的平衡。例如，在 COCONUT 中呈现锯齿状模式（交替进行信息巩固与探索），而在 CODI 中则表现为早期强引导、后期探索的模式。

---

5. 意义与结论 (Significance & Conclusion)

• 理论意义： 论文挑战了“更多随机性=更好推理”的直觉，指出在连续潜在空间中，受控的、可优化的采样比单纯的随机扰动更为关键。GTS 为潜在推理的 ITS 提供了一个基于概率分布的、原则性的框架。
• 实践价值：
  • 提供了一种轻量级的插件式模块，无需微调昂贵的骨干模型即可显著提升推理性能。
  • 解决了现有潜在推理模型在推理时扩展中难以控制探索范围的问题。
• 局限性： 目前主要局限于数学推理任务，采样策略局限于对角高斯分布，且缺乏对高维潜在空间探索的严格理论表征。

总结： 该论文通过引入可学习的高斯采样器，成功将启发式的潜在扰动转化为可优化的概率策略，显著提升了潜在推理模型在推理时的扩展能力，证明了**“更好的引导”比“更多的噪声”**对于连续推理空间的探索更为有效。