Title: Module-wise Adaptive Distillation for Multimodality Foundation Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Multimodality foundation models

knowledge distillation
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

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 [48] as the teacher model.

1 Introduction

Large pre-trained multimodality foundation models have demonstrated remarkable generalizability on a 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 or 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 to 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 a multimodal decoder on top of the dual-encoder architecture. The decoder learns a 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

Tianyi Zhou

accommodate new modalities. This poses a significant challenge for deployment in applications with latency and storage constraints.

Layerwise distillation is a powerful approach to compressing large models (i.e., teacher models) into small ones (i.e., student models) with minimal loss of performance [16, 30]. This approach trains a student to match the hidden representations of the teacher at each layer. Given that the teacher'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 is distilled from a teacher that has been fine-tuned on the target task [38, 17].

In task-specific distillation in multimodality models, however, matching the student's representations with those of the teacher at every layer does not always benefit the student's performance. Prior research has shown that improving the representations of a 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, we 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).

Img

Txt

Multi

Modules

50

60

70

Test Accuracy (

%

)

61.5

51.7

66.9

NLVR2

Img

Txt

Multi

Modules

60

80

100

120

CIDEr

105.1

75.6

59.5

COCO Caption

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 on their contributions. In this way, we can more frequently distill the module with a greater contribution. To evaluate the contribution of a module, we distill it for a 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, we 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 a limited budget for training steps, we face a significant challenge in balancing exploration and exploitation.

To address this challenge, we adopt a multi-armed bandit (MAB) approach [36, 41]. 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 as 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 are observed, we obtain a more accurate estimate of the underlying reward distribution for each module, which reflects the actual contribution of this module. Based on the reward distribution, we choose modules to maximize the cumulative loss decrement.

However, the actual contribution of a module may change with model updating. For example, distilling the multimodal decoder may become more beneficial when its input, i.e., the image and text representations, are sufficiently optimized [39]. To estimate such a non-stationary underlying reward

Multi-armed Bandit

Algorithm

Round t: *Distill

Round t+1: *Distill

Round t+2: “Distill

Image Encoder'

Multimodal Decoder"

Text Encoder'

Ltrain

Ltrain

Ltrain

Lhidn

Lattn I

+

Multimodal

Lhidn + Lattn

Decoder

Lhidn + Lattn

Image

Text

Encoder

Encoder

Student

Teacher

Image

Text

distribution, we need to count more on rewards observed in recent history . This makes it challenging for us to adopt classical MAB algorithms because most of them are designed for stationary underlying reward distributions [4, 10, 40]. Therefore, they average the observed rewards over the whole history to estimate the underlying reward distribution. In our case, since the model is unlikely to change drastically within a few steps, and so is the underlying reward distribution, we tailor these stationary algorithms by replacing the simple average with the exponential moving average of 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), a novel taskspecific distillation method for multimodality foundation models. We exemplify OPTIMA using the Thompson Sampling algorithm for its strong empirical and theoretical performance [32, 1, 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 image captioning benchmarks [14, 47, 37, 6]. Specifically, we distill a CoCa-Tiny 12 ( 102 M) and a CoCa-Tiny 6 ( 55 M) from a CoCa-Large teacher ( 672 M, [48]), a powerful pre-trained multimodality foundation model. For both students, OPTIMA substantially surpasses the layerwise distillation baseline. Moreover, CoCa-Tiny 12 outperforms CoCa-Base in 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.

2 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]). Dualencoder 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 an image and relevant text. This enables each encoder to learn both the unimodal representations and the crossmodal alignment knowledge. Encoder-decoder models contain a multimodal decoder on 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-attention mechanism.

Knowledge Distillation trains a small model (i.e., student model) to match the output prediction of a large and well-trained model (i.e., teacher model) by minimizing their prediction discrepancy [16]. Specifically, we denote the teacher model as f t ( θ t ) and the student model as f s ( θ s ) , and consider the

following optimization problem:

where L train ( θ s ) is the prediction loss on the target task, e.g., the cross-entropy loss for a classification task; D KL ( θ s , θ t ) is the KL-Divergence between the probability distributions over their predictions, i.e., KL ( f s ( θ s ) || f t ( θ t )) . 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, 1]. TS assumes a prior distribution on 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 of an arm is specified as

where µ 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 c > 1 modules. For example, they are CoCa models containing c = 3 modules in this work: an image encoder, a text encoder, and a multimodal decoder.

