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

1 Introduction

• CoT improves LLM multi-step problem solving, but long traces increase latency, memory use, and serving cost; token pruning asks which tokens to retain under a budget.The pruning setting shortens full traces into subsets of reasoning tokens that supervise models to generate shorter chains.Source: §1 · p. 1
• MIST treats token importance as the target model’s own reliance on reasoning-token information in the residual stream, rather than an indirect heuristic proxy.The introduction contrasts this perspective with TokenSkip-style methods that use auxiliary scorers whose rankings depend on their own training objectives, supervision, and domain assumptions (Xia et al., 2025; Pan et al., 2024).Source: §1 · pp. 1–2
• Necessity measures the answer-likelihood drop from removing a token’s contribution, while sufficiency measures the likelihood recovered by providing that contribution alone.Necessity captures dependence within the full chain but can favor low-information context; sufficiency captures no-chain recoverability but underperforms under aggressive compression, so combining them supports ranking across budgets.Source: Figure 1 · p. 1, §1 · p. 2
• MIST scores tokens through residual-stream erasure and patching interventions, linearized so one backward pass per axis yields all token scores.It combines necessity and sufficiency into a unified importance score and retains the highest-scoring tokens to create compressed traces for model adaptation.Source: §1 · p. 2
• Across four reasoning benchmarks and four models, MIST consistently improves CoT compression performance across settings.The evaluation spans mathematical reasoning benchmarks GSM8K and MATH, and general-domain benchmarks MMLU-Pro and BIG-Bench Hard, using Qwen2.5-1.5B-Instruct, Qwen2.5-7B-Instruct, Llama-3.1-8B-Instruct, and Mistral-7B-Instruct-v0.3.Source: §1 · p. 2

Efficient chain-of-thought reasoning

• CoT prompting improves multi-step reasoning but increases inference cost, motivating methods including concise generation, adaptive budgets, length-controllable fine-tuning, latent reasoning, and trace pruning.MIST focuses on pruning generated traces.Source: §2 · p. 3

Token-level CoT pruning

• MIST derives token importance from two residual-stream saliency measures, unlike external scoring, single-axis heuristics, or costly iterative deletion methods.TokenSkip uses an auxiliary token scorer, step-level methods use signals such as entropy, GoGI-Skip uses gradient norm, and likelihood-preserving greedy deletion repeatedly evaluates deletions.Source: §2 · p. 3

Model-internal attribution and interventions

• Mechanistic interpretability uses interventions to identify behaviorally relevant components, while residual-stream analysis and direct logit attribution decompose contributions to output logits.MIST adapts this model-internal perspective to measure reasoning-token importance for CoT compression.Source: §2 · p. 3

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

• MIST scores each chain token’s internal contribution to answer computation, retains top-scoring tokens under a retention budget, and uses compressed chains to adapt the model.Token importance is measured through residual-stream states, which the paper treats as the model’s internal computational substrate.Source: §3 · p. 3

3.1 Problem Formulation: A Model-Internal Notion of Token Importance

• Necessity measures the answer-likelihood drop when a token’s residual state is zeroed; sufficiency measures the gain when that state alone is patched into a no-chain pass.The two axes operationalize what the model loses when a contribution is removed and what it can recover when that contribution is supplied alone.Source: §3.1 · p. 3
• On 100 gold MATH chains, necessity rankings beat positional and random baselines, sufficiency recovered answer-relevant information, and the weakly correlated rankings captured distinct aspects of importance.Random-token and norm-matched controls did not recover comparable information with sufficiency; together, the findings motivate combining both axes.Source: §3.1 · p. 4

3.2 Quantifying Model-Internal Saliency in Practice

• Linearizing answer log-likelihood around unperturbed forward passes yields a first-order scorer computed with two backward passes per chain, independent of chain length.Naively evaluating each token requires an additional forward pass per token, for O(T) cost per chain.Source: §3.2 · p. 4

3.3 The Unified MIST Score

• MIST weights layer contributions using a logit-lens-derived measure of their push toward the gold answer, aggregates each axis across layers, and combines them with α.The resulting per-token scores form compressed chains used to supervise fine-tuning at retention budget γ.Source: §3.3 · p. 5

4.2 Main Results

• MIST outperforms TokenSkip across evaluated datasets and models, with smaller GSM8K accuracy drops and MATH-500 accuracy gains.On GSM8K, MIST drops 1.3–2.4 pp versus TokenSkip’s 2.1–4.3 pp for Qwen2.5-7B and Mistral-7B; on MATH-500, MIST gains 5.3 pp for Qwen2.5-1.5B and 1.2 pp for Llama-3.1-8B, while TokenSkip drops 4.1 pp and 2.3 pp, respectively.Source: Table 1 · p. 5, §4.2 · p. 6

