nabla-Reasoner : Raisonnement par LLM via une descente de gradient en phase de test dans l'espace latent

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

March 5, 2026

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

cs.AI

Résumé

L'augmentation de la capacité de calcul lors de l'inférence pour les grands modèles de langage (LLM) a débloqué des capacités de raisonnement sans précédent. Cependant, les méthodes existantes de mise à l'échelle à l'inférence reposent généralement sur des algorithmes de recherche discrets inefficaces et sous-optimaux ou sur des sollicitations par essais et erreurs pour améliorer la politique en ligne. Dans cet article, nous proposons nabla-Reasoner, un cadre de génération itératif qui intègre une optimisation différentiable sur les logits des tokens dans la boucle de décodage pour affiner la politique à la volée. Notre composant principal, l'Optimisation Textuelle Différentiable (DTO), exploite les signaux de gradient provenant à la fois de la vraisemblance du LLM et d'un modèle de récompense pour affiner les représentations textuelles. nabla-Reasoner intègre en outre un échantillonnage de rejet et une conception accélérée pour robustifier et accélérer le décodage. Théoriquement, nous montrons que effectuer une descente de gradient à l'inférence dans l'espace des échantillons pour maximiser la récompense est dual à l'alignement d'une politique de LLM via un apprentissage par renforcement régularisé par la divergence KL. Empiriquement, nabla-Reasoner permet une amélioration de la précision de plus de 20 % sur un benchmark de raisonnement mathématique difficile, tout en réduisant le nombre d'appels au modèle d'environ 10 à 40 % par rapport à des bases de référence solides. Dans l'ensemble, notre travail introduit un changement de paradigme, passant d'une recherche d'ordre zéro à une optimisation de premier ordre au moment du test, offrant ainsi une voie rentable pour amplifier le raisonnement des LLM.

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.