Alternating Compressed/DeCompressed

AC/DC: Alternating Compressed/DeCompressed Training of Deep Neural Networks

• 论文链接：https://arxiv.org/pdf/2106.12379.pdf
• 源码链接：https://github.com/IST-DASLab/ACDC

简介

深度神经网络在解决各种任务方面取得巨大进展，推动了研究和行业对部署这些模型的有效版本浓厚兴趣。为此，已经开发了完整的模型压缩家族，如剪枝和量化，这些方法现在都有硬件和软件支持。

神经网络剪枝是本文的关注领域，这是一种由最长历史的压缩方法。剪枝的基本目标是获得很多连接通过设置为0而移除的神经网络，同时维持网络的精度。已经提出了无数种剪枝方法，目前已经了解到，将连接数量而言，许多流形的网络可以被压缩一个数量级以上，而不会造成显著的精度损失。

很多的精确剪枝方法需要一个完全精确的，模型的密集变体，之后权重移除的对应模型。该方法主要缺点是由于压缩造成的内存和计算节省只在推断时可用（后训练阶段）而在训练时不可用。这一区分对于可以用百万甚至上亿参数的大尺度现代模型变得特别重要，这些模型的全密集由很高计算甚至非平凡的环境成本。

解决该问题的一个办法是稀疏训练，目标是在训练时尽可能早地从神经网络移除连接，同时匹配或至少接近全密集模型精度。例如，RigL技术在训练早期移除打分数地连接，之后继续在稀疏支持上优化，这由于稀疏后向传播提供了节约。周期性地，该方法在训练过程中重引入了一些权重，这是基于启发组合，即需要采用全梯度。该研究与很多近期地稀疏训练方法，已经经验性地显示出非平凡的计算节省，这通常由理论的FLOPs度量，这可以使用稀疏训练获得，另外优化过程可以公平地对稀疏支持鲁棒。

同时，这些研究留下了开放性问题：第一个是理论性的：就本文所致，没有在稀疏支持的提供训练加速的方法显示出收敛保证。第二个是实践性的，这关乎对密集训练模型与稀疏训练模型之间的更深入理解。主要问题即：1. 大部分稀疏训练方法仍然留下了与密集训练相比甚至是后训练稀疏化相比不可忽视的精度差距，2. 大多数现有的稀疏性研究都需要对训练流程进行重大更改，并侧重于最大化全局精度指标；因此，当涉及到共同训练稀疏和密集模型时，以及在个体预测层面上稀疏和密集模型之间的相关性时，本文缺乏理解。

本文朝着解决这些问题迈出一步。研究了一种用于神经网络稀疏训练的通用混合方法，本文称之为交替压缩/去压缩（Alternating Compressed/Decompressed，AC/DC）训练。AC/DC执行稀疏和密集模型的联合训练，并且可以返回准确的稀疏模型和密集模型。本文证明了在分析假设下，AC/DC的一个变体确保了一般非凸但光滑目标的收敛性。

AC/DC建立在用于稀疏恢复的经典迭代硬阈值（Iterative Hard Thresholding，IHT）系列方法基础上。顾名思义，AC/DC通过将标准密集训练阶段于稀疏阶段交替来执行，其中优化尽在固定的稀疏支持上执行，并且权重的子集及其梯度固定为0，从而节省了计算量。这与误差反馈算法形成对比。后者需要计算全密集梯度，即使权重本身可能是稀疏的。该过程使用相同的超参数，包括作为常规训练周期的数量，并且阶段的频率和长度可以安全地设置为标准值，例如5-10个周期。本文确保训练在稀疏阶段结束，并返回生成的稀疏模型，以及在密集阶段结束时获得最后一个密集模型。这种密集模型可能会短时间内进行进行额外的微调，从而产生与更准确的密集微调模型，本文通常会发生与密集baseline的精度相匹配。

相关研究

Dense-Sparse-Dense范式

稀疏和密集训练阶段交替的思想以前已经在神经网络背景下进行过研究，但目的是使用临时稀疏化作为正则化器。 DSD: Dense-sparse-dense training for deep neural networks提出首先将密集模型训练到完全精确，然后通过幅度对该模型进行稀疏化；接下来，在稀疏支持上执行优化，然后在全密集支持上执行额外的优化阶段。因此该过程被用于密集模型的正则化机制，这导致相对于原始密集模型相对较小但一致的精度提升。Training skinny deep neural networks with iterative hard thresholding methods提出了一种类似于DSD的方法，但在常规训练过程中交替使用稀疏阶段。生成过程类似于AC/DC，但重要的是他们程序目标时返回更准确的密集模型。为此作者使用相对较低的稀疏性级别，并在优化过程中逐渐增加稀疏性；他们观察到以增加训练周期总数为代价，得到的密集模型精度有所提高。相比之下，本文重点是获得准确的稀疏模型，同时降低计算成本，并执行密集训练配方。本文在更高稀疏性级别以及更大规模的数据集和模型上执行。此外，本文还证明了该方法适用于其他稀疏性模式，例如2：4非结构化模式。

本文方法

AC/DC从标准DNN训练流开始，使用标准优化器，如带动量的SGD或Adam，并保留所有标准训练超参数。它会定期修改对优化的支持。

本文将训练周期集合划分为压缩周期$\mathcal{C}$和去压缩周期$\mathcal{D}$。本文从$\Delta w$连续周期的密集预热周期开始，在此期间执行规则的密集（解压缩）训练。然而，本文开始交替压缩优化阶段，每个阶段长度为${\Delta }_c$周期，解压缩（常规）优化阶段长度为${\Delta }_d$周期。该过程在压缩微调阶段完成，返回精确的稀疏模型。或者，如果本文目标是返回与baseline精度匹配的密集模型，本文采用在交替过程中获得的最佳密集检查点，并在整个支持过程中对其进行微调。注意到允许更长的最终去压缩阶段周期${\Delta }_D>{\Delta }_d$以提升密集模型性能，通过允许它在微调后更好地恢复baseline精度。

本文实验中专注于压缩操作是非结构化或半结构化的情况。在这种情况下，在每个稀疏优化阶段开始时，本文在所有网络权重上应用top-k算子，以获得权重$\theta$上的掩码$M$。TOP-K算子全局应用于所有网络权重，并将表示稀疏支持，在该稀疏支持上将对当前稀疏阶段的其余部分执行优化。在稀疏阶段结束，掩码M被重置为全1，以至于后续的密集阶段将在全密集支持上优化。一旦所有权重被重引入，优化器的梯度动量项重置为0是有益的；这对于先前剪枝的权重是特别有用的，否则只有陈旧版本的梯度。

在DNN背景下，从IHT转向稳健的实现需要一些调整。首先，每个解压缩阶段豆可以映射到确定性/随机IHT步骤，其中执行几个随机步骤，而不是连续截断支持之间的单个梯度步骤。这些额外的步骤在实践中提高了方法的准确性，我们也可以在理论上限制它们的影响，尽管它们不一定提供更好的界限。