Buscar en la Oscuridad: Razonamiento mediante Gradiente de Política a Nivel de Instancia en Tiempo de Prueba en el Espacio Latente

Seek in the Dark: Reasoning via Test-Time Instance-Level Policy Gradient in Latent Space

May 19, 2025

Autores: Hengli Li, Chenxi Li, Tong Wu, Xuekai Zhu, Yuxuan Wang, Zhaoxin Yu, Eric Hanchen Jiang, Song-Chun Zhu, Zixia Jia, Ying Nian Wu, Zilong Zheng

cs.AI

Resumen

La capacidad de razonamiento, un componente fundamental de la inteligencia humana, sigue representando un desafío significativo para los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) en la búsqueda de la Inteligencia General Artificial (AGI). Aunque el rendimiento de los modelos ha mejorado bajo la ley de escalamiento del entrenamiento, persisten desafíos importantes, particularmente en relación con los algoritmos de entrenamiento, como el olvido catastrófico, y la disponibilidad limitada de datos de entrenamiento novedosos. Como alternativa, el escalamiento en tiempo de prueba mejora el rendimiento del razonamiento al incrementar la computación durante la prueba sin actualizar los parámetros. A diferencia de métodos previos en este paradigma centrados en el espacio de tokens, proponemos aprovechar el espacio latente para un razonamiento más efectivo y una mejor adherencia a la ley de escalamiento en tiempo de prueba. Presentamos LatentSeek, un marco novedoso que mejora el razonamiento de los LLMs mediante la Adaptación a Nivel de Instancia en Tiempo de Prueba (TTIA, por sus siglas en inglés) dentro del espacio latente del modelo. Específicamente, LatentSeek utiliza el gradiente de política para actualizar iterativamente las representaciones latentes, guiado por señales de recompensa autogeneradas. LatentSeek es evaluado en una variedad de benchmarks de razonamiento, incluyendo GSM8K, MATH-500 y AIME2024, en múltiples arquitecturas de LLMs. Los resultados muestran que LatentSeek supera consistentemente a métodos de referencia sólidos, como el prompting de Cadena de Pensamiento y métodos basados en ajuste fino. Además, nuestro análisis demuestra que LatentSeek es altamente eficiente, convergiendo típicamente en pocas iteraciones para problemas de complejidad media, mientras que también se beneficia de iteraciones adicionales, destacando así el potencial del escalamiento en tiempo de prueba en el espacio latente. Estos hallazgos posicionan a LatentSeek como una solución ligera, escalable y efectiva para mejorar las capacidades de razonamiento de los LLMs.

English

Reasoning ability, a core component of human intelligence, continues to pose a significant challenge for Large Language Models (LLMs) in the pursuit of AGI. Although model performance has improved under the training scaling law, significant challenges remain, particularly with respect to training algorithms, such as catastrophic forgetting, and the limited availability of novel training data. As an alternative, test-time scaling enhances reasoning performance by increasing test-time computation without parameter updating. Unlike prior methods in this paradigm focused on token space, we propose leveraging latent space for more effective reasoning and better adherence to the test-time scaling law. We introduce LatentSeek, a novel framework that enhances LLM reasoning through Test-Time Instance-level Adaptation (TTIA) within the model's latent space. Specifically, LatentSeek leverages policy gradient to iteratively update latent representations, guided by self-generated reward signals. LatentSeek is evaluated on a range of reasoning benchmarks, including GSM8K, MATH-500, and AIME2024, across multiple LLM architectures. Results show that LatentSeek consistently outperforms strong baselines, such as Chain-of-Thought prompting and fine-tuning-based methods. Furthermore, our analysis demonstrates that LatentSeek is highly efficient, typically converging within a few iterations for problems of average complexity, while also benefiting from additional iterations, thereby highlighting the potential of test-time scaling in the latent space. These findings position LatentSeek as a lightweight, scalable, and effective solution for enhancing the reasoning capabilities of LLMs.