OmniMAE: Single Model Masked Pretraining on Images and Videos

FAIR, Meta AI

• OmniMAE applies one plain Vision Transformer to images and videos represented with spatio-temporal patches.The encoder receives only a subset of patches during training.Source: Figure 1 · p. 1
• A decoder reconstructs missing-patch pixels from encoder outputs, and pixel reconstruction loss trains the encoder-decoder.Source: Figure 1 · p. 1
• After training, the plain Transformer encoder performs competitively with specialized architectures on downstream image and video recognition tasks.Source: Figure 1 · p. 1
• The authors marked their technical contributions as equal.Source: FAIR, Meta AI · p. 1

1. Introduction

• Prior unified models either underperformed single-modality models or required vision-tailored architectures, limiting vanilla Transformers’ flexibility for other domains.Specialized vision Transformers can perform better on visual modalities but are less general for domains such as text, speech, and 3D.Source: §1 · p. 1
• OmniMAE uses masked autoencoding to train one omnivorous visual encoder across modalities with a shared objective and no supervision.The approach builds on masked pretraining used for images, videos, and text, audio, and images.Source: §1 · p. 2
• Masked pretraining avoids labeling effort and modality-specific supervised heads, while its simple objective scales across visual modalities with minimal changes.The passage contrasts it with discriminative self-supervision, which can be difficult to scale in model and data size.Source: §1 · p. 2
• OmniMAE is presented as the first single self-supervised model with good performance on both modalities, using only 10% of image and 5% of video patches.The 650M-parameter models reduce compute and memory by approximately 7× for images and 11× for videos; repeating samples reduces data-loading time without loss in transfer performance, and a shallow shared decoder improves performance while reducing parameters by 2–4×.Source: §1 · p. 2

2. Related Work

• Joint-embedding self-supervision learns invariances and strong linear-probing and KNN descriptors, but can be unstable and costly to scale with multi-crop augmentations.Source: §2 · p. 2
• Masked prediction learns representations by predicting masked inputs, with computer-vision methods ranging from reconstruction objectives to discrete visual-word prediction.BEiT predicts discrete visual words; SimMIM and MaskedFeat predict pixels and HOG features using specialized backbones.Source: §2 · p. 2
• MAE uses 75% masking and a low-capacity decoder, while ST-MAE [76] reports strong video recognition with 90% masking.Patch dropping benefits vanilla ViTs rather than multiscale MViT and Swin architectures; OmniMAE therefore uses a simple ViT with 90–95% masking across images and videos.Source: §2 · pp. 2–3
• Earlier multimodal systems often used separate backbones and alignment, while Omnivore and PolyViT [53] used supervised learning and human annotations.BEVT uses a specialized Swin architecture, separate image and video decoders, and a pretrained tokenizer, and cannot drop patches during training.Source: §2 · p. 3

3. OmniMAE

• OmniMAE extends masked autoencoding to a vanilla Vision Transformer for unified image-and-video pretraining.Source: §3 · p. 3

3.1. Training OmniMAE jointly on images and videos

• Images are represented as single-frame videos, and both modalities are split into spatio-temporal patches of size t×h×w×3.The input is a tensor with temporal and spatial dimensions and three color channels.Source: §3.1 · p. 3
• One omnivorous encoder shares parameters across image and video inputs, processes variable numbers of patches, and outputs per-patch embeddings.Source: §3.1 · p. 3
• The encoder processes unmasked patches; the decoder receives those embeddings plus mask tokens and positional encodings, then predicts pixels for all patches.Training minimizes reconstruction error on masked patches against normalized pixel targets using an ℓ2 loss.Source: §3.1 · p. 3
• Figure 2 reports that OmniMAE generalizes to image and video benchmarks, unlike MAE’s poor video recognition and ST-MAE’s drops on image datasets.The comparison finetunes MAE, ST-MAE, and OmniMAE at ViT-B and ViT-L after 800 pretraining epochs.Source: Figure 2 · p. 4
• Repeating a data point within a mini-batch with different crops and masks improves runtime without affecting final transfer accuracy.Replicas still present different non-masked inputs because their random crops and masks differ.Source: §3.1 · p. 4

3.2. Implementation Details

• Implementation uses ViT-B, ViT-L, and ViT-H encoders, and feeds only 10% of image patches and 5% of video patches to the encoder.Video clips use 16 frames at 6 FPS; spatial input is resized to 224×224, with patch sizes varying by model size.Source: §3.2 · p. 4

4. Experiments

• Pretraining uses ImageNet and Something Something-v2; transfer evaluation includes iNaturalist-2018, Places-365, Kinetics-400, and EPIC-Kitchens-100.SSv2 is chosen for reproducibility and its challenging temporal reasoning; all transfer results use top-1 accuracy.Source: §4 · p. 4

4.1. OmniMAE on Image and Video Recognition

• OmniMAE matches video performance of ST-MAE and image performance of MAE for ViT-B and ViT-L, while larger models improve on both modalities.MAE and ST-MAE lose performance when transferred across modalities; Figure 4 shows reconstructions generalizing to unseen K400 and EK100 data, with less detail at higher masking ratios.Source: §4.1 · pp. 4–5, Figure 4 · p. 7

4.2. Analyzing OmniMAE

