[论文解读] Bi-directional Masks for Efficient N:M Sparse Training

Yuxin Zhang, Yiting Luo+5 more|arXiv (Cornell University)|Feb 13, 2023

Speech and Audio Processing被引用 6

一句话总结

本文提出双向掩码（Bi-Mask），一种新型的高效N:M稀疏训练方法，通过使用独立的掩码分别解耦前向与反向权重稀疏性，实现在N:M稀疏张量核心上无性能损失的完整加速。通过在反向掩码前应用高效的行重排，Bi-Mask在保持高模型精度的同时，相比现有方法实现了更优的训练效率。

ABSTRACT

We focus on addressing the dense backward propagation issue for training efficiency of N:M fine-grained sparsity that preserves at most N out of M consecutive weights and achieves practical speedups supported by the N:M sparse tensor core. Therefore, we present a novel method of Bi-directional Masks (Bi-Mask) with its two central innovations in: 1) Separate sparse masks in the two directions of forward and backward propagation to obtain training acceleration. It disentangles the forward and backward weight sparsity and overcomes the very dense gradient computation. 2) An efficient weight row permutation method to maintain performance. It picks up the permutation candidate with the most eligible N:M weight blocks in the backward to minimize the gradient gap between traditional uni-directional masks and our bi-directional masks. Compared with existing uni-directional scenario that applies a transposable mask and enables backward acceleration, our Bi-Mask is experimentally demonstrated to be more superior in performance. Also, our Bi-Mask performs on par with or even better than methods that fail to achieve backward acceleration. Project of this paper is available at \url{https://github.com/zyxxmu/Bi-Mask}.

研究动机与目标

• 解决尽管硬件支持稀疏张量核心，但N:M稀疏训练因反向传播密集化而导致的低效问题。
• 克服现有方法中刚性可转置掩码导致的性能下降，此类掩码限制了稀疏模式的灵活性。
• 在保持模型精度的前提下，实现在N:M稀疏训练中的前向与反向完全加速。
• 设计一种灵活且高效的机制，以最小化稀疏训练中前向与反向传播之间的梯度差距。
• 证明双向掩码结合智能重排可达到或超越非加速稀疏训练方法的性能。

提出的方法

• 引入两个独立的N:M稀疏掩码：一个用于前向传播（行），一个用于反向传播（列），实现稀疏模式的解耦。
• 在计算反向掩码前对权重矩阵进行高效的行重排，以最大化可接受的N:M块数量。
• 对重排后的权重使用基于大小的剪枝生成反向掩码，确保与N:M稀疏张量核心加速的兼容性。
• 通过将输出梯度的列顺序与重排后的权重矩阵对齐，保持梯度一致性，从而最小化梯度差距。
• 从一个小的候选集（例如12个）中选择最佳重排方案，基于N:M块的数量，避免昂贵的贪心搜索。
• 使用权重大小作为反向掩码的二值化准则，使其与前向掩码对齐，减少梯度差异。

(a) Vanilla N:M Mask

实验结果

研究问题

• RQ1 我们能否在不施加可转置掩码约束的前提下，实现在N:M稀疏训练中的前向与反向完全加速？
• RQ2 解耦前向与反向稀疏掩码对模型性能与训练效率有何影响？
• RQ3 高效的行重排策略是否能提升反向传播中有效N:M块的数量，同时最小化梯度差距？
• RQ4 Bi-Mask是否优于现有方法（这些方法或为加速而牺牲性能，或无法加速反向传播）？
• RQ5 不同二值化准则对反向掩码的有效性及整体模型精度有何影响？

主要发现

• Bi-Mask通过解耦前向与反向传播的稀疏模式，实现了在N:M稀疏张量核心上的前向与反向完全加速。
• 与可转置掩码方法相比，该方法显著减少了性能下降，尤其在高稀疏度（如1:8和1:16）下表现更优。
• 消融实验表明，添加反向掩码仅在ResNet-18上导致0.3%的Top-1精度下降，而重排更新使精度相比基线提升0.3%。
• 在不同二值化准则中，使用权重大小可获得最佳性能（在2:4稀疏度下，ResNet-18的Top-1精度达70.73%），优于基于梯度大小（70.56%）和随机选择（67.76%）。
• 采用有限候选集（如12个）的重排策略已足够实现接近最优的N:M块数量，计算效率高。
• Bi-Mask超越了依赖可转置掩码的SOTA方法，甚至在精度上超过非加速基线，证明了其在效率与有效性上的双重优势。

(b) Transposable N:M Mask