SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hassan Akbari* Dan Kondratyuk* Yin Cui  
Rachel Hornung Huisheng Wang Hartwig Adam
Google Research  
{hassanak, dankondratyuk, yincui,Rachelhornung, huishengw, hadam}@google.com
We present Integrated Multimodal Perception (IMP), a simple and scalable multimodal multi-task training and modeling approach. IMP integrates multimodal inputs including image, video, text, and audio into a single Transformer encoder with minimal modality-specific components. IMP makes use of a novel design that combines Alternating Gradient Descent (AGD) and Mixture-of-Experts (MoE) for efficient model & task scaling. We conduct extensive empirical studies and reveal the following key insights: 1) performing gradient descent updates by alternating on diverse modalities, loss functions, and tasks, with varying input resolutions, efficiently improves the model. 2) sparsification with MoE on a single modality-agnostic encoder substantially improves the performance, outperforming dense models that use modality-specific encoders or additional fusion layers and greatly mitigates the conflicts between modalities. IMP achieves competitive performance on a wide range of downstream tasks including video classification, image classification, image-text, and video-text retrieval. Most notably, we train a sparse IMP-MoE-L focusing on video tasks that achieves new state-of-the-art in zero-shot video classification:  $77.0\%$  on Kinetics-400,  $76.8\%$  on Kinetics-600, and  $68.3\%$  on Kinetics-700, improving the previous state-of-the-art by  $+5\%$ ,  $+6.7\%$ , and  $+5.8\%$ , respectively, while using only  $15\%$  of their total training computational cost.
The human perception system is profoundly multimodal. We perceive the world through the integration of a vast array of sensory systems across domains — visual, auditory, olfactory, somatic, etc. Neurons for multimodal integration have been found in both multisensory convergence zones Calvert [2001] and unimodal regions Driver and Noesselt [2008] in the human brain. Studies in developmental psychology also suggest that interrelating simultaneous multimodal sensations is key for perceptual learning Smith and Gasser [2005]. Inspired by these findings, we see an opportunity for combined multisensory learning in machine learning systems as well.
The rapid rise of large-scale multitask frameworks and models [Raffel et al., 2020, Roberts et al., 2022, Radford et al., 2021, Yu et al., 2022, Wang et al., 2022] provides foundations for integrating capabilities that unify many disparate tasks under one model. However, given the vast quantity of independent variables involved in designing such a system, achieving an integrated multimodal machine learning model still remains an open research direction. More specifically, designing a multi-task model that integrates many multimodal signals is challenging due to various reasons: i. Different modalities require structurally different I/O signatures to properly train. ii. When training across multiple datasets, some modalities or objectives may not exist or cannot be applied, depending
on the input data and the task to perform. iii. The presence of multiple input modalities calls for careful considerations on the architectural design and allocation of parameters to certain modalities, often requiring extensive hyperparameter tuning to find the best use of computational resources.
Intuitively, as we scale a model, it becomes increasingly expensive to redesign the architecture or search for a better training objective. The issue is exacerbated in multimodal multi-task modeling, where we need to consider the combination of input modalities or datasets, loss functions, and tasks at large scales. Therefore, we would like to find a training approach that can be scaled incrementally: for any new task or objective, regardless of its input shape or output loss, we should be able to add it to the existing pretraining without compromising the previous tasks.
We navigate this problem by exploring ways in which we could train one multimodal model such that it (1) leverages as many existing datasets as possible, (2) can train on any combination of tasks or loss functions, and (3) does not slow down with the addition of any new dataset, task, or loss function. By solving all of these points simultaneously, a multimodal model could scale with an increasingly diverse and rich set of training data without needing to redesign the training framework when new tasks get integrated.
We observe in our empirical results that the combination of diverse, heterogeneous tasks that have been previously established as strong objectives individually (e.g., supervised classification and self-supervised contrastive learning) across multiple modalities are not only complementary, but can offer better convergence than training on individual tasks. By implementing Alternating Gradient Descent (AGD) and Mixture-of-Experts (MoE) via recently developed JAX primitives, we enable our model to use a fraction of the computational cost and memory required by similar large-scale perception models [Radford et al., 2021, Jia et al., 2021, Yu et al., 2022], despite the addition of multiple modalities which would normally require  $2 - 8 \times$  compute at similar batch sizes.
Given this context, our contributions and findings are as follows:
1. We define an integrated modality-agnostic encoder model, and leverage a strong combination of image-text contrastive, video-text contrastive, video-audio contrastive, and image/video/audio classification losses during pretraining to create an Integrated Multimodal Perception (IMP) model, as shown in Figure 1.
2. Contrasting the conventional approach of summing the losses of multiple objectives, we show that alternating between objectives results in a design that allows seamless integration of virtually any number of tasks and datasets without significant memory overhead and results in better downstream evaluations.
3. We show that optimization between multiple heterogeneous multimodal tasks is complementary and results in a higher quality model than trained on any individual task.
4. To train on large batches of video and audio modalities without reducing training efficiency or loss of accuracy, we design a dynamic mixture of various resolutions, sequence lengths, and batch sizes throughout pretraining, and alternate training on all input variations.
5. We enable our model with MoE, showing strong performance gains compared to a more conventional multi-tower contrastive model, even when appplying MoE to both towers.
6. We scale our resulting MoE-IMP model to 2B sparse parameters with similar compute to ViT-L (300M parameters), resulting in state-of-the-art performance on several large-scale multimodal video understanding datasets.
The optimality of AGD optimization vs. averaging the losses (or gradients) has been explored in prior work [Jain et al., 2017, Pascal et al., 2021]. Alternating multimodal multi-task training with AGD has been explored in PolyViT [Likhosherstov et al., 2021], which analyzes different methods to combining heterogeneous task in a single model. The work reports similar findings to our own, that combining objectives can be mutually beneficial and alternating between datasets weighted by the number of examples provides one of the best methods for optimization. Our work extends this to a much more generic setup supporting virtually any combination of modalities, tasks, resolutions, etc.
The use of sparse MoE for multimodal modeling can be seen in recent works like LIMoE [Mustafa et al., 2022], which uses a single MoE encoder for image-text tasks, and VL-MoE [Shen et al., 2023],
(a) The IMP model architecture.
(b) The AGD-based multi-data multi-objective training overview.
Figure 1: An overview of the IMP Training and Architecture. A mixture of datasets with varying modalities, resolutions, and objectives are randomly sampled at each optimization step and fed into the model. The model is updated alternatingly given the data-objective pairs. We use jax.jit to compile and cache computation graphs to keep each step efficient while also allowing I/O shapes to change for every optimization step without requiring any costly padding or masking strategies.
which uses modality-specific experts for image-text modeling. Our work extends this concept further, introducing video and audio modalities with alternating training on multiple tasks and resolutions without requiring modality-specific experts.
Due to the inherent complexities of integrating modalities in one model, work has been done by simplifying the problem and focusing on a small set of universal objectives applicable to all modalities [Yu et al., 2022, Wang et al., 2022]. Alternatively, some works focused on applying padding or masking strategies to handle the incompatible I/O signatures from different tasks. We found in either case, this severely limits the ability for a model to leverage pre-existing large-scale datasets or to scale to entirely new modalities. The historical reasons are that: (i) most existing models are designed for specific input modalities e.g., language [Brown et al., 2020, Chowdhery et al., 2022], vision [Dosovitskiy et al., 2020], or audio [Baevski et al., 2020]; (ii) different modalities are typically delegated to separate network weights for the best performance [Radford et al., 2021, Jia et al., 2021]; (iii) optimization difficulty with multiple modalities [Wu et al., 2022, Chen et al., 2022].
3.1 Alternating Gradient Descent (AGD)
One of the core pillars of our approach to multimodal understanding is task scalability. I.e., different combinations of data and loss objectives should be interchangeable throughout training, while the addition of any new data or objective should not cause memory or computation overhead. We found that a common issue in training large-scale foundation models in a distributed setting is that input signatures and loss objectives need to be static to avoid major inefficiencies. Accelerated graph compilation APIs allow for low-level graph optimizations that maximize hardware FLOPs utilization on distributed devices such as GPUs and TPUs, but come at a cost of requiring static I/O signatures.
Algorithm 1 Accelerated Multimodal AGD Algorithm
Input:  $M$  (model),  $T$  (training steps),  $X$  (dataset-objective pairs),  $f$  (sampling function)  
Initialize: Model state  $S_{t}$   
while  $t \leq T$  do  
    Sample data-objective pair  $(D_{t}, L_{t}) \in X$  according to sampling function  $f(t, S_{t})$   
    Compute forward pass predictions  $P_{t} = \text{jit}(M(D_{t}))$   
    Compute backwards pass on loss  $jit(L_{t}(P_{t}, D_{t}))$   
    Update model state  $S_{t}$   