4.1 Experimental Setup

• Experiments cover four reasoning benchmarks and four models, with LoRA-adapted chains trained across retention budgets and decoded freely at inference.The benchmarks span GSM8K, MATH, MMLU-Pro, and BBH; the evaluated models are Qwen2.5-1.5B, Qwen2.5-7B, Llama-3.1-8B, and Mistral-7B. Because γ controls retained supervision tokens rather than imposing a hard inference-length constraint, actual compression is typically smaller than 1−γ.Source: §4.1 · pp. 5–6

4.3 Generalization Beyond Mathematical Reasoning

• MIST remains strongest on MMLU-Pro and BBH across evaluated models, extending its advantage beyond mathematical reasoning.On BBH, it stays closest to—and in some settings surpasses—the full-CoT reference, while heuristic baselines fall noticeably behind.Source: Figure 4 · p. 6, Figure 5 · p. 6, §4.3 · p. 7

4.2 Main Results

• MIST’s necessity-and-sufficiency scorer outperforms internal heuristic baselines, including GoGI, which uses a single-layer gradient norm as a proxy for token importance.On GSM8K, GoGI drops 11.9 pp with Qwen2.5-7B and 9.2 pp with Mistral-7B, versus MIST’s 2.4 pp and 1.3 pp; perplexity, attention rollout, and H2O reach drops as large as 39.5, 28.8, and 32.6 pp, respectively.Source: §4.2 · p. 6
• On GSM8K with Qwen2.5-7B, MIST reduces chain length by 21.3% with a 2.4-pp accuracy drop, versus TokenSkip’s 20.0% reduction and 4.3-pp drop.Perplexity, attention rollout, and H2O can achieve similar or higher compression rates but suffer much larger accuracy drops.Source: §4.2 · pp. 6–7

4.4 Ablation Study

• Full MIST performs best at every tested retention budget, while removing either saliency axis, layer weighting, or multi-layer aggregation reduces accuracy.On MATH-500 with Llama-3.1-8B, bψ-only drops up to 7.5 pp at γ=0.5 (37.5% vs. 45.0%); bφ-only loses up to 3.4 pp at γ=0.8 (46.2% vs. 49.6%), and MIST-uniform and MIST-single each lose 0.8–4.6 pp depending on γ.Source: Figure 7 · p. 7, §4.4 · p. 7

4.5 Token Type Analysis

• MIST allocates more of its GSM8K retention budget to quantitative tokens and retains the highest fraction of numbers and mathematical operators.It assigns 29% of its budget to numbers and 10% to operators, retaining 75% of all numbers and 66% of all operators; attention rollout and perplexity allocate 7% and 12% to numbers, retaining 22% and 33%, respectively.Source: Figure 8 · p. 8, §4.5 · p. 8

4.6 Case Study

• In a GSM8K case study, necessity and sufficiency saliency emphasize overlapping but distinct answer-relevant tokens, and their combination preserves arithmetic anchors and semantic context.Necessity highlights intermediate facts ($4, 30) and $120, while sufficiency highlights buying, spending, currency, and equality; the unified score retains $4 × 30 = $120 and “spends $120”.Source: §4.6 · p. 8

5 Conclusion

• MIST scores CoT tokens using residual-stream necessity and sufficiency interventions, with first-order Taylor approximations reducing tokenwise intervention cost to two backward passes per chain.The method was evaluated across diverse datasets and models.Source: §5 · p. 8

Limitations

• Experiments use 1.5B to 8B instruction-tuned models, and whether saliency remains similarly concentrated in larger models is not established.MIST also requires access to the target model’s internal activations and gradients, so it does not apply to black-box API-only systems.Source: Limitations · p. 9
• The four benchmarks cover mathematical and general-domain reasoning, leaving effectiveness in open-ended generation, code generation, and multi-turn dialogue for future validation.Source: Limitations · p. 9
• MIST’s first-order Taylor intervention scores omit higher-order interactions among tokens, layers, and heads, making them saliency estimates rather than exact causal decompositions.Transformer computations are nonlinear.Source: Limitations · p. 9

A Dataset Details

