Title: TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks
Venue: NeurIPS 2025
Accepted: yes
Keywords: large language models

network pruning

N:M sparsity

optimization

entropy regularization

optimal transport
Abstract: Network pruning reduces computational requirements of large neural networks, with N:M sparsity—retaining only N out of every M consecutive weights—offering a compelling balance between compressed model quality and hardware acceleration. However, N:M sparsity only accelerates forward-pass computations, as N:M patterns are not preserved during matrix transposition, limiting efficiency during training where both passes are computationally intensive. While transposable N:M sparsity has been proposed to address this limitation, existing methods for finding transposable N:M sparse masks either fail to scale to large models or are restricted to M=4 which results in suboptimal compression-accuracy trade-off. We introduce an efficient solver for transposable N:M masks that scales to billion-parameter models. We formulate mask generation as optimal transport problems and solve through entropy regularization and Dykstra's algorithm, followed by a rounding procedure. Our tensor-based implementation exploits GPU parallelism, achieving up to 100× speedup with only 1-10\% error compared to existing methods. Our approach can be integrated with layer-wise N:M pruning frameworks including Wanda, SparseGPT and ALPS to produce transposable N:M sparse models with arbitrary N:M values. Experiments show that LLaMA3.2-8B with transposable 16:32 sparsity maintains performance close to its standard N:M counterpart and outperforms standard 2:4 sparse model, showing the practical value of our approach.  Our code is available at https://github.com/mazumder-lab/TSENOR.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes TSENOR, a novel and efficient algorithm for generating binary masks with transposable N:M sparsity, enabling acceleration in both forward and backward passes during neural network training. The authors formulate the mask generation problem as an entropy-regularized optimal transport problem and solve it using Dykstra’s algorithm. A GPU-optimized rounding step, composed of greedy selection and local search, ensures constraint satisfaction and mask quality. The method is integrated with several pruning frameworks (e.g., Wanda, SparseGPT, ALPS) and achieves substantial speedup (10–100×) and low reconstruction error (1–10%) on  LLaMA3.2–8B model.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\n1. Reformulating the sparse mask search as an optimal transport problem is an insightful approach, grounded in theory and convex optimization.\n\n2. Figure 3 shows that TSENOR achieves superior solution quality with significantly lower relative error compared to competing methods like 2-approx.\n\n3. Table 1 and related experiments show significant acceleration due to GPU parallelism, making the method viable for large models.\n\n4. Comprehensive empirical validation: Evaluation includes mask quality metrics, perplexity comparisons, and inference speedups on multiple sparsity levels.\n\n# Weaknesses\n1. Transposable N:M sparsity remains a relatively niche topic; thus, the broader significance and audience interest may be limited.\n\n2. Alternative baselines could be enhanced: While the current comparisons are reasonable, further exploration of acceleration for classic heuristics like 2-approx on GPU may reinforce the case.\n\n3. Incomplete code availability: The paper mentions GPU-efficient implementations, but only partial code appears to be provided. Without full implementation and integration examples, reproducibility and adoption may be hindered.\n\nQUESTIONS:\n1. Why is Bi-NM used in Figure 5 as the primary comparison instead of 2-approx? In Figure 3, 2-approx shows better reconstruction error than Bi-NM, so it would be more convincing to include 2-approx as a baseline in Figure 5.\n\n2. Have you considered GPU-accelerated implementations of 2-approx for fairer runtime benchmarking?\nTSENOR is compared to 2-approx on CPU in Table 1. Leveraging NVIDIA’s cuOPT for general LP or other parallel kernels could offer a more apples-to-apples comparison with Network Flow.\n\nLIMITATIONS:\nWhile they do briefly mention limitations in two sentences at the end of section 6, it's not enough. I'd encourage the authors to add a dedicated section on limitations. A thorough discussion of what's missing or where the current work falls short is truly invaluable for paving the way for impactful future research.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces TSENOR, an algorithmic framework for efficiently generating transposable N:M sparse masks, essential for accelerating both forward and backward passes in neural network training. The key idea is formulating the mask generation problem as multiple optimal transport problems with entropy regularization, solved through Dykstra’s algorithm followed by a specialized rounding procedure. TSENOR is integrated into popular pruning methods such as Wanda, SparseGPT, and ALPS. The authors report significant computational speedups and reduced error relative to existing heuristics. Experimental results on large-scale models such as LLaMA demonstrate the practical efficacy of TSENOR, especially highlighting performance gains at larger M-values (e.g., 16:32 sparsity).\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths:\n1. Novel Optimization Formulation: Formulating mask selection as an optimal transport problem with entropy regularization is both elegant and theoretically appealing, significantly simplifying computational complexity compared to previous methods.\n2. Efficiency and Scalability: The GPU-optimized implementation of TSENOR demonstrates remarkable efficiency and scalability, achieving substantial speedups (up to 100×) over baseline methods, enabling its practical application to large-scale language models.\n3. Practical Integration with Popular Methods: The proposed method seamlessly integrates with existing pruning frameworks like Wanda, SparseGPT, and ALPS, allowing straightforward adoption within the broader neural network pruning community.\n4. Comprehensive Experimental Validation: Experiments are thorough and extensive, covering both single-matrix evaluations and full-scale pruning scenarios on LLaMA models, clearly demonstrating computational advantages and acceptable accuracy trade-offs.\n\n## Weaknesses:\n1. Clarity on Practical Gain: While the method clearly demonstrates computational efficiency, the practical gain in terms of the speed-accuracy trade-off in training and inference is not fully explicit. Specifically, results explicitly comparing accuracy versus computational speed between transposable and non-transposable N:M masks (or between regular dense matmul) are limited or unclear.\n2. Technical Novelty Moderation: Although the connection to optimal transport is interesting and practically beneficial, the overall approach leverages well-established optimization techniques (e.g., entropy regularization, Dykstra's algorithm) and therefore offers moderate conceptual novelty.\n3. Limited Detailed Discussion on Larger M Values: The claim regarding the computational benefits and accuracy at larger M-values (e.g., 16:32) seems insufficiently justified. It's unclear how these larger M-values maintain similar computational benefits as smaller structured masks, considering the tendency towards unstructured patterns at larger M.\n\n## Overall Evaluation:\nThe paper presents an interesting approach by leveraging optimal transport theory for efficient mask generation in neural network pruning, particularly relevant given the increasing prominence of large language models. While computationally impressive, the practical advantage regarding accuracy-speed trade-offs is not fully articulated. The technical novelty is moderate, considering the method largely relies on known optimization techniques.\nThus, I lean towards a weak rejection given the current state, primarily because the critical practical gains (accuracy-speed trade-off improvements) are not thoroughly demonstrated. The paper could significantly benefit from narrowing its scope, providing more explicit comparative results, and further justifying claims regarding large M-value effectiveness.\n\nQUESTIONS:\n1. Clarification of Computational Benefits at Larger M-values: Your claim— \"Our experiments show that networks with larger M values (e.g., 16:32) provide the same computational benefits as smaller M values but with much less impact on model accuracy\"— is unclear.\n    * Which specific experiments support this claim?\n    * Given that larger M-values inherently become less structured, it is counterintuitive that computational benefits remain constant. Could you clarify how these masks remain efficiently acceleratable?\n2. Accuracy-Speed Trade-off Analysis: If I understand correctly, transposable N:M sparsity should accelerate both training and inference. However, the current manuscript lacks explicit results that clearly illustrate improved accuracy-speed trade-offs.\n    * Could you clarify if and where such comparative results exist in the paper? If not, can you provide a rationale or consider including this comparison?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-23"}

