PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2603.04948v1cs.LG

TL;DR

Existing inference-time scaling methods rely on discrete search or trial-and-error prompting, motivating a more efficient way to improve reasoning policies at test time. ∇-Reasoner performs differentiable token-logit optimization through DTO, supplemented by rejection sampling and acceleration, and connects this procedure theoretically to KL-regularized RL. Empirically, it improves mathematical reasoning while reducing model calls relative to strong baselines.

  • Problem

    Existing inference-time scaling methods rely on inefficient discrete search or trial-and-error prompting to improve the online policy, while search struggles with sparse rewards and expanding reasoning spaces.

  • Method

    ∇-Reasoner iteratively refines base-policy token logits with DTO using LLM-likelihood and reward gradients, then applies rejection sampling and acceleration during decoding.

  • Results

    10-40% improvement in mathematical reasoning was reported across multiple models and benchmarks, with up to 40.2% fewer model calls than sampling-only methods.

  • Takeaways & Limitations

    The framework shifts inference-time reasoning from zeroth-order search toward first-order gradient optimization while reducing computation cost.

  • Takeaways & Limitations

    Performance remains bounded by the underlying base and reward models, especially under limited computation; both models must share a vocabulary for end-to-end logit optimization.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Inference-time scaling can improve LLM reasoning, but existing zeroth-order search methods struggle with sparse rewards and exponentially expanding reasoning spaces. ∇-Reasoner addresses this by refining token logits with gradients during iterative decoding, combining DTO with rejection sampling, acceleration, and a theoretical RL connection.

  • Motivation: Existing inference-time methods often saturate because they inadequately explore exponentially expanding reasoning spaces with sparse and noisy rewards.These approaches rely on sampling and evaluating candidate responses, with the LLM serving as a guiding heuristic.
  • Method: ∇-Reasoner applies inference-time gradient descent in sample space to refine a base policy before next-token prediction.Its iterative pipeline generates a completion with per-token logits, then DTO refines those logits via gradient descent.
  • Method: DTO directly leverages reward gradients and LLM likelihood signals to optimize textual representations under a combined objective.The method uses gradient descent on initial logit vectors and a straight-through estimator to support differentiable token representations.
  • Method: Rejection sampling accepts refined-policy tokens only when their continuations can achieve higher reward, while acceleration strategies skip unlikely-beneficial tokens and reuse rollouts.These components are interleaved with decoding to improve robustness and throughput.
  • Theory: DTO theoretically connects gradient-based sample refinement with RL-based policy alignment and enables bidirectional gradient propagation across sequences.The analysis characterizes sampling from an RL-optimized policy as equivalent to refining reference-model samples through DTO-induced gradient flow.
  • Results: 10-40% improvement in mathematical reasoning was reported across multiple models and benchmarks, with up to 40.2% fewer model calls than sampling-only methods.The method outperformed strong inference-time baselines and reached accuracy comparable to more costly training-based methods.

2 PRELIMINARIES

The paper formulates LLM reasoning as sequential decision-making over discrete token sequences, where the goal is to identify high-reward responses. Autoregressive structure reduces the problem to next-token decisions, but globally optimal choices require look-ahead and backtracking.

  • Language Models and Reward Models: An autoregressive language model assigns a question-answer pair probability by factorizing it into conditional next-token probabilities.The model also defines a categorical distribution for each next token given the question and preceding response prefix.
  • Language Models and Reward Models: A reward model evaluates response correctness, primarily through an overall sequence-level outcome reward in this work.The framework can also generalize to process rewards.
  • Reasoning as Decision Making: LLM reasoning is framed as a search problem over finite vocabulary sequences that seeks responses with high reward.The sequence space grows exponentially with response length, making direct optimization intractable.
  • Reasoning as Decision Making: Greedy decoding is not globally optimal because selecting the best next token requires look-ahead rollouts and backtracking.The recursive expected-reward formulation is identified with the Q-function.
  • Existing Approaches: Inference-time methods improve decoding without additional training, whereas training-time approaches include supervised fine-tuning and policy optimization.The paper focuses on inference-time methods for reasoning via decision making.

