Stream of Search (SoS): Leren zoeken in taal

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

Samenvatting

Taalmodellen krijgen tijdens de training zelden vruchtbare fouten te zien. Hierdoor hebben ze moeite om verder te kijken dan het volgende token, lijden ze onder een opeenstapeling van fouten en vinden ze het moeilijk om de gevolgen van hun acties meerdere stappen vooruit te voorspellen. In dit artikel laten we zien hoe taalmodellen kunnen worden geleerd om te zoeken door het zoekproces in taal te representeren, als een afgevlakte string – een stroom van zoekacties (Stream of Search, SoS). We stellen een uniforme taal voor zoeken voor die een reeks verschillende symbolische zoekstrategieën omvat. We demonstreren onze aanpak met het eenvoudige maar uitdagende spel Countdown, waarbij het doel is om invoergetallen met rekenkundige bewerkingen te combineren om een doelgetal te bereiken. We pretrainen een transformer-gebaseerd taalmodel vanaf nul op een dataset van zoekstromen die door heuristische oplossers zijn gegenereerd. We ontdekken dat SoS-pretraining de zoeknauwkeurigheid met 25% verhoogt ten opzichte van modellen die alleen zijn getraind om de optimale zoektrajectorie te voorspellen. We finetunen dit model verder met twee beleidsverbeteringsmethoden: Advantage-Induced Policy Alignment (APA) en Self-Taught Reasoner (STaR). De gefinetunde SoS-modellen lossen 36% van de voorheen onopgeloste problemen op, inclusief problemen die door geen van de heuristische oplossers kunnen worden opgelost. Onze resultaten geven aan dat taalmodellen kunnen leren om problemen op te lossen via zoeken, zichzelf kunnen verbeteren om flexibel verschillende zoekstrategieën te gebruiken en mogelijk nieuwe kunnen ontdekken.

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.