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

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

cs.CLAug 14, 2026

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

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

Organizations: The University of Tokyo

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.

Figures & tables

Figure 1: (Left) Top : existing adaptive pruning ( Liu et al., 2025 ) selects neurons with an offline-calibrated threshold, so the realized sparsity ratio drifts under the aggregated activation distribution and performance degrades. Middle : top- $k$ selection is unaffected by that shift, but performing it at every decoding step incurs computational overhead that erases the speedup. Bottom : our method updates the pruning mask by top- $k$ selection periodically and carries importance across updates through an activation memory, retaining reasoning performance while maintaining computational efficiency. (Right) Activation visualization of $\mathbf{Z}$ in an FFN layer of DeepSeek-R1-Distill-Qwen-7B, where the blue horizontal line indicates the input boundary. Neurons with large activation magnitudes (highlighted in red ) exhibit a high re-firing ratio and short re-firing period, motivating our periodic mask update strategy.

Task Type	Batch Size	Static Pruning	Existing Adaptive Pruning	Ours
Non-reasoning	1	✓	✓	✓
Non-reasoning	$\geq$ 2	✓	✓	✓
Reasoning	1	✗	✓	✓
Reasoning	$\geq$ 2	✗	✗	✓ ∗

Table 1: Summary of our experiments for various pruning methods across inference settings. Our method ( ∗ ) uniquely maintains consistent performance across all settings, whereas existing approaches degrade on reasoning tasks or batched inference.

Figure 2: Overview of our batch-wise adaptive pruning approach. (a) Importance scores are computed per sample and aggregated into a single score vector, since a batch must share one pruning mask; we use an element-wise maximum as the default aggregation operator. (b) The method operates in three phases: input processing, dense exploration, and periodic adaptive pruning, with an activation memory carrying importance scores across successive mask updates.

Figure 3: Average accuracy across the five reasoning tasks with varying batch sizes at 50% target sparsity. (Left) DeepSeek-R1-Distill-Qwen-7B, (Right) DeepSeek-R1-Distill-Llama-8B.

Model	Method	Tasks					AVG
Model	Method	GSM8K	MATH500	MINERVA	AMC23	GPQA-DIAMOND	AVG
DeepSeek- R1-Distill- Qwen-7B	Dense	92.0	91.8	39.7	87.5	52.5	72.7
	Wanda	59.0	24.4	7.4	12.5	15.7	23.8
	Griffin	22.0	14.8	11.0	10.0	15.2	14.6
	TEAL	30.0	11.2	4.8	10.0	15.7	14.3
	Ours	89.0	71.0	29.4	50.0	30.8	54.0

Table 2: Pruning performance comparison at 50% target sparsity with batch size 4. Our method significantly outperforms all pruning baselines (Wanda, Griffin, TEAL) across multiple reasoning benchmarks on both DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B models. Bold values indicate the highest performance among all pruning methods.

Model	Median Period	Re-firing (%)
DeepSeek-R1 Distill-Qwen-7B	22.8 $\pm$ 1.9	76.3 $\pm$ 2.8
DeepSeek-R1 Distill-Llama-8B	20.8 $\pm$ 2.9	71.9 $\pm$ 2.2

Table 3: Median firing period and re-firing ratio of important neurons (

\bar{\mathbf{Z}}>0.05

). Mean

\pm

std over 10 samples

\times

5 benchmarks.

Figure 4: Accuracy on MATH500 across sparsity ratios with batch size 4 on DeepSeek-R1-Distill-Qwen-7B. (Left) Accuracy vs. target sparsity, (Right) Accuracy vs. actual sparsity. Target sparsity refers to the intended pruning ratio set before inference, while actual sparsity indicates the proportion of parameters effectively pruned during generation.

Figure 5: Generation throughput speedup vs. actual sparsity at batch size 4. (Left) DeepSeek-R1-Distill-Qwen-7B, (Right) DeepSeek-R1-Distill-Llama-8B.

Setting	DeepSeek-R1 Distill-Qwen-7B	DeepSeek-R1 Distill-Llama-8B
Max (Ours)	54.0	34.5
Mean	54.8	31.4
Memory (Ours)	54.0	34.5
No Memory	45.3	25.5

