【论文笔记】ResRep: Lossless CNN Pruning via Decoupling Remembering and Forgetting

Abstract

提出了ResRep，一种无损通道修剪的新方法，它通过减少卷积层的宽度（输出通道数）来缩小CNN的尺寸。
建议将CNN重新参数化为记忆部分和遗忘部分，前者学习保持性能，后者学习修剪。通过对前者使用常规 SGD 进行训练，对后者使用带有惩罚梯度的新颖更新规则进行训练，实现了结构化稀疏性，然后等效地将记忆和遗忘部分合并到层数更窄的原始架构中。

github仓库

1 Introduction

压缩和加速卷积神经网络（CNN）的主流技术包括稀疏化、通道剪枝、量化、知识蒸馏等。通道剪枝通过减少卷积层的宽度（即输出通道数）来有效计算次数和内存占用，这与其他模型压缩方法相辅相成，因为它产生了原始架构的较薄模型，没有自定义的结构或操作。

然而，由于CNN的表示能力取决于卷积层的宽度，因此很难在不降低性能的情况下减小宽度。理想情况下，如果剪枝通道的参数足够小，剪枝后的模型可以提供与之前相同的性能（即训练后但剪枝前），称为完美剪枝。

由于训练和剪枝都可能降低性能，因此从两个方面评估基于训练的剪枝方法。

• 抗损性。训练阶段倾向于降低准确性（称为训练引起的损害），因为它引入了一些期望的属性，如结构化稀疏性，这可能是有害的，因为优化目标发生了改变，参数偏离了最优解。如果模型在训练期间性能保持高水平，则说它具有高抗损性。
• 可剪性。当训练后将模型剪枝成一个较小的模型时，获得的属性（例如，许多通道接近零）将减少剪枝引起的损害。如果模型承受高剪枝比例且性能下降较小，则说它具有高可剪性。

在本文中，提出了ResRep来解决上述问题，这受到了神经生物学对记忆和遗忘的研究的启发。

• 记忆需要大脑加强某些突触但减弱其他突触，这类似于CNN的训练，使一些参数变大而另一些变小。
• 通过萎缩或丢失突触来消除突触是一种经典的遗忘机制，作为生物神经网络中提高能量和空间效率的关键过程，类似于剪枝。
  神经生物学研究表明，记忆和遗忘分别由Rutabaga腺苷酸环化酶介导的记忆形成机制和Rac调节的突触萎缩机制独立控制，这表明通过两个解耦模块学习和剪枝更为合理。

受到这种独立性的启发，本文提出将“记忆”和“遗忘”解耦，传统范式中这两者是耦合在一起的，因为卷积参数同时参与了“记忆”（目标函数）和“遗忘”（惩罚损失）以达到折衷的目的。
传统方法强制每个通道“遗忘”，并删除了“遗忘最多”的通道。相比之下，本文的方法首先将原始模型重新参数化为“记忆部分”和“遗忘部分”，然后在前者上应用“记忆学习”（即使用原始目标的常规SGD）以保持“记忆”（原始性能），在后者上应用“遗忘学习”以“消除突触”（将通道置零）。

ResRep包括两个关键组成部分：

• 卷积重新参数化（Rep，解耦的方法和相应的等效转换）；
• 梯度重置（Res，用于“遗忘”的更新规则）。

具体而言，在希望剪枝的原始卷积层后插入一个压缩器，它是一个$1\times 1$卷积。在训练过程中，仅对压缩器添加惩罚梯度，选择一些压缩器通道并将它们从目标函数推导出的梯度置零。这样的训练过程使得一些压缩器通道的梯度非常接近于零，并且在没有剪枝引起的损害的情况下被移除。
然后，通过一些线性变换，将压缩器连同前面的卷积等效转换为具有较少通道的单个卷积。这种方法很容易推广到原始卷积后跟批量归一化（BN）的常见情况。在这种情况下，在BN后附加压缩器，并通过首先等效融合卷积-BN为带偏置的卷积（Eq. 4）来在训练后转换卷积-BN-压缩器序列。最终，结果模型具有与原始模型相同的架构（即没有压缩器），但是层次更窄。

Algorithm 1: Pipeline of ResRep channel pruning

输入：训练好的模型$\mathcal{W}$
使用压缩器构建重参数化的模型$\hat{\mathcal{W}}$。
用单位矩阵初始化压缩器，其他部分使用$\mathcal{W}$的原参数。
for i=0 to 最大迭代次数：

