PaperlayerView exact revision on arXiv

Source-linked AI summary

Every Token Leaves a Ripple in the Stream of Thought: Eliciting Model-Internal Token Saliency for Chain-of-Thought Compression

Tianyi Zhao, Yinhan He, Wendy Zheng, Chen Chen

arXiv:2608.31066v1cs.CL

TL;DR

Long CoT traces raise inference costs, while existing compression methods often select tokens using signals only indirectly related to the target model’s answer computation. MIST measures token importance through necessity and sufficiency residual-stream interventions, approximates them with two backward passes, and consistently improves compression across four benchmarks and four models.

  • Problem

    Token-level CoT compression needs to identify which reasoning tokens to retain, but existing methods often rely on external or heuristic signals indirectly tied to the target model’s internal answer computation.

  • Method

    MIST defines necessity and sufficiency through residual-stream interventions and combines first-order Taylor approximations of both axes into a unified token-importance score.

  • Results

    MIST consistently improves CoT compression performance across four reasoning benchmarks and four evaluated models.

  • Takeaways & Limitations

    Model-internal saliency provides a reliable proxy for identifying reasoning tokens relevant to the target model’s answer computation.

  • Takeaways & Limitations

    MIST requires access to target-model internal activations and gradients and does not apply to black-box API-only systems.

Abstract

from arXiv · show

Chain-of-thought (CoT) reasoning improves multi-step problem solving, but long reasoning traces inflate inference cost. Token-level CoT compression reduces this cost by pruning full reasoning chains into shorter traces for model adaptation, making token selection the central challenge. Existing methods often rely on external scorers or heuristic signals only indirectly tied to the model's internal answer computation. We instead adopt a model-internal perspective: as the model forms an answer, each reasoning token leaves a ripple in the residual stream, the model's \emph{stream of thought}, and the magnitude of this ripple reflects the token's contribution to the answer computation. Building on this view, we propose \textsc{MIST} (Model-Internal Saliency for Token-level CoT compression), which defines token importance along two complementary axes: \emph{necessity}, the drop in answer likelihood when a token's internal contribution is removed, and \emph{sufficiency}, the gain in answer likelihood when that contribution alone is provided. Combining the two yields a unified importance score for pruning. Across four reasoning benchmarks and four models, \textsc{MIST} consistently outperforms baseline methods, suggesting that model-internal saliency provides an effective proxy for reasoning-token importance.

1 Introduction

MIST reframes token-level CoT compression as selecting reasoning tokens according to their contributions to the target model’s answer computation. It measures complementary necessity and sufficiency through residual-stream interventions, approximates them efficiently, and reports consistent gains across datasets and models.

  • Motivation: Long CoT traces increase inference cost, motivating compression that preserves reasoning performance while shortening intermediate traces.The token-pruning setting retains a subset of reasoning tokens under a retention budget γ.
  • Motivation: Existing token-level methods use external scorers or heuristic signals only indirectly tied to the target model’s internal answer computation.MIST instead asks which tokens the target model itself relies on when forming the answer.
  • Model-Internal Saliency: Necessity measures the answer-likelihood drop after removing a token’s residual contribution, whereas sufficiency measures answer-likelihood recovery when that contribution alone is patched into a no-chain pass.The two axes capture full-chain dependence and no-chain recoverability, respectively.
  • MIST: MIST combines residual-stream interventions with first-order Taylor linearizations, reducing tokenwise scoring from O(T) intervention passes to two backward passes per chain.The unified score ranks tokens for retention under compression.
  • Evaluation: MIST is evaluated on four reasoning benchmarks with four instruction-tuned models, and the experiments report consistent improvements in CoT compression performance.The benchmarks span mathematical and general-domain reasoning.

2 Related Work

