[论文解读] Subclass Distillation

Rafael Rios Müller, Simon Kornblith+1 more|arXiv (Cornell University)|Feb 10, 2020

Process Optimization and Integration参考文献 17被引用 8

一句话总结

本文提出子类蒸馏（subclass distillation），一种通过训练教师网络将每个类别拆分为多个学习到的子类来增强低类别数据集知识蒸馏的方法。通过使用对比性辅助损失从这些子类概率中蒸馏知识，学生模型即使在未直接监督二元标签的情况下，也能实现更好的泛化能力和更快的收敛速度。

ABSTRACT

After a large "teacher" neural network has been trained on labeled data, the probabilities that the teacher assigns to incorrect classes reveal a lot of information about the way in which the teacher generalizes. By training a small "student" model to match these probabilities, it is possible to transfer most of the generalization ability of the teacher to the student, often producing a much better small model than directly training the student on the training data. The transfer works best when there are many possible classes because more is then revealed about the function learned by the teacher, but in cases where there are only a few possible classes we show that we can improve the transfer by forcing the teacher to divide each class into many subclasses that it invents during the supervised training. The student is then trained to match the subclass probabilities. For datasets where there are known, natural subclasses we demonstrate that the teacher learns similar subclasses and these improve distillation. For clickthrough datasets where the subclasses are unknown we demonstrate that subclass distillation allows the student to learn faster and better.

研究动机与目标

• 解决传统知识蒸馏在低类别场景（如二分类）中的低效问题，其中有限的类别信息会降低知识迁移的有效性。
• 通过利用数据中隐藏的子类结构，提升资源受限场景下的模型泛化能力和训练效率。
• 实现在仅提供二元标签而无需事先知晓子类结构的情况下，教师网络仍能学习语义上有意义的子类，并实现知识蒸馏。
• 证明学生可以通过模仿教师的类内子类概率分布，间接学习二分类任务，而无需对目标类别进行显式监督。

提出的方法

• 训练一个深层教师网络，为每个原始类别生成 s 个子类的 logits，最终类别概率通过求和子类概率得到。
• 引入一个辅助对比损失，通过在动量更新的队列中将每个样本与其自身进行对比，以增强子类的多样性与置信度。
• 对子类 logits 的 softmax 使用温度缩放，以强调子类之间的相对差异，同时保持类内关系。
• 使用两个交叉熵损失训练学生网络：一个基于原始类别标签，另一个基于匹配教师的子类概率分布。
• 使用动量更新机制对教师的动量编码器进行更新，以在训练过程中稳定子类表征的学习。
• 在已知（CIFAR-2x5, MNIST-2x5）和未知（CelebA, Criteo）子类结构的数据集上端到端应用该方法，以评估泛化能力和效率。

实验结果

研究问题

• RQ1 当仅提供二元标签时，教师网络是否能在二分类任务中发现语义上有意义的子类？
• RQ2 从子类概率中蒸馏是否能带来比传统蒸馏或倒数第二层特征蒸馏更好的泛化性能？
• RQ3 学生是否能通过模仿教师的类内子类概率分布，在未直接监督目标类别的情况下学习二分类任务？
• RQ4 当数据有限或子类结构未知时，子类蒸馏如何影响训练速度和性能？
• RQ5 所学习的子类在多大程度上与真实世界数据集中已知或自然存在的数据结构相关联？

主要发现

• 在 CIFAR-2x5 上，教师通过子类学习成功恢复了原始 CIFAR-10 类别，子类蒸馏达到 1.83% 的错误率，优于传统蒸馏和倒数第二层蒸馏。
• 在 CelebA 上，教师学习到的子类与从未训练过的其他二元标签表现出强相关性，表明其语义意义明确，并提升了蒸馏性能。
• 在 Criteo 点击预测数据集上，子类蒸馏实现了更快的收敛速度和更好的性能，尤其当学生对完整数据集访问受限时更为显著。
• 在 MNIST-2x5 上，学生仅通过学习相对子类概率，未见过任何二元标签，即实现了 2.06% ± 0.18 的错误率，证明了对目标任务的间接学习能力。
• 辅助对比损失成功促使子类预测更加多样化且更具置信度，且子类结构在多个数据集中均表现出语义一致性。
• 在低数据场景下，子类蒸馏比标准蒸馏传递了更多的信息量（bits per example），表明其信息迁移效率更高。