Table 4: Performance comparison of aggregation methods (mean/max) and activation memory (on/off) at 50% target sparsity with batch size 4. Values are averaged over 5 benchmarks. Bold indicates best per model.

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Ours	TEAL
DeepSeek-R1-Distill-Qwen-7B	2.2	96.8
DeepSeek-R1-Distill-Llama-8B	3.8	99.2
Qwen3-1.7B	7.4	98.8
Qwen3-8B	4.6	96.4

Appendix

Table 5: Rate of degenerate-repetition collapse (% of generations with word 4-gram repetition ratio

\geq 0.8

; lower is better) on MATH500 at 50% target sparsity with batch size 4.

Method	Last characters of the generation
TEAL	…Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11/36. Wait, 11 (cut off at the generation length limit; graded incorrect)
Ours	…we subtract this from 1: 1 - 25/36 = 36/36 - 25/36 = 11/36. Thus, the probability that the product of the two numbers is a multiple of 5 is \boxed{11/36}.

Appendix

Table 6: End of the generation for MATH500 problem 75 (level 4), “Two fair, 6-sided dice are thrown. What is the probability that the product of the two numbers is a multiple of 5?”, whose answer is 11/36, at 50% target sparsity with batch size 4 on DS-R1-Qwen-7B.

Model	Dense	Ours	R-KV + Ours
DeepSeek-R1-Distill-Qwen-7B	72.7	55.0	47.3
DeepSeek-R1-Distill-Llama-8B	71.3	37.6	31.4

Appendix

Table 7: Composition of our pruning method with R-KV KV-cache compression, at batch size 1 with a KV budget of 2048. Our method and the composition use 50% target sparsity. Values are average accuracy over the five reasoning benchmarks.

Model	Ours (BF16)	Ours (FP8)
DeepSeek-R1-Distill-Qwen-7B	54.0	54.3
DeepSeek-R1-Distill-Llama-8B	34.5	38.4

Appendix

Table 8: Composition of our pruning method with FP8 weight quantization at 50% target sparsity with batch size 4. Values are average accuracy over the five reasoning benchmarks.

Model	Method	Batch Size	Tasks					AVG
Model	Method	Batch Size	GSM8K	MATH500	MINERVA	AMC23	GPQA-DIAMOND	AVG
DeepSeek- R1-Distill- Qwen-7B	Dense	–	92.0	91.8	39.7	87.5	52.5	72.7
	Wanda	–	59.0	24.4	7.4	12.5	15.7	23.8
	Griffin	–	22.0	14.8	11.0	10.0	15.2	14.6
	TEAL	1	91.0	86.8	40.1	80.0	43.4	68.3
		4	30.0	11.2	4.8	10.0	15.7	14.3

Appendix

Table 9: Pruning performance comparison at 50% target sparsity across different batch sizes. Our method significantly outperforms all pruning baselines (Wanda, Griffin, TEAL). Dense, Wanda, and Griffin results are batch-size independent. Bold values indicate superior performance compared to TEAL at the same batch size.

Method	BS=1	BS=4	BS=8	BS=16
Dense	200.9 (1.00 $\times$ )	813.0 (1.00 $\times$ )	1596.0 (1.00 $\times$ )	3040.5 (1.00 $\times$ )
TEAL	290.7 (1.45 $\times$ )	971.1 (1.19 $\times$ )	1769.1 (1.11 $\times$ )	3335.9 (1.10 $\times$ )
GRIFFIN	323.5 (1.61 $\times$ )	1084.5 (1.33 $\times$ )	2111.0 (1.32 $\times$ )	3854.3 (1.27 $\times$ )
Ours	265.6 (1.32 $\times$ )	925.3 (1.14 $\times$ )	1794.7 (1.12 $\times$ )	3331.8 (1.10 $\times$ )

Appendix

Table 10: Generation throughput (tokens/sec) and speedup comparison on DeepSeek-R1-Distill-Qwen-7B at 50% target sparsity relative to Dense baseline across different batch sizes on NVIDIA H200 GPU.

