[论文解读] MultiMAE: Multi-modal Multi-task Masked Autoencoders

Roman Bachmann, David Mizrahi+2 more|arXiv (Cornell University)|Apr 4, 2022

Domain Adaptation and Few-Shot Learning被引用 10

一句话总结

MultiMAE 提出了一种多模态、多任务掩码自编码框架，通过伪标签数据在单个视觉 Transformer 上对 RGB 图像、深度图和语义分割进行预训练。通过在多个模态上掩码图像块，并使用特定任务的解码器重建这些块，MultiMAE 学习到跨模态和多任务表示，在仅进行最小架构修改和高效训练的前提下，实现了图像分类、分割和深度估计任务的最先进迁移性能。

ABSTRACT

We propose a pre-training strategy called Multi-modal Multi-task Masked Autoencoders (MultiMAE). It differs from standard Masked Autoencoding in two key aspects: I) it can optionally accept additional modalities of information in the input besides the RGB image (hence "multi-modal"), and II) its training objective accordingly includes predicting multiple outputs besides the RGB image (hence "multi-task"). We make use of masking (across image patches and input modalities) to make training MultiMAE tractable as well as to ensure cross-modality predictive coding is indeed learned by the network. We show this pre-training strategy leads to a flexible, simple, and efficient framework with improved transfer results to downstream tasks. In particular, the same exact pre-trained network can be flexibly used when additional information besides RGB images is available or when no information other than RGB is available - in all configurations yielding competitive to or significantly better results than the baselines. To avoid needing training datasets with multiple modalities and tasks, we train MultiMAE entirely using pseudo labeling, which makes the framework widely applicable to any RGB dataset. The experiments are performed on multiple transfer tasks (image classification, semantic segmentation, depth estimation) and datasets (ImageNet, ADE20K, Taskonomy, Hypersim, NYUv2). The results show an intriguingly impressive capability by the model in cross-modal/task predictive coding and transfer.

研究动机与目标

• 开发一种统一的预训练框架，利用多种模态（RGB、深度图、分割图）和多种任务，以提升下游迁移性能。
• 使同一预训练模型能够灵活微调于单模态和多模态下游任务。
• 仅使用 RGB 图像和伪标签数据进行高效训练，避免对多模态数据集的需求。
• 探究跨模态和多任务预测编码是否能提升表征学习和迁移性能。
• 证明单一预训练的 MultiMAE 模型可在多种下游任务中超越或匹配特定任务基线模型的性能。

提出的方法

• MultiMAE 使用视觉 Transformer 编码器处理从多个输入模态（RGB、深度图、分割图）中随机采样的图像块子集，且在图像块和模态上均应用掩码。
• 模型采用特定任务的解码器，利用可学习查询（包括掩码标记）的交叉注意力机制，重建每个任务的掩码图像块。
• 训练目标仅在掩码图像块上最小化重建损失，促使模型从多模态中稀疏输入图像块预测缺失区域。
• 使用伪标签生成深度图和分割图，基于 ImageNet-1K 上预训练的模型，实现在无需人工标注的情况下进行预训练。
• 框架采用固定掩码策略，使用狄利克雷集中参数（α=1）以在模态和空间位置上均匀采样图像块。
• 模型在 8 张 A100 GPU 上使用混合精度训练进行端到端训练，采用浅层解码器以在输入图像块数量增加的情况下仍保持训练效率。

实验结果

研究问题

• RQ1 是否能够通过联合在多个模态（RGB、深度图、分割图）上预训练的统一掩码自编码框架，提升下游迁移性能？
• RQ2 跨模态预测编码（即通过掩码图像块从一个模态重建另一个模态）是否能带来比单模态预训练更优的表征？
• RQ3 在图像分类、语义分割和深度估计等多样化下游任务中，MultiMAE 的性能与单任务 MAE 基线相比如何？
• RQ4 单一预训练的 MultiMAE 模型是否能在无需微调的情况下，在单模态和多模态下游任务中均取得具有竞争力的结果？
• RQ5 使用伪标签数据为额外模态提供支持，能在多大程度上实现有效的多模态预训练，而无需依赖人工标注的多模态数据集？

主要发现

• 在 ImageNet-1K 图像分类任务中，MultiMAE 达到 83.3% 的 top-1 准确率，与最佳单任务 MAE 水平相当，但具备多模态能力。
• 在 ADE20K 语义分割任务中，MultiMAE 实现 46.2 的 mIoU，优于最佳单任务 MAE（D8）的 0.3 个百分点，展现出强大的泛化能力。
• 在 NYUv2 深度估计任务中，MultiMAE 达到 85.1 的 δ1 准确率，显著优于单任务 MAE（D8）的 1.4 个百分点。
• 在 Taskonomy 数据集上，MultiMAE 在 RGB 到分割图迁移任务中实现最低的 L1 损失（1.040），优于所有基线模型，展现出强大的多任务泛化能力。
• 尽管未掩码图像块数量是 MAE 的两倍（98 vs. 49），但 MultiMAE 的预训练时间（6.0 分钟/epoch）仅比使用更深解码器的 MAE（5.0 分钟）略高，表明其具有极高的训练效率。
• 使用 PyTorch 训练的模型性能略低于原始 TensorFlow 实现，但差异微小，表明在不同代码库之间具有良好的可复现性。