Поиск в темноте: рассуждения через градиент политики на уровне экземпляров во время тестирования в латентном пространстве

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

May 19, 2025

Авторы: 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

Аннотация

Способность к рассуждению, являющаяся ключевым компонентом человеческого интеллекта, продолжает оставаться серьёзным вызовом для крупных языковых моделей (LLMs) на пути к созданию искусственного общего интеллекта (AGI). Несмотря на улучшение производительности моделей в соответствии с законом масштабирования обучения, сохраняются значительные трудности, особенно в отношении алгоритмов обучения, таких как катастрофическое забывание, и ограниченной доступности новых обучающих данных. В качестве альтернативы, масштабирование на этапе тестирования улучшает способность к рассуждению за счёт увеличения вычислительных ресурсов на этапе тестирования без обновления параметров. В отличие от предыдущих методов в этой парадигме, сосредоточенных на пространстве токенов, мы предлагаем использовать латентное пространство для более эффективного рассуждения и лучшего соответствия закону масштабирования на этапе тестирования. Мы представляем LatentSeek — новый фреймворк, который улучшает способность LLM к рассуждению через адаптацию на уровне экземпляров на этапе тестирования (TTIA) в латентном пространстве модели. В частности, LatentSeek использует градиент политики для итеративного обновления латентных представлений, руководствуясь самостоятельно генерируемыми сигналами вознаграждения. LatentSeek оценивается на ряде бенчмарков для рассуждений, включая GSM8K, MATH-500 и AIME2024, на различных архитектурах LLM. Результаты показывают, что LatentSeek стабильно превосходит сильные базовые методы, такие как prompting с цепочкой рассуждений (Chain-of-Thought) и методы, основанные на тонкой настройке. Кроме того, наш анализ демонстрирует, что LatentSeek обладает высокой эффективностью, обычно сходясь за несколько итераций для задач средней сложности, при этом также выигрывая от дополнительных итераций, что подчеркивает потенциал масштабирования на этапе тестирования в латентном пространстве. Эти результаты позиционируют LatentSeek как лёгкое, масштабируемое и эффективное решение для улучшения способности LLM к рассуждению.

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.