本文提出 LATENTSEEK 框架,通过在潜在空间中基于策略梯度的测试时实例级适应(TTIA),显著提升大型语言模型的推理能力,同时探索测试时扩展的新方向。
Large Language Model, Reasoning, Test Time, Reinforcement Learning, Representation Learning
Hengli Li, Chenxi Li, Tong Wu, Xuekai Zhu, Yuxuan Wang, Zhaoxin Yu, Eric Hanchen Jiang, Song-Chun Zhu, Zixia Jia, Ying Nian Wu, Zilong Zheng
北京大学人工智能研究所, 北京通用人工智能研究所 NLCo 实验室, 清华大学自动化系, 上海交通大学, 中国科学院自动化研究所, 加州大学洛杉矶分校
Generated by grok-3
Background Problem
大型语言模型(LLMs)在推理能力上仍面临挑战,尤其是在需要结构化思维和逐步分析的任务中。传统的训练方法(如监督微调和强化学习)虽然有所改进,但存在高计算成本、灾难性遗忘以及探索能力下降等问题。此外,新型训练数据的有限性也限制了模型性能的进一步提升。论文提出了一种测试时实例级适应(TTIA)范式,旨在通过增加测试时计算而非更新参数来增强推理能力,解决上述问题,并探索潜在空间中测试时扩展(test-time scaling)的潜力。
Method
论文提出了 LATENTSEEK 框架,通过在潜在空间中进行测试时实例级适应(TTIA)来提升 LLMs 的推理能力。具体方法如下:
- 核心思想:不更新模型参数,而是通过策略梯度方法在测试时针对每个问题实例优化潜在表示(latent representations),引导模型生成更好的推理路径。
- 实现步骤:
- 将输入问题的初始潜在表示通过预训练模型生成,并以 Chain-of-Thought (CoT) 初始化作为起点。
- 使用策略梯度(基于 REINFORCE 算法)迭代更新潜在表示,目标是最大化自生成的奖励信号(self-reward)。
- 在每次迭代中,将更新后的潜在表示解码为 token,计算奖励,并根据奖励调整潜在表示,直至奖励超过预设阈值或达到最大迭代次数。
- 采用部分序列优化(fractional sequence optimization),仅更新潜在表示序列的一部分(由超参数 ρ 控制),以降低计算成本并提高稳定性。
- 奖励函数:采用自奖励机制,依赖模型内部能力生成奖励信号,无需外部数据。
- 关键创新:在潜在空间而非 token 空间中进行优化,理论上通过多证明者交互证明(MIP)支持其表达能力,强调独立更新潜在表示仍能保持模型的强大推理能力。
批判性思考:虽然潜在空间优化避免了参数更新带来的成本,但其理论依据(MIP)与实际推理任务的关联性较弱,缺乏直接证据证明潜在空间的优化确实捕捉到了推理的本质而非特定任务的模式。此外,自奖励机制在复杂任务中可能因奖励信号不准确而导致优化陷入局部最优,论文未充分讨论这一风险。
Experiment
实验在多个推理基准数据集(GSM8K, MATH-500, AIME2024)上进行,测试了不同模型架构(Qwen2, Qwen2.5, LLaMA3.1, Mistral)和规模(1.5B 到 14B 参数)。
- 数据集与设置:GSM8K 和 MATH-500 用于评估一般数学推理能力,AIME2024 则测试复杂问题解决能力。实验采用两种提示方式(Prompt 1 和 Prompt 2),并对比了多种基线方法,包括 CoT、Few-Shot CoT、Best-of-N (BoN)、强化学习方法(如 Self-Rewarding, Genius)以及监督微调(SFT)。
- 结果:
- LATENTSEEK(自奖励)在 GSM8K 上平均提升 CoT 10.75%,在 MATH-500 上提升 3.93%,在 AIME2024 上提升 4.73%,优于大多数基线方法,尤其在 LLaMA3.1-8B-Instruct 上表现突出(GSM8K 提升 14.6%)。
- 使用完美稀疏奖励模型(PSRM)时,性能进一步提升,平均比 CoT 高 19.12%,显示潜在空间优化的潜力。
- 测试时扩展实验表明,增加迭代次数可提升性能,尤其在 PSRM 下,小模型(1.5B)性能接近甚至超过大模型(7B)。
- 分析与合理性:实验设置覆盖了多种模型和任务,较为全面,但对 PSRM 的依赖可能高估了实际应用效果,因为现实中难以获得完美奖励信号。自奖励机制在部分模型(如 Mistral-7B)上表现不佳,可能是由于其奖励生成能力不足,实验未深入探讨这一局限性。此外,迭代次数与性能的关系显示出测试时扩展的潜力,但平均迭代次数较少(GSM8K 为 0.86,MATH-500 为 1.23),可能未充分体现复杂任务的挑战。
- 批判性思考:虽然结果显示出改进,但实验缺乏对自奖励机制失败案例的详细分析,也未充分比较潜在空间优化与 token 空间优化的具体差异。此外,部分数据集(如 AIME2024)样本量较小(仅 30 个问题),可能导致结果的统计显著性不足。
Further Thoughts
LATENTSEEK 的潜在空间优化方法启发我们重新思考机器推理的本质,尤其是在与人类推理路径的差异上。论文中提到的机器推理路径可能更适合在潜在空间中导航的观点,与近期一些关于压缩思维链(Compressed Chain of Thought)的研究相呼应,这些研究也尝试通过连续表示而非离散文本来表达推理过程。未来可以探索潜在空间优化是否能与其他技术(如提示优化或多代理验证框架)结合,以进一步提升奖励信号的准确性和推理的稳定性。此外,测试时扩展的概念在计算资源受限的场景中可能面临挑战,如何在效率与性能之间找到平衡,特别是在更大规模模型(如 70B 或 405B 参数)上的应用,是一个值得深入研究的方向。另一个有趣的点是,潜在空间的优化是否能推广到其他任务领域(如多模态推理或代码生成),这可能需要重新设计奖励机制和优化策略,以适应不同任务的特性。