【论文阅读】Multi MAE：Multi-modal Multi-task Masked Autoencoder

论文地址：MultiMAE: Multi-modal Multi-task Masked Autoencoders

代码地址：https://github.com/EPFL-VILAB/MultiMAE

引言

MultiMAE 的设计基于传统的 Masked Autoencoding，但在两个关键方面有所不同：

1、除了 RGB 图像，它还可以选择接受输入中的附加模态信息（因此是“多模态”）

2、其训练目标相应地包括 预测除 RGB 图像之外的多个输出（因此称为“多任务”）。

架构

（左）MultiMAE 预训练：将来自多种模态（例如 RGB、深度和语义分割）的一小部分随机采样补丁被线性投影到具有固定维度的标记，并使用 Transformer 进行编码。特定于任务的解码器通过首先执行从查询到编码标记的cross-attention，然后是浅层 Transformer 来重建被屏蔽的补丁。查询由掩码标记（灰色）组成，在它们各自的位置添加了特定于任务的编码标记。 （右）微调：通过对多种模态进行预训练，MultiMAE 有助于对单模态和多模态下游任务进行微调。传输时不执行屏蔽。

Multi-model encoder

多模态 Transformer 编码器是一个 ViT ，但每个额外的输入模态都有patch projection layer。具体来说，每个模态的 16×16 块被投影到具有正确 Transformer 维度的标记，对每个模态使用不同的线性投影。投影的补丁被连接成一系列标记，并作为输入提供给同一个 Transformer 编码器。我们还添加了一个带有学习嵌入的额外全局标记，类似于 ViT 中使用的类标记。由于与 ViT 的架构相似，MultiMAE 预训练权重可以通过仅加载所需的输入投影并忽略其他投影来直接用于标准单模态 ViT。

仅将所有标记的随机抽样子集作为掩码自动编码目标的一部分传递给 Transformer 编码器。

Dncoder

为了从可见标记中重建被屏蔽的标记，我们为每个任务使用一个单独的解码器。每个解码器的输入是来自它正在重建的相应任务的完整可见标记集。与 MAE一样，这些可见标记与一组掩码标记一起解码，用作解码器的占位符以写入重建的补丁（如图 2 所示）。为了整合来自其他模态的编码标记的信息，我们在每个解码器中添加一个交叉注意层，使用这些标记作为查询，并将所有编码标记作为键/值。在此步骤之前，将正余弦位置嵌入和学习模态嵌入添加到标记中。然后是一个小的 MLP 和 Transformer 块。在 MAE 之后，我们只在mask token上计算损失。

由于每个任务都需要自己的解码器，因此解码器的计算成本与任务数量呈线性关系。

为了保持预训练的效率，我们使用低维（256 维）的浅层解码器（单个交叉注意层和 MLP，后跟两个 Transformer 块）。与编码器相比，这些解码器对整体计算成本的增加很少，它们在 ImageNet-1K 微调上的表现与更深层次的解码器相似。

Multi-modal masking strategies

为了使掩码自动编码正常工作，需要屏蔽掉大部分标记。 掩码采样策略的选择会对传输性能产生很大影响。更具体地说，对于 MultiMAE 和一般学习多任务表示，跨不同模态的掩码可确保模型开发跨不同模态的预测编码，除了不同的空间块。为了提高效率和简单性，我们为所有实验选择固定数量的可见标记，固定为 98。当使用尺寸为 224×224 像素和补丁大小为 16×16 的三种模态时，这相当于所有标记的 1/6（即掩码1/6，重建剩余的5/6）。

通过从所有标记中统一选择可见标记来调整 MAE 掩码采样策略将导致大多数模态以相似的程度表示。一种或多种模态的样本很少或没有样本的情况将非常罕见。我们提出了一个多模式令牌抽样策略,这种抽样策略允许更多样化的抽样方法。它可以分为两个步骤：首先，选择每个模态的标记数量，其次，为每个模态随机抽取标记集。

Pseudo labeled multi-task training dataset

我们使用我们在 ImageNet-1K [23] 上伪标记的三个任务对 MultiMAE 进行预训练。伪标签的优点是我们不需要具有对齐任务图像的大型多任务数据集。相反，为我们想要训练的任务访问一组好的预训练神经网络可能是有效的。伪标签可扩展到任意大小的 RGB 数据集，并且是一次性预处理步骤。与训练成本相比，如果并行化，此步骤计算成本低且速度快。

实验方法

为了执行语义分割补丁投影，论文的作者用学习的 64 维的类嵌入替换每个类索引。并且仅对可见标记的随机子集进行编码，这样可以显著的加速计算和减少内存使用，并且使用了具有三种密集输入模态的 MultiMAE 多模态预训练。 每个任务使用一个单独的解码器，因此解码器的计算随着任务的数量线性扩展，并且只增加了最小的成本。

在图像分类、语义分割和深度估计这三个任务上对 MultiMAE 进行了预训练，并在 ImageNet-1K 上进行伪标记，然后在 ImageNet、ADE20K、Taskonomy、Hypersim 和 NYUv2 数据集上进行微调。

• table1:仅在RGB上进行微调，性能提升

• table2：使用 RGB 和真实深度进行微调，MultiMAE 可以有效地利用其他模式，例如深度，而 MAE 不能。

• table3：使用 RGB 和伪标签进行微调，MultiMAE 比 MAE 从伪标记模态作为输入中受益更多。

结果表明，当 只使用RGB 进行微调时，MultiMAE 保留了常规 MAE 的优势，并且它还可以利用深度等其他模态，例如使用伪标记深度或语义分割来提高性能。 MultiMAE 预训练策略可以显着提高迁移性能。

• table4：消融实验，我们在表 4a 中研究了其他模式的影响，并在表 4b 中将 MultiMAE 与非掩码预训练进行了比较。虽然一些专门的预训练模型在某些下游任务中表现更好，但在其他任务中表现不佳。使用 RGB、深度和语义分割进行预训练的 MultiMAE 是一种更通用的模型，可以很好地转移到一系列下游任务。

结论

提出了多模态多任务掩码自动编码器 (MultiMAE)，这是一种有效且简单的 Vision Transformers 预训练策略。 MultiMAE 对来自多种模态的可见标记的一小部分随机子集进行编码，并经过训练以重建缺失的标记。通过仅对固定数量的非屏蔽标记进行编码，我们可以使 Transformer 编码器中的大部分计算保持不变，而只有浅层任务特定的解码器会随着任务数量的增加而扩展。掩蔽（跨图像块和输入模态）确保网络除了跨不同空间块之外，还学习跨不同模态执行预测编码。实验展示了 MultiMAE 在交叉模态编码方面的有趣功能，并证明了这种预训练策略可以在可选地提供额外的输入模式（作为基本事实或伪标签）时显着提高传输性能。