MultiMAE: Multi-modal Multi-task Masked Autoencoders

TL;DR

MultiMAE extends Masked Autoencoders to multi-modal and multi-task pre-training by using a ViT-based encoder with per-modality projections and shallow task-specific decoders. A Dirichlet-based masking strategy enforces cross-modal predictive coding, and pseudo labeling enables large-scale training on RGB data alone. Empirical results across ImageNet, ADE20K, Taskonomy, Hypersim, and NYUv2 show strong transfer improvements, especially when depth and segmentation modalities are available or pseudo-labeled. The approach yields a flexible, efficient framework that maintains RGB-only performance while significantly boosting multi-modal transfer, highlighting cross-modal information exchange learned during pre-training.

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.

Figures (10)

• Figure 1: MultiMAE pre-training objective. We randomly select 1/6 of all 16$\times$16 image patches from multiple modalities and learn to reconstruct the remaining 5/6 masked patches from them. The figure shows validation examples from ImageNet, where masked inputs (left), predictions (middle), and non-masked images (right) for RGB (top), depth (middle), and semantic segmentation (bottom) are provided. Since we do not compute a loss on non-masked patches, we overlay the input patches on the predictions.
• Figure 2: (Left) MultiMAE pre-training: A small subset of randomly sampled patches from multiple modalities (e.g., RGB, depth, and semantic segmentation) is linearly projected to tokens with a fixed dimension and encoded using a Transformer. Task-specific decoders reconstruct the masked-out patches by first performing a cross-attention step from queries to the encoded tokens, followed by a shallow Transformer. The queries consist of mask tokens (in gray), with the task-specific encoded tokens added at their respective positions. (Right) Fine-tuning: By pre-training on multiple modalities, MultiMAE lends itself to fine-tuning on single-modal and multi-modal downstream tasks. No masking is performed at transfer time.
• Figure 3: MultiMAE predictions for several randomly sampled masks. For each ImageNet validation image, we randomly sample three masks using Dirichlet concentration parameter $\alpha=1$. Only 1/6 of total patches are left unmasked. Even when very few tokens from one modality are visible, the resulting predictions are relatively stable and plausible due to cross-modal interaction. More examples are shown in the Appendix and on https://multimae.epfl.ch/#random-masks.
• Figure 4: Single-modal predictions. We visualize MultiMAE cross-modal predictions on ImageNet-1K validation images. Only a single, full modality is used as input. The predictions remain plausible despite the absence of input patches from other modalities.
• Figure 5: Demonstration of cross-modal interaction. The input is the full depth, only two RGB patches, and no semantic segmentation. By editing the hue of a single input patch, the color of the lizard (left) and oranges (right) changes, while keeping the background constant. More interactive examples are available on https://multimae.epfl.ch#hue-change.