Learning to Reason Efficiently with A* Post-Training
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.
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