3 REASONING WITH GRADIENT-DRIVEN DECODING

∇-Reasoner refines autoregressive decoding by applying gradient-based optimization to token logits, then uses resampling and rejection sampling to improve responses. DTO combines reward and language-model likelihood signals, while caching and selective refinement reduce redundant computation.

  • 3 REASONING WITH GRADIENT-DRIVEN DECODING: ∇-Reasoner iteratively optimizes base-model logits to maximize reward, treating refined first-token logits as an improved policy for decoding.The process repeats optimization, resampling, and prefix extension token by token.
  • 3.2 ITERATIVE DECODING WITH DTO: Rejection sampling retains a resampled token only when its regenerated continuation receives higher reward than the original response.If the candidate is unchanged, decoding proceeds without an additional effective policy update.
  • 3.1 DIFFERENTIABLE TEXTUAL OPTIMIZATION: DTO differentiates reward over token space by optimizing logits with an objective combining sequence-level reward and language-model log-likelihood.The log-likelihood regularizer discourages deviations from the language model’s distribution and helps preserve fluency and faithfulness.
  • 3.1 DIFFERENTIABLE TEXTUAL OPTIMIZATION: Gumbel-softmax straight-through parameterization enables gradient descent over discrete textual representations through continuous logit vectors.The logits are initialized from the LLM-generated response and updated using the gradient of the objective.
  • 3.1 DIFFERENTIABLE TEXTUAL OPTIMIZATION: DTO propagates information bidirectionally: preceding tokens regularize later predictions, while trailing tokens transmit sequence-level reward signals to earlier tokens.This supports global consistency and links future-context feedback to earlier decoding decisions.
  • 3.3 ACCELERATION: Acceleration uses rollout reuse, gradient caching, and entropy- or gradient-based token selection to avoid unnecessary optimization and generation.DTO is applied only when token-logit entropy and gradient magnitude exceed their respective thresholds.

4 THEORETICAL ANALYSIS

The theoretical analysis connects sample-space gradient optimization with KL-regularized policy optimization. It argues that DTO-driven trajectories can generate samples from the reward-optimized distribution without directly optimizing the entire policy.

  • 4 THEORETICAL ANALYSIS: The loss gradient decomposes into prefix, postfix, and reward components that respectively encode autoregressive context, future-context consistency, and sequence-level feedback.These components explain how DTO updates individual tokens using information from both directions and the reward model.
  • 4 THEORETICAL ANALYSIS: KL-regularized policy optimization minimizes negative expected reward plus the KL divergence between the optimized distribution and the base LLM.The KL term regularizes distributional discrepancy from π_LLM.
  • 4 THEORETICAL ANALYSIS: Theorem 4.1 states that samples from the optimal KL-regularized policy can be obtained by initializing from the reference LLM and following stochastic gradient flow in sample space.Thus, DTO provides a trajectory-based alternative to optimizing the entire policy distribution.
  • 4 THEORETICAL ANALYSIS: The paper interprets pre-training scaling and test-time scaling as complementary forms of statistical inference: parametric and non-parametric, respectively.The non-parametric view treats test-time optimization as particle-based inference in sample space.

5 EXPERIMENTS

