PaperlayerView exact revision on arXiv

Source-linked AI summary

MultiMAE: Multi-modal Multi-task Masked Autoencoders

Roman Bachmann, David Mizrahi, Andrei Atanov, Amir Zamir

arXiv:2204.01678v1cs.CVcs.LG

TL;DR

MultiMAE addresses the limitation of RGB-only, single-objective masked autoencoding when additional modalities and diverse downstream tasks are available. It masks patches across modalities, reconstructs multiple outputs with task-specific decoders, and trains from pseudo-labeled RGB data. The resulting network supports RGB-only, depth-only, and multimodal transfer, while the discussion identifies limited pretraining modality diversity and pixelwise reconstruction as boundaries.

  • Problem

    RGB-only masked autoencoding does not use optionally available modalities, while no single pretraining objective usually transfers best to all downstream tasks.

  • Method

    MultiMAE encodes sparse patches from multiple modalities, reconstructs masked outputs with task-specific decoders, and uses pseudo-labeled RGB data for training.

  • Results

    MultiMAE performs best across the reported tasks, matches MAE on ImageNet-1K classification and ADE20K segmentation, and outperforms it on other tasks and datasets while supporting RGB-only and additional-modalities transfer.

  • Takeaways & Limitations

    The same pretrained network can retain RGB-only transfer performance while using optional modalities when they are available, including depth-only and RGB-D configurations.

  • Takeaways & Limitations

    The experiments pretrain on three visual modalities, and pixelwise L1 or L2 reconstruction can produce blurry predictions for ambiguous masked regions.

Abstract

from arXiv · show

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.

1. Introduction

MultiMAE extends masked autoencoding beyond RGB by combining optional input modalities with multiple prediction tasks. It uses masking and pseudo-labeled data to learn representations intended to transfer across diverse downstream tasks.

  • Motivation: RGB-only masked autoencoding does not use optionally available modalities such as depth, although multiple information sources are often accessible.The paper notes that additional modalities can come from sensors or pseudo-labeling.
  • Motivation: A single pretraining objective does not usually transfer best to every downstream task, motivating diverse multitask pretraining.Modifying training objectives can steer the representation learned by the model.
  • Contribution: MultiMAE makes masked autoencoding multimodal and multitask by reconstructing multiple outputs from masked inputs across modalities.The model studies dense scene depth and segmentation maps alongside RGB.
  • Contribution: Pseudo-labeling ImageNet-1K depth and segmentation avoids requiring a large aligned dataset with multiple modalities and annotations.Only an unstructured RGB dataset and off-the-shelf neural networks are needed for pseudo-label generation.

2. Related Work

Prior work covers masked image prediction, multimodal learning, multitask learning, and self-training. MultiMAE combines these directions through masked predictive coding among optional visual inputs and outputs.

  • Masked image prediction: Masked image prediction pretrains vision models by reconstructing corrupted image content, using targets including pixels, discrete tokens, and learned features.These methods are presented as scalable self-supervised pretraining approaches for vision models.
  • Multimodal learning: Multimodal learning relates information from multiple sources through separate encoders or unified architectures such as Transformers.Prior modalities include images, text, video, audio, and depth.
  • Multimodal learning: MultiMAE differs from prior multimodal autoencoding by using a flexible masked-autoencoding architecture to learn cross-modal predictive coding among optional inputs.This distinction is tied to the model's masking strategy.
  • Multitask learning: Multitask learning commonly uses one encoder with task-specific decoders, whereas MultiMAE masks and predicts multiple input and output modalities.The related-work distinction concerns both the inputs and the outputs.
  • Multitask learning: Task diversity is used to learn representations intended to cover multiple downstream vision tasks more generally.The paper positions its pretraining method within this task-diversity literature.
  • Self-training: Self-training generates pseudo labels with supervised models, while MultiMAE additionally uses pseudo labels as masked input modalities.This extends pseudo-label use beyond output targets.

