Reservoir of Importance: Learning Semi-Structured Sparsity with Differentiable Subset Sampling

Ha Dinh; Xuan Duy Ta; Khoat Than; Khac-Hoai Nam Bui

Source-linked AI summary

Reservoir of Importance: Learning Semi-Structured Sparsity with Differentiable Subset Sampling

Ha Dinh, Xuan Duy Ta, Khoat Than, Khac-Hoai Nam Bui

arXiv:2608.23048v1Aug 24, 2026cs.LG

TL;DR

Existing learnable-mask methods model all feasible N:M patterns, creating substantial parameter and memory overhead. RoI uses differentiable subset sampling to learn compact masks, achieving lower WikiText-2 perplexity and competitive zero-shot accuracy across Qwen2.5 models from 0.5B to 7B parameters.

Problem
Existing learnable-mask approaches use combinatorial categorical parameterizations for N:M patterns, imposing substantial memory and parameter overhead as M increases.
Method
RoI learns N:M masks by sampling subsets without replacement from compact logits with differentiable Weighted Reservoir Sampling and Gumbel-Top-K relaxation.
Results
RoI consistently achieves lower WikiText-2 perplexity and competitive zero-shot accuracy across Qwen2.5 models ranging from 0.5B to 7B parameters.
Takeaways & Limitations
RoI provides a practical and scalable approach to semi-structured LLM pruning with reduced trainable parameters and memory overhead.
Takeaways & Limitations
RoI’s practical speedups depend on hardware kernels supporting the selected N:M sparsity pattern, especially native 2:4 support.

Abstract

from arXiv · show

Semi-structured $N$:$M$ sparsity has emerged as a practical direction for accelerating large language models (LLMs). However, existing learnable-mask approaches incur substantial parameter and memory overhead, limiting their scalability to large models and aggressive sparsity regimes. In this work, we revisit semi-structured pruning from a perspective that reconciles efficiency with scalability. We propose Reservoir of Importance (RoI), a lightweight semi-structured pruning framework that learns sparsity masks through differentiable subset sampling. Unlike prior methods that model full categorical distributions over all feasible $N$:$M$ patterns, RoI introduces a compact-logit parameterization for sparsity mask learning and performs sampling without replacement to select masks, thereby reducing trainable parameters from combinatorial complexity to $\mathcal{O}({M})$. As a result, RoI requires 1.5-8.75$\times$ fewer learnable parameters and significantly lower memory cost, while remaining fully aligned with hardware-friendly sparsity patterns. Extensive evaluations across multiple scales of the Qwen2.5 LLM family (0.5-7B parameters) demonstrate that RoI achieves competitive performance with strong memory efficiency, stability, and scalability to more aggressive $N$:$M$ sparsity patterns, offering a practical path toward efficient LLM deployment.

1 Introduction

Semi-structured N:M pruning balances flexible weight removal with hardware-friendly regularity, but learnable-mask methods based on full categorical distributions incur combinatorial overhead as M increases. RoI addresses this limitation by differentiably sampling N entries without replacement from compact O(M)-sized logits, enabling scalable and memory-efficient mask learning.

N:M sparsity removes redundant weights while preserving regular patterns compatible with modern accelerators, enabling a practical trade-off between flexibility and inference efficiency.
Learnable-mask methods such as MaskLLM (Fang et al., 2024) and HyperPrune model each group with a categorical distribution over all feasible N:M configurations.These methods achieve strong pruning performance and generalization, but their combinatorial parameterization creates substantial memory and parameter overhead as M increases.
RoI replaces full categorical modeling with differentiable subset sampling, selecting N entries without replacement from a compact logit vector of size M.This reduces the per-group parameterization to O(M) while retaining differentiability and hardware alignment.
RoI enables scalable, stable, and memory-efficient semi-structured mask learning by avoiding combinatorial distributions and using efficient subset-sampling dynamics.

2 Related Work

Semi-structured pruning balances hardware efficiency and accuracy through regular N:M patterns, while saliency-based and learning-based methods offer different approaches to selecting parameters. RoI addresses remaining scalability challenges by reducing mask-learning parameterization to O(M) per group with weighted reservoir sampling.

