[论文解读] Model compression for faster structural separation of macromolecules captured by Cellular Electron Cryo-Tomography

Jialiang Guo, Bo Zhou+3 more|arXiv (Cornell University)|Jan 31, 2018

Computational Physics and Python Applications参考文献 8被引用 3

一句话总结

本文提出基于知识蒸馏的模型压缩方法，以加速细胞电子冷冻断层成像（CECT）中的大分子分类。通过在深度教师网络（DSRF3D-v2）生成的软标签上训练紧凑的学生网络，该方法将模型大小减少最多113倍，推理时间减少2.85倍，同时保持原始准确率的93.8%，实现了在大规模CECT数据集上的高效部署。

ABSTRACT

Electron Cryo-Tomography (ECT) enables 3D visualization of macromolecule structure inside single cells. Macromolecule classification approaches based on convolutional neural networks (CNN) were developed to separate millions of macromolecules captured from ECT systematically. However, given the fast accumulation of ECT data, it will soon become necessary to use CNN models to efficiently and accurately separate substantially more macromolecules at the prediction stage, which requires additional computational costs. To speed up the prediction, we compress classification models into compact neural networks with little in accuracy for deployment. Specifically, we propose to perform model compression through knowledge distillation. Firstly, a complex teacher network is trained to generate soft labels with better classification feasibility followed by training of customized student networks with simple architectures using the soft label to compress model complexity. Our tests demonstrate that our compressed models significantly reduce the number of parameters and time cost while maintaining similar classification accuracy.

研究动机与目标

• 解决处理高通量CECT数据中数百万个亚断层图像带来的日益严重的计算瓶颈问题。
• 降低用于大分子分类的大型高精度3D-CNN模型所带来的高昂推理成本。
• 在计算资源有限的条件下，实现在大规模CECT数据集上对深度学习模型的高效部署。
• 在显著压缩模型的同时保持高分类准确率，以满足真实世界生物应用的需求。

提出的方法

• 在模拟的CECT亚断层图像上训练复杂的DSRF3D-v2教师网络，以生成高质量的软标签。
• 设计三个逐步缩小的学生成型网络（DSRF3D-v2-s1、s2、s3），采用简化架构。
• 使用知识蒸馏训练学生网络，其中教师网络的软标签指导学习过程。
• 采用温度缩放的Softmax函数（T=5）以增强教师到学生的知识迁移效果。
• 使用随机梯度下降优化学生模型，并采用带有Nesterov动量的分类交叉熵损失函数。
• 基于验证损失实施早停策略，以防止过拟合并确保模型泛化能力。

实验结果

研究问题

• RQ1 知识蒸馏能否在几乎不损失准确率的情况下，有效压缩用于CECT亚断层图像分类的大型3D-CNN模型？
• RQ2 在3D生物图像分析中，模型压缩率、推理速度与分类准确率之间的权衡关系如何？
• RQ3 与仅使用硬标签训练相比，从教师网络进行知识蒸馏如何提升小型学生网络的性能？
• RQ4 在该场景下，哪种学生网络架构在压缩率、速度与准确率之间实现了最佳平衡？
• RQ5 模型压缩在多大程度上能够实现深度学习模型在大规模CECT数据集上的实际部署？

主要发现

• DSRF3D-v2-s1学生模型实现了18.1倍的压缩率，将参数从1830万减少至100万，同时保持了教师模型93.8%的准确率。
• DSRF3D-v2-s1模型将每百万亚断层图像的推理时间从15.91秒减少至9.41秒，实现了1.69倍的速度提升。
• 与仅使用硬标签训练相比，使用教师网络生成的软标签训练学生网络，准确率最高提升了8.16个百分点。
• DSRF3D-v2-s3模型实现了最高的压缩率（113.3倍），但准确率下降最大，仅保留了教师模型89.83%的准确率。
• 知识蒸馏使学生模型的泛化能力优于仅使用硬标签训练的模型，证明了软标签指导的价值。
• 结果表明，通过蒸馏实现的模型压缩对CECT中的3D-CNN极为有效，可实现对大规模生物数据集的高效处理。