Experiments evaluate ∇-Reasoner across mathematical reasoning benchmarks, model families, baselines, reward models, computational cost, scaling curves, and rejection sampling. The method improves accuracy and efficiency while maintaining robustness to reward-model size and reducing rejection rates.

  • Experimental Setup: Experiments cover Qwen-2.5-math and Llama-3.1 models on MATH-500, AIME24, AIME25, and AMC, comparing test-time and training-based methods.BoN and SC use N = 8, while ∇-Reasoner uses Nmax = 8.
  • Performance Comparison: ∇-Reasoner achieves accuracy comparable to training-based methods while outperforming all test-time baselines in the reported comparisons.The comparison includes SFT and GRPO as training-based methods.
  • Performance Comparison: ∇-Reasoner leads test-time methods across models and benchmarks, reaching 80.4% on MATH-500 and 56.8% on AMC with Qwen-2.5-7B-Instruct.It also achieves 55.8% on MATH-500 and 28.9% on AMC with Llama-3.1-8B-Instruct.
  • Cost Comparison: 40.2% fewer model calls are used for instruction-tuned models than baselines, while base models outperform all baselines using about 90% of the model-call metric.The cost metric counts recurrent and parallel forward computations as one model call under ideal attention parallelization.
  • Test-Time Scaling Law: Across models and computational budgets, ∇-Reasoner’s test-time scaling curve remains above BoN and SC, indicating higher accuracy for a given number of calls.The method changes N for BoN and SC and Nmax for ∇-Reasoner.
  • Algorithm Analysis: Using the smaller Skywork-Reward-V2-Qwen3-4B model produces a performance gap below 1 point versus the 8B variant on MATH-500 and AMC.The result supports the original smaller reward-model choice and suggests potential efficiency benefits.
  • Algorithm Analysis: DTO reduces rejection rates by up to 30% relative to rejection sampling without DTO, compared with the baseline’s theoretical 66.0% rejection rate at N = 8.The analysis defines rejection as the percentage of DTO candidates rejected for failing to improve reward.

6 CONCLUSION AND LIMITATIONS

∇-Reasoner refines token logits during inference with DTO, combining gradient-based policy improvement, rejection sampling, and acceleration techniques. The framework is theoretically connected to reward alignment and is reported to improve performance while reducing computation cost, subject to model, vocabulary, and serving constraints.

  • 6 CONCLUSION AND LIMITATIONS: ∇-Reasoner introduces DTO to refine token logits through gradient-based optimization during inference.It combines gradients from the LLM likelihood and a reward model.
  • 6 CONCLUSION AND LIMITATIONS: Rejection sampling and speedup techniques improve the framework’s effectiveness and efficiency.The method accepts refined choices conditionally and incorporates acceleration techniques into iterative decoding.
  • 6 CONCLUSION AND LIMITATIONS: ∇-Reasoner frames reward alignment as equivalent to gradient-based optimization in the sample space.This provides the paper’s central theoretical connection between inference-time optimization and policy alignment.
  • 6 CONCLUSION AND LIMITATIONS: The framework’s performance remains bounded by the capabilities of its base and reward models, especially under limited computation.It also requires shared vocabularies and more careful system co-design for efficient serving.

B IMPLEMENTATION

The implementation supports outcome or process rewards and repeatedly integrates DTO into iterative decoding. The full algorithm also includes gradient-guided selection, rollout reuse, early stopping, and gradient caching to accelerate decoding.

  • B IMPLEMENTATION: The complete implementation adds confidence- and gradient-guided token selection, rollout reuse, early stopping, and gradient caching.These techniques are presented as acceleration components for iterative decoding and DTO.
  • B IMPLEMENTATION: The reward can be an outcome score for the complete response or a process score assigned to intermediate reasoning steps.The generalized total reward sums rewards over different subsequences.

C.1 GRADIENT DERIVATION OF L

The gradient derivation decomposes DTO’s loss by token and distinguishes its dependence on language-model likelihood and reward terms. The derivation emphasizes that DTO propagates information beyond prior context.

  • C.1 GRADIENT DERIVATION OF L: The generalized reward is written as a sum over rewards defined on different subsequences.The proposition derives gradients for this generalized reward formulation.
  • C.1 GRADIENT DERIVATION OF L: For each token index, the loss is decomposed into five components before calculating the relevant partial derivatives.The derivation identifies which components depend on the selected token.
  • C.1 GRADIENT DERIVATION OF L: DTO differs from prior controlled-generation methods because its computational graph does not detach the postfix sequence.Instead, the method allows information from later tokens to influence earlier-token optimization.

C.2 PROOF OF THEOREM 4.1

