nabla-Reasoner: Raciocínio em LLM via Descida de Gradiente no Espaço Latente em Tempo de Teste
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
Resumo
A escalagem do poder computacional durante a inferência para Modelos de Linguagem de Grande Porte (LLMs) desbloqueou capacidades de raciocínio sem precedentes. No entanto, os métodos de escalagem de inferência existentes geralmente dependem de algoritmos de busca discreta ineficientes e subóptimos ou de *prompts* de tentativa e erro para melhorar a política *online*. Neste artigo, propomos o nabla-Reasoner, uma estrutura de geração iterativa que integra otimização diferenciável sobre os *logits* dos *tokens* no ciclo de decodificação para refinar a política em tempo real. Nosso componente central, a Otimização Textual Diferenciável (DTO), aproveita sinais de gradiente tanto da verossimilhança do LLM quanto de um modelo de recompensa para refinar as representações textuais. O nabla-Reasoner incorpora ainda amostragem por rejeição e um projeto de aceleração para robustecer e acelerar a decodificação. Teoricamente, demonstramos que realizar gradiente descendente no espaço amostral durante a inferência para maximizar a recompensa é dual ao alinhamento de uma política de LLM via aprendizado por reforço regularizado por KL. Empiricamente, o nabla-Reasoner alcança uma melhoria de mais de 20% na precisão em um benchmark desafiador de raciocínio matemático, enquanto reduz o número de chamadas do modelo em aproximadamente 10-40% em comparação com bases de comparação fortes. No geral, nosso trabalho introduz uma mudança de paradigma da busca de ordem zero para a otimização de primeira ordem durante o teste, oferecendo um caminho economicamente viável para amplificar o raciocínio de 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.