向$\hat{\mathcal{W}}$前向传播一个小批量，使用原来的目标函数计算损失，求出梯度。
只对压缩器的梯度应用梯度重置。
使用压缩器重置后的梯度和其它参数的原始梯度更新$\hat{\mathcal{W}}$。
end for
删除$\hat{\mathcal{W}}$中接近于零的压缩器行（如范数小于等于$10^{-5}$）。将$\hat{\mathcal{W}}$的参数等效地转换到${\mathcal{W}}^{\prime }$中。此时${\mathcal{W}}^{\prime }$和$\mathcal{W}$具有相同的结构，但是具有更窄的层。
输出：剪枝后的模型${\mathcal{W}}^{\prime }$。

总结ResRep的特点：

• 高抗损性。为了维持性能，ResRep不改变原始模型（即卷积-BN部分）的损失函数、更新规则或任何训练超参数。
• 高可剪性。压缩器受惩罚梯度驱动，使得许多通道的值足够小，实现了完美的剪枝，即使惩罚强度较轻。
• 针对所需的FLOP全局减少比例，ResRep自动找到每层的最终宽度，无需先验知识，使其成为CNN结构优化的强大工具。
• 端到端训练和易于实现。

总结本文贡献：

• 受神经生物学研究启发，提出了将“记忆”和“遗忘”解耦以进行剪枝。
• 提出了两种技术，Rep和Res，以实现高抗损性和高可剪性。它们可以分开使用，而组合使用效果最佳。
• 在常见的基准模型上取得了最先进的结果，包括在ImageNet上对ResNet50进行了真正的无损剪枝，剪枝比例为54.5%。

3 ResRep for Lossless Channel Pruning

3.1 Formulation and Background

首先介绍卷积和通道剪枝的公式。

令$D$和$C$为输出和输入的通道数，$K$为卷积核的大小，$\mathbf{K}\in {\mathbb{R}}^{D\times C\times K\times K}$是卷积核的参数tensor，$\mathbf{b}\in {\mathbb{R}}^D$为可选的偏差，$\mathbf{I}\in {\mathbb{R}}^{N\times C\times H\times W}$和$\mathbf{O}\in {\mathbb{R}}^{N\times D\times H^{\prime }\times W}$是输入和输出，$⊛$是卷积操作，$B$是广播函数，将$\mathbf{b}$广播城$N\times D\times H^{\prime }\times W^{\prime }$的尺寸。有：
$$\mathbf{O}=\mathbf{I}⊛\mathbf{K}+B(\mathbf{b})\text{\hspace{1em}(1)}$$
对于一个不含偏差项但是后面跟着一个bn层（均值为$\mu$，标准差为$\sigma$，比例因子为$\gamma$，偏差为$\beta \in {\mathbb{R}}^D$）的卷积层，有：
$${\mathbf{O}}_{:,j,:,:}=((\mathbf{I}⊛{\mathbf{K}}_{:,j,:,:})-{\mu }_j)\frac{{\gamma }_j}{{\sigma }_j}+{\beta }_j,\forall 1\le j\le D\text{\hspace{1em}(2)}$$
令$i$为卷积层的编号，为了剪枝卷积层$i$，根据某种规则获得了被剪掉的通道集合 P ( i ) ⊂ { 1 , 2 , ⋯   , D } \mathcal{P}^{(i)}\subset\{1,2,\cdots,D\} P(i)⊂{1,2,⋯,D}，留下其补集 S ( i ) = { 1 , 2 , ⋯   , D } \ P ( i ) \mathcal{S}^{(i)}=\{1,2,\cdots,D\}\backslash\mathcal{P}^{(i)} S(i)={1,2,⋯,D}\P(i)。剪枝操作留下了第$i$个卷积层的${\mathcal{S}}^{(i)}$的输出通道和后续层（第$i+1$个卷积层）的相应输入通道，然后丢弃其他的。偏差或后续BN中的相应输入通道也被丢弃。
获得的新卷积核为：
$${{\mathbf{K}}^{(i)}}^{\prime }={\mathbf{K}}_{{\mathcal{S}}^{(i)},:,:,:}^{(i)},{{\mathbf{K}}^{(i+1)}}^{\prime }={\mathbf{K}}_{:,{\mathcal{S}}^{(i)},:,:}^{(i+1)}\text{\hspace{1em}(3)}$$

