SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 4: Model size variants: The model sizes and naming scheme follow T5 [86]. We exclude the size of the embedding layers in the parameter count.
Model	Encoder Blocks	Decoder Blocks	Model Dim	Num Heads	Total Params
4M-Ti	6	6	384	6	24M
4M-S	8	8	512	8	59M
4M-B	12	12	768	12	198M
4M-L	24	24	1024	16	705M
4M-XL	24	24	2048	32	2818M
consumption through the use of activation checkpointing along with optimizer state and gradient sharding (ZeRO-2 [87]), via PyTorch's Fully Sharded Data Parallel (FSDP).
Table 5: Pre-training settings. Training configuration for the 4M models used in the transfer experiments and generation results.
Configuration	4M-B	4M-L	4M-XL
Training length (n tokens)	500B
Warmup length (n tokens)	10B
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.95
Base learning rate [46]	1e-4	1e-4	2e-5
Batch size	8192
Weight decay	0.05
Gradient clipping	X	X	3.0
Learning rate schedule	Cosine decay
Feedforward activation (Appendix E.6)	SwiGLU [100]
Input token budget (Appendix E.4)	128
Target token budget (Appendix E.4)	128
Input and target α (Appendix E.4)	0.5, 0.5
Masking strategy (Appendix E.4)	RGB → All & All → All
Image resolution	2242
Augmentation	None (Center Crop)
Repeated sampling [38] (Appendix E.7)	4
Data type	float16 [13]
D Transfer Experiments Details
In this section, we report the fine-tuning details for all transfers shown in Section 3. Whenever possible, we follow commonly used settings from other papers performing such transfers. However, to avoid excessive computational costs, we adjust some transfer settings such that no single transfer requires more than  $20\%$  of the pre-training compute. These adjustments include reducing the training resolution when fine-tuning on COCO, and reducing the number of epochs for intermediate fine-tuning on ImageNet-21K.
D.1 Architectural differences with the baselines
4M's encoder and all baselines can be viewed as ViT-Base or ViT-Large backbones [31], and are nearly identical in terms of parameter count and model FLOPS. Despite this, a few minor architectural differences could potentially affect their performance. These include: (1) 4M lacks bias terms and uses SwiGLU [100] in the feedforward, while the baselines include bias terms and use GELU [49]. (2) All baselines include an additional learnable [CLS] token as input, which is absent in 4M. (3) DeiT III uses learnable positional embeddings, whereas all other methods including 4M use fixed sine-cosine positional embeddings. (4) When fine-tuning on COCO, MAE uses both absolute and relative positional embeddings (following ViTDet's approach [66]), while all other methods only use absolute positional embeddings.
D.2 Image classification on ImageNet-1K
We use settings inspired by [114] and first perform intermediate fine-tuning on ImageNet-21K [93], followed by full fine-tuning on ImageNet-1K [29]. Intermediate fine-tuning helps improve overall performance for models that were not originally trained on ImageNet-21K (all 4M models and baselines aside from DeiT III). Details are shown in Table 6.
Table 6: Image classification settings. Configuration for intermediate fine-tuning on ImageNet-21K and fine-tuning on ImageNet-1K.
Configuration	ImageNet-21K	ImageNet-1K
Base	Large	Base	Large
Fine-tuning epochs	20	50	20
Warmup epochs	2	2
Optimizer	AdamW [73]	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.95	β1, β2 = 0.9, 0.999
Base learning rate [46]	1e-4	1e-4
Batch size	4096	4096
Weight decay	0.05	0.05
Learning rate schedule	Cosine decay	Cosine decay
Layer-wise lr decay [26]	0.75	0.85	0.75	0.85
Drop path [55]	0.1	0.2	0.1	0.2
Input resolution	2242	2242
Augmentation	RandAug(9, 0.5) [27]	RandAug(9, 0.5) [27]
Random resized crop	(0.5, 1)	(0.08, 1)
Label smoothing ε	0.1	0.1
Mixup [129]	0.1	0.1
Cutmix [124]	1.0	1.0
D.3 Object detection and instance segmentation on COCO
We follow the settings from ViTDet [66] with a simplified Cascade Mask-RCNN head [47, 15] with two major changes: First, we reduce the image resolution from  $1024 \times 1024$  to  $512 \times 512$ , which reduces computational costs during training by over  $4 \times$  but lowers the performance. Second, we do not use window attention, and instead keep the attention layers unchanged (i.e. global) for all models and rely on activation checkpointing to reduce the GPU memory usage. Details are shown in Table 7.
Table 7: Object detection and instance segmentation settings. Configuration for object detection and instance segmentation fine-tuning on COCO, the settings follow ViTDet [66].
Configuration	COCO
Base	Large
Fine-tuning epochs	100
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Weight decay	0.1
Learning rate	0.0001
Learning rate schedule	Multi-step decay
Lr schedule milestones	[Epoch 89, Epoch 96]
Lr schedule decay values	[1.0, 0.1, 0.01]
Warmup epochs	0.01	0.007
Batch size	512	320
Layer-wise lr decay [26]	0.7	0.8
Drop path [55]	0.1	0.4
Input resolution	5122
Augmentation	Large-scale jitter (LSJ) [41]
D.4 Semantic segmentation on ADE20K
We follow the settings from MultiMAE [5] and use a ConvNext [72] prediction head of depth 4 on top of the encoder to perform semantic segmentation. Details are shown in Table 8.
Table 8: Semantic segmentation settings. Configuration for semantic segmentation fine-tuning on ADE20K, the settings follow MultiMAE [5].
Configuration	ADE20K
Base	Large
Fine-tuning epochs	64
Warmup epochs	1
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Learning rate	2e-4
Batch size	64
Weight decay	0.05
Learning rate schedule	Cosine decay
Layer-wise lr decay [26]	0.75	0.85
Drop path [55]	0.1	0.2
Input resolution	5122
Augmentation	Large-scale jitter (LSJ) [41]
Color jitter	✓
D.5 Depth estimation on NYUv2
We closely follow the settings from MultiMAE [5] but replace the DPT [91] prediction head by a ConvNeXt [72] prediction of depth 2. We find that this ConvNeXt head reaches comparable performance to DPT on the MultiMAE baseline, while also removing the need to extract and rescale features from intermediate layers of the network. Details are shown in Table 9.
Table 9: Depth estimation settings. Configuration for depth estimation fine-tuning on NYUv2, the settings follow MultiMAE [5].
Configuration	NYUv2
Base	Large
Fine-tuning epochs	1000
Warmup epochs	100
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Learning rate	1e-4
Batch size	128
Weight decay	1e-4
Learning rate schedule	Cosine decay
Layer-wise lr decay [26]	0.75	0.85
Drop path [55]	0.1	0.2
Input resolution	2562
Random crop	✓
Color jitter	✓
E Ablation Details & Results
E.1 Benchmark tasks
All transfer tasks in the ablation are cast as token-to-token prediction problems, similar to the way 4M was pre-trained. We perform several transfers with RGB (pixels) as the input, test how well 4M variants perform at adapting to unseen input modalities, and ablate their ability to make use of multiple input modalities.
All models are fine-tuned with a constant learning rate and evaluated using an exponential moving average of the weights. As different models may overfit to the transfer tasks at different speeds, we measure the validation loss 50 times per transfer and report the best loss, as in T5 [86].
For all tasks and modalities, we train new tokenizers on the respective training sets, which means that they are all either completely unseen, or a new instantiation of a training task. When transferring from an unseen modality (e.g., in  $\mathrm{X} \to \mathrm{Y}$  and  $\mathrm{X} + \mathrm{Y} \to \mathrm{Z}$  transfers), the input embeddings are not initialized, which can lead to instabilities and lower the overall performance. To address this problem, we initially freeze the Transformer encoder-decoder and only train the new input and output embedding layers. Afterwards, we unfreeze the Transformer and train the entire model.
As described in Section 5, we report the validation set cross-entropy performance instead of task-specific metrics. The reasons for this are twofold: First, downstream performance hinges on A) how well the tokens are able to represent the downstream tasks (i.e. tokenizer reconstruction error), and B) how well 4M is able to predict these tokens (i.e. cross-entropy loss). Since the tokenizers are the same across all settings, we abstract away aspect A for this ablation and only report B. Second, reporting the cross-entropy loss makes comparisons more uniform and avoids having to scale and average wildly different task-specific metrics such as mAP, mIoU, MSE, etc. that are not comparable with each other.
COCO transfers. We perform object detection on the COCO dataset [67] in the same manner as during pre-training (see Appendix B.2), i.e. by casting it as a sequence prediction problem [21]. The input modality for the transfer task is RGB (pixels). Exact training settings are provided in Table 10.
ADE20K transfers. We perform semantic segmentation on ADE20K [131] by fine-tuning the 4M models to predict VQ-VAE tokens of the targets. Since our semantic segmentation tokenizer during pre-training was trained with COCO classes, we train a new tokenizer for ADE20K in the same manner (see Appendix B.2). The tokenizer is trained for 1500 epochs with a warmup period of 150 epochs. In order to perform random crop augmentations at transfer time, we also pre-trained the tokenizer with a crop scale of (0.2, 1.0). The input modality for the transfer task is RGB (pixels). Exact training settings are provided in Table 11.
Table 10: COCO detection transfer settings.
Configuration	COCO Det.
Fine-tuning epochs	100
Warmup epochs	5
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Base learning rate [46]	1e-4
Batch size	256
Weight decay	0.05
Learning rate schedule	Constant
EMA decay	0.998
Eval. freq (epochs)	2
Input resolution	2242
Augmentation	Large-scale jitter (LSJ) [41]
Color jitter	✓
Table 11: ADE20K semantic segmentation transfer settings.
Configuration	ADE20K Seg.
Fine-tuning epochs	200
Warmup epochs	20
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Base learning rate [46]	1e-4
Batch size	256
Weight decay	0.5
Learning rate schedule	Constant
EMA decay	0.998
Eval. freq (epochs)	4
Input resolution	2242
Augmentation	Random crop
Crop scale	(0.2, 1.0)
Crop ratio	(0.75, 1.333)
Color jitter	x
Taskonomy-20K & Hypersim RGB  $\rightarrow$  X transfers. We perform several transfers on Hypersim [94] and a subset of Taskonomy [126] Tiny, denoted here as Taskonomy-20K. The subset contains 20'000 training images and 2000 validation images.
For Taskonomy-20K, we transfer to depth, principal curvature, reshading, occlusion edges, 2D edges, 2D keypoints, and 3D keypoints. For Hypersim, we transfer to surface normals, semantic segmentation, and 2D edges. The input modality for all transfer tasks is RGB (pixels). Exact training settings are provided in Table 12.
Taskonomy-20K & Hypersim  $\mathbf{X}\rightarrow \mathbf{Y}$  transfers. There are 72 possible transfers between the Taskonomy tasks, and 12 possible transfers between the Hypersim tasks. To reduce the computational cost of the ablations, we subsample these transfer tasks, trying to cast a diverse net. For Taskonomy-20K, we perform curvature  $\rightarrow$  2D keypoints, depth  $\rightarrow$  normals, 2D edges  $\rightarrow$  depth, surface normals  $\rightarrow$  2D edges, occlusion edges  $\rightarrow$  principal curvature, and RGB (tokens)  $\rightarrow$  surface normals. For Hypersim we perform RGB (tokens)  $\rightarrow$  segmentation, 2D edges  $\rightarrow$  segmentation, normals  $\rightarrow$  segmentation, segmentation  $\rightarrow$  2D edges, and segmentation  $\rightarrow$  normals. Exact training settings are provided in Table 13.
Taskonomy-20K & Hypersim  $\mathbf{X} + \mathbf{Y}\rightarrow \mathbf{Z}$  transfers. As before, we subsample all possible transfer tasks that use two modalities as the input, aiming for a diverse set of transfers. For Taskonomy-20K, we perform RGB (pixels) + normals → depth, RGB (pixels) + depth → reshading, reshading + 2D keypoints → normals, 2D edges + 3D keypoints → RGB (tokens), 2D edges + curvature → occlusion edges, and occlusion edges + 2D keypoints → curvature. For Hypersim we perform RGB (pixels) + depth → segmentation, RGB (pixels) + normals → segmentation, RGB (pixels) + segmentation → depth, depth + segmentation → RGB (tokens), 2D keypoints + normals → segmentation, and 2D edges + depth → segmentation. Exact training settings are provided in Table 13.
Table 12: RGB  $\rightarrow$  X transfer settings.
Configuration	Taskonomy-20K	Hypersim
Fine-tuning epochs	100	50
Warmup epochs	5	2
Frozen Transformer epochs	0	2
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Base learning rate [46]	1e-4
Frozen Transformer base lr	1e-4
Batch size	256
Weight decay	0.05
Learning rate schedule	Constant
EMA decay	0.998
Eval. freq (epochs)	2	1
Input resolution	2242
Augmentation	Random crop
Crop scale	(0.2, 1.0)
Crop ratio	(0.75, 1.333)
Color jitter	×	✓
Table 13:  $\mathbf{X} \rightarrow  \mathbf{Y}$  and  $\mathbf{X} + \mathbf{Y} \rightarrow  \mathbf{Z}$  transfer settings.
Configuration	Taskonomy-20K	Hypersim
Fine-tuning epochs	100	50
Warmup epochs	5	2
Frozen Transformer epochs	25	20
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.999
Base learning rate [46]	1e-4
Frozen Transformer base lr	2e-3
Batch size	256
Weight decay	0.05
Learning rate schedule	Constant
EMA decay	0.998
Eval. freq (epochs)	2	1
Input resolution	2242
Augmentation	Random crop
Crop scale	(0.2, 1.0)
Crop ratio	(0.75, 1.333)
Color jitter	Only if RGB is in the input
E.2 Reference model
In this section, we specify the reference model and its training settings. These choices are made by following common practices and performing educated guesses. The settings include model size, number of encoding and decoding layers, additional modifications on the architecture (e.g. no bias), choice of Dirichlet sampling parameter  $\alpha$ , input and target modalities, and training length. Following common practice for the base model size [30, 86], our base 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 seen, corresponding to roughly 400M masked samples seen for the base setting. We set the number of randomly sampled input and target tokens both to 128, and sample each using a symmetric Dirichlet distribution with parameter  $\alpha = 0.2$ . Full details are shown in Table 14.
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 distinct reference models, each trained with different seeds. In the subsequent ablation results, the average performance of the reference model setting is indicated by the symbol  $\triangleright$ . Furthermore, we highlight transfer results falling within two standard deviations (shown in Table 15) of the lowest achieved losses, enabling us to estimate which settings are significant. Table 15 further contains from-scratch results (training a randomly initialized model) for all the transfer tasks. This serves as a lower bound for the performance
that can be achieved on a given transfer task, and allows us to calibrate how difficult each transfer is, and by how much pre-training with a given setting improves upon random initialization.
Table 14: Reference model pre-training settings. Training configuration for the reference 4M model used in the ablation. All other models in ablation are obtained by changing just one of these hyperparameters.
Configuration	Reference model
Model size	Base (see Table 4)
Training length (n tokens)	100B
Warmup length (n tokens)	10B
Optimizer	AdamW [73]
Opt. momentum	β1, β2 = 0.9, 0.95
Base learning rate [46]	1e-4
Batch size	4096
Weight decay	0.05
Gradient clipping	x
Learning rate schedule	Cosine decay
Feedforward activation	GELU [49]
Input token budget	128
Target token budget	128
Input and target α	0.2, 0.2
Masking strategy	All → All
Image resolution	2242
Augmentation	None (Center Crop)
Repeated sampling [38]	4
Table 15: Reference model results: We report the average and standard deviation of the transfer results achieved by our reference 4M model. In addition, we show performance achieved by a randomly initialized model. All results are reported on the validation sets of each dataset.
COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss
▷Reference model avg.	3.06	6.11	5.07	6.75	5.80	5.36
No pre-training	3.31	6.59	6.03	7.52	6.62	6.01
Reference model std. dev.	0.008	0.008	0.021	0.018	0.016	0.010
E.3 Input modalities and target tasks
Importance of target tasks for representation learning. 4M is both a multimodal and a multitask pre-training scheme and the choice of target task(s) is a powerful way of steering what representation the model learns [10, 109]. Furthermore, multitask pre-training has been shown to improve downstream task transfers [42, 107, 5]. In this ablation, we would like to learn how the choice of target tasks affects transfers to various downstream tasks. For that, we fix the input modalities to be either RGB or all the modalities (denoted as "All"), and vary just the target tasks. The results in Table 16 (or equivalently, 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. That said, pre-training on all target modalities outperforms the other single-task and multitask alternatives in terms of average loss, no matter what input modalities were used during pre-training. Note that some of the settings in Table 16 are conceptually similar to well-known pre-training methods, and can serve as baselines for such methods, while abstracting away implementation details that would make them difficult to compare otherwise. E.g.,  $\mathrm{RBG} \rightarrow$  RGB (first row) is similar to MAE [48], while  $\mathrm{RGB} \rightarrow \mathrm{CLIP}$  (fifth row) is similar to BEITv2 [82].
Importance of multimodal pre-training for transferring to new modalities. Just as multitask pre-training improves transfers to new target tasks, we would hope that, likewise, multimodal pretraining improves transfers to new input modalities. Indeed, pre-training with multiple modalities
that may be available at transfer time has been shown to be beneficial [5], but it is not as clear how well multimodal models transfer to completely new, unseen input modalities. Table 16 shows that multimodal pre-training can significantly help with transferring to new input modalities (X→Y transfers), but comes at a performance loss at transfers that use RGB as the sole input modality. In Appendix E.4 we will explore pre-training using mixtures of different masking strategies, which enables us to train models that perform well in both regimes.
Table 16: 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. The average rank and best losses are shown for "RGB" and "All" inputs separately. Here, Geometric = RGB + Depth + Normals and Semantic = RGB + Segmentation + CLIP + Detection + Captions. 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. Note that the results shown here are identical to those of Table 2, and are included for completeness.
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
E.4 Multimodal masking strategy
Multimodal masking is at the core of 4M, so in this section, we ablate how modality tokens should be sampled and mixed, and how many tokens we should encode and decode. When performing multimodal training, it can be computationally challenging to deal with the large number of tokens from all modalities. For that reason, the choice of the number of input and target tokens has the potential to greatly reduce the computational burden, and setting them as low as possible is desirable.
Mask sampling parameter. The number of tokens to sample from each modality is determined by a symmetric Dirichlet distribution with parameter  $\alpha$  [5]. If  $\alpha$  is low, the sampling procedure will often choose 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. We ablate here various choices of  $\alpha$  values, both for the input and target sampling. Table 17 shows that models trained with higher  $\alpha$  values generally transfer better to RGB tasks, but methods trained with lower input  $\alpha$  values perform better at transferring to novel input modalities.
Table 17: Mask sampling parameter ablation: The choice of Dirichlet sampling parameter  $\alpha$  influences how many tokens are sampled from each modality. Low values lead to many samples consisting of entire modalities, while high values result in equal numbers of tokens to be sampled from each modality. We ablate various input and target  $\alpha$  values.
Input α	Target α	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
0.1	0.1	3.11	6.14	5.18	6.81	5.89	5.43	4.80
▷0.2	0.2	3.06	6.11	5.07	6.75	5.80	5.36	3.20
0.5	0.5	3.02	6.09	4.98	6.76	5.80	5.33	2.60
1.0	1.0	3.00	6.07	4.95	6.80	5.79	5.32	1.80
∞	∞	3.00	6.08	4.93	6.84	5.83	5.34	2.60
Input masking budget. The difficulty of the multimodal masked modeling task is largely determined by the number of visible (non-masked) input tokens, with fewer tokens used making the task more challenging. Indeed, since the modalities contain a lot of spatial information about each other, it is necessary to lower the number of visible tokens to keep the objective difficult enough. Furthermore, encoding only a small set of visible tokens (as opposed to encoding both visible and mask tokens [30, 7]) can significantly improve training efficiency [48]. This is especially important when training on multiple modalities [5], as the input length would otherwise grow too large.
Table 18: Input masking budget ablation: The number of visible input tokens decides the difficulty of the masked modeling task (the fewer given, the harder the task), and influences the computational cost of the encoder (the fewer, the less expensive). We fix the target masking budget at 128 tokens and vary the number of input tokens.
Input Tokens	COCO Det.	ADE20K Seg.	RGB→X	X→Y	X+Y→Z	Avg. Loss	Avg. Rank
64	3.15	6.21	5.27	6.89	5.99	5.50	4.00
▷ 128	3.06	6.11	5.07	6.75	5.80	5.36	2.00
256	3.03	6.08	5.04	6.71	5.75	5.32	1.00
512	3.07	6.12	5.11	6.78	5.83	5.38	3.00