Semi-structured pruning: Semi-structured pruning uses regular N:M patterns to enable hardware acceleration while preserving model accuracy, bridging structured and unstructured pruning.The constraint retains exactly N non-zero weights within each group of M parameters.
Saliency-based pruning: Saliency-based methods, including SparseGPT and Wanda, use small calibration datasets to estimate weight importance from magnitude, gradients, or second-order statistics before pruning.Their criteria include Hessian approximations inspired by Optimal Brain Damage and Optimal Brain Surgeon.
Learning-based pruning: Learning-based semi-structured pruning optimizes pruning decisions through gradient-based training under the N:M constraint, but large-scale learning efficiency and performance generalization remain open challenges.ProxSparse reduces training complexity through regularized optimization and smaller learning samples.
RoI: RoI reduces mask-learning cost to O(M) parameters per group by improving the categorical formulation with weighted reservoir sampling.The design provides significant memory savings while retaining high pruning accuracy for large-scale semi-structured pruning.

3 Preliminaries

These preliminaries define weighted reservoir sampling as weighted sampling without replacement and present Gumbel-Max/Top-K as a reparameterized, differentiable route to subset sampling. Gumbel-Top-K uses sequential softmax relaxations to produce soft K-hot representations suitable for gradient-based learning.

Weighted Reservoir Sampling: Weighted Reservoir Sampling samples K elements from L candidates without replacement, preferentially selecting items with larger relative weights.The probability of a subset is proportional to its item weights. (Efraimidis and Spirakis, 2006; Vitter, 1985)
Gumbel-Max: Gumbel-Max reparameterizes categorical sampling by adding independent Gumbel noise to logits and selecting the item with the largest perturbed key.The selected index is obtained by argmax over the perturbed keys.
Gumbel-Top-K: Gumbel-Top-K generalizes Gumbel-Max to select K candidates without replacement from a categorical distribution.It selects the K items with the highest Gumbel-perturbed keys. (Xie and Ermon, 2019)
Differentiable Relaxation: Replacing hard top-K selection with sequential softmax relaxations yields differentiable relaxed one-hot vectors whose sum is a soft K-hot representation for gradient-based training.The relaxation uses temperature τ and produces an ordered collection after K iterations. (Plötz and Roth, 2018)

4 Methodology

RoI reformulates N:M mask learning as stochastic variational optimization and uses weighted sampling without replacement with differentiable Gumbel-Top-K relaxation. This compact parameterization reduces mask-learning complexity from combinatorial scale to O(M) while preserving exactly N selected weights per group.

Problem formulation: The original mask-selection problem is NP-hard because each group must choose an N-hot length-M mask, while the massive number of groups makes exhaustive search intractable.RoI therefore reformulates the objective as a tractable stochastic variational optimization problem.
Variational optimization: RoI models masks as independent variational factors and optimizes their parameters through a differentiable relaxation of stochastic mask sampling.The joint distribution factorizes across groups, enabling gradient-based learning of the mask distributions.
Compact mask parameterization: RoI replaces full categorical distributions over N-hot masks with sequential weighted sampling without replacement, requiring exactly M parameters independent of N and reducing complexity to O(M).This yields an exponential improvement in memory efficiency while maintaining valid N:M masks.
Differentiable subset sampling: Gumbel-Top-K perturbs compact logits and samples N items without replacement, while summing relaxed one-hot selections produces a continuous N-hot mask for gradient optimization.The relaxation uses a temperature τ > 0 to control the sharpness of the soft selections.
Annealed sampling: RoI separately controls approximation sharpness and sampling randomness with τ and λ, annealing both during training to progressively guide mask learning.Larger λ induces higher sampling randomness through scaled logits.

5 Experiment

Across Qwen2.5 models from 0.5B to 7B under 2:4 sparsity, RoI achieves competitive or superior pruning quality while substantially reducing training overhead. Stress tests further show that learnable masks remain essential under 2:8 sparsity and that RoI scales better with additional training data.

