PaperlayerView exact revision on arXiv

Source-linked AI summary

Pruning the Unsurprising: Efficient LLM Reasoning via First-Token Surprisal

Wenhao Zeng, Yaoning Wang, Chao Hu, Yuling Shi, Chengcheng Wan, Hongyu Zhang, Xiaodong Gu

arXiv:2508.05988v2cs.LGcs.SE

TL;DR

Long Chain-of-Thought traces make LRM reasoning costly, while existing token- and step-level compression methods can disrupt coherence or dilute logical importance. ASAP combines anchor-guided structural pruning with first-token-surprisal refinement, then distills compact CoTs for inference. Across benchmarks it achieves state-of-the-art accuracy with substantially lower computation, including 36.19% Pass@1 on LiveCodeBench v4_v5 alongside 23.5% fewer tokens and 43.5% lower latency than the strongest baseline.

  • Problem

    Long reasoning traces impose substantial training and inference costs, while existing compression methods may disrupt coherence or miss logically critical steps.

    Source: §1 · p. 1

  • Method

    ASAP uses anchor-guided pruning to remove structural redundancy, first-token surprisal to retain critical reasoning steps, and distillation to generate compact CoTs.

    Source: §1 · p. 2, §3.2 · p. 3, §3.3 · p. 4

  • Results

    ASAP achieves state-of-the-art accuracy across multiple benchmarks while reducing computational costs; on LiveCodeBench v4_v5 it reaches 36.19% Pass@1 with 23.5% fewer tokens and 43.5% lower latency than the strongest baseline.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    Information-theoretic signals concentrated at reasoning-step onset can support efficient CoT compression while preserving benchmark performance.

    Source: §6 · p. 8

    ASAP depends on a capable LLM for accurate Stage 1 anchors, and its effectiveness beyond code generation and mathematical reasoning remains unverified.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large Reasoning Models (LRMs) have demonstrated remarkable capabilities by scaling up the length of Chain-of-Thought (CoT). However, excessively long reasoning traces pose substantial challenges for training cost and inference latency. While various CoT compression approaches have emerged to address this challenge, they face inherent trade-offs: token-level methods often disrupt syntactic and logical coherence, while step-level methods based on perplexity fail to reliably capture the logically critical reasoning steps because of the dilution of logical information. In this paper, we propose ASAP (Anchor-guided, SurprisAl-based Pruning), a novel coarse-to-fine framework for CoT compression. ASAP first performs anchor-guided pruning to preserve the core reasoning structure, which efficiently reduces the search space for subsequent processing. Leveraging the insight that logical branching choices are concentrated at the onset of reasoning steps, it then enables logic-aware pruning by selecting logically essential reasoning steps based on a novel first-token surprisal metric. Finally, ASAP distills the models to autonomously generate and leverage these concise CoTs at inference time, enabling efficient reasoning. Experiments show that ASAP achieves state-of-the-art accuracy across multiple benchmarks while substantially reducing training and inference costs.

1 Introduction

ASAP targets structural and logical redundancy in long Chain-of-Thought traces by combining anchor-guided pruning with first-token surprisal. Across benchmarks, it preserves accuracy while reducing reasoning cost and latency.

Motivation

  • Long LRM reasoning traces create substantial latency and memory overhead because they contain digressive branches, over-explanation, and syntactic fillers.These redundancies motivate retaining only the cognitive pivots of reasoning.Source: §1 · pp. 1–2

Method

  • ASAP distills compact, logic-dense CoTs into a target model so it can generate efficient reasoning chains at inference time.The framework addresses both structural redundancy and logical redundancy in a coarse-to-fine cascade.Source: §1 · p. 2

Results

  • 36.19% Pass@1 on LiveCodeBench v4_v5 accompanies 23.5% fewer generated tokens and 43.5% lower inference latency than the strongest baseline.Experiments use DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B across diverse domains.Source: §1 · p. 2

Motivation

  • First-token surprisal identifies logically critical reasoning steps more reliably than perplexity, whose sentence-level signal can dilute logical leaps with predictable syntax.The approach focuses on uncertainty concentrated at the beginning of reasoning steps.Source: §1 · p. 2

