Поток поиска (SoS): Обучение поиску в языке

Stream of Search (SoS): Learning to Search in Language

April 1, 2024

Авторы: Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, Noah D. Goodman

cs.AI

Аннотация

Языковые модели редко показывают плодотворные ошибки во время обучения. Затем они испытывают трудности в просмотре за пределами следующего токена, страдая от нарастания ошибок и затрудняясь предсказать последствия своих действий на несколько шагов вперед. В данной статье мы показываем, как языковые модели могут быть обучены к поиску, представляя процесс поиска на языке как сплошную строку - поток поиска (SoS). Мы предлагаем объединенный язык для поиска, который охватывает ряд различных символьных стратегий поиска. Мы демонстрируем наш подход, используя простую, но сложную игру Countdown, где целью является комбинирование входных чисел с арифметическими операциями для достижения целевого числа. Мы предварительно обучаем языковую модель на основе трансформера с нуля на наборе данных потоков поиска, сгенерированных эвристическими решателями. Мы обнаруживаем, что предварительное обучение SoS увеличивает точность поиска на 25% по сравнению с моделями, обученными предсказывать только оптимальную траекторию поиска. Затем мы донастраиваем эту модель двумя методами улучшения политики: Advantage-Induced Policy Alignment (APA) и Self-Taught Reasoner (STaR). Доработанные модели SoS решают 36% ранее нерешенных проблем, включая проблемы, которые не могут быть решены ни одним из эвристических решателей. Наши результаты показывают, что языковые модели могут научиться решать проблемы с помощью поиска, самосовершенствоваться для гибкого использования различных стратегий поиска и, возможно, открывать новые.

English

Language models are rarely shown fruitful mistakes while training. They then struggle to look beyond the next token, suffering from a snowballing of errors and struggling to predict the consequence of their actions several steps ahead. In this paper, we show how language models can be taught to search by representing the process of search in language, as a flattened string -- a stream of search (SoS). We propose a unified language for search that captures an array of different symbolic search strategies. We demonstrate our approach using the simple yet difficult game of Countdown, where the goal is to combine input numbers with arithmetic operations to reach a target number. We pretrain a transformer-based language model from scratch on a dataset of streams of search generated by heuristic solvers. We find that SoS pretraining increases search accuracy by 25% over models trained to predict only the optimal search trajectory. We further finetune this model with two policy improvement methods: Advantage-Induced Policy Alignment (APA) and Self-Taught Reasoner (STaR). The finetuned SoS models solve 36% of previously unsolved problems, including problems that cannot be solved by any of the heuristic solvers. Our results indicate that language models can learn to solve problems via search, self-improve to flexibly use different search strategies, and potentially discover new ones.