We construct K = 2 c -1 arms, each associated with a unique, non-empty subset of modules selected from the c modules. The k -th arm is formed as a set of model layers, denoted as S k , in its associated subset. We set T = T ′ P rounds, where T ′ is the total number of training steps and P ≥ 1 is the number of steps within each round. In this way, choosing the k -th arm in the t -th round is equivalent to distilling S k for P steps in the range of (( t -1) P, t · P ] .

For each arm k ∈ [ K ] , we denote its reward distribution as D k and set it as a Gaussian distribution specified as N ( µ k , 1 n k +1 ) following Eq. 2. To set a proper prior, we first randomly choose each arm for T 0 ≥ 0 rounds 1 , then initialize µ k = µ 0 k and n k = T 0 , where µ 0 k is the average of the observed rewards over T 0 rounds. We explain why setting the reward distribution as a Gaussian distribution is feasible in Appendix A.1.

Remark 3.1 We consider all possible combinations of modules because, at a 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, sampling 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.

for all arm k ∈ [ K ] .

At the t -th round, we sample a reward ̂ r k for the k -th arm from its reward distribution D k for all k ∈ [ K ] . The arm with the highest sampled reward is then chosen for this round, denoted as a t . In this way, the arm is chosen according to the belief of it being the best arm.

We then play the a t -th arm by distilling S a t for P steps. At any step t ′ ∈ (( t -1) P, tP ] , we first compute the distillation losses on S a t . Specifically, we denote the hidden representations at the i -th layer of the teacher and the student as H i t ∈ R | x |× d t and H i s ∈ R | x |× d s , where d t and d s denote the hidden dimension and | x | denotes the sequence length. Then the distillation loss on hidden representations is defined as 2 :

where MSE ( · , · ) is the mean-squared error, and W i hidn ∈ R d s × d t is a randomly initialized and learnable linear projection that projects H i s into the same space as H i t . Similarly, the distillation loss on attention scores is defined as 3 :

where A i t , A i s ∈ R | x |×| x | are the attention score matrices averaged over the multiple attention heads at the i -th layer. The student model is then optimized based on a weighted sum of all distillation losses on S a t , i.e.,

where L train and D KL are defined in Eq. 1 and α 1 , α 2 , α 3 ≥ 0 are hyper-parameters. Specifically, the student model is updated with an SGD-type algorithm [21]:

where η > 0 is the learning rate. The same procedure is then repeated for P steps.

We then compute the reward r a t for playing the a t -th arm. Specifically, we design the reward as the averaged ratios of decrements over all types of distillation losses:

where U = {D KL ( · ) , L hidn ( · ) , L attn ( · ) } . Note that L hidn ( · ) and L attn ( · ) are defined in Eq. 3 and Eq. 4, respectively, but by replacing S a t with the set of all layers. They measure the change in the distillation loss on 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 L ( θ (( t -1) P ) s ) and L ( θ ( tP ) s ) . 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 a t to the distillation performance of the full model.

We consider both the changes in output prediction discrepancy ( D KL ( · ) ) and the layerwise representation distances ( L hidn ( · ) and L attn ( · ) ). This improves the robustness of the reward over different tasks because a specific distance metric may not 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 r a t , we update the reward distribution D a t . Specifically, we update the distribution mean, µ a t , as the exponential moving average of the observed rewards in the past:

where γ ∈ (0 , 1) is a hyper-parameter. We then increase n a t by one and update D a t following Eq. 2. Since the exponential moving average discounts the old rewards, µ a t tends to reflect the average of the rewards within recent rounds. Recall that the actual contribution of each module changes dynamically. By using µ a t as the distribution mean, D a t can capture the changing contribution and provides up-to-date guidance for module selection. The complete algorithm is shown in Alg. 1 4 .

Algorithm 1 OPTIMA: Module Adaptive Distillation

4 Experiments

We verify the effectiveness of OPTIMA on 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 is to answer an input question with a sentence based on an input image. It is formulated as a classification problem with 3129 classes where each class corresponds to an answer. SNLI-VE is a three-way classification problem to predict whether an input sentence entails the input image, contradicts it, or is neutral to it. NLVR2 is a binary classification problem to predict whether an 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-based multimodality foundation model pre-trained on web-scale alt-texts [18] and image-text pairs [50]. Since CoCa consists of an image encoder, a text encoder, and a multimodal decoder, it provides a sufficiently large action space for OPTIMA. Furthermore, the decoder produces multimodal representations, which allows it to be directly fine-tuned without any extra step of multimodal pre-training, commonly needed for dual-encoder models [35, 49, 9].

