PaperlayerView exact revision on arXiv

Source-linked AI summary

Batch-wise Adaptive Pruning: Periodic Neuron Activation-Aware Weight Pruning for Language Reasoning Model

Yongmin Kim, Shota Takashiro, Yusuke Iwasawa, Takeshi Kojima, Yutaka Matsuo

arXiv:2608.14003v1cs.CL

TL;DR

Batched inference exposes a failure in existing adaptive pruning methods for reasoning models because shared masks and calibrated thresholds do not match aggregated activations. The paper introduces periodic top-k selection with activation memory, outperforming the prior state of the art by 39.7 percentage points in average accuracy at batch size 4 and achieving 1.40× speedup at 50% actual sparsity.

  • Problem

    Existing adaptive pruning methods degrade under batched inference because shared masks require aggregation, shifting activations away from offline-calibrated threshold distributions.

  • Method

    The method uses periodic top-k selection over aggregated importance scores and activation memory to retain neurons that re-activate across reasoning phases.

  • Results

    Across multiple reasoning benchmarks, the method outperforms the previous state-of-the-art adaptive pruning method by 39.7 percentage points in average accuracy at batch size 4 with 50% target sparsity.

  • Takeaways & Limitations

    The results establish batch-wise adaptive pruning as a viable approach for retaining accuracy in high-sparsity batched reasoning inference.

  • Takeaways & Limitations

    Below roughly 40% target sparsity, TEAL remains more accurate on MATH500 and faster on DS-R1-Qwen-7B, while it also leads at batch size 1.

Abstract

from arXiv · show

Large Reasoning Models (LRMs) achieve strong performance on complex tasks through extended chain-of-thought generation, but incur substantial computational costs during inference. In production settings, batched inference is essential for high throughput, yet the existing training-free adaptive pruning methods we evaluate severely degrade in this regime. Because a batch must share a single pruning mask, these methods aggregate activations across samples and then apply threshold-based selection; the threshold, calibrated offline on unaggregated activations, no longer matches the aggregated distribution, so the realized sparsity ratio drifts and accuracy on reasoning tasks collapses under batched inference. In this work, we propose a training-free adaptive pruning method designed specifically for batched inference in LRMs, built on two components. First, we replace threshold-based selection with periodic top-k selection over the aggregated importance scores, which is unaffected by the shift that aggregation induces in the activation distribution, and which runs selection once per update period rather than at every token, preserving the speedup. Second, based on the observation that important neurons re-fire periodically during long reasoning generation, we introduce an activation memory that accumulates importance across update phases so that recurring neurons are retained. Experiments on diverse reasoning benchmarks demonstrate that our method outperforms the previous state-of-the-art adaptive pruning method by 39.7 percentage points in average accuracy at batch size 4 with 50% target sparsity on DeepSeek-R1-Distill-Qwen-7B, and reaches 1.40x speedup over dense inference at 50% actual sparsity.

1 Introduction

Large Reasoning Models use extended chain-of-thought generation but incur substantial inference costs, and existing pruning methods degrade under batched reasoning inference. The proposed training-free method uses periodic top-k mask updates and activation memory to preserve performance while reducing computation.

  • Motivation: Large Reasoning Models solve complex problems with extended chain-of-thought processes, resulting in substantial computational costs.This motivates pruning to reduce inference cost.
  • Method: The method replaces threshold selection with periodic top-k selection over aggregated importance scores, avoiding distribution-shift effects while reducing selection overhead.Top-k is performed once per update period rather than at every decoding step.
  • Batched inference limitation: Existing adaptive pruning methods degrade in batched inference because one shared mask requires aggregated activations, shifting the distribution relative to offline-calibrated thresholds.Threshold selection therefore causes realized sparsity drift and reasoning-task accuracy collapse.
  • Method: Activation memory accumulates importance across update phases so neurons that re-fire periodically during long reasoning generation remain retained.The memory complements periodic mask updates by preserving recurring important neurons.
  • Results: 39.7 percentage points higher average accuracy is achieved than the previous state-of-the-art adaptive pruning method at batch size 4 with 50% target sparsity on DeepSeek-R1-Distill-Qwen-7B.The method also achieves 1.40× speedup over dense inference at 50% actual sparsity.

