# STEP: Learning N:M Structured Sparsity Masks from Scratch with Precondition

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

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 51 references that are in the map.
- First posted on arXiv 2 Feb 2023.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Pruning Filters for Efficient ConvNets](https://mlsanity.com/paper/1608.08710/) (2016)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [SNIP: Single-shot Network Pruning based on Connection Sensitivity](https://mlsanity.com/paper/1810.02340/) (2018)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Learning Structured Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1608.03665/) (2016)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Accelerating Sparse Deep Neural Networks](https://mlsanity.com/paper/2104.08378/) (2021)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Learning N:M Fine-grained Structured Sparse Neural Networks From Scratch](https://mlsanity.com/paper/2102.04010/) (2021)
- [The Break-Even Point on Optimization Trajectories of Deep Neural Networks](https://mlsanity.com/paper/2002.09572/) (2020)
- [Adam Can Converge Without Any Modification On Update Rules](https://mlsanity.com/paper/2208.09632/) (2022)
- [Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks](https://mlsanity.com/paper/2102.08124/) (2021)
- [On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length](https://mlsanity.com/paper/1807.05031/) (2018)
- [1-bit Adam: Communication Efficient Large-Scale Training with Adam's Convergence Speed](https://mlsanity.com/paper/2102.02888/) (2021)
- [An Algorithm-Hardware Co-Optimized Framework for Accelerating N:M Sparse Transformers](https://mlsanity.com/paper/2208.06118/) (2022)
- [1-bit LAMB: Communication Efficient Large-Scale Large-Batch Training with LAMB's Convergence Speed](https://mlsanity.com/paper/2104.06069/) (2021)

## Built on by (most cited first, 26 of 26)

- [MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models](https://mlsanity.com/paper/2409.17481/) (2024)
- [Scaling Laws for Sparsely-Connected Foundation Models](https://mlsanity.com/paper/2309.08520/) (2023)
- [Lookahead: An Inference Acceleration Framework for Large Language Model with Lossless Generation Accuracy](https://mlsanity.com/paper/2312.12728/) (2023)
- [SLoPe: Double-Pruned Sparse Plus Lazy Low-Rank Adapter Pretraining of LLMs](https://mlsanity.com/paper/2405.16325/) (2024)
- [Pruning Large Language Models with Semi-Structural Adaptive Sparse Training](https://mlsanity.com/paper/2407.20584/) (2024)
- [Accelerating Transformer Pre-training with 2:4 Sparsity](https://mlsanity.com/paper/2404.01847/) (2024)
- [Effective Interplay between Sparsity and Quantization: From Theory to Practice](https://mlsanity.com/paper/2405.20935/) (2024)
- [Progressive Gradient Flow for Robust N:M Sparsity Training in Transformers](https://mlsanity.com/paper/2402.04744/) (2024)
- [SLiM: One-shot Quantization and Sparsity with Low-rank Approximation for LLM Weight Compression](https://mlsanity.com/paper/2410.09615/) (2024)
- [Accurate Neural Network Pruning Requires Rethinking Sparse Optimization](https://mlsanity.com/paper/2308.02060/) (2023)
- [SparseLoRA: Accelerating LLM Fine-Tuning with Contextual Sparsity](https://mlsanity.com/paper/2506.16500/) (2025)
- [S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training](https://mlsanity.com/paper/2409.09099/) (2024)
- [Cut Less, Fold More: Model Compression through the Lens of Projection Geometry](https://mlsanity.com/paper/2602.18116/) (2026)
- [ELSA: Exploiting Layer-wise N:M Sparsity for Vision Transformer Acceleration](https://mlsanity.com/paper/2409.09708/) (2024)
- [Beyond 2:4: exploring V:N:M sparsity for efficient transformer inference on GPUs](https://mlsanity.com/paper/2410.16135/) (2024)
- [The Journey Matters: Average Parameter Count over Pre-training Unifies Sparse and Dense Scaling Laws](https://mlsanity.com/paper/2501.12486/) (2025)
- [L3: Large Lookup Layers](https://mlsanity.com/paper/2601.21461/) (2026)
- [MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs](https://mlsanity.com/paper/2506.12876/) (2025)
- [PermLLM: Learnable Channel Permutation for N:M Sparse Large Language Models](https://mlsanity.com/paper/2510.10136/) (2025)
- [Sparse by Rule: Probability-Based N:M Pruning for Spiking Neural Networks](https://mlsanity.com/paper/2511.12097/) (2025)
- [Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective](https://mlsanity.com/paper/2607.03784/) (2026)
- [TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks](https://mlsanity.com/paper/2505.23949/) (2025)
- [Extreme Model Compression with Structured Sparsity at Low Precision](https://mlsanity.com/paper/2511.08360/) (2025)
- [Learnable Permutation for Structured Sparsity on Transformer Models](https://mlsanity.com/paper/2601.22980/) (2026)
- [ELAS: Efficient Pre-Training of Low-Rank Large Language Models via 2:4 Activation Sparsity](https://mlsanity.com/paper/2605.03667/) (2026)
- [Compression Trinity: Exploring Sparsity, Quantization, and Low-Rank Approximations for LLM Compression](https://mlsanity.com/paper/2608.24070/) (2026)

## Abstract (arXiv)

Recent innovations on hardware (e.g. Nvidia A100) have motivated learning N:M structured sparsity masks from scratch for fast model inference. However, state-of-the-art learning recipes in this regime (e.g. SR-STE) are proposed for non-adaptive optimizers like momentum SGD, while incurring non-trivial accuracy drop for Adam-trained models like attention-based LLMs. In this paper, we first demonstrate such gap origins from poorly estimated second moment (i.e. variance) in Adam states given by the masked weights. We conjecture that learning N:M masks with Adam should take the critical regime of variance estimation into account. In light of this, we propose STEP, an Adam-aware recipe that learns N:M masks with two phases: first, STEP calculates a reliable variance estimate (precondition phase) and subsequently, the variance remains fixed and is used as a precondition to learn N:M masks (mask-learning phase). STEP automatically identifies the switching point of two phases by dynamically sampling variance changes over the training trajectory and testing the sample concentration. Empirically, we evaluate STEP and other baselines such as ASP and SR-STE on multiple tasks including CIFAR classification, machine translation and LLM fine-tuning (BERT-Base, GPT-2). We show STEP mitigates the accuracy drop of baseline recipes and is robust to aggressive structured sparsity ratios.