Simple Unsupervised Knowledge Distillation With Space Similarity

📄 arXiv: 2409.13939v1 📥 PDF

作者: Aditya Singh, Haohan Wang

分类: cs.AI, cs.CV

发布日期: 2024-09-20

💡 一句话要点

提出基于空间相似性的无监督知识蒸馏方法，提升小模型性能

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

关键词: 无监督学习 知识蒸馏 模型压缩 嵌入空间 空间相似性

📋 核心要点

  1. 自监督学习难以直接应用于小型网络架构，无监督知识蒸馏是解决此问题的有效途径。
  2. 该方法的核心思想是让学生网络模仿教师网络的嵌入流形，而非手工设计样本关系。
  3. 提出的“空间相似性”损失函数，促使学生网络学习教师网络特征空间的维度信息，实验证明有效。

📝 摘要（中文）

本文提出一种简单的无监督知识蒸馏(UKD)方法，旨在解决自监督学习(SSL)难以直接应用于小型架构的问题。现有UKD方法通常手工设计教师网络和学生网络之间具有代表性的样本间/内关系，但可能忽略教师网络映射中的其他关键关系。本文不采用启发式构建关系的方法，而是直接促使学生网络学习教师网络的嵌入流形。如果映射的流形相似，则所有样本间/内关系都会被间接保留。研究表明，现有方法由于过度依赖$L_2$归一化嵌入特征，无法保留教师网络的潜在流形。因此，本文提出一种名为“空间相似性”的损失分量，旨在捕捉归一化造成的损失信息，促使学生网络特征空间的每个维度与教师网络对应维度相似。大量实验表明，该方法在多个基准测试中表现出色。

🔬 方法详解

问题定义：现有无监督知识蒸馏方法依赖于手工设计的样本间或样本内关系，来指导学生网络的学习。这种方法的缺点在于，人工设计的关系可能无法完全捕捉教师网络所学习到的所有重要信息，特别是嵌入空间中的细微结构和维度关系。此外，现有方法过度依赖L2归一化，导致信息损失。

核心思路：本文的核心思路是让学生网络直接学习教师网络的嵌入流形，而不是仅仅关注样本之间的关系。如果学生网络能够成功地模仿教师网络的嵌入流形，那么样本之间的关系自然也会被保留。这种方法避免了人工设计关系的局限性，能够更全面地学习教师网络的信息。

技术框架：该方法主要包含两个网络：教师网络和学生网络。教师网络可以是预训练好的模型，也可以是更大的网络。学生网络是需要训练的小型网络。训练过程中，首先将无标签数据输入到教师网络和学生网络中，得到它们的嵌入表示。然后，使用提出的“空间相似性”损失函数来衡量两个嵌入表示之间的相似度，并以此来指导学生网络的学习。

关键创新：该方法最重要的创新点在于提出了“空间相似性”损失函数。与传统的基于样本关系的损失函数不同，“空间相似性”损失函数直接比较教师网络和学生网络嵌入空间的每个维度，促使学生网络学习教师网络特征空间的维度信息。这种方法能够更全面地学习教师网络的信息，避免了人工设计关系的局限性。

关键设计：空间相似性损失函数的设计是关键。具体来说，对于教师网络和学生网络的嵌入表示，该损失函数计算它们对应维度之间的相似度，例如可以使用余弦相似度或L1距离。然后，将所有维度的相似度加权平均，得到最终的损失值。此外，论文还探讨了不同的权重分配策略，以进一步提高性能。网络结构方面，可以使用常见的卷积神经网络或Transformer网络。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，提出的方法在多个图像分类数据集上取得了显著的性能提升。例如，在CIFAR-10数据集上，使用ResNet-18作为教师网络，MobileNetV2作为学生网络，该方法可以将学生网络的准确率提升到与教师网络相近的水平，甚至超过一些传统的知识蒸馏方法。

🎯 应用场景

该研究成果可应用于各种需要模型压缩和加速的场景，例如移动设备上的图像识别、边缘计算等。通过无监督知识蒸馏，可以将大型模型的知识迁移到小型模型上，从而在资源受限的环境中实现高性能的推理。

📄 摘要（原文）

As per recent studies, Self-supervised learning (SSL) does not readily extend to smaller architectures. One direction to mitigate this shortcoming while simultaneously training a smaller network without labels is to adopt unsupervised knowledge distillation (UKD). Existing UKD approaches handcraft preservation worthy inter/intra sample relationships between the teacher and its student. However, this may overlook/ignore other key relationships present in the mapping of a teacher. In this paper, instead of heuristically constructing preservation worthy relationships between samples, we directly motivate the student to model the teacher's embedding manifold. If the mapped manifold is similar, all inter/intra sample relationships are indirectly conserved. We first demonstrate that prior methods cannot preserve teacher's latent manifold due to their sole reliance on $L_2$ normalised embedding features. Subsequently, we propose a simple objective to capture the lost information due to normalisation. Our proposed loss component, termed \textbf{space similarity}, motivates each dimension of a student's feature space to be similar to the corresponding dimension of its teacher. We perform extensive experiments demonstrating strong performance of our proposed approach on various benchmarks.