SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weighting of different modalities ("classifier-free guidance") for added steering power to guide edits
Weakly condition on masked RGB and
Strongly on caption & bounding box
Weakly condition on RGB and strongly
on caption
Equally condition on caption and
masked RGB
Figure 4: Multimodal generation and editing. 4M's in-painting and any-to-any prediction abilities unlock a suite of multimodal generation and editing capabilities, which allow for fine-grained creative control. We show several key capabilities such as grounding the generation in predicted geometry, performing semantic edits, and being able to control how much certain input modalities influence the generation via weighting.
4M implicitly infers object context
4M can perform semantic manipulations
Figure 5: Probing the learned representation through manipulation: (Left): All of the conditioning bounding boxes are fixed except the marked "potted plant" ones that are changing. Depending on where the "potted plant" bounding boxes are placed, 4M infers how they can be part of the scene (e.g., as a painting or real plant) in a geometrically and physically plausible manner. (Right): Changing a single semantic class accordingly affects how 4M predicts the overall image. See more examples in Appendix A.4 and interactive visualizations on our website.
the input and keeping the remainder fixed [5]. Figure 5 shows a series of manipulations, in which 4M shows intriguing capabilities at predicting geometrically and physically plausible arrangements while taking into account the semantic context of the scene.
Multimodal editing. By combining the multimodal conditional generation and in-painting capabilities of 4M, we can perform various multimodal editing tasks, such as performing semantic edits or in-painting grounded by geometric conditioning (see Figure 4 top and middle). Drawing parallels to ControlNet [130], these allow for steering the generation using more than just text. However, 4M is able to perform these tasks with just a single network and condition on multiple (partial) modalities - individually or simultaneously. The conditions can either be hand-specified, or extracted from an image using 4M itself, thereby removing the need for specialist models to create the conditions.
Multimodal weighted guidance. Classifier-free guidance [50] has been shown to improve image fidelity in token-based generative models [40, 123, 18]. Inspired by Liu et al. [68] that perform compositional generation on multiple text conditions, we can guide our generation by weighting different (parts of) modalities by different continuous amounts - even negatively. This unlocks further
multimodal editing capabilities (see Figure 4 bottom), such as being able to weakly condition on certain modalities, or using negative weighing to avoid a certain concept in the generation. Multimodal guidance can be achieved by computing a weighted sum of the logits of an unconditional and each conditional case:  $\log_{\mathrm{guided}} = \log_{\mathrm{uncond}} + \sum_{i=1}^{n} w_i$  ( $\log_{\mathrm{cond}, i} - \log_{\mathrm{uncond}}$ ).
We provide additional visualizations covering 4M's wide range of generative capabilities on our website and in Appendix A.4.
5 Ablations
Pre-training on a large set of modalities using a masking objective creates a large design space that raise questions such as: what modalities should we pre-train on, what is the optimal masking ratio, or how do we select the number of tokens to mask from each modality? By performing a thorough ablation of these design parameters, we aim to find out which ones matter the most for multimodal pre-training and which ones do not. We used these findings to decide on the settings for the models shown in Section 3 and 4.
We first choose a reference setting and measure the deviation of the model performance as a result of the one aspect under study while keeping the rest fixed. Performance is measured through transferring the models to a large set of downstream tasks and measuring the validation set performance. The aim of the benchmarks is to measure how good a certain instantiation of 4M is at transferring both to new target tasks, but also to unseen input modalities. For that, we include tasks that use RGB (pixels) as inputs and transfer to a new target, such as COCO [67] object detection, ADE20K [131] semantic segmentation and ten transfers from RGB to dense tasks in Taskonomy [126] and Hypersim [94] (e.g. depth, curvature, segmentation). Furthermore, we include eleven single-modal and twelve multimodal transfers from some Taskonomy or Hypersim modalities to others (e.g. curvature  $\rightarrow$  occlusion edges, or RGB + depth  $\rightarrow$  segmentation).
For simplicity, and similar to pre-training, we model each of the benchmark tasks as predicting one set of tokens from another. This is achieved by training tokenizers on these modalities and tasks in the same way as we did for pre-training. All transfers are performed at  $224 \times 224$  resolution. For comparability across modalities and tasks, we measure the validation set cross-entropy performance rather than task-specific metrics. Note that in the interest of space, we aggregate the results of the Taskonomy and Hypersim tasks, and denote them as  $RGB \rightarrow X$ ,  $X \rightarrow Y$  and  $X + Y \rightarrow Z$ . See Appendix E for further training details.
5.1 Reference model
Following common practice for the reference model size [30, 86], our model consists of 12 encoder and 12 decoder layers. We train it on CC12M [19] using all modalities as both inputs and targets, at resolution  $224 \times 224$  pixels (corresponding to  $14 \times 14$  tokens per dense modality), and using no augmentations such as cropping nor color augmentations. The total training length is fixed at 100B tokens, corresponding to roughly 400M masked samples. We set the number of randomly sampled input and target tokens to 12 each, and sample each using a symmetric Dirichlet distribution with parameter  $\alpha = 0.2$ . To determine the significance of various modeling choices, we adopt the approach used by Raffel et al. [86] that calculates the standard deviation of the transfer results of ten independently trained reference models.
5.2 Input modalities and target tasks
Importance of target tasks for representation learning. 4M is both a multimodal and a multitask training scheme and the choice of target task(s) is a powerful way of steering what representation the model learns [10, 109, 42, 107, 5]. To ablate this for 4M, we fix the input modalities to be either RGB or all the modalities (denoted as "All"), and vary the target tasks. The results in Table 2 mirror the findings of Sax et al. [98], that the optimal choice of pre-training setting depends highly on the type of transfer that is performed, and that there is no single pre-training task that performs best on all transfers. However, pre-training on all target modalities consistently outperforms the other single-task and multitask alternatives in terms of average loss, no matter what input modalities were used during pre-training. This makes it the preferred configuration for generalist models, especially when their future applications are unknown or varied.
Importance of multimodal pre-training for transferring to new input modalities. Table 2 shows that multimodal pre-training can significantly help with transferring to new input modalities  $(\mathrm{X} \rightarrow \mathrm{Y}$  and  $\mathrm{X} + \mathrm{Y} \rightarrow \mathrm{Z}$  transfers), but comes at a performance loss at transfers that use RGB as the sole input modality. In Appendix E.4, we explore pre-training using mixtures of different masking strategies, which enables us to train models that perform well in both regimes.
Table 2: Pre-training input and target modalities ablation: The choice of pre-training tasks and modalities influences what representations the model learns, and how well it can be transferred to novel tasks and modalities. Here, Geometric = RGB + Depth + Normals and Semantic = RGB + Segmentation + CLIP + Detection + Captions. We show the average losses (↓) for several task categories and compute the average rank and best losses for "RGB" and "All" inputs separately. The reference model setting is indicated by  $\triangleright$  and results that lie within two reference model standard deviations of the best result are bolded. Performing 4M pre-training on all input and target modalities is the most versatile choice, if the optimal set of pre-training modalities for any given downstream task is unknown.
Pre-training inputs	Pre-training targets	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
RGB	RGB	3.14	6.21	5.03	6.94	6.15	5.49	8.00
RGB	Depth	3.11	6.06	4.72	6.89	5.84	5.32	4.00
RGB	Normals	3.12	6.02	4.66	6.83	5.87	5.30	3.20
RGB	Segmentation	3.17	5.94	4.84	6.86	5.89	5.34	4.60
RGB	CLIP	3.07	6.11	4.83	6.85	5.94	5.36	4.80
RGB	Detection	2.78	6.11	5.03	7.07	6.24	5.45	7.20
RGB	Captions	3.45	6.55	5.92	7.35	6.86	6.03	10.00
RGB	Geometric	3.11	6.08	4.70	6.88	5.85	5.32	3.80
RGB	Semantic	2.88	5.99	4.86	6.97	6.06	5.35	5.40
RGB	All	2.90	5.99	4.74	6.91	5.93	5.29	4.00
All	RGB	3.21	6.20	5.07	6.75	5.85	5.42	3.80
All	CLIP	3.19	6.18	5.06	6.80	5.88	5.42	3.80
All	Geometric	3.20	6.13	4.98	6.72	5.76	5.36	1.80
All	Semantic	3.05	6.13	5.16	6.77	5.87	5.39	3.40
All	All	3.06	6.11	5.07	6.75	5.80	5.36	2.20
Figure 6: Ablations results: We ablate several key design choices of the multimodal masking objective (in blue), and study how well 4M scales (in green). We show the overall average losses  $(\downarrow)$  and highlight the reference model setting in blue / green. A detailed breakdown of the task losses is provided in Appendix E.4 and Appendix E.5.
5.3 Multimodal masking strategy
Multimodal masking is at the core of 4M, so in this section, we ablate how modality tokens should be sampled, and how many tokens we should encode and decode.
Modality proportions in masking. We ablate various choices of Dirichlet parameter  $\alpha$ , both for the input and target sampling. If  $\alpha$  is low, the sampling procedure will often select cases where most of the tokens are sampled from only one modality. If  $\alpha$  is high, however, most samples will contain tokens from all modalities to equal proportions. Results in Figure 6 (a) show that uniformly sampling over the simplex performs best on average, but not by a large margin.
Input masking budget. The difficulty of the multimodal masked modeling task is largely determined by the number of visible (non-masked) input tokens. Encoding only the small set of visible tokens can significantly improve training efficiency [48]. Figure 6 (b) shows that with a fixed training token budget, training with 128-256 input tokens performs well.
Target masking budget. We can decide to decode merely a small random subset of all remaining masked-out tokens, which is especially important for enabling efficient multimodal training with large decoders. As Figure 6 (c) shows, decoding only a small random subset of all targets performs well (for a fixed number of total training tokens), while also reducing computational costs.
5.4 How well does 4M scale?
Scalability is a key property that models and training objectives should have. We therefore ablate the following three axes: To ablate the dataset size, we train 4M models on various subsets of CC12M, down to 1/64th of the full dataset. To ablate the training length, we vary the total number of tokens seen. We define tokens seen as the total number of both input tokens and target tokens the model was
trained on. To ablate the model size, we train different sizes of 4M models, ranging from Tiny (4M-Ti) with 24M parameters to Large variants (4M-L) with 705M parameters. Exact model specifications are given in Appendix C.1. Figure 6 (d), (e), (f) show that 4M scales with dataset size, training length, and model size, respectively.
For additional ablations on the architectural and training design choices of 4M, see Appendix E.
6 Related Work
Large language models have been demonstrated to be capable of performing a diverse range of tasks out of the box [86, 12, 81, 25, 45] by training on large datasets with simple objectives [30, 84, 86, 106]. In vision, however, many scaling efforts have instead focused on training specialized models on a single task and modality, such as predicting masked RGB pixels [20, 31, 4, 48, 118, 34], discrete tokens [7, 132], or other (deep) features [117, 6, 82, 114, 36, 70] from RGB inputs. Training models instead on multiple tasks [16, 33, 63, 42, 92, 11] and modalities [85, 122, 76, 59, 133, 1, 3, 57, 43], or both [54, 103, 114, 5, 44] usually requires modality-specific modeling choices, making it difficult to extend these methods.
While some recent works aim to consolidate various modalities and tasks by representing them as images [77, 8, 115], these approaches have limitations when dealing with modalities that cannot be readily converted into images, such as text or neural network feature maps. Instead, 4M adopts the approach of Pix2Seq [21, 22] and Unified-IO [74] which addresses these issues by unifying the representation space on which models are trained through tokenization [110, 35, 30]. However, unlike methods like Unified-IO which operate on a single RGB image tokenizer, 4M's ability to work with multiple modality-specific tokenizers enables scaling to visual modalities beyond those that can be represented as images, such as neural network feature maps. 4M also builds upon the multimodal masking approach of MultiMAE [5] and extends it beyond image-like modalities.
Both token-based generative models [88, 123, 17, 18, 65] and diffusion models [89, 79, 95, 97] have been mostly limited text-to-image generation. While there are works that enable a greater amount of control by conditioning on additional modalities, they are either very limited in the number of ways they can be conditioned on [58, 40, 113, 119, 62, 9, 128], or require training a separate model for each new modality [130]. 4M flexibly allows for conditioning on any subset of the training modalities, and can, likewise, generate all these modalities, unlocking powerful generative editing capabilities.
7 Conclusion and Limitations
4M is a generalist framework for training multimodal and multitask models that not only perform many key vision tasks out of the box, but also demonstrate strong transfer results to a wide range of downstream tasks. 4M's in-painting and any-to-any generation capabilities enable it to perform a wide range of multimodal generative and expressive editing tasks - all using a single model. In the following, we discuss some limitations of our approach and potential future work.
Additional modalities. While 4M already includes a number of modalities, from semantics-based and geometry-based to text-based, bringing additional modalities could vastly improve its usability. For example, training on features extracted from a large language model has been shown to significantly boost text-to-image generation capabilities [97, 123, 18]. Introducing modalities like edges, sketches, or human poses has the potential to greatly improve the expressiveness [130] of 4M, but it may also be expanded to videos or multi-view imagery to unlock spatial-temporal generation and editing capabilities. We anticipate 4M to conveniently extend to such new modalities.
Tokenizer quality. 4M can benefit from better tokenizers, both in terms of generation and transfer results, but there are limits to the amount of information that can be encoded in tokenized patches. Operating on tokens that cover a smaller image region, or operating on higher-resolution images may improve image quality, but is expected to be computationally more expensive. Generally, improvements along this direction are expected to directly boost the performance of 4M.
Dataset size and quality. While our binding dataset choice of CC12M is standard, 4M can benefit from training on significantly larger datasets [99, 14, 39]. Web-scraped image-text datasets contain many low-quality images, as well as captions that are not related to the image content. Fine-tuning 4M on a more curated dataset like LAION-Aesthetics V2 [99], or tuning the model using reinforcement learning [83] could significantly improve generation quality and diversity. We leave this for the future.
Acknowledgements. We thank Hanlin Goh and Elmira Amirloo Abolfathi for their valuable feedback on earlier versions of this manuscript.
[131] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ADE20K dataset. In Conference on Computer Vision and Pattern Recognition, 2017. 6, 8, 36
[132] Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. iBoT: Image BERT pre-training with online tokenizer. In International Conference on Learning Representations, 2022. 10
[133] Xizhou Zhu, Jinguo Zhu, Hao Li, Xiaoshi Wu, Xiaogang Wang, Hongsheng Li, Xiaohua Wang, and Jifeng Dai. Uni-Perceiver: Pre-training unified architecture for generic perception for zero-shot and few-shot tasks. In Conference on Computer Vision and Pattern Recognition, 2022. 10
A Generative Capabilities & Probing the Learned Representation 20
A.1 Token super-resolution 20
A.2 Generation-specific specializations of 4M 20
A.3 Generation procedure details 21
A.4 Additional visualizations 22
B Multimodal Dataset & Tokenization 30
B.1 Pseudo labeled multimodal training dataset details 30
B.2 Tokenization of captions & bounding boxes 30
B.3 Tokenization of dense modalities 30
C Method & Training Details 32
C.1 Additional 4M architecture details 32
C.2 Training details 32
D Transfer Experiments Details 33
D.1 Architectural differences with the baselines 33
D.2 Image classification on ImageNet-1K 34
D.3 Object detection and instance segmentation on COCO 34
D.4 Semantic segmentation on ADE20K 35
D.5 Depth estimation on NYUv2 35
E Ablation Details & Results 36
E.1 Benchmark tasks 36
E.2 Reference model 37
E.3 Input modalities and target tasks 38
E.4 Multimodal masking strategy 39
E.5 How well does 4M scale? 41
E.6 Architectural design choices 42
E.7 Training design choices 43
E.8 Self-baselines 43
E.9 Comparison to the final models 44
F Additional Evaluations 44
F.1 Out of the box (zero-shot) performance 44
F.2 Text-to-image performance 45
G Broader Impact 46
G.1 Computational costs 46
G.2 Social impact 46
A Generative Capabilities & Probing the Learned Representation
A.1 Token super-resolution
For generating high-resolution images, it is common to first generate images at a lower resolution, followed by a super-resolution step [52, 18, 123]. These two-stage approaches often use a large model to generate a low-resolution image that contains most of the semantically and geometrically important information, followed by one or more smaller super-resolution models that fill in higher resolution details. Directly generating high-resolution images with Transformers is both computationally expensive due to the long sequence lengths, but has also been shown to perform worse at generating semantically coherent images [18] compared to cascade approaches. 4M is trained at a base resolution of  $224 \times 224$  pixels, which corresponds to  $14 \times 14$  tokens. At this low image resolution, fine details are difficult to model. We therefore train a super-resolution model, 4M-SR, starting with 4M-L as the base and fine-tune it to map from low-resolution image tokens to higher-resolution tokens. All images shown are generated at the base resolution using 4M-XL, and subsequently up-sampled to double the resolution using 4M-SR.
Training details. For the high-resolution tokens, each  $448 \times 448$  pixel image is represented by  $28 \times 28$  tokens. The super resolution training-scheme is very similar to the base resolution 4M pre-training scheme, but the inputs to the model consist of a random subset of low-resolution and high-resolution tokens of all modalities, while the target is a random subset of all high-resolution tokens. At inference-time, one or multiple full low-resolution modalities (and sequence-like modalities) are used as conditioning to decode the high-resolution targets step-by-step.
To train 4M-SR, we start from the pre-trained 4M-L and continue training it for an additional 100B tokens (20% of the pre-training length). For each sample, we randomly sample the input budget uniformly between 64 and 1024 to better accommodate for the varying number of tokens throughout the generation process, and keep the target budget fixed at 1024. As full captions are otherwise rarely seen during training, we train on a mixture of masking strategies where 1/4 of the input samples are heavily skewed towards fully unmasked captions ( $\alpha = 5.0$  for captions with  $p_{\mathrm{mask}} = 0$ , and  $\alpha = 0.05$  for all other modalities).
Figure 7: Token super-resolution: We specialize a version of 4M-L for performing super-resolution in the space of tokens ( $14^{2} \rightarrow 28^{2}$  tokens) to double the resolution ( $224^{2} \rightarrow 448^{2}$  pixels) and in the process add more details to the generated images. Just like the base model, it can generate any high-resolution modality conditioned on any low-resolution modality, but we use it to super-resolve modalities we generated at the base resolution. Best viewed when zoomed in.
A.2 Generation-specific specializations of 4M
Adapting 4M to specific use-cases by fine-tuning it either on a subset of pre-training modalities, or fine-tuning it with a different token sampling distribution can improve performance on select downstream tasks. In a similar manner, to optimize a 4M model for text-to-image generation, we can fine-tune it by skewing the token sampling distribution to include more captions as input. To train this caption-skewed model, we start from a pre-trained 4M model and continue training it for an additional 50B tokens (10% of the pre-training length). For each sample in a batch, we randomly sample the input budget uniformly between 64 and 256 to better accommodate for the varying number of tokens throughout the generation process, and keep the target budget fixed at 256. To skew the model towards full captions, we train on a mixture of masking strategies where 1/3 of the input
samples are heavily skewed towards fully unmasked captions ( $\alpha = 5.0$  for captions with  $p_{\mathrm{mask}} = 0$ , and  $\alpha = 0.05$  for all other modalities).
All text-to-image results in this section were generated using a caption-skewed 4M specialization.
A.3 Generation procedure details
Token-based masked image models that were trained with variable masking rates can be directly used for image generation [17, 65, 18]. These models can be seen as order-agnostic autoregressive models [53] in which generation is performed by iteratively decoding tokens in any order. Unlike traditional autoregressive models which need to generate tokens one-by-one in a predetermined order, these models are able to speed up inference by parallelizing the decoding process because the distribution over each masked token can be predicted at the same time.
Generation schedule. The generation process is guided by a generation schedule that outlines a sequence of generation steps. Each step in the sequence determines the target modality, the decoding scheme (described below), the number of tokens to decode, the temperature, and the top-k and top-p sampling values.
Decoding schemes. A crucial feature of our method is the flexibility to incorporate different decoding schemes, such as for different target modalities. We describe these decoding schemes below: