SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\beta = (0.9, 0.999)$  and a learning rate of  $1 \times 10^{-3}$  with a linear decay schedule. We match the  $i$ -th layer of a student's module with the  $\lceil gi \rceil$ -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_{\mathrm{hid}} \in \mathbb{R}^{d_s \times d_t}$ . We set  $\alpha_1 = 0$ ,  $\alpha_2 = 1$  and  $\alpha_3 = 1 \times 10^{-2}$  for all tasks. For OPTIMA, we set  $\gamma = 0.98$ ,  $T_0 = 10$ ,  $P = 100$  and  $T = \frac{T'}{P} = 1k$ . Full details are deferred to Appendix A.4.
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-ViLp [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-ViLp 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 task-specific 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<sup>5</sup>. On both CoCa-Tiny<sub>6</sub> and CoCa-Tiny<sub>12</sub>, 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<sub>12</sub> 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<sub>6</sub> than CoCa-Tiny<sub>12</sub>. This might be because CoCa-Tiny<sub>6</sub> converges slower, and that specific module is less likely to over-fit quickly and interfere with others.
5The standard deviations are reported in Appendix A.5.
Compared with pre-trained VLMs at the same scale, CoCa-Tiny<sub>12</sub> (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.
Table 1: The evaluation results of OPTIMA and baseline methods on VQA, SNLI-VE, NLVR2 and COCO Caption. The results of all baselines are 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-standard	SNLI-VE test	NLVR2 test-p	COCO CIDEr
UNITER-Base [7]	-	146.5	-	72.9	78.3	77.9	-
OSCAR-Base [24]	-	146.5	-	73.4	-	-	137.6
ViLT-Base [20]	-	85.6	-	-	76.4	76.1	-
VILLA-Base [12]	-	146.5	-	73.7	79.0	79.3	-
CLIP-ViLp[35]	-	187.7	-	74.1	79.0	-	127.9
ALBEF [23]	-	241.0	-	74.7	80.3	80.5	-
CoCa-Base [48]	-	293.1	-	69.2	83.6	80.2	126.7
CoCa-Large [48]	-	672.1	-	75.3	85.6	82.6	132.3
MiniVLM12×384 [43]	OSCAR-Base	30.0	3.1×	68.1	-	-	115.0
DistilVLM12×384 [11]	OSCAR-Base	30.0	3.1×	69.2	-	-	117.1
DIDE12 [45]	ViLT-Base	171.3	3.0×	69.2	76.3	75.6	-
CoCa-Tiny6	CoCa-Large	54.5	6.2×	68.7	82.0	76.5	112.2
CoCa-Tiny6 (OPTIMA)	CoCa-Large	54.5	6.2×	69.0	82.3	77.0	113.5
CoCa-Tiny12	CoCa-Large	101.8	3.4×	71.2	83.9	80.4	116.8
CoCa-Tiny12 (OPTIMA)	CoCa-Large	101.8	3.4×	71.6	84.3	81.3	121.2
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.,  $\mathcal{L}_{\mathrm{train}}$  and  $\mathcal{D}_{\mathrm{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.
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.
Figure 4: The prediction loss and the output layer prediction discrepancy (i.e.,  $\mathcal{L}_{\mathrm{train}}$  and  $\mathcal{D}_{\mathrm{KL}}$  defined Eq. 1, respectively) in OPTIMA and layerwise distillation (denoted by "LWD").
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.
Figure 5: Left: The means of the reward distributions of all arms (i.e.,  $\{\mu_k\}_{k=1}^K$ ) 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.
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.
Figure 6: A visualization of the frequency of each arm being chosen through training. A deeper color represents a higher frequency.
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.
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.
This project was completed during a full-time internship at Google Research. We thank all collaborators 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.
[49] 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
[50] 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
[51] 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.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 83k images and 444k questions for training, 41k images, and 214k 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. It contains 672M parameters in Transformer layers and contains a 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 a 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 hyper-parameters for all tasks, as listed in Table 2.
Table 2: Hyper-parameters for fine-tuning CoCa-Large teacher models.
Hyper-parameters	VQA	SNLI-VE	NLVR2	COCO Caption
Optimizer	Adafactor with Decoupled Weight Decay
Adam βs	(0.9, 0.999)
Gradient Clipping	1.0
Learning Rate Schedule	Linear Schedule Decaying to Zero
Warm-up Steps	1k
Weight Decay Rate	0.1
Pooling Learning Rate	5 × 10-4	1 × 10-3	5 × 10-3	N/A
Encoder Learning Rate	2 × 10-5	5 × 10-5	2 × 10-5	1 × 10-5
RandAugment	1,10	1,10	None	None
Training Steps	100k	50k	50k	50k
Batch Size	64	128	64	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 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 CoCa-Large teacher on that task. CoCa-Tiny $_{12}$  contains 102M parameters in the Transformer layers and contains a total of 152M parameters including the embedding size. CoCa-Tiny $_{6}$  contains 55M parameters in the Transformer layers and contains a 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, we use a sentence-piece model [33, 22] with a vocabulary size of 64k 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 \cdot 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	NLVR2	COCO Caption
T0	10
P	100
T	1000
γ	0.98
αs	(0.0, 1.0, 1 × 10-2)
Optimizer	Adafactor with Decoupled Weight Decay
Adam βs	(0.9, 0.999)
Gradient Clipping	1.0
Learning Rate Schedule	Linear Schedule Decaying to Zero
Learning Rate	1 × 10-3
Warm-up Steps	1k
Weight Decay Rate	0.1
RandAugment	1, 10	1, 10	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-Tiny $_{12}$  and CoCa-Tiny $_{6}$ . 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	VQAAcc	SNLI-VEAcc	NLVR2Acc	COCO CaptionCIDEr
CoCa-Tiny6	0.20	0.25	0.11	0.15
CoCa-Tiny6 (OPTIMA)	0.22	0.13	0.35	0.15
CoCa-Tiny12	0.17	0.23	0.15	0.88
CoCa-Tiny12 (OPTIMA)	0.08	0.15	0.30	0.32
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:  $\mathcal{D}_{\mathrm{KL}}$  (Eq. 1),  $\mathcal{L}_{\mathrm{hidn}}$  (Eq. 3) and  $\mathcal{L}_{\mathrm{attn}}$  (Eq. 4). Figure 7 compares ours with two variants: 1)  $r_{\mathrm{KD}}$ : the ratio of loss decrement of  $\mathcal{D}_{\mathrm{KL}}$ ; 2)  $r_{\mathrm{LWD}}$ : the averaged ratio of loss decrements over  $\mathcal{L}_{\mathrm{hidn}}$  and  $\mathcal{L}_{\mathrm{attn}}$ . We can observe that  $r_{\mathrm{KD}}$  performs better than  $r_{\mathrm{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_{\mathrm{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  $-\frac{1}{\log_2\gamma}$ . 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  $\gamma$  can correctly track the changing contribution.
Figure 7: A comparison of three variants of reward: 1)  $r_{\mathrm{KD}}$ : the ratio of decrement of  $\mathcal{D}_{\mathrm{KL}}$ ; 2)  $r_{\mathrm{LWD}}$ : the averaged ratio of decrement over  $\mathcal{L}_{\mathrm{hidn}}$  and  $\mathcal{L}_{\mathrm{attn}}$ ; 3)  $r_{\mathrm{OPTIMA}}$ .
Figure 8: A hyper-parameter study on the halflife of the EMA, computed as  $-\frac{1}{\log_2\gamma}$ .