Do LLMs Encode Functional Importance of Reasoning Tokens?

TL;DR

The paper investigates whether LLMs encode token-level functional importance in their reasoning and introduces greedy pruning, a likelihood-preserving deletion method that iteratively removes tokens to produce a ranking of functional importance and length-controlled reasoning. It formalizes two pruning objectives, $L^{Ans}_{\theta_{\mathcal{P}}}$ and $L^{Joint}_{\theta_{\mathcal{P}}}$, and evaluates the approach via a teacher–pruner–student distillation framework across GSM8K, MMLU-Pro, and MATH, showing that pruned reasoning preserves answer-relevant information and enhances distillation at matched lengths. Across benchmarks, students trained on greedily pruned reasoning outperform baselines, indicating that important tokens for answer generation are retained while less critical content is pruned. Analyses reveal a stable functional structure with symbolic computation preferentially preserved and non-symbolic content pruned earlier, dynamic re-evaluation of pruning ranks as context contracts, and strong predictability of pruning ranks from attention signals, highlighting internal cues of token importance. Overall, greedy pruning provides a principled diagnostic tool for uncovering token-level functional organization in model-generated reasoning and informs efficient reasoning-space exploration.

Abstract

Large language models solve complex tasks by generating long reasoning chains, achieving higher accuracy at the cost of increased computational cost and reduced ability to isolate functionally relevant reasoning. Prior work on compact reasoning shortens such chains through probabilistic sampling, heuristics, or supervision from frontier models, but offers limited insight into whether models internally encode token-level functional importance for answer generation. We address this gap diagnostically and propose greedy pruning, a likelihood-preserving deletion procedure that iteratively removes reasoning tokens whose removal minimally degrades model likelihood under a specified objective, yielding length-controlled reasoning chains. We evaluate pruned reasoning in a distillation framework and show that students trained on pruned chains outperform a frontier-model-supervised compression baseline at matched reasoning lengths. Finally, our analysis reveals systematic pruning patterns and shows that attention scores can predict greedy pruning ranks, further suggesting that models encode a nontrivial functional importance structure over reasoning tokens.

Figures (8)

• Figure 1: Greedy pruning as a diagnostic probe.A. A teacher model generates a full reasoning chain for a given question. B. A greedy pruning step scores candidate token deletions by post-deletion likelihood $L^{del}$ and removes the token whose deletion best preserves likelihood. C. Iterating this procedure over decreasing keep fractions $\rho$ yields length-controlled chains and induces a pruning order $\pi$, where earlier-ranked tokens are safer to prune. D. This order enables analysis of compression–quality tradeoffs and the structure of encoded functional importance.
• Figure 2: Distillation under reasoning token pruning. Accuracy of a Llama2-7B student trained on pruned reasoning at varying keep fractions, teacher, pruner, and dataset; dashed lines indicate zero-shot performance. Greedy pruning achieves the strongest performance at matched lengths, indicating preservation of important tokens.
• Figure 3: Functional structure under greedy pruning. Each curve shows the fraction of tokens retained per category at a given keep fraction. Panels vary teacher, pruner, and pruning objective; the dashed line indicates uniform pruning. (a) Pruning preferentially preserves symbolic computation while removing referential, descriptive, and linguistic scaffolding. (b) Excluding reasoning likelihood in pruning objective softens induced functional structure. (c) A weaker pruner preserves symbolic computation but disrupts the balance of non-symbolic structure.
• Figure 4: Effect of pruning objective and pruner strength on distillation. The base greedy setting uses $\mathcal{T},\mathcal{P}=\text{Qwen2.5-7B}$ and $\mathcal{S}=\text{Llama2-7B}$ with the Joint objective. We ablate pruner strength ($\mathcal{P}=\text{Llama2-7B}$) and pruning objective (answer-only, Ans). Both weaken student performance, with pruner strength having the larger effect.
• Figure 5: Dynamics of pruning ranks.Hit@|S| alignment between tokens removed at keep fraction $\rho_{\text{curr}}$ and $L^{del}$-based local ranks at the previous pruning stage. Dynamic rankings ($\rho_{\text{prev}}=\rho_{\text{curr}}+0.1$) consistently outperform frozen ($\rho_{\text{prev}}{=}1.0$) and random baselines across keep fractions, indicating that greedy pruning re-evaluates token importance as context contracts.