非结构化剪枝算法的巧妙利用：非结构化掩码的后处理

2023-02-01 15:02

©PaperWeekly 原创 · 作者 | GlobalTrack

经典彩票假说算法通过 IMP 剪枝算法寻找稀疏子网。通过 IMP 剪枝得到的掩码中非 0 元素一般是分散的，这种非结构化掩码对硬件加速是非常不利的。本文介绍近期提出的两种通过特殊的后处理算法将彩票假说得到的非结构稀疏掩码转变为合适的结构化稀疏掩码。这里两种结构化剪枝关注于通道剪枝、分组剪枝与 N:M 剪枝。

Structure-LTH

论文标题：

Coarsening the Granularity: Towards Structurally Sparse Lottery Tickets

收录会议：

ICML 2022

论文链接：

https://proceedings.mlr.press/v162/chen22a/chen22a.pdf

代码链接：

https://github.com/VITA-Group/Structure-LTH

1.1 简介

彩票假说（Lottery Ticket Hypothesis）及其变种已经有了很多研究，然而有一个问题仍然很有挑战性：由于获胜的彩票网络由 IMP 剪枝算法发现，剪枝后的子网络只有非结构化稀疏度。由于高度不规则的稀疏掩码对硬件加速非常不友好，非结构化剪枝一般是很难在实际上应用的，这使得彩票假说获得的稀疏模型难以运用到实际项目。另外如果直接将 IMP 中的非结构化剪枝替换为结构化剪枝，这回带来严重的性能下降。

本文给出了第一个关于结构化彩票假说的研究，探索通过彩票假说s相关理论运用于寻找硬件友好的稀疏度（这里研究通道剪枝和组剪枝）。本文的核心思想是通过后处理机制，重新填充（Refill）部分连接得到通道级别的结构化稀疏度。另外本文提出一种基于超图划分（hypergraph partition）的重分组（regroup）机制实现组级别的结构化稀疏度。

1.2 方法

1.2.1 重新填充机制

核心步骤：

1. 根据选择的评价指标在非结构剪枝后的子网中寻找重要的通道；

2. 重要的元素重新生长（Regrow），具体数值选择为随机初始化或类似彩票假说选择早期未剪枝前对应位置的权重数值；

3. 剩下在不重要的通道的元素会被移除。

评估指标多样，可以是通道内剩下权重数量，通道内权重统计特征，特征统计特征或显著性分数。

这里给出了另一个较弱的版本，Refilling+，对清楚所有剩余通道元素做出补偿。具体方法是重新激活额外的一部分通道用于减缓网络容量下降速度。

1.2.2 重新填充机制

重新分组机制主要目标即在未剪枝的元素中寻找密集块（dense block）。与原始的稀疏矩阵比较，密集块尺度小得多。

具体方法是：寻找相似的行和列，将他们聚在一起。这里使用 Jaccard 相似度计算非零元素所在行的相似度。比如 kernel 1，kernel 2 在八个非零列中有三个相同的，则他们的相似度为 3/8。

这里使用超图划分算法生成密集块。这里将每一行作为超图的一个节点，每一列作为超图一条边。这里的超图划分使用 hMETIS（http://glaros.dtc.umn.edu/gkhome/metis/hmetis/overview）算法。划分后将缺失的元素重新填充，剩余不包括在密集块的元素将被置 0。

1.3 实验

本文方法在 Tiny-ImageNet 和 ImageNet 数据集和 ResNet 模型上的实验结果。在 Tiny-ImageNet 上获得的结果：

1. 结构化获胜彩票网络有 IMP-Refill 和 IMP-Regroup 算法获得 60% 的结构化通道稀疏度和 74% 的组稀疏度，该结果验证了本文方法的有效性；

2. 从相同的非结构化子网出发，在高稀疏度条件下（>50%），IMP-Refill+超过了 IMP-Refill。然而在大部分稀疏度条件下，IMP-Refill 超过了 IMP-Refill+性能。可能的一个解释是 IMP-Refill+带入了不希望的通道；

3. IMP-Regroup 在高稀疏度条件下获得了表现较好；

4. 极端的通道（45%-50%）/组稀疏度（74%）条件下，取得了 57.53% ∼ 61.79% / 64.84% 的 GPU 推断时间减少。

ImageNet 数据集上结果：

1. Refill 和 Refill+策略几乎没有差别，且不能找到通道级别的获胜彩票子网。

2. 组级别的获胜彩票子网在 31% 稀疏度仍然存在，另外低稀疏度之带来 1% 推断时间节省。

CIFAR 数据集上结果：

1. 在 (WRN-32-2,C10), (WRN-32-2,C100), (MBNet-v1,C10), (MBNet-v1,C100), (VGG-16,C10), (VGG-16,C100), (RN-18,C10), (RN-18,C100) 的数据集和深度模型组合上分别在 {53%, 28%, 67%, 0%, 60%, 40%, 50%, 0%} 的通道稀疏度与 {66%, 36%, 72%, 56%, 80%, 80%, 78%, 78%} 的组稀疏度上获得了获胜彩票子网；

