Source-linked AI summary
Multimodal Masked Autoencoders Learn Transferable Representations
Xinyang Geng, Hao Liu, Lisa Lee, Dale Schuurmans, Sergey Levine, Pieter Abbeel
TL;DR
Multimodal pretraining needs transferable representations while overcoming contrastive learning’s dependence on paired data and modality-specific encoders. M3AE uses masked prediction with a unified vision-language encoder, and learns representations that transfer well while supporting paired and unpaired data.
Problem
Contrastive multimodal pretraining requires paired image-text data and separate modality encoders, limiting access to unpaired data and joint image-language understanding.
Method
M3AE combines image patches and language tokens in one sequence and learns a unified representation by reconstructing masked pixels and text.
Results
M3AE learns transferable representations that perform well across downstream tasks, with 75% language masking performing best and multimodal pretraining outperforming image-only MAE on ImageNet-1k linear classification.
Takeaways & Limitations
M3AE’s unified masked-token objective supports scalable training on both paired image-text and unpaired data.
Takeaways & Limitations
Large-scale M3AE pretraining consumes significant energy.
Abstract
from arXiv · showhide
Building scalable models to learn from diverse, multimodal data remains an open challenge. For vision-language data, the dominant approaches are based on contrastive learning objectives that train a separate encoder for each modality. While effective, contrastive learning approaches introduce sampling bias depending on the data augmentations used, which can degrade performance on downstream tasks. Moreover, these methods are limited to paired image-text data, and cannot leverage widely-available unpaired data. In this paper, we investigate whether a large multimodal model trained purely via masked token prediction, without using modality-specific encoders or contrastive learning, can learn transferable representations for downstream tasks. We propose a simple and scalable network architecture, the Multimodal Masked Autoencoder (M3AE), which learns a unified encoder for both vision and language data via masked token prediction. We provide an empirical study of M3AE trained on a large-scale image-text dataset, and find that M3AE is able to learn generalizable representations that transfer well to downstream tasks. Surprisingly, we find that M3AE benefits from a higher text mask ratio (50-90%), in contrast to BERT whose standard masking ratio is 15%, due to the joint training of two data modalities. We also provide qualitative analysis showing that the learned representation incorporates meaningful information from both image and language. Lastly, we demonstrate the scalability of M3AE with larger model size and training time, and its flexibility to train on both paired image-text data as well as unpaired data.
1 Introduction
M3AE addresses limits of multimodal contrastive pretraining with a unified masked-token model for image and language. It learns transferable representations and benefits from high text masking during joint training.
- Contrastive multimodal methods require paired image-text data and separate encoders, limiting unpaired-data use and joint image-language understanding.
- M3AE uses one unified model that masks image patches and language tokens, then reconstructs masked pixels and text through masked token prediction.
- M3AE trained on CC12M learns representations that transfer well to image classification and out-of-distribution detection.
- 75% language masking performs best for M3AE, versus BERT’s conventional 15% ratio, while the learned representation incorporates information from both modalities.
- M3AE scales with model size and training time and can train on paired image-text data or unpaired data.
2 Related work
Prior work develops self-supervised representations through reconstruction, contrastive objectives, and joint vision-language learning. These approaches motivate unified multimodal representation learning but face alignment challenges across modalities.
- Self-supervised representation learning via reconstruction: Reconstruction-based self-supervision pretrains language and vision models by predicting removed data and has produced transferable representations.
- Self-supervised representation learning via contrastive objectives: Contrastive learning models similarity and dissimilarity between multiple image or text views, often relying heavily on data augmentation.
- Joint learning for language and image: Joint vision-language representation learning is challenging because the modalities contain differing information and require semantic alignment.
3 MultiModal Masked Autoencoder (M3AE)
M3AE concatenates image patches and text tokens, masks both modalities, and uses a shared transformer encoder with a reconstruction decoder. Its unified loss supports paired and unpaired training data.
- Image-language masking: M3AE combines image patches and text tokens into one sequence, masks random subsets, and reconstructs the missing content.High masking reduces redundancy and makes prediction difficult to solve from neighboring visible inputs alone.
- M3AE encoder: The encoder maps visible image and language embeddings into a shared representation using one transformer with modality type and positional encodings.
- M3AE decoder: The decoder processes visible encoded tokens and learned mask tokens, then predicts image pixels and language token logits with separate output heads.
- Self-supervised training of M3AE: M3AE trains with masked-image MSE and masked-language cross-entropy losses combined as a weighted sum.
- Self-supervised training of M3AE: The same training loss supports mixtures of paired image-text data and unpaired data without changing the procedure.
4 Experiments
M3AE is evaluated for transferable representations across classification, fine-tuning, out-of-distribution detection, scaling, masking, and qualitative multimodal analyses. It consistently outperforms MAE, supports paired and unpaired data, benefits from high text masking, and captures cross-modal semantic relationships.
- ImageNet Linear Classification: M3AE significantly outperforms MAE in ImageNet linear classification while flexibly using paired image-text and unpaired image data.With 10% text annotations, M3AE reaches 53.3% versus MAE’s 45.2%.
- ImageNet Fine-tuning and Model Scaling: M3AE outperforms MAE under partial fine-tuning and across ViT-S, ViT-B, and ViT-L model sizes.The performance gap narrows when more layers are fine-tuned, possibly because of the domain gap between CC12M and ImageNet.
- Out-of-Distribution Detection: M3AE achieves higher near-OOD detection performance than MAE using both Mahalanobis and max-over-softmax scores.The reported AUROCs are 71.4% versus 69.0% for Mahalanobis scoring and 78.5% versus 77.2% for max-over-softmax scoring.
- Representation Analysis: Attention visualizations show that M3AE attends to corresponding objects and words across image patches and text tokens.The model’s representations also produce holistic reconstructions and more semantically coherent ImageNet clusters than MAE.
5 Conclusion
M3AE learns multimodal representations from image and language data without contrastive objectives and generalizes to downstream tasks. Its scalability and flexibility support training on large-scale multimodal data, while pretraining has substantial energy costs.
- M3AE learns multimodal representations from image and language data without contrastive objectives.
- M3AE supports masked image reconstruction and produces representations that cluster more strongly with ImageNet semantic labels than MAE.
- M3AE representations generalize well to downstream tasks and may support visual reasoning, dialog systems, and language-guided image generation.
- Large-scale M3AE pretraining consumes significant energy.
A.1 Pre-training datasets
CC12M contains approximately 12M image-text pairs collected from internet image URLs. The available dataset is incomplete because some URLs expired and some captions are non-English.
- CC12M contains approximately 12M image-text pairs whose images are provided through internet URLs.
- The collected dataset is incomplete because some image URLs expired and some captions are non-English.
A.2 Downstream datasets
The downstream evaluation uses ImageNet for image-encoder transferability and CIFAR-100 versus CIFAR-10 for out-of-distribution detection. The dataset details are summarized in Table 1.
- ImageNet evaluates image-encoder transferability, while CIFAR-100 and CIFAR-10 evaluate out-of-distribution detection.
- Table 1 provides detailed information about the downstream datasets.
A.3 Network architectures
M3AE uses ViT-based Transformer encoders in three model sizes and a lightweight decoder. The architecture follows MAE and supports image and text processing.
- M3AE uses ViT-S/16, ViT-B/16, and ViT-L/16 encoder architectures for image and text.ViT-S/16 is comparable to ResNet-50 in FLOPs and parameters.
- The model uses standard Transformer blocks, while its decoder is lightweight with 8 blocks and width 512.
A.4 Pre-training hyperparameters
M3AE and MAE pre-training largely follow the original MAE hyperparameters, with M3AE adding text-specific masking and loss-weight settings.
- M3AE and MAE use the original MAE optimizer, learning rate, and weight decay for pre-training.The same hyperparameters are used across network sizes and training epochs.
- M3AE adds text token mask ratio and text token classification loss weight as its only additional pre-training hyperparameters.
- Table 2 lists the pre-training hyperparameters for M3AE and MAE on CC12M.
A.5 Downstream evaluation hyperparameters
Downstream evaluation uses shared hyperparameters for M3AE and MAE across classification, fine-tuning, and OOD detection, with large-scale TPU-based pre-training infrastructure documented separately.
- M3AE and MAE use the same hyperparameters for ImageNet linear classification and CIFAR OOD detection.The corresponding settings are listed in Tables 3 and 5.
- The evaluation configuration includes separate hyperparameter tables for ImageNet linear classification, ImageNet fine-tuning, and CIFAR10 fine-tuning.
- All experiments run on Google Cloud TPUs using JAX, data parallelism, and batch size 4096 for pre-training.
- Table 6 reports TPU pod size and compute hours for 50 epochs of M3AE and MAE pre-training on CC12M.