3.2 Convolutional Re-parameterization

对于每一个需要修剪的带有BN层（如果有的话）的卷积层（这个卷积层被称为目标层），附加一个$1\times 1$卷积的压缩器，卷积核$\mathbf{Q}\in {\mathbb{R}}^{D\times D}$。给一个训练好的模型$\mathcal{W}$，构建一个重参数化的模型${\mathcal{W}}^{\prime }$，方法是通过$\mathcal{W}$的原始参数初始化卷积层和bn层，$\mathbf{Q}$用单位矩阵初始化，以便重新参数化的模型产生与原始模型相同的输出。
如果目标层没有后续的BN，那么BN可以被看做偏差。
梯度重置(3.3)训练后，修剪压缩器产生的接近于0的通道并将模型转换为${\mathcal{W}}^{\prime }$，其具有与$\mathcal{W}$相同的结构，但是具有更窄的层。
具体来说，对于具有卷积核$\mathbf{Q}$的特定压缩器，修剪范数小于阈值$ϵ$的通道。
形式上，获得要被剪枝的通道集合 P = { j ∣   ∣ ∣ Q j , : ∣ ∣ 2 < ϵ } \mathcal{P}=\{j|\ ||\mathbf{Q}_{j,:}||_2<\epsilon\} P={j∣ ∣∣Qj,:​∣∣2​<ϵ}或者保留通道集合 S = { j ∣   ∣ ∣ Q j , : ∣ ∣ 2 ≥ ϵ } \mathcal{S}=\{j|\ ||\mathbf{Q}_{j,:}||_2\geq\epsilon\} S={j∣ ∣∣Qj,:​∣∣2​≥ϵ}。与Eq.3相似，通过${\mathbf{Q}}^{\prime }={\mathbf{Q}}_{\mathcal{S},:}$来剪枝$\mathbf{Q}$。实验中设置$ϵ=10^{-5}$，这个阈值足够剪枝获得不错的效果了。剪枝后，压缩器的行数会小于列数，如${\mathbf{Q}}^{\prime }\in {\mathbb{R}}^{D^{\prime }\times D},D^{\prime }=|\mathcal{S}|$。
为了将$\hat{\mathcal{W}}$转变为${\mathcal{W}}^{\prime }$，将每个卷积层-bn层-压缩器序列转换为具有${\mathbf{K}}^{\prime }\in {\mathbb{R}}^{D^{\prime }\times C\times K\times K}$和偏差${\mathbf{b}}^{\prime }\in {\mathbb{R}}^{D^{\prime }}$。
首先，等效地将卷积层-bn层融合到用于推理的卷积层中，这个卷积层会产生和原始相同的输出。对于一个具有$\mathbf{K},\mu ,\sigma ,\gamma ,\beta$的卷积层-bn层，用如下的方法构建一个具有$\bar{\mathbf{K}}$卷积核、$\bar{\mathbf{b}}$偏差的卷积核。对于$1\le j\le D$，有：
$${\bar{\mathbf{K}}}_{j,:,:,:}=\frac{{\gamma }_j}{{\sigma }_j}{\mathbf{K}}_{j,:,:,:},{\bar{\mathbf{b}}}_j=-\frac{{\mu }_j{\gamma }_j}{{\sigma }_j}+{\beta }_j\text{\hspace{1em}(4)}$$
找出一种方法构建${\mathbf{K}}^{\prime }$和${\mathbf{b}}^{\prime }$，使得：
$$(\mathbf{I}⊛\bar{\mathbf{K}}+B(\bar{\mathbf{b}}))⊛{\mathbf{Q}}^{\prime }=\mathbf{I}⊛{\mathbf{K}}^{\prime }+B({\mathbf{b}}^{\prime })\text{\hspace{1em}(6)}$$
通过卷积的可加性，得到：
$$\mathbf{I}⊛\bar{\mathbf{K}}⊛{\mathbf{Q}}^{\prime }+B(\bar{\mathbf{b}})⊛{\mathbf{Q}}^{\prime }=I⊛\overline{{\mathbf{K}}^{\prime }}+B({\mathbf{b}}^{\prime })$$
由于$B(\bar{\mathbf{b}})$的每个通道都是常值矩阵，因此$B(\bar{\mathbf{b}}⊛{\mathbf{Q}}^{\prime })$也应当是常值矩阵。由于$\mathbf{I}⊛\bar{\mathbf{K}}$上使用${\mathbf{Q}}^{\prime }$进行的$1\times 1$卷积只是通道级别的重组，因此可以通过重组$\bar{\mathbf{K}}$的通道的方式将${\mathbf{Q}}^{\prime }$合并到$\bar{\mathbf{K}}$中。设转置函数为$T(\bar{\mathbf{K}})$，如$T(\bar{\mathbf{K}})$的尺寸是$C\times D\times K\times K$，于是构造${\mathbf{K}}^{\prime }$和${\mathbf{b}}^{\prime }$的公式可以写为：
$${\mathbf{K}}^{\prime }=T(T(\bar{\mathbf{K}}⊛{\mathbf{Q}}^{\prime }))\text{\hspace{1em}(8)}$$
$${\mathbf{b}}_j^{\prime }=\bar{\mathbf{b}}\cdot {\mathbf{Q}}_{j,:}^{\prime },\forall 1\le j\le D^{\prime }\text{\hspace{1em}(9)}$$
在实践中，转换并保存训练好的重新参数化模型的权重，构建一个具有原始架构但层数更窄且没有BN的模型，并使用保存的权重进行测试和部署。

