[论文解读] Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks

Itay Hubara, Brian Chmiel+4 more|arXiv (Cornell University)|Feb 16, 2021

Advanced Neural Network Applications被引用 25

一句话总结

本文提出一种可转置的 N:M 细粒度稀疏性掩码以加速前向和后向传播，将掩码优化表述为最小成本流问题（带有快速的2近似算法），并在视觉与语言模型上展示2x训练加速，且精度损失可忽略，同时给出通过 AdaPrune 将无结构稀疏性适配到 N:M 掩码的方法。

ABSTRACT

Unstructured pruning reduces the memory footprint in deep neural networks (DNNs). Recently, researchers proposed different types of structural pruning intending to reduce also the computation complexity. In this work, we first suggest a new measure called mask-diversity which correlates with the expected accuracy of the different types of structural pruning. We focus on the recently suggested N:M fine-grained block sparsity mask, in which for each block of M weights, we have at least N zeros. While N:M fine-grained block sparsity allows acceleration in actual modern hardware, it can be used only to accelerate the inference phase. In order to allow for similar accelerations in the training phase, we suggest a novel transposable fine-grained sparsity mask, where the same mask can be used for both forward and backward passes. Our transposable mask guarantees that both the weight matrix and its transpose follow the same sparsity pattern; thus, the matrix multiplication required for passing the error backward can also be accelerated. We formulate the problem of finding the optimal transposable-mask as a minimum-cost flow problem. Additionally, to speed up the minimum-cost flow computation, we also introduce a fast linear-time approximation that can be used when the masks dynamically change during training. Our experiments suggest a 2x speed-up in the matrix multiplications with no accuracy degradation over vision and language models. Finally, to solve the problem of switching between different structure constraints, we suggest a method to convert a pre-trained model with unstructured sparsity to an N:M fine-grained block sparsity model with little to no training. A reference implementation can be found at https://github.com/papers-submission/structured_transposable_masks.

研究动机与目标

• 通过一种新的度量方法称为掩码多样性（MD），对不同稀疏掩码如何影响模型准确性进行排序与理解。
• 开发一种可转置的 N:M 稀疏掩码，在训练期间同时加速前向和后向传播。
• 提供高效算法以寻找最优掩码（最小成本流及快速2近似），并实现无需重新训练的掩码切换（AdaPrune）。
• 在视觉和语言任务中演示矩阵乘法的2x加速，且精度下降可忽略。

提出的方法

• 将掩码多样性（MD）定义并证明其意义，即在结构约束下给定稀疏度时有效掩码的数量。
• 将可转置的 N:M 掩码设计表述为优化问题，目标是在每个 MxM 块中最大化保留的 L1 范数，且每行/每列有 N 个非零。
• 将最优掩码问题化简为对预训练密集模型的最小成本流问题。
• 引入从头开始训练时的快速2近似贪心算法，时间几乎线性于块大小。
• 提出 AdaPrune，将无结构稀疏转化为 N:M 结构掩码，且仅需最少的再训练。
• 提供基于 AdaPrune 的标定方法，在切换稀疏结构时精度损失很小。

实验结果

研究问题

• RQ1 掩码多样性能否预测稀疏结构对模型准确性的影响？
• RQ2 我们能否设计一种可转置的 N:M 稀疏掩码，同时加速前向和后向传播？
• RQ3 在训练过程中寻找最优或近似最优的可转置掩码的高效算法是什么？
• RQ4 我们能否在尽量少的再训练下将无结构稀疏转换为 N:M 稀疏（AdaPrune）？

主要发现

• 4:8 可转置掩码在精度上可与 2:4 掩码媲美，但通过在前向和后向传播中使用稀疏张量核心实现训练加速。
• 在视觉和语言模型中展示了矩阵乘法的2x加速，并且没有精度下降。
• 精确的最小成本流公式为预训练密集模型提供最优的可转置掩码，而2近似算法在近线性时间内运行，解的效果在最优解的2倍内。
• AdaPrune 可以将无结构稀疏模型转变为 4:8 的可转置稀疏模型，在某些设置下精度下降小于1%，实现硬件友好部署且无需完全重新训练。
• 经验结果显示训练过程中稀疏张量核心利用率达到2x，与仅加速前向传播的先前方法相比，精度相当。