2 Related Work

Related pruning methods differ in granularity, input dependency, and adaptivity, with trade-offs involving hardware requirements, batch compatibility, and responsiveness to changing activations. The proposed method addresses batched-inference limitations using periodic top-k selection and activation memory.

  • Method Taxonomy: Pruning methods are categorized by granularity, input dependency, and adaptivity, with Table 1 summarizing differences from the proposed approach.These dimensions organize comparisons among existing pruning methods and ours.
  • Pruning Granularity: Unstructured and semi-structured pruning can require specialized hardware, whereas structured pruning typically requires post-training.Unstructured pruning removes individual weights, while semi-structured pruning includes N:M sparsity.
  • Input Dependency: Input-independent pruning is batch-compatible but cannot exploit input-specific activations, whereas input-dependent pruning dynamically selects neurons from runtime activations but introduces challenges.The supplied passage truncates the specific challenges associated with input-dependent methods.
  • Adaptivity: Static pruning reuses one fixed mask throughout generation, while adaptive pruning updates masks during decoding to respond to changing activation patterns.Static pruning can degrade on reasoning models with long outputs; existing adaptive methods use threshold-based selection that suffers under batched aggregation.
  • Proposed Approach: Periodic top-k selection and activation memory address batched-inference limitations by updating masks periodically and carrying importance across mask updates.The method is contrasted with ProSparse and other individual-sample approaches that make batched inference challenging.

3 Method

The method is a training-free, real-time pruning approach for batched LLM inference that targets gated-MLP neurons using activation-derived importance scores. It combines batch aggregation, top-k shared-mask selection, and alternating sparse–dense phases with activation memory for periodic adaptation.

  • Method overview: The training-free method makes pruning decisions in real time from model activations observed during inference, targeting gated-MLP structures.Gated MLPs account for most parameters in modern Transformer-based LLMs and are the pruning target.
  • Importance scoring: Importance scores aggregate row-wise ℓ2-normalized activation magnitudes across tokens to capture each neuron’s relative task importance.Normalization prevents tokens with larger overall activation magnitudes from dominating the scores.
  • Batch aggregation and mask selection: For batched inference, element-wise maximum aggregation preserves a neuron’s high score for any sample, and top-k selection retains k = ⌊(1 −ρ) · DFF⌋ neurons under one shared mask.Top-k depends on the relative ordering of aggregated scores rather than their absolute scale; padding and EOS tokens are excluded.
  • Generation phases: The three phases are input scoring, dense exploration for Tinit steps, and periodic adaptive pruning that alternates sparse pruning with dense activation-collection exploration.Dense exploration establishes the initial reasoning direction and gathers statistics reflecting generation context.
  • Structured sparse computation: Structured pruning retains I = {i : M[i] = 1} with |I| = (1 −ρ) · DFF, reducing gated-MLP matrices to k neuron dimensions for smaller GPU computations.The retained rows produce pruned matrices c Wg, c W1 ∈Rk×D and c W2 ∈RD×k.

4 Experiment

Experiments on diverse reasoning benchmarks show that the proposed batch-wise adaptive pruning method substantially outperforms training-free baselines under batched inference, while preserving performance and achieving practical speedups across sparsity levels. Ablations further show that periodic activation memory improves robustness by retaining neurons that re-activate during reasoning.

  • Evaluation setup: The method is evaluated on DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B using diverse mathematical and scientific reasoning benchmarks.Additional evaluations on Qwen3-1.7B and Qwen3-8B are reported in Appendix P.
  • Batched performance: At 50% target sparsity and batch size 4, accuracy reaches 54.0% versus TEAL’s 14.3% on DS-R1-Qwen-7B, a 39.7-point improvement.The proposed method significantly outperforms Wanda, Griffin, and TEAL across multiple reasoning benchmarks on both evaluated models.
  • Batched performance: At this operating point, TEAL produces degenerate repetition loops on 96–99% of MATH500 samples, compared with 2.2–7.4% for the proposed method.The proposed method otherwise maintains coherent chain-of-thought generation.
  • Batch-size robustness: Across batch sizes at 50% target sparsity, the proposed method maintains robust accuracy while TEAL degrades as aggregated activations diverge from single-sample calibration.The degradation is attributed to distribution shift caused by batch aggregation.
  • Speed and sparsity: At approximately 50% actual sparsity, throughput reaches 1.40× on DS-R1-Qwen-7B and 1.29× on DS-R1-Llama-8B at batch size 4.Speed is measured end-to-end on an NVIDIA H100 GPU using the GPT-Fast benchmarking setup.
  • Ablations: Enabling activation memory improves average accuracy by 9.0 points for DS-R1-Llama-8B and 8.7 points for DS-R1-Qwen-7B.The comparison is against a memory-free variant that reconstructs the pruning mask independently at every transition step.

