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.

翻译：摘要：硬件领域的近期创新（如Nvidia A100）激发了从零学习N:M结构化稀疏掩码以实现快速模型推理的研究。然而，当前最先进的学习方法（如SR-STE）主要针对非自适应优化器（如动量SGD）设计，而在基于Adam训练的模型（如注意力大语言模型）中会导致显著的精度下降。本文首先揭示此类差距源于被掩码权重导致的Adam状态中第二矩（即方差）估计不当。我们认为，使用Adam学习N:M掩码时应充分考虑方差估计的关键阶段。基于此，我们提出STEP——一种感知Adam的N:M掩码学习方法，包含两个阶段：首先，STEP计算可靠的方差估计（预条件阶段）；随后，固定该方差作为预条件来学习N:M掩码（掩码学习阶段）。STEP通过沿训练轨迹动态采样方差变化并测试样本集中度，自动识别两阶段的切换点。实验部分，我们在CIFAR分类、机器翻译及大语言模型微调（BERT-Base、GPT-2）等多项任务中评估了STEP与ASP、SR-STE等基准方法。结果表明，STEP有效缓解了基准方法的精度下降问题，并对高结构化稀疏比具有鲁棒性。