ResRep: Lossless CNN Pruning via Decoupling Remembering and Forgetting
一、论文基本信息论文题目ResRep: Lossless CNN Pruning via Decoupling Remembering and Forgetting作者Xiaohan Ding、Tianxiang Hao、Jianchao Tan、Ji Liu、Jungong Han、Yuchen Guo、Guiguang Ding发表信息ICCV 2021论文链接ICCV Open Access / arXiv官方代码DingXiaoH/ResRep这篇论文发表于ICCV 2021ICCV Open Access 页面显示论文收录于Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021, pp. 4510–4520。论文提出ResRep目标是实现lossless channel pruning也就是在大幅降低 CNN FLOPs 的同时尽量不损失精度。论文摘要中明确指出ResRep 将 CNN 重新参数化为remembering parts和forgetting parts前者负责保持性能后者负责学习剪枝。官方代码仓库为DingXiaoH/ResRepREADME 中说明该仓库对应 ICCV 2021 论文并提供 ResNet-50 / ImageNet 的剪枝流程仓库介绍中也强调ResRep 通过在训练中对 remembering parts 使用普通 SGD、对 forgetting parts 使用带 penalty gradients 的特殊更新规则从而产生结构化稀疏并最终转换回更窄的原始架构。二、论文要解决的问题传统 learning-based channel pruning 常用的做法是对卷积核、BN scale 或 gate 加稀疏正则让一部分通道逐渐变小然后删除这些小通道。典型形式是任务损失 稀疏惩罚例如CrossEntropy λ × Group Lasso CrossEntropy λ × L1 CrossEntropy λ × BN gamma sparsity这种范式的问题是同一组参数既要负责“记住任务信息”又要负责“被稀疏正则推向 0”。论文把这个矛盾称为resistance-prunability trade-off强稀疏惩罚 通道更容易变成 0prunability 高 但重要参数也会被压小性能受损。 弱稀疏惩罚 对性能破坏小resistance 高 但通道很难真正接近 0剪枝不彻底。论文明确指出传统 penalty-based 方法中performance-related loss 和 penalty loss 会相互竞争如果惩罚过强剩余通道也会被压得过小表示能力受损如果惩罚过弱大多数通道只是“变小一点”但达不到可以完美剪掉的程度。所以 ResRep 要解决的问题是能不能让一部分参数专门负责保持性能 另一部分参数专门负责产生可剪枝结构 从而把 remembering 和 forgetting 解耦这就是标题里Decoupling Remembering and Forgetting的含义。三、思想ResRep 的核心思想可以概括为一句话在待剪卷积层后插入一个可被结构化稀疏化的 1×1 compactor让原卷积继续“记住”任务信息让 compactor 负责“忘记”冗余通道训练结束后再把原卷积、BN 和 compactor 等价合并成一个更窄的普通卷积网络。它包含两个核心组件Rep: Convolutional Re-parameterization 用 1×1 compactor 把原始网络重参数化。 Res: Gradient Resetting 只对 compactor 使用特殊梯度更新让部分通道稳定趋近 0。论文自己也明确写到ResRep 包含两个关键组件Convolutional Re-parameterization和Gradient Resetting具体做法是在要剪枝的原卷积层后插入一个1×1 conv compactor。直观理解如下原始卷积 负责 remembering即保持模型原有能力。 1×1 compactor 负责 forgetting即学习哪些通道应该被压到 0。 训练结束 把 Conv-BN-Compactor 等价融合成更窄的普通 Conv。这点和 RepVGG 一类结构重参数化方法很像训练时有额外结构部署时把额外结构合并掉最终模型没有额外推理开销。四、方法细节4.1 普通通道剪枝形式设某个卷积层输入通道数为 (C)输出通道数为 (D)卷积核为输入为输出为普通卷积可以写成其中 B(b)表示将 bias broadcast 到输出 feature map。论文第 3.1 节给出了这个卷积和通道剪枝的基本形式。如果要剪掉第 (i) 层的一些输出通道保留通道集合记为那么不仅要删除当前层输出通道还要删除下一层对应输入通道论文也明确说明剪枝会保留当前层的 surviving output channels并保留下一个卷积层对应的 input channels对应的 bias 或 BN 参数也要删除。4.2 Convolutional Re-parameterizationResRep 在每个目标卷积层后插入一个1×1 compactor如果原卷积输出为 (D) 个通道那么 compactor 也是的 1×1 卷积。初始化时也就是单位矩阵。这样插入 compactor 后模型输出和原模型完全相同原始模型 Conv-BN 重参数化模型 Conv-BN-Compactor 如果 Compactor 初始化为单位矩阵 两者输出完全一致论文明确说明给定训练好的模型 (W)ResRep 构造重参数化模型 (\hat{W})用原模型参数初始化 conv-BN用 identity matrix 初始化 (Q)因此重参数化模型产生与原始模型相同的输出。这一步非常关键因为它保证插入 compactor 不会破坏原模型性能 后续训练可以从原模型等价点开始。4.3 compactor 为什么能剪枝compactor 是一个 1×1 卷积本质上做通道线性组合。如果 (Q) 的某一行接近 0那么该 compactor 输出的某个通道就接近 0。这个输出通道可以被删除。所以 ResRep 的剪枝目标变成不是直接让原卷积 K 的某些 filters 变成 0 而是让 compactor Q 的某些输出通道变成 0。这样做的好处是原始卷积 K 继续负责保持性能。 Compactor Q 负责产生可剪结构。这就是 remembering / forgetting 的解耦。4.4 Conv-BN-Compactor 如何等价合并训练结束后ResRep 不能把 compactor 留在模型里否则会多一个 1×1 卷积增加推理开销。因此论文需要把Conv BN 1×1 Compactor等价转换成一个更窄的 Conv论文先把 Conv-BN 融合成一个带 bias 的卷积然后再把 1×1 compactor 合并进原卷积权重中。论文公式中给出经过 BN 融合后的卷积权重和 bias 可以再与 (Q) 合并最终构造出新的 (K) 和 (b)。其中核心公式为这里 T 表示转置操作表示融合 BN 后的卷积权重和 bias。论文明确说明1×1 compactor 只做 cross-channel recombination因此可以被合并进原卷积核实际部署时作者保存转换后的权重并构造一个没有 BN、具有更窄层宽度的原始架构模型用于测试和部署。所以最终模型没有 compactor也没有额外 1×1 分支。4.5 传统 penalty-based pruning 的问题传统方法通常写成如果使用 Group Lasso则其梯度为这个式子说明每个 filter 同时受到两种力量影响任务梯度 想让 filter 保持任务性能。 稀疏惩罚梯度 想把 filter 推向 0。论文指出重要通道会被任务梯度“拉回来”不重要通道则更容易被惩罚项推向 0但这会带来前面说的 dilemma强惩罚伤性能弱惩罚剪不干净。4.6 Gradient Resetting梯度重置ResRep 的关键创新是Gradient Resetting。它不把稀疏惩罚直接加进 loss而是正常前向和反向计算然后对 compactor 的梯度做手动修改当 (m0) 时也就是说这个通道不再接收任务损失梯度只被惩罚梯度持续推向 0。论文明确说明ResRep 不向目标函数添加额外项而是正常求出梯度后手动应用 mask、添加 penalty gradients再用修改后的梯度进行 SGD 更新当 (m0) 时performance gradient 不再和 penalty gradient 竞争因此即使很小也能让该通道持续向 0 移动。直观理解m 1 这个 compactor 通道还在“记任务”。 m 0 这个 compactor 通道开始“忘记”只往 0 走。4.7 为什么不能直接对原卷积做 Gradient Resetting如果直接对原卷积 (K) 做 Gradient Resetting被 mask 的通道会失去任务梯度。这意味着原始卷积中的某些参数被迫忘记 它们不再接收监督信息 可能损害模型性能。ResRep 用 compactor 解决这个问题原卷积 K 仍然接收正常任务梯度 继续 remembering。 compactor Q 被选择的通道执行 Gradient Resetting 负责 forgetting。论文明确解释直接把 Res 用在 conv kernels 上会丢弃一些用于保持性能的监督梯度而 Rep 正好解决这个问题因为 ResRep 只强迫 compactors “forget”其他层仍然专注于 remembering。这就是 ResRep 的核心逻辑。4.8 如何决定哪些 compactor 通道开始 forgetting在训练重参数化模型时ResRep 会给 compactors 加 Lasso gradients。经过若干 epoch 后某些 compactor 通道的范数会变小。论文认为||Q_j,:|| 越小 第 j 个通道越适合被剪。训练中会逐步选择 (t) 值最低的通道把它们的 mask 设置为 0然后这些通道在 Gradient Resetting 下稳定趋近于 0。在实验设置中论文对 ResNet-50 / MobileNet 使用相同超参数(\lambda10^{-4})batch size 为 256初始学习率 0.01cosine annealing 训练 180 epochs先 warm-up 5 epochs然后选取 4 个最低 (t) 值的通道每 200 个 batches 再多选 4 个直到达到目标 FLOPs reduction。这说明 ResRep 不是一次性剪枝而是先 warm-up ↓ 逐步选择低重要性 compactor 通道 ↓ Gradient Resetting 让它们趋近 0 ↓ 达到目标 FLOPs 后转换为窄模型五、完整算法流程ResRep 的整体流程可以写成输入 一个训练好的 CNN 目标 FLOPs reduction 目标剪枝层集合 Step 1: 对每个目标卷积层插入 1×1 compactor。 compactor 初始化为单位矩阵。 Step 2: 原卷积和 BN 继承预训练模型参数。 此时重参数化模型和原模型输出完全一致。 Step 3: 正常训练模型使用原任务 loss。 原卷积、BN、分类头等正常 SGD 更新。 Step 4: 对 compactors 应用 Gradient Resetting 对未选中通道 使用任务梯度 penalty gradient。 对选中待剪通道 重置任务梯度只保留 penalty gradient。 Step 5: 根据 compactor 通道范数逐步选择待剪通道。 被选中的通道 mask 设为 0。 Step 6: 训练继续进行待剪 compactor 通道逐渐接近 0。 Step 7: 达到目标 FLOPs 后删除 close-to-zero 通道。 Step 8: 将 Conv-BN-Compactor 等价转换为更窄的普通卷积。 Step 9: 得到最终部署模型。论文 Algorithm 1 也描述了类似过程输入 well-trained model构造带 compactors 的重参数化模型compactors 初始化为 identity前向计算原始 objective随后只对 compactors 应用 Gradient Resetting再更新模型最后转换为更窄模型。六、关键公式6.1 原始卷积6.2 Conv-BN 融合对于带 BN 的卷积层6.3 通道剪枝6.4 传统 penalty-based pruning6.5 Group Lasso6.6 传统 penalty 梯度6.7 Gradient Resetting6.8 Conv-BN-Compactor 合并七、实验设置论文主要在ImageNet和CIFAR-10上验证。ImageNet 实验包括ResNet-50和MobileNetCIFAR-10 实验包括ResNet-56和ResNet-110。论文使用 torchvision 官方 ResNet-50 作为基线其 Top-1 为 76.15%MobileNet 从头训练Top-1 为 70.78%CIFAR-10 使用标准数据增强包括 padding 到 (40\times40)、random crop 和 flip。FLOPs 统计方面论文将 ResNet-50 记为 4.09G FLOPsMobileNet 为 569MResNet-56 / ResNet-110 分别为 126M / 253M。官方代码仓库中提供 ResNet-50 / ImageNet 剪枝命令说明流程包括下载并转换 torchvision ResNet-50 权重、运行rr/exp_resrep.py、保存剪枝后的finish_converted.hdf5并可显示剪枝模型的权重名称和形状。八、实验结果解读8.1 ResNet-50 / ImageNet54.54% FLOPs reduction无 Top-1 精度下降ResRep 最重要的实验结果是 ResNet-50 / ImageNet。论文报告原始 ResNet-50 Top-1 76.15% Top-5 92.87% ResRep Pruned Top-1 76.15 ± 0.01% Pruned Top-5 92.89 ± 0.04% FLOPs reduction 54.54%也就是说在减少54.54% FLOPs的情况下Top-1 没有下降Top-5 还略高。论文摘要也强调它把标准 ResNet-50 从 76.15% ImageNet accuracy 剪到只剩约 45% FLOPs且没有 accuracy drop。这就是论文标题中Lossless CNN Pruning的主要依据。8.2 更高压缩率下仍优于 HRank / LFPC 等方法ResRep 还报告了更高 FLOPs reduction 的结果。在 ResNet-50 / ImageNet 上56.11% FLOPs reduction Top-1 75.97 ± 0.02% Top-5 92.75 ± 0.01% 62.10% FLOPs reduction Top-1 75.30 ± 0.01% Top-5 92.47 ± 0.01%论文表 1 中显示在 62.10% FLOPs reduction 下HRank 的 Top-1 为 71.98%而 ResRep 为 75.30%在 60.8% FLOPs reduction 下LFPC 的 Top-1 为 74.46%ResRep 在相近甚至更高压缩率下仍有更好结果。这说明 ResRep 的优势不仅在轻度剪枝在中高压缩率下也很明显。8.3 MobileNet / ImageNet轻量网络上仍有效MobileNet 本身已经是轻量网络因此剪枝更难。论文表 1 中MetaPruning 在 73.81% FLOPs reduction 下 Top-1 为 66.1%ResRep 在 73.91% FLOPs reduction 下 Top-1 达到 68.02%比 MetaPruning 高 1.77%。论文正文也明确指出在 MobileNet 上 ResRep outperform MetaPruning by 1.77%。这点很重要因为很多剪枝方法在 VGG / ResNet 上有效但在 MobileNet 这类高度压缩结构上很难继续剪。8.4 ResNet-56 / CIFAR-1052.91% FLOPs reduction0% 精度下降在 CIFAR-10 上ResRep 测试了 ResNet-56 / ResNet-110。ResNet-56Base Top-1 93.71% Pruned Top-1 93.71 ± 0.02% FLOPs reduction 52.91% Top-1 drop 0.00%ResNet-110Base Top-1 94.64% Pruned Top-1 94.62 ± 0.04% FLOPs reduction 58.21% Top-1 drop 0.02%论文表 2 中显示ResRep 在 ResNet-56 / ResNet-110 上也取得了非常小的精度下降并优于 FPGM、SFP、LFPC、HRank 等多个对比方法。8.5 自动发现不同数据集上的层宽结构ResRep 不需要预先指定每层剪枝比例。给定全局 FLOPs 目标它会通过 compactor 的稀疏化过程自动决定各层最终宽度。论文 Figure 3 和正文指出ResRep 在 ResNet-50 / MobileNet 上倾向于保留更多高层通道而在 ResNet-56 / CIFAR-10 上会更激进地剪最后几个 blocks。作者解释说ImageNet 任务更难高层语义特征对拟合能力很重要而 ResNet-56 在 CIFAR-10 上可能存在过拟合因此后层可以更激进剪枝。这说明 ResRep 不是简单均匀缩小网络而是能根据任务和网络自动形成非均匀宽度结构。8.6 消融实验Rep 和 Res 都重要论文做了多组消融Only Rep: 只使用 re-parameterization但不用 Gradient Resetting。 Only Res: 直接对原卷积做 Gradient Resetting。 ResRep: 同时使用 Rep 和 Res。结果显示完整 ResRep 最好。论文表 3 中还比较了Base model fine-tuned Uniformly shrunk baseline Pruned-finetuned baseline Vector re-parameterization Momentum on compactors 0.9其中 uniformly shrunk baseline 在 55.4% FLOPs reduction 下 Top-1 为 74.39%pruned-finetuned baseline 在 54.5% FLOPs reduction 下 Top-1 为 74.66%而 ResRep 在相近压缩率下能达到 76.15% 或 75.97%。这说明 ResRep 的效果不是简单来自训练设置也不是只来自更长 fine-tuning而是来自Compactor re-parameterization Gradient Resetting的组合。九、方法优点9.1 记忆和遗忘解耦ResRep 最大贡献是把传统稀疏训练中耦合在同一组参数上的两个目标分开原卷积 remembering负责保持性能。 compactor forgetting负责产生剪枝结构。这直接缓解了传统 penalty-based pruning 的 resistance-prunability trade-off。9.2 结构重参数化思想清晰训练时插入 compactor部署时等价合并回更窄的普通卷积。最终模型没有额外 1×1 卷积没有特殊 mask也没有额外推理分支。论文明确说明转换后模型具有原始架构形式但层更窄可以直接用于测试和部署。9.3 不需要手动指定每层剪枝率ResRep 只需要全局 FLOPs reduction target然后通过 compactor 通道选择自动形成每层最终宽度。论文指出ResRep 在没有任何 prior knowledge 的情况下发现合适的最终结构。9.4 对高压缩率 ResNet-50 实现近似无损在 ResNet-50 / ImageNet 上54.54% FLOPs reduction 下 Top-1 0.00% drop这是 ResRep 最有代表性的结果。9.5 对 MobileNet 等轻量模型也有效MobileNet 本身冗余更少但 ResRep 在 73.91% FLOPs reduction 下仍显著优于 MetaPruning。十、方法局限10.1 训练流程比简单剪枝复杂ResRep 需要插入 compactor 初始化 identity 实现 Gradient Resetting 动态选择待剪通道 维护 mask 训练后做等价转换相比 L1 / FPGM / HRank 这类排序后剪枝方法工程实现复杂度更高。10.2 依赖训练过程不是后训练剪枝ResRep 不是直接拿一个预训练模型算重要性并剪掉而是需要进行较长的 re-parameterized training。例如 ResNet-50 / MobileNet 实验中使用 180 epochs 的训练流程。因此它不适合“完全无训练成本”的剪枝需求。10.3 对不同结构需要设计目标层论文中 ResNet-50 主要剪 residual block 中第一个 (1\times1) 和第二个 (3\times3) 卷积层MobileNet 剪 non-depthwise convResNet-56/110 剪 residual block 的第一层。也就是说不同网络结构仍需要决定哪些层适合插入 compactor 和剪枝。10.4 “Lossless” 依赖实验设置ResRep 的 lossless 主要指在 ResNet-50 / ImageNet 的 54.54% FLOPs reduction 下 Top-1 无下降。更高压缩率下仍会出现精度下降例如 62.10% FLOPs reduction 时 Top-1 drop 为 0.85%。所以不能泛化理解为“任何模型任何压缩率都无损”。10.5 对 Transformer / LLM 不能直接照搬ResRep 原始方法面向 CNN 通道剪枝依赖卷积层后插入 1×1 compactor并通过卷积重参数化合并。对于 ViT、LLM、VLM剪枝对象可能是attention heads MLP hidden neurons tokens layers vision tokens KV cache channels但它的思想可以迁移给结构单元插入可合并的 compactor / adapter ↓ 原模块负责 remembering ↓ compactor 负责 forgetting ↓ 训练后把 compactor 等价合并或删除这对后续结构重参数化剪枝、head pruning、MLP neuron pruning 都有启发。十二、一句话总结《ResRep: Lossless CNN Pruning via Decoupling Remembering and Forgetting》提出通过结构重参数化在卷积层后插入 1×1 compactor让原卷积负责保持性能、compactor 负责产生可剪稀疏结构并通过 Gradient Resetting 让待剪 compactor 通道稳定趋近于 0训练结束后Conv-BN-Compactor 被等价合并为更窄的普通卷积网络从而在 ResNet-50 / ImageNet 上实现 54.54% FLOPs reduction 且 Top-1 无下降。