À la recherche dans l'obscurité : Raisonnement via un gradient de politique au niveau des instances en espace latent pendant le test

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

May 19, 2025

Auteurs: 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

Résumé

La capacité de raisonnement, composante essentielle de l'intelligence humaine, continue de représenter un défi majeur pour les modèles de langage à grande échelle (LLMs) dans la quête de l'intelligence artificielle générale (AGI). Bien que les performances des modèles se soient améliorées grâce à la loi d'échelle de l'entraînement, des défis importants subsistent, notamment en ce qui concerne les algorithmes d'entraînement, tels que l'oubli catastrophique, et la disponibilité limitée de nouvelles données d'entraînement. Comme alternative, l'échelle au moment du test améliore les performances de raisonnement en augmentant le calcul au moment du test sans mise à jour des paramètres. Contrairement aux méthodes précédentes dans ce paradigme axées sur l'espace des tokens, nous proposons d'exploiter l'espace latent pour un raisonnement plus efficace et une meilleure adhérence à la loi d'échelle au moment du test. Nous introduisons LatentSeek, un nouveau cadre qui améliore le raisonnement des LLMs grâce à l'Adaptation au Niveau de l'Instance au Moment du Test (TTIA) dans l'espace latent du modèle. Plus précisément, LatentSeek utilise le gradient de politique pour mettre à jour itérativement les représentations latentes, guidé par des signaux de récompense auto-générés. LatentSeek est évalué sur une série de benchmarks de raisonnement, notamment GSM8K, MATH-500 et AIME2024, à travers plusieurs architectures de LLMs. Les résultats montrent que LatentSeek surpasse systématiquement des bases de référence solides, telles que l'incitation en chaîne de pensée et les méthodes basées sur le fine-tuning. De plus, notre analyse démontre que LatentSeek est très efficace, convergeant généralement en quelques itérations pour des problèmes de complexité moyenne, tout en bénéficiant d'itérations supplémentaires, mettant ainsi en lumière le potentiel de l'échelle au moment du test dans l'espace latent. Ces résultats positionnent LatentSeek comme une solution légère, évolutive et efficace pour améliorer les capacités de raisonnement des 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.