3.3 Gradient Resetting

本节描述了如何在保持准确性的同时在压缩器中产生结构化稀疏性，首先讨论特定内核$\mathbf{K}$上惩罚的传统用法，以使某些通道的大小更小，即$||{\mathbf{K}}_{\mathcal{P},:,:,:}||\to 0$。令$\Theta$为通用参数集，$X,Y$是数据示例和标签，$L_{\text{perf}}(X,Y,\Theta )$是与性能相关的目标函数（如：用于分类的交叉熵）。传统范式通过预定义的强度因子$\lambda$来添加惩罚项$P(\mathbf{K})$：
$$L_{\text{total}}(X,Y,\Theta )=L_{\text{perf}}(X,Y,\Theta )+\lambda P(\mathbf{K})\text{\hspace{1em}(10)}$$其中$P$可以是$L1$，$L2$，或group Lasso。尤其是group Lasso可以有效地产生通道方式的结构化稀疏性。下面的讨论中用${\mathbf{F}}^{(j)}={\mathbf{K}}_{j,:,:,:}$来表示$\mathbf{K}$中的特定通道。于是group Lasso损失可以表述为：
$$P_{\text{Lasso}}(\mathbf{K})=\sum _{j=1}^D||{\mathbf{F}}^{(j)}|{|}_E\text{\hspace{1em}(11)}$$
其中$||{\mathbf{F}}^{(j)}|{|}_E$是欧几里得范数
$$||\mathbf{F}|{|}_E=\sqrt{\sum _{c=1}^C\sum _{p=1}^K\sum _{q=1}^K{\mathbf{F}}_{c,p,q}^2}\text{\hspace{1em}(12)}$$
以$G(\mathbf{F})$作为梯度，对其进行求导
$$G(\mathbf{F})=\frac{\partial L_{\text{total}}(X,Y,\Theta )}{\partial \mathbf{F}}=\frac{\partial L_{\text{perf}}(X,Y,\Theta )}{\partial \mathbf{F}}+\lambda \frac{\mathbf{F}}{||\mathbf{F}|{|}_E}\text{\hspace{1em}(13)}$$
特定通道$\mathbf{F}$的训练动态非常直接。从一个训练良好的模型开始，$\mathbf{F}$位于局部最优解附近，因此Eq.13的第一项接近于0，但第二项不接近0，因此$\mathbf{F}$被推向0。如果$\mathbf{F}$对性能很重要，目标函数将试图保持其幅度，即第一个梯度项将与第二个梯度项相竞争，因此$\mathbf{F}$最终会比原来更小，这取决于$\lambda$。否则，以极端情况为例，如果$\mathbf{F}$根本不影响$L_{\text{perf}}$，第一项将为0，因此$\mathbf{F}$将通过第二项持续向0增长。换句话说，性能相关的损失和惩罚损失相互竞争，使得$\mathbf{F}$的最终值将反映其重要性，为了方便起见，称之为基于能力的重要性评估。

