A Training Data Recipe to Accelerate A* Search with Language Models (2024)

Abstract: Combining Large Language Models (LLMs) with heuristic search algorithms like A* holds the promise of enhanced LLM reasoning and scalable inference. To accelerate training and reduce computational demands, we investigate the coreset selection problem for the training data of LLM heuristic learning. Few methods to learn the heuristic functions consider the interaction between the search algorithm and the machine learning model. In this work, we empirically disentangle the requirements of A* search algorithm from the requirements of the LLM to generalise on this task. Surprisingly, we find an overlap between their requirements; A* requires more accurate predictions on search nodes near the goal, and LLMs need the same set of nodes for effective generalisation. With these insights, we derive a data-selection distribution for learning LM-based heuristics. On three classical planning domains, maze navigation, Sokoban and sliding tile puzzles, our technique reduces the number of iterations required to find the solutions by up to 15x, with a wall-clock speed-up of search up to 5x. The code has been made available at https://github.com/devaansh100/a_star.

## introduction contrary to the view that large language models (llms) serve as a monolithic paradigm for intelligence, the dual-process theory of cognitive science @xcite posits that human cognition consists of two closely collaborating systems, system 1 and system 2. system 1 exhibits typical traits of statistical learning such as fast inference and slow adaptation to novel problems. in comparison, system 2 can solve novel problems and excels at logical reasoning, but performs slow inference. recent analyses analogize llms to system 1 @xcite , as llms perform poorly at novel, out-of-distribution problem formulations @xcite or problems that require planning and reasoning @xcite @xcite . on the other hand, treesearch methods like a* @xcite ) and variants (e.g., @xcite , provide classic solutions to logical reasoning and planning, but they are unable to learn from past experiences and limited in speed due to sequential dependencies. though it has been speculated that artificial general intelligence requires both system 1 and system 2 capabilities @xcite , how to fruitfully combining llms with search techniques remains an open problem. we study the problem of using llms to learn a* heuristics, which are functions that estimate the distance from a search node to the goal state. however, it can be computational demanding to train llms and to generate training data, as groundtruth labels for training can only be obtained from successfully solved problems. with this paper, we aim to improve the efficiency of heuristic learning by selecting a small subset of training data, known as the coreset, which would lead to near-identical a* performance as the whole dataset. to the best of our knowledge, no previous work investigated the coreset selection problem for a* heuristic learning. a complication of coreset selection in the a* + llm setup is that the two algorithms may impose different requirements on training data. in this work, we attempt to disentangle and individually quantify the requirements of the two algorithms. we empirically test how different training data would change the generalization of the llm, and how a* reacts to generalization errors in different positions of the search trajectory. we divide the training trajectory into three equally sized portions: the beginning, the middle, and the end. first, we evaluate their effectiveness as training data. this is inspired by research us-ing training data difficulty as a metric for coreset selection @xcite . a natural definition for difficulty in a* is the distance to goal, which indicates how many decisions must be made before reaching the goal. intuitively, it should be more difficult to guess the exact distance to goal at a given search node if the search node is in fact farther away from the goal. further, to simulate the effect of llm noise on a*, we inject random errors into oracle heuristic values in the three portions and observe effects on the search length. we obtain interesting and unexpected findings. for the llm, training on the last portion, where the search node is closest to the goal and the distances are easiest to fit, leads to the best generalization among the three portions. unexpectedly, a* demonstrates a similar behavior; erroneous predictions on the end portion are the most detrimental to search efficiency, even though one might expect earlier decisions to be more important in pruning search nodes. these observations suggest that we should prioritize training data from the last portion, which would lead to overall good llm generalization and best accuracy on the end portion, which in turn accelerates search. accordingly, we devise a planner-aware sampling strategy for training data, which prioritizes search nodes near the end. in addition, this sampling strategy is general enough to be combined with other coreset selection methods. the proposed strategy incurs, on average, 9.5% fewer a* search steps than uninformed baselines and, in some cases, outperforms models trained with double the amount of data. our contributions can be summarized as follows, 1. to the best of our knowledge, we are the first work to study the coreset selection problem for a* heuristic learning. further, we propose a mathematical criterion to select training data based on their distance to goal. ## related works we review several research directions related to our work. for a tabular summary of the works, see subsection a.5. ## proposed solution the utility of a node in accelerating search we propose to quantify the utility of a node in reducing the search length as, c(•) assigns higher values to nodes closer to the goal. while there can be nodes with g(n) ≥ |π * |, since they are never added to the tree, c(•) is not defined for them. considerations and other choices for c(•) are discussed in subsection a.4. ## conclusion in this work, we study the training data requirements to learn a strong heuristic for a* search. we find that accurate prediction of heuristics for nodes close to the goal are the most important for a* speed. similarly, generalization of the llm heuristic requires training on nodes near the goal. based on these insights, we propose a mathematical model to select search nodes as training data. this results in substantially reduced search lengths and significant wall-clock speedups on hard problems. our study lays the groundwork for bootstrapped heuristic learning, which learns heuristic functions for increasingly larger problems using solved problems of smaller sizes. referred to as the data flywheel, such techniques hold promise to scale up the capabilities of llm + tree search foot_2 .