Prior efficient-reasoning work compresses or avoids explicit CoT through concise generation, adaptive budgets, latent computation, and pruning. Related pruning methods use auxiliary scorers, entropy, gradient norms, or repeated deletion, while MIST draws on model-internal attribution and residual-stream interventions.

  • Efficient Chain-of-Thought Reasoning: Efficient-reasoning methods reduce CoT cost through concise generation, adaptive decoding budgets, length-controllable fine-tuning, latent reasoning, and trace pruning.This paper focuses on pruning generated reasoning traces.
  • Token-Level CoT Pruning: TokenSkip uses an auxiliary token scorer, while step-level pruning methods can use entropy to decide which generation steps to skip.Other direct approaches use gradient norms or repeated deletion evaluations.
  • Model-Internal Attribution and Interventions: Mechanistic interpretability identifies behaviorally relevant neurons and attention heads through causal mediation, causal tracing, and activation patching.Residual-stream analysis and direct logit attribution motivate measuring token contributions internally.

3 MIST: Model-Internal Saliency for Token-level CoT Compression

MIST frames token-level CoT compression as model-internal saliency estimation, scoring tokens by how their residual-stream information supports answer computation. It combines necessity and sufficiency interventions, then uses first-order linearization and layer aggregation to produce a unified pruning score.

  • Problem formulation: MIST assigns saliency to reasoning tokens from the target model’s internal states and retains the highest-scoring tokens under a retention budget.The resulting compressed chains supervise fine-tuning for compact reasoning-trace generation.
  • Two saliency axes: Necessity measures the answer-likelihood drop after zeroing token i’s residual state in the full chain.Sufficiency measures the answer-likelihood gain when only that state is patched into a no-chain forward pass.
  • Empirical motivation: On 100 gold MATH chains, exact-necessity ranking beats positional and random baselines, while sufficient-token patching recovers answer information beyond random and norm-matched controls.The two rankings are individually effective but only weakly correlated, motivating their combination.
  • Efficient scoring: Naively evaluating either intervention costs O(T) additional forward passes per chain, making direct tokenwise scoring impractical at scale.MIST instead linearizes answer log-likelihood around unperturbed passes, yielding a scorer computable with two backward passes per chain, independent of T.
  • Unified score: MIST aggregates per-layer necessity and sufficiency contributions using logit-lens-induced weights tied to each layer’s push toward the gold answer token.It combines the weighted axes as α · bϕ_i + (1 − α) · bψ_i, with α ∈ [0, 1], to form the unified per-token score.

4 Experiments

MIST is evaluated across four benchmarks and four instruction-tuned models, with comparisons spanning external scorers, heuristics, compression trade-offs, general-domain reasoning, sensitivity, ablations, and retained-token patterns.

  • Experimental Setup: The evaluation covers GSM8K, MATH, MMLU-Pro, and BIG-Bench Hard across Qwen2.5, Llama-3.1, and Mistral instruction-tuned models.The setup also includes multiple model scales and three model families.
  • Main Results: MIST outperforms TokenSkip across evaluated datasets and models, with smaller accuracy losses on GSM8K and gains on MATH-500.On GSM8K, MIST drops 1.3-2.4 pp versus TokenSkip’s 2.1-4.3 pp; on MATH-500, MIST gains 5.3 pp and 1.2 pp where TokenSkip drops 4.1 pp and 2.3 pp.
  • Main Results: MIST combines compression and accuracy more favorably than TokenSkip on GSM8K with Qwen2.5-7B.It reduces chain length by 21.3% versus 20.0% while incurring a 2.4-pp accuracy drop versus 4.3 pp.
  • General-Domain Results: MIST remains strongest on MMLU-Pro and BIG-Bench Hard across evaluated models, staying close to or sometimes surpassing full CoT on BIG-Bench Hard.Heuristic baselines fall behind, supporting effectiveness beyond mathematical reasoning.
  • Sensitivity and Ablation: MIST’s performance is stable for α ∈{0.4, 0.5, 0.6, 0.7}, and full MIST outperforms its component ablations at every tested retention budget.The main experiments fix α = 0.6; on MATH-500, bψ-only loses up to 7.5 pp and bφ-only up to 3.4 pp.
  • Token Analysis: MIST retains more quantitative tokens than heuristic baselines while preserving content and function words, aligning selection with answer-relevant computation.It allocates 29% of its budget to numbers and 10% to mathematical operators; heuristics retain frequent or high-surprisal but information-light tokens.

