Alternating Gradient Descent and Mixture-of-Experts for Integrated Multimodal Perception
Authors: Hassan Akbari
Dan Kondratyuk
Yin Cui
Rachel Hornung
Huisheng Wang
Hartwig Adam
Venue: NeurIPS 2023 poster (2023)
Keywords: Alternating Gradient Descent
Multimodal
Mixture of Experts
AGD
MoE
Deep Learning
Optimization
TL;DR: We combine Alternating Gradient Descent and Mixture-of-Experts to push the boundaries of efficient and scalable multimodal perception.

Abstract: 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 mitigating 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.

Alternating Gradient Descent and Mixture-of-Experts for Integrated Multimodal Perception
Hassan Akbari"
Dan Kondratyuk"
Yin Cui Hartwig Adam
Rachel Hornung
Huisheng Wang
Google Research
(hassanak, dankondratyuk, yincui, rachelhormung, huishengu, hadam) Ogoogle, com
Abstract
We present Integrated Multimodal Perception (IMP), a simple and scalable mul- timodal multi-task training and modeling approach. IMP integrates multimodal inputs including image, video, text, and audio into a single Transformer encoder with minimal mppelity-specific components. IMP makes of novel design that combines Alternating Gradient Descent (AGD) and Mixture-of-Expeer efficient model & task scaling. We conduct extensive empirical studies and reveal the following key insights: 1) performing gradient descent updates by alternating diverse modalities, loss functions, and tasks, with varying input resolutions, efficiently improves the model. 2) sparsification with MoE. on 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 a wide range of downstream tasks including video classification, image clas- sification, image-text, and video-text retrieval. Most notably, we train a sparse IMP-MoE-L. focusing on video tasks that achieves new tfte-off-hee-at art in zero-shot video classification: 77.0% on Kinetics-400, 76.8% on Kinetics-600, and 88.3% on Kinetics-70M improving the previous mate-of-theearr by +5% +6.7%, and +5.8%, respectively, while using only 15% of their total training computational cost.
(MoE) for
Introduction
The human perception system is profoundly multimodal. We perceive the world through the integra- tion of vast array sensory systems across domains -viuual, 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 perceptutal learning Smith and Gasser (2005). Inspired by these findings, we see an opportunity for combined multisensory learning machine learning systems as well.
The rapid rise of large- scale multitask frameworks and models [Raffel et 2020 Roberts 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. When training across multiple datasets, some modalities objectives may not exist cannot be applied, depending
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

on the imput 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 redesign the architecture or search for a better training objective. The issue is exacerbated in multimodal multi-task modeling, where we need consider the combination of input modalities or datasets, loss functions, and tasks at large scales. Therefore, we would like to find 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 to the existing pretraining without compromising the previous tasks.
We navigate this problem by exploring ways in which 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 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-supervisee 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-Expert (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. 202 1 Yu et al., [002], despite the addition of multiple modalities which would normally require 2-8x x 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 combi- nation of image-text contrastive, video-text contrastive, video-audio contrastive, and im- age/video/audio classification losses during pretraining to create an Integrated Multimodal Perception (IMP) model, as shown in Figure [
2. Contrasting the conventional approach of summing the losses of multiple objectives, we show that alternating between objectives results in design that allows seamless integration 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 comple- mentary 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.
Related Work
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 ViT 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 MMustafa al. 2022). which uses single MoE encoder for image-text tasks, and VL-MoE (Shen et al. 2023

---

Embed Flatten
McË Encoder
Heads
Objectives
Projection
1D Text
XJ2 Sparse Layers
Linear Classi
Binary Crass Crass-Entrog Entropy
Pos
FFN
FFN
JFT
Patch Projection
Waveform
Router
WTS
Softmax Cross Entropy
AudioSet
Pos Encedo
Seif-Atention
Spectrogram Patch Projection
Noise- Contrastive Estimation (NCE)
Fos Encode
L/2 Dense Layers
Common Space
V-+
V-A
30 Vision Patch Projection
T-
3D Pos Encodo
Self Attention
(a) The IMP model architecture.
Gradient Step: N
Objective:
Dota: N
IMP Model
Optimizer
Oojectiver
Loss Value:
Data:
Gracient
Step:
Sample Data-Objective Pairs
Perform Inference and Loss Computation
Perform seta-OOjective-Speeeee Gradient Step
(Any Subset)
(h) The AGD-hased multi-data multi-ohjective training overview.
Figure 1: An overview of the IMP Training and Architecture. mixture datasets with varying modalities, resolutions, and objectives randomly sampled at each optimization step and fed into model. The model 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-speeiff 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-specife experts.
Due to the inherent complexities of integrating modalities in one model, work has been done by simplifying the problem and focusing on small set of universal objectives applicable 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 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 al, 2022], vision (Dosovitskiy et al. 2020), or audio [Baevski al. [020] (ii) different modalities
are ty sypically delegated to separate network weights for the best performance [Rafford et al. [001] Jia 2022].
al. 2021] (iii) optimization difficulty with multiple modalities [Wu et al., 1022Chen
Method
3.1
Alternating Gradient Descent (AGD)
One of the core pillars of our approach to multimodal understanding is task scalability. Lee. 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 computation overhead. We found that a common issue 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 distributed devices such GPUs and TPUs, but come at 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 St
while t - T do
Sample data-objective pair (Dr, Le) E X according to sampling function f((,, St) Compute forward pass predictions P. (i((M(D)) Compute backwards pass on loss (P., Dz))
Update model state St
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 masked accordingly at each training step. However, this comes at 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 batch is divided to multiple mini-batches with their corresponding objective functions. The gradients for each mini-batch and objective function pair 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 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, 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 tai 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 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-T -n-Time (JIT) compilation with the jax. jit APA which compiles the graph-of 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 We note that we carefully design the loop be as agnostic as possible with respect to the fata-modellobjectiin 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. F.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 al., 2022). In our default setup, we sample each unique task on a given step from (single trial) Multinomial distribution with probabilities directly proportional to the number of examples 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 XLA and similar APls. 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

---

native rematerialization API, jax. checkpoint 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.
a method which rolls all of the
Transformer layers into a single layer that 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.pji pjit AP 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, choose the most established supervised and unsupervised learning objectives: i. Supervised Classification using Softmax/Binary Cross-Entropy Cross-Entropy(SCEE Entropy(SCE/BCE) Costrascnee Estimiation (NCE). Unless otherwise specified, we do not
ii Cross-Modal 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 [shows a high-level overview of the architecture of IMP, which consists three main modules: i. The Embedder, which accepts specific modalities and embeds them in a shared 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-speciff space. We briefly explain each module here and provide more details in Appendix.
One design decision important to multimodal modeling how to allocate parameters to each modality. As seen in works like BASIC [Pham 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 tuming to find the optimal parameterization. As 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-speciff pre- and post-encoder layers is more optimal than 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/ID positional encodings to the embeddings and project them to 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 standard Transformer architecture equipped with Mixture-of-Exxe -xpeertsss FFN layers. We follow V-MoE [Riquelme et al., 2021] and LIMoE [Mustafa et al., [022] for expert allocation. This inductive bias, allowing each expert to be allocated to multiple modalities if the
can be optimization benefits. One immediate benefit is that the addition of new modalities for fine-tuning does not need any specific changes the encoder, unlike modality-specific experts which require additional modifications and input handling Wang et al. 2022 Shen et al.) 2023
We apply podaalty-sspecifi heads 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-pooles outputs. For noise-contrastive estimation (NCE), we closely follow the CLIP architecture, applying separate feedforward heads for each
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.
5

---

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 x Tr x Th x Tw, which is factorized by the batch size B, frame tokens Tr, height tokens Th. and width tokens Tw 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 Drop'ooken [Akbari et al., [021] as an additional method to reduce tokens per batch by randomly dropping a fixed ratio of tokens per example. We find that for Tp temporal frame tokens, we can randomly drop ratio of 1 + tokens per example to match the same tokens per batch as images. For certain objectives we find that a different mix of trade-off's is more optimal. For example, contrastive objectives favor large batch sizes, so we reduce the resolution or apply Drop'Token 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., (021]. WIT [Srinivasan et al.
202T], CC12M [Chnnpiinyo et al. 202T]. and VCC [Nagrani et al. 2022] for vision-text contrastive
learning; JFT-3B [Zhai et al. 2022]. , 12KK [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]); HT 100M [Miech et al., (019) and AudioSet [Gemmeke et al. 2017] for vision-audio-tex triplet contrastive loss (similar to VATT [Akbari et al., 2021]).
We use 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 that the number of frame tokens are always equal to 4. This will result in the base tokens per video batch being exactly 4x of image's, 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 4x, iii. Apply Drop'Token d = + 0.75. For image datasets, we also apply similar strategy but for the purpose of high-resolution learning. In addition to the base resolution, have two extra variants: i. Reduce the batch size by 4x and double each spatial dimension, ii. Apply DropToken d'
0.75.
Training Parameters. For our final experiments, we train with a patch size of 4x 161x6x 66 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 1 IBB 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., 16x512x512 12, and use total of 8 clips per video at approximately 12.5 fps on all evaluated datasets.

---

MODEL
PPT
TPU-DAYS
INIK
C100
K400
K600
K700
UCFI0L
HMDB51
ESC-50
Radford et
2021
400M
762
CoCa-B5 Yir th 2024
380M
1.8k
82.6
X-CLIP
65.2
72.0
BIKE
Wo 20244
230M
68.5
80.8
52.8
Vexi
AudioCLIP
Crowul
Lrowal cl.a. 2021
430M
69.4
IMP-B
82.4
63.6
39.7
IMP-MoE-BB
90M
150
83.2
84.9
68.2
65.7
52.1
88.7
46.6
47.8
IMP-MoE-L
350M
1.5k
83,9
87.0
77.0
76.8
91.5
59.1
Large-scale models
L.IMF [Mustafa al. 2022
680M
84.1
LIT VIL Chen er 2024
2B
84.5
83.6
86.3
VideoCoCS Tan et at 2022
2B
10.5k
72.0
70.1
62.5
86.6
58.6
Table 1: Zero-Shot Classification (top-1) 1) results on image, video, and audio datasets. IMP achieves a new state-of-the-ar on zero-shot video action recognition by a wide margin with significantly low training cost. Considering the total number of Parameters Per Token (PPT), IMP also significantly outperforms comparable models on zero-shot image classification.
4.2
Main Results
We scale up and tune IMP for best performance on video datasets and evaluate it on a variety of downstream tasks and datasets understand how it generalizes to other modalities. Table shows the sero classification capabilities of the model on several image, video, and audio datasets.
We note that IMP significantly outperforms previous state-of- ofttt--o--he-rrr the- regardless of the model size and sets new record on Kinetics (Kay et al., [017 Carreira et al., 2018 [201)] UCF101 (Soomro et al., 2012]. and HMDB-51 Kuehnne et al. 2011 top-1 accuracy. Compared to the previous offte-of-the-arr VideoCoCa Yan al., 2022], we train IMP-MoE-L on 256 TPU chips for 6 days, representing only 15% of the total training cost of VideoCoCa. Considering the total number of parameters per token (PPT), IMP also outperforms the previous comparable state-of-the-ard model, CoCa, on ImageNet [Russakovsky et al.) 20155 and CIFAR-100 with relatively large margin. However, observe that the model falls behind the state-of-the-art in zero-shot audio classification on ESC- 50 PPizzak 2013]. This might be explained by the fact that the total training examples for audio modality are almost negligible compared to image and video. Hence, the model has a very strong performance video and image. We argue that this could be resolved by simply introducing more samples and a more balanced train scheduling method, which we differ to future studies.
4.3 Ablation
In this section, we hightlight experimentation with some key results which motivate the chosen set of features for our final IMP model. We refer the reader to Appendix for more ablation several other aspects of the model. The experiments in this section use IMP-S or IMP-B trained for 250k steps with base batch size of 8192. We set fixed video/image resolution of 16x224x22444x22xx222 using a patch size of 4x 16. Unless otherwise specified, we do not apply multi-scale resolution.
Combined objectives are mutually beneficial. We train the model on ImageNet21k with two objectives: Noise-Contrastive Estimation (NCE) and Softmax Cross-Entropy (SCE) and explore the following: i. train on the objectives separately, ii. combine the objectives by summing them, and ii. alternating (AGD) between the objectives on each step. In the case of alternating, for a fair comparison so that training time is equivalent, we fix the same number of steps (250k) so that each objective only optimizes 50% of the total steps (125k). We evaluate on ImageNet k and CIFAR-101 by image-to-text retrieval and linear probing on the frozen model's features and report the results in Table 2 It is not surprise that classification objective benefits fine-tuning evals the most, while contrastive objective benefits open vocabulary classification. However, we observe that combining both objectives is better than optimizing on them individually. And alternating between the objectives is better than non-AGD objective mixing. These results are similar to the findings of ViT [Likhosherstov et al. 2021], which report optimal performance on alternating the objectives, weighted by the size of each dataset. This motivates us to fix one objective per training step and alternate optimization between them.
7

---

IMAGENETIK K
CIFAR- 100
OBJECTIVE
LINEAR
I-T
LINEAR I-T
NCE
39.8
46.7
84,0
49.8
Softmax
41.1
82.6
NCE Softmax, Sum
47.6
46.7
82.4
51.6
NCE + Softmax, Alternating
49.9
48.0
84.1
52.4
Table 2:
Combining multiple objectives during pre-training. Alternating between both objectives offer the best performance, despite training on the same number of total steps.
CC
121K (Softmax)
CC 121K (NCtm)) LAJON
CC 121K INCE)
121K (NCE4
50
70
60
50-
$15
40
10
ImageNerz2 25
CIFAR10O Liner
DIFAR10O 25
Fiicr T-1
1-T
COCO T
COCO I-T
Combining multiple datasets and objectives using AGD. We integrate CC, 121k, and LAION with NCE and SCE using AGD and observe consistent improvement in downstream results. We also observe that NCE and SCE mutually beneficial. Further optimality is provided by adding larger and more diverse datasets like LAION.
Figure 2:
224x224
224x224
60-
320x3320 (drop 0.5)
949192
224x222
50
0.5x)
40
320x320 (batch
224x224
2
320x.320 (eatchh 0.5x)
320x320 (drop 0.5)
160x160 (batch 2x)
ImageNet 25
Flie T-4
FlickT-T
COCO T-4
COCO 1-T
Figure 3: Combining multiple input sizes using AGD. We sample each variant with equal probability. Downstream results improve as we add more variants to the training mixture.
Multi-task multi-dataset AGD is also mutually beneficial. In Figure 2 compare the result of adding additional datasets to the pretraining mixture. We additionally compare results across Flickr30k (Young al. [014] and COCO [Lin et al.) 2014) datasets. We start with CC12M dataset and gradually add new datasets and objectives. Most notably, we compare the addition of 121K dataset, showing complementary improvement when combining NCE and SCE objectives. Similar to 121K isolated experiments, adding SCE benefits the entire pretraining mixture. While SCE benefits zero-shot results, NCE benefits linear probing results too. Certain dataset combinations (CC+121 K, CC+LAION) cause instability at the beginning of training. Adding classification objective has stabilizing effect, significantly reducing the chance of slow convergence. Optimizing on LAION directly is difficult, but benefits training a lot more when mixed in with other datasets. This motivates us to further integrate a larger set of diverse datasets and objectives.
Multi-scale resolution provides universal improvement. Figure shows comparison of using different combinations of resolution, batch size, and DropToken as input. In all settings, fix the total tokens per batch, and we ensure that all training runs use the same number of total steps. We see that certain types of datasets respond well DropToken while others may not. CC with double the batch size and DropTokenn 0.5 improves zero-shot image classification. Droptoken + 320x320 image on 121K SCE pretrain is better for linear probing and image-text retrieval. Adding multiple versions of smaller batch size higher res, Drop'Token + higher res, larger batch size + lower res, can significantly improve downstream results. We find that dynamic mixtures of resolution, batch size, and DroeToken are always helpful.
8

---

33
CI Linee
SPAA1D0
Aic T4
Flcur
COCO
COCO Tl
Figure 4: IMP with Mixture-of- -ffe-rerts.. Results show that using modest 4 experts increases the model's 's accuracy substantially on ImageNet zero-shot evaluation. When we scale up the experts to 16, we consistent and significant improvement across all downstream evaluations.
(image only)
IMP-McE-8, experts (image only)
experts (image
IMP-8 (image video)
IMP-MEE-B, 8, 16 experts (image video)
90
80
XLire
70
60
50
MIN
ImageNet 25
CIFAR1OB Linear CIFAR100 25
Flickr T-+1
Flickr T-T
COCO T-+l
COCO T-T
Kinetics400 25
Figure 5: Improved multimodal perception using MoE. Results show significant improvement on diverse multimodal perception when we use the MoE variant of IMP. The addition of audio reduces accuracy on image and video metrics across the board, but is much less prominent when using MoE.
Towrr M)
EM)
(172M)
Single Tower Nof., expets
12:
93
50
40
315
2a
OFARLOD Unear
CIFARLOO
COCO
COCO
Figure 6: Comparison of single-tower vs. multi-tower designs on IMP-B. A single-tower MoE model is both parameter efficient and compute efficient compared to multi-tower dense variants.
MoE provides universal improvement across modalities, and resolves the single-tower encoder parameter bottleneck. The main challenge in designing a unified encoder tower as we have described is that parameters must be split between multiple modalities, thus harming accuracy. Compared to two-tower contrastive model, the encoder of unified image-text model contains half the parameters, while keeping training efficiency the same. One direction we explore whether large increase in parameters from MoE is sufficient to resolve parameter bottlenecks. In Figure 4 we observe that simply replacing a dense model with an equivalent MoE model with just 4 experts, we can provide large gain in accuracy, especially for zero-shot metrics. This provides promising indication that MoEs can be used to bridge the multimodal gap. We observe that with the addition of MoE, we can significantly close the gap between multiple modalities as seen in Figure 5 Since experts are free to choose which tokens are allocated to different experts, we observe strong alignment between experts and modalities.
Single-tower MoE outperforms multi-tower dense variants. We find that out all the variants we tested, unified MoË encoder provided the most parameter and compute efficient design, while significantly outerforming a multi-encoder modality-specife model in downstream results, as seen in Figure 6 When comparing two-tower models, can either split the parameters to be roughly

---

equal in size to a single tower, or duplicate the towers to double the parameter count while providing equivalent computation. We observe that multi-tower dense models are more compute efficient than single-tower dense models, but less parameter efficient. However, single-tower MoF model both more compute and parameter efficient than all variants, showing improved generalization and using fewer parameters with the same compute budget as multi-tower dense model. These results show universal superior parameter and compute efficiency and higher accuracy by using just 4 experts. This observation suggests that our method can be used for integrated multimodal multi-task modeling without worrying about the complications of modality-specine design choices or downstream performance degradation observed in previous modality -agnosti designs [Akbari et al. 2021].
5 Conclusion
In this paper we presented an integrated training and modeling approach for multimodal perception using AGD and MoE. We observed that AGD enables task scalability and multi-resolution training. which improves the training convergence and the model's generalization capabilities. On the other hand, we found that MoE can play very important role in integrating multiple modalities into one unified model. Given these findings, we scaled the model with hyperparamters tuned specifically for video understanding and achieved state-of-the-ar performance zero-shot video action recognition with a significant margin. Furthermore, we observed that the model also generalizes on other modalities and achieves competitive downstream results. In a nutshell, IMP opens a door to data (e.g., modality, resolution, etc.) and task scalability two important directions that have been neglected in many multimodal understanding works due to the inherent technical limitations. Due to the vast range of elements involved in this system, defer multiple directions to be explored in future work: generative objectives and model architectures, 2. causal MoE for generation, sophisticated methods for data-objective sampling, 4. more downstream evaluations.
Acknowledgments and Disclosure of Funding
We would like thank Joan Puigcerver, Carlos Riquelme, and Basil Mustafa for their advice on MoE implementation and analysis; Anselm Levskaya for his help with advanced core JAX and Flax implementation; the T5X team for their support for scalable model partitioning, and Erica Moreira and Victor Gomez for their help with resource allocation.
References
[1] Hassan Akbari, Liangzhe Yuan, Rui Qian, Wei-Hong Chuang, Shih-Fu Chang, Yin Cui, and Boqing Gong. Vatt: Transformers for multimodal self-supervised learning from raw video, audio and text. NeurIPS, 2021.
[2] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli wav2vec 2.0: A framework for self-supervised learning of speech representations. NeurIPS, 2020.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. NeurIPS, 2020.
[4] Gemma A Calvert. Crossmodal processing in the human brain: insights from functional neuroimaging studies. Cerebral cortex, 2001.
[5] Joao Carreira, Eric Noland, Andras Banki- Morkit, Chloe Hillier, and Andrew Zisserman. A short note about kinetics-600. arXiv preprint arXiv:1808.00 2018.
[6] Joao Carreira, Eric Noland, Chloe Hillier, and Andrew Zisserman. A short note on the kinetics- 2019.
700 human action dataset. arXiv preprint
[7] Soravit Changpinyo, Piyush Sharma, Nan Ding, and Radu Sorieut. Conceptual 12m: Pushing web-scale image-text pre-training to recognize long-tail tail visual concepts. In CVPR, 2021.
10

---

[8] Xi Chen. Xiao Wang, Soravit Changpinyo, AJ Piergiovanni, Piotr Padlewski. Daniel Salz, Sebastian Goodman, Adam Gryener, Basil Mustafa. Lucas Beyer, al. Pali: A jointly-scaeed multilingual language-image model. arXiv preprint 20Xiv:2009.06794, 2022.
[9] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprini urXiv:2204.00 0.00010,0...... 2022.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenbom, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 2020.
[11] Jon Driver and Toemme Noesselt. Multisensory interplay reveals crossmodal influences on "sensory-ppecifie brain regions, neural responses, and judgments. Neuron, 2008.
[12] Jort F Gemmeke, Daniel PW Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, Channing Moore, Manoj Plakal, and Marvin Ritter. Audio set: An ontology and human-labeled dataset for audio events. In ICASSP, 2017.
[13] Shreyank N Gowda, Mareus Rohrbach, and Laura Sevilla-Lara. Smart frame selection for action recognition. In AAAI, 2021.
[14] Po-Ya Yao Huang, Hu Xu, Juncheng Li, Alexei Baevski, Michael Auli, Wojciech Galuba, Florian Metze, and Christoph Feichtenhofer. Masked autoencoders that listen. NeurIPS, 2022.
[15] Prateek Jain, Purushottam Kar, et al. Non-convex optimization for machine learning. Founda- tions and Trends@ in Machine Learning, 2017.
[16] Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoe Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In /CML. PMLR, 2021.
[17] Will Kay, Joao Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijaya- narasimhan, Fabio Viola, tim Green, Trevor Back, Paul Natsev, et al. The kinetics human action video dataset. arXiv preprint arXiv:1705.06600 2017.
[18] Hildegard Kuehne, Hueihan Jhuang, Estíbaliz Garrote, Tomaso Poggio, and Thomas Serre. Hmdb: a large video database for human motion recognition. In ICCV, 2011.
[19] Valerii Likhosherstov, Anurag Arnab, Krzysztof Choromanski, Mario Lucic, Yi Tay, Adrian Weller, and Mostafa Dehghani. Polyvit: Co-training vision transformers on images, videos and audio. arXiv preprint 1rXiv:2111.1299: 2021.
[20] Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. ECCV, 2014.
[21] Antoine Miech, Dimitri Zhukov, Jean-Baptiste Alayrac, Makarand Tapaswi, Ivan Laptev, and Josef Sivic. Howto 00m: Learning a text-video embedding by watching hundred million narrated video clips. In /CCV, 2019.
[22] Sören Mindermann, Jan M Brauner, Muhammed T Razzak, Mrinank Sharma, Andreas Kirsch, Winnie Xu, Benedikt Höltgen, Aidan N Gomez, Adrien Morisot, Sebastian Farquhar, et al. Prioritized training on points that are learnable, worth learning, and not yet learnt. In ICML. PMLR, 2022.
[23] Basil Mustafa, Carlos Riquelme, Joan Puigcerver, Rodolphe Jenatton, and Neil Houlsby. Mul- timodal contrastive learning with limoe: the language-image mixture of experts. NeurIPS, 2022.
[24] Arsha Nagrani, Paul Hongsuck Seo, Bryan Seybold, Anja Hauth, Santiago Manen, Chen Sun, and Cordelia Schmid. Learning audio-video modalities from image captions. arXiv preprint MrXiv2224.00677 2022.

---

[25] Bolin Ni, Houwen Peng, Minghao Chen, Songyang Zhang, Gaofeng Meng, Jianlong Fu, Shiming Xiang, and Haibin Ling. Expanding manguage-image pretrained models for general video recognition. In ECCV, 2022.
[26] Lucas Pascal, Pietro Michiardi, Xavier Bost, Benoit Huet, and Maria A Zuluaga. Improved optimization strategies for deep multi-task networks. arXiv preprint arXiv:2109.11628 2021.
[27] Hieu Pham, Zihang Dai, Golnaz Ghiasi, Kenji Kawaguchi, Hanxiao Liu, Adams Wei Yu, Jiahui Yu, Yi-Ting Chen, Minh-Thang Luong, Yonghui Wu, al. Combined scaling for open-vocabulary image classification. arXiv preprint arXiv: 2111 10050, 2021.
[28] Karol J Piczak. Esc: Dataset for environmental sound classification. In ACM MM, 2015.
[29] AJ Piergiovanni, Weicheng Kuo, Wei Li, and Anelia Angelova. Dynamic pretraining of vision-language models, 2023. "ttt.tttts://pperee///e /et.:e//rrormm///efcfcccc.e..
[30] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In JCML. PMLR, 2021.
[31] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J Liu, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR, 2020.
[32] Tal Ridnik, Emanuel for the masses. arXiv preprint arXiv:2104. 10972, 2021.
Asaf Noy, and Lihi Zelnik-Manor Imagenet-21k pretraining
[33] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. NeurIPS, 2021.
[34] Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, et al. Scaling up models and data with t5x and seqio. arXiv preprint arXiv: 17189 2022.
[35] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma. Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. IJCV, 2015.
[36] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarezyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-400m: Open dataset of fillfeee 400 million image-text pairs. arXiv preprint arXiv:2111.0211 2021.
[37] Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, and Yuxiong He. Scaling vision-language models with sparse mixture of experts. arXiv preprint arXiv:2303.07226 2023.
[38] Linda Smith and Michael Gasser. The development of embodied cognition: Six lessons from babies. Artificial life. 2005.
[39] Khurram Soomro, Amir Roshan Zamir, and Mubarak Shah. Ucf101: dataset of 101 human actions classes from videos in the wild. arXiv preprint arXiv:1212.04022 2012.
[40] Krishna Srinivasan, Karthik Raman, Jiecao Chen, Michael Bendersky, and Mare Najork. Wit: Wikipedia-based image text dataset for multimodal multilingual machine learning. In SIGIR, 2021.
[41] Jonathan C Stroud, Zhichao Lu, Chen Sun, Jia Deng, Rahul Sukthankar, Cordelia Schmid, and David A Ross. Learning video representations from textual web supervision. arXiv preprint arXiv:2007. 14937. 2020.
[42] Wenhui Wang, Hangbo Bao, Li Dong, Johan Bjorek, Zhiliang Peng, Qiang Liu, Kriti Aggarwal, Owais Khan Mohammed, Saksham Singhal, Subhojit Som, et al. Image as foreign language: Beit pretraining for all vision and vasion-language tasks. arXiv preprint arXiv:2208. 10442, 2022.
12

---

[43] Junru Wu, Yi Liang, Feng Han, Hassan Akbari, Zhangyang Wang, and Cong Yu. Scaling multimodal pre-training via moos-tt gradient harmonization. In NeurIPS, 2022.
[44] Wenhao Wu, Zhun Sun, and Wanli Ouyang. Transferring textual knowledge for visual recogni- tion. arXiv preprini arXiv:2207.01297, 2022.
[45] Wenhao Wu, Xiaohan Wang. Haipeng Luo, Jingdong Wang, Yi Yang, and Wanli Ouyang. Bidirectional cross-modal knowledge exploration for video recognition with pre-trained vision- language models. arXiv preprint 230.::3300100/822. 2022.
[46] Shen Yan, Tao Zhu, Zirui Wang. Yuan Cao, Mi Zhang, Soham Ghosh, Yonghui Wu, and Jiahui Yu. Video-text modeling with zero-shot transfer from contrastive captioners. arXiv preprint arXiv:2212.04979, 2022.
[47] Peter Young, Alice Lai, Micah Hodosh, and Julia Hockenmaier. From image descriptions to visual denotations: New similarity metrics for semantic inference over event descriptions. Transactions of the Association for Computational Linguistics, 2014.
[48] Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. arXiv preprint arXiv:2205.01917. 2022.
[49] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- In CVPR, 2022.
13


Reviews:
[{"review_id": "t3iPOexFyF", "reviewer": "Reviewer_tAnb", "summary": "The paper proposed Integrated Multimodal Perception (IMP) for multi-modal multi-task learning. IMP consists of a shared Mixture-of-Experts (MoE) Encoder as well as modality-specific embedding layers and heads re-projecting representations to modality-specific space. Optimizing towards both supervised classification losses and self-supervised contrastive learning NCE losses, the authors proposed to adopt Alternative Gradient Descent (AGD) during training to accommodate different inputs, outputs and learning objectives without bringing too much memory / computation overhead when incorporating more tasks. Furthermore, a multi-resolution strategy was also proposed to train IMP on large batches of video and audio modalities without losing efficiency or accuracy, where various resolutions, sequence lengths, and batch sizes are used dynamically. Experiments on a good amount of public datasets show that the proposed method can achieve better or competitive performance on various downstream zero-shot classification tasks except audio classification on ESC-50. Ablation studies support some of the design choices.", "questions": "While the paper's major claim/contribution is a new method for multi-modal multi-task learning which can integrate any number of modalities/tasks with arbitrary inputs/outputs/objectives, Weaknesses 1 and 2 are critical and should be addressed before an accept decision can be made, in my opinion. Please address them accordingly.\n", "limitations": "N/A", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper introduced a possible way to integrate arbitrary number of modalities into one model. The AGD + unified encoder + MoE design seems an interesting solution without losing efficiency and accuracy when incorporating more modalities **in some situations**.\n\n2. The proposed dynamic resolution strategy is somewhat novel and effective.\n\n3. Experiments on considerably various public datasets demonstrate the effectiveness of the proposed method on zero-shot classification when comparing with previous state-of-the-arts\n\n4. Ablation studies on important model designs were conducted to give a more comprehensive understanding of the methods  ", "weaknesses": "1. One critical problem of IMP is that it seems to suffer from performance drop when incorporating new modalities as per to Fig. 5, possibly due to the shared encoder. Although MoE can alleviate it to some extent in some cases, it is not always working. For example, in Fig. 5, for COCO Text->Image and Image->Text, and K400 zero-shot, IMP can achieve better or similar performance when new modalities of videos are introduced, while IMP-MoE-16Exp obtained worse accuracy. This defect somewhat goes against the major claim of this paper.\n\n2. Another problem is that the model performance on different tasks subjects to change depending on training sample distribution, according to Tab. 1 and Line 261-264. While tuning the distribution of samples of different modalities can improve audio classification, it is not sure whether or not performance on other modality classification will drop, based on Weakness 1. It is also not quite sure how robust the proposed method is against any changes of sample distribution, which may hinder the incorporation of more datasets in the future, going against the major claim again.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review, we will try to address any concerns encountered in the paper.\n\n>One critical problem of IMP is that it seems to suffer from performance drop when incorporating new modalities as per to Fig. 5\n\nWe would like to note that this problem of performance drop when integrating new modalities is fairly universal across any model, assuming a fixed computation or parameter budget. The model is forced to balance all modalities across the available computation/parameter budget, so it's no surprise that a model could decrease in average performance with so many diverse modalities competing against each other. Without the integration of our techniques such as MoE, multi-scale data, multi-objective training, and accelerated AGD, we would observe a much steeper decline across all of our benchmarks. For example, we show in Figure 2 of our supplementary material a fairly steady reduction in accuracy after each new modality is integrated, while with MoEs in Figure 5 in the main paper this distribution is remarkably different. In fact, the zero-shot image performance improves with the presence of video data that did not happen without MoEs. So while we may see, e.g., COCO results look lower when comparing image only vs. image+video, our main claim is that, *on average*, the combination of all of our methods provide significantly better improvement in integrating all of these modalities into one model than previous approaches (while also enabling cross-modal knowledge transfer on four separate modalities). These relative reductions in accuracy tend to apply on the small scale but become less prominent when scaling the MoE model as seen in Table 1 (see also section C the supplementary material). We observe fluctuations may occur on very specific benchmarks but the average shows a steady improvement.\n\nWe also note that a two-tower or multi-tower model may incur additional memory consumption as the parameters are replicated for each modality (see supplementary material), so we can typically scale our model larger than existing approaches. In Figure 5, we compare our model against itself so it mainly serves as as a way to gauge the influence of various modalities, but because our method also incorporates various aspects of training, when we compare vs. other models, the intersection of all of our various improvements (multi-task multi-objective multi-resolution data, MoEs, etc.) are all very important for surpassing the existing state-of-the-art. We are comparing a highly general multimodal model (one that trains on 4 modalities at the same time) against other larger models more specialized to one or two modalities, but despite that, our methods consume fewer resources and scale across these modalities much more easily.\n\n>Another problem is that the model performance on different tasks subjects to change depending on training sample distribution, according to Tab. 1 and Line 261-264\n\nAgain, performance differences across changes in sample distribution is another universal problem for multimodal models, we would see very similar types of concerns for other models as well. Our paper is not specifically about providing a comprehensive solution to this type of problem but rather finding a method that is better than other competing methods. Our models in table 1 are all trained on the same data, so this comparison is not subject to change in the sample distribution, and we see that the use of MoEs and increase in model size tend to bridge large gaps that might be caused by the data distribution despite using the same data for training. We do provide results in Figure 2 and Figure 5 (and Figure 2 in the supplementary material) to suggest that on especially image and video datasets, these modalities provide mutual improvement when integrated in conjunction with all of our techniques, so scaling one larger would help the other. Audio deserves special consideration, as this modality is sufficiently different from the other modalities that we've observed any scale of audio data could cause accuracy drops on the other modalities. This is one reason why video+text multimodal works typically avoid modeling audio, as it requires special handling like a separate audio network. Instead, we show a much simpler approach of using a combined MoE tower to mitigate this. Therefore, we would argue that the approach is not negatively impacted by the sample distribution in the same way as competing works, as AGD in conjunction with MoEs and multi-scale data provide a robust way to maximize performance on all modalities even in the presence of highly unbalanced audio to video/image data.\n"}, {"review_id": "M9Bnyq5iXX", "reviewer": "Reviewer_tYu8", "summary": "This paper proposes a scalable multimodal multitasking approach. It combines alternating gradient descent and mixture-of-experts to train a unified model. The extensive experiments verify the effectiveness of the proposed method. By scaling up the model, this method sets up a new state-of-the-art in zero-shot video classification.", "questions": "- In Figure 5, the audio modality contributes less to the performance improvement on the image (or video)-text tasks. How about the importance of image (or video) to audio-related tasks?\n", "limitations": "Yes, the authors have adequately addressed the limitations.\n", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method achieves excellent performance in zero-shot video classification. It further confirms improving data and task scalability in multimodal learning is promising.\n\n- This paper is well written and easy to understand.\n", "weaknesses": "- Using alternating gradient descent for efficient multimodal multitasking learning is not new, which has been explored in PolyViT. It seems that the proposed method is an extension with more engineering improvement. Also, the success of MoE architecture has been validated on image-text tasks. The contribution of this paper is to incrementally extend this architecture to video and audio modalities.\n\n- Compared with other methods, the proposed methods use more training data. Does the improvement mainly come from the scale of the training data? If PolyViT is pretrained on a similar scale of the data, does the proposed method still have advantages?\n\n- typo. Line 264, \"differ\" -> \"defer\"\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for providing a detailed review of our paper, here we try to address some of the comments of the paper.\n\n>Using alternating gradient descent for efficient multimodal multitasking learning is not new, which has been explored in PolyViT. It seems that the proposed method is an extension with more engineering improvement.\n\nWhile AGD itself has been used in PolyViT, we propose a novel generalization of the approach which has not been explored in any other work. Due to technical limitations of previous implementations of the approach, we were unable to find other methods training efficiently on many modalities each with their own dynamic input shapes; the most popular techniques tend to waste computation on padding which we found to be a significant source of inefficiency. By carefully constructing computation graphs that leverage primitives such as a compilation cache in a modality agnostic way, we can maximize effective training FLOPs utilization across distributed accelerators. The addition of this along with multi-scale representations of each modality further boosts accuracy and efficiency at the same time, as shown in table 1. We also experiment with both classification and contrastive objectives, where we might have any combination of optional paired or unpaired data, something that PolyViT does not try. \n\nAGD when coupled with cross-modal learning is something we found to be very important to our results, as we can see knowledge transfer from videos to images and vice versa. We have results in Figure 2 in the supplementary material which indicates that integrating multiple modalities like image and video into a dense model tends to harm the accuracy of the other modalities. However, with Figure 5 in the main paper, we see that MoEs reverse this trend and improve. Most notably, the addition of video data actually improves the zero-shot performance of the image datasets, which was not true of the dense model. These results show that a good diversity of tasks and modalities that are enabled by our more flexible version of AGD, in combination with expanding the model capacity, are important for multimodal learning. We would not have been able to demonstrate such capability without leveraging the dynamic graph compilation and dataset-objective pair task sampling as described.\n\n>the success of MoE architecture has been validated on image-text tasks\n\nWhile concurrent work such as image-text has been shown on MoEs, we note that the addition of video and audio data is significantly more challenging to show improvement across all four modalities, and has been the primary focus of the paper, i.e., to demonstrate which scaling techniques are the most useful to integrate these modalities into a single model.\n\n> Does the improvement mainly come from the scale of the training data?\n\nNot all of our results can be explained by scaling up our data. We note that CoCa-B achieves 82.6% on ImageNet1k zero-shot classification while IMP-B, a similarly sized encoder (ViT-B) is 80.5%. However, IMP-MoE-B trained on the same data closes the gap with  83.2%, showing the practical effects of using the MoE architecture. Similarly, our ablations show that alternating training on multi-scale data is useful (figure 3), and multiple objectives (table 2, figure 2) in tandem also help. We admit the addition of more data is certainly helpful, but it is only the combination of all of our approaches (accelerated AGD, multi-scale data, multi-objective, MoE) that we see improvement that is capable of surpassing the state-of-the-art results across multiple multimodal benchmarks. If we simply train a PolyViT architecture on more data, our improvements would be fairly marginal (in figure 2, addition of the larger scale LAION dataset helps some benchmarks but on average the combination of multiple objectives help more).\n\n> In Figure 5, the audio modality contributes less to the performance improvement on the image (or video)-text tasks. How about the importance of image (or video) to audio-related tasks?\n\nWe found that one significant challenge in training is in the integration of the audio modality. Especially without our AGD data mixture and MoE improvements, the performance declines significantly due to the difference in how audio data is learned vs. other modalities (see figure 2 in our supplementary material). Despite this, figure 5 shows a minimal penalty of our model on a few tasks, with the benefit that the model now has zero-shot audio capability that previously did not exist. \n"}, {"review_id": "MRJTefdxpz", "reviewer": "Reviewer_SrPj", "summary": "This paper proposes Integrated Multimodal Perception, which can use image, video, and audio datasets for multimodal and multitask training. The method is scalable, benefiting from the alternating gradient descent as it can alternate diverse modalities, loss functions, and datasets to perform gradient descent. Moreover, using a mixture of experts can help handle different modalities with only a modality-agnostic encoder. The multi-resolution training accelerates the training by using different spatial resolutions, sequence lengths, and batch sizes.", "questions": "1. Algorithm 1 shows the sampling uses step t and model state S_t, but Lines 142-143 say that sampling is directly proportional to the number of examples in each dataset. That is, the more samples a task has, the more likely the dataset gets sampled. It would be better to make the algorithm consistent with the descriptions. Moreover, in Algorithm 1, should there be another mini-batch sampling given a chosen task? Should the sampling function f be only for task/dataset sampling?\n\n2. Line 227 mentions filtering examples close to the downstream tasks. How is the filtering conducted?\n\n3. Why the ablation study in Section 4.3 uses spatial resolution 224 while the main results in Table 1 use resolution 256 during training?\n\n4. The paper mentions using JAX in several places, like jax.jit compiling graph-of-interest at runtime, jax.checkpoint saving memory by not checkpointing any of the 152 intermediate model states, and jax.lax.scan accelerating the graph compilation. I'd like to know whether Pytorch can also do so. I know Pytorch also has checkpointing functionality, and Pytorch 2.0 provides torch.compile() API.\n\n5. In Figure 3, a 224x224 image can result in 196 tokens, while a 320x320 image corresponds to 400 tokens. After dropping 50% tokens, it has 200 tokens longers than 196 tokens. Do you use padding to make them have the same length in experiments? \n\n6. Line 237 mentions a patch size of 4x16x16. What does the 4 represent? An image generally has 3 RGB channels. Shouldn't it be 3x16x16? Line 238 says image resolution is 4x256x256, which also confuses me.", "limitations": "The experiments seem not to show whether pretraining with one modality can help improve another modality's downstream performance. For example, whether the image pretraining datasets help improve the downstream video action recognition results and whether using the video pretraining datasets can boost the downstream performance on ImageNet-1K.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper provides a generic multimodal training solution involving input, model, and objective and successfully trained models with a combination of image, video, and audio datasets.\n\n2. It achieves SOTA zero-shot results on video action recognition.\n\n3. The paper is well-written and easy to follow.", "weaknesses": "1. Ablation studies show that alternating between the objectives on each step is better than combining them by summing them. Is this true for all cases or only for large-scale pretraining? Many papers on various topics sum up multiple losses for training. For example, in semi-supervised learning, the loss general consists of two parts: supervised loss and unsupervised loss. We may need to give a proper context for this conclusion.\n\n2. The method's two technical components, alternating gradient descent and mixture of experts, mainly follow previous works. This paper integrates them and applies them to a larger scale of pretraining.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for a very detailed review, we will try our best to address your comments.\n\n>Ablation studies show that alternating between the objectives on each step is better than combining them by summing them. Is this true for all cases or only for large-scale pretraining?\n\nThis is a good point, and we are careful to note that this experiment is limited to the case of multimodal multitask pretraining, especially in the context of classification and contrastive losses. However, we should also note that we provide a theoretical basis in addition for our empirical results (see lines 124-127), i.e., it is proven that the gradients of tasks applied across multiple optimization steps are either equal or more marginally convex than summing the gradients, which can reduce the difficulty of the optimization problem. This is what motivated our initial experiment and works very well for our model setup, but we do not claim that alternating training would result in better downstream performance than summing gradients on any possible combination of tasks. We would argue that this is not a weakness, but simply another broad direction which is out of scope of the paper.\n\n>The method's two technical components, alternating gradient descent and mixture of experts, mainly follow previous works. This paper integrates them and applies them to a larger scale of pretraining.\n\nWe propose a novel generalization of AGD which has not been explored in other works. Due to technical limitations of previous implementations of the approach, we were unable to find other methods training efficiently on many modalities each with their own dynamic input shapes; the most popular techniques tend to waste computation on padding which we found to be a significant source of inefficiency. By carefully constructing computation graphs that leverage primitives such as a compilation cache in a modality agnostic way, we can maximize effective training FLOPs utilization across distributed accelerators. Our accelerated AGD approach is also what unlocks the efficient use of multi-scale (multi-resolution) data, which to our knowledge, is a novel training approach applied to multimodal training not observed in other works.\n\n>Algorithm 1 shows the sampling uses step t and model state S_t, but Lines 142-143 say that sampling is directly proportional to the number of examples in each dataset. That is, the more samples a task has, the more likely the dataset gets sampled. It would be better to make the algorithm consistent with the descriptions.\n\nIn this case, f(t, S_t) would simply compute over a probability distribution independent of t and S_t. We intentionally left this open to try with alternative sampling algorithms, such as [29], but we found the global sampling strategy to work well enough in combination with our other improvements. If you would prefer we reword this, we can do so.\n\n>Moreover, in Algorithm 1, should there be another mini-batch sampling given a chosen task? Should the sampling function f be only for task/dataset sampling?\n\nWe consider the function f to only be for choosing the dataset-objective pair, the sampling function is operating on the current time step t and the optionally the previous model state. Once a task is chosen, then a minibatch is sampled from that dataset as is usually done.\n\n>Line 227 mentions filtering examples close to the downstream tasks. How is the filtering conducted?\n\nWe use perceptual hashing as used in similar works (e.g., [46]) to make sure that duplicate images or frames are removed if they appear in the validation set to avoid data leakage. We can update with this information.\n\n>Why the ablation study in Section 4.3 uses spatial resolution 224 while the main results in Table 1 use resolution 256 during training?\n\nThese ablations were mainly to try to closely match the resolutions and datasets as seen in literature so they can be more easily compared and reproduced. It is only our large-scale pretraining that we break from this kind of setup to try to achieve the best performance.\n\n>I'd like to know whether Pytorch can also do so.\n\nYes, actually. We mostly speak to jax because our implementation is based on it. But analagously, it is possible to use `torch.compile` to construct an AGD training pipeline to achieve a similar effect, along with our `rematerialization` and `scan-over-layers` optimizations we mentioned.\n\n>Do you use padding to make them have the same length in experiments?\n\nNo, we don't use padding. The sequence lengths are fairly close enough so we have roughly the same compute/memory usage to not be an issue. In practice these two inputs would be compiled separately before alternating training on them.\n\n>Line 237 mentions a patch size of 4x16x16. What does the 4 represent?\n\nThe 4 represents the temporal axis, so 4 frames are patched together in the same way the pixels are. But this also requires all inputs to be multiples of 4 frames, so in the case of images we inflate/tile them to 4 frames. This means that images are basically treated as 4-frame videos with no motion, from a modeling perspective.\n\n>The experiments seem not to show whether pretraining with one modality can help improve another modality's downstream performance. For example, whether the image pretraining datasets help improve the downstream video action recognition results and whether using the video pretraining datasets can boost the downstream performance on ImageNet-1K.\n\nWe have results in Figure 2 in the supplementary material which indicates that integrating multiple modalities like image and video into a dense model tends to harm the accuracy of the other modalities. However, with Figure 5 in the main paper, we see that MoEs reverse this trend and improve. Most notably, the addition of video data actually improves the zero-shot performance of the image datasets, which was not true of the dense model.\n"}, {"review_id": "qVILAGKz5y", "reviewer": "Reviewer_SjW7", "summary": "This work proposes a solution to train multimodal multi-task training and demonstrate the model trained via this strategy outperforms other approaches in zero-short learning problems. The proposed solution mainly relies on current solutions including AGD, JAX library, DropToken, etc. it seems that this work can be helpful in multimodal related works.  ", "questions": "1) Could the author demonstrate the major difference compared to AGD? In line 119, the author mentioned that it’s a more generic solution based on AGD, however, not clear about the innovation and difference.\n2) How can the proposed approach be applied on these cases where one or two modalities are missing? \n3) Will the proposed approach handle the situation like training with several different datasets together? And how to deal with the difference  among them? Is there any other additional steps needed? ", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1) The proposed training strategy and solution can help related work in multimodal research area. \n2) The training strategy demonstrated better performance on zero-short learning approach over other compared works. ", "weaknesses": "1) If this work focuses on the training strategy, more studies on downstream tasks are needed. \n2) Although the work focuses on the training solution, more analysis of multimodal experimental results analysis are needed. For example, the analysis of the improvement over other approaches,  especially what kind of cases the new approach help improve the performance? \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for providing a review of our paper, here we try to answer some of the concerns and comments about the paper.\n\n>If this work focuses on the training strategy, more studies on downstream tasks are needed.\n\nWhat kind of downstream tasks are you looking for? We provide evaluation of zero-shot ImageNet, CIFAR, Kinetics, UCF, HMDB, and ESC in our main results table, which should already provide a large number of downstream benchmarks. Our ablations furthermore look at additional benchmarks in context of the various modalities. We also do not solely focus on the training strategy, but also highlight the importance of using the MoE architecture as a way to integrate many modalities into one model without harming performance, and that is confirmed in our results across many downstream evaluations and ablations.\n\n>Although the work focuses on the training solution, more analysis of multimodal experimental results analysis are needed. For example, the analysis of the improvement over other approaches, especially what kind of cases the new approach help improve the performance?\n\nWe show specifically the improvement of our approach when comparing various aspects of the integration of different parts of our training strategy in various figures and tables (multimodal multitask AGD, Mixture-of-Experts, multi-scale resolution, etc.). We also compare with other similar models (for image-text MoE, we compare with LIMoE, for the previous SoTA on zero-shot video classification, we compare with VideoCoCa). See also our supplementary material for additional findings and results where we explain which methods are most important for our modeling and training strategies. If there's a more specific set of experiments or evaluations that is missing that weaken our claims, please report them so we can address them.`\n\n>Could the author demonstrate the major difference compared to AGD? In line 119, the author mentioned that it’s a more generic solution based on AGD, however, not clear about the innovation and difference.\n\nWe propose a novel generalization of the approach which has not been explored in any other work. Due to technical limitations of previous implementations of AGD, we were unable to find other methods training efficiently on many modalities each with their own dynamic input shapes; the most popular techniques tend to waste computation on padding which we found to be a significant source of inefficiency (and graph compilation is quite important to be able to fully saturate the computation on an accelerator). By carefully constructing computation graphs that leverage primitives such as a compilation cache in a modality agnostic way, we can maximize effective training FLOPs utilization across distributed accelerators. Previous instantiations of AGD only considered a special case where all of the inputs are fixed in shape/size. We show that this feature is quite important especially to be able to train efficiently on many different modalities at the same time. This is a point we may need to emphasize more strongly, we can provide more results on the training efficiency when using our accelerated AGD vs. prior approaches.\n\n>How can the proposed approach be applied on these cases where one or two modalities are missing?\n\nThis is the core part of our accelerated AGD technique. Because each step only optimizes on a single dataset and alternates training between them, we can define a dataset with only the modalities we want to specify for that step. For example, if we want to optimize on audio+video data, we can compile a computation graph on audio data as input and computes classification, another computation graph that has video data as input and also computes classification, and a third computation graph that computes audio-video contrastive loss. So long as there exists a dataset definition that matches the modalities/objectives of at least one of the computation graphs, then it can be used in training. So if we have any dataset with missing data/modalities, we just simply leverage the computation graph compatible with the losses where that modality is missing and sample that as a new task.\n\n>Will the proposed approach handle the situation like training with several different datasets together? And how to deal with the difference among them? Is there any other additional steps needed?\n\nYes, this is precisely our approach, and the major reason why we can obtain state-of-the-art zero-shot results. Without being able to train multiple datasets with different subsets of modalities we would not be able to make our approach work. We explain in section 3 how defining datasets in separate compiled computation graphs can allow for efficient multitasking, only computing on the data that is available, without relying on any expensive padding or gradient masking that other works typically rely on. For a concrete example, we might start training step 0 by computing the loss for a video-text-audio contrastive task on VideoCC, followed by an image classification task on JFT on training step 1, and then an image-text contrastive task on CC12M on step 2. We can schedule these datasets/objectives in any order because the computation graphs for each are cached after they are encountered for the first time, this is all automatic due to the `jax.jit` implementation (and this should also work with `torch.compile`, but we have not implemented this)."}]

Author rebuttal:

Meta-review: The paper received mixed reviews with 3 borderline accepts and 1 reject. Given the borderline nature of the reviews, the AC reviewed the whole paper carefully and would rank it as 7. After further discussions with the SAC, the AC recommends the paper for acceptance. Most reviewers (including the AC) agreed that the AGD + unified encoder + MoE design provides an interesting solution to multi-modal training of modality-agnostic encoders. The paper also provides extensive experimental evaluation on a variety of different benchmarks, achieving impressive results, especially on zero-shot action recognition. Given the strong results obtained using an under-studied type of architecture, the AC and SAC agree that the paper should be accepted.

With that being said, the paper needs to be revised in the final version. One important concern raised by tAnb is the fact that when a single model is used (instead of MoE), sharing the model across modalities often hurts performance. While MoEs address this issue, the result still raises questions about the use of modality-agnostic encoders. A thorough discussion of these tradeoffs (modality agnosticity vs. model size vs. performance) would benefit the paper. One suggestion from the AC would be to use hybrid architectures with modality-specific stems followed by modality-agnostic modules. The authors could then assess the model with increasing levels of modality-agnostic computation.
Finally, please take into account other comments and suggestions for the camera-ready version. For example, consider moving the experiments supporting contribution 5 (in the intro) to the main manuscript (as opposed to supplementary material), as this is a critical finding of this work.
Decision: Accept (poster)