SlowGuess/NeurIPS_2023_OCR huggingface dataset:
David Mizrahi $^{1,2*}$  Roman Bachmann $^{1*}$  Oğuzhan Fatih Kar $^{1}$
Teresa Yeo $^{1}$  Mingfei Gao $^{2}$  Afshin Dehghan $^{2}$  Amir Zamir $^{1}$
$^{1}$ Swiss Federal Institute of Technology Lausanne (EPFL)  $^{2}$ Apple
https://4m.epfl.ch
Current machine learning models for vision are often highly specialized and limited to a single modality and task. In contrast, recent large language models exhibit a wide range of capabilities, hinting at a possibility for similarly versatile models in computer vision. In this paper, we take a step in this direction and propose a multimodal training scheme called 4M. It consists of training a single unified Transformer encoder-decoder using a masked modeling objective across a wide range of input/output modalities – including text, images, geometric, and semantic modalities, as well as neural network feature maps. 4M achieves scalability by unifying the representation space of all modalities through mapping them into discrete tokens and performing multimodal masked modeling on a small randomized subset of tokens.
4M leads to models that exhibit several key capabilities: (1) they can perform a diverse set of vision tasks out of the box, (2) they excel when fine-tuned for unseen downstream tasks or new input modalities, and (3) they can function as a generative model that can be conditioned on arbitrary modalities, enabling a wide variety of expressive multimodal editing capabilities with remarkable flexibility.
Through experimental analyses, we demonstrate the potential of 4M for training versatile and scalable foundation models for vision tasks, setting the stage for further exploration in multimodal learning for vision and other domains.
In recent years, the field of natural language processing (NLP) has seen a shift toward training large language models (LLMs) that are inherently capable of performing a wide range of tasks without requiring extensive task-specific adaptations [12, 25]. While these models have demonstrated remarkable success in NLP, there remains a need to develop similarly versatile and scalable models for vision. A crucial aspect of scalability and versatility in vision is the ability to handle multiple (input) modalities and (output) tasks, as vision models must deal with a diverse range of sensory inputs, such as images, 3D, and text and solve a wide range of tasks. Unlike NLP, where language modeling on raw text has led to multitask capabilities [84, 86], training on only RGB images with a single objective has not exhibited the same behavior for vision. Therefore, it is deemed important to incorporate multiple modalities and tasks in training. It has been indeed suggested by psychophysical studies that multimodality is one key driver behind the development of biological intelligence [104].
To create a model that exhibits the desirable properties of foundation models in vision, it is important to consider three key aspects in terms of scalability: data, architecture, and training objective. For data, scalability means being able to benefit from more training samples toward improving performance.
For clarity, "modalities" usually denote the inputs to a model (e.g. sensory signals), and "tasks" usually denote the outputs (e.g. semantics). Our method enables a symmetric input-output structure, thus we use "modalities" and "tasks" interchangeably in this paper.
*Equal contribution & corresponding authors.
Figure 1: 4M enables training a versatile multimodal and multitask model, capable of performing a diverse set of vision tasks out of the box, as well as being able to perform multimodal conditional generation. This, coupled with the model's ability to perform in-painting, enables powerful image editing capabilities. This generalist model transfers well to a broad range of downstream tasks or to novel modalities, and can be easily fine-tuned into more specialized variants of itself.
In terms of architecture, scalability implies increased performance with growing model size and remaining stable when trainingd at large sizes. Lastly, a scalable training objective should efficiently handle a growing number of modalities without incurring excessive computational costs. In our approach, we target scalability across these three aspects while maintaining compatibility with multiple modalities.
We address these challenges by proposing a method consisting of training a single unified Transformer encoder-decoder using a multimodal masked modeling objective. We name this approach 4M (short for "Massively Multimodal Masked Modeling")\* to emphasize its ability to scale to many diverse modalities. Our method unifies the benefits of multimodal learning and masked modeling, such as (1) serving as an effective pre-training objective for learning rich representations [30, 48], (2) leading to strong cross-modal predictive coding abilities and shared scene representations [5, 44], (3) enabling models to be used for generative tasks through iterative sampling [17, 18]. Crucially, 4M combines these benefits while remaining efficient via a number of mechanisms.
To enable training a single Transformer on modalities with different formats, like text, bounding boxes, images, or neural network features, we choose to unify their representational spaces by mapping them into sets or sequences of discrete tokens [21, 22, 74, 64] using modality-specific tokenizers [110]. This tokenization approach enhances compatibility, scalability, and sharing by removing the need for task-specific encoders and heads, allowing the Transformer to be compatible with all modalities and maintain full parameter-sharing. Also, although 4M operates on a large set of modalities, it can train in a highly efficient manner through the use of input and target masking. This involves randomly selecting a small subset of tokens from all modalities as inputs to the model, while another small subset of the remaining tokens is treated as targets. Decoupling the number of input and target tokens from the number of modalities prevents the computational cost from rapidly escalating with increasing modalities, allowing for a scalable training objective.
Leveraging the availability of single-modal or text-image pair datasets, such as CC12M [19], we employ strong pseudo labeling networks to generate aligned binding data across various modalities. This pseudo labeling approach enables training on diverse and large-scale datasets without demanding them to come with multimodal/multitask annotations.
4M models are capable of performing many key vision tasks out of the box and can also be fine-tuned to achieve highly competitive performance on unseen downstream tasks and input modalities. In
$^{\ddagger}$ We sometimes also refer to the trained models as 4M (short for "Massively Multimodal Masked Model").
addition, training using a multimodal masked modeling objective leads to steerable generative models that can be conditioned on arbitrary modalities, enabling the user's intent to be expressed in a versatile manner (as depicted in Figure 4) as well as various multimodal editing tasks (see Figure 1).
We further perform an extensive ablation analysis that studies the factors affecting 4M's performance. This thorough examination, along with the simplicity and generality of our approach, demonstrates the potential of 4M for a wide range of vision tasks and further expansions.
Our main contributions and results can be summarized as follows:
1. Method: we introduce 4M, a framework for training versatile and scalable foundation models for vision tasks using a multimodal masked modeling objective. Our approach results in models that learn rich representations and perform well on a wide range of tasks without requiring task-specific adaptations.
2. Performance: we demonstrate the efficacy of our approach through extensive experiments and benchmarks, showcasing the ability of these models to perform many key vision tasks out of the box, as well as achieving highly competitive performance when fine-tuned on unseen downstream tasks.
3. Generative Capabilities: we showcase the flexible and steerable generative capabilities of models trained using 4M, enabling a variety of multimodal editing tasks utilizing conditioning on arbitrary modalities.
4. Experimental Study: we conduct an extensive ablation analysis to study the factors affecting 4M's performance, providing important insights into these models' behavior and design.
Code, models, and additional interactive visualizations are available at https://4m.epfl.ch.
2 Method Description
The 4M architecture and training objective (depicted in Figure 2) were designed with a focus on being as compatible and scalable as possible in terms of the number and type of modalities it accepts, while being conceptually simple and computationally efficient. We enable these through the conjunction of the following key aspects:
1. Tokenizing modalities: We abstract away modality-specific intricacies by mapping all modalities into sequences or sets of discrete tokens, whether they are images, text, sparse data, or neural network feature maps. This allows every possible mapping between modalities to be seen as predicting one sequence or set of tokens from another. In Section 2.1, we discuss what types of modalities we train on, how we generate the training data, and how we enable training a model on different modalities through tokenization.
2. Training a single compatible network on all modalities: Different tasks in vision, NLP and other domains traditionally required vastly different modeling choices, architectures, and losses, making the joint training on multiple modalities challenging. Tokenizing all modalities into a unified representation space allows us to train a single Transformer encoder-decoder (see Figure 2) to map between different modalities through (parallel or serialized autoregressive) token prediction. In Section 2.2, we provide more details on the 4M architecture.
3. Multimodal masked pre-training objective: Transformers have demonstrated excellent scalability with data and model size across a diverse set of tasks [60, 2], particularly when paired with a scalable pre-training objective such as masked reconstruction [30, 48, 56]. In Section 2.3, we detail our approach to training 4M using a multimodal masked modeling objective on randomized token subsets to learn strong cross-modal predictive coding abilities.
2.1 Modalities & data
Pre-training modalities. We train 4M models on a diverse set of modalities, namely RGB, captions, depth, surface normals, semantic segmentation maps, bounding boxes, and tokenized CLIP feature maps [85, 82, 114]. These modalities were chosen to cover several key aspects: First, they contain a mix of semantic information (captions, semantic segmentation, bounding boxes, CLIP), geometric information (depth, surface normals), and RGB. When used as input modalities, these modalities can be used as informative priors about the scene geometry and its semantic content [98, 69], and when used as target tasks, they allow us to steer what kind of representations are learned [42, 5, 82, 114]. Second, these modalities are diverse in terms of the format they use to encode information. They consist of dense visual modalities (RGB, depth, surface normals, semantic segmentation), sparse and/or sequence-base modalities (captions, bounding boxes), as well as neural network feature maps (CLIP [85]). Finally, these modalities allow for diverse and rich interaction with the model for
Figure 2: Method overview. (Left): 4M is a framework for training multimodal and multitask models that operate on tokenized versions of multiple image-like modalities (such as RGB, depth, etc.) and sequence modalities (such as captions and bounding boxes). (Right): The 4M pre-training objective consists of training a Transformer encoder-decoder to predict a randomly selected subset of tokens, which is sampled from all modalities, based on another random subset of tokens.
generative purposes. For example, captions, segmentation maps, and bounding boxes allow for semantically conditioned generation, while geometric modalities enable grounding the generation on 3D information. 4M's versatility in handling various modalities, its capacity to benefit from cross-training, and its ability to learn cross-modal predictive representations (as demonstrated in Sections 3 and 4) suggest its potential for extension to even more modalities.
Pseudo labeled multimodal training dataset. Training 4M models requires a large-scale and aligned multimodal/multitask dataset that contains all the above modalities/tasks and is sufficiently diverse. Most multimodal datasets, however, either do not contain all our pre-training modalities [32], are too small [94], or are not diverse enough [126]. For those reasons, we resort to pseudo labeling [42, 5] the publicly available Conceptual Captions 12M (CC12M) [19] as a binding dataset using powerful off-the-shelf models. Because this approach only requires access to a dataset of RGB images, it may scale to even larger web-scale image datasets [99, 14, 39].
Tokenization. All modalities are mapped to sets or sequences of discrete tokens (indices of a vocabulary) through the use of modality-specific tokenizers. Captions and bounding boxes are both treated as text and encoded using WordPiece [30]. For modeling bounding boxes, we follow the approach of Pix2Seq [21], that turns the task of object detection into a sequence prediction problem. RGB, depth, normals, semantic segmentation maps, and CLIP feature maps are tokenized using learned vector quantized autoencoders (VQ-VAE) [110]. Unlike Unified-IO [74] that represents all image-like modalities using an RGB pre-trained VQ-GAN [35], we instead use modality-specific tokenizers. This allows us to incorporate neural network feature maps that would otherwise be difficult to represent with existing image tokenizers. While mapping modalities to tokens and vice versa incurs a small computational overhead during inference, we avoid this overhead during pre-training by pre-computing the tokens while assembling our multimodal dataset.
We provide a detailed overview of the multimodal dataset, pseudo labeling procedure, and tokenization in Appendix B.
2.2 Multimodal Transformer
We design the architecture of 4M with efficiency, scalability, and simplicity in mind. 4M's architecture closely resembles a standard Transformer [112] encoder-decoder but includes a few crucial modifications to enable joint modeling of multiple different image-like modalities, such as RGB or semantic segmentation, but also of sequence modalities, such as captions or bounding boxes.
Multimodal encoder. The encoder is a standard Transformer encoder but features modality-specific learnable input embedding layers to map token indices to vectors. To each token of a specific modality, we add a learnable modality embedding and either 1D (for sequences) or 2D (for dense modalities) sine-cosine positional embeddings. To facilitate transfer learning, the encoder is additionally designed to accept RGB pixels using a learnable patch-wise linear projection, enabling it to double as a Vision Transformer [31] backbone.
Iteration 1 (generate / in-paint RGB with MaskGIT)
4M chained multimodal generation
Iteration 4 (generate caption autoregressively)
Detokenization
Iteration 6
Figure 3: Chained multimodal generation. This simplified example illustrates the generation of a full RGB image from a partial RGB and bounding box input using the MaskGIT [17] decoding scheme, followed by autoregressive generation of a caption. Note that through chaining (i.e. using fully generated modalities as conditioning when generating subsequent modalities), we can predict multiple modalities in a self-consistent manner. This is in contrast to independently generating each modality from the original conditioning, where each generated output is consistent with the input but not necessarily with other outputs. See Figures 9 and 10 for visual examples of chained generation. Generated tokens can be turned back into images, text, and other modalities, using the tokenizers.
Multimodal decoder. The decoder handles tokens from both dense image-like and sequence-like modalities, with each type requiring a different approach. However, two aspects are common to all tokens: First, they can all freely attend to any encoder tokens in the cross-attention layers, ensuring full access to the encoded information. Second, we employ attention masks to separate decoder tokens of different modalities. This ensures that the decoder produces consistent outputs for each specific modality, irrespective of what other outputs are being generated simultaneously. For dense image-like modalities, the decoder input consists of mask tokens along with modality and positional information. The decoder's role is to predict this masked content. For sequence-like modalities, the input to the decoder comprises modality, positional, and content information. The decoder is tasked to predict the next token in the sequence. To ensure that each token is only influenced by preceding tokens (and not by any future tokens), we apply a causal mask to the self-attention, as is standard in autoregressive models. Since all target tasks consist of discrete tokens, we can use the cross-entropy loss for all of them, which we found removes the need for task-specific loss balancing and improves training stability. Further details on the architecture are provided in Appendix C.
2.3 Multimodal masking strategy
For multimodal pre-training, we use a pre-training strategy similar to MultiMAE [5], in that we sample and encode a small set of visible tokens/patches from all modalities, and train the model to perform cross-modal predictive coding.
Input & target masking. Dropping masked-out tokens and only encoding the small set of visible ones when performing masked image modeling has been shown to yield significant increases in training efficiency [48], and is crucial when training on multiple modalities [5]. The imbalance of the usually low number of input tokens and the much higher number of target tokens can induce significant computational costs in the decoder, even if they are small. We propose to use target masking, meaning that we do not decode all masked-out tokens, but only a randomly sampled subset. By fixing the number of randomly sampled input and target tokens (see Figure 2), 4M enables pre-training on many modalities while keeping training costs low. Similar to MultiMAE [5], we sample the number of input tokens per modality using a symmetric Dirichlet distribution with concentration parameter  $\alpha$ . We follow the same approach to also sample the number of target tokens per modality. After sampling the per-modality number of input and target tokens, we sample tokens from dense modalities uniformly at random, and perform span masking [86] on sequence modalities. The experimental consequences of these design choices are studied in Section 5, and in more detail in Appendix E.4.
3 Transfer Experiments
To assess the effectiveness of 4M as a pre-training strategy, we train two models: a base version 4M-B with 86M encoder parameters and a large version 4M-L with 303M encoder parameters (for more details, see Appendix C). We then transfer these trained models to several common downstream tasks and compare their performance against relevant baselines. To better control for the dataset, augmenta
Table 1: Transfer learning study: We transfer 4M models to semantic and geometric downstream tasks and compare it to several baselines. For transfers to ImageNet-1K [29, 96], we first perform intermediate fine-tuning on ImageNet-21K [29, 93]. 4M outperforms the baselines on all tasks except for ImageNet-1K, surpassed by DeiT III which is a specialized model. In contrast to 4M, all of the baselines employed data augmentations to achieve their results. Best results per category are bolded.
Method	Pre-training data	Data aug.	Extra labels	ImageNet-1K Top-1 acc. ↑	COCO APbox↑	APmask↑	ADE20K mIoU ↑	NYU depth δ1 acc. ↑
MAE B [48]	IN-1K	✓	×	84.2	48.3	41.6	46.1	89.1
DeiT III B [108]	IN-21K	✓	×	85.4	46.1	38.5	49.0	87.4
MultiMAE B [5]	IN-1K	✓	✓	84.0	44.1	37.8	46.2	89.0
4M-B (RGB → RGB only)	CC12M	×	×	82.8	42.3	36.6	38.3	80.4
4M-B (RGB → CLIP only)	CC12M	×	✓	83.4	46.6	39.9	43.0	85.7
4M-B	CC12M	×	✓	84.5	49.7	42.7	50.1	92.0
MAE L [48]	IN-1K	✓	×	86.8	52.8	45.3	51.8	93.6
DeiT III L [108]	IN-21K	✓	×	87.0	48.7	41.1	52.0	89.6
4M-L	CC12M	×	✓	86.6	53.7	46.4	53.4	94.4
tions, model architecture, and compute, which can significantly affect downstream performance, we additionally show self-baselines that are conceptually similar to MAE [48] (Masked RGB  $\rightarrow$  RGB) and BEiT-v2 [82] (Masked RGB  $\rightarrow$  CLIP).
The transfer tasks include ImageNet-1K classification [29, 96], COCO detection and instance segmentation [67], ADE20K semantic segmentation [131], and NYUv2 depth estimation [102]. While some transfer tasks have similarities to our pseudo labeled tasks, they are different instantiations (e.g., ADE20K instead of COCO semantic classes, or absolute depth instead of relative depth).
To make 4M models comparable to other ViT backbones, we train all methods by attaching transfer task-specific heads (e.g. Cascade Mask R-CNN [15, 47]) to the encoder, and discard any decoders. Note that we can also choose to keep the decoder for transfer learning, which we explore in Section 5. For comparability, we perform the transfers of all 4M and baseline models in the same controlled manner, by closely following commonly used settings from other papers. Exact training details are provided in Appendix D.
Results in Table 1 show that 4M transfers exceptionally well to all downstream tasks, outperforming the baselines on detection, segmentation and depth estimation. While on ImageNet-1K, 4M is outperformed by more specialized models such as DeiT III, the results demonstrate 4M to be a versatile vision model that can strongly benefit from being pre-trained on multiple (pseudo labeled) tasks. We note that preliminary experiments with an even larger 4M-XL model with 2.7B parameters showed overfitting to the pseudo labeled CC12M dataset, resulting in limited additional improvement on downstream tasks. While a larger dataset or adding data augmentations are therefore necessary to fully benefit from larger 4M models, we still observe significant improvements in generation quality and use the 4M-XL model in the next section.
4 Generative Capabilities & Probing the Learned Representation
4M can directly be used for generation of all pre-training modalities, through iteratively decoding tokens [17, 65, 18], as illustrated in Figure 3. In addition, we enable several other generative capabilities by utilizing two key aspects of 4M: The first is the fact that 4M is crucially able to generate any of the training modalities, either unconditionally or conditioned on any other set of modalities (see Figure 4 top). The second is the fact that 4M is trained using masking, which enables (conditional) in-painting and out-painting (see Figs. 1 and 4). Combining these two key aspects enables several multimodal editing tasks, such as semantic editing, geometrically grounded generation, or guiding the generation with multiple strong and weak conditions (via weighting).
To improve image generation fidelity, we trained a 4M-XL version and all subsequent images were generated with it. While 4M can be directly used for generation, there are certain common improvements we perform to improve image fidelity in the results shown. These include specializing 4M-L into a super-resolution variant that maps tokens of low-resolution generations to a higher resolution [18, 123], and specializing 4M models by fine-tuning them to be more aligned with specific generative use-cases (e.g. for text-to-image, or in-painting) [95]. See Appendix A for more details on these specializations.
Probing learned representations through generation. In addition to performing transfers, we can get a glimpse into what kind of (predictive) representations 4M learned by manipulating one part of
Any-to-any generation
Grounded generation of image variations
Edit segmentation label
Mask RGB
A chain of detailed multimodal edits allows the user to better generate what they had in mind
Mask RGB