SHUFFLESPARSE: Learned Shuffles for Structured Sparse Networks

Structured weight sparsity accelerates training and inference on modern GPUs, but it trails unstructured dynamic sparse training (DST) in accuracy especially at extreme sparsity. We pinpoint the reason for this difference in performance to a lack of expressivity: a dense layer can implement any pattern of non-zero weights, whereas structured patterns are restricted to only a small set of weight configurations. We introduce SHUFFLESPARSE, a single permutation primitive that applies uniformly across DST-from-scratch and one-shot pruning, and across N:M, block etc. We close most of this gap by learning a single permutation matrix jointly with the structured weight matrix. When used on three different types of structures (block, N:M, and diagonal), SHUFFLESPARSE is able to reduce the structured-vs-unstructured accuracy gap on ViT-B16 (ImageNet-1K) and GPT-2 (WikiText-103) at 90-95% sparsity, while adding minimal inference overhead (< 8.7% inference overhead) and preserving any training acceleration the host structure provides. The same permutation formulation transfers to one-shot 2:4 pruning of pretrained LLMs, where it improves zero-shot accuracy by 4.6 points on LLaMA-2 7B. Together, these results establish learned permutations as a general tool for recovering unstructured-level accuracy from structured sparsity patterns.