5 Conclusion

MIST defines token importance through model-internal necessity and sufficiency interventions, then uses first-order Taylor approximations to make scoring tractable. Evaluations across diverse datasets and models demonstrate its effectiveness, while the method’s scope remains bounded by model scale, gradient access, evaluation domains, and approximation error.

  • MIST defines token importance through necessity and sufficiency interventions on the residual stream.Necessity removes a token’s internal contribution; sufficiency provides that contribution alone to measure answer-likelihood effects.
  • First-order Taylor approximations reduce tokenwise intervention scoring to two backward passes per chain.This makes the model-internal scorer tractable compared with separate intervention passes for every token.
  • Extensive empirical evaluations across diverse datasets and models demonstrate MIST’s effectiveness.
  • Whether MIST’s saliency signal remains similarly concentrated in larger models is not established.
  • MIST requires access to target-model internal activations and gradients, excluding black-box API-only systems.
  • The evaluation spans four reasoning benchmarks, while open-ended generation, code generation, and multi-turn dialogue remain future extensions.
  • First-order Taylor scores omit higher-order interactions among tokens, layers, and heads, so they estimate rather than exactly decompose causal contributions.

A Dataset Details

The appendix documents dataset construction and evaluation examples across GSM8K, MATH, MMLU-Pro, and BBH-MC. It also includes dataset split tables and a per-token MIST saliency visualization for a GSM8K chain.

  • Table 3 summarizes the per-dataset split sizes used throughout the evaluation pipeline.
  • GSM8K: GSM8K uses the standard train and test splits without further filtering.
  • MATH: MATH combines seven Hendrycks subjects for training and uses the canonical MATH-500 test subset.
  • GSM8K: The appendix illustrates GSM8K reasoning with a coffee-cost question whose answer is $120 for April.
  • Figure 9 displays per-token MIST saliency on a GSM8K chain from Qwen2.5-1.5B-Instruct with γ = 0.5.
  • MMLU-Pro: MMLU-Pro provides 12,032 test examples and 70 few-shot validation examples, partitioning the test set into 10,501 training and 1,531 evaluation examples.
  • BBH-MC: BBH-MC uses 17 letter-multiple-choice subtasks and partitions available examples into 3,258 training and 816 evaluation examples.

B Baseline Details

The appendix specifies the external, heuristic, and reference baselines used alongside MIST. These include attention-based, likelihood-based, gradient-based, random, prompting, and uncompressed or no-chain comparisons.

  • The appendix details nine baselines compared in Section 4.1.
  • Prompting and references: Prompt-reduce uses a frozen target model and a γ-aware reduction directive to compress chains in one forward pass.
  • Prompting and references: Full-chain fine-tunes on unmodified self-generated chains, whereas No-chain fine-tunes on answer spans with reasoning removed.
  • Token selection: Uniform retains a uniformly random subset of ⌈γT⌉ tokens at each retention budget.
  • Token selection: TokenSkip retains the top-⌈γT⌉ tokens scored by a pretrained LLMLingua-2 classifier.
  • Internal and heuristic scores: GoGI-ℓ1 scores tokens using the single-layer gradient norm at the last causally connected layer.
  • Internal and heuristic scores: Perplexity retains tokens with higher self-information, si = −log pM(ti | x, t<i).
  • Attention baselines: Attention rollout scores tokens using cross-layer aggregated attention and answer-position column sums.

C Training, Evaluation, and Implementation Details