For each target task, we fine-tune a pre-trained CoCa-Large as the teacher ( 672 M, 48 layers ( 24 / 12 / 12 layers in image/text/multimodal modules), d t = 1024 ). We then distill a task-specific student from the fine-tuned teacher. Specifically, we consider two architectures: CoCa-Tiny 12 ( 101 . 8 M, 12 layers ( 6 / 3 / 3 layers), d s = 768 ) and CoCa-Tiny 6 ( 54 . 5 M, 6 layers ( 3 / 1 / 2 layers), d s = 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 ( 293 M, 36 layers, d = 768 ).

4.3 Implementation Details

Teacher Initialization. We fine-tune a 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-entropy loss. For the image captioning task, we directly fine-tune CoCa with the captioning loss [48], i.e., L cap = ∑ | x | ℓ =1 log P θ t ( y ℓ | y <ℓ , x ) . We follow [48] for all fine-tuning hyper-parameters (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 ′ = 100 k steps. We use Adafactor with decoupled weight decay [34] as the optimizer with β = (0 . 9 , 0 . 999) and a learning rate of 1 × 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 g is the ratio of the number of layers of the two modules. We randomly initialize W hidn ∈ R d s × d t . We set α 1 = 0 , α 2 = 1 and α 3 = 1 × 10 -2 for all tasks. For OPTIMA, we set γ = 0 . 98 , T 0 = 10 , P = 100 and T = T ′ P = 1 k. 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], CLIP-ViL p [35] and ALBEF [23]. Different from foundation models, VLMs often contain a single backbone, e.g., a 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 p uses a 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 in the pre-training stage. Pre-training distillation significantly improves the student's generalizability but requires much more computational resources than taskspecific distillation. 2) All methods consider VLMs as the teachers. VLMs are much smaller than CoCa models, and a small teacher-student capacity gap often improves the effectiveness of distillation [26, 25]. However, VLMs are less generalizable and can only distill students on limited tasks. MiniVLM adopts knowledge distillation on additional unlabeled data (Eq. 1). DistilVLM adopts layerwise distillation. DIDE distills a dual-encoders student from a VLM teacher by aligning cross-modal attentions. We ignore concurrent methods with a teacher that significantly outperforms the CoCa-Large teacher.

4.5 Main Result

Table 1 summarizes the evaluation results of layerwise distillation and OPTIMA. We report the median over five random seeds 5 . On both CoCa-Tiny 6 and CoCa-Tiny 12 , 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-Tiny 12 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 6 than CoCa-Tiny 12 . This might be because CoCa-Tiny 6 converges slower, and that specific module is less likely to over-fit quickly and interfere with others.

Compared with pre-trained VLMs at the same scale, CoCa-Tiny 12 (OPTIMA) can achieve similar 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, DistilVLM and DIDE, OPTIMA can achieve significantly better performance on CoCa-Tiny 12 and similar performance on CoCa-Tiny 6 over all tasks. Recall that these baseline methods use web-scale pre-training data for distillation and use single-backbone VLMs as teachers. In contrast, OPTIMA uses limited target task data and faces a much larger capacity gap, e.g., 5 times larger for CoCa-Tiny 12 . Nevertheless, the performance gaps between CoCa-Tiny 12 (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.

5 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 nonstationary reward distribution in Appendix A.7.

5.1 OPTIMA Improves the Student's Performance

Figure 3 compares the performance of the students obtained by 1) fixed-arm distillation : distilling a 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+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 a better weighting for arm selection for each task.

Figure 4 shows the prediction loss and the output layer prediction discrepancy (i.e., L train and D KL 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.

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Random

OPTIMA

Arm Selection Methods

75

80

85

Test Accuracy (

%

)

78.6

73.8

78.3

82.2

84.0

78.3

83.9

82.8

84.3

SNLI-VE

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Random

OPTIMA

Arm Selection Methods

50

60

70

80

90

Test Accuracy (

%

)

61.5

51.7

66.9

70.9

81.1

67.6

80.4

76.9

81.3

NLVR2

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Random

OPTIMA

Arm Selection Methods

60

80

100

120

CIDEr

105.1

75.6

59.5

106.7

116.7

62.9

116.8

116.0

121.2

COCO Caption

50

100

Training Steps (k)

0.00

0.25

0.50

0.75

1.00

L

train

NLVR2

50

100

Training Steps (k)

2

4

6

D

KL

NLVR2

25

50

75

100

Training Steps (k)

1.0

1.1

L

train

COCO Caption

25

50

75

100

Training Steps (k)

1100

1200

1300

D

KL

COCO Caption

Methods

LWD

OPTIMA

5.2 Reward Distributions Reflect the Contributions of Arms

Figure 5 (Left) shows the means of the reward distributions of all arms through training. In all tasks, 'Img+Txt+Multi' quickly dominates the others, while 'Multi' and 'Txt+Multi' 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 can correctly reflect the changing contributions of modules. We can observe a strong and positive correlation between the mean of the reward distribution of an arm and the accuracy obtained by distilling this fixed arm through training.

0

50

100

Training Steps (k)

0.0

0.1

0.2

0.3

SNLI-VE

0

50

100

Training Steps (k)

0.0

0.1

0.2

0.3

NLVR2

0

50

100

Training Steps (k)

0.000

0.025

0.050

0.075

COCO Caption

25

50

75

100

Training Steps (k)

0.0

0.5

1.0

Corr.

SNLI-VE

NLVR2

COCO Caption

Arms

Img+Txt+Multi

Img+Multi

Img+Txt

Img

Txt

Multi

Txt+Multi

µk

Corr. between

µ

and Test Acc.

5.3 OPTIMA Choose Arms Based on Reward Distributions

Figure 6 visualizes the frequency of choosing each arm through training. In all tasks, we can observe that 'Img+Txt+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.

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Arms

20

40

60

80

100

Training Steps (k)

VQA

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Arms

SNLI-VE

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Arms

NLVR2

Img

Txt

Multi

Img+Txt

Img+Multi

Txt+Multi

Img+Txt+Multi

Arms

COCO Caption

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

6 Discussion

Application to Other Multi-module Multimodality Models. In this study, we primarily demonstrate the efficacy of OPTIMA when applied to models within the CoCa family. This initial success signifies that OPTIMA could be applicable in scenarios where one or more modules predominantly contribute to 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 a subset of modules. This efficiency enables the exploration of OPTIMA's applicability to 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. In this study, we divide the CoCa model coarsely into three modules based on modalities, and design the action space to encompass all possible module combinations. However, different layers within a single module could exhibit variable contributions to the target task. For example, lower layers tend to capture high-frequency signals crucial for feature extraction, while upper layers typically learn low-frequency signals, vital for semantic understanding. By further subdividing each module into groups of layers, we can refine the action space to encompass all possible combinations of layer groups, potentially enabling the finding of better arms. Additionally, this enhanced granularity enables the extension of our method to single-module VLMs and single-modality models, such as large language models (LLMs) and vision models.

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 a fine-grained 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 we designate such explorations for future research.

References

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

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 pattern recognition , pages 6077-6086, 2018. 7

Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning , 47(2):235-256, 2002. 3

Sébastien Bubeck and Che-Yu Liu. Prior-free and prior-dependent regret bounds for thompson sampling. Advances in neural information processing systems , 26, 2013. 3, 4

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, 15

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

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 arXiv:2010.11929 , 2020. 7

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-end vision-and-language transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 18166-18176, 2022. 6

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

Zhiyuan Fang, Jianfeng Wang, Xiaowei Hu, Lijuan Wang, Yezhou Yang, and Zicheng Liu. Compressing 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

Zhe Gan, Yen-Chun Chen, Linjie Li, Chen Zhu, Yu Cheng, and Jingjing Liu. Large-scale adversarial training for vision-and-language representation learning. Advances in Neural Information Processing Systems , 33:6616-6628, 2020. 7, 8

Aurélien Garivier and Eric Moulines. On upper-confidence bound policies for non-stationary bandit problems. arXiv preprint arXiv:0805.3415 , 2008. 3

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 IEEE conference on computer vision and pattern recognition , pages 6904-6913, 2017. 3, 6, 15

Neha Gupta, Ole-Christoffer Granmo, and Ashok Agrawala. Thompson sampling for dynamic multi-armed bandits. In 2011 10th International Conference on Machine Learning and Applications and Workshops , volume 1, pages 484-489. IEEE, 2011. 3

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531 , 2015. 2, 3

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

Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, YunHsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In International Conference on Machine Learning , pages 4904-4916. PMLR, 2021. 1, 3, 6

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:1909.10351 , 2019. 2, 4

Wonjae Kim, Bokyung Son, and Ildoo Kim. Vilt: Vision-and-language transformer without convolution or region supervision. In International Conference on Machine Learning , pages 5583-5594. PMLR, 2021. 7, 8

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014. 5

Taku Kudo. Subword regularization: Improving neural network translation models with multiple subword candidates. arXiv preprint arXiv:1804.10959 , 2018. 15

Junnan Li, Ramprasaath Selvaraju, Akhilesh Gotmare, Shafiq 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 , 34:9694-9705, 2021. 3, 7, 8

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 vision-language tasks. In European Conference on Computer Vision , pages 121-137. Springer, 2020. 7, 8

Chen Liang, Simiao Zuo, Qingru Zhang, Pengcheng He, Weizhu Chen, and Tuo Zhao. Less is more: Task-aware layer-wise distillation for language model compression. arXiv preprint arXiv:2210.01351 , 2022. 7

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

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. 1, 3, 7

Vishnu Raj and Sheetal Kalyani. Taming non-stationary bandits: A bayesian approach. arXiv preprint arXiv:1707.09727 , 2017. 3

Steven J Rennie, Etienne Marcheret, Youssef Mroueh, Jerret Ross, and Vaibhava Goel. Selfcritical sequence training for image captioning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 7008-7024, 2017. 15

Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:1412.6550 , 2014. 2, 4

Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research , 39(4):1221-1243, 2014. 3, 4

Steven L Scott. A modern bayesian look at the multi-armed bandit. Applied Stochastic Models in Business and Industry , 26(6):639-658, 2010. 3, 4

Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. arXiv preprint arXiv:1508.07909 , 2015. 15

Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning , pages 4596-4604. PMLR, 2018. 7

Sheng Shen, Liunian Harold Li, Hao Tan, Mohit Bansal, Anna Rohrbach, Kai-Wei Chang, Zhewei Yao, and Kurt Keutzer. How much can clip benefit vision-and-language tasks? arXiv preprint arXiv:2107.06383 , 2021. 6, 7, 8, 10

Aleksandrs Slivkins et al. Introduction to multi-armed bandits. Foundations and Trends® in Machine Learning , 12(1-2):1-286, 2019. 2

Alane Suhr, Stephanie Zhou, Ally Zhang, Iris Zhang, Huajun Bai, and Yoav Artzi. A corpus for reasoning about natural language grounded in photographs. arXiv preprint arXiv:1811.00491 , 2018. 2, 3, 6

Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprint arXiv:1908.09355 , 2019. 2

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.02984 , 2020. 2

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-294, 1933. 3, 4

Joannes Vermorel and Mehryar Mohri. Multi-armed bandit algorithms and empirical evaluation. In European conference on machine learning , pages 437-448. Springer, 2005. 2

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 3156-3164, 2015. 3

Jianfeng Wang, Xiaowei Hu, Pengchuan Zhang, Xiujun Li, Lijuan Wang, Lei Zhang, Jianfeng Gao, and Zicheng Liu. Minivlm: A smaller and faster vision-language model. arXiv preprint arXiv:2012.06946 , 2020. 7, 8

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 a simple sequence-to-sequence learning framework. arXiv preprint arXiv:2202.03052 , 2022. 3

Zekun Wang, Wenhui Wang, Haichao Zhu, Ming Liu, Bing Qin, and Furu Wei. Distilled dual-encoder model for vision-language understanding. arXiv preprint arXiv:2112.08723 , 2021. 2, 4, 7, 8

Zirui Wang, Jiahui Yu, Adams Wei Yu, Zihang Dai, Yulia Tsvetkov, and Yuan Cao. Simvlm: Simple visual language model pretraining with weak supervision. arXiv preprint arXiv:2108.10904 , 2021. 1, 3

Ning Xie, Farley Lai, Derek Doran, and Asim Kadav. Visual entailment: A novel task for fine-grained image understanding. arXiv preprint arXiv:1901.06706 , 2019. 3, 6

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. 1, 2, 3, 6, 7, 8, 15

Lu Yuan, Dongdong Chen, Yi-Ling Chen, Noel Codella, Xiyang Dai, Jianfeng Gao, Houdong Hu, Xuedong Huang, Boxin Li, Chunyuan Li, et al. Florence: A new foundation model for computer vision. arXiv preprint arXiv:2111.11432 , 2021. 1, 3, 6

Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 12104-12113, 2022. 6

Pengchuan Zhang, Xiujun Li, Xiaowei Hu, Jianwei Yang, Lei Zhang, Lijuan Wang, Yejin Choi, and Jianfeng Gao. Vinvl: Revisiting visual representations in vision-language models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 5579-5588, 2021. 2

A Appendix

A.1 Gaussian as the Posterior

Here we explain why setting the reward distribution as a Gaussian distribution is feasible. Recall that we consider a non-stationary environment. We design the reward distribution to model the rewards observed in recent history. Within a small number of steps, the model is 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 83 k images and 444 k questions for training, 41 k images, and 214 k questions for validation. For the image captioning task on COCO, we use [6] for training and testing. It contains 11 k images for training and 5 k images for validation and 5 k images for testing.

A.3 Teacher Model Implementation Details

We fine-tune a pre-trained CoCa-Large model [48] as the teacher. It contains 672 Mparameters in Transformer layers and contains a total of 787 Mparameters 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 a vocabulary size of 64 k. 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 hyper-parameters for all tasks, as listed in Table 2.

For multimodal understanding tasks, we follow [48] to apply an attentional pooler with a single query to 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-Tiny 12 student and a CoCa-Tiny 6 student from a fine-tuned CoCaLarge teacher on that task. CoCa-Tiny 12 contains 102 M parameters in the Transformer layers and contains a total of 152 M parameters including the embedding size. CoCa-Tiny 6 contains 55 Mparameters in the Transformer layers and contains a total of 105 Mparameters including the embedding size. We use 576 as the image resolution and 18 as the patch size for image inputs. To tokenize text input, we use a sentence-piece model [33, 22] with a vocabulary size of 64 k trained on the sampled pre-training dataset. We use 64 as the max sequence length for text inputs. We follow [48] for fine-tuning hyper-parameters for all tasks, as listed in Table 3.

We conduct a two-stage distillation for CoCa-Tiny 6 . We first distill CoCa-Tiny 12 from CoCa-Large, then use the distilled CoCa-Tiny 12 as the teacher to teach CoCa-Tiny 6 . Existing works have shown

that introducing an intermediate-sized teacher reduces the gap between the teacher and the student model, which allows the distillation to be more effective [26].

We distill the model for a total of T ′ steps, i.e., a total of T = T ′ /P rounds. Among the T rounds, the first T 0 · K rounds are used for initialization of the parameters of the reward distribution.

A.5 Statistics of Experimental Results

We report the median of five random seeds for experiment results on CoCa-Tiny 12 and CoCa-Tiny 6 . Table 4 show the standard deviations of the experimental results in Table 1.

A.6 Design of Reward

Recall that we design the reward (Eq. 6) as the averaged ratio of loss decrements over three types of distillation losses: D KL (Eq. 1), L hidn (Eq. 3) and L attn (Eq. 4). Figure 7 compares ours with two variants: 1) r KD : the ratio of loss decrement of D KL ; 2) r LWD : the averaged ratio of loss decrements over L hidn and L attn . We can observe that r KD performs better than r LWD in NLVR2 but reversely in COCO. Since captioning tasks often rely more on contextual knowledge in the layerwise representations than classification tasks, the layerwise representation distance may better characterize the distillation performance in COCO. By taking both distance metrics into consideration, r OPTIMA performs well on both tasks.