5 Conclusion

The paper presents a training-free batch-wise adaptive pruning method for batched inference in LRMs after finding that existing pruning baselines degrade severely on reasoning tasks when masks are shared. Experiments show improved accuracy in high-sparsity batched inference, with speedups continuing to grow as sparsity increases.

  • The work presents a training-free batch-wise adaptive pruning method for batched inference in LRMs.
  • 58.4 to 67.7 points of average accuracy are lost by TEAL relative to the dense model at 50% target sparsity under batched inference.Adaptive methods work at batch size 1 but collapse once a mask is shared, while static methods degrade at any batch size.
  • The method retains substantially more accuracy than prior methods in the high-sparsity batched regime, while speedup keeps growing with sparsity.Activation-skipping methods instead saturate in speedup as sparsity increases.

A Limitations

The method has important deployment and performance limitations: it requires exposed intermediate activations, is less competitive at low sparsity or batch size 1, and still trails dense inference at high sparsity. These constraints reflect both access requirements and overhead or accuracy challenges in batched reasoning.

  • Model access: The method requires direct access to intermediate activations, limiting applicability to closed-source or API-based models.Its internal states must be exposed for activation-aware pruning.
  • Low-sparsity regime: Below roughly 40% target sparsity, TEAL remains ahead in MATH500 accuracy, and its activation skipping is faster on DS-R1-Qwen-7B.Periodic updates and top-k selection add overhead that low sparsity does not yet justify.
  • Low-sparsity regime: At batch size 1, TEAL remains ahead because per-token threshold-mask updates are essentially free, whereas top-k selection imposes overhead.This limitation arises when the batching advantage does not offset the selection cost.
  • Dense-model gap: 54.0 versus 72.7 average accuracy at 50% target sparsity and batch size 4 shows a remaining dense-model gap on DS-R1-Qwen-7B.The averages cover five reasoning benchmarks.
  • Dense-model gap: 34.5 versus 71.3 average accuracy under the same setting shows a dense-model gap on DS-R1-Llama-8B.The paper attributes these gaps to the intrinsic difficulty of maintaining accuracy in batched reasoning at high sparsity.

B Reasoning Trajectory Quality … E Full Performance Comparison

At 50% target sparsity, the method avoids the severe reasoning-trajectory failures observed with TEAL under batched inference. It also composes with KV-cache compression and FP8 quantization, while maintaining stronger performance across batch sizes.

  • B Reasoning Trajectory Quality: 96.4–99.2% of TEAL generations collapse into degenerate repetition, versus below 10% with our method at batch size 4.Degeneracy is defined as a word 4-gram repetition ratio of at least 0.8 on MATH500 generations.
  • B Reasoning Trajectory Quality: 28,560 characters: TEAL loops until the length limit and is graded incorrect, whereas our method closes the same derivation in 13,103 characters.The example is MATH500 problem 75 on DS-R1-Qwen-7B at 50% target sparsity with batch size 4.
  • C Composition with KV-Cache Compression: R-KV and our method prune along different axes, enabling simultaneous use to combine KV-cache memory savings with model-weight pruning.R-KV targets KV-cache compression, while our method prunes neurons in model weights and accelerates inference independently of the serving environment.
  • C Composition with KV-Cache Compression: R-KV + Ours achieves accuracy close to our method alone while adding R-KV’s KV-cache memory savings.The composition’s accuracy drop is attributed to R-KV’s accuracy–memory trade-off; results are reported at batch size 1 with a KV budget of 2048.
  • D Composition with Weight Quantization: 3.9 points: FP8 quantization improves accuracy when composed with pruning on DS-R1-Llama-8B, while imposing essentially no accuracy cost on DS-R1-Qwen-7B.The composition is evaluated at 50% target sparsity with batch size 4.
  • E Full Performance Comparison: 68.3% and 63.5%: TEAL performs strongly at batch size 1 on DS-R1-Qwen-7B and DS-R1-Llama-8B, respectively, but degrades sharply as batch size increases.The passage attributes this degradation to distribution shift from aggregating activations across multiple samples.
  • E Full Performance Comparison: 14.3%, 4.1%, and 6.8%: TEAL accuracy on DS-R1-Qwen-7B falls to these values at batch sizes 4, 8, and 16; DS-R1-Llama-8B falls to 10.4%, 2.4%, and 1.3%.These are pruning comparisons at 50% target sparsity across batch sizes, where the method outperforms pruning baselines including TEAL.