• Table 3 summarizes the per-dataset split sizes used throughout the evaluation pipeline.The accompanying text introduces dataset-specific preprocessing and self-generation hyperparameters, but does not provide their values here.Source: A Dataset Details · p. 10, Table 3 · p. 10
• GSM8K uses its standard release train and test splits without further filtering.Source: A Dataset Details · p. 10
• MATH uses the concatenated training sets of seven Hendrycks et al. (2021) subjects and the canonical MATH-500 test subset.Source: A Dataset Details · p. 10
• MMLU-Pro's 12,032 examples are randomly partitioned into 10,501 training and 1,531 evaluation examples.The original release provides a 70-example few-shot validation set but no standard training split.Source: A Dataset Details · p. 11
• BBH-MC uses 17 letter-multiple-choice subtasks, excluding 10 binary or free-form subtasks, and partitions examples 3,258/816 for training/evaluation.The 80/20 split is drawn from the available examples because BBH provides no standard training split (Suzgun et al., 2023).Source: A Dataset Details · p. 11

• Prompt-reduce asks the frozen target model to reduce the CoT by (1−γ)×100% in one forward pass, without finetuning (Xia et al., 2025).It uses the reduction directive from TokenSkip, Table 1.Source: B Baseline Details · p. 11
• Full-chain finetunes on unmodified self-generated chains at γ = 1.0, while No-chain finetunes on answer spans with reasoning removed.Source: B Baseline Details · p. 11
• Uniform retains a random ⌈γT⌉ tokens, whereas TokenSkip retains the top ⌈γT⌉ tokens scored by a pretrained LLMLingua-2 classifier (Xia et al., 2025; Pan et al., 2024).Source: B Baseline Details · p. 11
• GoGI-ℓ1 uses a single-layer gradient norm at residual layer l⋆ = max(L − 2, 0), and log pM∥1 is its single-layer gradient-norm-only counterpart to MIST (Zhuang et al., 2025).Source: B Baseline Details · pp. 11–12
• Other scorers retain high-self-information tokens using perplexity (Jiang et al., 2023b), aggregate cross-layer attention rollout (Abnar and Zuidema, 2020), or score cumulative answer-position attention mass with H2O.Source: B Baseline Details · pp. 11–12

C Training, Evaluation, and Implementation Details

• Each model–dataset cell uses answer-matched self-generated chains shared across scorers, scores tokens, keeps the top-⌈γT⌉ per chain, fine-tunes with LoRA, and evaluates by greedy decoding.Reusing the same chains across scorers isolates differences attributable to scoring rather than supervision data.Source: C Training, Evaluation, and… · p. 12
• Training retention γ is the per-chain fraction of CoT tokens kept, whereas inference compression is the reduction in generated tokens against the full-chain LoRA baseline; the ratios need not match.Adapters can learn a length policy that produces chains shorter than the average retention floor, especially on smaller models.Source: C Training, Evaluation, and… · p. 12
• Adapters are trained once per model–dataset–scorer on a mixture of retention budgets, then reused at evaluation budgets with greedy decoding.Training uses rank-8 LoRA for 3 epochs; evaluation uses dataset-specific generation limits and records decoding time for inference-cost analysis.Source: C Training, Evaluation, and… · p. 13

D Taylor Expansions and Remainder Bounds

• Appendix D provides first-order Taylor expansions for per-layer necessity and sufficiency terms, together with explicit per-layer remainder bounds.These derivations complete technical material deferred from Section 3.2.Source: D Taylor Expansions and Remainder Bounds · p. 13

D.1 Setup and notation

• The derivation defines local source- and target-side functions for residual interventions and assumes local C2 regularity only to state Lagrange-form remainders; empirical scoring needs only first-order gradients.Source: §D.1 · p. 13

D.2 Necessity: Taylor expansion of ϕ(l)

• For necessity, the per-layer intervention replaces the chain-position residual with zero, and the Taylor expansion recovers the defined first-order term.Source: §D.2 · pp. 13–14

D.3 Sufficiency: Taylor expansion of ψ(l)

• For sufficiency, the intervention inserts a source chain-position residual at the target’s final position; its expansion recovers the first-order term, with a positive remainder sign because sufficiency is a likelihood gain.Source: §D.3 · p. 14

D.4 Recovering all per-layer gradients from a single backward pass

• Retain-gradient hooks expose all intermediate residual gradients in one backward pass, so necessity needs one source backward and sufficiency one target backward.Each also needs its corresponding forward pass, with activation-caching overhead.Source: §D.4 · p. 14

D.5 Empirical Tightness of the First-Order Term

