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

arXiv: https://arxiv.org/abs/2407.09985 · HTML: https://mlsanity.com/paper/2407.09985/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 26 of its 49 references that are in the map.
- First posted on arXiv 13 Jul 2024.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 26 of 26)

- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [Language Agent Tree Search Unifies Reasoning Acting and Planning in Language Models](https://mlsanity.com/paper/2310.04406/) (2023)
- [Deep Learning on a Data Diet: Finding Important Examples Early in Training](https://mlsanity.com/paper/2107.07075/) (2021)
- [Beyond neural scaling laws: beating power law scaling via data pruning](https://mlsanity.com/paper/2206.14486/) (2022)
- [LLM+P: Empowering Large Language Models with Optimal Planning Proficiency](https://mlsanity.com/paper/2304.11477/) (2023)
- [On the Planning Abilities of Large Language Models : A Critical Investigation](https://mlsanity.com/paper/2305.15771/) (2023)
- [SemDeDup: Data-efficient learning at web-scale through semantic deduplication](https://mlsanity.com/paper/2303.09540/) (2023)
- [Reasoning or Reciting? Exploring the Capabilities and Limitations of Language Models Through Counterfactual Tasks](https://mlsanity.com/paper/2307.02477/) (2023)
- [Differentiation of Blackbox Combinatorial Solvers](https://mlsanity.com/paper/1912.02175/) (2019)
- [Leveraging Pre-trained Large Language Models to Construct and Utilize World Models for Model-based Task Planning](https://mlsanity.com/paper/2305.14909/) (2023)
- [LLMs Can't Plan, But Can Help Planning in LLM-Modulo Frameworks](https://mlsanity.com/paper/2402.01817/) (2024)
- [Distilling System 2 into System 1](https://mlsanity.com/paper/2407.06023/) (2024)
- [Stream of Search (SoS): Learning to Search in Language](https://mlsanity.com/paper/2404.03683/) (2024)
- [Beyond A*: Better Planning with Transformers via Search Dynamics Bootstrapping](https://mlsanity.com/paper/2402.14083/) (2024)
- [An investigation of model-free planning](https://mlsanity.com/paper/1901.03559/) (2019)
- [Coupling Large Language Models with Logic Programming for Robust and General Reasoning from Text](https://mlsanity.com/paper/2307.07696/) (2023)
- [Path Planning using Neural A* Search](https://mlsanity.com/paper/2009.07476/) (2020)
- [When is Tree Search Useful for LLM Planning? It Depends on the Discriminator](https://mlsanity.com/paper/2402.10890/) (2024)
- [Learning Generalized Reactive Policies using Deep Neural Networks](https://mlsanity.com/paper/1708.07280/) (2017)
- [Learning Heuristic Selection with Dynamic Algorithm Configuration](https://mlsanity.com/paper/2006.08246/) (2020)
- [Policy-Guided Heuristic Search with Guarantees](https://mlsanity.com/paper/2103.11505/) (2021)
- [TransPath: Learning Heuristics For Grid-Based Pathfinding via Transformers](https://mlsanity.com/paper/2212.11730/) (2022)
- [What Are We Measuring When We Evaluate Large Vision-Language Models? An Analysis of Latent Factors and Biases](https://mlsanity.com/paper/2404.02415/) (2024)
- [Levin Tree Search with Context Models](https://mlsanity.com/paper/2305.16945/) (2023)
- [Heuristic Search Planning with Deep Neural Networks using Imitation, Attention and Curriculum Learning](https://mlsanity.com/paper/2112.01918/) (2021)

## Built on by (most cited first, 1 of 1)

- [Constraint-Aware Route Recommendation from Natural Language via Hierarchical LLM Agents](https://mlsanity.com/paper/2510.06078/) (2025)

## Abstract (arXiv)

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 LLM-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 codebase is at https://github.com/devaansh100/a_star.