Flujo de Búsqueda (SoS): Aprendiendo a Buscar en el Lenguaje

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

April 1, 2024

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

cs.AI

Resumen

Los modelos de lenguaje rara vez son expuestos a errores fructíferos durante el entrenamiento. Como resultado, les cuesta mirar más allá del siguiente token, sufren una acumulación de errores y tienen dificultades para predecir las consecuencias de sus acciones varios pasos adelante. En este artículo, demostramos cómo se puede enseñar a los modelos de lenguaje a buscar representando el proceso de búsqueda en lenguaje, como una cadena aplanada: un flujo de búsqueda (Stream of Search, SoS). Proponemos un lenguaje unificado para la búsqueda que captura una variedad de estrategias de búsqueda simbólica. Ilustramos nuestro enfoque utilizando el juego simple pero desafiante de Countdown, donde el objetivo es combinar números de entrada con operaciones aritméticas para alcanzar un número objetivo. Preentrenamos un modelo de lenguaje basado en transformers desde cero en un conjunto de datos de flujos de búsqueda generados por solucionadores heurísticos. Encontramos que el preentrenamiento con SoS aumenta la precisión de búsqueda en un 25 % en comparación con modelos entrenados para predecir solo la trayectoria de búsqueda óptima. Además, ajustamos este modelo con dos métodos de mejora de políticas: Alineación de Políticas Inducida por Ventaja (Advantage-Induced Policy Alignment, APA) y Razonador Autodidacta (Self-Taught Reasoner, STaR). Los modelos SoS ajustados resuelven el 36 % de los problemas previamente no resueltos, incluyendo problemas que no pueden ser resueltos por ninguno de los solucionadores heurísticos. Nuestros resultados indican que los modelos de lenguaje pueden aprender a resolver problemas mediante la búsqueda, auto-mejorarse para usar flexiblemente diferentes estrategias de búsqueda y potencialmente descubrir nuevas.

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.