Disentangling Masked Autoencoders for Unsupervised Domain Generalization

📄 arXiv: 2407.07544v1 📥 PDF

作者: An Zhang, Han Wang, Xiang Wang, Tat-Seng Chua

分类: cs.CV, cs.AI

发布日期: 2024-07-10

💡 一句话要点

提出DisMAE，通过解耦表征学习实现无监督领域泛化

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture)

关键词: 无监督领域泛化 解耦表征学习 掩码自编码器 领域对抗训练 领域不变特征

📋 核心要点

  1. 现有领域泛化方法依赖大量标注数据，但在无监督场景下表现受限，难以有效应对领域偏移。
  2. DisMAE通过解耦领域不变的语义特征和领域特定的变异特征，学习更鲁棒的表征，提升泛化能力。
  3. 实验表明，DisMAE在多个数据集上取得了与SOTA方法相当甚至更优的性能，验证了解耦表征的有效性。

📝 摘要（中文）

领域泛化(DG)旨在利用充分的监督信号学习领域不变性，从而增强分布外(OOD)泛化能力。然而，由于标注数据的稀缺，无监督领域泛化(UDG)应运而生，这是一个更重要但更具挑战性的任务，模型在不同的领域以无监督的方式进行训练，并最终在未见过的领域进行测试。UDG正迅速获得关注，但仍远未得到充分研究。为了弥合这一研究差距，我们提出了一种专为UDG设计的新型学习框架，称为解耦掩码自编码器(DisMAE)，旨在发现解耦的表征，忠实地揭示内在特征和表面变化，而无需访问类标签。其核心是提炼领域不变的语义特征（领域分类器无法区分），同时过滤掉不稳定和冗余的领域特定变化（例如，配色方案和纹理模式）。值得注意的是，DisMAE使用语义和轻量级变异编码器共同训练非对称双分支架构，提供动态数据操作和表征级别的增强能力。在四个基准数据集（即DomainNet、PACS、VLCS、Colored MNIST）上进行的DG和UDG任务的广泛实验表明，与最先进的DG和UDG基线相比，DisMAE可以实现有竞争力的OOD性能，这为利用大规模未标记数据提高泛化能力的潜在研究方向提供了启示。

🔬 方法详解

问题定义：论文旨在解决无监督领域泛化(UDG)问题。现有方法在有监督领域泛化(DG)中表现良好，但在缺乏标签的UDG场景下，难以有效学习领域不变特征，导致泛化能力下降。痛点在于如何从无标签数据中提取鲁棒的、与领域无关的特征表示。

核心思路：DisMAE的核心思路是将图像表征解耦为领域不变的语义特征和领域特定的变异特征。通过这种解耦，模型可以专注于学习更本质的、与领域无关的特征，从而提高在未见领域上的泛化能力。设计思路是利用掩码自编码器(MAE)的重建能力，结合领域对抗训练的思想，迫使模型学习到对领域变化不敏感的特征。

技术框架：DisMAE采用非对称双分支架构，包含一个语义编码器和一个变异编码器。语义编码器负责提取领域不变的语义特征，变异编码器负责提取领域特定的变异特征。两个编码器共同训练一个掩码自编码器，该自编码器需要从掩码后的图像中重建原始图像。此外，还引入了一个领域分类器，用于区分不同领域，并利用对抗训练的思想，使得语义编码器提取的特征难以被领域分类器区分。

关键创新：DisMAE的关键创新在于解耦表征学习框架，它能够有效地分离领域不变的语义特征和领域特定的变异特征。与传统的领域对抗训练方法不同，DisMAE不是直接最小化领域分类器的损失，而是通过掩码自编码器的重建任务和变异编码器的辅助，间接地学习领域不变特征。这种方法更加稳定和有效。

关键设计：DisMAE的关键设计包括：1) 非对称双分支架构，允许模型分别学习语义和变异特征；2) 掩码自编码器的重建任务，迫使模型学习图像的本质特征；3) 领域对抗训练，使得语义特征对领域变化不敏感；4) 轻量级的变异编码器，避免模型过度拟合领域特定的噪声。具体的损失函数包括重建损失、领域对抗损失和正则化损失。参数设置方面，掩码比例、学习率、batch size等都需要根据具体数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DisMAE在DomainNet、PACS、VLCS和Colored MNIST四个基准数据集上进行了广泛的实验，结果表明，DisMAE在DG和UDG任务上均取得了具有竞争力的性能。例如，在DomainNet数据集上，DisMAE的OOD性能优于多个SOTA的DG和UDG基线方法。实验结果验证了DisMAE解耦表征学习框架的有效性，并表明其在大规模未标记数据上具有良好的泛化能力。

🎯 应用场景

DisMAE在图像识别、目标检测、语义分割等领域具有广泛的应用前景，尤其是在训练数据缺乏标注或存在领域偏移的情况下。例如，在医疗图像分析中，不同医院的扫描设备和协议可能导致图像存在领域差异，DisMAE可以用于提高模型在不同医院数据上的泛化能力。此外，DisMAE还可以应用于自动驾驶、机器人等领域，提高模型在不同环境下的鲁棒性。

📄 摘要（原文）

Domain Generalization (DG), designed to enhance out-of-distribution (OOD) generalization, is all about learning invariance against domain shifts utilizing sufficient supervision signals. Yet, the scarcity of such labeled data has led to the rise of unsupervised domain generalization (UDG) - a more important yet challenging task in that models are trained across diverse domains in an unsupervised manner and eventually tested on unseen domains. UDG is fast gaining attention but is still far from well-studied. To close the research gap, we propose a novel learning framework designed for UDG, termed the Disentangled Masked Auto Encoder (DisMAE), aiming to discover the disentangled representations that faithfully reveal the intrinsic features and superficial variations without access to the class label. At its core is the distillation of domain-invariant semantic features, which cannot be distinguished by domain classifier, while filtering out the domain-specific variations (for example, color schemes and texture patterns) that are unstable and redundant. Notably, DisMAE co-trains the asymmetric dual-branch architecture with semantic and lightweight variation encoders, offering dynamic data manipulation and representation level augmentation capabilities. Extensive experiments on four benchmark datasets (i.e., DomainNet, PACS, VLCS, Colored MNIST) with both DG and UDG tasks demonstrate that DisMAE can achieve competitive OOD performance compared with the state-of-the-art DG and UDG baselines, which shed light on potential research line in improving the generalization ability with large-scale unlabeled data.