Do LLMs Encode Functional Importance of Reasoning Tokens?

AI summary

Overview

Research area: Natural Language Processing — efficient and interpretable reasoning in large language models (LLMs), specifically reasoning-chain compression and token-level attribution.

Technical level: Advanced. The paper assumes familiarity with LLM distillation, likelihood objectives, attention mechanics, and perturbation-based attribution.

Scope: The paper introduces "greedy pruning," a likelihood-preserving deletion procedure, and uses it both to compress reasoning chains for distillation and to diagnose whether LLMs internally encode token-level functional importance for answer generation.

What This Paper Is About

LLMs solve hard tasks by generating long reasoning chains, which improves accuracy but increases inference cost, memory use, and makes it hard to tell which parts of the reasoning actually matter for the answer. Existing compression methods (temperature sampling, heuristics, or supervision from frontier models) produce shorter chains but give limited insight into whether the model itself internally ranks reasoning tokens by functional importance. This paper reframes compression as a diagnostic question: can a model's own likelihood and attention signals reveal which reasoning tokens are functionally important for generating the answer?

Key Contributions

1. Greedy pruning. A likelihood-preserving deletion procedure, inspired by perturbation-based attribution and greedy decoding, that iteratively removes the reasoning token whose deletion minimally degrades a specified likelihood objective. This produces both a length-controlled reasoning chain and a monotonic ranking over tokens.

2. Two pruning objectives. An answer-only objective, Ans (maximize log P(A | Q, R_K)), and a joint objective, Joint (maximize log P(R_K, A | Q)), which induce different notions of token importance.

3. Distillation evaluation. A teacher–pruner–student framework in which students are trained on greedily pruned chains and compared against token-level importance baselines (TokenSkip, H2O, Surprisal, Uniform) at matched reasoning lengths.

4. Diagnostic analyses. A token-level functional taxonomy of pruning behavior across six categories, an analysis of how pruning ranks change dynamically as context contracts, and a test of whether attention patterns alone can predict pruning ranks.

Main Findings

• Greedy pruning wins at matched lengths. Across GSM8K, MMLU-Pro, and MATH, students trained on greedily pruned reasoning achieve the strongest performance among pruning-based methods at matched keep fractions, for both Llama3.1-8B and Qwen2.5-7B teacher–pruner pairs. This includes outperforming TokenSkip, which relies on token importance labels from frontier models.

• Graceful degradation under compression. Accuracy improves smoothly as the keep fraction increases. Uniform and Surprisal baselines drop sharply at aggressive pruning, while TokenSkip improves at higher keep fractions but stays below greedy pruning. Zero-shot student performance is consistently worse than all pruning-based distillation methods.

• Stable functional structure emerges. Annotating 1,000 randomly sampled GSM8K examples into six categories (SymbMath, MetaDisc, CoRef, EntName, VerbalMath, Grammar) shows SymbMath tokens (explicit equations and mathematical symbols) are strongly over-retained relative to a uniform-deletion baseline, while CoRef tokens (pronouns and referential expressions) are markedly under-retained until late stages. VerbalMath falls below the diagonal; Grammar, MetaDisc, and EntName sit closer to uniform.

• The pruning objective matters. Under the answer-only (Ans) objective, coarse functional ordering is preserved (SymbMath over-retained, CoRef pruned early), but separation across categories weakens — VerbalMath, MetaDisc, and EntName collapse toward uniform, and even SymbMath and CoRef move closer to uniform. This attenuation aligns with lower student accuracy.

• Pruner strength matters more. Replacing the Qwen2.5-7B pruner with a weaker Llama2-7B pruner keeps SymbMath and CoRef behavior stable, but shifts non-symbolic categories: EntName moves from under-retained to over-retained, Grammar is pruned more aggressively, and VerbalMath approaches uniform. This coincides with a larger drop in student accuracy than the objective ablation.

• Importance is dynamic, not static. Using Hit@|S| alignment, local rankings computed at each pruning stage consistently outperform a Frozen=1.0 baseline (rankings computed at ρ=1.0) and a random baseline, with the gap widening at intermediate keep fractions. Greedy pruning re-evaluates token importance as context contracts rather than exposing a fixed global ordering.