Method	BS=1	BS=4	BS=8	BS=16
Dense	175.0 (1.00 $\times$ )	755.2 (1.00 $\times$ )	1471.7 (1.00 $\times$ )	2713.0 (1.00 $\times$ )
TEAL	263.8 (1.51 $\times$ )	837.2 (1.11 $\times$ )	1618.2 (1.10 $\times$ )	2997.4 (1.10 $\times$ )
GRIFFIN	269.5 (1.54 $\times$ )	1010.5 (1.34 $\times$ )	1957.1 (1.33 $\times$ )	3747.6 (1.38 $\times$ )
Ours	227.3 (1.30 $\times$ )	868.1 (1.15 $\times$ )	1685.9 (1.15 $\times$ )	3124.9 (1.15 $\times$ )

Appendix

Table 11: Generation throughput (tokens/sec) and speedup comparison on DeepSeek-R1-Distill-Llama-8B at 50% target sparsity relative to Dense baseline across different batch sizes on NVIDIA H200 GPU.

Figure 6: Relative generation throughput speedup vs. actual sparsity across batch sizes on NVIDIA H100 GPU. Each row corresponds to a different batch size. (Left) DeepSeek-R1-Distill-Qwen-7B, (Right) DeepSeek-R1-Distill-Llama-8B.

Figure 7: Throughput versus average accuracy trade-off on DeepSeek-R1-Distill-Qwen-7B across different batch sizes at 50% target sparsity on NVIDIA H200 GPU. Each subplot shows the speed-accuracy relationship for a specific batch size. Our method maintains consistent accuracy across batch sizes while achieving competitive throughput, whereas TEAL’s accuracy degrades significantly as batch size increases.

Figure 8: Throughput versus average accuracy trade-off on DeepSeek-R1-Distill-Llama-8B across different batch sizes at 50% target sparsity on NVIDIA H200 GPU. Each subplot shows the speed-accuracy relationship for a specific batch size. Our method maintains consistent accuracy across batch sizes while achieving competitive throughput, whereas TEAL’s accuracy degrades significantly as batch size increases.

Model	Dense MLP	Ours MLP	Overhead	Reduction
DeepSeek-R1-Distill-Qwen-7B	11.41B	6.27B	212.18K	45.0%
DeepSeek-R1-Distill-Llama-8B	11.27B	6.20B	183.51K	45.0%

Appendix

Table 12: Per-token FLOPs analysis. Overhead refers to adaptive mask computation.

Model	Method	HellaSwag	PIQA	COPA	ARC-E	ARC-C	BoolQ	Avg
DeepSeek-R1- Distill-Qwen-7B	Dense	59.7	70.2	70.0	67.1	40.3	77.6	64.1
	TEAL	60.0	69.6	71.0	65.3	38.7	74.9	63.3
	Ours	59.6	70.2	71.0	65.2	39.6	76.1	63.6
DeepSeek-R1- Distill-Llama-8B	Dense	74.8	77.0	90.0	69.9	40.7	83.5	72.6
	TEAL	74.7	77.1	87.0	69.3	39.9	82.1	71.7
	Ours	74.7	76.7	88.0	68.8	40.9	82.3	71.9

Appendix

Table 13: Non-reasoning benchmark performance at 50% target sparsity. Bold indicates best among pruning methods.

Model	Median Period	Re-firing (%)
DeepSeek-R1-Distill-Qwen-7B	22.8 $\pm$ 1.9	76.3 $\pm$ 2.8
DeepSeek-R1-Distill-Llama-8B	20.8 $\pm$ 2.9	71.9 $\pm$ 2.2
Qwen3-1.7B	11.8 $\pm$ 0.8	80.7 $\pm$ 3.1
Qwen3-8B	17.8 $\pm$ 1.3	76.8 $\pm$ 2.9

Appendix

Table 14: Periodicity analysis of important neuron activations across all evaluated models (mean

\pm

std across 10 samples

\times

5 benchmarks).

Setting	GSM8K	MATH500	MINERVA	AMC23	GPQA-DIAMOND	AVG
DeepSeek-R1-Distill-Qwen-7B
Max (Ours)	89.0	71.0	29.4	50.0	30.8	54.0
Mean	83.0	74.0	29.4	57.5	30.3	54.8
Memory (Ours)	89.0	71.0	29.4	50.0	30.8	54.0
No Memory	76.0	63.6	21.7	37.5	27.8	45.3
DeepSeek-R1-Distill-Llama-8B