• 90% image and 95% video masking reduce ViT-H FLOPS by 7.2× and 11.3×, respectively, versus using all patches.The authors select these masking ratios for final models; they report 19 unmasked image patches and 78 unmasked video patches for patch size 2×16×16.Source: §4.2 · p. 5, Table 1 · p. 6
• Random and tube masking perform comparably on image and video tasks, while causal and frame masking yield weaker video recognition; a shared decoder performs better than separate decoders.Transfer accuracy is robust to decoder capacity, and oversampling SSv2 improves video transfer without reducing image transfer.Source: §4.2 · p. 6, Table 1 · p. 6
• Sample replication speeds training while maintaining or improving transfer accuracy; on video it yields 20% faster training even with an optimized dataloader.Figure 5 studies repeating samples within a fixed mini-batch size and number of training updates.Source: Figure 5 · p. 5, §4.2 · p. 7

4.3. Comparison to Prior Work

• OmniMAE performs competitively across image and video benchmarks; ViT-H is better or within the margin of error across all benchmarks and outperforms concurrent video-focused work.Table 2 compares the joint model with prior self-supervised methods; OmniMAE is pretrained for 1600 epochs, or 2400 for ViT-H.Source: Table 2 · p. 8, §4.3 · p. 8

• OmniMAE is a unified Transformer for images and videos pretrained with masked autoencoding, and is competitive with specialized models tailored to either modality.Source: §5 · p. 8
• The model has not been empirically validated on other visual modalities such as 3D or non-visual modalities such as audio.Source: §5 · p. 8
• The study uses masked-autoencoder pretraining; the authors suggest exploring other pretraining and joint finetuning strategies, while noting pixel reconstruction's limits for linearly separable features and similarity metrics.Those properties matter when finetuning with limited labeled data.Source: §5 · p. 8
• The authors do not claim suitability for production data and note that learned representations remain susceptible to biases in visual-data distributions.Source: §5 · p. 8

• ImageNet (IN1K) uses the ILSVRC 2012 subset, with 1.28M training and 50K validation images across 1,000 classes.Source: §A.1 · p. 13
• iNaturalist-2018 (iNat18) is a fine-grained plant and animal species dataset with 437K training and 24K validation images across 8,142 classes.Source: §A.1 · p. 13
• Places-365 (P365) evaluates indoor and outdoor scene recognition using 1.8M training and 36K validation images across 365 scene classes.Source: §A.1 · p. 13

• Something Something-v2 (SSv2) contains approximately 169K training and 25K validation clips across 174 actions, with classes requiring temporal modeling.Source: §A.2 · p. 13
• EPIC-Kitchens-100 (EK100) provides densely labeled egocentric human-object interaction clips and serves as a transfer task for the learned representation.It has approximately 67K training and 10K validation clips; the task recognizes 3,806 verb-noun pairs, not all of which occur in both training and test data.Source: §A.2 · p. 13
• Kinetics-400 (K400) contains approximately 240K training and 20K validation clips, but video removals make comparisons with prior work difficult, so SSv2 is used for primary comparisons.Source: §A.2 · p. 13

B.1. Details about Pretraining

• Table 3 gives pretraining hyperparameters with dataset-specific settings in separate columns and shared settings in the middle; the stated settings cover ViT-B, ViT-L, and ViT-H.The ViT-H model in Table 2 is pretrained for 2,400 epochs.Source: §B.1 · p. 13, Table 3 · p. 13
• Asynchronous video dataloading uses asyncio to issue multiple reads without blocking, reducing amortized data-reading time to effectively zero.Sample replication is faster even than this optimized dataloading setup.Source: §B.1 · p. 14
• Pretraining minimizes ℓ2 loss on pixel values normalized by each patch's mean and variance, using 32GB+ A100 or Volta GPUs.Decoder depth and dimension vary by model size: ViT-B and ViT-L use four layers, while ViT-H uses eight.Source: §B.1 · p. 14

B.2. Transfer Learning

• Image transfer learning evaluates ViT-B, ViT-L, and ViT-H on IN1K, iNat18, and P365, using shared settings apart from peak learning rates and total epochs.The reported measure is peak test accuracy observed during training.Source: §B.2 · p. 14
• Video transfer learning samples 16 frames per clip and evaluates five clips with three spatial crops at test time.SSv2 and EK100 use 2.7-second clips, while K400 uses 2-second clips; ViT-L and ViT-H share finetuning hyperparameters.Source: §B.2 · p. 14

B.3. Details about Ablations

• Video performance is best at 95% masking, while image performance remains stable across image masking ratios.The ablation starts from 75% masking for each modality and increases masking up to 90% for images and 95% for videos.Source: §B.3 · p. 14
• Replicating samples reduces data-reading and decoding I/O in proportion to the replication factor while preserving the total batch size.For dataset ratios other than the default 1:1 IN1K:SSv2, one dataset is replicated each epoch, leading to longer training wall-clock time.Source: §B.3 · p. 15
• Across six runs, measured variance was 0.3% on ImageNet and 0.7% on SSv2.The trial used three pretraining seeds and two finetuning runs.Source: §B.3 · p. 15

B.4. Visualization details

• For visualization, the model is trained without patch mean/variance loss normalization, and predicted pixels replace masked patches while unmasked patches retain ground truth.This setup allows direct visualization of predicted pixel values.Source: §B.4 · p. 15

C. Additional Visualizations

• Additional reconstruction visualizations use 90% image masking and 95% video masking.The visualizations cover IN1K images and different video datasets.Source: C. Additional Visualizations · pp. 15–16, Figure 7 · p. 16, Figure 6 · p. 16

D. Additional Ablations

• With matched IN1K images and SSv2 data, SSv2 video classification drops by almost 5% when training uses frames instead of video clips.The comparison uses the same data, epochs, parameter updates, and learning-rate schedule; frame training still outperforms using only IN1K images.Source: Table 7 · p. 16, D. Additional Ablations · p. 16