Meta提出OmniMAE，用统一架构预训练视频和文本，ImageNet上达86.5%，SSv2上达75.3%！...

关注公众号，发现CV技术之美

本篇分享论文『OmniMAE: Single Model Masked Pretraining on Images and Videos』，丢弃95%的视频patch！Meta AI提出OmniMAE，用统一架构预训练视频和文本，在ImageNet上达86.5%，在SSv2上达75.3%！代码已开源！

详细信息如下：

• 论文地址：https://arxiv.org/abs/2206.08356[1]

• 代码地址：https://github.com/facebookresearch/omnivore[2]

      01      

摘要

基于 Transformer 的架构已经在各种视觉领域变得具有竞争力，尤其是图像和视频。虽然之前的工作已经孤立地研究了这些模态，但拥有一个共同的架构表明可以为多种视觉模态训练一个统一的模型。先前统一建模的尝试通常使用为视觉任务量身定制的架构，或者与单一模态模型相比获得更差的性能。

在这项工作中，作者展示了掩蔽自动编码（masked autoencoding）可用于在图像和视频上训练一个简单的视觉Transformer，而不需要任何标记数据。这个单一模型在图像和视频基准上学习的视觉表示与单一模态表示相当或更好，同时使用更简单的架构。

本文的单个预训练模型可以微调，在 ImageNet 上达到 86.5%，在具有挑战性的 Something Something-v2 视频基准上达到 75.3%。此外，该模型可以通过丢弃 90% 的图像和 95% 的视频patch来学习，从而实现极快的训练。

      02      

Motivation

从图像到 3D和视频，Transformer 架构在计算机视觉中的不同视觉模态中迅速变得具有竞争力。这种向统一架构的融合自然表明应该能够训练一个跨不同视觉模态的单一模型。然而，与单一模态模型相比，最近训练统一模型的尝试要么导致性能更差，要么需要使用替代架构，即 Swin Transformer，使其具有针对视觉任务的归纳偏置。虽然用于视觉的专用 Transformer 架构可以为视觉模式提供更好的性能，但它们失去了普通 Transformer 的通用性和灵活性，使得以后对文本、语音、3D 等不同领域进行建模变得更加困难。

在这项工作中，作者训练了一个适用于图像和视频的普通 Transformer。为此，作者发现最近几项关于使用掩蔽预训练的工作能够极大地提高了 Transformer 在图像、视频或跨文本、音频和图像的性能。作者认为，这种掩蔽预训练是一种可行的策略，可以跨视觉模态预训练统一的Transformer。具体来说，作者考虑使用掩蔽自动编码（MAE）方法来训练 Omnivorous 视觉编码器。生成的 OmniMAE 模型从具有相同目标函数的所有模态中学习，并且不需要任何监督。

与监督目标或判别式自监督目标相比，使用掩蔽预训练目标具有几个优点。首先，与监督目标相反，通用无监督损失不需要任何人工标记工作。因此，它对由一组预定义的标签引入的偏置具有鲁棒性。此外，它不需要多头架构来结合来自与每个模态相对应的每个标签空间的监督，这很难维护和扩展新模态。其次，尽管与重建目标相比，判别式自监督方法产生了更好的冻结特征，但它们在模型和数据大小上的规模化并非易事。本文的掩蔽预训练目标简单、训练效率高，并且可以以最小的变化扩展到不同的视觉模态。

本文的工作建立在将 Transformer 扩展到图像的基础上，称为 Vision Transformer (ViT)，并通过使用时空patch将其扩展到图像和视频。由于 ViT只能处理一小部分输入，即通常只覆盖图像 10% 和视频 5% 的非掩码输入，因此使用 ViT 的重建损失在计算上是高效的。作者还表明，在小batch中简单地重复样本可以减少数据加载时间，而不会损失最终迁移性能。样本复制对于掩码预训练特别有用，因为即使是复制样本，未掩码的patch在副本之间也是不同的。

本文预训练的 OmniMAE 模型使用相同的模型参数和简单的 ViT 架构在图像和视频识别基准上显示出具有竞争力的性能，甚至优于针对单个模态进行调整和训练的专门架构。在微调本文的单个模型后，作者在六个图像和视频识别基准上获得了强大的性能，包括 ImageNet 上 86.5% 的准确率和Something Something-v2 视频分类任务上的 75.3% 准确率。

      03      

方法

本文的目标是为图像和视频预训练一个统一的模型。作者没有使用为视觉模态量身定制的专门架构，而是建立在 vanilla Vision Transformer (ViT)架构之上，该架构具有有限的视觉归纳偏置。对于预训练，作者扩展了简单的自监督掩蔽自动编码 (MAE) 方法。原始架构和预训练方法仅在图像上进行测试，作者展示了统一模型的简单设计决策。

3.1 Training OmniMAE jointly on images and videos

作者在上图中说明了本文的方法。对于预训练，本文使用编码器-解码器架构，其中编码器仅对输入的“non-masked”子集进行操作。解码器预测整个输入的像素值，即mask和non-masked像素。该模型被训练以最小化输入的被mask部分的重建误差。预训练后，作者通过迁移学习评估编码器（解码器被丢弃）。

Images and videos as spatio-temporal patches

输入图像或视频可以表示为形状为的 4D 张量，其中 T 是时间维度，H、W 是空间维度，3 表示颜色通道。将图像视为 T = 1 的单帧视频。输入被分成 N 个时空patch，每个块的大小为。

Omnivorous visual encoder

