Efficient Dynamic Structured Sparse Training with Learned Shuffles

TL;DR

This work tackles the efficiency gap between structured sparsity and unstructured dynamic sparse training by introducing Permutation-Augmented DST (PA-DST), which jointly learns a permutation matrix per layer with a structured sparse weight matrix. The authors provide a combinatorial expressivity analysis based on linear regions to show how mixing via permutations restores depth-dependent expressivity that structured patterns alone limit, and they validate the theory with experiments on Vision (ImageNet-1K, ViT) and Language (WikiText-103, GPT-2). Empirically, PA-DST achieves accuracy comparable to unstructured DST at 90–95% sparsity while delivering up to 1.21× training and 2.9× inference speedups, highlighting structure plus learned permutation as a practical compromise between accuracy and hardware efficiency. The approach integrates a differentiable permutation estimator (soft-to-hard via the Birkhoff polytope) and re-indexing tricks to avoid extra matmuls, enabling efficient end-to-end training and inference in transformer-like architectures.

Abstract

Structured sparsity accelerates training and inference on modern GPUs, yet it still trails unstructured dynamic sparse training (DST) in accuracy. The shortfall stems from a loss of expressivity: whereas a dense layer can realize every possible mask obtained by choosing any $w$ active weights out of $n$, a fixed block or N:M layout explores only a subset of those possibilities. We propose to close this gap by learning, for each layer, a single permutation matrix jointly with the structured weight matrix. Applied to three canonical structures -- block, N:M, and diagonals -- we show that permutation-augmented DST (PA-DST) matches unstructured baselines (RigL, SET) at 90--95\% sparsity on ImageNet-1K (ViT-B/16) and WikiText-103 (GPT-2), yet trains up to $1.21\times$ and infers up to $2.9\times$ faster. The results position structure + learned permutation as a sweet spot between accuracy and efficiency.

Figures (6)

• Figure 1: Overview of the PA-DST training approach. The training starts by initializing the target weight matrices with the desired sparse structure and a soft permutation matrix. As the training progresses, the soft permutation matrix starts moving towards a permutation matrix, and by the end of the training, we obtain the final sparse weight matrix and the permutation matrix. The permutation matrix is then used to reindex the activations of the layers so that an explicit matrix multiplication operation can be avoided during inference. More details can be found in Sec. \ref{['sec:method:train']}.
• Figure 2: Accuracy vs sparsity of networks listed in Sec. \ref{['sec:exp:setup']}. The plots in blue, green, and red correspond to unstructured, structured, and structured + permutations-based (PA-DST) methods, respectively. We see that a permutation helps improve the generalization performance (red plots are above the green plots). Moreover, permutations help bridge the gap between unstructured and structured sparse methods (structures + permutations, in red plot, reach close to the unstructured methods, in blue).
• Figure 3: Impact of permutations on the training and inference time for a ViT-B/16 model. We observe that there is a 3.16% - 8.69% overhead related to permutations during inference for the obtained accuracy gains. Even for training, learning permutations increases the overall training time, but at higher sparsity levels, we still obtain speedup compared to a dense model for all structured sparsities. Training time for GPT models can be found in Tbl. \ref{['tab:overhead_gpt2_medium_diag']}.
• Figure 4: Distance of permutations to an identity matrix (no permutation) for a ViT-B/16 at 90% sparsity trained with DynaDiag. We see that with depth, the learned permutations move closer to an identity matrix, signalling less shuffling with depth. A: Attention block; F: Fully Connected; b:blocks
• Figure 5: Tracking the loss of permutation matrices based on Eqn. \ref{['eq:permloss']}. We plot the loss for every 10 training epochs.