• The test scaled one chain token’s residual stream at one block while leaving others unperturbed, then measured the answer-log-likelihood drop.This directly evaluates how tightly the first-order term predicts real-chain ablation drops, beyond worst-case bounds with an unspecified Hessian constant.Source: §D.5 · p. 14
• Across 8 GSM8K chains, the per-layer first-order term’s empirical median tracked the answer-log-likelihood drop along y = ε to within ≲10% overall.Measurements used Qwen2.5-1.5B-Instruct in fp32 over functionally active token-layer cells.Source: Figure 10 · p. 15, §D.5 · p. 15

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

• Necessity achieves top-γ agreement of 0.72 at γ = 0.3 and 0.81 at γ = 0.5; sufficiency achieves 0.68 and 0.76, respectively, above chance.The evaluation uses GSM8K with Qwen2.5-1.5B-Instruct; the corresponding chance levels are 0.3 and 0.5.Source: Table 5 · p. 15, §D.6 · p. 15

E Licenses and Terms of Use

• The paper states that external artifacts and released artifacts are used consistently with their intended research purposes.All artifacts are publicly available under research-friendly terms; the paper says it does not redistribute the underlying datasets.Source: E Licenses and Terms of Use · p. 15
• GSM8K, MATH, and MMLU-Pro use the MIT License, while BIG-Bench Hard uses Apache License 2.0.The experiments load datasets from their original sources and use them for benchmarking and training LoRA adapters on released training splits.Source: E Licenses and Terms of Use · p. 15
• Qwen2.5-1.5B/7B use the Qwen License, Llama-3.1-8B uses the Llama 3.1 Community License, and Mistral-7B uses Apache License 2.0.The stated terms permit research use, fine-tuning, and release of derived artifacts such as LoRA adapters.Source: E Licenses and Terms of Use · p. 15
• Adapters inherit their base models’ access conditions, and Llama-3.1 adapters are restricted to use with Llama and not for improving models outside its family.Adapters built on Qwen2.5 and Mistral carry no such derivative restrictions; the trained adapters are not deployed in commercial products.Source: E Licenses and Terms of Use · pp. 15–16

F Additional Main Results on BBH and MMLU-Pro

• Table 6 reports per-model, per-scorer compression results for BBH and MMLU-Pro in ∆Acc/compression-rate format, underlying Figs. 4 and 5.The format matches Table 1.Source: F Additional Main Results on BBH and… · p. 16

G Direct inference-cost measurement.

• MIST’s γ=0.5 adapter is compared with the same adapter at γ=1.0 using mean wall-clock decode time per GSM8K problem and the γ=1.0-to-γ=0.5 latency ratio.The measurement uses all 1,319 GSM8K test examples, batched greedy decoding, batch size 16, bfloat16, and SDPA.Source: G Direct inference-cost measurement · p. 16, Table 7 · p. 16

H Empirical Per-Layer Logit-Lens Weight

• Figure 11 shows answer-directed layer weights concentrated in late layers, with many middle layers near zero, across Qwen2.5-1.5B-Instruct and Llama-3.1-8B-Instruct.The mean is computed across 30 BBH chains; this pattern supports layer-weighted aggregation over uniform per-layer weighting.Source: Figure 11 · p. 16, H Empirical Per-Layer Logit-Lens Weight · p. 17
• The layer weight ¯cl is the chain-averaged inner product of layer l’s residual-stream update with the answer unembedding direction.The average is taken over chain tokens.Source: H Empirical Per-Layer Logit-Lens Weight · pp. 16–17
• MIST-uniform loses up to 4.6 pp because uniform weighting spends selection budget on layer-token cells whose updates are essentially answer-neutral.Late-layer dominance is consistent across the two model families, indicating the aggregation is not specific to one architecture.Source: H Empirical Per-Layer Logit-Lens Weight · p. 17

I Per-Axis Normalization Before Combining

• Necessity and sufficiency scores are inner products of log-likelihood gradients with different reference activations, so they have different scales across chains.Per-chain standardization makes mixing coefficient α in Eq. (8) comparable across chains and benchmarks.Source: I Per-Axis Normalization Before… · p. 17
• Necessity is standardized in log space, while sufficiency uses a standard z-score before the axes are blended.The log transformation suits per-chain magnitudes spanning several orders of magnitude and preserves token ordering within a chain.Source: I Per-Axis Normalization Before… · p. 17
• The per-chain normalization uses token-index statistics and ε = 10^-12 for numerical stability before computing the unified score on normalized necessity scores.The supplied text identifies these normalization details but does not provide the full equations.Source: I Per-Axis Normalization Before… · p. 17