MultiMAE: Multi-modal Multi-task Masked Autoencoders

Standard Masked Autoencoding is an effective pre-training strategy, but so far has been limited to a single modality, namely RGB images. Practically speaking, often more than just a single modality is available (either from sensors or pseudo labeled), and pre-training strategies should be able to optionally make use of them. In addition, pre-training on more than just a single task is a powerful tool to steer what representations a model will learn and has been shown to improve performance on transfer tasks.

The MultiMAE pre-training objective is to reconstruct masked-out patches of multiple modalities. To do that, a small set of visible patches is sampled from all modalities and processed jointly using a Transformer encoder. Task-specific decoders reconstruct the masked-out patches by first performing a single cross-attention step from queries to all 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. By keeping the number of encoded patches constant, the bulk of the computational cost does not scale with the number of modalities, making MultiMAE a highly efficient pre-training method.

We pre-train MultiMAE on two additional tasks besides RGB images, namely scene depth and semantic segmentation. Using powerful pre-trained networks, we pseudo label the entire ImageNet-1K dataset with these tasks. This has the advantage that no aligned multi-task dataset is required – all one needs is a large unlabeled RGB dataset and off-the-shelf pseudo labeling networks.

In the RGB-only transfer case (left part of figure), the encoder of a pre-trained MultiMAE can directly be used as a ViT with a newly initialized decoder head for the given downstream task. Since MultiMAE was pre-trained on more than just RGB, it can optionally accept additional modalities at transfer time (right part of figure).

MultiMAE is an effective pre-training strategy even when RGB is the only available downstream modality. It retains the benefits of MAE, an already powerful pre-training strategy, and can sometimes even outperform it. Notice how MAE and MultiMAE can both significantly surpass supervised ImageNet-1K pre-training.

Making use of additionally available modalities during fine-tuning has the potential to significantly increase performance. On NYUv2 and Hypersim semantic segmentation, using either RGB-only or depth-only inputs works well, but when combining the two, we see large increases in performance – especially so when the ground-truth depth is of high quality, as in Hypersim.

In case ground truth depth is not available at transfer time, MultiMAE can also make use of pseudo labeled depth (and even pseudo labeled semantic segmentation 😉 ). Here, pD = pseudo depth, and pS = pseudo semantic segmentation. While the improvement is not as large as with ground-truth depth, adding pseudo labeled depth and semantic segmentation still provides a slight performance increase.

MultiMAE demonstrates intriguing capabilities for cross-modal predictive coding since not only was it trained to predict each modality from any other modality, but also to combine information from multiple modalities. In the following interactive demonstrations, you can explore some of these capabilities.

How do MultiMAE predictions change if you give it more RGB patches than depth patches, or even only semantic information? In this interactive demo, you can use the slider to change the proportion of visible tokens per modality. Notice how semantically stable the predictions stay as tokens are completely removed from certain modalities and added to others, suggesting that MultiMAE has learned a shared representation for these modalities.

In this demo, we give as input the full depth map and two RGB patches, and we show the full reconstructed RGB image. No semantic maps are given as input. Using the slider, you can change the hue of a single RGB patch, while the other is locked. Notice how only the color of the selected objects change, while the rest stays constant. MultiMAE effectively propagates color information from two RGB patches to only the relevant objects in the image, which it extracts from the depth map. In the images showing fruit and vegetables, color information is not only propagated to the same object, but to all objects that MultiMAE implicitly assigns to the same semantic class – even over larger distances.

Similar to the above demonstration, we now again give a full depth map, but this time change the hue of the three middle RGB patches. Depending on the given color, MultiMAE predicts the objects as either oranges, apples, or sports balls.

Let's see how MultiMAE reconstructs incomplete RGB inputs given full semantic segmentation maps. When editing the "road" class to "grass" or "snow", the RGB predictions change accordingly. Notice how even the trees in the background are much greener when there is no snow in the semantic segmentation inputs, even though their class label is the same for all examples!