YC Root AccessImproving Small Language Model Reasoning With A* Search
CHAPTERS
- 0:07 – 0:45
Why small language models need better reasoning at test time
Alex frames the core problem: small language models (SLMs) are attractive for cost and latency, but they lag in multi-step reasoning. He introduces the guiding question of how to add high-quality reasoning under tight compute constraints.
- •SLMs are preferred in many products due to lower cost/latency
- •SLMs often lack the reasoning depth of large LLMs
- •Research goal: improve reasoning without upgrading model size
- •Focus is explicitly on test-time improvements
- 0:45 – 1:30
Limitations of prior approaches: distillation and MCTS-style search
He outlines two common solution paths and why they can be impractical. Distillation depends on access to a large teacher and expensive training, while tree-search methods often require external reward models for evaluation.
- •Distillation transfers capability from a large teacher to a smaller student
- •Distillation requires teacher access + training compute/time
- •MCTS/tree search samples multiple reasoning paths
- •Evaluation often needs external reward models that are hard to train/host
- 1:30 – 1:56
Test-time A* search: reasoning as a tree search problem
Alex introduces their method: use an A*-inspired tree search at inference time to reach the correct answer node. The key shift is replacing external evaluators with the model’s own self-critique as a heuristic.
- •Model reasoning is treated as navigating a search tree
- •Goal is to find the destination/correct answer node
- •A* search structure guides which nodes to expand
- •Self-critique is used as the heuristic instead of an external reward model
- 1:56 – 2:25
What “self-critique” means and how it becomes a heuristic
He defines self-critique as the model scoring the quality/correctness of its own proposed answer. That score is then incorporated into the search cost to prioritize promising branches.
- •Ask the same model to rate an answer (e.g., 0–100)
- •High-quality answers should receive higher self-scores
- •Incorrect/hallucinated answers should be scored lower
- •Self-critique feeds into the A*-style value/cost function
- 2:25 – 2:52
Algorithm loop: expand nodes, compute cost, pick lowest-cost frontier
Alex walks through the high-level mechanics: start at a root response, iteratively expand selected nodes into children, score them, and always choose the lowest-cost node to explore next. The process stops once a solution node is reached and returned.
- •Start from the question and generate a root node
- •Iteratively generate children for a selected node
- •Compute an A*-inspired cost using critique-informed scoring
- •Select the lowest-cost node for continued exploration
- •Stop when a solution/correct node is found and output it
- 2:52 – 2:55
Intuition: self-critique helps prune bad branches early
He explains why this works: when two nodes are at similar depth, self-critique can differentiate them and steer exploration away from low-quality or hallucinated paths. This effectively prunes the tree and concentrates compute on better reasoning trajectories.
- •Depth being equal makes critique the key differentiator
- •Low self-critique signals likely hallucination or wrong direction
- •Search deprioritizes high-cost/low-confidence subtrees
- •Pruning improves efficiency and solution quality
- 2:55 – 3:43
Walkthrough on a math word problem: branching + critique-guided traversal
Using a concrete example, Alex describes generating a root answer, branching into two children at each step via prompting, and injecting self-critique into the prompts. The method traverses toward the correct solution path, shown as the bottom-left branch in the visualization.
- •Root node is produced by a standard LLM call
- •Each step expands into two branches (children)
- •Prompts include both reasoning context and self-critique instructions
- •Traversal follows the critique-weighted cost to the correct leaf
- 3:43 – 4:20
Benchmark results on GSM8K and MATH500: accuracy gains from A* self-critique search
Alex reports experimental results on math reasoning benchmarks, highlighting that their method (shown in red) improves accuracy over other test-time scaling approaches. He clarifies GSM8K as a high-school-level math reasoning dataset and notes comparisons are done using the same base model.
- •Evaluated on GSM8K and MATH500 reasoning benchmarks
- •Method outperforms other test-time scaling algorithms in accuracy
- •GSM8K: high-school-level style math word problems
- •Comparisons use the same underlying model (Qwen 3 4B)
- 4:20 – 4:58
Efficiency ablation: accuracy vs tokens and wall-clock time
Because test-time scaling increases compute, Alex presents an ablation measuring gains at fixed token and time budgets. He claims their method delivers ~3–4 percentage point improvements at comparable token/time usage, suggesting the gains aren’t merely from spending more compute.
- •Test-time scaling raises concerns about extra compute/time
- •Ablation compares performance at matched token budgets
- •Ablation compares performance at matched time budgets
- •Reported +3–4 points at similar tokens and ~+3 points at similar time
- 4:58 – 5:42
Ongoing direction: causal reinforcement learning for better credit assignment in agents
Shifting to new work, Alex describes interest in RL for LLM reasoning—specifically causal RL—arguing correlation-based learning struggles with credit assignment in complex agent systems. He motivates counterfactual reasoning to better identify which actions truly caused outcomes.
- •New work separate from the A* test-time scaling paper
- •Motivation: RL credit assignment is hard for LLM agents with many moving parts
- •Causal RL aims to move from correlation to cause-and-effect understanding
- •Counterfactual reasoning proposed to isolate what decisions caused outcomes
- 5:42 – 6:11
Wrap-up: related prior papers and contact info
Alex briefly references two earlier papers—one systems-focused and one on hallucinations—then shares his website and LinkedIn. The host closes the session with thanks.
- •Mentions prior systems paper and hallucination paper
- •Provides website and LinkedIn for follow-up
- •Session concludes with acknowledgments
