Inference-Time Chain-of-Thought Pruning with Latent Informativeness Signals

AI summary

Overview

Research area: Inference-time scaling for large language model reasoning, specifically efficient chain-of-thought (CoT) decoding and branch pruning.

Technical level: Intermediate. Readers need familiarity with sampling-based decoding (Best-of-N, temperature/top-k/top-p), Kullback–Leibler (KL) divergence, entropy, and confidence metrics, but the high-level idea is accessible.

Scope: The paper introduces a training-free, inference-time pruning algorithm called KAPPA (KL-Adjusted Pruned Path Algorithm) that scores parallel reasoning branches using uncertainty signals and progressively eliminates the weakest ones, evaluated on two math benchmarks with two open-source models.

What This Paper Is About

Large language models can reason more accurately when they generate many candidate solutions and pick the best one, but standard Best-of-N (BoN) sampling pays for this by fully generating every branch, which is expensive in memory and tokens. Prior work such as Self-Truncation Best-of-N (ST-BoN) cuts branches early, but decides which to keep using consistency heuristics rather than any direct measure of branch quality. This paper proposes KAPPA, which scores branches at every decoding step using KL divergence from an unconditional reference distribution, token confidence, and entropy, then prunes the lowest scorers on a linear schedule so only one branch is generated to completion.

Key Contributions

1. A principled scoring function for branch quality. KAPPA combines three signals — an EMA-smoothed information-change term derived from KL divergence against an unconditional reference distribution, token confidence (max softmax probability), and entropy — into a single weighted score, avoiding external reward models and requiring no training.

2. A progressive, linear pruning schedule. Branches are eliminated one at a time according to the schedule R_t = N − floor(((t − c + 1)N)/τ), leaving exactly one survivor after the horizon τ, which is then decoded to completion.

3. Robust signal processing. Information-change signals are stabilized with a median-of-means (MoM) estimate over a sliding window of size w split into m buckets, then smoothed with a bias-corrected exponential moving average at rate α.

4. Empirical evaluation of the memory/accuracy tradeoff. Experiments on GSM8K and MATH500 with DeepSeek-R1-Distill-Qwen-1.5B and Qwen2.5-7B-Instruct compare KAPPA against Full-BoN (primary baseline) and ST-BoN across sampling sizes N = 5, 10, and 20.

Main Findings

• Reduced peak memory: Peak GPU memory falls by roughly 4% to 60% relative to BoN across datasets and models. The largest reported gap is DeepSeek-R1-Distill-Qwen-1.5B on MATH500 at N = 20, where KAPPA uses 6495.25 MB against BoN's 16239.977 MB.

• Reduced token generation: The results section reports total token reductions of roughly 65% to 90% relative to BoN. The maximum reported gap is again DeepSeek-R1-Distill-Qwen-1.5B on MATH500 at N = 20: 2113.162 tokens for KAPPA versus 20053.28 for BoN. The introduction states a larger figure of "up to 97.3%" token reduction, while the abstract states "up to 90%"; these three figures as written in the paper do not agree with one another.

• Accuracy gains on the smaller model: On DeepSeek-R1-Distill-Qwen-1.5B, KAPPA improves accuracy by 1–2% across both GSM8K and MATH500 relative to Full-BoN. In the appendix, for example, BoN on DeepSeek GSM8K scores 0.705, 0.697, and 0.704 at N = 5, 10, 20, while KAPPA scores 0.716, 0.713, and 0.707.

• Mixed results on the larger model: On Qwen2.5-7B-Instruct, the authors do not observe the same consistent improvement as baseline performance strengthens. The paper attributes this to over-pruning: for larger models the lowest-scoring branch early in the scoring phase is often higher quality than the weakest branches of smaller models, so a linear schedule can eliminate promising branches prematurely.

• Accuracy can degrade as N grows: For some model/dataset combinations, KAPPA's accuracy decreases as the sampling size N increases, which the authors also suspect is caused by over-pruning.

• Headline accuracy figure: The introduction states KAPPA achieves 72.2% on MATH500, which corresponds in the appendix table to Qwen2.5-7B-Instruct at N = 5 (0.722).

• Reported wall-clock time is variable: Appendix A lists per-configuration times, but the main text does not analyze them; in some configurations (for example DeepSeek-R1-Distill-Qwen-1.5B on MATH500 at N = 20) BoN's reported time of 14.49 s is lower than KAPPA's 24.899 s, while in others KAPPA is faster (Qwen2.5-7B-Instruct GSM8K at N = 5: KAPPA 6.345 s versus BoN 15.661 s).

Methodology in Plain English

KAPPA runs in three phases.

Draft (exploration). The model generates N candidate branches in parallel up to a cutoff timestep c, defined, following ST-BoN, as the earliest point at which all branches are pairwise inconsistent. This preserves diversity before any pruning happens.

