Fluxo de Busca (SoS): Aprendendo a Buscar na Linguagem

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

Resumo

Modelos de linguagem raramente são expostos a erros frutíferos durante o treinamento. Eles, então, lutam para enxergar além do próximo token, sofrendo com um efeito bola de neve de erros e tendo dificuldade para prever as consequências de suas ações várias etapas à frente. Neste artigo, mostramos como modelos de linguagem podem ser ensinados a buscar, representando o processo de busca em linguagem, como uma string achatada — um fluxo de busca (Stream of Search, SoS). Propomos uma linguagem unificada para busca que captura uma variedade de diferentes estratégias de busca simbólica. Demonstramos nossa abordagem usando o jogo simples, porém desafiador, de Countdown, onde o objetivo é combinar números de entrada com operações aritméticas para alcançar um número alvo. Pré-treinamos um modelo de linguagem baseado em transformer do zero em um conjunto de dados de fluxos de busca gerados por solucionadores heurísticos. Descobrimos que o pré-treinamento com SoS aumenta a precisão da busca em 25% em comparação com modelos treinados para prever apenas a trajetória de busca ideal. Ajustamos ainda mais esse modelo com dois métodos de melhoria de política: Alinhamento de Política Induzido por Vantagem (Advantage-Induced Policy Alignment, APA) e Raciocinador Autoensinado (Self-Taught Reasoner, STaR). Os modelos SoS ajustados resolvem 36% dos problemas anteriormente não resolvidos, incluindo problemas que não podem ser resolvidos por nenhum dos solucionadores heurísticos. Nossos resultados indicam que modelos de linguagem podem aprender a resolver problemas por meio de busca, autoaperfeiçoar-se para usar estrategicamente diferentes estratégias de busca e, potencialmente, descobrir novas.

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.