Unsqueeze [CLS] Bottleneck to Learn Rich Representations

📄 arXiv: 2407.17671v2 📥 PDF

作者: Qing Su, Shihao Ji

分类: cs.CV, cs.LG

发布日期: 2024-07-24 (更新: 2024-07-26)

备注: ECCV 2024

🔗 代码/项目: GITHUB

💡 一句话要点

提出UDI：一种解压缩[CLS]瓶颈的自监督学习方法，提升表征能力。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自监督学习 蒸馏学习 表征学习 多模态预测 分层抽样 目标检测 图像分割

📋 核心要点

  1. 传统蒸馏自监督学习方法易导致表征压缩，损失输入信息，限制了模型性能。
  2. UDI通过解压缩[CLS]瓶颈，利用分层抽样和多模态预测，丰富学习到的表征。
  3. 实验表明，UDI在图像分类、目标检测、分割和小样本学习任务上均有提升。

📝 摘要（中文）

基于蒸馏的自监督学习通常由于其激进的聚类过程和更锐利的目标分布而导致更压缩的表征。为了克服这个限制并保留来自输入的更多信息，我们引入了UDI，概念化为基于解压缩蒸馏的自监督学习（SSL）。UDI通过鼓励从分层抽样获得的局部预测的整合概况中提取的多模态预测来丰富学习到的表征。我们的评估表明，UDI不仅在实例级别促进了语义上有意义的表征，在图像分类中提供了优于或可与最先进的SSL方法相媲美的结果，而且有效地保留了输入的干扰，从而在包括目标检测和分割在内的密集预测任务中产生了显著的改进。此外，UDI在小样本图像分类中表现出竞争力，提高了联合嵌入管道的可扩展性。我们提供了各种可视化和消融研究，以进一步阐明UDI背后的机制。我们的源代码可在https://github.com/ISL-CV/udi获得。

🔬 方法详解

问题定义：现有的基于蒸馏的自监督学习方法，由于其激进的聚类过程和更锐利的目标分布，往往会产生过于压缩的表征。这种压缩会导致输入信息的丢失，从而限制了模型在下游任务中的性能，尤其是在需要保留输入细节的任务中，例如目标检测和分割。

核心思路：UDI的核心思路是通过“解压缩”[CLS]瓶颈来保留更多来自输入的信息。具体来说，UDI不再直接预测一个单一的类别标签，而是鼓励模型预测一个多模态的分布，该分布是从多个局部预测的整合概况中蒸馏得到的。这种方式可以保留更多关于输入数据的细节和不确定性，从而产生更丰富的表征。

技术框架：UDI的整体框架包括以下几个主要步骤：1) 输入图像经过数据增强后，输入到编码器网络中；2) 编码器输出的[CLS] token经过分层抽样，得到多个局部表征；3) 对每个局部表征进行预测，得到多个局部预测；4) 将这些局部预测整合为一个概况；5) 使用蒸馏损失，鼓励模型预测与该概况相似的多模态分布。

关键创新：UDI的关键创新在于其解压缩[CLS]瓶颈的思想，以及通过分层抽样和多模态预测来实现这一思想的具体方法。与传统的蒸馏方法不同，UDI不是简单地将知识从一个模型转移到另一个模型，而是通过鼓励模型预测一个更丰富的分布来保留更多信息。

关键设计：UDI的关键设计包括：1) 分层抽样策略，用于生成多个具有代表性的局部表征；2) 多模态预测，用于保留输入数据的细节和不确定性；3) 蒸馏损失函数，用于鼓励模型预测与整合概况相似的多模态分布。具体的损失函数选择和参数设置需要在实验中进行调整，以达到最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

UDI在图像分类任务上取得了与SOTA方法相当或更优的结果。更重要的是，UDI在目标检测和分割等密集预测任务上取得了显著的提升，例如在COCO数据集上，使用Mask R-CNN作为下游任务时，UDI相比于其他自监督学习方法，AP指标提升了多个百分点。此外，UDI在小样本图像分类任务上也表现出竞争力，提高了联合嵌入管道的可扩展性。

🎯 应用场景

UDI具有广泛的应用前景，可用于提升各种计算机视觉任务的性能，尤其是在需要保留输入细节的任务中，例如目标检测、分割和图像生成。此外，UDI在小样本学习中的良好表现也使其适用于数据稀缺的场景。该方法还可应用于其他自监督学习领域，例如自然语言处理。

📄 摘要（原文）

Distillation-based self-supervised learning typically leads to more compressed representations due to its radical clustering process and the implementation of a sharper target distribution. To overcome this limitation and preserve more information from input, we introduce UDI, conceptualized as Unsqueezed Distillation-based self-supervised learning (SSL). UDI enriches the learned representation by encouraging multimodal prediction distilled from a consolidated profile of local predictions that are derived via stratified sampling. Our evaluations show that UDI not only promotes semantically meaningful representations at instance level, delivering superior or competitive results to state-of-the-art SSL methods in image classification, but also effectively preserves the nuisance of input, which yields significant improvement in dense prediction tasks, including object detection and segmentation. Additionally, UDI performs competitively in low-shot image classification, improving the scalability of joint-embedding pipelines. Various visualizations and ablation studies are presented to further elucidate the mechanisms behind UDI. Our source code is available at https://github.com/ISL-CV/udi.