技术摘要：基于潜层思维向量的推理时重构（Inference-Time Rethinking）用于数学推理

1. 问题陈述

标准的大语言模型（LLM）思维链（CoT）推理通过单次前向传播运行，模型会对生成的每个标记（token）做出不可撤销的承诺。这种架构将陈述性知识（推理计划的内容）与程序性执行（标记的句法生成）纠缠在一起。因此，生成过程中的早期错误会由于模型缺乏在最终输出前对推理策略进行回溯、批判或优化的独立机制，而导致错误无法得到及时纠正并持续传播。

作者认为，有效的推理需要实现“思考什么”（陈述性）与“如何将其语言化”（程序性）之间的功能分离，这模仿了人类认知中语义处理（如 Wernicke 区）与句法产出（如 Broca 区）之间的区别。

2. 方法论

2.1 带有潜层思维向量的生成模型

所提出的框架将推理分解为两个不同的组成部分：

1. 潜层思维向量 ($z$)： 一种连续的、陈述性的推理策略表示（“思考什么”）。
2. 程序性解码器： 一个以 $z$ 为条件的自回归 Transformer，用于生成标记序列（“如何表达”）。

联合模型定义为：
$$p_\theta(x, z) = p_\alpha(z) p_\beta(x|z)$$
其中 $x$ 是地面标记序列（问题 + 推理轨迹），$\theta = (\alpha, \beta)$ 代表所有模型参数。

• 先验模型 ($p_\alpha(z)$)： 不同于使用固定先验的标准 VAE，该模型使用了一个可学习的传输映射（transport map）。编码器 ($U_\alpha$) 将无结构的噪声标记 $z_0 \sim \mathcal{N}(0, I)$ 转换为结构化的潜层思维向量 $z$。这创建了一个学习到的流形，其几何结构反映了有效的推理模式。
• 思维引导生成器 ($p_\beta(x|z)$)： 一个 Transformer 解码器，其中潜层向量 $z$ 对每一层进行交叉注意力（cross-attend）。至关重要的是，解码器使用了短上下文窗口（$w=64$）。这种架构上的瓶颈迫使模型通过 $z$ 来路由长程依赖，从而强制要求潜层向量编码全局的、序列级的结构，而非冗余的局部上下文。

2.2 学习：双速率优化

训练采用了一种双速率方案，以分离快速的单实例适配与缓慢的全局学习：

• 快速局部更新： 对于每个训练序列，通过高学习率使用梯度下降（Adam）直接优化噪声标记 $z_0$ 的变分后验参数 $(\mu, \sigma^2)$。不训练推理网络；每个实例都维护其各自优化的潜层表示。
• 缓慢全局更新： 使用标准的较低学习率更新模型参数 $\theta$（编码器和解码器权重）。

这种方法避免了后验坍缩（posterior collapse），并允许模型在保持单实例精确度的同时，积累通用的模式（$\theta$）。

2.3 推理时重构

在测试阶段，模型采用一种 Gibbs 式迭代程序来精炼推理策略。给定问题 $x_q$，目标是通过交替进行生成与反思来逼近边际分布 $p_\theta(x_r|x_q)$：

1. 生成： 在当前潜层思维 $z^{(t-1)}$ 的条件下，采样一个候选推理轨迹 $x_r^{(t)}$。
   $$x_r^{(t)} \sim p_\beta(x_r | U_\alpha(z_0^{(t-1)}), x_q)$$
2. 反思： 优化潜层噪声向量 $z_0$，以最大化生成轨迹的似然度（有效地最小化变分后验与真实后验之间的 KL 散度）。
   $$z_0^{(t)} = \arg \min_{z_0} D_{KL}(q(z_0|x_r^{(t)}, x_q) \parallel p_\theta(z_0|x_r^{(t)}, x_q))$$

这个循环允许模型进行“重构”：“推理轨迹中的错误会反馈给精炼后的潜层向量，进而引导下一轮生成更正后的轨迹。该过程在理论上被证明是基于变分下界（ELBO）的坐标上升优化，保证了在近似联合后验方面的单调改进。

3. 核心贡献

1. 生成式框架： 一种新颖的架构，将陈述性内容（潜层向量）与程序性生成解耦，使得小型模型能够超越大型单体架构。
2. 推理时重构： 一种 Gibbs 式程序，在测试时迭代精炼潜层思维，无需额外的训练数据或模型重参数化即可实现自我修正。
3. 经验验证： 证明了推理时计算是一个极其有效的缩放维度，与参数量是正交的。

4. 实验结果

作者从头开始在 GSM8K-Aug 数据集（385K 个样本）上训练了一个 0.2B 参数的模型，并在三个基准测试上进行了评估：GSM8K（域内）、SVAMP（域外表面变化）和 MultiArith（多步推理）。

• 性能： 经过 30 次重构迭代的 0.2B 模型（Rethink-30）取得了：
  • 在 GSM8K 上达到 31.54%（超过了 3B 参数的 CoT-SFT 模型 22.70% 和 2B 参数的 MARCoS 模型 24.11%）。
  • 在 SVAMP 上达到 51.50%（相比之下，2B 参数的 MARCoS 为 27.77%）。
  • 在 MultiArith 上达到 68.00%（相比之下，2B 参数的 MARCoS 为 42.33%）。
• 效率： 即使是单次通过（Rethink-1）也优于几乎所有的基准模型，包括那些参数量大 10–15 倍的模型。
• 泛化能力： 该方法在域外基准测试中表现出显著的鲁棒性，这归功于潜层向量能够抽象掉表层的标记变体。

5. 意义与主张

本文声称，有效的数学推理可以从复杂的推理时计算中涌现，而不仅仅依赖于庞大的参数量。

• 正交缩放： 结果表明，“思考更久”（投入计算资源进行迭代精炼）是一个强大的缩放维度，它可以与增加模型规模相辅相成。
• 认知对齐： 通过将陈述性规划与程序性执行分离，该模型模拟了人类的认知机制，允许在做出最终决定之前对思维进行检查和修订。
• 似然度作为代理： 作者假设，当在接近专家的演示数据上进行训练时，似然度景观自然地倾向于有效的逻辑路径。因此，通过优化潜层向量以最大化生成轨迹的似然度，实际上是在引导模型走向正确的推理策略。

作者指出了一些局限性，特别是依赖似然度作为正确性的代理，这假设了高质量的训练数据；在存在噪声的环境中，这种相关性可能会减弱，因此未来的工作需要引入外部验证信号。