SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Target masking budget. Just like how we could improve the training efficiency by only decoding a small set of visible tokens, so can we decide to decode merely a subset of all remaining masked-out tokens. Decoding all masked-out tokens, like MAE [48], or in the multimodal case MultiMAE [5], can quickly become infeasible as the number of modalities grows and the masking ratio is kept high. While they attempt to address this issue by decreasing the size of the decoder, such a strategy might ultimately hurt generation quality if the decoder is used for generative tasks. As Table 19 shows, decoding only a small random subset of all targets performs better (for a fixed training duration) than decoding a larger number of tokens, while also significantly reducing the computational costs.
Table 19: Target masking budget ablation: Similar to passing only visible tokens to the encoder, we can decide to only decode a subset of masked-out tokens, which improves computational efficiency considerably. We fix the input masking budget at 128 and vary the number of target tokens.
Target Tokens	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
64	3.06	6.11	5.09	6.73	5.79	5.36	1.40
▷ 128	3.06	6.11	5.07	6.75	5.80	5.36	1.60
256	3.10	6.16	5.15	6.81	5.91	5.43	3.00
512	3.16	6.20	5.27	6.84	5.95	5.48	4.00
Mixture of masking strategies. As shown in Appendix E.3, pre-training using RGB as the sole input modality instead of training on all of them performs significantly better at transfers that, likewise, use RGB as input. On the flip-side, pre-training with all modalities performs significantly better when transferring to unseen input modalities. This difference can be explained by the fact that these two mask sampling strategies are on opposite extremes – the number of RGB inputs in the multimodal pre-training case is relatively small, while the RGB-only case is never trained on any other modalities. We can find a compromise by sampling batch elements using these two strategies uniformly at random, creating a pre-training strategy, where approximately half the time input tokens are RGB-only, and half the time they are sampled from all modalities. Table 20 shows that while the mixture approach does not perform better than either individual mask sampling strategy at any of the transfers, it is a good compromise if the use-case of the model is open at the time of pre-training.
Table 20: Mixture of masking strategies ablation: Mixing different masking strategies by randomly sampling between the two provides a good compromise in downstream performance.
Masking strategy	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
RGB → All	2.90	5.99	4.74	6.91	5.93	5.29	1.80
▷ All → All	3.06	6.11	5.07	6.75	5.80	5.36	2.20
RGB → All & All → All	2.95	6.00	4.86	6.86	5.87	5.31	2.00
E.5 How well does 4M scale?
We argue that scalability is a key property that models and training objectives should have. To ablate how well 4M scales, we ablate the following three axes: dataset size, training length, and model size.
Does 4M benefit from more data? We train the base configuration on various subsets of CC12M, down to 1/64th of the full dataset. The results in Table 21 show that 4M is able to scale with the pre-training dataset size.
Table 21: Dataset size ablation: To estimate how well 4M scales with pre-training dataset size, we train it on different subsets of CC12M ranging from  $\frac{1}{64}$  of the dataset to the entire dataset..
Dataset fraction	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
1/64	3.22	6.23	5.31	6.86	5.94	5.51	4.00
1/16	3.15	6.17	5.23	6.80	5.87	5.44	3.00
1/4	3.08	6.13	5.11	6.79	5.84	5.39	2.00
▷ 1	3.06	6.11	5.07	6.75	5.80	5.36	1.00
Does 4M benefit from longer training? We train the base configuration for different amounts of total tokens seen. We define tokens seen as the total number of both input tokens and target tokens the model was trained on. Table 22 shows that 4M scales to longer training schedules.
Table 22: Training duration ablation: To see whether 4M benefits from longer training schedules, we train it for up to 400B tokens seen.
Train tokens [B]	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
50	3.13	6.17	5.23	6.86	5.93	5.46	4.00
▷ 100	3.06	6.11	5.07	6.75	5.80	5.36	3.00
200	3.00	6.05	4.93	6.71	5.74	5.29	2.00
400	2.96	6.02	4.84	6.70	5.71	5.25	1.00
Does 4M scale with model size? We train 4M models of different sizes, ranging from Tiny (4M-Ti) to Large variants (4M-L). Exact model specifications are given Table 4. Table 23 shows that 4M scales with model size.
Table 23: Model size ablation: To see how 4M scales with model size, we train Tiny, Small, Base and Large versions.
Model size	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
4M-Ti	3.28	6.38	5.57	6.93	6.15	5.66	4.00
4M-S	3.16	6.25	5.34	6.84	5.97	5.51	3.00
▷ 4M-B	3.06	6.11	5.07	6.75	5.80	5.36	2.00
4M-L	2.93	5.97	4.86	6.63	5.75	5.23	1.00
E.6 Architectural design choices
In the following, we ablate several design choices related to the Transformer architecture.
Number of encoder & decoder layers. We ablate whether one should train 4M using a balanced allocation of encoder and decoder layers or rather go with a deeper or more shallow decoder. We fix the total number of Transformer layers to 24 and ablate three different configurations in Table 24. The results show that all configurations perform comparably, with the balanced and encoder-heavy settings having a slight edge over the decoder-heavy one.
Table 24: Encoder & decoder depth ablation: We ablate three different ways of allocating 24 Transformer layers between the encoder and decoder.
Enc. depth	Dec. depth	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
8	16	3.10	6.14	5.09	6.75	5.82	5.38	2.80
▷ 12	12	3.06	6.11	5.07	6.75	5.80	5.36	1.40
16	8	3.05	6.11	5.11	6.75	5.78	5.36	1.80
Per-token vs. per-modality loss. Since we predict multiple modalities at the same time, there are several ways we can compute their losses. We ablate the following two loss weighting strategies: (1) We treat every predicted token the same and average all their losses. We call this setting per-token loss. This setting is biased against modalities that don't contain a lot of tokens, such as captions. (2) We first average the loss for every target modality individually and then average those. We call this setting per-modality loss. Computing the loss per-modality noticeably outperforms the per-token loss, as shown in Table 25.
Table 25: Multitask loss aggregation ablation: We ablate weighting the loss contribution of every modality equally vs. weighting every token equally.
Loss type	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
▷Per-modality loss	3.06	6.11	5.07	6.75	5.80	5.36	1.00
Per-token loss	3.07	6.16	5.12	6.78	5.86	5.40	2.00
Transformer architecture modifications. We ablate several modifications to the Transformer architecture, and show results in Table 26. (1) Following Shazeer [100], we ablate the use of the SwiGLU activation function in the feed-forward layers, and reduce the feed-forward dimension by a factor of  $\frac{2}{3}$  (from  $4d$  to  $\frac{2}{3}4d$ ) to keep the parameter count and amount of computation constant. (2) Following T5 [86] and PaLM [25], we ablate removing all bias terms from the Transformer. (3) Following ViT-22B [28], we ablate the use of query/key normalization, which has been shown to help the stability of very large Vision Transformers trained to perform image classification.
We find SwiGLU slightly outperforms GELU while also removing the need for bias terms, and therefore use it to train the final model. As we do not observe noticeable stability issues with our training objective with current model scales, we refrain from using query/key normalization as we find that it negatively impacts the transfer performance.
Table 26: Architecture modifications ablation: We ablate the use of different activations in the feed-forward [49, 100], removing bias terms [86, 25], and performing query-key normalization [28].
FFN act.	Bias	QK Norm.	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
▷ GELU	✓	X	3.06	6.11	5.07	6.75	5.80	5.36	2.20
GELU	X	X	3.07	6.11	5.11	6.80	5.86	5.39	4.00
GELU	X	✓	3.08	6.13	5.16	6.92	5.96	5.45	5.20
SwiGLU	✓	X	3.08	6.13	5.16	6.94	6.26	5.51	5.80
SwiGLU	X	X	3.07	6.11	5.03	6.73	5.79	5.34	1.80
SwiGLU	X	✓	3.06	6.09	5.03	6.79	5.80	5.35	2.00
E.7 Training design choices
Base learning rate. We ablate several base learning rates [46] and show results in Table 27.
Table 27: Base learning rate ablation. We ablate the choice of learning rate used during training. 4M is not too sensitive to the choice of learning rate, although we observe some performance degradation and instabilities when strongly increasing the learning rate.
Base Ir [46]	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
5e-5	3.06	6.11	5.07	6.78	5.82	5.37	2.80
1e-4	3.06	6.11	5.07	6.75	5.80	5.36	1.20
2e-4	3.05	6.11	5.10	6.75	5.86	5.37	2.20
3e-4	3.07	6.11	5.11	6.77	5.88	5.39	3.80
5e-4	3.10	6.14	5.17	6.81	5.89	5.42	5.00
Training with repeated sampling. Data loading can be a significant bottleneck when training efficient masked models. For that reason, we use webdataset [116] to load tar files consisting of 1000 samples instead of loading individual samples. We keep a buffer of loaded images in RAM and randomly sample from it repeatedly – each time applying different random masking. Whenever an element in the buffer has been sampled more than the specified number of repeats, we replace it with a fresh sample. Using repeated sampling [38] has the potential to significantly improve training efficiency at no loss in performance (see Table 28).
Table 28: Repeated sampling ablation: To improve training efficiency, we reuse samples from a buffer multiple times, at no significant loss in performance.
Num. repeats	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
1	3.06	6.11	5.06	6.73	5.80	5.35	1.40
▷ 4	3.06	6.11	5.07	6.75	5.80	5.36	1.60
8	3.07	6.12	5.10	6.76	5.81	5.37	3.00
E.8 Self-baselines
When comparing 4M to other pre-training methods in the transfer learning study, we also report the performance of two self-baselines to better control for the dataset, augmentations, model architecture, compute, and use of tokenizers, all of which can significantly affect downstream performance. These two self-baselines were chosen to be conceptually similar to MAE [48] (Masked RGB  $\rightarrow$  RGB) and BEiT-v2 [82] (Masked RGB  $\rightarrow$  CLIP). Following these methods, we train these self-baselines with a relatively high masking ratio (in our case,  $50\%$ ) and ensure that the target tokens do not spatially overlap with the inputs. In this section, we ablate the masking strategy for these self-baselines by also training a version with a lower masking ratio and with spatial overlap between the input and the targets (as in 4M), and a version without any masking and where all targets are predicted (standard RGB  $\rightarrow$  X training).
Results are shown in Table 29. We find that while the RGB  $\rightarrow$  RGB setting does benefit from higher masking ratios and no spatial overlap (as shown in MAE), RGB  $\rightarrow$  CLIP benefits from a lower masking ratio and from allowing targets to spatially overlap with the input.
Table 29: Self-baseline ablation: The RGB  $\rightarrow$  RGB self-baseline performs best with a higher masking ratio and no spatial overlap between inputs and targets, while RGB  $\rightarrow$  CLIP benefits from a lower masking ratio and spatial overlap.
Target mod.	Input tok.	Target tok.	Spatial overlap	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
RGB	98	98	X	3.12	6.16	5.01	6.91	6.14	5.47	1.40
RGB	128	128	✓	3.14	6.21	5.03	6.94	6.15	5.49	2.40
RGB	196	196	✓	3.24	6.37	5.42	6.87	6.13	5.61	2.20
CLIP	98	98	X	3.08	6.09	4.93	6.94	6.03	5.41	2.40
CLIP	128	128	✓	3.07	6.11	4.83	6.85	5.94	5.36	1.20
CLIP	196	196	✓	3.12	6.20	4.93	6.88	6.01	5.43	2.40
E.9 Comparison to the final models
We gather insights from the ablation to train the "final" 4M models described in Appendix C. Here, we compare these models to the reference models used throughout the ablation. Unsurprisingly, we find that these models perform significantly better than the reference models on the benchmark tasks. The large performance gap for tasks that take as input RGB images can in large part be attributed to the mixture of masking strategies described in Appendix E.4.
Table 30: Comparison to the final models: The final models obtained by combining insights from the ablations significantly outperform the reference models on all benchmark tasks.
Size	Setting	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss
▷ Base	Reference (Table 14)	3.06	6.11	5.07	6.75	5.80	5.36
Base	Final (Table 5)	2.85	5.91	4.69	6.74	5.76	5.19
Large	Reference (Table 14)	2.93	5.97	4.86	6.63	5.75	5.23
Large	Final (Table 5)	2.68	5.76	4.48	6.67	5.63	5.04
F Additional Evaluations
F.1 Out of the box (zero-shot) performance
In this section, we evaluate the out of the box capabilities of 4M models in three RGB  $\rightarrow$  X tasks: surface normals estimation, depth estimation, and semantic segmentation. The results are shown in Table 31. Our findings demonstrate that 4M performs competitively in these tasks even without fine-tuning, when compared to both task-specific baselines and the pseudo labelers used for 4M training. Additionally, the tokenization quality, as shown in Table 31 and Figure 20, does not seem to limit the overall performance of 4M in these zero-shot scenarios. We therefore anticipate that the zero-shot performance of 4M can be further improved if not limited by the current pseudo labels, either through the use of ground truth data or stronger pseudo labeling networks.
Table 31: Out of the box (zero-shot) performance: We report the out of the box performance of 4M models and of baselines on several tasks, all at  $224 \times 224$  resolution. We use the DIODE [111] validation set for normals and depth, and COCO validation set for semantic segmentation. [P] denotes the pseudo labeler used to train 4M. Best results are bolded, second best are underlined. 4M's zero-shot performance matches or outperforms strong baselines such as OASIS [23] and MiDaS [90], and is competitive with pseudo labelers. * To estimate the performance upper bound of 4M due to tokenization, we also report the tokenizer reconstruction quality on validation images from CC12M.
Method	Surface Normals (mean angle error)↓	Depth (standardized L1 error)↓	Semantic Segmentation (mean IoU)↑
OASIS	34.3	-	-
MiDaS DPT Hybrid	-	0.73	-
Omnidata+3DCC [P]	22.5	0.68	-
Mask2Former Swin-S	-	-	44.6
Mask2Former Swin-B [P]	-	-	45.7
Mask2Former Swin-L	-	-	48.0
4M-B	21.9	0.71	41.3
4M-L	21.4	0.69	45.8
4M-XL	20.8	0.68	46.5
Tokenizer reconstruction*	4.0	0.06	90.5
Figure 20: Tokenizer reconstructions: We show sample reconstructions of surface normals, depth, and semantic segmentation pseudo labels from the CC12M validation set at  $224 \times 224$  resolution ( $14 \times 14$  tokens). Quantitative evaluations are provided in Table 31 (last row), confirming that the reconstruction quality does not bottleneck  $4\mathrm{M}$ 's out of the box performance.
F.2 Text-to-image performance
In this section, we evaluate 4M's text-to-image capabilities, with detailed quantitative results shown in Table 32. To perform a controlled comparison, we train a pure text-to-image variant of 4M-B, akin to Muse [18], for a total of 300B tokens on CC12M, and using the same RGB tokenizer as the 4M models. 4M trained on all modalities achieves comparable FID and CLIP scores to this specialist model, and at the same time can be conditioned on any pre-training modality and can solve several common vision tasks out of the box. We also compare against the  $512 \times 512$  base model of Stable Diffusion 2.1 (SD-2.1) [95] and observe a considerable gap to SOTA generative models on OOD data. It is important to highlight that SD-2.1 was trained on datasets that were two orders of magnitude larger and with a compute budget one order of magnitude greater than 4M-XL. When taking into account the scaling behaviors of comparable token-based text-to-image models such as Muse or MAGE, we anticipate a substantial improvement in the generation quality of 4M models if they were trained under a similar data and computational regime, alongside the use of improved image tokenizers.
Table 32: Quantitative evaluation of text-to-image generative capabilities: We evaluate the text-to-image capabilities of 4M and compare them with a controlled text-to-image version of 4M (similar to Muse), and Stable Diffusion 2.1-base. We compute FID and CLIP-L/14 scores on 30K images of the CC12M and COCO validation sets after resizing the generations to 256x256. All runs use guidance scale 3.0. In-domain, 4M can approach SD-2.1 and matches the Muse baseline in a controlled setting. On COCO, we observe a larger gap to SD-2.1 which was trained on orders of magnitude more data and compute.
Model	Res.	Dataset	A100-hours	Text encoder	CC12M val (30K)	COCO val (30K)
FID ↓	CLIP score ↑	FID ↓	CLIP score ↑
4M-B	224	CC12M	2.3k	-	15.9	20.6	37.5	21.4
4M-L	224	CC12M	9.2k	-	11.9	18.9	30.1	23.1
4M-XL	224	CC12M	24.5k	-	10.7	22.1	27.0	23.7
4M-B "Muse"	224	CC12M	1.6k	-	18.3	18.8	39.8	20.3
SD-2.1-base	512	Curated LAION-5B	>200k	CLIP ViT-L/14	9.1	25.5	10.1	25.5
G Broader Impact
G.1 Computational costs
The training time of the 4M models used in the transfer and generative results depends on their size: 4M-B was trained in 1.5 days on 64 A100 GPUs, 4M-L was trained in 3 days on 128 A100 GPUs, and training 4M-XL took 8 days on 128 A100 GPUs.
The reference model used in the ablations can be trained in 12 hours on 32 A100 GPUs, or 2 days on 8 A100 GPUs. Transferring a pre-trained model to the 35 benchmark tasks used for the ablations takes approximately 3 days on 8 A100 GPUs.
G.2 Social impact
We developed a framework for training general-purpose foundation models that, as demonstrated, can be conveniently re-purposed for various tasks a practitioner may be interested in. We also committed to open-sourcing our code and models. These actions support the public with the democratization of the tools and the possibility of transparent inspection and safeguarding. While our model is not particularly poised for negative use compared to the alternatives, it should be noted that powerful generative models are a general tool and have the potential to be used in ways authors did not intent. In addition, the data they are trained on may incorporate various societal biases or contain samples gathered in different ways from the internet. We trained our models on CC12M [19] which is an open-sourced dataset and has been curated to some degree (e.g., people's names are redacted), yet, due to the imperfections in this process, we still advise caution when using the models for generative purposes.