F Generation Throughput at 50% Target Sparsity … J Computational Overhead

At 50% target sparsity, throughput depends strongly on batch size: fixed-mask Griffin leads some settings, while the proposed method becomes more competitive as batching increases and preserves accuracy better than TEAL. The evaluation covers specified reasoning datasets and models, while adaptive mask computation adds negligible FLOPs relative to dense MLP computation.

  • F Generation Throughput at 50% Target Sparsity: Griffin achieves the highest speedup across all settings, including 1.61× at BS=1 on DS-R1-Qwen-7B, but its reasoning-task accuracy is severely degraded.At BS=1, TEAL exceeds the proposed method on DS-R1-Qwen-7B, with 1.45× versus 1.32× speedup.
  • F Generation Throughput at 50% Target Sparsity: At BS=16 on DS-R1-Qwen-7B, TEAL, Griffin, and the proposed method achieve 1.10×, 1.27×, and 1.10× speedup, respectively.The H200 experiments compare generation throughput at 50% target sparsity relative to dense inference.
  • H Speed-Accuracy Trade-off: At BS=16, TEAL’s average accuracy drops to 6.8% on DS-R1-Qwen-7B and 1.3% on DS-R1-Llama-8B, whereas the proposed method maintains consistent accuracy across batch sizes.At BS=1, TEAL records 68.3% and 63.5% average accuracy on the Qwen and Llama models, respectively.
  • I.1 Dataset Details: The study evaluates TinyGSM8K, MATH500, MINERVA Math, AMC23, and GPQA-DIAMOND using their reported test splits and example counts.The listed licenses are MIT for the first four datasets except AMC23, whose license is unspecified, and CC BY 4.0 for GPQA-DIAMOND.
  • I.2 Model Details: Evaluation uses DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Llama-8B, Qwen3-1.7B, and Qwen3-8B under zero-shot greedy decoding with maximum generation length of 16,000 tokens.The DeepSeek models use the official DeepSeek-R1 chat template, while Qwen3 models use the Qwen3 template with enable thinking=True.
  • J Computational Overhead: The proposed method’s adaptive mask computation adds less than 0.002% of dense MLP FLOPs while reducing approximately 45% of MLP computation.All per-token FLOPs measurements are conducted in Float16.

K Non-Reasoning Task Performance

At 50% target sparsity, the method largely preserves non-reasoning performance relative to dense inference across four models. It also outperforms TEAL on average, with negligible degradation on shorter-output tasks.

  • Model results: 0.5% average accuracy gap is reported for DS-R1-Qwen-7B, with 63.6% versus 64.1% for dense inference.This is the smallest reported gap among the four models.
  • Model results: 1.2% average accuracy gap is reported for Qwen3-1.7B, with 65.2% versus 66.4% for dense inference.The other reported gaps are 0.7% for DS-R1-Llama-8B and 1.0% for Qwen3-8B.
  • Comparison with TEAL: The method outperforms TEAL on average across all four models, while both methods show negligible degradation relative to dense inference on shorter-output tasks.These results concern non-reasoning benchmarks at 50% target sparsity.

L Extended Periodicity Analysis

