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

TL;DR

This paper introduces KAPPA, a new inference-time pruning method for large language models that uses a principled scoring function to reduce computational costs while maintaining reasoning accuracy.

Contribution

KAPPA combines KL divergence, confidence, and entropy into a scoring function for effective branch pruning during inference, improving efficiency over existing heuristics.

Findings

KAPPA stabilizes performance in smaller models.

Achieves up to 60% memory reduction.

Reduces total token generation by 90%.

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.