Task-customized Masked AutoEncoder via Mixture of Cluster-conditional Experts

📄 arXiv: 2402.05382v1 📥 PDF

作者: Zhili Liu, Kai Chen, Jianhua Han, Lanqing Hong, Hang Xu, Zhenguo Li, James T. Kwok

分类: cs.CV, cs.LG

发布日期: 2024-02-08

备注: Accepted by ICLR 2023

💡 一句话要点

提出基于混合聚类条件专家的MAE以解决下游任务定制问题

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

关键词: 自监督学习 Masked Autoencoder 混合专家 聚类条件门 模型定制化 计算机视觉 目标检测 图像分割

📋 核心要点

  1. 现有的MAE方法在面对不同数据分布的下游任务时，可能会因为负迁移而影响性能。
  2. 本文提出的MoCE方法通过聚类条件门为每个专家分配语义相关的图像，从而实现定制化的预训练。
  3. 在11个下游任务的实验中，MoCE平均性能提升2.45%，并在检测和分割任务上取得了新的自监督学习的最优结果。

📝 摘要（中文）

Masked Autoencoder（MAE）是一种流行的自监督学习方法，在模型预训练中取得了良好的效果。然而，当下游任务的数据分布与预训练数据不同，语义上不相关的预训练信息可能导致负迁移，阻碍MAE的可扩展性。为了解决这一问题，本文提出了一种新颖的基于MAE的预训练范式——混合聚类条件专家（MoCE），该方法可以一次训练但为多样化的下游任务提供定制的预训练模型。与混合专家（MoE）不同，MoCE通过使用聚类条件门仅用语义相关的图像训练每个专家，从而使每个下游任务能够分配到其定制的模型，该模型是用与下游数据最相似的数据预训练的。在11个下游任务的实验中，MoCE的平均性能比传统MAE提高了2.45%。

🔬 方法详解

问题定义：本文旨在解决MAE在下游任务中因数据分布差异导致的负迁移问题。现有方法在不同任务中可能会引入不相关的预训练信息，影响模型性能。

核心思路：提出混合聚类条件专家（MoCE）方法，通过聚类条件门为每个专家分配语义相关的图像，从而实现针对不同下游任务的定制化预训练。这样的设计可以有效避免负迁移，提升模型的适应性。

技术框架：MoCE的整体架构包括多个专家，每个专家专注于特定的任务和数据集。通过聚类条件门，模型能够选择与当前任务最相关的图像进行训练，确保每个专家的训练数据具有高度的语义一致性。

关键创新：MoCE的核心创新在于使用聚类条件门来选择训练数据，这与传统的混合专家方法不同，后者通常使用固定的专家组合。MoCE能够根据任务需求动态调整专家的训练数据，从而实现更高的灵活性和性能。

关键设计：在模型设计中，聚类条件门的实现是关键，确保每个专家仅接收与其任务相关的图像。此外，损失函数的设计也考虑了不同任务的特性，以优化模型在特定任务上的表现。整体网络结构经过精心设计，以支持多任务学习和专家的动态选择。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在11个下游任务的实验中，MoCE的平均性能提升了2.45%，显著优于传统的MAE方法。此外，MoCE在检测和分割任务上也取得了新的自监督学习的最优结果，展示了其在实际应用中的有效性和优势。

🎯 应用场景

该研究的潜在应用领域包括计算机视觉中的图像分类、目标检测和图像分割等任务。通过提供定制化的预训练模型，MoCE能够在多样化的下游任务中提升模型性能，具有广泛的实际价值。未来，该方法可能推动自监督学习在更多领域的应用，尤其是在数据稀缺的场景中。

📄 摘要（原文）

Masked Autoencoder~(MAE) is a prevailing self-supervised learning method that achieves promising results in model pre-training. However, when the various downstream tasks have data distributions different from the pre-training data, the semantically irrelevant pre-training information might result in negative transfer, impeding MAE's scalability. To address this issue, we propose a novel MAE-based pre-training paradigm, Mixture of Cluster-conditional Experts (MoCE), which can be trained once but provides customized pre-training models for diverse downstream tasks. Different from the mixture of experts (MoE), our MoCE trains each expert only with semantically relevant images by using cluster-conditional gates. Thus, each downstream task can be allocated to its customized model pre-trained with data most similar to the downstream data. Experiments on a collection of 11 downstream tasks show that MoCE outperforms the vanilla MAE by 2.45\% on average. It also obtains new state-of-the-art self-supervised learning results on detection and segmentation.