English

Learning to Reason Efficiently with A* Post-Training

Artificial Intelligence 2026-05-26 v1 Computation and Language Machine Learning

Abstract

Many applications of large language models (LLMs) require deductive reasoning, yet models frequently produce incorrect or redundant inference steps. We frame natural language inference as a search problem where the final answer is the valid proof itself, requiring a reasoning procedure in which intermediate inferences are correct. Specifically, we investigate whether LLMs can learn to generate correct and efficient proofs with guidance from A* search -- an algorithm that guarantees an optimally efficient path to a goal. We explore two training techniques: supervised fine-tuning on execution traces from A* and reinforcement learning with A*-informed process reward models. Empirically, we find that Llama-3.2 models in the 1B--3B range benefit substantially from A* post training, going from near-zero accuracy to outperforming DeepSeek-V3.2 -- a much larger model. Our analysis uncovers a trade-off: while simple correctness rewards maximize accuracy, A*-informed signals strike a balance between accuracy and efficiency. Furthermore, we find that on larger search spaces, models trained with imperfect heuristics exhibit superior accuracy. Our results demonstrate a promising direction towards reasoning guided by principles derived from classical search algorithms.

Keywords

Cite

@article{arxiv.2605.24597,
  title  = {Learning to Reason Efficiently with A* Post-Training},
  author = {Andreas Opedal and Francesco Ignazio Re and Abulhair Saparov and Mrinmaya Sachan and Bernhard Schölkopf and Ryan Cotterell},
  journal= {arXiv preprint arXiv:2605.24597},
  year   = {2026}
}

Comments

Preprint