Do LLMs Encode Functional Importance of Reasoning Tokens?

1. Problem Statement

Large Language Models (LLMs) solve complex tasks (e.g., math, science) by generating long chains of thought (reasoning). While effective, these chains incur high computational costs (latency, memory) and make it difficult to isolate which specific tokens are functionally responsible for the final answer.

Existing approaches to "compact reasoning" (shortening chains) rely on:

• External exploration: Sampling multiple chains and selecting shorter ones.
• Heuristics: Using probability thresholds or post-processing rules.
• Frontier model supervision: Using stronger models (e.g., GPT-4) to label or generate compact chains.

The Gap: These methods offer limited insight into whether the reasoning model internally encodes a token-level functional importance structure. It remains unclear if a model's own likelihood and attention mechanisms can identify which reasoning tokens are critical for answer generation without external supervision.

2. Methodology: Greedy Pruning

The authors propose Greedy Pruning, a diagnostic procedure that treats reasoning compression as a likelihood-preserving deletion task.

Core Mechanism

Instead of generating new chains, Greedy Pruning starts with a complete reasoning chain generated by a "Teacher" model and iteratively removes tokens.

1. Objective: The goal is to remove a token such that the degradation of the model's likelihood (under a specific objective) is minimized.
2. Process:
   • Start with the full chain $R$.
   • For every token $i$ in the current chain, calculate the post-deletion likelihood $L_{del}^i$ (the likelihood of the remaining sequence + answer).
   • Identify the token $i^*$ whose removal results in the highest remaining likelihood (i.e., the least damage to the model's predictive behavior).
   • Remove $i^*$ and record its "pruning rank."
   • Repeat until the desired keep fraction ($\rho$) is reached.
3. Pruning Objectives:
   • $L_{JOINT}$: Preserves the likelihood of both the reasoning chain and the answer ($\log P(R_K, A | Q)$). This is more restrictive and preserves the reasoning trajectory.
   • $L_{ANS}$: Preserves only the likelihood of the answer ($\log P(A | Q, R_K)$). This allows for more aggressive pruning of intermediate steps.

Framework

The authors utilize a Teacher-Pruner-Student distillation framework:

• Teacher: Generates the initial reasoning chain.
• Pruner: Applies Greedy Pruning to create a ranked list of tokens and compressed chains.
• Student: A smaller model trained via Supervised Fine-Tuning (SFT) on the pruned chains to see if it can learn the compressed reasoning effectively.

3. Key Contributions

1. Diagnostic Tool: Introduced Greedy Pruning as a method to probe internal token-level functional importance without external labels.
2. Distillation Superiority: Demonstrated that students trained on greedily pruned chains outperform students trained on chains pruned by frontier-model-supervised baselines (like TokenSkip) at matched lengths.
3. Functional Structure Discovery: Revealed that LLMs encode a non-trivial functional structure where symbolic computation (equations, numbers) is prioritized for retention, while referential (pronouns), descriptive, and grammatical scaffolding are pruned first.
4. Dynamic Importance: Showed that token importance is dynamic, not static. The importance of a token changes as the context contracts, meaning early pruning decisions cannot predict later ones perfectly.
5. Attention Correlation: Proved that pruning ranks are highly predictable (Pearson correlation $\approx 0.88$) using only attention scores from the teacher model, suggesting importance signals are accessible via internal mechanisms.

4. Experimental Results

Datasets and Models

• Benchmarks: GSM8K (arithmetic), MATH (olympiad math), MMLU-Pro (multi-domain).
• Models:
  • Teachers/Pruners: Llama-3.1-8B, Qwen-2.5-7B.
  • Students: Llama-2-7B, Mistral-7B.
• Baselines: TokenSkip (frontier-supervised), H2O (attention-based cache eviction), Surprisal (prediction probability), and Uniform (random).

Performance Findings

• Distillation Accuracy: Greedy pruning consistently yielded the highest student accuracy across all datasets and keep fractions ($\rho$). For example, on GSM8K, greedy-pruned students significantly outperformed TokenSkip and H2O.
• Compression Quality: Even at aggressive pruning (keeping only 30-40% of tokens), the pruned chains retained sufficient semantic and mathematical structure to train effective students.
• Pruner Strength: Using a weaker model as the pruner (e.g., Llama-2-7B instead of Qwen-2.5-7B) degraded the functional structure and student performance, indicating that the quality of the importance signal depends on the pruner's capacity.

Functional Structure Analysis

By categorizing tokens into six functional types (SYMBMATH, METADISC, COREF, ENTNAME, VERBALMATH, GRAMMAR), the authors found:

• SYMBMATH (Symbols/Math): Strongly over-retained. Essential for the answer.
• COREF (Pronouns) & GRAMMAR: Pruned earliest. These are "scaffolding" that the model deems less critical for the final likelihood.
• Objective Impact: The $L_{JOINT}$ objective preserved a clearer functional structure than $L_{ANS}$, leading to better distillation results.

Dynamics and Predictability

• Dynamic Rankings: A "frozen" ranking (computed once at the start) performed significantly worse than dynamic rankings (re-evaluated at each step). This confirms that token importance is context-dependent.
• Attention Prediction: A simple MLP trained on attention scores could predict the greedy pruning order with high accuracy, suggesting that attention mechanisms inherently encode functional importance.

5. Significance and Implications

• Internal Representation: The paper provides strong evidence that LLMs do not treat all reasoning tokens equally; they encode a latent hierarchy of importance that aligns with the functional necessity of tokens for answer generation.
• Efficiency without External Supervision: Greedy pruning offers a principled, self-supervised way to compress reasoning chains, potentially reducing the need for expensive frontier models to generate training data for smaller models.
• Interpretability: The method serves as a "probe" to understand how models organize reasoning, distinguishing between computational core and linguistic fluff.
• Future Directions: The authors suggest that these findings could lead to inference-time pruning (removing tokens on the fly) and better training curricula that focus on functional tokens rather than raw sequence length.

Limitations

• Computational Cost: The naive implementation of Greedy Pruning is $O(n^2)$ or $O(n^4)$ depending on implementation, making it expensive for very long chains. However, the authors note this is an offline preprocessing step.
• Readability: Aggressive pruning can degrade surface-level fluency (missing punctuation, broken grammar), though the mathematical logic remains intact.
• Scope: Experiments were limited to moderately sized models; performance on massive reasoning-focused models (e.g., DeepSeek-R1) was not tested due to resource constraints.