Method

  • ASAP combines structural alignment with information-guided refinement to compress Chain-of-Thought reasoning.Stage 1 removes irrelevant branches using a logical backbone, while Stage 2 filters logic-sparse steps and retains high-surprisal cognitive pivots.Source: §1 · p. 2

2 Empirical Analysis

An analysis of 10 million reasoning tokens finds that uncertainty is concentrated in the first tokens of reasoning steps rather than their deterministic elaborations. High-entropy first tokens distinguish cognitive pivots from predictable syntactic connectors.

Experimental Setup

  • The analysis uses 10 million tokens from DeepSeek-R1-Distill-Qwen-32B across AIME and LiveCodeBench reasoning benchmarks.Entropy is used to quantify uncertainty about the next state transition.Source: §2 · p. 2

Information Concentration

  • First tokens have substantially higher entropy than subsequent tokens, showing that logical branching is concentrated at reasoning-step onset.Once a step begins, later elaboration is largely deterministic and syntactically driven.Source: §2 · pp. 2–3

Identifying Cognitive Pivots

  • High-entropy filtering suppresses predictable connectors such as “So” and “Then” while highlighting pivots such as “Alternative,” “Because,” and “Perhaps.”These pivots represent exploration, causality, and self-correction.Source: §2 · p. 3, Figure 2 · p. 3

Implication

  • First-token surprisal operationalizes high entropy as a proxy for identifying and preserving critical reasoning hops.The metric is designed to distinguish logical transitions from predictable elaborations.Source: §2 · p. 3

3 Methodology

ASAP compresses reasoning in a coarse-to-fine pipeline: it first anchors pruning to a concise logical backbone, then removes low-information steps using first-token surprisal, and finally fine-tunes the model on the resulting CoTs.

3.1 Overall Framework

  • ASAP combines anchor-guided pruning, surprisal-based refining, and supervised fine-tuning to produce concise reasoning traces.The anchor stage reduces structural redundancy, the refining stage removes logical redundancy, and the pruned traces train the target model.Source: §3.1 · p. 3, Figure 3 · p. 4

3.2 Anchor-guided Pruning

  • A Direct Thought generated from the question-answer pair acts as a logical anchor for extracting an ordered, semantically aligned subsequence from the original CoT.Pattern matching validates that retained steps correspond to the original CoT, exceed threshold τ, and preserve order.Source: §3.2 · p. 4

3.3 Surprisal-based Refining

  • First-token surprisal estimates each step’s logical importance because the first token carries most uncertainty, distinguishing high-information transitions from deterministic fillers.The method selects steps under a length budget by iteratively removing the lowest-surprisal step while preserving relative order.Source: §3.3 · pp. 4–5

3.4 Supervised Fine-tuning

  • The pruned CoT and final answer are concatenated into target responses, and the LRM is fine-tuned with standard token-level negative log-likelihood.This distills compact, logically salient examples so the model can internalize efficient reasoning patterns.Source: §3.4 · p. 5

4 Experiments

Experiments across code and mathematical reasoning benchmarks show that ASAP preserves or improves accuracy while reducing inference and training costs. Ablations, architecture transfer, token-budget tests, and training measurements support the complementary value and practical efficiency of its pruning stages.

4.2 Main Results

  • ASAP achieves the best accuracy–efficiency trade-off across the reported benchmarks, generating the fewest tokens and lowest latency.Tables 1 and 2 report the cross-benchmark comparison for code and mathematical reasoning tasks.Source: §4.2 · p. 6, Table 1 · p. 6, Table 2 · p. 6
  • Token-level pruning degrades performance by disrupting syntactic and semantic coherence, whereas step-level pruning better preserves sentence-level integrity.SPIRIT improves efficiency with comparable accuracy, while ASAP improves both efficiency and accuracy.Source: §4.2 · p. 6
  • 23.5% fewer tokens, 43.5% lower latency, and 7.8% higher Pass@1 distinguish ASAP from the Original baseline on LiveCodeBench v4_v5.Generated tokens decrease from 7892 to 6035, latency from 4.62s to 2.61s, and Pass@1 rises from 33.58% to 36.19%.Source: §4.2 · p. 6

