Source-linked AI summary
Entropy-Gated Branching for Efficient Test-Time Reasoning
Xianzhi Li, Ethan Callanan, Abdellah Ghassel, Xiaodan Zhu
TL;DR
Test-time reasoning methods can demand substantial computation, including uniform exploration of branches where the model is confident. EGB allocates branching to high-uncertainty steps and uses external feedback to steer search. Across mathematical reasoning benchmarks, the authors report improved accuracy and reduced computational cost compared with other test-time search methods.
Problem
Method
EGB uses predicted-logit entropy to trigger branching at uncertain steps, combines branching with external feedback, and lets confident beams continue with standard sampling.
Source: §1 · p. 2Results
EGB consistently delivers speedups while maintaining superior performance over other test-time search algorithms across mathematical reasoning benchmarks.
Source: §6 · p. 8Takeaways & Limitations
Selective resource allocation during inference achieves improved accuracy with reduced computational cost across the evaluated mathematical reasoning benchmarks.
Source: §6 · p. 8Evaluation focuses on domains with well-established process reward models; extension to domains with less mature verifiers remains future work.
Source: Limitations · p. 9
Abstract
from arXiv · showhide
Test-time compute methods can significantly improve the reasoning capabilities and problem-solving accuracy of large language models (LLMs). However, these approaches require substantially more computational resources, with most compute wasted on exploring low-diversity branches where the model already exhibits high confidence. We observe that a small subset of uncertain reasoning steps has a disproportionately large impact on final prediction accuracy, and branching at these critical junctures tends to yield more diverse and higher-quality candidate reasoning steps. We propose Entropy-Gated Branching (EGB), which branches only at high-uncertainty steps and prunes expansions with a lightweight verifier. On mathematical and financial reasoning benchmarks, EGB improves accuracy by 22.6% over standard inference while operating 31%-75% faster across math benchmarks than test-time beam search with higher performance. Our results show that dynamic resource allocation during inference can substantially improve both efficiency and effectiveness, offering a more scalable pathway to enhanced LLM reasoning capabilities.
1 Introduction
Language models still make reasoning errors that are especially consequential in high-stakes domains such as finance. The introduction motivates selective test-time computation and proposes EGB, which branches at high-uncertainty steps rather than expanding at every step.
1 Introduction
- Reasoning errors persist despite language models’ capabilities and are especially detrimental in high-stakes domains such as finance.The introduction links this concern to the need for outputs that are verifiable and interpretable.Source: §1 · p. 1
- Test-time compute can improve performance without changing model weights, but methods that resample or search branches may require substantial time and computation.The introduction notes that exhaustive test-time search can take minutes, creating a barrier to real-world deployment.Source: §1 · p. 1
- High-entropy spikes coincided with flawed reasoning in 39 of 50 observed cases, motivating branching at uncertain steps instead of every step.The introduction frames selective expansion as a way to improve efficiency while retaining beam search’s performance benefits.Source: §1 · pp. 1–2
- EGB uses predicted-logit entropy to choose branching points and external feedback to steer search, while confident beams continue with standard sampling.The approach concentrates computation on high-impact uncertainty moments to avoid unnecessary overhead.Source: §1 · p. 2
2 Related Work
Related work addresses uncertainty estimation, decoding diversity and test-time computation, while process reward models provide step-wise feedback. EGB combines these threads by using uncertainty to decide when to branch and a PRM to rank candidate branches.
Uncertainty Quantification
- Token-probability methods estimate uncertainty without repeated resampling, and EGB uses token-level entropy to guide branching during generation rather than estimate uncertainty post hoc.Sampling-based approaches can detect inconsistencies but require multiple iterations and cannot assess a single response.Source: §2 · p. 2
Diversity and Efficiency in Decoding
- Standard beam search explores uniformly and can generate redundant alternatives; diverse beam search encourages dissimilarity, while SEGBS still expands at every timestep regardless of confidence.EGB differs by making branching itself uncertainty-aware.Source: §2 · p. 2
Test-Time Computation
- Test-time computation improves performance without retraining, but methods such as repeated sampling, tree search and guided beam search typically apply effort uniformly across generation steps.EGB uses entropy-based gating to identify where additional computation is most beneficial.Source: §2 · pp. 2–3
Process Reward Models and Feedback
- Process reward models assess intermediate reasoning steps, and EGB uses PRM scores to rank candidate branches generated at uncertainty spikes.The cited work characterizes PRMs as evaluating steps rather than only final answers [Lightman et al., 2023; Uesato et al., 2022].Source: §2 · p. 3
3 Methodology
EGB replaces uniform beam expansion with selective branching driven by uncertainty. It rolls back to the earliest high-entropy point, forms a smaller candidate pool, and uses process reward scores to keep the strongest continuations.
3 Methodology
- For each active beam, standard step-wise beam search proposes W continuations, scores K × W candidates with a PRM, and retains the top K, regardless of confidence.When entropy crosses the threshold during generation, EGB rolls back to the earliest such position and branches there.Source: §3 · p. 3, §3.1 · p. 4
3.1 Entropy-Gated Branching
- EGB monitors token-distribution entropy and branches only when uncertainty exceeds a threshold; below it, beams continue with default sampling.The threshold is domain-specific and depends on model calibration and task, so the authors tune it on held-out validation data.Source: §3.1 · p. 3
3.2 Formalizing Selective Expansion
- Selective expansion generates at most K + (W − 1)|U_t| candidates, no more than the K × W produced by uniform beam search.Its candidate-generation complexity is O(|U_t| · W · V), compared with O(K · W · V) for uniform expansion.Source: §3.2 · p. 4
3.3 Beam Evaluation
- EGB removes exact token-level duplicate continuations, scores the remaining candidates with a process reward model, and keeps the top K beams active.The PRM evaluates reasoning-step quality, including logical consistency and mathematical accuracy.Source: §3.3 · pp. 4–5
4 Experiments
The experiments cover mathematical and financial reasoning tasks, multiple Qwen3 and Llama3 model sizes, and test-time reasoning baselines. They use a process reward model for feedback and specify a common branching configuration for the reported Table 1 results.
Datasets
- Evaluation spans AIME, MATH-500, GSM8K and CFA financial-math exams covering two program levels.The CFA setting is included as a domain-specific, high-stakes financial reasoning task.Source: §4 · p. 5
Language Models
- Experiments use three Qwen3 sizes and three Llama3 sizes; Qwen3 is tested in non-thinking mode, with thinking-mode decoding explored separately.Thinking mode produces longer reasoning traces and different efficiency–accuracy trade-offs.Source: §4 · p. 5
Feedback Models
- The primary feedback model is Qwen2.5-Math-PRM-7B, with an alternative PRM used to assess robustness to verifier choice.The PRM assigns numerical scores used to rank branches.Source: §4 · p. 5
Baselines and Settings
- Baselines include self-consistency, self-evaluation-guided stochastic beam search and step-wise beam search with external PRM feedback.Table 1 uses beam size K=4, beam width W=4, and a sample budget of 16 for self-consistency; experiments run on one H100 GPU.Source: §4 · p. 5
5 Results
Across five benchmarks, EGB improves on standard inference while remaining competitive or superior to beam search; it also reduces runtime and memory use.
5 Results
- EGB gains an average 18.4% over standard base inference across models and benchmarks, with absolute accuracy gains of 3–8 percentage points in most settings.Table 1 compares EGB with beam search, SEGBS, self-consistency, and standard decoding on finance and math benchmarks.Source: §5 · p. 5, Table 1 · p. 6
5.1 Main results
- Financial tasks retain robust EGB gains as model scale increases while general-math gains diminish; financial gains fit at approximately twice the rate.The authors suggest this pattern reflects financial problems’ combination of numerical computation, contextual interpretation, and domain-specific procedural knowledge.Source: §5.1 · p. 6
- Llama models gain an average 31.2% over standard decoding with EGB, compared with Qwen models’ 12.8%, despite Qwen’s stronger baseline performance.For example, Qwen3-1.7B scores 48.78% on CFA Level I versus 30.44% for Llama-3.2-1B.Source: §5.1 · pp. 6–7
- EGB’s advantage over beam search grows with model size: 8B models win 70% of comparisons (+5.8% relative), while 1–1.7B models show mixed results.The authors suggest larger models use strategic exploration more effectively by generating and evaluating higher-quality alternatives.Source: §5.1 · p. 7
- At budget=32, EGB reaches 85.71% accuracy, outperforming beam search by 3.02pp, SEGBS by 5.51pp, and self-consistency by 6.47pp.At budget=4, EGB surpasses beam search by 3.0pp after trailing it by 0.8pp at budget=2.Source: §5.1 · p. 7, Figure 3 · p. 8
5.2 Computational Efficiency Analysis
- EGB runs 31–75% faster than beam search across evaluated benchmarks and uses fewer active candidates, reducing memory for intermediate states, attention caches, and token embeddings.The memory savings compound over longer sequences.Source: §5.2 · p. 8
6 Conclusion
EGB allocates inference compute according to model uncertainty, branching at high-uncertainty points. Across mathematical and financial reasoning benchmarks, the authors report improved accuracy alongside reduced computational cost.
6 Conclusion
- EGB dynamically allocates computational resources using model uncertainty, branching only at high-uncertainty decision points.The conclusion frames this as addressing the assumption that every generation step warrants equal computational investment.Source: §6 · p. 8
- Across mathematical and financial reasoning problems, EGB delivers speedups while maintaining superior performance over other test-time search algorithms.The authors report particular effectiveness on moderately challenging tasks where models have relevant knowledge but struggle to find reliable reasoning paths.Source: §6 · p. 8
- The authors present EGB as a pragmatic path to robust, cost-aware test-time reasoning at scale.Their stated conclusion is to focus compute where it matters most.Source: §6 · p. 8
Limitations
The authors identify deployment and evaluation boundaries involving threshold selection, verifier maturity, and the additional cost of thinking-mode decoding.
Limitations
- Automated threshold selection could make deployment easier, although sensitivity analysis found broad safe zones where performance is stable.The limitation is the lack of an automated selection mechanism, not reported instability across thresholds.Source: Limitations · p. 9
- Evaluation focuses on domains with established process reward models; extending EGB to open-ended generation or commonsense reasoning remains future work.Source: Limitations · p. 9
- Combining EGB with thinking-mode decoding costs more because reasoning traces are longer, requiring a latency–accuracy trade-off in resource-sensitive applications.Source: Limitations · p. 9
A.1 Algorithm
EGB allocates extra branches to uncertain decoding steps while confident beams continue with standard sampling, then uses a process reward model to guide candidate selection.
A.1 Algorithm
- EGB branches selectively at uncertain steps while confident beams continue with standard sampling, allocating computation to high-impact uncertainty moments.The algorithm takes a model, problem, beam size K, beam width W, entropy threshold τ, process reward model, and maximum steps as inputs, and returns the best-scoring solution.Source: §A.1 · p. 12
- Candidates are deduplicated, ranked by process reward model scores, and the highest-scoring beam is returned.Figure 6 provides an example PRM response.Source: §A.1 · p. 12
- Final-answer correctness is assessed by Gemini 2.5 Flash using the question, generated response, and ground-truth solution, rather than exact string matching.The judge evaluates semantic correctness to avoid penalizing minor stylistic or formatting variations.Source: §A.1 · pp. 12–13
A.2 Ablation Studies
The ablations examine beam size and width, verifier choice, entropy-threshold sensitivity, and thinking-mode decoding. They report benefits from broader branching and selective thresholds, alongside verifier- and model-dependent trade-offs.
Beam Size K
- 36.2% to 48.8% accuracy for Llama-3.1-8B as beam size increases from K=2 to K=16, with gains plateauing at K=16.The largest gains occur between K=2 and K=8; the authors suggest PRM ranking may favor similar partial solutions at larger beam sizes.Source: §A.2 · p. 13
Beam Width W
- 1–2 percentage-point gains as beam width W increases from 2 to 16, a more modest effect than increasing beam size.The authors interpret this as evidence that maintaining diverse active branches matters more than adding candidates at each individual branching point.Source: §A.2 · p. 13
PRM Quality Sensitivity and Domain Generality
- EGB’s average gain for Qwen3-1.7B is +2.84pp with Llama3.1-8B-PRM and +2.60pp with Qwen2.5-Math-PRM across four benchmarks.For Llama-1B, the average is +1.67pp with Llama3.1-8B-PRM and -0.49pp with Qwen2.5-Math-PRM; verifier choice affects absolute accuracy.Source: §A.2 · pp. 13–14
Entropy-threshold sensitivity
- Accuracy typically varies only 1–3 percentage points when τ shifts by ±0.5–1.0 around a reasonable value, and useful ranges tend to transfer between MATH-500 and CFA-L2.The sweep finds a practical safe zone rather than a sharply tuned optimum; feasible validation tuning can yield additional gains often on the order of 1–3 points.Source: §A.2 · p. 14
Thinking-mode models
- 72.97% accuracy with ~280s average runtime for Qwen3-8B thinking+EGB, versus 61.36% with ~60s for non-thinking+EGB.Thinking mode benefits the larger models under EGB, while Qwen3-1.7B changes from 54.64% to 54.00%.Source: §A.2 · p. 14
A.3 Reduced Uncertainty
An example contrasts entropy-aware branching with argmax decoding, while the accompanying evaluations examine threshold sensitivity and EGB with thinking-mode decoding.
A.3 Reduced Uncertainty
- In one CFA Level II example, entropy-aware branching yields a correct answer while argmax sampling is incorrect; after branching, entropy and varentropy decline, unlike the elevated, fluctuating argmax traces.The figure marks two branching points; the first selects the same token as argmax, while the later branch precedes the decline in both uncertainty measures.Source: §A.3 · pp. 14–15, Figure 5 · p. 16
- Table 3 sweeps τ on CFA-L2 and MATH-500, comparing selective-branching settings with τ=0 (beam search) and τ=∞ (self-consistency).The caption identifies settings that outperform both endpoints for the same model and task, but the supplied passages give no specific table values.Source: Table 3 · p. 15
- Table 4 compares CFA Level II accuracy for Qwen3 models with thinking mode enabled or disabled, including base decoding and EGB gains over beam search and base decoding.The supplied caption identifies the comparisons but does not provide specific table values.Source: Table 4 · p. 15