Module-wise Adaptive Distillation for Multimodality Foundation Models
Authors: Chen Liang
Jiahui Yu
Ming-Hsuan Yang
Matthew Brown
Yin Cui
Tuo Zhao
Boqing Gong
Tianyi Zhou
Venue: NeurIPS 2023 poster (2023)
Keywords: Multimodality foundation models
knowledge distillation
TL;DR: We propose a module-wise adaptive distillation method for compressing multimodality models.

Abstract: Pre-trained multimodal foundation models have demonstrated remarkable generalizability but pose challenges for deployment due to their large sizes. One effective approach to reducing their sizes is layerwise distillation, wherein small student models are trained to match the hidden representations of large teacher models at each layer. Motivated by our observation that certain architecture components, referred to as modules, contribute more significantly to the student's performance than others, we propose to track the contributions of individual modules by recording the loss decrement after distillation each module and choose the module with a greater contribution to distill more frequently. Such an approach can be naturally formulated as a multi-armed bandit (MAB) problem, where modules and loss decrements are considered as arms and rewards, respectively. We then develop a modified-Thompson sampling algorithm named OPTIMA to address the nonstationarity of module contributions resulting from model updating. Specifically, we leverage the observed contributions in recent history to estimate the changing contribution of each module and select modules based on these estimations to maximize the cumulative contribution. We evaluate the effectiveness of OPTIMA through distillation experiments on various multimodal understanding and image captioning tasks, using the CoCa-Large model \citep{yu2022coca} as the teacher model.

Module-wise Adaptive Distillation for Multimodality Foundation Models
Chen Liang
Jiahui Yu Google Research j jiahiiyu@oogge.. com
Georgia Tech cli aliann73@gatech.. edu
Ming-Hsuan Yang UC Merced, Google Research minghsuan@google. com
Matthew Brown Google Research
Yin Cui
NVIDIA Research
mtbr@google.com com
ri richardaeen@@gaiilcc com
Tuo Zhao
Boqing Gong Google Research bgongogoogle.com com
Tianyi Zhou
Georgia Tech
University of Maryland, College Park t tiany senddimdd..... edu
tourzhao@gatech.edd
Abstract
Pre-trained multimodal foundation models have demonstrated remarkable general- izability but pose challenges for deployment due to their large sizes. One effective approach to reducing their sizes is layerwise distillation, wherein small student models trained to match the hidden representations of large teacher models at each layer. Motivated by our observation that certain architecture components, referred to modules, contribute more significantly the student's performance than others, we propose to track the contributions of individual modules by record- ing the loss decrement after distillation each module and choose the module with a greater contribution to distill more frequently. Such an approach can be naturally formulated multi-armed bandit (MAB) problem, where modules and loss decrements are considered arms and rewards, respectively. We then develop modified-Thompponr sampling algorithm named OPTIMA address the non- stationarity of module contributions resulting from model updating. Specifically, we leverage the observed contributions recent history to estimate the changing contribution of each module and select modules based on these estimations to maximize the cumulative contribution. We evaluate the effectiveness of OPTIMA through distillation experiments various multimodal understanding and image captioning tasks, using the CoCa-Large model [48] as the teacher model.
Introduction
Large pre-trained multimodality foundation models have demonstrated remarkable generalizability on wide range of tasks, e.g., image classification, image-text retrieval, and visual question answering [27, 18, 49, 46, 48). These models often contain several architectural components (referred to as modules), each acquiring knowledge of one more modalities through pre-training. For example, (27) introduce a dual-encoder architecture composed of an image encoder and a text encoder, which are jointly pre-trained align an image with the corresponding text. This process equips each encoder with both the unimodal representation power and the crossmodal alignment capability. [48] further add multimodal decoder on top of the dual-encoder architecture. The decoder learns joint visual and textual representation and obtains the multimodal understanding knowledge. However, the sizes of these models have reached billions of parameters, and the number of modules will further scale to
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