作者使用统一视觉编码器，它使用相同的参数处理图像和视频。编码器对来自图像和视频的 N 个时空patch进行操作。编码器可以自然地处理来自图像和视频的可变数量 N 个patch，因为它使用了 Transformer 架构。编码器为图像和视频输入共享相同的参数，并以相同的方式处理它们以输出每个patch嵌入。

Pretraining

作者将输入转换为 N 个时空patch并随机mask  M 个块。将未屏蔽的 N - M 个patch（及其位置）馈送到编码器以生成每个patch嵌入。编码器输出与可学习的“mask token”的 M 个可学习token拼接以获得 N 个嵌入。作者将适当的位置编码添加到 N 个嵌入中，并将它们输入解码器，解码器输出像素值。

Loss function and optimization

作者最小化解码器预测和输入像素值之间的重建误差。输入像素值被归一化为零均值和单位方差以获得损失的目标。作者最小化了 M 个mask patch上的预测和目标之间的距离。在训练的每个点，作者都会对一个batch图像或视频进行采样，并使用解码器预测来计算损失。

Dataset sample replication

由于本文对输入patch的一个非常小的子集（M ≪N）进行操作，本文方法的数据读取速度无法跟上能够在单个优化步骤中处理的数据点数量。这个问题对于视频来说更加明显，必须花费大量的计算来读取和解码一个完整的视频，而只能丢弃超过 90% 的视频。

受RepeatAugment 成功的启发，这表明在mini-batch中复制样本是一种提高泛化能力的有效技术，作者复制加载的数据点并每次对其应用不同的掩码。即使使用复制的样本，模型的非掩蔽输入也是不同的，因为随机裁剪和掩蔽在副本中是不同的。

3.2 Implementation Details

上表展示了本文方法在预训练时候的一些参数细节。

上表展示了本文方法在图片数据集上微调时候的一些参数细节。

上表展示了本文方法在视频数据集上微调时候的一些参数细节。

Architecture

作者将 ViT架构用于编码器，并对其 ViT-B、ViT-L 和 ViT-H 变体进行实验。作者在 ViT 模型中不使用 [CLS] token，在运行时产生了小的改进，而性能没有任何损失。作者使用由 4 层（ViT-H 为 8 层）组成的 Transformer 解码器，ViT-B、ViT-L 和 ViT-H 的嵌入维度分别为 384、512 和 512。解码器为所有输入patch中的像素输出 RGB 颜色。作者对编码器和解码器中的patch使用正弦位置编码。

Training details

作者使用分布式训练以 2048 的 mini-batch 大小训练本文的模型。作者将输入图像和视频的空间大小调整为 224×224 像素。对于视频，作者以 6 FPS 采样 T = 16 帧的片段。作者对 ViT-B 和 ViT-L 使用 2×16×16 的patch大小，对 ViT-H 使用 2×14×14 的patch大小。图像被临时复制以满足patch大小。

Masking the input

与之前的工作相比，本文在预训练中使用了极高的掩蔽率，并且只有 10% 和 5% 的图像和视频patch被馈送到编码器进行预训练。

      04      

实验

上图展示了所有模型在图像和视频基准上的评估结果。特定于模态的 MAE 和 VideoMAE 模型在转移到与预训练模态匹配的数据集时实现了良好的转移性能。然而，当转移到不同的视觉模态时，这两种模型都表现出性能下降。

MAE 的视频识别性能较低，尤其是在来自 EK100的视频上，其中较小的 ViT-B MAE 模型比 VideoMAE 差 18.4%。同样，与 MAE 相比，VideoMAE 在 iNat18 的细粒度分类任务上差 8.7%。而OmniMAE在这两个视觉模态上都取得了非常好的结果。

作者重新训练 OmniMAE 而不对像素目标进行归一化以获得易于可视化的 RGB 重建。上图展示了使用训练期间未见过的样本（验证集或完全不同的数据集）可视化图。

上表展示了本文模型的消融实验结果。（a）中可以看出本文的方法可以在视频和图片上达到非常高的掩蔽率；（b）中展示了不同掩蔽策略（示意图如下所示）的影响；（c）展示了不同解码器结果的影响。

上图展示了图片和视频样本复制以及数据集比率对实验结果的影响。

在上表中，作者将 OmniMAE 在图像和视频分类基准上的表示与其他最先进的自监督方法进行了比较。

      05      

总结

在本文中，作者了 OmniMAE，这是一个用于图像和视频的统一Transformer，可以使用掩蔽自动编码进行预训练。OmniMAE 使用简单的架构，视觉特定的变化最小，并且与为图像和视频量身定制的专业架构和模型具有竞争力。

OmniMAE 在图像和视频上进行了测试，但尚未在其他视觉模态（如 3D）或非视觉模态（如音频）上进行实验验证。然而，OmniMAE 的简单性和通用性很可能使未来的多模态系统能够使用共享参数对多模态进行建模。

作者只研究了基于掩蔽自动编码器的 OmniMAE 预训练策略，因为它简单且易于扩展。然而，像素重建方法不能学习良好的线性可分特征或相似性度量。当使用带有有限token数据的预训练模型时，这些属性是必不可少的。在联合建模的背景下，探索这种预训练策略和联合微调策略可能会提高性能。

参考资料

[1]https://arxiv.org/abs/2206.08356
[2]https://github.com/facebookresearch/omnivore

▊ 作者简介

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。

知乎/公众号：FightingCV

END

加入「计算机视觉」交流群👇备注：CV