超越模仿：面向主动潜在规划的强化学习方法 Beyond Imitation: Reinforcement Learning for Active Latent Planning

ATP-Latent的完整方法分为SFT阶段和RL阶段两个主要步骤。**SFT阶段**：(1) 使用Coconut的课程学习策略，在第k个课程阶段移除CoT标签的前k个步骤，由LLM生成$c \times k$个latent token $l_1, ..., l_{c \times k}$来替代。(2) 编码器损失$\mathcal{L}_{Enc}$：最大化剩余语言步骤在给定latent token条件下的对数似然。(3) 添加MLP作为latent head，预测每个latent token的均值$\mu_t$和标准差$\sigma_t$，使latent token服从高斯分布$l_t \sim \mathcal{N}(\mu_t, \text{diag}(\sigma_t^2))$，而非确定性值。(4) 解码器损失$\mathcal{L}_{Dec}$：训练另一个LLM作为VAE解码器，从latent token恢复原始语言标签。(5) Stop head损失$\mathcal{L}_{Stop}$：训练MLP判断何时停止生成latent token，确保每个latent token编码的信息量相当。(6) KL散度损失$\mathcal{L}_{KL}$：正则化潜在空间接近标准正态分布。总损失$\mathcal{L}_{SFT} = \mathcal{L}_{Enc} + \mathcal{L}_{Dec} + \mathcal{L}_{Stop} + \beta \mathcal{L}_{KL}$，其中$\beta = 10^{-3}/d$以优先保证推理过程。**RL阶段**：(1) 基于GRPO框架，对每个问题生成G个不同的latent CoT。(2) 计算coherence奖励$R_{Coh}(L)$：将latent token通过VAE解码器翻译为语言步骤，检查方程结果是否在后续步骤或答案中出现。(3) 定义总奖励$f(L, A) = (1 + 0.1 \cdot R_{Coh}(L)) R_{Correct}(A) + 0.5 \cdot R_{Format}(A)$，优先选择答案正确且latent CoT更一致的样本。(4) 使用GRPO的目标函数更新模型参数，引入stop policy的概率计算。

ATP-Latent在技术上具有以下新颖性：(1) **VAE建模的latent token训练**：首次将latent reasoning的模仿过程建模为条件VAE，相比之前的确定性映射，获得了更平滑的潜在空间表示，为RL探索提供了更好的基础。之前的方法（如Coconut、SIM-CoT）学习的是确定性的latent token，导致潜在空间"尖锐化"，限制了RL的探索能力。实验表明（Figure 3），在SIM-CoT上直接应用GRPO效果不佳，但在Coconut上则有效，论文将此归因于SIM-CoT对潜在空间的"尖锐化"。(2) **Coherence奖励**：提出了一种全新的无监督奖励信号，利用VAE解码器将latent token翻译回语言，检查解码步骤的逻辑一致性。这与现有RL方法仅依赖答案正确性（二值奖励）形成鲜明对比，提供了更细粒度的训练信号。(3) **Stop head机制**：引入可学习的停止策略，替代固定长度的latent token生成，使每个latent token编码的信息量更加均匀，进一步平滑化潜在空间。(4) **主动规划vs被动模仿**：整体框架从"被动模仿单一标签"转向"主动探索最优策略"，这是latent reasoning方法论上的范式转变。

从独立分析的角度来看，ATP-Latent存在以下弱点和改进方向：(1) **Coherence奖励的领域局限性**：当前coherence奖励的设计（方程RHS出现在后续LHS中）明显是为数学推理量身定制的。在自然语言推理或常识推理任务中，这种一致性检查标准不适用。改进方向是设计更通用的一致性度量，例如利用语言模型评估解码步骤之间的语义连贯性，或训练一个学习的coherence评价模型。(2) **VAE解码器的额外计算负担**：SFT阶段需要同时训练一个完整的LLM作为解码器（论文使用另一个LLaMA-1B），这实际上将模型参数量翻倍。在RL阶段，每次生成候选路径后还需要用解码器翻译latent token来计算coherence，增加了推理开销。改进方向可以探索更轻量的解码器架构，如小模型或专门训练的MLP解码器。(3) **Stop head的学习信号不够直接**：stop head只在第K步有明确的监督信号（停止），其他步骤的监督来自BCE损失，可能导致学习不够稳定。可以考虑引入强化学习来直接优化停止策略的长期收益。(4) **β值的选择过于保守**：论文选择$\beta = 10^{-3}/d$使得KL损失几乎可以忽略，这意味着VAE的正则化效果有限，潜在空间可能不够平滑。虽然论文解释这是因为latent token类别无限多，但这可能导致VAE退化为普通AE，失去概率建模的优势。(5) **缺少对latent token维度$d$的消融**：latent token的维度$d$是决定信息密度的关键参数，但论文未探索不同$d$值的影响。(6) **实验规模较小**：仅在1B模型上实验，缺乏对更大模型和更多任务的验证。

基于本文的成果，未来研究可以从以下方向展开：(1) **扩展到更多推理任务**：将ATP-Latent推广到代码生成、逻辑推理、常识推理等任务，设计通用的coherence奖励机制。可以探索利用LLM-as-a-judge来评估latent推理过程的质量。(2) **更大模型的验证**：在7B、13B甚至70B模型上验证ATP-Latent的效果，研究模型规模对latent planning能力的影响。更大的模型可能具有更丰富的潜在表示空间，可能从active planning中获益更多。(3) **多模态latent reasoning**：论文的参考文献中已有多模态latent reasoning的工作，将ATP-Latent的VAE+RL框架应用到视觉-语言模型的推理任务中是一个自然的扩展方向。(4) **与test-time compute scaling结合**：Pass@K的分析表明ATP-Latent具有更好的多样性，可以与best-of-N选择、多数投票等test-time scaling策略结合，进一步提升推理性能。(5) **探索混合策略**：将latent token和语言token混合使用，对简单的推理步骤使用latent token，对需要精确语言表达的步骤使用语言token。(6) **改进coherence度量**：设计更通用、更细粒度的coherence度量，例如基于信息论的度量，或者利用外部知识库验证推理步骤的正确性。

ATP-Latent的复现条件评估如下：(1) **开源代码**：论文提供了GitHub仓库https://github.com/zz1358m/ATP-Latent-master，包含完整的训练和评估代码，这对复现非常有利。(2) **数据集**：使用公开数据集GSM8K-Aug（385k实例，HuggingFace上可下载），以及GSM8K测试集、GSM-Hard、SVAMP、MultiArith等标准benchmark，数据获取无障碍。(3) **基础模型**：使用LLaMA-3.2-1B-Instruct，这是一个公开可用的模型（需遵守Llama 3.2 License），参数量较小，便于复现。(4) **计算资源要求**：论文使用8×NVIDIA H200 GPU（141GB显存），SFT阶段需要15个epoch，RL阶段需要1个epoch。对于1B模型来说这个资源要求不算高，但也不算低。单张高端消费级GPU（如RTX 4090 24GB）可能无法满足批量大小256的SFT训练，需要减小batch size或使用梯度累积。(5) **超参数**：论文在Table C中提供了详细的超参数设置，包括学习率、batch size、KL系数等，便于复现。(6) **复现难度**：中等偏低。代码开源、数据公开、模型公开、超参数完整，主要挑战在于计算资源和对latent reasoning框架的理解。整体来说，这是一个复现性较好的工作。