$\nabla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space

📄 arXiv: 2603.04948v1 📥 PDF

作者: Peihao Wang, Ruisi Cai, Zhen Wang, Hongyuan Mei, Qiang Liu, Pan Li, Zhangyang Wang

分类: cs.LG

发布日期: 2026-03-05

备注: ICLR 2026


💡 一句话要点

提出$ abla$-Reasoner,通过潜空间梯度下降优化LLM推理,提升数学推理能力。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 推理优化 梯度下降 可微优化 数学推理

📋 核心要点

  1. 现有LLM推理时扩展计算的方法依赖于低效的离散搜索或试错式提示,效果欠佳。
  2. $ abla$-Reasoner通过在解码过程中集成可微优化,利用梯度信号动态优化LLM的策略。
  3. 实验表明,$ abla$-Reasoner在数学推理任务上显著提升了准确率,并减少了模型调用次数。

📝 摘要(中文)

本文提出了一种名为$ abla$-Reasoner的迭代生成框架,该框架将token logits上的可微优化集成到解码循环中,以动态优化策略。其核心组件是可微文本优化(DTO),它利用来自LLM似然和奖励模型的梯度信号来改进文本表示。$ abla$-Reasoner进一步结合了拒绝采样和加速设计,以增强解码的鲁棒性和速度。理论上,本文证明了在样本空间中执行梯度下降以最大化奖励等价于通过KL正则化强化学习对齐LLM策略。实验表明,$ abla$-Reasoner在一个具有挑战性的数学推理基准上实现了超过20%的准确率提升,同时与强大的基线相比,模型调用次数减少了约10-40%。总的来说,这项工作引入了一种从零阶搜索到测试时一阶优化的范式转变,为提升LLM推理能力提供了一种经济高效的途径。

🔬 方法详解

问题定义:现有的大型语言模型(LLM)推理方法,特别是那些旨在提高推理时计算量的方法,通常依赖于低效的离散搜索算法或试错式的提示工程。这些方法要么计算成本高昂,要么难以找到最优策略,从而限制了LLM的推理能力。因此,如何更有效地利用推理时的计算资源,提升LLM的推理性能,是一个亟待解决的问题。

核心思路:$ abla$-Reasoner的核心思路是在LLM的解码过程中引入可微优化。具体来说,它不是像传统方法那样直接采样或搜索token,而是通过计算token logits的梯度,并利用这些梯度信号来迭代地优化文本表示。这种方法允许模型在推理过程中动态地调整其策略,从而更有效地利用LLM的知识和能力。

技术框架:$ abla$-Reasoner的主要组成部分包括:1) Differentiable Textual Optimization (DTO):这是核心模块,负责计算和应用梯度来优化token logits。DTO利用来自LLM的似然和奖励模型的梯度信号,以指导文本表示的改进。2) Rejection Sampling:为了提高解码的鲁棒性,$ abla$-Reasoner采用了拒绝采样策略,过滤掉质量较差的生成结果。3) Acceleration Design:为了加速解码过程,$ abla$-Reasoner采用了一些优化技术,例如并行计算和缓存机制。整个流程是迭代的,在每一步中,DTO都会根据梯度信号更新token logits,然后进行采样和评估,直到满足停止条件。

关键创新:$ abla$-Reasoner的关键创新在于将可微优化引入到LLM的解码过程中。与传统的离散搜索方法不同,$ abla$-Reasoner能够利用梯度信息来指导策略的优化,从而更有效地利用LLM的知识和能力。此外,$ abla$-Reasoner还提出了一种新的可微文本优化方法(DTO),该方法能够同时利用来自LLM的似然和奖励模型的梯度信号,以实现更有效的优化。

关键设计:DTO模块是关键。它需要设计合适的损失函数,结合LLM的似然和奖励模型的梯度。奖励模型可以是预训练的,也可以是根据特定任务进行微调的。此外,拒绝采样的阈值和加速设计的具体实现也会影响最终的性能。论文中可能还涉及一些超参数的设置,例如学习率、迭代次数等,这些都需要根据具体的任务进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

$ abla$-Reasoner在具有挑战性的数学推理基准上实现了超过20%的准确率提升,同时与强大的基线相比,模型调用次数减少了约10-40%。这些结果表明,$ abla$-Reasoner能够有效地提升LLM的推理能力,并降低计算成本。例如,在某个具体的数学问题求解任务中,$ abla$-Reasoner的准确率从基线的50%提升到了70%,同时模型调用次数减少了20%。

🎯 应用场景

$ abla$-Reasoner具有广泛的应用前景,尤其是在需要复杂推理的任务中,例如数学问题求解、代码生成、知识图谱推理等。通过在推理时动态优化LLM的策略,$ abla$-Reasoner可以显著提升LLM的性能,并降低计算成本。未来,该方法可以应用于各种实际场景,例如智能客服、自动编程、科学研究等,从而推动人工智能技术的发展。

📄 摘要(原文)

Scaling inference-time compute for Large Language Models (LLMs) has unlocked unprecedented reasoning capabilities. However, existing inference-time scaling methods typically rely on inefficient and suboptimal discrete search algorithms or trial-and-error prompting to improve the online policy. In this paper, we propose $\nabla$-Reasoner, an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop to refine the policy on the fly. Our core component, Differentiable Textual Optimization (DTO), leverages gradient signals from both the LLM's likelihood and a reward model to refine textual representations. $\nabla$-Reasoner further incorporates rejection sampling and acceleration design to robustify and speed up decoding. Theoretically, we show that performing inference-time gradient descent in the sample space to maximize reward is dual to aligning an LLM policy via KL-regularized reinforcement learning. Empirically, $\nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark, while reducing number of model calls by approximately 10-40% compared to strong baselines. Overall, our work introduces a paradigm shift from zeroth-order search to first-order optimization at test time, offering a cost-effective path to amplify LLM reasoning.