Appendix

Table 15: Per-benchmark performance comparison of aggregation methods (mean/max) and activation memory (on/off) at 50% target sparsity with batch size 4. Bold indicates best per model.

$T_{\text{init}}$	$T_{E}$	$T_{\text{trans}}$	Accuracy	Speed (token/sec)
0	2	20	68.6	925.3
32	2	20	70.8	925.3
64	2	20	71.0	925.3
128	2	20	71.4	925.3
64	1	20	67.4	941.0
64	2	20	71.0	925.3

Appendix

Table 16: Hyperparameter ablation study on DeepSeek-R1-Distill-Qwen-7B using the MATH500 dataset with batch size 4 on NVIDIA H200 GPU, examining the effects of

T_{\text{init}}

(initial dense steps),

T_{E}

(exploration steps), and

T_{\text{trans}}

(update period) on accuracy and inference speed.

Model	Method	GSM8K	MATH500	MINERVA	AMC23	GPQA-DIAMOND	AVG
DeepSeek-R1- Distill-Qwen-7B	Dense	92.0	91.8	39.7	87.5	52.5	72.7
	Ours (top- $k$ )	89.0	71.0	29.4	50.0	30.8	54.0
	Ours (threshold)	3.0	3.2	2.9	0.0	1.5	2.1
DeepSeek-R1- Distill-Llama-8B	Dense	96.0	91.0	33.8	90.0	45.5	71.3
	Ours (top- $k$ )	75.0	38.8	12.1	25.0	21.7	34.5
	Ours (threshold)	3.0	2.8	0.7	2.5	0.0	1.8

Appendix

Table 17: Threshold-based pruning at 50% target sparsity with batch size 4. The activation threshold is calibrated on C4.

Model	Method	GSM8K	MATH500	MINERVA	AMC23	GPQA-DIAMOND	AVG
DeepSeek-R1- Distill-Qwen-7B	Dense	92.0	91.8	39.7	87.5	52.5	72.7
	TEAL	30.0	11.2	4.8	10.0	15.7	14.3
	Ours	89.0	71.0	29.4	50.0	30.8	54.0
DeepSeek-R1- Distill-Llama-8B	Dense	96.0	91.0	33.8	90.0	45.5	71.3
	TEAL	28.0	5.4	2.2	0.0	16.2	10.4
	Ours	75.0	38.8	12.1	25.0	21.7	34.5

Appendix

Table 18: Pruning performance comparison at 50% target sparsity with batch size 4. Bold indicates best among pruning methods.

Figure 9: Visualization of feedforward activations $\mathbf{Z}$ across all layers ( $\bar{\mathbf{Z}}>0.05$ ) of DeepSeek-R1-Distill-Qwen-7B on MATH500. The x-axis represents neurons and the y-axis represents token indices from top to bottom. The blue horizontal line indicates the last position of the prompt.

Figure 10: Visualization of feedforward activations $\mathbf{Z}$ across all layers ( $\bar{\mathbf{Z}}>0.05$ ) of DeepSeek-R1-Distill-Qwen-7B on GPQA-DIAMOND. The x-axis represents neurons and the y-axis represents token indices from top to bottom. The blue horizontal line indicates the last position of the prompt.

Figure 11: Visualization of feedforward activations $\mathbf{Z}$ across all layers ( $\bar{\mathbf{Z}}>0.05$ ) of DeepSeek-R1-Distill-Llama-8B on MATH500. The x-axis represents neurons and the y-axis represents token indices from top to bottom. The blue horizontal line indicates the last position of the prompt.

Figure 12: Visualization of feedforward activations $\mathbf{Z}$ across all layers ( $\bar{\mathbf{Z}}>0.05$ ) of DeepSeek-R1-Distill-Llama-8B on GPQA-DIAMOND. The x-axis represents neurons and the y-axis represents token indices from top to bottom. The blue horizontal line indicates the last position of the prompt.

LLM Pruning

Large language model pruning removes less useful weights or components to reduce memory and inference costs while preserving model capabilities.

See all 168 papers

Oct 14, 2026

To Reason but Not to Stop: Understanding and Recovering Pruned Large Reasoning Models