• Attention predicts pruning ranks. A two-layer MLP surrogate trained on only 200 GSM8K examples, using features of dimension [#layers × #heads] built from average attention received from subsequent tokens, achieves a Pearson correlation of 0.88 with true post-deletion likelihoods on a held-out set of 1,000 GSM8K examples (Llama3.1-8B as teacher).

• Compression has modest readability cost. Even at 30% token reduction, degradation of pruned traces is primarily surface-level, with semantic coherence and mathematical state largely preserved. Students trained on pruned data also learn to generate proportionally shorter reasoning traces.

• Zero-shot context. Teacher models: Qwen2.5-7B (GSM8K 90.98, MMLU-Pro 56.25, MATH-500 65.80) and Llama3.1-8B (82.34, 45.21, 43.60). Students: Mistral-7B (31.46, 30.28, 13.40) and Llama2-7B (25.02, 15.58, 4.60).

Methodology in Plain English

The setup has three roles. A teacher LLM generates a full reasoning chain for a question. A pruner LLM scores candidate deletions by likelihood. A student model is fine-tuned on the compressed chains.

Greedy pruning starts from the complete chain and works backwards. At each step, the pruner tries deleting each remaining token, measures how much the specified likelihood objective drops, and removes the token whose deletion hurts least. Repeating this down to a target keep fraction (ρ) produces both a compressed chain and a ranking of tokens by how safe they were to remove — earlier ranks mean the token was less important to the model's likelihood.

Crucially, this is not leave-one-out attribution: because the context shrinks after each deletion, token importance is re-evaluated in a changing context, capturing interactions between tokens. The authors note it is a myopic approximation and does not guarantee optimal subsets; the claim is about the existence and accessibility of internal ranking signals, not subset optimality.

Reasoning chains are generated by the teacher using rejection sampling with temperature 0.7, and unless stated otherwise the pruner equals the teacher with the Joint objective. Students are trained with supervised fine-tuning. Token counts are measured with the student tokenizer so all methods are compared at equal budget.

For the functional analysis, each reasoning token in 1,000 sampled GSM8K examples is annotated with one of six functional roles using gpt-5-mini, validated by manual inspection and stability checks. To test attention as a predictor, a small surrogate network is trained to predict each token's post-deletion likelihood from aggregated attention features.

Why This Matters

Impact on research. The paper reframes reasoning compression from a purely engineering problem into a diagnostic one. It provides evidence that LLMs encode a nontrivial token-level functional structure over their own reasoning, and that this structure is partly recoverable from attention patterns — a bridge between attribution research and practical distillation.

Real-world applications:

• Serving and inference cost. Compressed reasoning chains reduce latency and memory for deployed reasoning models.
• Distillation pipelines. Compact chains from stronger models can be used as cheaper training supervision for smaller models.
• Debugging and auditing. A ranking of functionally important reasoning tokens helps identify whether a model reasoned correctly or got lucky, and which steps carried the answer.
• Context management. Attention-predictable importance rankings hint at inference-time pruning of completed reasoning segments to control context length.

Industry relevance. Reasoning-token cost is a direct operating expense for any product running chain-of-thought models. A likelihood-preserving, offline preprocessing step that shortens chains while outperforming frontier-model-supervised compression at matched lengths is directly applicable to training pipelines, and the attention predictability result suggests a path toward cheaper, online approximations.

Future Directions

1. Reducing computational cost. Greedy pruning scales superlinearly with sequence length because it repeatedly evaluates candidate deletions. The paper points to prefix caching for shared prefixes, removing multiple tokens per iteration given the observed stability of local ranks, and attention-based or learned surrogate approximations.

2. Extending beyond deletion. The authors frame greedy pruning as one instance of a broader class of "reasoning-space operators," suggesting token insertion or replacement under alternative objectives.

3. Inference-time and training-time use. Applying likelihood-preserving objectives to inference-time pruning of completed segments, or embedding pruning signals into training curricula, is described as a natural next step.

4. Open questions. How pruning behaves on incorrect or low-quality reasoning (where likelihood may misalign with functional importance), how the structure varies across model scales and training procedures, and where the threshold lies between compressed-but-readable and uninterpretable reasoning traces.

Target Audience

Researchers and practitioners working on LLM reasoning efficiency, chain-of-thought distillation, and mechanistic interpretability of reasoning. The paper is most useful to readers already comfortable with likelihood objectives, supervised fine-tuning, and attention-based attribution, though the core diagnostic question — does the model itself know which reasoning tokens matter — is accessible to a broader audience interested in how LLMs organize their reasoning.

Authors’ 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.