A Training Data Recipe to Accelerate
A* Search with Large Language Models

A* is a tree-based search algorithm that aims to find a path between a start node and any goal node by building a tree $\mathcal{T}$. The algorithm is presented as Algorithm 1. The set of all tree nodes is denoted as $\mathcal{N}$. For each node $n$, A* search keeps track of two values, (i) historical cost $g(n)$, which is the distance between the start node and $n$ and (ii) heuristic $h(n)$ which is an estimate of the true distance $h^{*}(n)$ between $n$ and the closest goal node. Each node may be associated with a state $s(n)$. An action modifies the state, causing a transition to a new node. For the search, A* maintains two lists, the frontier list $P_{\text{frt}}$ and the closed list $P_{\text{cld}}$. At the beginning, the closed list is empty and the frontier list is initialised with the start node. The search is terminated when either a goal state is encountered, or $P_{\text{frt}}$ is empty. Each iteration performs two steps, described below.

Our goal is train a language model $\theta$, that, given a node $n$, can predict the residual $d^{*}(n)=h^{*}(n)-h(n)$ between the perfect heuristic $h^{*}(n)$ and a quick estimate $h(n)$. Given a series of similar problem instances, we derive training data from their A* search trees after a search is complete. For each tree node $n$, computing the ground-truth $d^{*}(n)$ would require running A* starting from node $n$, which quickly becomes prohibitively expensive as the problem size grows. Following Chrestien et al. (2021); ús Virseda et al. (2013), we only consider nodes on the optimal path. After the first A* run, their $h^{*}(\cdot)$ is trivial: for any node $n_{j}\in\pi^{*}$, $h^{*}(n_{j})=|\pi^{*}|-j$. Formally, the training sequences $\mathcal{X}$ are given by $\mathcal{X}=\bigcup_{\pi^{*}_{i}\sim\Gamma_{i=0}^{N}}\{(n_{j},d^{*}(n_{j})),n_{j}\in\pi^{*}_{i}\}$.

We train the LLM with the L2 loss

where $f_{\theta}$ represents a forward pass of the LLM. We use encoder-decoder transformers and add a regression head $\phi_{L2}$ on the decoder that predicts $d^{*}(n)$ given the $\langle BoS\rangle$ token as the input.

Additionally, since the LLM can be trained in a text-to-text setting, we train a separate model with the canonical autoregressive loss, given by:

With $\mathcal{L}_{LM}$, the pre-trained language model head $\phi_{LM}$ is used.

Inference involves leveraging the trained LLM in A* search. During the expansion step, children nodes to be evaluated are converted into an LLM prompt, from which the LLM predicts $d(n)$. This value is added to the quick estimate of $h(n)$. Notably, only a single forward pass is performed per expansion as we collate all children nodes as one batch. Additionally, we cache these prompts, such that if a state is revisited in another node $m$, $d(m)$ can simply be retrieved.

For $\theta$ trained with $\mathcal{L}_{LM}$, we perform top-k decoding, with $k=5$²²2This value was arbitrarily chosen and fixed for all experiments. It allows the LLM to make additional choices, without straying too much from the greedy one, along with self-consistency Wang et al. (2022), predicting 3 sequences, as this works slightly better in practice.

The exact prompt inputs for the encoder have been provided in Section A.2.

We conduct our experiments on three problems domains. Each domain comprises of the in-distribution (IID) and out-of-distribution (OOD) test sets for a total of six datasets.

We adopt several metrics defined by Lehnert et al. (2024), (i) inverse-length-ratio (ILR) to measure the differences in the search length, (ii) success weighted by cost (SWC) to measure the differences in solution length and (iii) optimal %, to measure the percentage of problems solved optimally. ILR measures the average inverse ratio between the search length $\tilde{\mathcal{S}}$ of an A* solution, to the optimal reference $\mathcal{S}^{*}$. It is computed as

ILR can be averaged over various sets. ILR-on-solved is averaged over all puzzles in the test set and ILR-on-optimal is averaged over all puzzles whose solutions are optimal. Suboptimal solutions, found with inadmissible heuristics, are often discovered before optimal ones, leading to a lower $\mathcal{S}$, but a higher ILR; due to this, ILR-on-optimal allows us to measure the informativeness of the heuristic on equal, minimum length solutions.

SWC measures the average inverse ratio between the start-to-goal path length $|\tilde{\pi}|$ of an A* solution, to that of an optimal reference, denoted by $|\pi^{*}|$.

To measure computational cost, we propose a new metric, inverse time ratio, which is defined as the average inverse ratio between the wall-clock time of an A* solution $\tilde{WT}$ and a reference solution $WT^{*}$,

Prediction errors by the LLM in the learned heuristic function are inevitable. In this section, we aim to examine two research questions: (i) how the prediction errors in the learned heuristic function affects the search length $\mathcal{S}$, and (ii) how they affect optimality of the solutions.

Specifically, we start with the oracle heuristic $h^{*}(n)$ and artificially introduce error in different sections of the search trajectory in order to observe effects on $\mathcal{S}$ and optimality. The search tree is divided into three sets—initial, middle and end. A node $n$ is placed in the initial set if its cost places itself in the first third of the optimal path: $g(n)<|\pi^{*}|/3$. Alternatively, it may be placed in the middle set if $|\pi^{*}|/3\leq g(n)<2|\pi^{*}|/3$, and in the end set if $g(n)\geq 2|\pi^{*}|/3$. We introduce zero-mean Gaussian error by drawing a random value from $\mathcal{N}(0,\sigma)$ and adding it to $h^{*}(n)$. In each experiment, we introduce errors in two of three sections and use the oracle in one section. We use maze as the domain of experiment and obtain the oracle heuristic $h^{*}(\cdot)$ by running Dijkstra’s algorithm on the maze, starting from the goal.

In this section, we explore how training on pairs of (node, distance-to-goal) affects the generalisation of the heuristic-learning LLM. We create four training splits by uniformly sampling nodes on the optimal path from the Initial, Middle, and End sections of the path. The All set contains nodes uniformly sampled from all three sections. Additionally, we also create exclusion sets, which excludes one of the three sections, and these sets are denoted as $\sim$Initial, $\sim$Middle and $\sim$End. For instance, $\sim$Initial contains only data sampled from the Middle and End sets. All training splits have the same size.

We adopt the following evaluation metrics: (i) mean absolute error (MAE) on validation splits containing nodes from each of the aforementioned splits, and (ii) ILR achieved by applying the trained models as heuristic functions for A*. While (i) directly evaluates the generalisation of the model, (ii) provides a more realistic test of how well the trained model works with A*.

Each training split contains 12k and 8k nodes in total for maze and Sokoban, respectively. All models are initialized with code-t5-small. Hyperparameter details are mentioned in Section A.3.