💬 NLP

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

This paper proposes a training-free, batch-wise adaptive pruning method for Large Reasoning Models that utilizes periodic top-k selection and an activation memory mechanism to overcome the accuracy degradation of existing methods in batched inference, achieving significant speedups while maintaining high reasoning performance.

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

Published 2026-08-17

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a really hard puzzle, like a complex math problem or a tricky logic riddle. You have a brilliant but very hungry assistant (a Large Reasoning Model) who can think through the steps to find the answer. The catch? This assistant is so thorough that it writes out a massive, step-by-step diary of its thoughts before giving you the final answer. While this "chain of thought" makes the assistant incredibly smart, it also makes it slow and expensive to run, like trying to power a supercomputer with a single AA battery.

To make this assistant faster, scientists often try to "prune" it—basically, they tell the assistant to ignore certain parts of its brain (neurons) that it doesn't seem to need at the moment. Think of it like a chef who, while cooking a huge banquet, decides to stop chopping vegetables for dishes that haven't been ordered yet. The problem is, when you try to cook for a whole crowd at once (batched inference), the old methods of deciding what to chop fall apart. They use a fixed rule that works for one person but gets confused when applied to a group, causing the assistant to forget how to think and start repeating nonsense. This paper introduces a new, smarter way to decide what to chop, ensuring the assistant stays sharp even when serving a crowd.

The Problem: The "One-Size-Fits-All" Mistake

Imagine you are a conductor leading an orchestra. If you have just one violinist, you can easily tell them, "Play loudly now, then play softly later." But if you have a whole section of violinists playing together (a batch), and you give them all the exact same instruction based on what one violinist usually does, the music might turn into a mess.

This is exactly what happens with current AI models when they try to process multiple requests at once. Existing methods for speeding up these models use a "threshold" rule. It's like saying, "If a neuron's activity is above 5, keep it; if it's below 5, turn it off." This works fine when the AI is thinking alone. But when you feed it a batch of different questions, the "activity levels" of the neurons get mixed together, shifting the average. The old rule (the threshold of 5) no longer matches the new reality. The result? The AI accidentally turns off the wrong neurons, its "brain" goes fuzzy, and it starts failing at reasoning tasks. The paper shows that when you try to run these models on a batch of 4 requests, the best existing methods lose nearly all their smarts, dropping their accuracy by huge margins (sometimes over 50 points).

The Solution: A Periodic "Top-K" Strategy with a Memory

The authors propose a new method called Batch-wise Adaptive Pruning. Instead of using a rigid "pass/fail" line (a threshold), they use a smarter, more flexible approach with two main tricks.

First, they stop using the fixed line and start using a "Top-K" selection. Imagine you have a group of 100 runners (neurons) and you only have room for 50 on the bus. Instead of saying, "Only run if you are faster than 10 seconds," you simply say, "The 50 fastest runners get on the bus." This doesn't matter if the runners are generally faster or slower today; you always pick the best 50. This solves the problem of the "shifting distribution" that broke the old methods.

Second, they realized that important neurons don't just fire once and disappear; they have a rhythm. The paper observed that during long reasoning tasks, the most important neurons tend to "re-fire" (wake up and do work) in a regular pattern, roughly every 22.8 tokens (chunks of text). To catch this rhythm, the new method updates its pruning mask only periodically (every 20 steps) rather than every single word. This saves time because the computer doesn't have to stop and recalculate which neurons to keep at every tiny moment.

But here is the cleverest part: they added an Activation Memory. Think of this as a "highlight reel" or a sticky note. If a neuron was super important in the first batch of thoughts, the memory keeps a record of that. Even if the neuron is quiet for a few steps, the memory ensures it doesn't get kicked out of the bus just because it took a short nap. This way, the AI retains the neurons that are consistently important over time, preventing it from forgetting the logic it started with.

The Results: Fast, Smart, and Ready for the Crowd

The paper tested this new method on powerful reasoning models (specifically DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B) using difficult math and science benchmarks. The results were striking.

When running with a batch size of 4 (processing four questions at once) and aiming to cut the model size by 50%:

The old best method (TEAL) crashed, achieving only about 14.3% average accuracy.
The new method maintained a strong 54.0% average accuracy.
This is a massive improvement of 39.7 percentage points.

Furthermore, the new method actually made the models faster. By cutting out the unnecessary work, it achieved a 1.40x speedup compared to running the full, unpruned model. The paper notes that this speedup is especially helpful when the computer is busy (compute-bound), which happens when you are processing many requests at once.

What This Means

The authors are careful to note that this method is "training-free," meaning it doesn't require re-teaching the model; it just tweaks how the model runs in real-time. They also point out that while other methods might work for simple tasks, they fail spectacularly on complex reasoning when batched. This new approach, by using a periodic update and a memory of what's important, keeps the AI's reasoning sharp while making it efficient enough to handle real-world workloads where many users ask questions at the same time. It's a practical fix that lets these giant, smart models run faster without losing their ability to think deeply.

Technical Summary: Batch-wise Adaptive Pruning for Large Reasoning Models