3. Method Description

MultiMAE uses a multimodal ViT encoder, task-specific decoders, modality-aware masking, and pseudo-labeled RGB data to reconstruct multiple visual tasks. Its design keeps encoder computation efficient while supporting flexible inputs and outputs.

  • 3.1. Multi-modal encoder: Each modality’s 16×16 patches receives a separate linear projection before concatenation into one Transformer token sequence.The architecture also includes a learned global token and can load selected input projections into a standard single-modal ViT.
  • 3.1. Multi-modal encoder: Only a small random subset of visible tokens enters the encoder, reducing the quadratic self-attention cost as modalities increase.This enables efficient pretraining with three dense input modalities.
  • 3.2. Decoders: Each task has a decoder that combines visible task tokens with mask tokens and cross-attends to encoded tokens from other modalities.A small MLP and Transformer blocks follow the cross-attention layer, with losses computed only on masked tokens.
  • 3.2. Decoders: Decoder computation scales linearly with the number of tasks, so shallow low-dimensional decoders preserve pretraining efficiency.The stated decoder design uses one cross-attention layer, an MLP, and two Transformer blocks.
  • 3.3. Multi-modal token sampling: Masking across modalities is intended to make the model learn cross-modal predictive coding in addition to spatial patch prediction.The sampling strategy first selects modality proportions and then randomly samples tokens within each modality.
  • 3.4. Pseudo labeled multi-task training dataset: Three pretraining tasks are pseudo-labeled on ImageNet-1K, eliminating the need for a large aligned multitask dataset.Pseudo-labeling is described as scalable to RGB datasets of arbitrary size and computationally cheap relative to training.
  • 3.4. Pseudo labeled multi-task training dataset: The task set is designed to cover low-level, geometric, and semantic vision categories identified by Taskonomy.The paper selects one representative task from each category.
  • 3.4. Pseudo labeled multi-task training dataset: RGB outputs use pixel reconstruction losses, while depth uses standardized values and an L1 loss and segmentation uses a cross-entropy loss.The segmentation target contains 133 semantic classes derived from pseudo-label predictions.

4. Experiments

Experiments evaluate MultiMAE across RGB-only, multimodal, pseudo-labeled, and pre-training-ablation transfer settings. MultiMAE remains competitive with RGB-only input, leverages additional modalities, and provides strong cross-modal predictions.

  • RGB-only transfers: MultiMAE performs best across the evaluated RGB-only transfer tasks, matching MAE on ImageNet-1K classification and ADE20K segmentation while exceeding it elsewhere.The RGB-only setting evaluates classification, semantic segmentation, and depth estimation across ImageNet-1K, ADE20K, Hypersim, and NYUv2.
  • Multimodal transfers: Ground-truth depth substantially improves MultiMAE transfer performance, while standard MAE cannot sufficiently exploit the additional modality.MultiMAE also transfers effectively from depth alone; on Hypersim, depth-only transfer surpasses RGB-only, and RGB-D performs best among the compared inputs.
  • Pseudo-labeled transfers: Pseudo-labeled depth or semantic segmentation boosts MultiMAE beyond RGB-only transfer, with smaller gains than real depth; combining both pseudo-labeled modalities performs best.These experiments use pseudo labels when ground-truth modalities are unavailable.
  • Pre-training task choices: Pre-training with RGB, depth, and semantic segmentation yields the best average Taskonomy rank and a more generalist transfer model, despite some specialized models leading on individual tasks.RGB-S slightly improves ImageNet and semantic segmentation over RGB-D-S, but sacrifices flexibility for using depth later.
  • Pre-training task choices: MultiMAE outperforms supervised single-task and multitask baselines on nearly all categories in comparisons with non-masked pre-training.The ablation compares models trained for 400 epochs using ImageNet classification, NYUv2 segmentation and depth, and average Taskonomy rank.
  • Cross-modal exchange of information: Cross-modal predictions remain plausible from a single full modality and under substantial changes in visible-patch count, while edited RGB patches alter corresponding colors without changing the background.These observations indicate cross-modal predictive coding and interaction in the learned representation.

