EPSD: Early Pruning with Self-Distillation for Efficient Model Compression

📄 arXiv: 2402.00084v1 📥 PDF

作者: Dong Chen, Ning Liu, Yichen Zhu, Zhengping Che, Rui Ma, Fachao Zhang, Xiaofeng Mou, Yi Chang, Jian Tang

分类: cs.LG, cs.AI, cs.CV

发布日期: 2024-01-31

备注: The first two authors are with equal contributions. Paper accepted by AAAI 2024

💡 一句话要点

提出EPSD以实现高效模型压缩

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 模型压缩 知识蒸馏 网络剪枝 自蒸馏 深度学习 高效训练 计算机视觉

📋 核心要点

  1. 现有的剪枝与知识蒸馏方法效率低下，需大量预训练和复杂步骤。
  2. 提出EPSD框架，通过早期剪枝与自蒸馏结合，保留可蒸馏权重以提高训练效果。
  3. EPSD在CIFAR-10/100、Tiny-ImageNet等多个数据集上表现优异，超越了现有技术。

📝 摘要（中文）

神经网络压缩技术如知识蒸馏(KD)和网络剪枝受到越来越多的关注。近期研究表明，经过剪枝的教师网络可以提升KD的性能。然而，传统的教师-学生流程需要繁琐的预训练和复杂的压缩步骤，使得KD与剪枝的结合效率较低。为了解决这一问题，本文提出了早期剪枝与自蒸馏(EPSD)框架，旨在高效地结合早期剪枝和自蒸馏，通过识别和保留可蒸馏权重来提高剪枝网络的训练效果。实验结果表明，EPSD在多个基准数据集上超越了先进的剪枝和蒸馏技术。

🔬 方法详解

问题定义：本文旨在解决传统剪枝与知识蒸馏结合效率低的问题，现有方法需要大量的预训练和复杂的步骤，导致计算成本高。

核心思路：提出EPSD框架，通过早期剪枝与自蒸馏相结合，识别并保留可蒸馏的权重，从而提高剪枝网络的训练效果，确保更好的蒸馏性能。

技术框架：EPSD的整体流程分为两个主要阶段：第一阶段进行早期剪枝，识别可蒸馏权重；第二阶段进行自蒸馏训练，优化剪枝后的网络。

关键创新：EPSD的创新在于通过早期剪枝与自蒸馏的有效结合，保持更多的可蒸馏权重，从而提升剪枝网络的训练效率，与传统方法相比显著降低了计算成本。

关键设计：在设计中，EPSD采用了特定的损失函数以优化蒸馏过程，并在剪枝过程中设置了权重保留策略，以确保剪枝网络的可训练性和蒸馏效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

EPSD在CIFAR-10/100、Tiny-ImageNet等多个基准数据集上表现优异，显著超越了现有的剪枝和自蒸馏技术，展示了更高的训练效率和更好的模型性能，具体提升幅度未知。

🎯 应用场景

该研究在深度学习模型压缩领域具有重要应用潜力，尤其适用于资源受限的设备和实时系统。通过提高模型的压缩效率，EPSD可广泛应用于移动设备、边缘计算和嵌入式系统等场景，推动智能设备的普及与发展。

📄 摘要（原文）

Neural network compression techniques, such as knowledge distillation (KD) and network pruning, have received increasing attention. Recent work `Prune, then Distill' reveals that a pruned student-friendly teacher network can benefit the performance of KD. However, the conventional teacher-student pipeline, which entails cumbersome pre-training of the teacher and complicated compression steps, makes pruning with KD less efficient. In addition to compressing models, recent compression techniques also emphasize the aspect of efficiency. Early pruning demands significantly less computational cost in comparison to the conventional pruning methods as it does not require a large pre-trained model. Likewise, a special case of KD, known as self-distillation (SD), is more efficient since it requires no pre-training or student-teacher pair selection. This inspires us to collaborate early pruning with SD for efficient model compression. In this work, we propose the framework named Early Pruning with Self-Distillation (EPSD), which identifies and preserves distillable weights in early pruning for a given SD task. EPSD efficiently combines early pruning and self-distillation in a two-step process, maintaining the pruned network's trainability for compression. Instead of a simple combination of pruning and SD, EPSD enables the pruned network to favor SD by keeping more distillable weights before training to ensure better distillation of the pruned network. We demonstrated that EPSD improves the training of pruned networks, supported by visual and quantitative analyses. Our evaluation covered diverse benchmarks (CIFAR-10/100, Tiny-ImageNet, full ImageNet, CUB-200-2011, and Pascal VOC), with EPSD outperforming advanced pruning and SD techniques.