nabla-Reasoner: Razonamiento en LLM mediante Descenso de Gradiente en el Espacio Latente en Tiempo de Prueba
nabla-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space
March 5, 2026
Autores: Peihao Wang, Ruisi Cai, Zhen Wang, Hongyuan Mei, Qiang Liu, Pan Li, Zhangyang Wang
cs.AI
Resumen
La escalado del cómputo en tiempo de inferencia para los Modelos de Lenguaje a Gran Escala (LLMs) ha desbloqueado capacidades de razonamiento sin precedentes. Sin embargo, los métodos de escalado en tiempo de inferencia existentes suelen depender de algoritmos de búsqueda discreta ineficientes y subóptimos o de técnicas de prueba y error mediante *prompts* para mejorar la política en línea. En este artículo, proponemos nabla-Reasoner, un marco de generación iterativa que integra optimización diferenciable sobre los *logits* de los tokens en el bucle de decodificación para refinar la política sobre la marcha. Nuestro componente central, la Optimización Textual Diferenciable (DTO), aprovecha las señales de gradiente tanto de la verosimilitud del LLM como de un modelo de recompensa para refinar las representaciones textuales. nabla-Reasoner incorpora además muestreo por rechazo y un diseño de aceleración para robustecer y agilizar la decodificación. Teóricamente, demostramos que realizar descenso de gradiente en tiempo de inferencia en el espacio de muestras para maximizar la recompensa es dual a alinear una política de un LLM mediante aprendizaje por refuerzo regularizado con divergencia de Kullback-Leibler (KL). Empíricamente, nabla-Reasoner logra una mejora de más del 20% en precisión en un benchmark desafiante de razonamiento matemático, al tiempo que reduce el número de llamadas al modelo en aproximadamente un 10-40% en comparación con líneas de base sólidas. En general, nuestro trabajo introduce un cambio de paradigma desde la búsqueda de orden cero hacia la optimización de primer orden durante la fase de prueba, ofreciendo una ruta rentable para amplificar el razonamiento de los LLMs.
English
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.