5. Discussion

MultiMAE extends masked autoencoding across modalities and tasks while keeping encoder computation fixed through a constant number of visible tokens. The discussion identifies broader modalities, datasets, probabilistic modeling, and masking policies as directions for improvement.

  • Core strategy: MultiMAE encodes a small random subset of visible tokens from multiple modalities and reconstructs missing tokens with task-specific decoders.Masking across patches and modalities promotes predictive coding across both spatial regions and modalities.
  • Efficiency: Keeping the number of visible tokens fixed keeps most Transformer-encoder computation constant, while only shallow task-specific decoders scale with task count.
  • Future directions: The current pre-training uses three visual modalities, leaving videos, text, bounding boxes, sparse depth, and feature maps as possible extensions.A larger modality set could provide more optional inputs and broader cross-modal predictive coding.
  • Future directions: The experiments train on pseudo-labeled ImageNet-1K, while the pseudo-labeling setup can also support datasets used for RGB-only self-supervised learning.Future improvements in architectures, training, or supervised datasets could improve pseudo-label quality.
  • Future directions: Pixelwise L1 or L2 reconstruction losses can produce blurry predictions for ambiguous masked regions, motivating probabilistic or generative output modeling.The paper notes that better visual fidelity may not necessarily improve downstream representations.
  • Future directions: Random unbiased token sampling works well but may not be optimal; biasing masks toward selected modalities or spatial locations remains open.

A. Additional pre-training implementation details

The implementation details describe MultiMAE’s token, decoder, classification, segmentation, and depth-transfer configurations. Multi-modal segmentation requires a simplified output-token head and an evaluation procedure adapted to whole-image modality statistics.

  • Pre-training: The default encoder receives 49 visible tokens for one modality and 98 for two or three modalities, with segmentation inputs downsampled by four.Segmentation patches remain constant using 4×4 patches after downsampling.
  • Decoder: Each decoder linearly projects encoder tokens, adds positional and modality embeddings, applies cross-attention, then uses an MLP and two Transformer blocks before image reconstruction.In the illustration, each token expands into four pixels.
  • Classification: ImageNet-1K classification replaces the decoders with average pooling over encoded tokens, LayerNorm, and a linear projection, following MAE fine-tuning.
  • Semantic segmentation: The standard segmentation fine-tuning approach is unsuitable for multi-modal inputs because both the segmentation head and evaluation procedure require modification.
  • Semantic segmentation: The proposed ConvNeXt segmentation head projects output tokens, reshapes them into a feature map, applies four ConvNeXt blocks, and bilinearly upsamples to full resolution.It slightly outperforms Segmenter on average across ADE20K, Hypersim, and NYUv2.
  • Depth estimation: Depth estimation uses a DPT head on ViT-B layers 3, 6, 9, and 12, reverse Huber training loss, and δ1 evaluation on NYUv2.The δ1 metric measures the percentage of pixels whose error is below 1.25.

B.4. Taskonomy dense regression tasks

Taskonomy transfer evaluates MultiMAE on eight dense regression tasks using a small training subset, DPT decoding, and L1 losses.

  • Tasks: The Taskonomy transfer covers principal curvature, zbuffer depth, texture edges, occlusion edges, 2D and 3D keypoints, surface normals, and reshading.
  • Data split: Transfers use 800 training and 200 validation images from Taskonomy-tiny, with testing on the full 54,514-image test split.The checkpoint with the lowest validation loss is evaluated.
  • Training and decoding: All images are resized to 384 × 384 without further augmentation, and each task is trained with an L1 loss.The DPT head accesses ViT-B encoder layers 3, 6, 9, and 12.

C. Mask sampling strategies