The proof connects DTO’s stochastic gradient flow in sample space with Wasserstein gradient flow over distributions. Under the stated boundary and limiting conditions, samples initialized from the LLM converge in distribution to the optimized target.

  • C.2 PROOF OF THEOREM 4.1: The theorem considers a Wasserstein gradient flow starting from the LLM policy and converging to an optimized distribution.The boundary conditions are ρ0 = πLLM and ρ∞ = ρ* = arg minρ L_PPO(ρ).
  • C.2 PROOF OF THEOREM 4.1: The proof derives the distributional evolution through first variations, a partial differential equation, and the Fokker–Planck equation.These steps connect the sample dynamics to the time evolution of the density.
  • C.2 PROOF OF THEOREM 4.1: Samples from the optimized distribution can be obtained by initializing from the reference LLM and simulating the theorem’s stochastic gradient flow.The result establishes the distributional equivalence used to connect DTO with reward-aligned optimization.
  • C.2 PROOF OF THEOREM 4.1: DTO parameterizes token distributions with softmax and skips highly confident tokens because their logits are difficult to update effectively.The gradient magnitude is proportional to the post-softmax probability, so low-probability alternatives receive weak updates.

D.1 EXPERIMENT DETAILS

The experiments use fixed decoding settings across baselines and methods, with dataset-specific generation limits and repeated runs on smaller datasets. ∇-Reasoner uses fixed optimization hyperparameters, dataset subsets for training baselines, and family-specific reward models.

  • Evaluation settings: Temperature 0.5, top-p 0.95, and maximum generation lengths of 1024 for AMC and MATH-500 and 3072 for AIME were used across methods.Performance on AMC, AIME-24, and AIME-25 was averaged across 4 independent runs.
  • Training data: SFT used a random 10k Open-thoughts subset, while GRPO used a random 35k Numina-Math subset.
  • ∇-Reasoner settings: ∇-Reasoner used ϵ_ent = 0.25, ϵ_grad = 8, learning rate 0.01, and 20 optimization iterations in all experiments.The optimizer was Adam-W with a cosine schedule and minimum learning rate 0.001.
  • Reward models: The reward model was Skywork-Reward-V2-Qwen3-4B for the Qwen family and Skywork-Reward-V2-Llama-3.1-8B for Llama models.
  • Efficiency analysis: Three proposed decoding-efficiency techniques were evaluated by measuring how each reduced the cost of its applied stage.Direct runtime ablation was avoided because removing any component could make execution exceed several hours per sample.

D.3 WALL-CLOCK TIME BENCHMARKING UNDER IDEALIZED SETTINGS

The wall-clock benchmark compares ∇-Reasoner with BoN under an idealized deployment that fully batches generation and optimization. Execution traces are replayed to estimate parallel runtime and FLOPs on AMC.

  • Benchmark design: The benchmark compares wall-clock execution time for 83 AMC prompts between ∇-Reasoner and BoN.
  • Deployment setting: Generation and optimization requests are treated as batched operations with all prompts processed in parallel under the idealized serving assumption.
  • Measurement procedure: Recorded reasoning traces were replayed through a vLLM server to reconstruct parallel generation, batched gradient computation, wall-clock time, and computational FLOPs.The experiment used Qwen-2.5-math-Instruct-7B on 8x 140GB NVIDIA H200 GPUs.

D.4 A SIMPLIFIED EXAMPLE ON GSM8K

A GSM8K-style house-sale problem illustrates how greedy decoding miscomputes a percentage increase, while DTO-based refinement revises the reasoning structure and calculation to obtain the correct profit.

  • Greedy Decoding Baseline (Failure Analysis): Greedy decoding interprets “increased by 150%” as 80,000 × 1.5 = 120,000 and reports an incorrect profit of -10,000.
  • Initial Input and Refinement Round 1 (Operator Correction): DTO-based refinement corrects the multiplication operator by changing “×” to “+” in the new-value calculation.Gradient descent revises the logical structure to begin with 80,000 + … .
  • Refinement Round 2 (Calculation Correction): Gradient optimization then replaces the low-scoring token 120 with 200, completing the corrected new house value as 200,000.
  • Final output: After two refinement rounds, the method derives a final profit of 70,000 from 200,000 − 80,000 − 50,000.
Loading 2603.04948v1…