A* 搜索让小模型推理能力大幅跃升,做推理优化或小模型部署的团队值得关注,可以直接参考其训练方法。
研究者将自然语言推理建模为搜索问题,利用 A* 搜索算法指导 LLM 生成正确且高效的推理步骤。通过监督微调(基于 A* 执行轨迹)和强化学习(结合 A* 过程奖励模型),Llama-3.2 1B-3B 模型从近乎零准确率提升至超越 DeepSeek-V3.2。研究发现,简单正确性奖励最大化准确率,而 A* 信号能平衡准确率与效率。在更大搜索空间下,基于不完美启发式的训练反而带来更优准确率。这项工作展示了经典搜索算法指导 LLM 推理的潜力。
Learning to Reason Efficiently with A* Post-Training
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.