4.3 Ablation and Analysis

  • Removing either pruning stage harms ASAP: omitting anchor-guided pruning increases latency by 76.2%, while omitting surprisal refinement lowers Pass@1 from 36.19% to 31.72%.The ablation results characterize the two stages as mutually complementary for accuracy and efficiency.Source: §4.3 · p. 7, Table 3 · p. 7
  • ASAP generalizes to DeepSeek-R1-Distill-Llama-8B, surpassing Original accuracy on LiveCodeBench while generating 49.1% fewer tokens and reducing latency by over 3x.On LiveCodeBench, accuracy is 32.84% versus 31.34%, and latency falls from 8.60s to 2.69s.Source: §4.3 · p. 7, Table 4 · p. 7
  • ASAP scales smoothly with token budget and achieves higher accuracy at 8K tokens than SPIRIT and Original at 12K tokens.The comparison is reported for LiveCodeBench v4_v5 across budgets from 2K to 12K.Source: §4.3 · pp. 7–8, Figure 4 · p. 7
  • ASAP reduces training tokens by 75.6% and training time by 60.7% relative to the Original baseline on CodeForces-CoTs.These savings exceed those of the other reported baselines.Source: §4.3 · p. 8, Table 5 · p. 8

5 Related Work

Prior work improves reasoning efficiency through CoT scaling, context compression, and fine-tuning-based compression. ASAP addresses limitations of token removal and perplexity-based step selection by combining structural pruning with first-token surprisal.

Chain-of-Thought and Advanced Reasoning

  • Large reasoning models scale test-time compute by generating lengthy CoT traces, motivating methods that remove redundancy without compromising reasoning performance.This work contrasts pruning redundancy with prior approaches that primarily increase CoT length.Source: §5 · p. 8

Context Compression for LLMs

  • Token-oriented context compression methods often disrupt the syntactic and logical coherence required for valid reasoning when applied directly to CoT.These methods typically treat the input as unstructured text.Source: §5 · p. 8

Efficient Reasoning via Fine-Tuning

  • TokenSkip and SPIRIT use token filtering or perplexity shifts, whereas ASAP combines anchor-guided structural pruning with first-token surprisal as a proxy for cognitive pivots.The paper presents this combination as a more robust alternative for distinguishing logical necessity from syntactic fluency.Source: §5 · p. 8

6 Conclusion

ASAP addresses structural and logical redundancy in Chain-of-Thought reasoning with a coarse-to-fine, information-theoretic pruning framework. The paper reports state-of-the-art accuracy–efficiency trade-offs and identifies online inference acceleration as future work.

6 Conclusion

  • ASAP aligns reasoning with a logical anchor, then refines it using First-Token Surprisal to address structural and logical redundancy.The framework is grounded in the reported concentration of logical uncertainty at reasoning-step onsets.Source: §6 · p. 8
  • ASAP outperforms existing baselines across multiple benchmarks, establishing a new state-of-the-art Pareto frontier between accuracy and efficiency.Source: §6 · p. 8
  • Future work will explore applying ASAP to online inference acceleration.Source: §6 · p. 9

Limitations

The paper reports limitations concerning anchor quality and task coverage, then specifies the pattern-matching and surprisal-pruning algorithms and the prompts used to construct direct and coarse-grained CoTs.

Limitations

  • ASAP depends on a capable LLM generating high-quality Direct Thoughts, and anchor errors or hallucinations may misguide later pruning despite pattern-matching mitigation.Experiments focus primarily on code generation and mathematical reasoning; effectiveness for creative writing and commonsense reasoning remains unverified.Source: Limitations · p. 9

A Algorithms

  • Pattern Matching validates coarse pruning by splitting original and coarse CoTs into steps and requiring sequential Gestalt similarity above threshold τ.The procedure returns false when a coarse step lacks a matching original step.Source: A Algorithms · p. 11
  • Iterative pruning scores coarse-grained steps by first-token surprisal and removes steps until the specified maximum token budget is reached.If the coarse CoT already fits the budget, it is returned unchanged.Source: A Algorithms · p. 11

