PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models

AI summary

Overview

Research area: Model compression for large language models — specifically N:M semi-structured pruning and channel permutation.

Technical level: Advanced. The paper assumes familiarity with pruning pipelines, permutation matrices, doubly stochastic matrices, the Hungarian algorithm, and straight-through gradient estimators.

Scope: This paper proposes PermLLM, a post-training pruning framework that replaces handcrafted channel-permutation heuristics with an end-to-end learnable, block-wise permutation optimized to minimize the output gap between dense and N:M sparse models.

What This Paper Is About

N:M sparsity (such as 2:4) forces a model to keep only N of every M consecutive weights, and the accuracy of the resulting sparse model depends heavily on which channels end up grouped together. Prior methods (notably RIA) reorder channels using a handcrafted quality metric — the sum of retained weight importance — but the authors show this proxy can be maximized while the actual output error gets worse (Figure 1 gives an example where direct 2:4 sparsity yields an output loss of 12.375, and maximizing the score metric does not reliably reduce it). PermLLM's goal is to learn the channel permutation directly against the true objective: minimizing the discrepancy between dense-model outputs and sparse-model outputs.

Key Contributions

1. First learnable channel permutation (LCP) for N:M sparsity. The authors replace handcrafted quality metrics with an end-to-end optimization that directly minimizes the dense-to-sparse output error, rather than a proxy score.
2. Sinkhorn-based relaxation of the permutation matrix. Discrete, non-differentiable permutation matrices are relaxed into differentiable soft permutation matrices via Sinkhorn normalization, hardened into strict permutations with the Hungarian algorithm in the forward pass, and trained with a straight-through estimator (STE) in the backward pass.
3. Block-wise permutation strategy. Restricting permutation to within blocks reduces the learnable parameter count from C_in² to C_in × B (a factor of B/C_in of the original) and reduces the hardening complexity from O(C_in³) to O(C_in · B²).
4. A customized CUDA kernel for channel permutation, reported to accelerate the operation by an average of 84× over a PyTorch implementation on LLaMA-2 7B, plus compatibility as a plugin for existing one-shot pruners such as Wanda and RIA.

Main Findings

• Perplexity on Wikitext2 (2:4 sparsity, lower is better). PermLLM_Wanda reaches 14.27 / 9.41 / 8.06 / 9.39 / 8.20 / 14.03 / 13.58 on OPT 6.7B, LLaMA 7B, LLaMA 13B, LLaMA-2 7B, LLaMA-2 13B, LLaMA-3.1 8B, and Qwen-2.5 7B respectively, versus Wanda+CP at 15.28 / 11.07 / 8.69 / 11.00 / 8.51 / 21.09 / 18.76. PermLLM_RIA records 14.23 / 9.95 / 7.81 / 9.60 / 7.97 / 15.79 / 15.93 versus RIA+CP at 15.13 / 10.99 / 8.15 / 10.26 / 8.08 / 19.80 / 17.58.
• Handcrafted permutation helps, but insufficiently. Channel permutation consistently improves Wanda and RIA, yet on newer models such as LLaMA-3.1 8B and Qwen-2.5 7B, Wanda+CP and RIA+CP still lag behind SparseGPT (whose Wikitext2 perplexities are 16.62 and 14.34). PermLLM surpasses SparseGPT on these models (14.03 and 13.58).
• Zero-shot accuracy. PermLLM_Wanda achieves the best average in Table 2 on every model shown: 41.81 (OPT 6.7B), 45.67 (LLaMA 7B), 46.59 (LLaMA-2 7B), 43.33 (LLaMA-3.1 8B), and 52.17 (Qwen-2.5 7B).
• Inference runtime. On LLaMA-2 7B with 2048 tokens, the custom CUDA kernel cuts channel-permutation time from 3.288 ms to 0.039 ms (84× speedup). Per-layer speedups are 1.632× (Q/K/V/O_proj), 1.708× (Up/Gate_proj), and 1.703× (Down_proj), for roughly 1.67× overall acceleration across all linear layers.
• Sinkhorn iterations matter. With 0 iterations (soft matrix furthest from a doubly stochastic matrix) versus 5 iterations, Qwen-2.5 7B averages 42.96 vs 43.33 and LLaMA-3.1 8B averages 49.18 vs 52.17, supporting the use of a doubly stochastic matrix as the relaxation.
• Calibration-data robustness. Using Pile, Wikitext2, or C4 (128 samples each) yields averages of 44.74, 44.61, and 46.59 for LLaMA-2 7B, indicating the learned permutation is not tied to a single calibration set.
• Block size trade-off. Block sizes 32 / 64 / 128 give averages of 43.58 / 46.59 / 47.09 and runtimes of 2h / 2.5h / 6h on LLaMA-2 7B; 64 is chosen as the default balance.
• Partial PermLLM. Applying learnable permutation only to the last six decoder layers of LLaMA-2 7B reduces the runtime to 0.4 hours on a single GPU and still beats RIA+CP (43.78 vs 43.42 average), though it falls short of full PermLLM (44.30).
• Not limited to 2:4. On 4:8 sparse LLaMA-2 7B, PermLLM_Wanda achieves the best average (47.97) and Wikitext2 perplexity (7.96), ahead of Wanda+CP (47.64 / 8.26), Wanda (45.82 / 8.63), and SparseGPT (46.25 / 8.56).
• Cost context. Pruning takes about 2.5 hours for the 7B model on 4 GPUs and 5.5 hours for the 13B model on 8 GPUs. Linear layers targeted by the method constitute approximately 99% of total LLM parameters.