end while
One approach to handle the issue with the static input signature would be to use mixed batching, where all possible inputs are constructed, and inapplicable inputs for a given dataset are padded and outputs are masked accordingly at each training step. However, this comes at a great efficiency cost, since the more tasks that are added the more time is spent computing on padded inputs. The issue with having multiple objective functions is usually resolved by mixed mini-batching, where a batch is divided to multiple mini-batches with their corresponding objective functions. The gradients for each mini-batch and objective function pair are accumulated across multiple mini-batches and the model weights are updated once using an aggregated gradient. However, this approach is also difficult to scale since the gradients across multiple mini-batches are accumulated in memory and per-task batch size naturally reduces as we add more tasks.
We propose a more generic solution based on AGD [Jain et al., 2017], allowing any changes to the three elements of the optimization system: inputs, model, and objective. AGD can be seen as a superset of the conventional SGD, where at each gradient step, a different objective may be optimized given different sets of model weights and/or input modalities. More specifically, any input modality with arbitrary shape could consume any subset of the model while focusing on minimizing any specific combination of the objective functions. According to Jain et al. [2017] it is proven that if each of such optimization steps are convex individually, an alternation between them leads to a guaranteed convergence; as opposed to mini-batching where gradient accumulation could result in sub-optimal results. From a technical point of view, this approach requires compiling multiple computation graphs, one for each unique optimization step. To enable efficient execution of multiple computation graphs, JAX offers native Just-in-Time (JIT) compilation with the jax.jit API<sup>2</sup>, which compiles the graph-of-interest at runtime and compiles a new graph if a change in structure is seen in any of the next optimization steps. Graphs themselves are cached by JAX in an in-memory lookup table so that tasks only need to be compiled once. We empirically observe that graph memory consumption constitute a negligible portion of the entire training. In our experiments, we tested up to 20 unique task structures with no significant reduction in training speed, with compilation taking only  $0.34\%$  of the total training time, up to the largest scales.
The main AGD algorithm for our multimodal framework is provided in Algorithm 1. We note that we carefully design the loop to be as agnostic as possible with respect to the data-model-objective triplet. This allows for greater flexibility in defining logic inside the model to handle the processing of individual modalities and input shapes. The sampling function can also incorporate state from the optimization process itself. E.g., the loss value of a given step can be used as a reward signal to affect the sampling behavior in the next samplings [Piergiovanni et al., 2023, Mindermann et al., 2022]. In our default setup, we sample each unique task on a given step from a (single trial) Multinomial distribution with probabilities directly proportional to the number of examples in each task. We defer more complicated or reward-based settings to future studies.
3.1.1 AGD-Specific Efficiency Considerations
We notice that in each forward call, certain model states such as activations are stored by default to be used later in the backward call for gradient calculation. This is an established compute optimization at the low-level graph compilation in XLA and similar APIs. Although this trick helps a lot with the training time, it significantly consumes memory and creates memory overhead if more than one graph is compiled and used during training. To reduce memory overhead, we use JAX's
$^{2}$ https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html
native rematerialization API, jax. checkpoint<sup>3</sup> to save memory by not checkpointing any of the intermediate model states. In our experiments, we observe an average reduction of 70-80% TPU HBM usage while resulting in only 18-20% longer step times.
We also notice that large models with many different objectives may still incur long compilation times. Therefore, we apply scan-over-layers with jax.lax.scan, a method which rolls all of the Transformer layers into a single layer that is called multiple times using different weights (instead of compiling the same function multiple times). This alone results in 15-30x faster compilation time depending on the model length. We observe increasing relative time savings with larger model sizes.
Furthermore, we accomplish these across distributed accelerators through the jax.pjit API<sup>4</sup>, which distributes JIT compilation across multiple accelerators.
3.2 Objectives
Our goal in designing the IMP model is to reuse objectives that have been shown to be robust for learning each modality. Hence, we choose the two most established supervised and unsupervised learning objectives: i. Supervised Classification using Softmax/Binary Cross-Entropy(SCE/BCE), ii. Cross-Modal Noise-Contrastive Estimation (NCE). Unless otherwise specified, we do not sum any of the above losses or accumulate gradients as would be done traditionally. Instead we apply backprop on each objective individually with AGD.
3.3 Architecture
Figure 1 shows a high-level overview of the architecture of IMP, which consists of three main modules: i. The Encoder, which accepts specific modalities and embeds them in a shared modality-agnostic space; ii. The MoE Encoder, which computes semantic contextual embeddings from the embedded tokens; iii. The Heads, which produce all the final predictions from the Encoder by re-projecting its embeddings back into a modality-specific space. We briefly explain each module here and provide more details in Appendix.
One design decision important to multimodal modeling is how to allocate parameters to each modality. As seen in works like BASIC [Pham et al., 2021], an asymmetric modality-specific design can be more optimal than using a similar-sized model for each modality. However, this comes at the cost of requiring additional hyperparameter tuning to find the optimal parameterization. As we show later in the next section, we observe that through the use of model sparsification with MoE, a unified encoder design coupled with certain modality-specific pre- and post-encoder layers is more optimal than a traditional multi-encoder setup as seen in CLIP model variants.
We follow VATT [Akbari et al., 2021], AudioMAE [Huang et al., 2022], and T5 [Raffel et al., 2020] to extract the vision, audio, and text embeddings, respectively. We add learnable 3D/2D/1D positional encodings to the embeddings and project them to a space with the same dimensionality as the model's. We pass these embedded inputs regardless of modality as-is through the shared encoder, which is a standard Transformer architecture equipped with Mixture-of-Experts FFN layers. We follow V-MoE [Riquelme et al., 2021] and LIMoE [Mustafa et al., 2022] for expert allocation. This can be seen as an inductive bias, allowing each expert to be allocated to multiple modalities if the optimization benefits. One immediate benefit is that the addition of new modalities for fine-tuning does not need any specific changes to the encoder, unlike modality-specific experts which require additional modifications and input handling [Wang et al., 2022, Shen et al., 2023].
We apply modality-specific heads on the encoder representations to produce the final outputs for loss and prediction calculations. For classification objectives, we apply a dataset-specific linear classifier to the average-pooled outputs. For noise-contrastive estimation (NCE), we closely follow the CLIP architecture, applying separate feedforward heads for each modality-to-common-space projection.
3.4 Multi-Resolution Training
One major issue when training Transformers on video data is that computation and memory efficiency are usually bottlenecked due to Transformer's quadratic complexity as a function of the input length.
<sup>3</sup>https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html
4 https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
To counteract this, we propose to adjust batch size or resolution to compensate the additional temporal tokens, hence achieving a similar total number of input tokens compared to a single-frame still image.
We first fix a set tokens per batch  $T = B \times T_{F} \times T_{H} \times T_{W}$ , which is factorized by the batch size  $B$ , frame tokens  $T_{F}$ , height tokens  $T_{H}$ , and width tokens  $T_{W}$  representing each patchified video. We observe that we can further factorize each batch by trading off different dimensions such that the total number of input tokens per step are roughly equal so that peak memory usage is preserved. For example, we can halve the spatial resolution while quadrupling the number of frames. This can increase convergence especially at the start of training, and provide a more memory efficient encoding of each objective. Furthermore, we leverage DropToken [Akbari et al., 2021] as an additional method to reduce tokens per batch by randomly dropping a fixed ratio of tokens per example. We find that for  $T_{F}$  temporal frame tokens, we can randomly drop a ratio of  $1 - \frac{1}{T_{F}}$  tokens per example to match the same tokens per batch as images. For certain objectives we find that a different mix of trade-offs is more optimal. For example, contrastive objectives favor large batch sizes, so we reduce the resolution or apply DropToken to be more memory efficient. On the other hand, classification objectives do not need as large batch sizes for optimal convergence, hence we reduce the batch size while increasing the spatiotemporal tokens.
4 Experiments and Results
4.1 Training Setup
Datasets. Our datasets consist of a diverse set of learnable signals across multiple modalities. We use WebLI [Chen et al., 2022], LAION-400M [Schuhmann et al., 2021], WIT [Srinivasan et al., 2021], CC12M [Changpinyo et al., 2021], and VCC [Nagrani et al., 2022] for vision-text contrastive learning; JFT-3B [Zhai et al., 2022], I21K [Ridnik et al., 2021], and WTS-70M [Stroud et al., 2020] for both supervised classification and label-based vision-text contrastive estimation (similar to BASIC [Pham et al., 2021]); HT100M [Miech et al., 2019] and AudioSet [Gemmeke et al., 2017] for vision-audio-text triplet contrastive loss (similar to VATT [Akbari et al., 2021]).
We use a proportionally weighted sampling algorithm, executing each task in succession. To ensure that datasets are evenly sampled, we weight each task by the number of examples, normalized to a probability distribution. For each dataset variant with different resolution sizes, we apply the same weight. For a fair evaluation on downstream tasks, we filter all near-domain examples from our pretraining datasets (about 5M examples total).
Multi-Resolution Strategy. In our experiments, we always configure the input parameters so that the number of frame tokens are always equal to 4. This will result in the base tokens per video batch being exactly  $4\mathrm{x}$  of image's. For video datasets, we construct three variants and uniformly sample from each variant during training: i. Reduce the resolution by half in each dimension, ii. Reduce the batch size by  $4\mathrm{x}$ , iii. Apply DropToken  $d = 1 - \frac{1}{T_F} = 0.75$ . For image datasets, we also apply a similar strategy but for the purpose of high-resolution learning. In addition to the base resolution, we have two extra variants: i. Reduce the batch size by  $4\mathrm{x}$  and double each spatial dimension, ii. Apply DropToken  $d = 1 - \frac{1}{4} = 0.75$ .
Training Parameters. For our final experiments, we train with a patch size of 4x16x16 on base input resolutions of 16x256x256 and 4x256x256 on video and image modalities respectively, resulting in a total of 1024 and 256 patches per sample. The text inputs in ImageNet21K and JFT are truncated to 16 tokens to improve step efficiency with no loss of information, while keeping the text length of the rest of the datasets to a maximum of 256 tokens. We use a base batch size of 65536 and train using the Adam optimizer, a peak learning rate of 1e-3 with a cosine schedule, and apply no weight decay. For MoE parameters, we apply experts-choose routing with a top- $c$  capacity factor of 1.0 and do not apply any jittering to the routing or other auxiliary losses. Training results in roughly 16B examples seen, or about 5T tokens. Taken together, these datasets represent about 11B unique image-text and video-audio-text examples.
During inference, we evaluate on the largest available resolution that the model was trained on, i.e.,  $16 \times 512 \times 512$ , and use a total of 8 clips per video at approximately 12.5 fps on all evaluated datasets.