B Prompt Templates and Generated Examples

  • Direct Thoughts are concise, step-by-step solutions with a single explicit final answer, guiding coarse pruning to remove irrelevant paths while preserving aligned original wording and order.The pruning prompt forbids rewriting, reordering, or adding new words and concepts.Source: B Prompt Templates and Generated… · pp. 11–12

C Implementation Details

The implementation uses memory- and throughput-oriented software settings, fixed fine-tuning and evaluation protocols, standardized baselines, and specified hyperparameters for ASAP's pruning stages.

C Implementation Details

  • Experiments use unsloth for memory-efficient fine-tuning, vLLM for inference throughput, NVIDIA H20 GPUs, and Intel Xeon Platinum 8480+ CPUs.Source: C Implementation Details · p. 12
  • Fine-tuning uses full-parameter bf16 training for 10 epochs, an effective batch size of 16, AdamW with cosine scheduling, and fixed seed 42.The maximum sequence length is 16384, with model-specific peak learning rates of 4 × 10−5 and 2 × 10−5.Source: C Implementation Details · p. 13
  • Inference uses deterministic bfloat16 vLLM decoding with temperature 0.0, disabled prefix caching, and task-specific generation budgets from 2K to 10K tokens.Source: C Implementation Details · p. 13
  • Baseline comparisons preserve questions and final answers while compressing only CoT steps, use mostly 0.5 target compression, and modify SPIRIT for scalable stepwise evaluation.ASAP uses DeepSeek-V3 for coarse pruning, τ = 0.6 for pattern matching, and a 4096-token fine-grained budget.Source: C Implementation Details · p. 13

D Effect of Different Components.

The component study evaluates whether ASAP's coarse- and fine-grained pruning stages are necessary across three programming benchmarks, reporting accuracy, token count, and latency.

D Effect of Different Components

E Generalization to Different Architectures

The paper evaluates ASAP across architectures and inference-time token budgets using code-generation and mathematical-reasoning benchmarks. These experiments compare ASAP with Zero-shot, Original, and SPIRIT under the stated protocols.

E Generalization to Different Architectures

  • ASAP is evaluated on DeepSeek-R1-Distill-Llama-8B to test generalizability across architectures.The comparison uses Zero-shot, Original, and SPIRIT on four code-generation and four mathematical-reasoning benchmarks.Source: E Generalization to Different… · p. 14

F Performance under Different Token Budgets

  • The study compares ASAP with SPIRIT, Original, and Zero-shot under varying inference-time token budgets.Simpler benchmarks use 1K–6K-token budgets, while more complex benchmarks use 2K–12K budgets.Source: F Performance under Different Token… · p. 14
  • The budget analysis spans HumanEval+, LiveCodeBench v1_v3, LiveCodeBench v4_v5, LeetcodeDataset, GSM8K, MATH500, AIME24, and AIME25.Results are reported across eight benchmark-specific tables covering code generation and mathematical reasoning.Source: F Performance under Different Token… · p. 14

G Training Efficiency

The training-efficiency evaluation measures sequence length and training time across CoT datasets. OpenR1-Math results report reductions relative to the Original baseline.

G Training Efficiency

  • Training efficiency is quantified using average tokens per sample and average training time in seconds per step.These metrics are reported for CodeForces-CoTs and OpenR1-Math.Source: G Training Efficiency · p. 14
  • The OpenR1-Math comparison reports average tokens per sample and seconds per training step for each method.Percentages indicate reductions relative to the Original baseline.Source: Table 19 · p. 14
  • The reported experimental tables cover accuracy, generated-token counts, and generation latency across code-generation and mathematical-reasoning benchmarks.These measures appear in the architecture-generalization and token-budget evaluations.Source: Table 9 · p. 15, Table 10 · p. 15, Table 11 · p. 15

Related papers

Loading 2508.05988v2…