Scoring and gating. From step c onward, each surviving branch is scored at every step. The core signal is how much a branch's next-token distribution p_t^i differs from a reference distribution q obtained by generating unconditional logits from the beginning-of-sentence token. The KL divergence D_t^i = D_KL(p_t^i ‖ q) measures this, and the change from the previous step, ΔI_t^i, is treated as an information-gain signal. To reduce noise and outlier sensitivity, this change is robustified with a median-of-means estimate over the last w steps divided into m buckets, then smoothed with a bias-corrected EMA at rate α. Two further signals — confidence (the maximum softmax probability) and entropy — are computed from the same distribution. All three signals are converted to z-scores across the currently alive branches at each timestep and clipped to the range [−3, 3]. A weighted sum s_t^i = w_KL · EMA_t^i + w_C · C_t^i + w_H · H_t^i becomes the instantaneous score, and a trajectory score weights recent steps more heavily (weights proportional to the timestep), since later steps are more diagnostic of the final answer. Signals are read from intermediate logits projected onto the vocabulary space. At each step, the lowest-scoring branches are pruned so the number alive follows the linear schedule down to one.

Continuation (exploitation). The single surviving branch is decoded to completion (until an end-of-sequence token or the maximum length) and its answer is returned.

Experimental setup. Models are DeepSeek-R1-Distill-Qwen-1.5B and Qwen2.5-7B-Instruct; datasets are GSM8K and MATH500. Sampling uses top-k = 20, top-p = 0.95, temperature = 0.7, max new tokens = 1024, run on 80G A100 GPUs. Baselines use HuggingFace model.generate(), with Full-BoN selecting answers by negative perplexity. KAPPA hyperparameters are EMA rate α = 0.5, window w = 16, MoM buckets m = 4, and weights (w_KL, w_C, w_H) = (0.7, 0.2, 0.1), chosen by grid search on a subset of the data. Prompts instruct the model to reason step by step and put the final answer in \boxed{}. Evaluation uses exact-match accuracy (N_m / N_t) and a memory cost ratio M_peak / M_peak^greedy against greedy decoding.

Why This Matters

Impact on research. The paper tests whether an intrinsic, training-free uncertainty signal — the KL divergence between a branch's predictions and an unconditional reference — can substitute for consistency heuristics or external reward models in deciding when to cut a reasoning path. It also isolates a clean failure mode for aggressive pruning: smaller models benefit, while larger models with higher-quality branches can be hurt, which is a directly useful result for anyone designing inference-time compute-allocation schemes.

Real-world applications:

• Deploying reasoning models under tight GPU memory budgets, where holding many live branches simultaneously is the binding constraint.
• Reducing serving cost for multi-sample reasoning pipelines (for example, math tutoring, code generation, or agentic workflows) where per-query token spend dominates.
• Edge or on-premise inference with small distilled reasoning models, the regime where the paper reports 1–2% accuracy gains alongside memory savings.
• Batch evaluation or dataset labeling with reasoning models, where total tokens generated is the principal cost.

Industry relevance. Serving LLMs with Best-of-N or self-consistency multiplies both memory and generate-time cost by the sampling factor N. An inference-time pruning layer that reports roughly 4–60% lower peak memory and roughly 65–90% fewer tokens (as stated in the results section) at comparable accuracy is directly relevant to inference providers and to anyone paying per-token, subject to the caveat that the paper's headline reduction figures vary between the abstract, introduction, and results.

Future Directions

• Dynamic pruning schedules. The paper suggests less aggressive schedules such as a cosine schedule to avoid eliminating promising branches early, especially for larger models where low-scoring branches remain high quality.
• Adaptive horizon and draft length. Varying the pruning horizon τ or extending the draft phase based on problem complexity is proposed as a way to trade accuracy against cost more intelligently, rather than fixing τ.
• Broader generalization testing. The authors note that the evaluation covers only specific model scales and datasets, and call for experiments on other models and on commonsense reasoning datasets.
• Hyperparameter robustness. Lowering the EMA rate α, increasing the MoM window and bucket sizes, and raising the confidence and entropy weights are suggested to reduce sensitivity to noisy signals. The paper also notes that 1024 max new tokens was a compute-imposed cap that specifically affected the MATH500 experiments on DeepSeek, since that model often needed more tokens to reach a final answer.

Target Audience

Researchers and engineers working on inference-time scaling, test-time compute allocation, and efficient LLM decoding — particularly those already familiar with Best-of-N, self-consistency, and ST-BoN. It is also relevant to practitioners deploying reasoning models under memory or token-budget constraints, and to readers interested in using intrinsic model uncertainty (KL divergence, entropy, confidence) as a self-supervised signal rather than relying on trained reward models.

Authors’ abstract

Large language models (LLMs) improve reasoning accuracy when generating multiple candidate solutions at test time, but standard methods like Best-of-N (BoN) incur high computational cost by fully generating all branches. Self-Truncation Best-of-N (ST-BoN) mitigates this by truncating unpromising paths early, but its reliance on consistency-based heuristics is a limitation as it does not directly evaluate branch quality. We present KL-Adjusted Pruned Path Algorithm (KAPPA), an inference-time method that combines Kullback-Leibler divergence, confidence, and entropy into a principled scoring function to guide progressive pruning. By promoting diversity during exploration and selectively eliminating low-scoring branches, KAPPA maintains accuracy while substantially reducing memory and token usage. Experiments on GSM8K and MATH500 with DeepSeek-R1-Distill-Qwen-1.5B and Qwen2.5-7B-Instruct demonstrate that KAPPA stabilizes performance in smaller models and achieves up to ~60% reduction in peak memory and ~90% reduction in total token generation relative to BoN, with minimal impact on accuracy.