Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A*-Decoding: Token-Efficient Inference Scaling

Giannis Chatziveroglou

TL;DR

This work addresses token-inefficient inference in large language models by reframing autoregressive decoding as a state-space search guided by the A* algorithm. It introduces A*-decoding, which uses a process-supervised heuristic from external signals (PRMs) and a budget-aware expansion bound to prioritize promising partial trajectories, achieving high exact-match accuracy with substantially fewer tokens and PRM passes. Empirical results on math benchmarks (MATH500, AIME 2024) show that small to mid-size models can match or exceed the performance of much larger models, demonstrating strong token efficiency and scalable reasoning. The approach offers a practical, interpretable alternative to brute-force sampling and scale-driven gains, enabling more efficient deployment of reasoning-enabled language models.

Abstract

Inference-time scaling has emerged as a powerful alternative to parameter scaling for improving language model performance on complex reasoning tasks. While existing methods have shown strong performance gains under fixed compute budgets, there has been little focus on optimally utilizing that budget during inference. In this work, we introduce A*-decoding, a search-based inference-time strategy that builds on the A* search algorithm to optimally utilize a fixed compute budget by prioritizing high-quality reasoning paths during generation. We frame language model decoding as a structured search in a state space of partial solutions, applying the A* transition model to identify promising continuations guided by an external process supervision signal. In our experiments, A*-decoding reaches the performance levels of strong inference scaling baselines like best-of-N and particle filtering while using up to 3x fewer tokens and 30% fewer PRM passes under equivalent compute budgets. On the MATH500 and AIME 2024 benchmarks, A*-decoding enables Llama-3.2-1B-Instruct to match the performance of the 70x larger Llama-3.1-70B-Instruct, and allows Qwen3-1.7B to reach o1-like reasoning accuracy. These results highlight the power of structured search in decoding, offering an alternative to brute-force sampling or scale-driven gains. Our work demonstrates how thoughtful inference-time strategies can enhance reasoning in SLMs, pointing toward future advances in more efficient and scalable language model deployment.

A*-Decoding: Token-Efficient Inference Scaling

TL;DR

This work addresses token-inefficient inference in large language models by reframing autoregressive decoding as a state-space search guided by the A* algorithm. It introduces A*-decoding, which uses a process-supervised heuristic from external signals (PRMs) and a budget-aware expansion bound to prioritize promising partial trajectories, achieving high exact-match accuracy with substantially fewer tokens and PRM passes. Empirical results on math benchmarks (MATH500, AIME 2024) show that small to mid-size models can match or exceed the performance of much larger models, demonstrating strong token efficiency and scalable reasoning. The approach offers a practical, interpretable alternative to brute-force sampling and scale-driven gains, enabling more efficient deployment of reasoning-enabled language models.

Abstract

Inference-time scaling has emerged as a powerful alternative to parameter scaling for improving language model performance on complex reasoning tasks. While existing methods have shown strong performance gains under fixed compute budgets, there has been little focus on optimally utilizing that budget during inference. In this work, we introduce A*-decoding, a search-based inference-time strategy that builds on the A* search algorithm to optimally utilize a fixed compute budget by prioritizing high-quality reasoning paths during generation. We frame language model decoding as a structured search in a state space of partial solutions, applying the A* transition model to identify promising continuations guided by an external process supervision signal. In our experiments, A*-decoding reaches the performance levels of strong inference scaling baselines like best-of-N and particle filtering while using up to 3x fewer tokens and 30% fewer PRM passes under equivalent compute budgets. On the MATH500 and AIME 2024 benchmarks, A*-decoding enables Llama-3.2-1B-Instruct to match the performance of the 70x larger Llama-3.1-70B-Instruct, and allows Qwen3-1.7B to reach o1-like reasoning accuracy. These results highlight the power of structured search in decoding, offering an alternative to brute-force sampling or scale-driven gains. Our work demonstrates how thoughtful inference-time strategies can enhance reasoning in SLMs, pointing toward future advances in more efficient and scalable language model deployment.
Paper Structure (20 sections, 3 equations, 4 figures, 1 table, 1 algorithm)

This paper contains 20 sections, 3 equations, 4 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology: Inference-Time Scaling with A*
  4. A* Search
  5. Language Model Inference as State Space Search
  6. A process supervision definition for the A* heuristic
  7. Bounding Search Expansion for Tractable Inference
  8. Experiments
  9. Experimental setup
  10. Main results
  11. Scaling inference-time exploration
  12. Ablations
  13. Discussion
  14. Appendix
  15. Language Model Inference as State Space Search (Extended)
  16. ...and 5 more sections

Figures (4)

  • Figure 1: Top: Results with Llama models. Bottom: Results with Qwen models. Left: Token cost-performance frontier on MATH500. A*-decoding (ours) consistently achieves higher or equal accuracy compared to alternatives with lower token usage. Right: PRM cost-performance frontier on MATH500. A*-decoding optimally allocates PRM inference budget by focusing compute on promising partial trajectories. Points show results for 4–64 sampled generations; A*-decoding uses up to 16 for Llama, and 32 for Qwen.
  • Figure 2: Illustration of A*-decoding applied to a math question involving function composition. At each step, the decoder expands candidate states, computes heuristic values, and selects the most promising path according to the A* total estimated cost. The final solution path, highlighted in green from the lowest green state up to the root, shows how A*-decoding efficiently navigates the solution space to find a correct answer with optimal computation.
  • Figure 3: Left: Llama-3.2-1B-Instruct performance on a 100-problem subset of MATH500 with different sampling temperature settings (0.2, 0.4, 0.6, 0.8, 1.0). Right: Llama-3.2-1B-Instruct performance on a 100-problem subset of MATH500 with process supervision signals from different models (Qwen2.5-Math-PRM-7B, Skywork-o1-Open-PRM-Qwen-2.5-1.5B).
  • Figure 4: Llama-3.2-1B-Instruct performance on MATH500 with different max breadth $b_{max}$ settings (5, 10, 20). Setiting $b_{max} = 5$ achieves the highest token-efficiency and overall accuracy compared to other settings.