The implementation uses a three-stage pipeline: generate and filter chains, score tokens, then fine-tune and evaluate LoRA adapters at multiple retention budgets. Technical details cover retention definitions, precision, hardware, gradient extraction, and Taylor approximations for necessity and sufficiency.

  • Pipeline: The pipeline self-generates answer-correct chains, scores tokens, retains top tokens at γ ∈ {0.5, 0.6, 0.7, 0.8, 0.9}, and fine-tunes LoRA adapters.
  • Retention and compression: Training retention γ is the per-chain fraction of CoT tokens kept for supervision, whereas inference compression rate measures reduction against the full-chain LoRA baseline.
  • Pipeline: Self-generation uses greedy decoding and dataset-specific maximum lengths of 512 for GSM8K, 768 for MMLU-Pro and BBH-MC, and 1,024 for MATH.
  • Token scoring: MIST scoring uses source and target backward passes with retain-grad hooks exposing gradients at every transformer block and chain position.
  • Fine-tuning: LoRA fine-tuning follows TokenSkip’s recipe with rank r = 8, scaling α = 16, zero dropout, and all linear projections targeted.
  • Evaluation: Evaluation uses greedy decoding with temperature 0, top-p 1.0, batch size 16, dataset-specific generation limits, SDPA, and bfloat16 precision.
  • Hardware and software: Experiments run on two NVIDIA A100 80 GB GPUs using PyTorch 2.9, Transformers, and PEFT, with separate costs for generation, scoring, training, and evaluation.
  • Taylor formalization: The Taylor formalization defines local source- and target-side functions for residual-stream interventions and derives first-order terms with remainder bounds.

D.5 Empirical Tightness of the First-Order Term

The paper tests whether the first-order Taylor term predicts actual answer-likelihood changes caused by residual-stream token interventions. On active cells, its empirical behavior closely follows the first-order prediction.

  • The experiment scales one token's residual stream at one residual block while leaving all other activations unchanged, then measures the answer-log-likelihood drop.
  • The per-layer Taylor main term is evaluated as the first-order approximation to the intervention effect.
  • The empirical median tracks the first-order prediction y = ε within ≲10% overall on 8 GSM8K chains.Measurements use Qwen2.5-1.5B-Instruct in fp32 over functionally active token-layer cells.

D.6 Empirical Ranking Fidelity of the Taylor Proxies Against the Exact Interventions

The Taylor-based necessity and sufficiency proxies are compared with exact interventions for token selection on GSM8K. Both proxies substantially exceed chance agreement at the tested selection budgets.

  • Necessity reaches top-γ agreement of 0.72 at γ = 0.3 and 0.81 at γ = 0.5 with its exact intervention.
  • Sufficiency reaches top-γ agreement of 0.68 at γ = 0.3 and 0.76 at γ = 0.5 with its exact intervention.
  • Both proxies exceed the corresponding chance agreements of 0.3 and 0.5, indicating reliable preservation of exact-intervention token selections.

F Additional Main Results on BBH and MMLU-Pro

This section provides additional BBH and MMLU-Pro results and reports realized serving latency for MIST under compressed and uncompressed decoding conditions.

  • Additional Main Results on BBH and MMLU-Pro: Table 6 reports per-model and scorer compression results on BBH and MMLU-Pro using the same accuracy-change and compression-rate format as Table 1.
  • Latency Measurement: Wall-clock decoding latency is measured on all 1,319 GSM8K test examples with batched greedy decoding, batch size 16, bfloat16, and SDPA.
  • Latency Measurement: Table 7 compares the MIST adapter at γ = 0.5 with the same adapter at γ = 1.0, reporting mean latency per problem and the resulting speedup.

H Empirical Per-Layer Logit-Lens Weight

The paper measures how answer-unembedding-aligned residual-stream updates vary across layers and finds that their magnitude is concentrated in late layers rather than distributed uniformly.

  • Empirical Distribution: Across 30 BBH chains, mean layer weights for Qwen2.5-1.5B-Instruct and Llama-3.1-8B-Instruct are overwhelmingly concentrated in late layers, while many middle layers are near zero.
  • Layer Weight Definition: The layer weight ¯c_l is the chain-averaged inner product between layer l's residual-stream update and the answer unembedding direction.
  • Implication for Aggregation: Uniform weighting loses up to 4.6 pp because it spends selection budget on layer-token cells whose updates are essentially answer-neutral.
  • Score Standardization: The necessity and sufficiency axes are standardized per chain before blending because their inner-product scores have different scales across chains.
  • Score Standardization: Necessity uses log-space standardization, while sufficiency uses a standard z-score with ε = 10^-12 for numerical stability.
Loading 2608.31066v1…