A.7 Design of the Reward Distribution

Recall that we design the mean of the reward distribution (Eq. 7) as the exponential moving average (EMA) of the past rewards. Figure 8 shows a hyper-parameter study on the halflife of the EMA, computed as -1 log 2 γ . 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 instantaneous rewards can both be harmful to the student's performance. This corroborates that the actual contribution of each module is non-stationary in the long term and stationary in the short term, and using EMA with an appropriate γ can correctly track the changing contribution.

r

KD

r

LWD

r

OPTIMA

Reward Types

80

81

82

Test Accuracy (

%

)

81.0

80.3

81.2

NLVR2

r

KD

r

LWD

r

OPTIMA

Reward Types

100

110

120

130

CIDEr

108.4

121.0

121.2

COCO Caption

0

20

40

60

Halflife of EMA (Rounds)

83.5

84.0

Test Accuracy (

%

)

SNLI-VE

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nMinor: citations in the Appendix are not compiled correctly.\n\nLIMITATIONS:\nAuthor provided discussion on the limitation.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\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\nQUESTIONS:\nThe questions are listed in the \"weaknesses\" section.\n\nLIMITATIONS:\nThe authors have discussed the limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\n\nSTRENGTHS:\n1.\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\nWEAKNESSES:\n1.\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\nQUESTIONS:\nPleae see the weakness part\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}

comments: 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.

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.

QUESTIONS:
No.

LIMITATIONS:
No.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-21

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\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.\n\nWEAKNESSES:\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.\n\nQUESTIONS:\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?\n\nLIMITATIONS:\nLimitations have been addressed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-08-03"}