Problem Statement
Large Reasoning Models (LRMs) achieve high performance on complex tasks through extended chain-of-thought generation but incur substantial computational costs. In production environments, high throughput requires batched inference (processing multiple samples simultaneously). However, existing training-free adaptive pruning methods, such as TEAL and CATS, fail significantly in this regime.

The core failure mechanism is twofold:

Distribution Shift: Existing methods use a single pruning mask for the entire batch. They aggregate per-sample activations (e.g., via mean) and apply a threshold calibrated offline on single-sample distributions. This aggregation shifts the activation distribution, causing the threshold to no longer match the data, leading to a drift in the realized sparsity ratio and a collapse in reasoning accuracy.
Static vs. Dynamic Trade-off: Static pruning methods (fixed patterns) cannot track evolving activation patterns during long reasoning chains. Conversely, selecting a new mask at every token step is computationally prohibitive, negating any speedup.

Methodology
The authors propose Batch-wise Adaptive Pruning, a training-free method designed specifically for batched inference in LRMs. The approach relies on two primary components:

Periodic Top-k Selection:
- Instead of using a calibrated threshold (which is sensitive to distribution shifts), the method employs top-k selection on aggregated importance scores. Top-k selection depends only on the relative ordering of scores, making it robust to the distribution shifts caused by batch aggregation.
- To avoid the computational overhead of updating the mask at every token, the pruning mask is updated periodically (every $T_{trans}$ steps) rather than at every decoding step.
- Batch Aggregation: For a batch of $B$ samples, per-sample importance scores are aggregated into a single vector using an element-wise maximum (rather than a mean). This conservative approach ensures that a neuron critical to at least one sample retains a high score, preventing it from being pruned due to low activation in other samples within the same batch.
Activation Memory:
- Observations of LLM activation patterns reveal that important neurons exhibit periodic re-firing during long reasoning generations (e.g., re-firing every ~22 tokens with a ~76% re-firing ratio).
- To capture this, the method maintains an activation memory that accumulates importance scores across update phases. The memory is updated via an element-wise maximum ( $m_n \leftarrow \max(m_{n-1}, s_n)$ ).
- This ensures that neurons identified as important in earlier phases remain candidates for retention, even if their activation dips temporarily, allowing the mask to adapt to evolving patterns without losing critical components.

Algorithm Flow
The method operates in three phases during autoregressive generation:

Input Phase: Processes the initial prompt to compute initial importance scores.
Dense Exploration Phase: Performs a set number of dense steps ( $T_{init}$ ) to establish the reasoning direction and collect context-specific activation statistics. The memory is updated, and the first pruning mask is generated.
Periodic Adaptive Pruning Phase: Alternates between sparse computation ( $T_p$ steps) and dense exploration ( $T_E$ steps). At the end of each dense exploration stage, the activation memory is updated, and a new top-k mask is generated based on the accumulated memory.

Key Contributions

Empirical Identification of Batched Collapse: The paper is the first to empirically demonstrate that existing training-free pruning methods (both static and adaptive) severely degrade on reasoning tasks under batched inference. Specifically, adaptive methods like TEAL lose 58.4 to 67.7 percentage points of average accuracy at batch size 4 compared to dense models.
Robust Batched Pruning Framework: The authors introduce a method that maintains accuracy in batched settings by replacing threshold-based selection with periodic top-k selection and incorporating an activation memory to track recurring neuron importance.
Calibration-Free Operation: Unlike methods requiring offline calibration on external corpora (e.g., Wanda, TEAL), this method relies solely on runtime activations, making it calibration-free.

Experimental Results
Experiments were conducted on DeepSeek-R1-Distill-Qwen-7B and DeepSeek-R1-Distill-Llama-8B across five reasoning benchmarks (GSM8K, MATH500, MINERVA, AMC23, GPQA-DIAMOND).

Accuracy: At 50% target sparsity and batch size 4, the proposed method outperforms the previous state-of-the-art (TEAL) by 39.7 percentage points in average accuracy on the Qwen-7B model (54.0% vs. 14.3%) and by 24.1 points on the Llama-8B model.
Stability: Unlike TEAL, which shows high performance at batch size 1 but collapses as batch size increases, the proposed method maintains consistent performance across all batch sizes (1, 4, 8, 16).
Speedup: The method achieves a 1.40× speedup over dense inference at 50% actual sparsity on the Qwen-7B model. The speedup is more pronounced in batched settings because the method structurally reduces weight matrix dimensions, directly lowering computation, whereas methods relying on skipping zero activations become less effective as workloads shift from memory-bound to compute-bound.
High Sparsity: The method sustains strong performance even at high sparsity levels (e.g., 71.0% accuracy at 50% target sparsity), whereas baselines degrade sharply beyond 40% sparsity.

Significance and Claims
The paper claims that its method addresses a critical bottleneck in deploying Large Reasoning Models: the incompatibility of existing adaptive pruning techniques with batched inference. By solving the distribution shift problem inherent in batch aggregation and leveraging the periodic nature of neuron re-firing, the method enables efficient, high-throughput inference without the performance collapse seen in prior approaches. The authors emphasize that their gains are driven by the periodic top-k selection and activation memory, rather than the specific choice of aggregation operator (max vs. mean), which showed only marginal differences in performance.

Technical Summary: Batch-wise Adaptive Pruning for Large Reasoning Models

More like this