Stream of Search (SoS) : Apprendre à rechercher dans le langage

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

April 1, 2024

Auteurs: Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, Noah D. Goodman

cs.AI

Résumé

Les modèles de langage sont rarement exposés à des erreurs productives pendant leur entraînement. Ils ont ensuite du mal à voir au-delà du prochain jeton, souffrant d'un effet boule de neige des erreurs et peinant à prédire les conséquences de leurs actions plusieurs étapes à l'avance. Dans cet article, nous montrons comment les modèles de langage peuvent être enseignés à effectuer des recherches en représentant le processus de recherche dans le langage, sous forme de chaîne aplatie - un flux de recherche (Stream of Search, SoS). Nous proposons un langage unifié pour la recherche qui capture un éventail de différentes stratégies de recherche symboliques. Nous démontrons notre approche en utilisant le jeu simple mais difficile de Countdown, où l'objectif est de combiner des nombres d'entrée avec des opérations arithmétiques pour atteindre un nombre cible. Nous pré-entraînons un modèle de langage basé sur les transformateurs à partir de zéro sur un ensemble de données de flux de recherche générés par des solveurs heuristiques. Nous constatons que le pré-entraînement SoS augmente la précision de la recherche de 25 % par rapport aux modèles entraînés à prédire uniquement la trajectoire de recherche optimale. Nous affinons ensuite ce modèle avec deux méthodes d'amélioration de la politique : l'Alignement de Politique Induit par l'Avantage (Advantage-Induced Policy Alignment, APA) et le Raisonneur Auto-Enseigné (Self-Taught Reasoner, STaR). Les modèles SoS affinés résolvent 36 % des problèmes précédemment non résolus, y compris des problèmes qui ne peuvent être résolus par aucun des solveurs heuristiques. Nos résultats indiquent que les modèles de langage peuvent apprendre à résoudre des problèmes via la recherche, s'améliorer de manière autonome pour utiliser de manière flexible différentes stratégies de recherche, et potentiellement en découvrir de nouvelles.

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.