# Learnable Permutation for Structured Sparsity on Transformer Models

arXiv: https://arxiv.org/abs/2601.22980 · HTML: https://mlsanity.com/paper/2601.22980/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 40 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Builds on (most cited first, 28 of 28)

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Qwen2.5-VL Technical Report](https://mlsanity.com/paper/2502.13923/) (2025)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [HellaSwag: Can a Machine Really Finish Your Sentence?](https://mlsanity.com/paper/1905.07830/) (2019)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questions](https://mlsanity.com/paper/1905.10044/) (2019)
- [Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering](https://mlsanity.com/paper/1809.02789/) (2018)
- [MMMU: A Massive Multi-discipline Multimodal Understanding and Reasoning Benchmark for Expert AGI](https://mlsanity.com/paper/2311.16502/) (2023)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [Towards VQA Models That Can Read](https://mlsanity.com/paper/1904.08920/) (2019)
- [SparseGPT: Massive Language Models Can Be Accurately Pruned in One-Shot](https://mlsanity.com/paper/2301.00774/) (2023)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch](https://mlsanity.com/paper/2102.04010/) (2021)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [STEP: Learning N:M Structured Sparsity Masks from Scratch with Precondition](https://mlsanity.com/paper/2302.01172/) (2023)
- [ProxSparse: Regularized Learning of Semi-Structured Sparsity Masks for Pretrained LLMs](https://mlsanity.com/paper/2502.00258/) (2025)
- [Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding](https://mlsanity.com/paper/2503.10135/) (2025)
- [S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training](https://mlsanity.com/paper/2409.09099/) (2024)
- [SparK: Query-Aware Unstructured Sparsity with Recoverable KV Cache Channel Pruning](https://mlsanity.com/paper/2508.15212/) (2025)

## Abstract (arXiv)

Structured sparsity has emerged as a popular model pruning technique, widely adopted in various architectures, including CNNs, Transformer models, and especially large language models (LLMs) in recent years. A promising direction to further improve post-pruning performance is weight permutation, which reorders model weights into patterns more amenable to pruning. However, the exponential growth of the permutation search space with the scale of Transformer architectures forces most methods to rely on greedy or heuristic algorithms, limiting the effectiveness of reordering. In this work, we propose a novel end-to-end learnable permutation framework. Our method introduces a learnable permutation cost matrix to quantify the cost of swapping any two input channels of a given weight matrix, a differentiable bipartite matching solver to obtain the optimal binary permutation matrix given a cost matrix, and a sparsity optimization loss function to directly optimize the permutation operator. We extensively validate our approach on vision and language Transformers, demonstrating that our method achieves state-of-the-art permutation results for structured sparsity.