[论文解读] Task-customized Masked AutoEncoder via Mixture of Cluster-conditional Experts

Zhili Liu, Kai Chen+5 more|arXiv (Cornell University)|Feb 8, 2024

Advanced Clustering Algorithms Research被引用 6

一句话总结

简述：提出 MoCE，一种基于聚类引导的专家混合扩展到 MAE 的任务定制自监督预训练，在11个下游任务上实现更好的迁移，并在检测/分割方面达到最先进的结果。

ABSTRACT

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.

研究动机与目标

• 在 transferring 到语义差异较大的下游任务时，演示标准 MAE 的负迁移现象。
• 识别在自监督设置中直接将 TokenMoE 应用于 MAE 的局限性。
• 引入 MoCE 以在语义相似的图像簇上训练专家。
• 证明 MoCE 在多个下游任务与部署中的迁移性能优于基线。
• 就部署任务定制的 MoCE 模型以实现高效推断提供指南。

提出的方法

• 使用来自预训练 MAE 的特征对预训练数据进行聚类，形成语义连贯的组。
• 用 MoCE 门替换选定的 MLP 层，将图像令牌路由到簇条件专家。
• 在其分配簇中的数据上对每个专家进行训练，并使用蒸馏损失来稳定学习。
• 在部署时，利用聚类模块选择与下游任务最相关的专家。
• 引入一种不平衡/正则化损失以鼓励自信的门控路由。

实验结果

研究问题

• RQ1 MAE 是否会在语义上多样的下游任务上产生负迁移？
• RQ2 将 TokenMoE 直接应用于 MAE 是否会提升或降低迁移性能？
• RQ3 MoCE 是否通过簇条件路由实现比原生 MAE 和 TokenMoE 更好的任务特异性预训练？
• RQ4 MoCE 在包含检测与分割等广泛下游任务上的表现如何？

主要发现

• MoCE 在11个下游任务上相较 MAE* 的 Top-1 准确率平均提升可达 2.45%。
• MoCE 在检测与分割基准上实现了自监督的最先进结果。
• TokenMoE 对 MAE 没有改善，且因语义不可知路由可能表现不佳。
• MoCE 的簇条件专家学得语义连贯的专门化（如服装、鸟类、狗类）。
• 在测试阶段若仅有一个激活的 MoCE 专家，部署时参数量可减半并提升效率。