Weight pruning can degrade the accuracy of large reasoning models (LRMs), yet it remains unclear whether it damages reasoning itself. We find that much of this degradation stems from a failure to stop rather than to reason. As sparsity increases, pruned LRMs increasingly fail to terminate within the token budget, often after the correct answer has appeared in their reasoning, instead falling into repetitive loops, even though they still recognize where the dense model would stop. We further find that the hidden-state difference between terminating and looping continuations is consistent across pruned checkpoints, suggesting a shared direction associated with committing to an answer. Building on this, we propose Termination-Guided Activation Steering (TGAS), a training-free method that adds this direction to the residual stream at inference time. Across six LRMs, three pruning pipelines, and math and code benchmarks, TGAS improves termination and overall accuracy while reducing repetitive reasoning and generation length, showing that this failure is actionable.

Minh-Anh Nguyen, Nguyen T. Ha Lan, Xiyan Fu +2

Nanyang Technological University, Singapore · VinUniversity, Vietnam

Repetition in Language Generation LLM Reasoning Large Reasoning Models LLM Pruning Activation Steering

Ask PDF

Oct 13, 2026

Unlocking Structured Pruning for Reasoning Language Models via Attention-Preserving MLP Trimming

Structured pruning reduces LLM inference costs by removing model components to produce smaller dense models. However, existing methods can cause reasoning performance to collapse on long-form reasoning tasks such as AIME. We propose LayerTrim, a training-free structured pruning method that identifies redundant decoder layers and selectively trims their MLP features while preserving attention and decoder depth. LayerTrim is the first training-free structured pruning method designed to preserve long-form reasoning without gradients or recovery training. Across three reasoning language models and five benchmarks, LayerTrim outperforms existing structured pruning methods in most evaluated settings, improving average accuracy by up to 21.4 percentage points over RESP. On Qwen3-14B, LayerTrim achieves 1.24

\times

faster prefill and 19.7% lower peak GPU memory than the dense model, while better preserving reasoning ability. It also achieves the shortest average generation length among evaluated pruned models, while retaining the highest average accuracy. These results demonstrate that LayerTrim improves reasoning efficiency through both structural compression and reduced token generation.

Woosang Lim, Haneul Yoo, Vincent-Daniel Yun +1

Seoul National University, Republic of Korea · New York University, United States · University of Southern California, United States

LLM Pruning LLM Compression Efficient Language Model Reasoning LLM Inference Acceleration

Ask PDF

Jul 13, 2026

OS-Pruner: Pruning Chains-of-Thought of Reasoning Models via Optimal Stopping

Large Language Models (LLMs) have achieved remarkable success in complex reasoning tasks through Chain-of-Thought (CoT) prompting. However, these models often exhibit "computational overthinking," generating redundant reasoning steps that increase latency and cost without improving accuracy. Recent studies suggest that CoT trajectories can be significantly pruned, yet existing methods often rely on forcing a static thinking budget, heuristic filtering, sub-optimal early exit via classification, or expensive re-training. In this paper, we introduce OS-Pruner, a lightweight plug-in framework that formulates chain-of-thought pruning as an optimal stopping problem. Given a reasoning prefix, OS-Pruner learns whether further reasoning is worth its token cost by optimizing an explicit utility that trades off final-answer accuracy against generated length. Our novel formulation enables the model to dynamically assess the sufficient point of termination for a reasoning chain. OS-Pruner is designed to be lightweight during both training and inference, and to provide users with fine-grained control over the reasoning-effort vs. accuracy trade-off. On diverse reasoning benchmarks and base models, OS-Pruner achieves 20-60% reduction in generation length with minimal accuracy sacrifice.

Mohammed Ehab, Aymane El Gadarri, Vivek F. Farias +2

MIT · Columbia Business School

Overthinking in Language Models Sequential Decision Making Optimal Stopping Efficient Language Model Reasoning Efficient Language Model Inference Large Reasoning Models

Ask PDF

Abstract

LLM Pruning

Explore similar work

To Reason but Not to Stop: Understanding and Recovering Pruned Large Reasoning Models

Unlocking Structured Pruning for Reasoning Language Models via Attention-Preserving MLP Trimming

OS-Pruner: Pruning Chains-of-Thought of Reasoning Models via Optimal Stopping