[论文解读] On Compressing U-net Using Knowledge Distillation

Karttikeya Mangalam, Mathieu Salzamann|arXiv (Cornell University)|Dec 1, 2018

Energy Efficient Wireless Sensor Networks参考文献 11被引用 7

一句话总结

本文提出一种知识蒸馏方法，将U-Net架构的参数量压缩超过1000倍，参数量减少至原始大小的0.1%，同时保持近乎相同的分割精度。通过在蒸馏过程中将批量归一化和类别重加权集成到学生网络中，该方法克服了训练不稳定性和类别不平衡问题，实现了标准蒸馏失败时仍能可靠压缩模型。

ABSTRACT

We study the use of knowledge distillation to compress the U-net architecture. We show that, while standard distillation is not sufficient to reliably train a compressed U-net, introducing other regularization methods, such as batch normalization and class re-weighting, in knowledge distillation significantly improves the training process. This allows us to compress a U-net by over 1000x, i.e., to 0.1% of its original number of parameters, at a negligible decrease in performance.

研究动机与目标

• 为解决在移动设备等资源受限环境部署大型U-Net模型的挑战。
• 探究知识蒸馏是否能有效将U-Net架构压缩至极小尺寸。
• 识别并解决标准蒸馏在压缩小型U-Net变体时的失败原因。
• 通过引入架构和损失函数的改进，提升小型学生网络蒸馏的可靠性。

提出的方法

• 在学生U-Net的编码路径中，每个卷积层均引入批量归一化层，以减少内部协变量偏移并稳定训练过程。
• 在交叉熵损失中应用类别重加权，基于背景与前景像素比例设定前景类别权重为17.8，以缓解类别不平衡问题。
• 采用混合知识蒸馏，利用来自4-Unet教师网络的温度缩放软标签损失，结合T=5时的软标签与硬标签。
• 使用蒸馏损失（软标签）和原始交叉熵损失（硬标签）联合训练学生网络，最后对输出进行温度软化处理。
• 以从头训练的4-Unet作为教师模型（使用硬标签），并将其蒸馏至参数极少的2-Unet中。
• 在训练过程中将软标签损失按T²缩放，以保持不同温度下的梯度幅值一致。

实验结果

研究问题

• RQ1 标准知识蒸馏能否可靠地将U-Net压缩至原始参数量的1%以下？
• RQ2 为何标准蒸馏在将U-Net压缩至极小尺寸（如初始通道数为2）时会失败？
• RQ3 批量归一化和类别重加权如何改善压缩U-Net学生网络的训练？
• RQ4 在结合正则化技术增强的蒸馏训练下，2-Unet能否实现与64-Unet相当的分割性能？

主要发现

• 仅含122,394个参数的4-Unet在测试集上损失为0.0974，IoU为0.807，与原始64-Unet的0.804 IoU相当，表明不使用蒸馏的参数压缩同样有效。
• 标准蒸馏（原始或混合）无法训练2-Unet，如表2所示，测试损失达0.505–0.507，表明泛化能力差且训练不稳定。
• 引入批量归一化和类别重加权后，2-Unet在混合蒸馏下实现IoU 0.759和交叉熵损失0.135，优于未使用蒸馏的2-Unet（IoU 0.752，损失0.134）。
• 最终的2-Unet模型仅使用30,902个参数，占原始64-Unet（3100万参数）的0.1%，同时保持近乎相同的性能。
• 若不使用蒸馏但采用改进方法训练2-Unet，测试损失为0.307，证实蒸馏对性能提升仍至关重要。
• 批量归一化与类别重加权的结合使低参数量下的蒸馏成功成为可能，而标准蒸馏在此类场景下会失败。