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

2026-03-05

AI Summary

This work proposes $\nabla$-Reasoner, a novel framework that introduces first-order gradient-based optimization at test time to enhance reasoning in large language models (LLMs). Unlike conventional approaches relying on inefficient discrete search or trial-and-error prompting, $\nabla$-Reasoner dynamically refines the generation strategy by performing differentiable optimization over token logits during decoding, jointly leveraging the language model’s likelihood and gradient signals from a reward model. The method integrates rejection sampling with decoding acceleration mechanisms to improve both robustness and computational efficiency. Theoretically, it establishes a duality between this approach and KL-regularized reinforcement learning. Experimental results demonstrate significant gains, achieving over 20% absolute accuracy improvement on challenging mathematical reasoning benchmarks while reducing model invocation counts by 10–40%.

Technology Category

• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• inference-time scaling
• LLM reasoning
• discrete search
• trial-and-error prompting
• online policy optimization

Innovation

• Differentiable Textual Optimization
• test-time gradient descent
• latent space optimization
• LLM reasoning
• first-order optimization