MiGrATe: Mixed-Policy GRPO for Adaptation at Test-Time

2508.08641v1 cs.LG, cs.AI, cs.CL 2025-08-14
Авторы:

Peter Phan, Dhruv Agarwal, Kavitha Srinivas, Horst Samulowitz, Pavan Kapanipathi, Andrew McCallum

Резюме на русском

#### Контекст Много платформенных языковых моделей (LLMs) становятся важной частью решения задач смешанного типа, таких как программная синтезирование, оптимизация молекул и поиск абстрактных решений. Однако применение этих моделей к задачам поиска многоугольников часто сталкивается с проблемами баланса между исследованием новых решений и использованием высокооплачиваемых решений. Традиционные подходы, такие как in-context learning, либо неэффективны в этих задачах, либо требуют большого количества ресурсов для поиска индивидуально настроенных решений. Дополнительно, некоторые методы, такие как test-time training (TTT), широко используются для улучшения решений, но часто требуют ручной подготовки специальной данных, что снижает их выгоду в процессе интеграции. Наша мотивация заключается в создании метода, который использует LLMs для решения сложных задач без необходимости подготовки данных заранее. #### Метод Мы предлагаем метод **MiGrATe** (Mixed-Policy GRPO for Adaptation at Test-Time), который работает как гибкий поисковый алгоритм, используя LLMs в качестве ядра поиска. Метод MiGrATe оперирует с помощью процедуры сборки групп микшированной политики (mixed-policy group construction procedure), которая сочетает в себе два типа сбора данных: greedy sampling (выбор лучших решений из прошлых итераций) и neighborhood sampling (поиск решений, структурно похожих на высокооплачиваемые). Эти две техники дают возможность применить больше внимания к высокооплачиваемым решениям и поддерживать эксплорaцию новых решений при использовании LLMs. Метод MiGrATe также включает в себя применение GRPO (Group Policy Optimization) для оптимизации групповой политики в реальном времени. #### Результаты Мы провели эксперименты на трех реальных задачах: поиск слов, оптимизация молекул и индукция программ+абстракций на датасетом ARC. Метод MiGrATe показал значительное улучшение качества решений в сравнении с традиционными TTT-методами и методами, не использовавшими тестируемую данную. Благодаря использованию смешанных политик в поиске, MiGrATe удалось эффективно оптимизировать решения без необходимости дополнительной подготовки данных, что демонстрирует мощь TTT-подходов для решения задач поиска без внешнего наблюдения. #### Значимость Метод MiGrATe может быть применен в сложных сетевых задачах, таких как оптимизация программ, поиск решений в лабораториях и молекул, а также в области моделирования абстрактных решений. Он предлагает значительные преимущества, включая более эффективное применение LLMs для решения задач без ручной подготовки данных, а также гибкость применения к разнообразным задачам. Метод MiGrATe может иметь будущее в широком спектре приложений, включая лаборатории, где требуется эффективное

Abstract

Large language models (LLMs) are increasingly being applied to black-box optimization tasks, from program synthesis to molecule design. Prior work typically leverages in-context learning to iteratively guide the model towards better solutions. Such methods, however, often struggle to balance exploration of new solution spaces with exploitation of high-reward ones. Recently, test-time training (TTT) with synthetic data has shown promise in improving solution quality. However, the need for hand-crafted training data tailored to each task limits feasibility and scalability across domains. To address this problem, we introduce MiGrATe-a method for online TTT that uses GRPO as a search algorithm to adapt LLMs at inference without requiring external training data. MiGrATe operates via a mixed-policy group construction procedure that combines on-policy sampling with two off-policy data selection techniques: greedy sampling, which selects top-performing past completions, and neighborhood sampling (NS), which generates completions structurally similar to high-reward ones. Together, these components bias the policy gradient towards exploitation of promising regions in solution space, while preserving exploration through on-policy sampling. We evaluate MiGrATe on three challenging domains-word search, molecule optimization, and hypothesis+program induction on the Abstraction and Reasoning Corpus (ARC)-and find that it consistently outperforms both inference-only and TTT baselines, demonstrating the potential of online TTT as a solution for complex search tasks without external supervision.

Ссылки и действия