accommodate new modalities. This poses significant challenge for deployment in applications with latency and storage constraints.
Layerwise distillation is powerful approach to compressing large models (i.e., teacher models) into small ones (i.e., student models) with minimal loss of performance [16, 30J. This approach trains a student to match the hidden representations of the teacher at each layer. Given that the teacher's s layerwise representations often contain rich semantic knowledge, they can significantly improve the student's generalizability [19, 45, 11]. Many existing works have demonstrated the effectiveness of this layerwise strategy in task- specific distillation, where a student distilled from a teacher that has been fine-tumed on the target task [38, 17].
In task- distillation in multimodality models, however, matching the student's representations with those of the teacher at every layer does not always henefit the student's performance. Prior research has shown that improving the representations of specific modality (i.e., image, text, or multimodality) tends to yield better fine-tuning results than others [51]. Therefore, we hypothesize that the representations of such a modality may contain more relevant knowledge to the target task. If the student focuses on mimicking these representations, it is more likely to achieve better performance on the target task. To verify this hypothesis, distill multiple students from a CoCa-Large teacher [48], each only matching the teacher's ' layerwise representations in a single module (i.e., image encoder, text encoder, or multimodal decoder). As shown in Figure 1, distilling a specific module yields significantly better performance than distilling others. If all modules are distilled with equal weighting, interference from other modules may affect the training of this specific module, thereby compromising the student's fitting on the target task (Figure 3).
NLVR2
coco Caption
120
Test Accuracy (%%
66.9
105.1
61.5
CIDE
100
60-
80
75.6
50
59.5
Img
Txt
Multi
Img
Txt Modules
Multi
Modules
Figure I: Performances of CoCa- 1oCy-1 students distilled from a CoCa-Large teacher [48] on NLVR2 [37] and Microsoft COCO Caption [6]. Each bar represents the performance obtained by matching only the layerwise representations in a single module. See Details in Section 4.
To mitigate this issue, we propose to track the individual modules' contributions to the distillation process on the target task and choose a module to distill at each step based their contributions. In this way, we can more frequently distill the module with greater contribution. To evaluate the contribution of module, we distill it for specific number of steps and observe the resulting ratio of decrement in the distillation loss. To track its contribution continually, we repeat this procedure throughout training. Hence, need to explore different modules by repeatedly evaluating their contributions and, meanwhile, exploit the module with the greatest contribution by distilling it more frequently. Given limited budget for training steps, face significant challenge in balancing exploration and exploitation.
To address this challenge, we adopt multi-armed bandit (MAB) approach [36, 41J. MAB targets an online decision-making problem that chooses an action (i.e., arm) and observes a reward in each round to maximize the cumulative reward over rounds. MAB algorithms are designed for balancing the evaluation of all arms and choosing the arm that maximizes the cumulative reward, hence can be leveraged to evaluate and choose modules. As illustrated in Figure 2, we consider each module an arm and every P steps as one round. In each round, we choose a module to distill for P steps and evaluate its contribution as the reward. As more rewards observed, we obtain more accurate estimate of the underlying reward distribution for each module, which reflects the actual contribution of this module. Based on the reward distribution, choose modules to maximize the cumulative loss decrement.
However, the actual contribution of module may change with model updating. For example, distilling the multimodal decoder may become more beneficial when input, i.e., the image and text representations, are sufficiently optimized [39]. To estimate such stattotoonaay underlying reward

---

Multi-armed Bandit Algorithm
Round "Distill
Round 1:; "Distill Multimodal Decoder"
Round 2: DDistill Text Encoder"
Image Encoder"
Laaai
Cinain
Cunin
Chidn
Lustn
Multimodal Decoder
Dhidh Caxan
Crida
Cuaza
Image Encoder
Text Encoder
Student
+
Teacher
Image
Figure 2: An illustration of OPTIMA.
distribution, we need to count more on rewards observed in recent history. This makes it challenging for adopt classical MAB algorithms because of them designed for stationary underlying reward distributions [4, 10, 40J. Therefore, they average the observed rewards over the whole history estimate the underlying reward distribution. In our case, since the model is unlikely to change drastically within few steps, and so is the underlying reward distribution, we tailor these stationary algorithms by replacing the simple average with the exponential moving avvrrage the observed rewards. By discounting old rewards, the reward distribution can track the changing contribution in recent history and provide stable and up-to-date guidance for module selection [13, 15, 28). By adapting MAB algorithms to the non-stationary environment of multimodality model distillation, we propose OPTIMA (Module-wise Adaptive Distillation with Multi-arm Bandit), novel task- specific distillation method for multimodality foundation models. We exemplify OPTIMA using the Thompson Sampling algorithm for its strong empirical and theoretical performance [32, 5, 31, 2]. We remark that OPTIMA can be generically combined with other MAB algorithms.
We demonstrate the effectiveness of OPTIMA on various multimodality understanding and im- age captioning benchmarks [14, 47, 37, 6]. Specifically, we distill ToCa-Tinyya (102M) and CoCa-Tinys (55M) from a CoCa-Large teacher (672M. [48]). powerful pre-trained multimodality foundation model. For both students, OPTIMA substantially surpasses the layerwise distillation baseline. Moreover, CoCa-Tiny12 outperforms CoCa-Bac Base three out of four tasks with only 1/3 its number of layers. Extensive analysis verifies that OPTIMA can track the changing contributions of different modules and choose modules based on their contributions.
Preliminaries
Architectures of Multimodality Foundation Models can be generally categorized into dual-encoder [27, 18, 49, 23], encoder-decoder [42, 46, 44], and Contrastive Captioners (CoCa, [48]). Dual- encoder models contain an image encoder and a text encoder. The two encoders take images and texts as inputs, respectively, and are jointly pre-trained to align image and relevant text. This enables each encoder to Iearn both the unimodal representations and the crossmodal alignment knowledge. Encoder-decoder models contain multimodal decoder top of an image encoder. The decoder learns a joint visual and textual representation and obtains the multimodal understanding knowledge. To integrate the knowledge from the image encoder, the text encoder, and the multimodal decoder, CoCa combines all modules in a single Transformer and trains them jointly. Specifically, the multimodal decoder takes the representations generated by both the unimodal encoders and learns multimodal representations with the cross-attentiog mechanism.
Knowledge Distillation trains small model (i.e., student model) to match the output prediction of large and well-trained model (i.e., teacher model) by minimizing their prediction discrepancy [16]. Specifically, denote the teacher model as f(:) and the student model as (0.), and consider the

---

following optimization problem:
min (0. + (0,, 0c),
0.
where Ltrain (fz) is the prediction loss on the target task, c.g., the cross-entropy loss for a classification task; (0..0... 81) is the KL-Divergence between the probability distributions over their predictions, ie., ( (0=) f((),)) In large Transformer based models, distilling knowledge from only the output predictions neglects the rich semantic knowledge in the intermediate layers. To leverage such knowledge, researchers have proposed to match the hidden representations and attention scores at each layer of the teacher and the student [30, 19, 11, 45).
Thompson Sampling (TS) is a widely-used MAB algorithm with Bayesian assumption [32, 40, 31, 5, I]. TS assumes prior distribution the parameters of the underlying reward distribution for each arm. To estimate the underlying reward distribution, TS updates a posterior reward distribution (often referred to as reward distribution) for each arm based on the statistics of the observed rewards in history. TS draws an arm from the posterior in each round to maximize the cumulative rewards in the long term. A typical example of TS is to set both the prior distribution and the reward distribution to be Gaussian distributions. In this case, the reward distribution an arm is specified as
D =
n +++1
where j is often set as the average of the observed rewards in history and n is set as the number of rounds the arm has been chosen.
3 Method
Problem Formulation. We consider both the teacher and the student models to be multimodality Transformers containing > 1 modules. For example, they are CoCa models containing e 3 modules in this work: an image encoder, a text encoder, and multimodal decoder.
We construct K = 2 - arms, each associated with a unique, non-empty subset of modules selected from the e modules. The k-th arm is formed subset. We set T T" rounds, where T' is the total number of training steps and 2 is the number of steps within each round. In this way, choosing the k-th arm in the t-th round is equivalent distilling S₁ for steps in the range of ((t-1) P) P].
set of model layers, denoted as Sx, in its associated
For each arm k E [K], we denote its reward distribution as Dr and set it as a Gaussian distribution specified as (Hk,, nT+ following Eq. 2. To set proper prior, first randomly choose each arm for To 2 0 rounds', then initialize rewards over To rounds. We explain why setting the reward distribution as Gaussian distribution is feasible in Appendix A.I.
S4 and nı = To. where μo is the average of the observed
Remark 3.1 We consider all possible combinations of modules because, at certain training stage, there might exist a combination such that distilling it leads to a higher reward than distilling any single module.
Remark 3.2 Since the dependencies among modules are unclear, we should assume all arms to be dependent and model the reward distribution as a multivariate Gaussian over all arms. However, sam- pling from such a distribution requires decomposing the covariance matrix, which is computationally costly. To improve efficiency, we ignore the dependency assumption. In practice, we observe no clear disadvantage in the model performance.
OPTIMA Algorithm. We initialize
n4 = To and D3 = N ((M
1
. pkt
nh +
for all arm & E [K].
At the t-th round, we sample reward Pse for the k-th arm from its reward distribution Dj for all k E [K]. The arm with the highest sampled reward is then chosen for this round, denoted as ar. In this way, the arm is chosen according the belief of it being the best arm.
'2 is included in T.

---

We then play the a,-th arm by distilling Sa, for P steps. At any step E ((t t)P, we first compute the distillation losses on Sa. Specifically, we denote the hidden representations at the i-th layer of the teacher and the student as Hi E TR fgl and H¹ E Rl where d, and d, denote the hidden dimension and |r|| denotes the sequence length. Then the distillation loss on hidden representations is defined as
hidn (0()) 0.)
Σ [Sail KESa.
MSE H; H.Whhin)
where MSE(.,)) -) is the mean-squared error, and W hidu Rd,xdd linear projection that projects Hi into the same space Hi. Similarly, the distillation loss on attention scores is defined as²
xde is randomly initialized and learnable
Cain (0,)
Σ iES
MSE(A) 4/,A()
S
where Ai, Ag E R|x] are the attention score matrices averaged over the multiple attention heads the i-th layer. The student model then optimized based on weighted sum of all distillation losses on Say i.e.,
votal
Ltrain + 0 DRL. + a22 + 03 Cai
(5)
where Lirain and DKL are defined in Eq. and a1, a2, 0g 2 0 are hyper Specifically, the student model is updated with an SGD-ttype algorithm [21]:
g 1) + 8(() Vo, Sto (0()) 0,).
where 1) > 0 the learning rate. The same procedure is then repeated for P steps.
We then compute the reward r for playing the a,-th arm. Specifically, we design the reward as the averaged ratios of decrements over all types of distillation losses:
Σ max (0.5103 C(0((-11))) ())()))))))))))))) ((0!P))
Faz
U
C(0((C-1))))
(6)
E()JEU
where U = (P),()), Luidn (-), Cattn (-)}. Note that Lhida (-) and Laut (-) are defined in Eq. 3 and Eq. 4, respectively, but by replacing Sa, with the set of all layers. They measure the change in the distillation loss the full model; otherwise, the reward definition changes with the pulled arm and thus cannot provide an evaluation metric consistent across all arms.
The numerator is the difference between C(0)((-1))) and C(O! This loss change is accumulated for P steps to reduce the uncertainty caused by noises in the stochastic gradients. As a result, it can reliably reflect the contribution of S., to the distillation performance of the full model.
We consider both the changes in output prediction discrepancy (-KL(-)) and the layerwise represen- tation distances (Chidn(-) and (.).. This improves the robustness of the reward over different tasks because a specific distance metric may consistently well reflect the distillation performance on all tasks (Figure 7).
For all types of losses, we compute the ratio of the loss change to the loss in the previous round. Normalizing the loss changes prevents the reward from being biased by the scales of different losses. We further clip each ratio by zero to bound the reward between (0, 1), which is desirable for achieving a good performance guarantee in the Thompson Sampling (TS) algorithm. Finally, we average the rewards over all types of losses.
After computing the reward Ya., we update the reward distribution Da. Specifically, we update the distribution mean, Ma., as the exponential moving average of the observed rewards in the past:
Ma: f 7/a0c + (1 - 7)
We omit the "(t')" superscript on (H), W Whddheee hes., to simplify the notations. We assume the teacher and student have same number layers. The case of having different layers will elaborated Section 4.3. We omit the "(t')" superscript on (At)es, to simplify the notations.
5

---

where y E (0, 1) is hhppr-yprameeee parameter. We then increase Ru, by one and update Da, following Eq. 2. Since the exponential moving average discounts the old rewards, Pa. tends to reflect the average of the rewards within recent rounds. Recall that the actual contribution of each module changes dynamically. By using Ma, as the distribution mean, Da, can capture the changing contribution and provides up-to-date guidance for module selection. The complete algorithm is shown in Alg. 1
Algorithm 1 OPTIMA: Module Adaptive Distillation
Input: T: the number of total rounds; P: the number of steps each round; K = 2 1: the number of arms; Ty: discounted factor. 0 0₃: the initialized student and the teacher models.
2: Output: 0!²))
3: I' y0=0
0. 72g = 1kk = 0'4k = I,..,, K. T do
4: for
5: for k = L,.., K do
6: Sample reward for arm Fx ~ N (ui: m²+1)
7: end for
8: Select arm O₁ f- arg max; fk.
9:
10:
for p porp=1-- 1, P do
( (0(-), 02) following Eq. 5.
11:
Compute Ciotal
0. 0(1'+1) ( Vo, Coo (0.") 2,(,,,,) ,0)) - 1'+1.
12:
13:
14:
end for
15:
16:
Compute Ta, following Eq. 6.
17:
Update Ha, following Eq. 7.
18:
T& +- 7kk + L.
19:
end for
4
Experiments
We verify the effectiveness of OPTIMA popular multimodal understanding and image captioning benchmarks.
4.1 Data
We conduct task-specific distillation on three multimodal understanding tasks: visual question answering (VQA, [14]), visual entailment (SNLI-VE. [47]), and visual reasoning (NLVR2, [37]). VQA to answer an input question with a sentence based on an input image. is formulated as classification problem with 3129 classes where each class corresponds to answer. SNLI-VE a thyee-way classification problem to predict whether an input sentence entails the input image, contradicts it, is neutral it. NLVR2 is binary classification problem predict whether input sentence is true or false about an input image pair. We further train and evaluate the model using the Microsoft COCO Caption dataset [6] and the Karpathy-test split, respectively. The task is to describe an input image with a sentence. See Appendix A.2 for details.
4.2 Models
We evaluate OPTIMA on CoCa [48]. It is a powerful Transformer-basee multimodality foundation model pre-trained web-scale alt-texts [18] and image-text pairs [50]. Since CoCa consists of an image encoder, a text encoder, and a multimodal decoder, provides a sufficiently large action space for OPTIMA. Furthermore, the decoder produces multimodal representations, which allows to be directly fine-tuned without any extra step of multimodal pre-traine training, commonly needed for dual-encoder models [35, 49,
We assume To 0 for simplicity of demonstration.

---

For each target task, we fine-tune a pre-trained CoCa-Large as the teacher (672M, 48 layers (24/12/12 layers in image/tex/mmliimdaal modules), d, 1024). We then distill task-specific student from the fine-tuned teacher. Specifically, we consider two architectures: CoCa-Tiny (101 .8M, 12 layers (6/3/3 layers), d, = 768) and CoCa-Tinys (54.5M, layers (3/1/2 layers), d,- 768). The layers in each module of the student are initialized with the layers uniformly sampled from the corresponding module of a pre-trained CoCa-Base (293M, 36 layers, 768).
4.3 Implementation Details
Teacher Initialization. We fine-tune pre-trained CoCa-Large as the teacher for each target task. For multimodal understanding tasks, we attach a randomly initialized task-specific linear classifier on top of the multimodal decoder for answer prediction. We fine-tune both the model and the classifier using the cross- ensropy loss. For the image captioning task, we directly fine-tune CoCa with the captioning
loss [48], i... Leap = > log Pe. (ye (yelMee,, I). We follow [48] for all fine-tuning Hyper-parametert (See Appendix A.3 for details).
Task-Specific Distillation. We fix the fine-tuned teacher and distill a student using OPTIMA on each task. For all tasks, we train the student for T 100k steps. We use Adafactor with decoupled weight decay [34] as the optimizer with 3 = (0.9, 0.999) and a learning rate of 1 x 10 3 with a linear decay schedule. We match the i-th layer of a student's module with the [gi] -th layer of the corresponding teacher's module, where is the ratio of the number of layers of the two modules. We randomly initialize Whide hidn E RR. xd. We set an 0, 02 1 and 0g = 1 10-2 for all tasks. For OPTIMA, set y 0.98, Tu 10, P = 100 and T T 1k. Full details are deferred to Appendix A.4.
4.4 Baselines
Pre-trained Vision-Language Models (VLMs). To compare with models with similar scales, we present the fine-tuning performance of existing pre-trained VLMs: UNITER [7], OSCAR [24], ViLT [20], VILLA [12), CLI.- p [35] and ALBEF [23]. Different from foundation models, VL.Ms often contain single backbone, c.g., 12- 12-layer Transformer. The backbone takes both image features and text embeddings as inputs and learns the multimodal representations. They then use an auxiliary model to predict the input image features, e.g., CLIP-ViL. ViLp uses pre-trained CLIP encoder [27], ALBEF uses a pre-trained ViT-B/16 [8], and the rest use a pre-trained Faster R-CNN [3]. We exclude methods that outperform the CoCa- Large teacher.
Multimodality Distillation Methods. We further compare OPTIMA with existing multimodality distillation methods: MiniVLM [43], DistilVLM [11] and DIDE [45]. Different from OPTIMA: 1) All methods conduct distillation the pre-training stage. Pre-training distillation significantly improves the student's s generalizability but requires much more computational resources than task- specific distillation. 2) All methods consider VLMs as the teachers. VLMs are much smaller than CoCa models, and small teacher-student capacity gap often improves the effectiveness of distillation [26, 25J. However, VLMs are less generalizable and can only distill students on limited tasks. MiniVLM adopts knowledge distillation additional unlabeled data (Eq. 1). DistilVLMM adopts layerwise distillation. DIDE distills a dual-encoder student from a VLM teacher by aligning cross- modal attentions. We ignore concurrent methods with teacher that significantly outperforms the CoCa- Large teacher.
4.5
Main Result
Table summarizes the evaluation results of layerwise distillation and OPTIMA. We report the median over five random seeds". On both CoCa-Tinys and CoCa-Tiny12. OPTIMA can achieve consistent and notable improvements upon layerwise distillation over all tasks. The gains are most prominent in NLVR2 and COCO, e.g., the gains on CoCa- CoCa-Tiny12 are 0.9 and 4.4, respectively. This may be explained by Figure 3, which shows the contributions of modules in these tasks exhibit high variances and the training of a specific module can impair the training of others. The gain is slightly smaller in CoCa- Tiny, than CoCa-Tiny:n This might be because CoCa-Tiny converges slower, and that specific module is less likely to over- fit quickly and interfere with others.
'The standard deviations reported in Appendix A.5.
7

---

Compared with pre-trained VLMs at the same scale, CoCa-Tinyı2 (OPTIMA) can achieve simi- lar performance with a smaller size. Compared with CoCa-Base OPTIMA can achieve notable improvements on three out of four tasks with only 1/3 its number of layers.
Compared with MiniVLM, DistilVIM and DIDE, OPTIMA achieve significantly better perfor- mance on CoCa- CoCa-Tiny12 and similar performance on CoCa-Tinye over all tasks. Recall that these baseline methods use web-scale pre-training data for distillation and single-backbone VLMs as teachers. In contrast, OPTIMA uses limited target task data and faces much larger capacity gap, e.g.. 5 times larger for CoCa-Ting Tiny12. Nevertheless, the performance gaps between CoCa-Tiny12 (OPTIMA) and CoCa-Large are comparable to those in the baseline methods. Furthermore, while baseline methods demonstrate gains on limited tasks, OPTIMA shows well-rounded gains on all tasks.
Table The evaluation results of OPTIMA and baseline methods on VQA, SNLI-VE, NLVR2 and COCO Caption. The results of all baselines reported from the original papers. The inference speedup is computed with respect to the batch-averaged inference time of the corresponding teacher model averaged across all tasks. We present the sizes of backbone parameters because the embedding sizes vary largely across methods depending on the implementation details. All sizes are counted based on the released checkpoints from the authors.
Method
Teacher
Params. (million)
Inference Speedup
VQA test-std
SNLI-VE test
NLVR2
COCO
test-p
CIDEr
UNITER-Bass [7]
146.5
72.9
78.3
77.9
OSCAR-Bass [24]
146.5
73.4
137.6
ViLT-Base [20]
85.6
76.4
76.1
VILLA-Buse [12]
146.5
73.7
79.0
79.3
CLIP-ViL, [35]
187.7
74.1
79.0
127.9
ALBEF [23]
241.0
74.7
80.3
80.5
CoCa-Base [48]
69.2
83.6
80.2
126.7
CoCa-L.arge [48]
672.1
75.3
85.6
82.6
132.3
MiniVI.M12x3s [43]
ASCAR-Base
30.0
3.1x
68.1
115.0
DssilVLM22xx88 [11]
OSCAR-Base
30.0
3.1x
69.2
117.1
DIDE,2 [45]
ViLT-Base
171.3
3.0x
69.2
76.3
75.6
CoCa-Tiny:
CoCa-Large
54.5
6.2x
68.7
82.0
76.5
112.2
CoCa-Ting Tinys (OPTIMA)
CoCa-Large
54.5
6,2x
69.0
82.3
77.0
113.5
CoCa-Tiny12
CoCa-Large
101.8
3,4x
71.2
83.9
80.4
116.8
CoCa-Tinyia (OPTIMA)
CoCa-Large
101.8
3.4x
71.6
84.3
81.3
121.2
Analysis
In this section, we verify that OPTIMA improves the student's performance, tracks the actual contributions of all modules, and chooses arms based on their contributions. We further investigate the design of the reward function in Appendix A.6 and study the discounted factor of the non- stationary reward distribution in Appendix A.7.
5.1 OPTIMA Improves the Student's Performance
Figure compares the performance of the students obtained by 1) fxed--am distillation: distilling fixed arm over rounds, 2) random-arm distillation: randomly choosing an arm to distill in each round, and 3) OPTIMA. The performance of fixed-arm distillation varies largely across tasks. Distilling an arm that contains more modules than other arms does not always improve the student performance, e.g., "Img+Txte "Img+Txt+Multi" achieves 80.4 while "Img+Multi" achieves 81.1 in NLVR2. In contrast, OPTIMA demonstrates superiority over all fixed-arm and random-arm distillation baselines. This suggests that OPTIMA can find better weighting for arm selection for each task.
Figure shows the prediction loss and the output layer prediction discrepancy (i.e., Lerain and DKL. defined Eq. 1) on the dev set are both lower in OPTIMA than in layerwise distillation. This suggests that the student achieves better generalization performance on both the distillation and the prediction of the target task.

