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

2026-08-14

AI Summary

This study addresses the accuracy collapse inherent in adaptive pruning during batch inference for large reasoning models by proposing a training-free, batch-aware adaptive pruning method. Innovatively replacing conventional threshold-based approaches with periodic Top-k selection, the framework incorporates an activation memory mechanism to cumulatively retain critical neurons, thereby effectively mitigating aggregated activation distribution shifts. Experimental results demonstrate that at 50% sparsity, this approach outperforms existing state-of-the-art methods by 39.7% in accuracy while achieving a 1.4× inference speedup. Consequently, the proposed method significantly enhances both efficiency and accuracy in high-sparsity batch inference scenarios, offering a robust solution to the trade-off between model compression and performance degradation in large-scale reasoning tasks.

Abstract

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.

Problem

• Large Reasoning Models
• Batched Inference
• Adaptive Pruning
• Activation Distribution Shift
• Training-free Compression

Innovation

• Batch-wise Adaptive Pruning
• Periodic Top-k Selection
• Activation Memory
• Training-free Pruning
• Large Reasoning Models