The extended periodicity analysis covers all evaluated models and sets model-specific transition periods from observed median activation periods. It uses Tinit = 64 and TE = 2 throughout.

  • L Extended Periodicity Analysis: Observed median activation periods yield Ttrans = 20 for DS-R1-Qwen-7B, DS-R1-Llama-8B, and Qwen3-8B, versus Ttrans = 10 for Qwen3-1.7B.These transition periods are set using the extended analysis across all evaluated models.
  • L Extended Periodicity Analysis: The periodicity analysis reports mean ± std across 10 samples × 5 benchmarks for all evaluated models.This extends the main-text analysis presented in Table 3.

M Extended Ablation

The ablation compares max versus mean batch aggregation and activation memory at 50% target sparsity with batch size 4. Aggregation strengths are task-dependent, with only a 1.2-point cross-model average gap, insufficient to establish a universally superior operator.

  • Activation memory: The ablation also evaluates activation memory by comparing memory on versus off.The supplied passage identifies this comparison but does not provide its complete result.
  • Aggregation methods: Max outperforms mean on GSM8K (82.0 vs. 73.5) and AMC23 (37.5 vs. 35.0), whereas mean leads on MATH500 (58.1 vs. 54.9) and GPQA (27.8 vs. 26.3).The aggregation winner varies by benchmark in the cross-model average.
  • Aggregation methods: 1.2 points separate the cross-model averages of max and mean (44.3 vs. 43.1), which is insufficient to claim either operator is superior.The per-model winner is likewise split.
  • Ablation setup: Table 15 evaluates max versus mean aggregation and activation memory on or off at 50% target sparsity with batch size 4.The table reports per-benchmark performance, with bold marking the best per model.

N Hyperparameter Ablation … R Use of Large Language Models

The ablations identify accuracy–throughput trade-offs among initialization, exploration, and update periods, while threshold pruning fails severely. Across additional models, periodic pruning remains effective because important neurons re-fire, and language models were used only for polishing, copy-editing, and code assistance.

  • N Hyperparameter Ablation: Increasing Tinit from 0 to 128 improves accuracy by 2.8 points with minimal speed impact, capturing post-prompt activation shifts.Tinit primarily affects accuracy without speed penalty.
  • N Hyperparameter Ablation: Increasing TE from 1 to 4 yields a 9.8-point accuracy gain but reduces throughput from 941.0 to 896.8 tokens/sec.Extending Ttrans from 10 to 30 instead degrades accuracy by 7.8 points while increasing throughput from 806.3 to 958.7 tokens/sec.
  • N Hyperparameter Ablation: 88.2 accuracy on DS-R1-Qwen-7B and 77.6 on DS-R1-Llama-8B are achieved by Tinit = 0, TE = 1, Ttrans = 2 on MATH500.This setting recomputes the top-k mask every other step, incurs selection costs more often, and is not adopted.
  • O Threshold-Based Pruning: Near-zero accuracy across all benchmarks results from fixed thresholds calibrated on 20,248 C4 tokens at 50% target sparsity, confirming static thresholds are inadequate.The threshold-based approach severely over-prunes both models.
  • P Extended Pruning Performance: 44.0% average accuracy versus TEAL’s 2.2% on Qwen3-1.7B represents a 41.8-point improvement at 50% target sparsity and batch size 4.On Qwen3-8B, the method reaches 52.7% versus TEAL’s 7.2%, a 45.5-point improvement.
  • P Extended Pruning Performance: 0.0% accuracy on AMC23 and GPQA for Qwen3-1.7B illustrates TEAL’s near-complete collapse, whereas the proposed method retains meaningful accuracy across benchmarks.Batch size 4 results for DS-R1-Qwen-7B and DS-R1-Llama-8B are also included in the full comparison.
  • Q Extensive Activation Visualization: Visualizations across MATH500 and GPQA-DIAMOND demonstrate periodic re-firing of important neurons during autoregressive generation, supporting periodic mask updates.The visualizations cover DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B across all layers.
  • R Use of Large Language Models: ChatGPT (“GPT-5” and Anthropic Claude) assisted with English-language polishing, light copy-editing, and experimental code implementation, not research ideas or experimental design.The authors authored and verified the technical content, claims, and experimental code.
Loading 2608.14003v1…