MultiMAE samples visible tokens across modalities with a Dirichlet distribution, exposing the model to varied modality compositions. α = 1 performs best on ADE20K and nearly best on ImageNet-1K among tested strategies.

  • Sampling: The main experiments sample per-modality visible-token counts from a Dirichlet distribution with concentration parameter α = 1.This choice exposes the model to a broad diversity of masks.
  • Transfer results: α = 1 performs best on ADE20K and is a close second on ImageNet-1K, behind α = ∞, among the evaluated mask strategies.The comparison uses α ∈ {0.2, 0.5, 1.0, ∞} and 400-epoch pre-training without the additional standardized RGB decoder.
  • Transfer results: Smaller α values do not improve these RGB-only transfers despite exposing models to more single-modality token samples.Biasing sampling toward modalities used during transfer is proposed as future work.
  • Sampling: Small α concentrates samples in single modalities, whereas large α produces more equal modality representation.The figure illustrates samples of nine tokens drawn from three modalities.

D. Detailed Taskonomy transfer results

MultiMAE’s Taskonomy transfer results show that pre-training with task-relevant modalities improves related downstream tasks without sacrificing performance when semantic segmentation is added. Its performance is comparable to a full-input single-task RGB→D baseline, while the choice of pre-training tasks remains important.

  • Depth-inclusive MultiMAE transfers better to depth-related Taskonomy tasks, such as surface normals, than MAE or RGB-S MultiMAE.
  • Adding semantic segmentation to RGB and depth does not degrade performance on depth-related tasks.
  • MultiMAE performs similarly to the single-task RGB→D baseline trained using full RGB inputs.
  • The choice of pre-training tasks is crucial because RGB→S baselines perform consistently worse than depth-inclusive baselines and RGB-S MultiMAE.

E. Robustness evaluation on ImageNet

The study evaluates ImageNet-fine-tuned models on four ImageNet-like validation sets containing domain shifts and corruptions. MultiMAE outperforms same-size baselines on ImageNet-R and ImageNet-S, but its advantage is dataset-dependent.

  • MultiMAE performs better than all ViT-B baselines on ImageNet-R and ImageNet-S.
  • MultiMAE outperforms MAE on ImageNet-C but trails DINO and MoCo-v3.
  • On ImageNet-A, MultiMAE performs worse than DINO and MAE but better than supervised and MoCo-v3 baselines.
  • The robustness evaluation uses ImageNet-R, ImageNet-S, ImageNet-C, and ImageNet-A variants containing domain shifts and corruptions.

F. Comparison of MAE variants

The MAE-variant comparison examines decoder configurations, implementation effects, and transfer metrics across classification, segmentation, and depth. Deeper decoders can improve MAE performance but increase pre-training cost, while implementation differences produce minor transfer changes.

  • Deeper MAE decoders produce a stark performance increase, suggesting MultiMAE could also benefit from wider and deeper decoders.
  • The comparison includes MAE decoder depth 2 and width 256 versus depth 8 and width 512 configurations.
  • The deeper-decoder benefit would significantly increase MultiMAE’s pre-training time.
  • MAE models from the original TensorFlow implementation slightly outperform PyTorch models on all transfer tasks.
  • MAE variants are compared using ImageNet-1K top-1 accuracy, ADE20K, Hypersim, and NYUv2 segmentation mIoU, and NYUv2 depth δ1 accuracy.

G. Comparison of pre-training time

MultiMAE maintains pre-training efficiency through shallow decoders and sparse visible-token processing despite handling multiple dense modalities. Masked-token removal from the encoder is presented as crucial for enabling this setup.

  • MultiMAE’s pre-training time is comparable to MAE with a depth-8 decoder despite twice as many unmasked tokens and multiple decoders.
  • Removing masked tokens from the encoder is crucial for pre-training on multiple dense modalities.
  • MultiMAE predictions use 98 visible patches sampled with Dirichlet concentration parameter α = 1.
  • Additional visualizations sample three random masks for each ImageNet validation image.
Loading 2204.01678v1…