[论文解读] LIT: Block-wise Intermediate Representation Training for Model Compression

Animesh Koratana, Daniel Kang+2 more|arXiv (Cornell University)|Oct 2, 2018

Advanced Neural Network Applications参考文献 25被引用 13

一句话总结

LIT 是一种新颖的模型压缩技术，通过逐块匹配中间表征（IR）来训练更窄的教师网络，使其与更宽、更深的教师网络对齐。通过将教师的中间表征作为每一学生模块的输入，LIT 稳定了训练过程，并实现了对现代架构（如 ResNeXt 和 VDCNN）高达 5.5 倍的压缩，且无准确率损失，优于知识蒸馏和提示训练方法，首次将学生/教师训练扩展至 GAN 生成器。

ABSTRACT

Knowledge distillation (KD) is a popular method for reducing the computational overhead of deep network inference, in which the output of a teacher model is used to train a smaller, faster student model. Hint training (i.e., FitNets) extends KD by regressing a student model's intermediate representation to a teacher model's intermediate representation. In this work, we introduce bLock-wise Intermediate representation Training (LIT), a novel model compression technique that extends the use of intermediate representations in deep network compression, outperforming KD and hint training. LIT has two key ideas: 1) LIT trains a student of the same width (but shallower depth) as the teacher by directly comparing the intermediate representations, and 2) LIT uses the intermediate representation from the previous block in the teacher model as an input to the current student block during training, avoiding unstable intermediate representations in the student network. We show that LIT provides substantial reductions in network depth without loss in accuracy -- for example, LIT can compress a ResNeXt-110 to a ResNeXt-20 (5.5x) on CIFAR10 and a VDCNN-29 to a VDCNN-9 (3.2x) on Amazon Reviews without loss in accuracy, outperforming KD and hint training in network size for a given accuracy. We also show that applying LIT to identical student/teacher architectures increases the accuracy of the student model above the teacher model, outperforming the recently-proposed Born Again Networks procedure on ResNet, ResNeXt, and VDCNN. Finally, we show that LIT can effectively compress GAN generators, which are not supported in the KD framework because GANs output pixels as opposed to probabilities.

研究动机与目标

• 解决知识蒸馏过程中深度学生网络中间表征（IR）不稳定的挑战。
• 实现对现代高度结构化深度网络（如 ResNeXt、VDCNN）的有效压缩，这些网络在标准提示训练或 KD 下难以被有效压缩。
• 将学生/教师训练扩展至 GAN 生成器，因其像素级输出与标准 KD 不兼容。
• 通过直接复制教师网络的部分结构并使用教师的中间表征作为学生模块的输入，实现高准确率、紧凑的学生模型。
• 证明当学生与教师架构完全相同时，LIT 可使学生模型的准确率超越教师模型。

提出的方法

• 训练一个与教师网络宽度相同但深度更小的学生网络，采用逐块中间表征（IR）匹配方法。
• 在训练过程中，将教师前一模块的 IR 作为当前学生模块的输入，以稳定内部表征。
• 通过加权插值因子 β 将知识蒸馏（KD）损失与 IR 匹配损失结合，以平衡输出 logits 和 IR 监督。
• 在学生与教师网络的对应模块之间应用 IR 损失，确保宽度和特征图尺寸的兼容性。
• 通过完整复制教师网络中的模块至学生网络，实现选择性压缩，支持模块化模型压缩。
• 通过在残差模块中匹配 IR，实现 GAN 生成器的压缩，绕过 KD 与像素级输出不兼容的问题。

实验结果

研究问题

• RQ1 中间表征匹配能否有效扩展至现代深层网络（如 ResNeXt 和 VDCNN），这些网络因中间表征不稳定而使标准提示训练失效？
• RQ2 与标准 KD 和提示训练相比，使用教师提供的 IR 作为输入的逐块 IR 训练是否能提升学生模型的准确率和泛化能力？
• RQ3 LIT 是否能够压缩 GAN 生成器，而这些生成器因像素级输出与标准知识蒸馏不兼容？
• RQ4 当学生与教师架构完全相同时，LIT 是否能使学生模型的准确率超越教师模型？
• RQ5 LIT 的性能对超参数（如 β（IR 损失权重）和 α（KD 损失权重））的敏感度如何？最优设置是什么？

主要发现

• LIT 在 CIFAR10 上将 ResNeXt-110 压缩为 ResNeXt-20（5.5 倍压缩），准确率无损失，优于 KD 和提示训练。
• LIT 在 Amazon Reviews 数据集上将 VDCNN-29 压缩为 VDCNN-9（3.2 倍压缩），准确率无下降，证明其在 NLP 任务中的有效性。
• 在压缩相同架构时，LIT 在 ResNet、ResNeXt 和 VDCNN 上使学生模型准确率超越教师模型，优于 Born Again Networks。
• LIT 通过在残差模块中匹配 IR，实现了 GAN 生成器的压缩，是首个适用于 GAN 的学生/教师方法。
• 使用平滑 L1 损失进行 IR 匹配会降低准确率，而 L1 和 L2 损失在性能上无显著差异。
• 混合精度训练对 LIT 性能影响极小，ResNeXt 上仅导致 0.06% 的准确率下降，证实其与现代训练实践的兼容性。