Methodology in Plain English

The starting point is a simple observation: N:M sparsity is sensitive to channel ordering, but previous methods pick the ordering by maximizing an easy-to-compute importance score rather than by minimizing the actual error the pruning causes.

PermLLM instead treats the permutation itself as a trainable object. The problem is that a permutation matrix is made of hard 0s and 1s, so gradients cannot flow through it. The authors work around this in three steps. First, they keep a learnable real-valued matrix of the same shape and pass it through Sinkhorn normalization — a repeated row-and-column normalization that turns any non-negative matrix into a doubly stochastic one, where every row and column sums to 1. A temperature parameter (linearly decayed from 1 to 0.1) controls how close these soft entries get to hard 0/1 values, and the normalization is truncated to 5 iterations.

Second, because a soft matrix also scales weight values (not just reorders them), the soft matrix is converted back into a strict permutation during the forward pass by solving a linear sum assignment problem with the Hungarian algorithm — picking the true permutation that best aligns with the soft one. Because that hardening step is not differentiable, a straight-through estimator passes gradients back as if the step were the identity.

Third, to keep the problem tractable, permutation is only allowed within blocks of consecutive channels rather than across the whole channel dimension. This turns one giant permutation into several small independent ones, forming a block-diagonal permutation matrix. The soft mask used in the backward pass is a softmax over each group of M weights, again paired with a straight-through estimator for the forward pass's argmax.

The only learnable parameters are these small permutation score matrices; the model weights stay frozen. Training happens for 50 iterations per layer with AdamW, using a cosine similarity loss between the dense model's output and the pruned model's output, on 128 C4 samples of 1024 tokens each. After training, weights are permuted and pruned, and the corresponding rows of the preceding layer are reordered to keep input activations aligned (a row-wise operation that preserves N:M sparsity).

Why This Matters

Impact on research. The paper reframes channel permutation from a combinatorial optimization guided by heuristics into a differentiable learning problem, and empirically demonstrates that the widely used "sum of retained importance" criterion is an imperfect proxy for pruning error. This challenges a common assumption in the N:M pruning literature and suggests similar learnable reformulations may apply to quantization ordering and other layout-dependent compression problems.

Real-world applications:

• Deploying LLMs on inference hardware with Sparse Tensor Core support, where 2:4 sparsity can theoretically double compute throughput for sparse matrix multiplication.
• Serving newer open-weight models (LLaMA-3.1, Qwen-2.5) where one-shot pruners degrade sharply without good channel ordering.
• Memory-constrained edge or on-premise deployment where model size reduction matters more than training budget.
• Post-training compression pipelines that can afford a few GPU-hours (about 2.5 hours for 7B on 4 GPUs) but not full retraining.

Industry relevance. The method is designed as a drop-in plugin for existing one-shot pruners (Wanda, RIA) rather than a replacement, and the customized CUDA kernel addresses the main practical objection to permutation-based pruning — runtime overhead. The 84× kernel speedup and the partial variant that fits on a single GPU both target deployment practicality rather than benchmark-only gains.

Future Directions

• Scaling the learned permutations beyond blocks. The block size of 64 is a deliberate efficiency compromise; whether larger or full-matrix permutations can be made affordable with better implementations remains open.
• Which layers deserve learnable permutation. The partial PermLLM experiment (last six decoder layers only) shows the full method is not always necessary, but a principled criterion for selecting layers is not established.
• Extension to other sparsity patterns and compression axes. The paper demonstrates 2:4 and 4:8; whether LCP transfers to less regular patterns or combines with quantization is untested here.
• Combining with weight-update methods. SparseGPT updates remaining unpruned weights, while PermLLM and the other compared methods do not modify weight values — whether learnable permutation and weight updating are complementary is not reported.

Target Audience

Researchers and engineers working on LLM compression, particularly those focused on semi-structured N:M pruning for GPU acceleration, as well as practitioners who already use one-shot pruners like Wanda or RIA and want to improve sparse-model accuracy without retraining. Readers should be comfortable with linear algebra (permutation and doubly stochastic matrices) and gradient-based optimization; the combinatorial and Sinkhorn machinery is explained, but the paper is not an introductory treatment.

Authors’ abstract

Channel permutation is a powerful technique for enhancing the accuracy of N:M sparse models by reordering the channels of weight matrices to prioritize the retention of important weights. However, traditional channel permutation methods rely on handcrafted quality metrics, which often fail to accurately capture the true impact of pruning on model performance. To address this limitation, we propose PermLLM, a novel post-training pruning framework that introduces learnable channel permutation (LCP) for N:M sparsity. LCP leverages Sinkhorn normalization to transform discrete permutation matrices into differentiable soft permutation matrices, enabling end-to-end optimization. Additionally, PermLLM incorporates an efficient block-wise channel permutation strategy, which significantly reduces the number of learnable parameters and computational complexity. PermLLM seamlessly integrates with existing one-shot pruning methods to adaptively optimize channel permutations, effectively mitigating pruning-induced errors. Extensive experiments on the LLaMA series, Qwen, and OPT models demonstrate that PermLLM achieves superior performance in optimizing N:M sparse models. The code is available at https://github.com/lanchengzou/PermLLM.