SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- MaskGIT [17] (for image-like modalities). This parallel decoding scheme generates the entire image in a pre-determined number of steps (see Figure 3). At every prediction step, we encode all visible tokens and decode all masked out tokens. We then sample from their predicted distributions and choose the  $n$  most confident tokens, where  $n$  is determined by the generation schedule. Finally, we add the  $n$  predicted tokens to the input and repeat the previous steps.
- Random order autoregressive (ROAR) (for image-like modalities). This scheme is conceptually similar to the decoding scheme proposed in XLNet [120]. Unlike MaskGIT, we do not decode all masked-out tokens at every step, but instead randomly select  $n$  tokens to decode.
- Left-to-right autoregressive (for sequences). Besides images, we are able to generate sequence modalities such as captions and bounding boxes by autoregressively decoding them with the Transformer decoder of 4M (see Figure 3).
Chained generation. Because any generated modality can be used as a conditioning too, we can perform chained generation of several modalities, one after another, with each fully generated one being added to the conditioning of the next (see Figure 3). Performing generation in this chained manner results in each additional modality being generated in a consistent manner, as shown in Figures 9 and 10. In addition, we found that for certain generative tasks, such as caption-to-RGB, generating intermediate modalities such as CLIP tokens can further improve image fidelity (see Figure 9).
Classifier-free guidance. Classifier-free guidance is crucial for improving both image fidelity and how well the generation matches the conditioning. It is most commonly used in diffusion models [50], but can be applied in token-based models as well [40, 123, 18]. We perform classifier-free guidance by computing a weighted combinations of the logits of a forward pass with the conditioning and one without the conditioning:
$$
\log_ {i} \left( \right.\log_ {i} \left( \right.\log_ {i} \left( \right.\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {i} \left(\log_ {1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1\right)\right)\right)\right)\right)\right)\right)\right.\right.\right.
$$
Here,  $w$  is the guidance scale. When performing chained generation, we add each fully generated modality to the set of guided modalities.
Multimodal guidance. While guidance has been shown to significantly improve image quality, it can still happen that generative models ignore parts of the input, unpredictably focus on some parts more than others, or generate undesired concepts. Negative prompting [78] is a popular way of keeping the model from generating undesired concepts. Liu et al. [68] show that performing compositional guidance on multiple conditions can further improve text-image similarity. In a similar way, we can perform compositional generation by weighting different (parts of) modalities by different continuous amounts – even negatively. We can do this by computing a weighted sum of the logits of an unconditional case and the logits of each conditional case:
$$
\log i t s _ {\text {g u i d e d}} = \log i t s _ {\text {u n c o n d}} + \sum_ {i = 1} ^ {n} w _ {i} \left(\log i t s _ {\text {c o n d , i}} - \log i t s _ {\text {u n c o n d}}\right).
$$
Here,  $w_{i}$  are the guidance scales for the different conditions. For example, this allows 4M to generate semantically or geometrically similar variants of images by weakly conditioning on their extracted segmentation, normal, or depth maps (see Figure 13). It can further be used for fine-grained steerability of multimodal edits (see Figs. 16, 14, and 15), or to use negative weighting to avoid certain concepts from being generated (see Figure 17).
A.4 Additional visualizations
$\mathbf{RGB} \rightarrow \mathbf{X}$ . The model can solve several common  $\mathbf{RGB} \rightarrow \mathbf{X}$  vision tasks out of the box, such as predicting normals, depth, semantic segmentation maps, or performing object detection and captioning. In Figure 8, we show examples of this functionality.
Figure 8: RGB  $\rightarrow$  X: 4M can perform several common vision tasks, such as predicting surface normals, depth, semantic segmentation maps, or performing object detection and captioning.
Chained generation. For certain generation tasks, like text-to-image, we observed that generating intermediate modalities, such as CLIP, before generating the RGB image can improve image fidelity. This leads to a form of progressive self-conditioning, which our model particularly enables as it includes several modalities. We demonstrate this phenomenon in Figure 9.
Figure 9: Chained generation: First, we generate intermediate modalities from the caption (CLIP, depth, etc.). Then we generate the RGB images by conditioning on both the caption and the intermediate modality, i.e., caption  $\rightarrow$  intermediate modality  $\rightarrow$  RGB. We compare this to directly predicting the RGB image from the caption, i.e., caption  $\rightarrow$  RGB, in the 'none' column. We observe that certain intermediate modalities like CLIP can improve image fidelity. As 4M includes multiple modalities, it particularly enables this form of progressive self-conditioning and engineering a sequence of modalities for a better final generation.
Any-to-any generative modeling. 4M can be used to generate any modality from any other (and partial) subset of modalities. Figure 10 displays examples where all modalities are generated from a single input, and Figure 11 shows several examples of generating different modalities using two full inputs. Subsequent figures will make use of this functionality and show masked predictions and examples with 1-3 different conditionings.
Figure 10: One-to-all generation: Starting from a single input modality, we use 4M to generate the corresponding outputs for all others. Each row shows the outputs from a different starting modality, all from the same original sample. Chained generation is used to ensure consistent outputs.
Figure 11: Any-to-any generation:  $4\mathrm{M}$  can perform any-to-any generation, meaning that any modality can be generated by conditioning on any subset of modalities.
Conditional in-painting. By virtue of being trained on a masking objective, 4M can be used for conditional and unconditional in-painting (see Figure 12). In-painting coupled with 4M's ability to predict any modality from any other and flexible multimodal conditioning is the basis for several multimodal editing capabilities (see Figure 4 middle and bottom).
Figure 12: Conditional in-painting: 4M can be used for conditional in-painting by masking out a selected image region and generating the remaining tokens using the masked image and a caption (or other modalities) as conditions.
Grounded generation of image variations. 4M can be used to generate different versions of images by grounding the generation in some aspect extracted from a reference input image (see Figure 13). For example, we can take an image and generate variations that respect the geometry of the original image but may contain completely different colors by conditioning only on extracted surface normals. This is conceptually similar to ControlNet [130], but 4M is able to perform this using a single model trained in one single run, and 4M is able to flexibly use any combination of input modalities (see Figure 11). In addition, unlike ControlNet which requires external networks to extract the conditioning modalities, 4M models are able to generate all of them by themselves.
Figure 13: Grounded generation of image variations: 4M can be used to generate variations of a reference image by conditioning on a modality extracted from the original image. This allows for generation of diverse images that are similar geometrically, semantically, or in another (user-specified) aspect.
Multimodal guidance. Multimodal guidance, as explained in Appendix A.3, can be used for various fine-grained editing tasks. We show here several demonstrations of how it can be used to (1) use weak/strong weighting to generate images that loosely/strongly match a certain input modality, and (2) avoid the generation of certain undesired concepts via negative weighting.
Figure 14: Multimodal guidance: By weakly conditioning on geometry (here surface normals) or semantics (here semantic segmentation maps), and strongly conditioning on a caption, we can generate image variations that are only loosely grounded in some aspect of the original image.
Figure 15: Multimodal guidance: Similar to Figure 14, we show here multimodal guidance with captions and either a normal or semantic segmentation map. We keep the weight of the caption fixed and vary the weight of the other modality. The examples demonstrate how a user can effectively control to which degree the generation should use a certain conditioning.
Figure 16: Multimodal guidance: In this example, we weakly condition on a masked RGB image, and strongly on a caption and bounding box. The generations follow the caption and bounding box precisely, while showing some degree of variation for the rest of the generated image.
Figure 17: Multimodal guidance: Multi captioning with negative weighting. (Top of each figure): Standard text-to-image generation can create pleasing results, but the user might want greater control over the generation to counteract dataset biases (e.g., most generated buses are yellow), or may have envisioned a different image. (Bottom of each figure): Enabled by the process described earlier, by adding an additional caption with a negative guidance weight, we can steer which concepts are generated, and which are not.
Probing the learned representation. We can probe what representations 4M models learned either by performing transfers to various downstream tasks or manipulating parts of the input and observe how 4M predicts some readout modality. Concretely, we perform these representation probes by keeping the entire input fixed and manipulating only a single aspect of it.
Figure 18: Probing 4M representations: In this example, we condition the generation on several bounding boxes. For each of the generations inside the blue cells, we keep all bounding boxes fixed, except for one (marked in red shade), which we move around and resize. Depending on where the bounding box is placed, 4M implicitly infers the semantic and geometric properties of the object. For example, placing a bicycle bounding box in front of a bed bounding box generates a real bicycle, while placing it above generates pictures of bicycles.
Figure 19: Probing 4M representations: Here, we condition the generation on segmentation maps in which we manipulate a single segment and keep all others fixed. 4M creates plausible generations, no matter what class the segment was changed to.
B Multimodal Dataset & Tokenization
We pseudo label a large aligned multimodal dataset using CC12M [19] and several pre-trained networks to map RGB images to pseudo labels. To enable training a single Transformer on a large number of vastly different modalities, we resort to tokenization. Concretely, we train modality-specific tokenizers, allowing us to map different modalities to sequences or sets of tokens.
B.1 Pseudo labeled multimodal training dataset details
We trained all 4M models using an aligned multimodal dataset that we created by pseudo labeling CC12M. CC12M already contains aligned RGB images and captions, and we pseudo labeled the remaining tasks using powerful off-the-shelf models.
Surface normals & depth. To introduce geometric priors and to enable generation and transfer learning conditioned on measured or estimated geometric modalities, we pseudo labeled both surface normals and scene depth. We used DPT-Hybrid [91] models trained on Omnidata [32] using cross-task consistency [125] and 3D data augmentations [61] to predict normals and depth from the CC12M RGB images. To encode images using DPT, we resized their heights and widths to the closest multiple of 32, while keeping the maximum side length below 768 to avoid prediction degradation from too large image sizes. We found that even though the Omnidata DPT models were mostly trained on simulated and indoor scenes, they generalize well to a wide range of scenes.
Semantic segmentation. For introducing dense semantic priors, semantic segmentation maps can be useful as it allows for fine-grained creative control when conditioned during generation [40]. We pseudo labeled semantic segmentation labels using a Mask2Former [24] with a SwinB [71] backbone trained on COCO [67] panoptic segmentation. We select the labels by taking the argmax over the predicted semantic classes.
Bounding boxes. To add object bounding boxes to the training data, we use a ViTDet ViT-H model [66] initialized from MAE weights [48] and fine-tuned on COCO [67]. We filter the detected bounding boxes by removing all instances with a confidence score below 0.6.
CLIP feature maps. Tokenized CLIP [85] feature maps have been demonstrated to be powerful targets for masked image models [82, 37, 36]. We use the ViT-B/16 visual backbone of a CLIP-B16 model to extract dense feature maps from its last Transformer layer. To visualize a  $H \times W \times D$  dimensional CLIP feature map, we project all  $H * W$  entries onto the first three principal components computed from said feature map, interpreting the normalized values as R, G, and B colors [126].
B.2 Tokenization of captions & bounding boxes
We follow Pix2Seq [21, 22] and treat detection as a sequence prediction problem. For example, a scene containing a bounding box around a cat and a plant could be parameterized as xmin=0.15 ymin=0.3 xmax=0.65 ymax=0.5 cat xmin=0.75 ymax=0.3 xmax=0.95 ymax=0.8 potted plant [EOS]. In practice, we model the corner coordinates (i.e. the minimum and maximum  $x$  and  $y$  coordinates) using a resolution of 1000 special tokens per each of those cases - namely we represent bounding boxes using xmin=0, ..., xmin=999, ymin=0, ..., ymin=999, xmax=0, ..., xmax=999, ymax=0, ..., ymax=999. In contrast to Pix2Seq, our method involves masking parts of the object sequence, which requires two modifications to the sequence construction pipeline to make this task more manageable: First, we order the objects in the sequence by their distance to the origin. Second, unlike in Pix2Seq where corner coordinates share tokens, we assign separate tokens for each of the corner coordinates, resulting in a vocabulary size 4 times larger than in Pix2Seq. Assigning separate tokens for each corner coordinate ensures that the meaning of each token is not made ambiguous when surrounding coordinates are masked out.
We jointly fit a WordPiece [30] text tokenizer on all captions and the 4000 special bounding box tokens. COCO class labels for the bounding boxes are treated as special tokens, meaning that if they appear in any caption, they get mapped to the same tokens. The joint text and bounding box vocabulary has a size of  $30\mathrm{K}$ .
B.3 Tokenization of dense modalities
We tokenize all dense modalities using variants of vector-quantized autoencoders (VQ-VAEs) [110]. To avoid blurry and unrealistic images, we train the tokenizers for RGB, normals and depth using diffusion decoders. All tokenizers are first trained on 100 epochs of an ImageNet-1K version of the pseudo labeled dataset (see Appendix B.1), after which they continued training for 15 epochs on the
pseudo labeled CC12M dataset. We follow up the training at the base resolution of  $224^2$  with a short multi-resolution fine-tuning step to adapt the tokenizers up to resolution  $448^2$ .
The base resolution training specifics of our tokenizers are shown in Table 3. We employ several common practices: Following Yu et al. [121], we heavily reduce the dimensionality of the codebook and  $l_{2}$ -normalize the codes and encoded vectors to improve codebook utilization and reconstruction quality. In addition, to avoid the joint VQ-VAE and diffusion training to collapse to degenerate solution, we found it to be crucial to restart stale codebook entries, similar to Zeghidour et al. [127]. To that end, we count the number of encoded vectors in a batch that map to a given codebook entry after every iteration, and replace (randomly from the batch) any codes that have an exponential moving average (EMA) count less than a specified threshold threshreplace. Because this value depends on the total batch size  $B$ , number of tokens per image  $N_{\mathrm{tokens}}$ , and codebook vocabulary size  $N_{\mathrm{vocab}}$ , we specify this EMA threshold using a coefficient  $c_{\mathrm{replace}}$  by:
$$
\operatorname {t h r e s h} _ {\text {r e p l a c e}} = \frac {B N _ {\text {t o k e n s}}}{c _ {\text {r e p l a c e}} N _ {\text {v o c a b}}}
$$
Table 3: Tokenizer training settings. Base resolution  $(224^{2})$  training configuration for the tokenizers used for training 4M. The base learning rate is specified for batch size 256.
Configuration	RGB	Normals	Depth	Segmentation	CLIP
Codebook size	16384	8192	8192	4096	8192
Patch size	16 × 16
Latent dimension	32
EMA stale code coef. creplace	32
Codebook EMA	0.99
l2-normalized codes [121]	✓
Codebook weight	1.0
Commitment weight	1.0
Encoder architecture	ViT-B/16
Decoder architecture	Patch-UNet	Patch-UNet	Patch-UNet	ViT-B/16	ViT-B/16
Diffusion decoder	✓	✓	✓	×	×
Loss function	MSE	MSE	Smooth L1	Cross-entropy	Smooth L1
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.95
Weight decay	0.05
Base learning rate [46]	5e-5
Batch size	640	640	640	640	1024
Learning rate sched.	Cosine decay
Training epochs	100 (IN1K) + 15 (CC12M)
Warmup epochs	5 (IN1K) + 1 (CC12M)
Resolution	2242
Random crop scale	(0.8, 1.0)	(0.8, 1.0)	(0.8, 1.0)	(0.8, 1.0)	(0.2, 1.0)
Random crop ratio	(0.75, 1.3333)
Horizontal flip	✓
Data type	float16	float16	float16	bfloat16	bfloat16
Semantic segmentation & CLIP feature maps. We tokenize segmentation maps and CLIP feature maps using ViT-B [31] encoders and decoders (without the patch projection on the latter). The reconstruction loss for segmentation maps is cross-entropy, while for CLIP we choose the smooth L1 loss. While we do not use any data augmentations for training 4M, we decided to train the CLIP tokenizer using a minimum crop scale of 0.2 to keep open the option of performing fine-tuning with CLIP targets and strong augmentations [82]. For all other tokenizers, we use a more moderate minimum crop scale of 0.8, to avoid training them on very low-resolution but upscaled images that could affect image fidelity negatively.
RGB, normals & depth. Training a VQ-VAE by simply minimizing a reconstruction loss can lead to blurry and unrealistic looking results [88], which is why VQ-GAN [35] proposes to use a discriminator to get more realistic results. We found VQ-GAN training to be unstable and decided to instead train a diffusion model as the decoder, similar to DiVAE [101]. Unlike DiVAE, which trains a diffusion model decoder using a frozen and pre-trained VQ-GAN encoder, we train both the encoder and diffusion decoder end-to-end and from scratch. As the decoder, we use a UNet [80] with four down and up layers, each consisting of three residual blocks and one up/down sampling block, with self-attention blocks at the two smallest down and up layers.
For improved training and inference efficiency, we process  $C$ -channel image patches of size  $C \times 4 \times 4$ , rather than individual pixels, similar to Patched Diffusion Models [75]. We condition the diffusion decoder by concatenating the 32-dimensional codebook entries of the  $14 \times 14$  encoded tokens with the noised input after patching, i.e. we concatenate the noised and patched image of shape  $16C \times 56 \times 56$  with the upsampled latent tensor of shape  $32 \times 56 \times 56$ . The output of the UNet gets reshaped back to  $C \times 224 \times 224$  before computing the loss.
We found that predicting the noise leads to undesirable color shifts during inference, which is why we predict the clean image instead. In addition, we found that to avoid the training to collapse to a degenerate solution, it is crucial to restart dead / unused codebook entries [127]. We train the diffusion decoder using 1000 DDPM [51] steps and a linear noise schedule. At inference time, we sample from the decoder using 25 DDIM [105] steps.
Multi-resolution training. To train super-resolution specializations of 4M (see Appendix A.1), we need to be able to use the tokenizers both at the base resolution of  $224 \times 224$ , and at our chosen higher resolution of  $448 \times 448$ . Unlike convolutional VQ-VAEs, however, our ViT based tokenizers do not perform well on resolutions different from the training resolution. For that reason, we perform multi-resolution training as a short fine-tuning step, and initialize from the weights trained at the base resolution. During training, the resolution of every batch is sampled randomly between 224 and 448, in increments of 32. All tokenizers were fine-tuned for one CC12M epoch, except for RGB, which was fine-tuned for five epochs. We lowered the batch size per GPU to 16, except for semantic segmentation, for which we lowered it to 10. The remaining settings are the same as in Table 3.
C Method & Training Details
C.1 Additional 4M architecture details
For both the ablations and the longer training runs we train 4M models of different sizes. Table 4 lists their respective encoder and decoder depths, model dimension, number of attention heads, as well as the number of trainable parameters (excluding the embedding layers).
**Embedding details.** Modality embeddings are shared between the encoder and decoder. Additionally, as is commonly done for autoregressive models, the parameters of the input and output embeddings (i.e., the final linear layer) of the decoder are shared for sequence modalities. Similar to MAE [48] and MultiMAE [5], the positional and modality embeddings are re-injected into the encoder output before it is passed on to the decoder.
RGB pixels as input. All 4M models are trained to accept both RGB pixels and tokenized RGB as input. This is implemented by treating RGB pixels and tokenized RGB as two entirely separate modalities, with RGB pixels being an input-only modality (as it isn't tokenized). This simple approach makes it easy to incorporate additional non-tokenized modalities as input in the future.
Span masking for sequences. We perform span masking [86] as follows: Given the probability of masking a token  $(p_{\mathrm{mask}})$ , we randomly mask out tokens in the sequence and replace each consecutive span of masked-out tokens by a sentinel token (e.g., [S_1], [S_2], [S_3],....). The target sequence then consists of the masked-out spans delimited by the sentinel tokens, followed by a final sentinel token to signal the end of the sequence.
Unlike dense modalities, it is not possible to strictly respect the token budget when sampling from sequences due their variable length. Instead, we treat the token budget as a strict upper bound and mask sequences as follows: For the input, we sample a masking probability  $p_{\mathrm{mask}}$  from a uniform distribution and use it for span masking. If the sequence length after masking is greater than the input budget, we progressively increase  $p_{\mathrm{mask}}$  until it fits within the assigned budget. For the target, if the sequence does not fit within the budget, we randomly truncate it while also ensuring that the first token of the truncated sequence is a sentinel token.
C.2 Training details
The training details for the 4M models used for the transfer experiments (Section 3) and the generation results (Section 4) are shown in Table 5. These hyperparameters were chosen using insights from the ablation study (Section 5), more detailed ablation results are shown in Appendix E. All 4M models were trained using bfloat16 mixed-precision [13]. Training durations varied with model size; 4M-B was trained in 1.5 days on 64 A100 GPUs, 4M-L was trained in 3 days on 128 A100 GPUs, while training 4M-XL took 8 days on 128 A100 GPUs. For 4M-XL, we reduced GPU memory