2. 在较低的稀疏度条件下（<50%），IMP-Regroup 与 IMP-Refill 性能相近。IMP-Regroup 在高稀疏度条件下精度较高。

下图给出了使用不同剪枝算法得到初始非结构化掩码的比较实验。

Shuffled Block Sparse Training

论文标题：

Exposing and Exploiting Fine-Grained Block Structures for Fast and Accurate Sparse Training

收录会议：

NeurIPS 2022

论文链接：

https://openreview.net/pdf?id=sFapsu4hYo

代码链接：

https://github.com/HPC-Research-Lab/shuffled-block-sparse-training

2.1 简介

目前的稀疏学习一般可以分为静态方法与动态方法。动态稀疏学习只需要随机初始化的参数，经过少量训练周期后，丢弃一部分不重要的元素并激活一部分新参数。通过探索更大的参数组合，动态稀疏学习一般可以比静态方法取得更高的性能。

本文是一种基于动态稀疏学习的网络剪枝算法。之前动态稀疏学习算法关注于非结构化剪枝任务，剪枝后的参数分散且没有结构化约束。尽管非结构化剪枝一般可以在较少的参数条件下获得更高的精度，然而难以获得硬件上的加速。

本文方法关注于探索在动态稀疏学习算法中寻找细粒度块（Fine-grained Block）结构。主要思想是将参数分组为若干洗牌块（Shuffled Block）。由于洗牌块和他们的转置可以很好地由硬件处理，硬件可以加速权重的向前传递和向后传递。

2.2 相关论文

动态稀疏学习：包含两个主要操作：丢弃和再生（drop, grow）。丢弃即舍弃模型中一些权重，再生即重新加回一些权重。稀疏学习算法一般在丢弃和再生阶段使用不同的权重重要性评价指标。

RigL（Rigging the lottery: Making all tickets winners）丢弃较小幅度值的权重，再生阶段选择较大梯度的权重。Top-kast（Top-kast: Top-k always sparse training.）选择一部分一直激活的权重，根据幅度值选择丢弃和再生的权重。

2.3 方法

本文算法主要步骤：

洗牌分块阶段（Shuffled Blocking）： 假定激活状态的权重数为，块大小为，分块阶段将权重分配在个块中。与上一篇论文 Regroup 策略类似，本文也根据行的相似度进行分组，这里使用层次聚类算法。另外注意到大部分行元素与其他元素只有较弱的相似度，重新排序行元素只能带来较小的密集度提升。

这里只选择在相同列有超过 q 比例元素的行进行聚类。在行完成重排后，将行划分为组。在每一个行组中，选择形成大小为的块。为了确保大幅度值的权重能被保留，这里往往选择有最大范数的列，并返回有最大范数的块。

块感知丢弃策略：在周期训练阶段，丢弃的参数量计算如下：

是初始丢弃率。

与 RiGL 等动态稀疏学习算法类似的，本文也选择基于幅度值的评价指标。丢弃权重可以描述为：。另外这里需要检查是否与洗牌快的行或列有超过一般元素。对于每一个在洗牌块的元素，如果在相同行或列没有足够的非 0 元素，该元素不会被丢弃。主要原因是丢弃后有很大可能在分块阶段会被加回。

一般的，如果较多小幅度值的元素聚集在相同的行或列，本文方法会执行较多的模型适应（model adaption，丢弃重生机制）。如果较多的小幅度值元素分散在不同的行和列，本文方法会执行较弱的模型适应。

块感知重生机制：之前 RigL 算法需要计算每一个权重梯度，计算量较大。本文使用输出行梯度和输入行计算重要性指标。这里的梯度定义为：。这里的重生机制将输入通道有较大幅度值的通道与有较大梯度值的通道连接的权重是相对重要的。

考虑到计算成本，不需要计算梯度的准确值。这里用表示组内对应行的元素，用表示组内对应行的元素。这里选择为梯度幅度值的估计。

2.4 实验

ResNet50 和 ResNet18 模型与不同剪枝策略在 ImageNet 数据集上的实验结果：

本文方法（DSB）可以在 50% 稀疏度条件下恢复模型精度。类似的方法 SSLT（第一个介绍的算法）在类似稀疏度条件下精度较低。另外本文方法在 75% 稀疏度条件下精度也比 70% 稀疏度条件下的 SSLT 高。

图 7 给出了 WideResNet, ResNet 和 VGG 与不同剪枝策略在 CIFAR10 数据集上的性能。可以看出本文方法 DSB 在 50% 和 60% 稀疏度条件下也能恢复原始模型精度。类似算法 SSLT 有明显精度下降。图 8 给出了在 CIFAR100 数据集上的性能。类似地，本文方法 DSB 也在 50% 和 60% 稀疏度条件下恢复了原始模型地精度。另外精度差异也随着稀疏度增大明显增大。