然而这面临了两个问题：
问题A：惩罚将每个通道的参数偏离了目标函数的最优解。值得注意的是，轻微的偏差可能不会带来负面影响，例如，$L2$正则化也可以被视为轻微的偏差。然而，如果惩罚过于强大，尽管一些通道被置零以进行剪枝，但剩余的通道也变得太小，无法保持表征能力，这是一个不希望出现的副作用。
问题B：通过轻微的惩罚实现高抗损性，我们无法实现高可剪性，因为大多数通道仅仅变得比原来更接近于0，但不足以实现完美的剪枝。

提出通过对由目标函数导出的梯度进行重置，以实现对轻微惩罚的高可剪性。
引入二进制掩码 m ∈ { 0 , 1 } m\in\{0,1\} m∈{0,1}，表示是否希望将$\mathbf{F}$归零。为了便于实现， 没有向目标函数添加任何项（即$L_{\text{total}}=L_{\text{perf}}$），只和往常一样求梯度，然后手动应用掩码，添加惩罚梯度并用所得梯度进行SGD更新：
$$G(\mathbf{F})\leftarrow \frac{\partial L_{\text{perf}}(X,Y,\Theta )}{\partial \mathbf{F}}m+\lambda \frac{\mathbf{F}}{||\mathbf{F}|{|}_E}\text{\hspace{1em}(14)}$$
下一节描述如何决定将哪些通道置零（即为多个通道设置掩码值）。通过这种方式，已经解决了上述两个问题。A) 尽管我们对每个通道的目标相关梯度添加了Lasso梯度，这等价于通过向原始损失添加Lasso损失来偏离最优解，但偏差是轻微的（在实验中$\lambda =10^{-4}$），因此对性能没有害处。B) 当$m=0$时，第一项不再存在与第二项竞争，因此即使λ很小，$\mathbf{F}$也会稳步向0移动。

3.4 The Remembering Parts Remember Always, the Forgetting Parts Forget Progressively

记住的部分永远记住，忘记的部分逐渐忘记。

如果直接应用于卷积核，Res会带来一个问题：一些与目标相关的梯度编码了保持性能所需的监督信息，但被丢弃了。直觉上，参数被强制“遗忘”了一些有用的信息（梯度）。幸运的是，Rep正是解决这个问题的方法，它允许只剪枝压缩器，而不是原始的卷积层。ResRep只强迫压缩器“遗忘”，而所有其他层仍然专注于“记忆”，因此不会失去编码在原始卷积核梯度中的信息。

要将Res与Rep结合起来，需要决定哪些$\mathbf{Q}$的通道将被置零。在训练重新参数化的模型时，只对压缩器添加Lasso梯度。
几个epochs后，$||{\mathbf{Q}}_{j,:}||$会反应通道$j$的重要性（3.3中的基于能力的重要性评估），因此根据$\mathbf{Q}$的值执行通道选择。
令$n$为压缩器的数量，$m^{(i)}$（一个$D^{(i)}$维的二进制向量）是第$i$个压缩器的掩码，定义$t^{(i)}\in {\mathbb{R}}^{D^{(i)}}$是度量向量，有
$$t_j^{(i)}=||{\mathbf{Q}}_{j,:}^{(i)}|{|}_2,\forall 1\le j\le D^{(i)}\text{\hspace{1em}(15)}$$
每次通道选择，对每个压缩器的每个通道，计算其度量值，并将其组织为映射$\mathcal{M}=\{(i,j)\to t_j^{(i)}|\forall 1\le i\le n,1\le j\le D^{(i)}\}$。然后对$\mathcal{M}$的值进行升序排序，依次选择最小的，将其对应$m_j^{(i)}$设为0。当减少的FLOPS达到目标，或已经选择了$\theta$（选择上限）个通道时，停止选择。
未选择的通道的掩码值设置为1。动机很简单：根据基于能力的重要性评估的讨论，就像传统的惩罚损失用于与原始损失竞争并选择具有较小范数的通道一样，使用惩罚梯度来与原始梯度竞争。 更好的是，所有指标值在开始时都是1（因为每个压缩器核都初始化为单位矩阵），这使得在不同层之间进行比较更加公平。将$\theta$初始化为一个小数，每隔几次迭代增加$\theta$，并逐渐重新选择通道以逐渐“遗忘”，避免一次性将太多通道置零。那些掩码为0的通道将变得非常接近0，因此压缩器中出现了结构化稀疏性。