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

TL;DR

This paper introduces STEP, a novel Adam-aware method for learning N:M structured sparsity masks from scratch, which improves accuracy and robustness in sparse models across various tasks by accounting for variance estimation in Adam optimizer.

Contribution

STEP is the first adaptive recipe that incorporates variance estimation to effectively learn N:M masks with Adam, addressing accuracy issues in prior methods.

Findings

STEP reduces accuracy drop in sparse models.

It is effective across classification, translation, and language models.

STEP maintains robustness at high sparsity ratios.

Abstract

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.