5.1 Experimental setup: Experiments cover Qwen2.5 models from 0.5B to 7B using 2:4 sparsity, 2,000 training steps, batch size 256, sequence length 4096, and approximately 2B training tokens.The 2:4 pattern is compatible with NVIDIA hardware acceleration, while baselines include Wanda, SparseGPT, ProxSparse, and MaskLLM.
5.1 Zero-shot evaluation: RoI’s main advantage is efficiency: its trainable-parameter and memory costs are substantially lower than MaskLLM’s, with the gap widening as model size and sparsity increase.At 2:8 sparsity for the 7B model, MaskLLM requires approximately 22.75B more parameters than RoI and causes a 3.5× increase in parameters participating in memory allocation, backpropagation, optimizer, and scheduler states.
5.2 Stress test under 2:8 sparsity: Under aggressive 2:8 sparsity, only mask-learning methods maintain reasonable perplexity, whereas saliency-based baselines collapse to perplexities exceeding 100.RoI sustains language-modeling capability using only 28.57% of the trainable parameters required by the categorical parameterization; 2:8 is evaluated for scalability because mainstream accelerators do not broadly support it.
5.3.2 Performance across data scales: RoI and MaskLLM continue improving monotonically as training tokens increase, while ProxSparse rapidly plateaus, demonstrating stronger practical data scalability for RoI.The comparison uses WikiText-2 perplexity on Qwen2.5-1.5B; RoI and MaskLLM optimize the same variational objective, but RoI scales more effectively in practice.

6 Conclusion

RoI uses differentiable subset sampling to efficiently derive N:M sparsity masks with fewer trainable parameters and lower memory overhead, while maintaining strong downstream performance across Qwen2.5 models from 0.5 to 7B parameters.

6 Conclusion: RoI employs differentiable subset sampling to efficiently derive N:M sparsity masks while substantially reducing trainable parameters and memory overhead.This enables effective compression of large language models.
6 Conclusion: Across Qwen2.5 models ranging from 0.5 to 7B parameters, RoI consistently achieves lower perplexity on WikiText-2 than existing pruning approaches.The conclusion also reports strong downstream performance across these model scales.

Limitations

RoI’s practical speedups remain hardware-dependent, with substantial gains currently limited to platforms that natively accelerate 2:4 sparsity. Although RoI can in principle extend to arbitrary N:M patterns, the supplied passage is truncated before detailing this limitation.

Hardware dependence: RoI’s substantial throughput gains currently depend on hardware platforms with native kernel-level support for 2:4 structured sparsity.The passage identifies AMD ROCm and certain NVIDIA Ampere and Hopper GPUs as examples of supported platforms.
Scope of sparsity patterns: Although RoI can in principle extend to arbitrary N:M sparsity patterns, the passage does not provide the associated practical qualification because it is truncated.

A Appendix · A.1 WRS Yields Equivalent Variational Objective

Theorem 1 establishes that sampling N:M masks from their target distributions yields the same expected loss as parameterizing them through ordered subsets sampled by the corresponding restricted WRS distributions.

A.1 WRS Yields Equivalent Variational Objective: Theorem 1 states that target-distribution mask sampling and restricted WRS subset sampling produce equal expected losses.The WRS parameterization represents each mask as a sum of elements from an ordered sampled subset.
A.1 WRS Yields Equivalent Variational Objective: The target mask distribution is expressed by summing restricted WRS probabilities over subsets whose elements sum to the mask.This provides the distributional identity underlying the equivalent variational objective.
A.1 WRS Yields Equivalent Variational Objective: The objective function f(m) is evaluated under the target mask distribution Pr(m|ϕ).Equation (14) introduces f(m) alongside the target distribution notation.
A.1 WRS Yields Equivalent Variational Objective: The WRS representation sums over all subsets S_m whose elements combine to form each mask m.This subset enumeration defines the correspondence used in the variational equivalence.
A.1 WRS Yields Equivalent Variational Objective: For any set of binary masks satisfying N:M sparsity, the proof begins from the definition of the target distribution Pr(m | ϕ).The argument applies to an arbitrary collection of valid N:M masks.
A.1 WRS Yields Equivalent Variational Objective: The proof reorganizes the resulting summations to connect mask-based expectations with subset-based expectations.Reordering the sums is the key algebraic step shown in the proof.