---

SNLI-VE
COCO Caption
90
Test Accuracy (95)
$4.0
$3.9
$4.3
81.1
30.4
$1,3
L16.7
116.816.0
2212
82.2
Accuracy
76.0
105.1
CIDE:
80-
78.6
70-
1000
60-
61.5
80
75.6
30
51.7
60-
62.9
log.
Txt
Molic
Img Imgthdoli
Tat
Tax Molis
Random
OPTIMA
Txt
Muli
Imgehhuli
TartMalis
TarMiiii
Random
OPTMM
Txt
Muli
Imse Thd
Exe-bduli
Rondon
OPTIMA
Img
Img
Iose
img-
Arm Selection Metbodk
Arm Scloction Methods
Arm Selection Methods
Figure 3: A comparison of the performance of the students by 1) distilling a fixed arm over rounds (shown in solid colors, where each arm is denoted by its associated subset of modules); 2) randomly choosing an arm to distill in each round (denoted by "Random"); and 3) OPTIMA.
NLVR2
COCO Caption
Caption
1.00
Methods
0.75
LWD
1300
OPTIMA
Drgt
Curin
0.50
1200
0.25
50
100
50
100
25
50
75
100
25
50
75
100
Steps
Training Steps
Trinning Steps
Training
Figure 4: The prediction loss and the output layer prediction discrepancy (i.e., Ltrain and DKL defined Eq. I, respectively) OPTIMA and layerwise distillation (denoted by "LWD").
5.2 Reward Distributions Reflect the Contributions of Arms
Figure (Left) shows the means of the reward distributions of all arms through training. In all tasks, "Imm+Tx++Mulli" quickly dominates the others, while "Multi" and "Tx++Muld" remain incompetent. The reward distributions of all arms slowly evolve through training. For example, in COCO, the arms containing the image encoder all show a non-increasing trend in the later stage of training, while "Txt" shows an increasing trend.
Figure 5 (Right) verifies such evolving reward distributions correctly reflect the changing contributions of modules. We can observe a strong and positive correlation between the mean of the reward distribution of and the accuracy obtained by distilling this fixed arm through training.
SNLL-VE
Caption
Act.
0.3
0.5-
0.075-
Ings Txa Muln
1.0
0.2-
0.000
IngeMalti
Inse
a.5
0.1-
0.025-
Ing
NLVR2
SNLI-VE
Caption
0.0-
0
50
100
0.0
30
100
0.000
50
100
MuliM
25
$0
75
100
Training Steps
Training Steps
Training Sleps (k)
Training Sleps (k)
Figure 5: Left: The means of the reward distributions of all arms (i.c., { through training. Right: The correlation between the means of the reward distributions of individual arms and the test accuracy obtained by distilling these arms through training.
K
5.3 OPTIMA Choose Arms Based on Reward Distributions
Figure 6 visualizes the frequency of choosing each arm through training. In all tasks, we observe that "Img+T Mxt+Multi" is the most frequently chosen arm. In COCO, the frequency distributions across arms are flatter than in SNLI-VE and NLVR2, which is consistent with Figure 5 that the means

---

of the reward distributions are more concentrated across arms. This suggests OPTIMA can choose arms based on the reward distributions.
100
VQA
SNL--VE
NEVR2
COCD Caption
50
Training Staps
60
40
004
-04
20
0.2
02
02
0.2
Img-
Img Txx
Imp for (Malk
Imp
long Tak-
bog logghMas
TariMale-
Imp Te MrrMa
Imn
Tue
Mall
Iogg Tak-
Iogg
Tariaal:
Img Img-TTa brrMak
Img
Imge Ta
long logrMMal..
Tor Matee
lins Tho (aaat.
Arms
Ams
Arms
Figure 6: A visualization of the frequency of each arm being chosen through training. A deeper color represents a higher frequency.
6 Discussion
Application to Other Multi-module Multimodality Models. In this study, we primarily demonstrate the efficacy of OPTIMA when applied models within the CoCa family. This initial success signifies that OPTIMA could be applicable in scenarios where one or more modules predominantly contribute the target task, a phenomenon also observed in other multi-module models [35]. The exploration of the application of OPTIMA to other multi-module models is left for future research.
Furthermore, OPTIMA does not incur computational and memory costs exceeding those of layerwise distillation, as it calculates only the necessary layerwise losses and gradients for distilling subset of modules. This efficiency enables the exploration of OPTIMA's applicability larger foundational models- - those incorporating diverse modules to accommodate not only image and text modalities but also others such as audio and video.
Design Fine-grained Action Space. this study, we divide the CoCa model coarsely into three modules based modalities, and design the action space encompass all possible module com- binations. However, different layers within a single module could exhibit variable contributions to the target task. For example, lower layers tend to capture high- freeueryy signals crucial for feature extraction, while upper layers typically learn low- www-frequency signals, vital for semantic understand- ing. By further subdividing each module into groups of layers, encompass all possible combinations of layer groups, potentially enabling the finding of better arms. Additionally, this enhanced granularity enables the extension of method to kingle-module VLMS and single-modality models, such as large language models (LLMs) and vision models.
refine the action space to
However, directly applying OPTIMA encounters challenges in this fine-grained scenario. Firstly, the dependencies among layer groups can substantially exceed those among modalities, conflicting with the current sampling strategy which assumes minimal dependencies between arms (See Remark 3.2 in Section 3). Secondly, expanding the action space prolongs exploration time and, consequently, increases training costs. Given these challenges, adapting OPTIMA to such fine-gra aineggrieed scenario necessitates modifications in the sampling strategy and explorations into more efficient reinforcement learning algorithms, aspects we designate for future research.
7 Conclusion
We propose OPTIMA, a module-wise adaptive distillation method for multimodality models. We employ a simple MAB algorithm to demonstrate that distillation based on the contributions of modalities represents a promising and intriguing research direction. We posit that adopting more sophisticated reinforcement learning algorithms can yield greater improvements, and designate such explorations for future research.
10

---

8 Acknowledgements
This project was completed during a full-time internship at Google Research. We thank all collabora- tors from the Brain team and the Perception team for their insightful feedback and intellectual support. We also extend our thanks to the TPU team for providing abundant computational infrastructure and resources.
References
[1] Shipra Agrawal and Navin Goyal. Analysis of thompson sampling for the multi-armed bandit problem. In Conference on learning theory, pages 39-1. JMLR Workshop and Conference Proceedings, 2012. 3,4
[2] Shipra Agrawal and Navin Goyal. Near-optimal regret bounds for thompson sampling. Journal of the ACM (JACM), 64(5): 1-24, 2017.
[3] Peter Anderson, Xiaodong He, Chris Buehler, Damien Teney, Mark Johnson, Stephen Gould, and Lei Zhang. Bottom-up and top-down attention for image captioning and visual question answering. In Proceedings of the IEEE conference on computer vision and puttern recognition,
pages 6077-6086, 2018.
[4] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-256, 2002.
[5] Sébastien Bubeck and Che- Yu Liu. Prior-free and dreorddep regret bounds for thompson sampling. Advances in neural information processing systems, 26, 2013. 3, 4
[6] Xinlei Chen, Hao Fang, Tsung- Yi Lin, Ramakrishna Vedantam, Saurabh Gupta, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco captions: Data collection and evaluation server. arXiv preprint arXiv: 1504. 00325, 2015. 2, 3, 6,
[7] Yen-Chun Chen, Linjie Li, Licheng Yu, Ahmed El Kholy, Faisal Ahmed, Zhe Gan, Yu Cheng, and Jingjing Liu. Uniter: Universal image-text representation learning. In European conference on computer vision, pages 104-120. Springer, 2020. 7,8
[8] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, 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 ar XIXv:2010.11929, 2020.
[9] Zi-Yi Dou, Yichong Xu, Zhe Gan, Jianfeng Wang, Shuohang Wang. Lijuan Wang, Chenguang Zhu, Pengchuan Zhang, Lu Yuan, Nanyun Peng. et al. An empirical study of training end-to-de sssion-and-lane Lad-lanu transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 18166-1811 1666--8176666 2022. 6
[10] Eyal Even-Dar, Shie Mannor, and Yishay Mansour. Pac bounds for multi-armed bandit and markov decision processes. In International Conference on Computational Learning Theory, pages 255-270. Springer, 2002. 3
[11] Zhiyuan Fang, Jianfeng Wang, Xiaowei Hu, Lijuan Wang, Yezhou Yang, and Zicheng Liu. Com- pressing visual-linguistic model via knowledge distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 1428-1438, 2021. 2, 4, 7,8
[12] Zhe Gan, Yen-Chun Chen, Linjie Li, Chen Zhu, Yu Cheng, and Jingjing Liu. Large-scale adversarial training for vision-and and-language representation learning. Advances in Neural Information Processing Systems, 33:66 16 6-6628, 2020. 7,8
[13] Aurélien Garivier and Eric Moulines. On upper-confidence bound policies for non-stationary bandit problems. arXiv preprint arXiv:0805..341 2008. 3
[14] Yash Goyal, Tejas Khot, Douglas Summers-Stay, Dhruv Batra, and Devi Parikh. Making the v in vqa matter: Elevating the role of image understanding in visual question answering. In Proceedings of the IEFE conference on computer vision and puttern recognition, pages 1904-69133 2017. 3, 3,6,1 15
11

---

[15] Neha Gupta, Ole-Christoffer Granmo, and Ashok Agrawala. Thompson sampling for dy- namic auuli--rmedd bandits. In 2011 10th International Conference on Machine Learning and Applications and Workshops, volume 1, pages 484-489, IEEF, 2.11.
[16] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in neural network. arXiv preprini arXiv:/503.000 1503.02531, 2015. 2,
[17] Lu Hou, Zhiqi Huang, Lifeng Shang. Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic BERT with adaptive width and depth. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. 2
[18] 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 International Conference Machine Learning, pages 4904-4916. PMLR, 2021. 3,
[19] Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv:1909100351, 2019. 2,4
[20] Wonjae Kim, Bokyung Son, and Ildoo Kim. Vilt: Vision-and-languuae transformer without convolution or region supervision. In International Conference on Machine Learning, pages 5583 5594. PMLR, 2021. 7.8
[21] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. 5
[22] Taku Kudo. Subword regularization: Improving neural network translation models with multiple subword candidates. arXiv preprint arXiv: 10959, 2018. 15
[23] Junnan Li, Ramprasaath Selvaraju, Akhilesh Gotmare, Shatiq Joty, Caiming Xiong, and Steven Chu Hong Hoi. Align before fuse: Vision and language representation learning with momentum distillation. Advances in neural information processing systems,
2021. 3, 7, 8
[24] Xiujun Li, Xi Yin, Chunyuan Li, Pengchuan Zhang, Xiaowei Hu, Lei Zhang, Lijuan Wang, Houdong Hu, Li Dong, Furu Wei, et al. Oscar: Object-semantics aligned pre-training for languaee tasks. In European Conference on Computer Vision, pages 121-137. Springer, 2020. 7, 8
[25] Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao. Less more: Task-aware layer-wise distillation for language model compression. arXiv preprint arXiv:2210.01351, 2022. 7
[26] Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. Improved knowledge distillation via teacher assistant. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 5191-5198, 2020. 7, 16
[27] 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 International Conference on Machine Learning, pages 8748-8763. PMLR, 2021. 3,
[28] Vishnu Raj and Sheetal Kalyani. Taming non-stationary bandits: A bayesian approach. arXiv preprint arXiv:1707..9777 2017.
[29] Steven J Rennie, Etienne Marcheret, Youssef Mroueh, Jerret Ross, and Vaibhava Goel. Self- critical sequence training for image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 7008-7024, 2017. 15
[30] Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and 2014. 2, 4
Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint
12

---

[31] Daniel Russo and Benjamin Van Roy, Learning to optimize via posterior sampling. Mathematics of Operations Research, 39(4):1211-1243, 2014. 3,
[32] Steven Scott. A modern bayesian look at the multi-armed bandit. Applied Stochastic Models
Business and Industry, 26(6):639 2010. 3,4 4
[33] Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint urXiv: urXiv:1508.07909, 2015. 15
[34] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604, PMLR, 2018.
[35] Sheng Shen, Liumian Harold Li, Hao Tan, Mohit Bansal, Anna Rohrbach, Kai- Chang, Zhewei Yao, and Kurt Keutzer. How much can clip benefit vision-and sision-and-lanne and-language tasks? arXiv preprint arXiv:2107.06383, 2021. 6, 7, 8, 10
[36] Aleksandrs Slivkins et al. Introduction to multi-armed bandits. Foundations and Trends in Machine Learning, 12(1-2):1-286, 2019. 2
[37] Alane Suhr, Stephanie Zhou, Ally Zhang, Iris Zhang, Huajun Bai, and Yoay Artzi. A corpus for reasoning about natural language grounded in photographs. arXiv preprint arXiv: /81 1a:::8114400 1.00491, 2018. 2, 6 6
[38] Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint arXiv:1908.09:555 2019. 2
[39] Zhiqing Sun, Hongkun Yu, Xiaodan Song. Renjie Liu, Yiming Yang. and Denny Zhou, Mobilebert: a compact task-agnostic bert for resource-limited devices. arXiv preprint arXiv:2004.02940 029844, 2020.
[40] William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-94 1933. 3, 4
[41] Joannes Vermorel and Mehryar Mohri. Multi-armed bandit algorithms and empirical evaluation.
In European conference on machine learning, pages 437-448. Springer, 2005. 2
[42] Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan. Show and tell: A neural image caption generator. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 315 3156-3164, 2015.
[43] Jianfeng Wang, Xiaowei Hu, Pengchuan Zhang, Xiujun Li, Lijuan Wang. Lei Zhang, Jianfeng Gao, and Zicheng Liu. Minivlm: A smaller and faster vision- model. arXiv preprint arXiv:2012.069440 2020. 7,8
[44] Peng Wang, An Yang. Rui Men, Junyang Lin, Shuai Bai, Zhikang Li, Jianxin Ma, Chang Zhou, Jingren Zhou, and Hongxia Yang. Unifying architectures. tasks, and modalities through simple sequence-to-sequence learning framework. arXiv preprint arXiv:2202.03052. 2022. 3
[45] Zekun Wang, Wenhui Wang, Haichao Zhu, Ming Liu, Bing Qin, and Furu Wei. Distilled dual-encoder model for vision-language understanding. arXiv preprint 1rXiv:2112.08723, 2021.
2, 4, 7.8
[46] Zirui Wang, Jiahui Yu, Adams Wei Yu, Zihang Dai. Yulia Tsvetkov, and Yuan Cao. Simvlm: Sim- plc visual language model pretraining with weak supervision. arXiv preprint arX9v:2108.10904, 2021. 3.3
[47] Ning Xie, Farley Lai, Derek Doran, and Asim Kadav. Visual entailment: A novel task for fine-grained image understanding. arXiv preprint 1990iv/901.000 067006, 2019.
[48] Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. arXiv preprint 275.v2200.001917, 2022. 2, 6, 15
13

---

[49] Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, al. Florence: A new foundation model for computer vision. arXiv preprint 2r
2021. 1,3, 6
[50] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transform- ers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12104-12113, 2022. 6
[51] Pengchuan Zhang, Xiujun Li, Xiaowei Hu, Jianwei Yang, Lei Zhang, Lijuan Wang. Yejin Choi, and Jianfeng Gao. Vinvi: Revisiting visual representations in vision-language models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 5579-5588, 2021. 2
14

---

A Appendix
A.1
Gaussian the Posterior
Here we explain why setting the reward distribution as a Gaussian distribution is feasible. Recall that consider non-stationary environment. We design the reward distribution model the rewards observed in recent history. Within small number of steps, the model unlikely to change drastically, Therefore, the observed rewards, which are computed as the ratios of decrement in loss, are unlikely to be skewed. As a result, Gaussian is a feasible choice.
A.2 Data
For the VQA task, we conduct downstream fine-tuning and testing on the VQA 2.0 dataset [14], which consists of 83k images and 444k questions for training, 41k images, and 21 4k questions for validation. For the image captioning task on COCO, we use [6] for training and testing. It contains 11k images for training and 5k images for validation and 5k images for testing.
A.3 Teacher Model Implementation Details
We fine-tune a pre-trained CoCa-Large model [48] as the teacher. contains 672M parameters in Transformer layers and contains total of 787M parameters including the embedding size. It contains 24 layers in the image encoder, and 12 layers in the text encoder, and 12 layers in the multimodal decoder. We use vocabulary size of 64k. We use 576 as the image resolution and 18 as the patch size for image inputs. We use 64 as the max sequence length for text inputs. We follow [48] for fine-tuning parammearrmeeers for all tasks, as listed in Table 2.
Table 2: Hyper-parameters parameters for fine-tuning CoCa-Large teacher models.
Hyper-parameters
VQA
SNLI-VE NLVR2 COCO Caption
Optimizer
Adafactor with Decoupled Weight Decay
Adam A8
(0.9. 0.999)
Gradient Clipping
1.0
Learning Rate Schedule
Linear Schedule Decaying to Zero
Warm-up Steps
Ik
Weight Decay Rate
0.1
Pooler Learning Rate
5 x 10-4
x 10-3
5 10-3
N/A
Encoder Learning Rate
2 x 10-5
x 10-5
2 x 10-5
x 10-5
RandAugment
1,10
1,10
None
None
Training Steps
100k
50k
50k
50k
Batch Size
128
128
Dropout of Task Layer
0.5
0.5
0.5
N/A
For multimodal understanding tasks, we follow [48] to apply an attentional pooler with a single query extract embedding from the decoder output, and train a linear classifier on top of the pooled embedding. For NLVR2, we construct two input sequences, each containing the concatenation of the description and one image. The two output representations are further concatenated as the input to the classifier. For image captioning, we apply the captioning loss proposed in [48]. We do not use the CIDEr metric-specific optimization [29]. We use a greedy strategy for decoding.
A.4 Distillation Implementation Details
For each task, we distill a CoCa-Tiny12 student and a CoCa-Tinya student from a fine-tuned CoCa- Large teacher that task. CoCa-Tiny12 contains 102M parameters in the Transformer layers and contains a total of 152M parameters including the embedding size. CoCa-Tinya contains 55M parameters in the Transformer layers and contains total of 105M parameters including the embedding size. We use 576 as the image resolution and 18 as the patch size for image inputs. To tokenize text input, the sampled pre-training dataset. We use 64 as the max sequence length for text inputs. We follow [48] for fine-tuning parameparreeters for all tasks, listed in Table 3.
sentence-pies piece model [33, 22] with vocabulary size of 64k trained on
We conduct a two-stage distillation for CoCa-Tinye. We first distill CoCa- Tiny12 from CoCa-Large then use the distilled CoCa- 1 i the teacher to teach CoCa-Tinyy -inyy- Existing works have shown
15

---

that introducing intermediate-sion sized teacher reduces the gap between the teacher and the student model, which allows the distillation be more effective [26].
We distill the model for a total of T' steps, i.e., a total of 7 = T /P rounds. Among the 7 rounds, the first To K rounds are used for initialization of the parameters of the reward distribution.
Table 3: Hyper-parameters for distilling CoCa-Tiny student models.
Hyper-parameters
VQA SNLI-VE
NIVR2 COCO Caption
To
10
P
100
1000
0.98
as
(0.0, 1.0, 1 x 10-3)
Optimizer
Adafactor with Decoupled Weight Decay
Adam 5s
(0.9. 0.999)
Gradient Clipping
1.0
Learning Rate Schedule
Lincar Schedule Decaying to Zero
Learning Rate
10-3
Warm-up Steps
Weight Decay Rate
0.1
RandAngment
1,
1,
None
None
Training Steps (T')
125k
100k
100k
100k
Batch Size
128
384
256
256
Dropout of Task Layer
0.5
0.5
0.5
N/A
A.5 Statistics of Experimental Results
We report the median of five random seeds for experiment results on CoCa-Tiny12 and CoCa-Tinye. Table 4 show the standard deviations of the experimental results in Table 1.
Table 4: The standard deviation of the experimental results in Table 1.
Method
VQA Acc
SNLI-VE
NLVR2
COCO Caption
Acc
Acc
CIDEr
CoCa-Tmye
0.20
0.25
0.11
0.15
CoCu-Tinye (OPTIMA)
0.22
0.13
0.35
0.15
CoCa-Tinny12
0.17
0.23
0.15
0.88
CoCa -o (OPTIMA)
0.08
0.15
0.30
0.32
A.6 Design of Reward
Recall that design the reward (Eq. 6) the averaged ratio of loss decrements three types of distillation losses: DKL (Eq. 1), Luidn (Eq. 3) and Lastn (Eq. 4). Figure 7 compares ours with two variants: 1) FKD: the ratio of loss decrement of DKL; 2) FLWD: the averaged ratio of loss decrements over Chida and Lattn- We can observe that FKD performs better than FLWD in NLVR2 but reversely in COCO. Since captioning tasks often rely more on contextual knowledge the layerwise representations than classification tasks, the layerwise representation distance may better characterize the distillation performance in COCO. By taking both distance metries into consideration, TOPTIMA performs well on both tasks.
A.7 Design of the Reward Distribution
Recall that we design the mean of the reward distribution (Eq. 7) the exponential moving average (EMA) of the past rewards. Figure 8 shows a hyper-parameter study on the halflife of the EMA, computed as logg"" > Halflife is the number of rounds the EMA decays by one-half. We can observe that a too-large or too-small halflife, meaning that counting too many old rewards or counting only instantancous rewards can both be harmful to the student's performance. This corroborates that the actual contribution each module non-stationary in the long term and stationary in the short term, and using EMA with appropriate y correctly track the changing contribution.
16

---

NLVR2
Caption
SNLI-VE
Test Accuracy (%)
82
130
Test Accurccy (%5
12
121.2
81-
8
CIDEr
120-
8q.3
110
108.4
83.5
80-
100
TeD
Reward Types
FLWD FOРТIА
Reward Types
TLAND TOPTIA
Hllliiife EMA (Rounds)
20
60
Figure 7: A comparison of three variants of reward: 1) TKD: the ratio of decrement of DKL.; 2) FIWD: the averaged ratio of decrement over Lhidu and Lattn; 3) TOPTIMA-
Figure 8: A the halflife of the EMA, computed as logg7
study on
17


Reviews:
[{"review_id": "7AZ1JLVIqz", "reviewer": "Reviewer_1KcH", "summary": "This paper proposes a Module-wise Adaptive Distillation method, OPTIMA, for a multimodal foundation model based on a modified multi-arm bandit algorithm. Comparisons with baseline multimodal pre-trainings and multimodal distillation methods on four different multimodal datasets demonstrate the effectiveness of the proposed method. ", "questions": "\n\n\n- In Fig.4, the loss on the dev set of the NLVR2 dataset is increasing, yet the performance of the distilled model is good (Table 1, the performance of CoCa-Tiny12 (OPTIMA) is close to the one of CoCa-Large). Could you provide a more in-depth discussion concerning this interesting behavior?\n- In Fig.5, the behavior of rewards seems to be somehow inconsistent with the results in Fig.3. For example, on the NLVR2 dataset, the reward of \"Img+Txt\" is always higher than \"Img\"; however, the test accuracy of \"Img+Txt\" is considerably lower than the one of \"Img.\" Do you have an explanation?\n- In this paper, the modules are coarse-grained, corresponding to the individual modality or the fusion. Would considering fine-grained modules be beneficial? ", "limitations": "Limitations have been addressed in the paper.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "\n\n- the paper is well-written and the idea is clearly presented.\n- the idea of using multi-arm bandit to guide distillation is interesting and the effective, as shown in the experiments.", "weaknesses": "\n- Concerning equation 6, I understand that for the sake of fair comparison, the authors choose to compare the loss change of the full model. However, for example, is that reasonable to use the change in $\\mathcal{L}_{hidn}$ of \"Img\" to evaluate the reward of the \"Txt\" module in the case when we only distill the \"Txt\" module in the current round.\n- the discussion should be elaborated; please look at the Questions below.\n- it is mentioned in the introduction that one primary motivation of this paper is that \"interference from other modules may affect the training of this specific module.\" However, later results did not reflect this point. In most cases, the best way is to distill all three modules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for the detailed comments and valuable suggestions. We hope our response addresses your concerns. Should you have additional questions, we are readily available for further discussion.\n\n**Is that reasonable to use the change in $L_{hidn}$ of \"Img\" to evaluate the reward of the \"Txt\" module in the case when we only distill the \"Txt\" module in the current round?**\n\nPulling one arm will influence the full model performance. Since our goal is to achieve a good full model performance, it is reasonable to design the reward of pulling one arm to reflect its influence on all modules. \n\nIn the reviewer’s example, distilling (pulling) “Txt” will change $L_{train}$ and $L_{kd}$. The change will further influence the update of parameters in “Img”, so $L_{hidn}$ of “Img” will change accordingly. \n\nSuppose distilling “Txt” increases $L_{hidn}$ of “Img” (i.e., causes negative interferences to “Img”) and further degrades the full model performance. In such a case, the reward of distilling “Txt” should be decreased as it negatively influences the full model performance. If we fail to take into account the increment in $L_{hidn}$ of “Img” in the reward computation, the reward would remain high. Consequently, the algorithm would keep distilling “Txt”, which may further degrade the full model performance.\n\n**One primary motivation of this paper is that \"interference from other modules may affect the training of this specific module.\" However, later results did not reflect this point.**\n\nWe provide an analysis for Figure 3 in lines 272-276: “Distilling an arm that contains more modules than other arms does not always improve the student performance, e.g., “Img+Txt+Multi'' achieves 80.4 while “Img+Multi” achieves 81.1 in NLVR2, …”. This implies that the “Txt” module causes some negative interferences to the training of other modules.\n\n**In Fig.4, the loss on the dev set of the NLVR2 dataset is increasing, yet the performance of the distilled model is good. Could you provide a more in-depth discussion?**\n\nWhile the prediction loss decreases quickly in the early stage of training (which implies that the student can quickly overfit the final-layer labels), the layerwise distillation loss gradually dominates. Hence, the optimization prefers to minimize the dominating layerwise loss, and therefore reduces the fitting of the student model to the final-layer labels, resulting in an increasing prediction loss. \n\nLayerwise distillation enforces the student to mimic all intermediate-layer semantics of the teacher. Such a strong regularization provides the student with more generalizable knowledge, leading to a good performance in Table 1.\n\n**The reward of \"Img+Txt\" is always higher than \"Img\"; however, the test accuracy of \"Img+Txt\" is considerably lower than the one of \"Img.**\n\nWe believe that the reviewer misread Figure 3. “Img+Txt” (orange) has a higher accuracy than “Img” (light blue). We will make the presentation clearer in the next version.\n\n**In this paper, the modules are coarse-grained, corresponding to the individual modality or the fusion. Would considering fine-grained modules be beneficial?**\n\nWe believe that using fine-grained arms would be beneficial because it may help us find better arm(s) through time. However, this idea requires further investigation. The dependency across arms will increase for fine-grained modules, but the current OPTIMA assumes a low dependency across arms for sampling efficiency. To adapt OPTIMA to more fine-grained arms, alteration on the sampling strategy is needed. Moreover, increasing the number of arms also leads to more expensive exploration costs. We will leave this idea for future research.\n\n"}, {"review_id": "0290Hxv6vH", "reviewer": "Reviewer_zaxt", "summary": "They use layer-wise distillation to reduce model size for multi-modal foundation models. They develop a modified-Thompson sampling algorithm named OPTIMA to address the non-stationarity of module contributions resulting from model updating.", "questions": "No.", "limitations": "No.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The problem is valuable to reduce size for large models. The paper is good in writing. The analysis is interesting. ", "weaknesses": "The result is based on COCA, experiments on other foundation models are needed. The improvement over baseline is not large.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**The result is based on COCA, experiments on other foundation models are needed.**\n\nThe commonly adopted multimodal foundation models are often computationally expensive to train and test. Given our limited computing resources, we were constrained to allocate them solely to a specific family of models, namely CoCa. Despite this limitation, we present a comprehensive analysis and results focused on CoCa.\n\nRegrettably, we have since lost access to the requisite computing resources, and are unable to provide further results for other model families given the short response/discussion window. While we agree with the reviewers that exploring the application of the module-adaptive idea to other models is of great interest to the community, we need to leave such investigations to future research. This limitation will be further highlighted in our paper.\n\nNevertheless, CoCa has demonstrated powerful and consistent performance over a diverse set of downstream tasks, thereby implying the broad applicability and efficacy of our method. Furthermore, the improvements on CoCa also indicate the promise of our idea on other multi-module models. We have shown that our idea is effective when a specific module is more useful to a given task than others, which is a shared observation in other multi-module models [1].\n\n[1] Shen, Sheng, et al. \"How much can clip benefit vision-and-language tasks?.\"\n\n**The improvement over baseline is not large.**\n\nThe improvements are both non-trivial and consistent. For example, our experiments on CoCa-Tiny are conducted using five random seeds. Compared with the layerwise distillation baseline (CoCa-Tiny in Table 1), we achieve a median of 0.9 points of gain out of 80.4 on NLVR2 and 4.4 points of gain out of 116.8 on COCO. On VQA with over 3k classes, we improve the accuracy by 0.4 points out of 71.2. Furthermore, we provide analysis to verify that our method works promisingly as designed: the proposed reward can correctly reflect the contribution of the arms, and the arms are chosen based on the reward distribution. \n\nWe remark that our major contribution is to demonstrate that module-wise adaptive distillation, which distills multi-modality models based on the contributions of modalities, is a promising and interesting research direction. We use a MAB algorithm to demonstrate this point, because MAB is simple, efficient, and good at balancing exploration and exploitation. While we believe that greater improvements can be achieved with more advanced RL algorithms, we will leave the algorithmic explorations for future works."}, {"review_id": "peVDPs4Vun", "reviewer": "Reviewer_8iuR", "summary": "This work proposed a new knowledge distillation method for multimodule multimodality models. This work pointed out that different modality modules contribute differently to the model performance, and it proposed to sample different modality modules for distillation during training. The paper proposed a multi-armed badit based method to adaptively select modules during distillation. The method is applied to CoCa model and shows better performance than baselines on multiple datasets", "questions": "Pleae see the weakness part", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1.\tThe paper is well motivated. It makes sense and is inspiring to show that different modules have different contributions in the knowledge distillation process. Adaptively selecting different modules for distillation could lead to better performance.\n\n2.\tThe proposed method shows better performance than the baselines. There are thorough ablation study to show the performance of different module selection strategies, which helps the reader to better understand the effectiveness of the method.\n", "weaknesses": "1.\tI have some questions about some details. In Table 1, how is the CoCa-Tiny baselines finetuned? Are they finetuned by Image+Text+Multi distillation?\n\n2.\tIt is hard to find a fair comparison in Table 1. The methods are based on different structures and different teacher models.\n\n3.\tThe only fair comparison seems to be the comparison with CoCa-Tiny at the third section in Table 1. However, as explained in Q1, one thing is that it is not clear how these beselines are obtained, another is that the improvements seem to be marginal, mostly within 0.5%.\n\n4.\tIn Figure 3, it also shows that the proposed method is only slightly better than the trivial Image+Multi selection strategy, while being much more complicated.\n\n5.\tThe method is only applied to CoCa. It would be better to show the effectiveness on more models, such BLIP2.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**In Table 1, how is the CoCa-Tiny baselines finetuned? Are they finetuned by Image+Text+Multi distillation?**\n\nYes, CoCa-Tiny baselines are fine-tuned by “Image+Text+Multi” distillation, where all modules are simultaneously distilled.\n\n**It is hard to find a fair comparison in Table 1. The methods are based on different structures and different teacher models.**\n\nThe comparison with CoCa-Tiny at the third section in Table 1 is strictly fair. The comparisons with MiniVLM, DistilVLM and DIDE are not fair. The results are from their original papers, which are achieved under very different experimental settings from ours (e.g., different model architectures and training data). \n\nWe have re-implemented MiniVLM to align with our experimental setting[1] and achieved 48.5 on VQA, 72.1 on SNLI-VE, 52.0 on NLVR2 and 79.6 on COCO. We can observe that OPTIMA achieves consistent gains over MiniVLM. \n\nDistilVLM adopts the layerwise distillation method, and the re-implementation of DistilVLM in our experimental setting will become the CoCa-Tiny baseline at the third section in Table 1. \n\nWe are unable to re-implement DIDE to align with our experimental setting because DIDE distills a dual-encoder's student from a single-module teacher. The method is architecture-specific and thus cannot use CoCa-Large as the teacher. \n\nWe choose to present the results from their original papers because: First, it is practically infeasible to align DIDE with our experimental setting. Second, it is unfair to present the result of a method outside its original experimental setting when that setting is one major contribution of this method. For example, MiniVLM proposed to collect external unlabeled data for distillation, which is one major contribution to its performance. In terms of algorithm, MiniVLM simply adopts knowledge distillation.\n\nAlthough the results are not fair to compare with, we present these methods because we believe it is important to acknowledge relevant methods. In Section 4.4, we describe their methodologies and summarize their strengths and weaknesses.\n\n[1] First, we replace the teacher and student in MiniVLM with the same teacher and student as in OPTIMA. Second, the original MiniVLM conducts distillation on additional unlabeled pre-training data (Open Images V6) in the vision-language pre-training stage. As our goal is to compare the distillation in the fine-tuning stage, for a better controlled experiment, we only use pre-trained models without training on the unlabelled pre-training data. \n\n**The improvements seem to be marginal, mostly within 0.5%.**\n\nThe improvements are both non-trivial and consistent. For example, our experiments on CoCa-Tiny are conducted using five random seeds. Compared with the layerwise distillation baseline (CoCa-Tiny in Table 1), we achieve a median of 0.9 points of gain out of 80.4 on NLVR2 and 4.4 points of gain out of 116.8 on COCO. On VQA with over 3k classes, we improve the accuracy by 0.4 points out of 71.2. Furthermore, we provide analysis to verify that our method works promisingly as designed: the proposed reward can correctly reflect the contribution of the arms, and the arms are chosen based on the reward distribution. \n\nWe remark that our major contribution is to demonstrate that module-wise adaptive distillation, which distills multi-modality models based on the contributions of modalities, is a promising and interesting research direction. We use a MAB algorithm to demonstrate this point, because MAB is simple, efficient, and good at balancing exploration and exploitation. While we believe that greater improvements can be achieved with more advanced RL algorithms, we will leave the algorithmic explorations for future works.\n\n**In Figure 3, it also shows that the proposed method is only slightly better than the trivial Image+Multi selection strategy, while being much more complicated.**\n\n“Image+Multi” is only slightly worse than OPTIMA as it is one of the best arms. Although directly distilling the best arm is less complicated than OPTIMA, it requires us to know in advance which arm is the best. Searching for the best arm requires an exhaustive search over all arms because the best arm for one task does not necessarily perform uniformly well on all other tasks. Such an exhaustive search introduces significant complications and computational overheads.\n\nIn contrast, OPTIMA is a clean, task-adaptive solution. It requires only a single run (much cheaper than the exhaustive search) and yields a better performance.\n\n**The method is only applied to CoCa. It would be better to show the effectiveness on more models, such BLIP2.**\n\nThe commonly adopted multimodal foundation models are often computationally expensive to train and test. Given our limited computing resources, we were constrained to allocate them solely to a specific family of models, namely CoCa. Despite this limitation, we present a comprehensive analysis and results focused on CoCa.\n\nRegrettably, we have since lost access to the requisite computing resources, and are unable to provide further results for other model families given the short response/discussion window. While we agree with the reviewers that exploring the application of the module-adaptive idea to other models is of great interest to the community, we need to leave such investigations to future research. This limitation will be further highlighted in our paper.\n\nNevertheless, CoCa has demonstrated powerful and consistent performance over a diverse set of downstream tasks, thereby implying the broad applicability and efficacy of our method. Furthermore, the improvements on CoCa also indicate the promise of our idea on other multi-module models. We have shown that our idea is effective when a specific module is more useful to a given task than others, which is a shared observation in other multi-module models [1].\n\n[1] Shen, Sheng, et al. \"How much can clip benefit vision-and-language tasks?.\""}, {"review_id": "hIIQQmq0P4", "reviewer": "Reviewer_fQyt", "summary": "This paper proposes a novel knowledge distillation (KD) method for multimodal foundation models. Motivated by the observation that distilling different modules yields distinct performance, the authors propose to dynamically choose an optimal combination of the modules for KD. Specifically, the choice of the optimal combination is formulated as a multi-armed bandit (MAB) problem. Experimental results demonstrate that the proposed method outperforms other KD baselines.", "questions": "The questions are listed in the \"weaknesses\" section.", "limitations": "The authors have discussed the limitations of the proposed method.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors introduce an interesting topic on the selection of submodules for KD with multimodal models. Recently, one similar approach has been proposed for distilling image classification models [1].\n2. The paper is clearly written and easy to follow.\n3. The ablation study appears to be sufficient.\n4. The experimental results shows the effectiveness of the propsed method.", "weaknesses": "\n1. **Limitation:** The model architectures adopted in this paper are limited to CoCa-base models. However, several multimodal large models have been proposed recently, such as [2-4]. I believe that if the proposed method can be applied to other foundation models, it will be of great interest to the community. Besides, I am curious to know if the proposed method can be applied to settings where the teacher and the student have different architectures.\n2. **Experiment:** What is the training cost of the proposed method compared to other baselines, e.g., training time, GPU memory usage, etc.?\n\n> [1] Song, Jie, et al. \"Spot-adaptive knowledge distillation.\" IEEE Transactions on Image Processing 31 (2022): 3359-3370.\n> \n> [2] Yang, Ziyi, et al. \"i-code: An integrative and composable multimodal learning framework.\" Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 9. 2023.\n>\n> [3] Lu, Jiasen, et al. \"Unified-io: A unified model for vision, language, and multi-modal tasks.\" arXiv preprint arXiv:2206.08916 (2022).\n>\n> [4] Shen, Yongliang, et al. \"Hugginggpt: Solving ai tasks with chatgpt and its friends in huggingface.\" arXiv preprint arXiv:2303.17580 (2023).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for the valuable suggestions. We hope our response addresses your concerns. Should you have additional questions, we are readily available for further discussion.\n\n**The model architectures adopted in this paper are limited to CoCa-base models.**\n\nThe commonly adopted multimodal foundation models are often computationally expensive to train and test. Given our limited computing resources, we were constrained to allocate them solely to a specific family of models, namely CoCa. Despite this limitation, we present a comprehensive analysis and results focused on CoCa.\n\nRegrettably, we have since lost access to the requisite computing resources, and are unable to provide further results for other model families given the short response/discussion window. While we agree with the reviewers that exploring the application of the module-adaptive idea to other models is of great interest to the community, we need to leave such investigations to future research. This limitation will be further highlighted in our paper.\n\nNevertheless, CoCa has demonstrated powerful and consistent performance over a diverse set of downstream tasks, thereby implying the broad applicability and efficacy of our method. Furthermore, the improvements on CoCa also indicate the promise of our idea on other multi-module models. We have shown that our idea is effective when a specific module is more useful to a given task than others, which is a shared observation in other multi-module models [1].\n\n[1] Shen, Sheng, et al. \"How much can clip benefit vision-and-language tasks?.\"\n\n**I am curious to know if the proposed method can be applied to settings where the teacher and the student have different architectures.**\n\nOPTIMA can be easily applied to settings where the teacher and student models have different widths and heights. This can be achieved by modifying the dimension of $W_{hidn}$ and layer mapping function. \n\nThe cases where the teacher and student models contain different numbers of modules require further investigation. In such cases, we might need to re-define the arms. For example, if the teacher is a CoCa model and the student is dual-encoder VLM, we may consider the following module mapping as arms: Img->Img, Txt->Txt, Img+Txt->Img+Txt, Img+Multi->Img, Txt+Multi->Txt, Img+Txt+Multi->Img+Txt. Due to the lack of access to computing resources and codebases, we are unable to supplement experimental results under such cases. We will leave the investigations to future works.\n\n**What is the training cost of the proposed method compared to other baselines, e.g., training time, GPU memory usage, etc.?**\n\nOPTIMA incurs a computational cost no more than the layerwise distillation baseline (CoCa-Tiny in Table 1) because OPTIMA only needs to compute the layerwise distillation losses and the corresponding gradients with respect to a subset of modules at each training iteration. OPTIMA does not introduce GPU memory overheads.\n\nAs mentioned in line 239-241, all other distillation baselines conduct distillation in the pre-training stage. In contrast, OPTIMA conducts distillation in the fine-tuning stage, which requires much less computational resources."}, {"review_id": "uFExQNQq3B", "reviewer": "Reviewer_3XdA", "summary": "The paper proposes a method called OPTIMA for reducing the size of pre-trained multimodal models. It uses layerwise distillation to train smaller student models to match the hidden representations of larger teacher models. The authors track the contributions of individual modules and choose modules with greater contributions to distill more frequently. OPTIMA leverages a modified-Thompson sampling algorithm to estimate changing module contributions and selects modules to maximize cumulative contribution. Experimental results on multimodal understanding and image captioning tasks demonstrate the effectiveness of OPTIMA.", "questions": "Minor: citations in the Appendix are not compiled correctly.", "limitations": "Author provided discussion on the limitation.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper works on an important problem: distilling large-scale multi-modal foundation models. The authors explored a relatively new perspective toward this goal, which is Module-wise Adaptive Distillation.  A smart and automatic algorithm, OPTIMA, is proposed to assist in the process, which enabled a balance between the distillation of different model components, defined by which modality (or both) it operates on.  This is overall a new direction for a meaningful topic.\n\nIn the presented results, the proposed method outperforms the vanilla layer-wise distillation.  \n\nThe paper is well written in general. Experimental details are provided, as well as ablation studies on the design choices in OPTIMA.", "weaknesses": "The authors did not provide enough evidence/motivation for why OPTIMA is only applied to CoCa. The authors are expected to make the limitation of the proposed method clear. Especially, from the results shown in Table 1, we see that when comparing with other baseline distillation methods, either the DistilVLM12×384 achieves the highest score for the two tasks reported, with the least number of parameters. This leaves doubts about the significance of the work. \n\nAlso, the authors chose to report the Median among the 5 repetitions as results in Table 1, which is not a common choice. I'd hope to see a comparison under the mean + std.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for the insightful comments. We hope our response addresses your concerns. Should you have additional questions, we are readily available for further discussion.\n\n**The authors did not provide enough evidence/motivation for why OPTIMA is only applied to CoCa. The authors are expected to make the limitation of the proposed method clear.**\n\nThe commonly adopted multimodal foundation models are often computationally expensive to train and test. Given our limited computing resources, we were constrained to allocate them solely to a specific family of models, namely CoCa. Despite this limitation, we present a comprehensive analysis and results focused on CoCa.\n\nRegrettably, we have since lost access to the requisite computing resources, and are unable to provide further results for other model families given the short response/discussion window. While we agree with the reviewers that exploring the application of the module-adaptive idea to other models is of great interest to the community, we need to leave such investigations to future research. This limitation will be further highlighted in our paper.\n\nNevertheless, CoCa has demonstrated powerful and consistent performance over a diverse set of downstream tasks, thereby implying the broad applicability and efficacy of our method. Furthermore, the improvements on CoCa also indicate the promise of our idea on other multi-module models. We have shown that our idea is effective when a specific module is more useful to a given task than others, which is a shared observation in other multi-module models [1].\n\n[1] Shen, Sheng, et al. \"How much can clip benefit vision-and-language tasks?.\"\n\n**In Table 1, we see that when comparing with other baseline distillation methods, either the DistilVLM12×384 achieves the highest score for the two tasks reported, with the least number of parameters. This leaves doubts about the significance of the work.**\n\nThe scores reported by DistilVLM are better than OPTIMA because DistilVLM adopts a different experimental setting from OPTIMA: \n\nAs mentioned in line 237-247, 1) DistilVLM conducts distillation in the pre-training stage while OPTIMA conducts distillation in the fine-tuning stage. Pre-training distillation significantly improves the student generalizability but requires much more computations. 2) DistilVLM initializes both the student and the teacher with smaller, but similarly strong models compared with the student and the teacher used in OPTIMA. A smaller teacher-student gap can significantly improve the effectiveness of distillation.\n\nAside from adopting a different experimental setting, DistilVLM essentially applies layerwise distillation algorithm. As a result, a fair comparison would be between OPTIMA and CoCa-Tiny in Table 1, which is the layerwise distillation baseline under our experimental setting. We can observe that OPTIMA achieves a consistent improvement upon this baseline.\n\nWe remark that our main significance is to demonstrate that module-wise adaptive distillation, which distills multi-modality models based on the contributions of modalities, is a promising and interesting research direction. We use a MAB algorithm to demonstrate this point, because MAB is simple, efficient, and good at balancing exploration and exploitation. While we believe that greater improvements can be achieved with more advanced RL algorithms, we will leave the algorithmic explorations for future works.\n\n**Also, the authors chose to report the Median among the 5 repetitions as results in Table 1, which is not a common choice. I'd hope to see a comparison under the mean + std.**\n\nThe mean statistics are reported below and are consistent with median results. The std statistics are reported in Table 4 of the paper. \n\n|                  | VQA | SNLI-VE | NLVR2 | CoCo |\n|------------------|-----|---------|-------|------|\n| CoCa-Tiny_6      |  68.7   |    82.0     |    76.5   |   112.1   |\n| CoCa-Tiny_6 (OPTIMA) |   69.0  |  82.3   |   77.0    |  113.3    |\n| CoCa-Tiny_12     |  71.2   |    83.9     |   80.4    |   117.1   |\n| CoCa-Tiny_12 (OPTIMA) |  71.6   |  84.2   |    81.3   |   121.3   |"}]

Author rebuttal:

Meta-review: Thank you for your submission and the subsequent clarifications provided during the review process.

The paper addresses a significant challenge in the realm of distilling large-scale multi-modal foundation models. The introduction of Module-wise Adaptive Distillation offers a fresh perspective on this topic. The OPTIMA algorithm, as proposed, is both innovative and automatic, effectively striking a balance in the distillation of various model components based on their operational modality. This approach not only presents a novel direction but also underscores the importance of the topic at hand. Furthermore, the superior performance of the method compared to the baselines, coupled with the exhaustive ablation study, provides readers with a deeper understanding of the method's effectiveness.

However, it's worth noting a concern raised by one of the reviewers regarding the fairness of comparisons and the method's applicability to other multi-modal models. While the reviewer acknowledges the challenges you faced and the efforts made to maintain a fair comparison, it would be beneficial to address these concerns in future work or iterations of this research.

Given the overall positive feedback from all reviewers and the constructive discussions that ensued, I am pleased to inform you that your paper has been accepted. We believe that your work will make a valuable contribution to the field and look forward to its presentation.

Area Chair, NeurIPS.
Decision: Accept (poster)