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

2025-11-01

AI Summary

To address the trade-off between computational overhead and inference accuracy in large language models, this paper proposes KAPPA: a dynamic branch scoring framework for fine-grained chain-of-thought (CoT) path pruning during inference. KAPPA computes scores based on KL divergence, confidence, and entropy—leveraging latent information signals rather than heuristic consistency—as guidance for progressive, quality-aware pruning. This enables more accurate branch evaluation and significantly improves inference stability, especially for smaller models. By integrating diversity-aware exploration with adaptive pruning, KAPPA achieves up to 60% peak memory reduction and 90% reduction in total generated tokens on GSM8K and MATH500, while incurring negligible accuracy degradation (<0.5%).

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Calibration & Uncertainty Quantification
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Reduces computational cost of multi-branch reasoning in LLMs
• Improves pruning by evaluating branch quality with scoring metrics
• Maintains accuracy while cutting memory and token usage significantly

Innovation

• Uses KL divergence for branch pruning
• Combines confidence and entropy in scoring
• Reduces memory and token usage significantly