A.2 Baseline Methods

The baseline methods include data-free magnitude pruning and Wanda, which incorporates layerwise activation statistics to improve pruning quality under high sparsity while retaining computational efficiency.

Baseline Methods: Magnitude pruning removes parameters with the smallest absolute values without data, but often produces subpar results because it ignores parameter sensitivity and model dynamics.It is simple to implement but limited by its data-free criterion.
Baseline Methods: Wanda (Sun et al., 2024) combines parameter magnitudes with layerwise activation statistics, improving over magnitude pruning especially at higher sparsity constraints while maintaining favorable computational efficiency.Its criterion uses activation information at each layer alongside parameter magnitudes.

A.3 Evaluation Metrics and Benchmark Datasets … A.13 Additional Results

Across diverse benchmarks, sparsity regimes, model scales, and deployment settings, RoI combines competitive or superior pruning quality with lower training cost and effective hardware acceleration. Its differentiable, item-level sampling improves optimization efficiency, while ablations and transfer experiments support the robustness of the approach.

A.3 Evaluation Metrics and Benchmark Datasets / RoI / A.11 Mask Difference Analysis: ACC and PPL evaluate pruning across zero- or few-shot NLP tasks and language modeling, using WikiText-2 plus ARC, BoolQ, HellaSwag, PIQA, SciQ, and RACE.ACC measures task performance, while lower PPL indicates better next-word prediction; dataset statistics are reported for zero-shot evaluation.
A.6 Training Time Comparison: 46.10 GPU hours lets RoI train faster than ProxSparse at 52.41 hours and MaskLLM at 55.39 hours on the 7B model.RoI requires less training time than both methods across all evaluated Qwen2.5 model sizes under the same 2:4 setting and training budget.
A.5 Annealing Schedule Ablation: Exponential annealing consistently outperforms linear schedules, with (0.001, 0.05) best for average accuracy and (0.002, 0.05) lowest for perplexity.The schedule reduces temperatures aggressively early and refines them near convergence on Qwen2.5-1.5B.
A.7 Optimization Efficiency / A.8 Gumbel-Top-K Algorithm / A.4 Hyperparameter Setting: RoI’s item-level logits and weighted sampling without replacement require M logits, share evidence across masks, and reduce memory traffic, optimizer-state overhead, and estimation noise.Gumbel-Top-K provides differentiable K-item sampling without replacement, enabling gradient-based mask learning; the method is evaluated with frozen weights and annealed temperatures.
A.9 Comparison with Straight-Through Gumbel-Top-K: 21.27 perplexity and 55.04% accuracy show that soft RoI outperforms straight-through Gumbel-Top-K on Qwen2.5-1.5B.The straight-through variant obtains 30.58 perplexity and 53.10% accuracy; on Qwen2.5-3B, perplexity is 23.57 versus RoI’s 19.37.
A.10 Throughput: 1.3× and 1.27× speedups are achieved by 2:4 sparse Qwen2.5-3B and Qwen2.5-7B models at 1024/4096 input/output tokens.Sparse variants consistently outperform dense counterparts, with similar gains under 2048/4096 configurations.
A.12 Accuracy Results for 2:8 Sparsity: Under 2:8 sparsity, explicit mask learners RoI and MaskLLM substantially outperform heuristic methods, while all pruned models remain below random guessing on ARC-Challenge.RoI remains competitive or superior to MaskLLM at a fraction of the computational cost and across tasks and sparsity regimes.
A.13 Additional Results: 16.30 perplexity and 64.08 average zero-shot accuracy make RoI best on Gemma2-9B, with best or tied-best results on six of seven tasks.RoI improves MaskLLM perplexity by 0.95 points and exceeds its average accuracy of 63.73; SparseGPT reaches 60.62.