STEP：从零学习带预条件的 N:M 结构化稀疏掩码
机器学习
2023-02-03 v1
摘要
近期硬件创新（如 Nvidia A100）推动了从零学习 N:M 结构化稀疏掩码以实现快速模型推理。然而，该 regime 下最先进的学习方案（如 SR-STE）是针对动量 SGD 等非自适应优化器提出的，而对于像基于注意力的 LLM 这类 Adam 训练模型会带来不小的精度下降。本文中，我们首先证明这种差距源于被掩码权重给出的 Adam 状态中估计不佳的第二矩（即方差）。我们推测，用 Adam 学习 N:M 掩码应考虑方差估计的关键 regime。鉴于此，我们提出 STEP，一种 Adam 感知的方案，分两阶段学习 N:M 掩码：首先，STEP 计算可靠的方差估计（预条件阶段）；随后，方差保持固定并用作学习 N:M 掩码的预条件（掩码学习阶段）。STEP 通过动态采样训练轨迹上的方差变化并检验样本集中度，自动识别两阶段的切换点。在经验上，我们在多项任务（包括 CIFAR 分类、机器翻译和 LLM 微调（BERT-Base、GPT-2））上评估 STEP 与 ASP、SR-STE 等基线。我们展示 STEP 缓解了基线方案的精度下降，并对激进的结构化稀疏比具有鲁棒性。
引用
@article{arxiv.2302.01172,
title = {STEP: Learning N:M Structured Sparsity Masks from Scratch with Precondition},
author = {Yucheng Lu and Shivani Agrawal and Suvinay Subramanian and Oleg Rybakov and Christopher De Sa and Amir Yazdanbakhsh},
journal= {arXiv preprint arXiv:2302.01172},
year = {2023}
}