comments: SUMMARY:
TSENOR is a work providing high quality N:M masks that scales to large models with efficient algorithms and exploitation of GPUs speedup in tensor multiplication. TSENOR generates weight mask by adapting Dykstra’s algorithm to solve an optimal transport problem. A greedy rounding method is further applied to create transposable binary mask.

STRENGTHS AND WEAKNESSES:
Strengths:

TSENOR well balances the effectiveness and efficiency of N:M sparsity, which enables large M for higher quality solutions while keeping solving time within a feasible range. 
The paper is well-written, algorithms are motivated in a clear way and explained with good examples. 
Ablation study in the appendix covered most concerns of mine. For different features of TSENOR, there are evidence to justify their usage.

Weaknesses:

For exploiting the GPU tensor multiplication, it would be nice to have a usable framework for TSENOR.
Source code is not provided, and it is difficult for anyone tends to reproduce the results even though hyperparameters are provided.
Fine-tuning performance of TSENOR is not very well presented.

QUESTIONS:
If one cares only about inference, can TSENOR be easily adapted to provide non-transposable masks?

It seems that optimal solution is more reachable when M is smaller than 8. Is it a general conclusion that larger M is always a significant improvements in solution quality (as shown in Upper Figure 4) in different types of networks?

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-27

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents TSENOR, an efficient algorithm for generating transposable N:M sparse masks that enable both forward and backward pass acceleration in large neural networks like LLMs. TSENOR uses an entropy-regularized optimal transport formulation solved with GPU-parallelizable Bregman projections to find high-quality fractional masks, then refines them into valid binary masks with a fast rounding procedure. Integrated with popular pruning frameworks like SparseGPT, Wandda and ALPS, TSENOR scales to billion-parameter models, supports flexible N:M ratios, and achieves substantial speedup with minimal performance loss. This work has showed it can achieve siganifically less relative error compared to other transposable pruning methods. Furthermore, it integrates with ALPS and achieves comparable performance with non-transposable ones on LLMs like LLaMA3.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper has formulated transposable N: M sparsity cleanly as an entropy-regularized optimal transport problem with capacity constraints, which is mathematically well-founded. \n2. Dykstra’s algorithm and the proposed rounding algorithm with the local search are highly efficient, which can achieve the lowest relative error compared to the optimal solution while speeding it up by 100 times. \n3. This method looks like an effective plug-in for the ALPS method. And for one-shot pruning, this transposable N: M sparsity can achieve comparable performance with non-transposable ones on LLMs.  \n\nWeaknesses:\n1. As transposable N: M sparsity is designed to speed up both forward and backward, it is best to show practical benefit on more training or fine-tuning settings. While this paper has demonstrated it by fine-tuning LLaMA3.2-8B, although with much lower relative error compared to other transposable methods, the final performance improvement over Bi-NM seems limited.\n\nQUESTIONS:\n1. In section 5.2.1, it evaluates the tradeoff between speedup and performance loss, which seems to be one-shot pruning without fine-tuning.  If I understand correctly, it compares the one-shot pruning performance and estimates the speedup of fine-tuning. A more straightforward way might be to evaluate both on fine-tuning settings. \n2. What are the fine-tuning settings? If fine-tuning hard, will the transposable method like Bi-NM also work very well?\n\nLIMITATIONS:\nMy main concern is that although the paper proposes an efficient method for generating transposable N: M sparse masks, the practical benefit may be limited if, after fine-tuning, models perform similarly to those produced by other transposable sparsity methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-04"}