SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chen Liang
cliang73@gatech.edu
Jiahui Yu
jiahuiyu@google.com
Ming-Hsuan Yang
UC Merced, Google Research
minghsuan@google.com
mtbr@google.com
Yin Cui
NVIDIA Research
richardaecn@gmail.com
Tuo Zhao
tourzhao@gatech.edu
Boqing Gong
bgong@google.com
Tianyi Zhou
tianyi@umd.edu
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.
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
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).
Figure 1: Performances of CoCa-Tiny $_{12}$  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 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
Figure 2: An illustration of OPTIMA.
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 task-specific 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}$  (102M) and a CoCa-Tiny $_{6}$  (55M) from a CoCa-Large teacher (672M, [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.
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 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}(\theta_{t})$  and the student model as  $f_{s}(\theta_{s})$ , and consider the
following optimization problem:
$$
\min  _ {\theta_ {s}} \mathcal {L} _ {\text {t r a i n}} \left(\theta_ {s}\right) + \mathcal {D} _ {\mathrm {K L}} \left(\theta_ {s}, \theta_ {t}\right), \tag {1}
$$
where  $\mathcal{L}_{\mathrm{train}}(\theta_s)$  is the prediction loss on the target task, e.g., the cross-entropy loss for a classification task;  $\mathcal{D}_{\mathrm{KL}}(\theta_s,\theta_t)$  is the KL-Divergence between the probability distributions over their predictions, i.e.,  $\mathrm{KL}(f_s(\theta_s)||f_t(\theta_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
$$
\mathcal {D} := \mathcal {N} \left(\mu , \frac {1}{n + 1}\right), \tag {2}
$$
where  $\mu$  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.
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 = \frac{T'}{P}$  rounds, where  $T'$  is the total number of training steps and  $P \geq 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 \cdot P]$ .
For each arm  $k \in [K]$ , we denote its reward distribution as  $\mathcal{D}_k$  and set it as a Gaussian distribution specified as  $\mathcal{N}(\mu_k, \frac{1}{n_k + 1})$  following Eq. 2. To set a proper prior, we first randomly choose each arm for  $T_0 \geq 0$  rounds $^1$ , then initialize  $\mu_k = \mu_k^0$  and  $n_k = T_0$ , where  $\mu_k^0$  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.
OPTIMA Algorithm. We initialize
$$
\mu_ {k} = \mu_ {k} ^ {0}, \quad n _ {k} = T _ {0} \quad \text {a n d} \quad \mathcal {D} _ {k} = \mathcal {N} \left(\mu_ {k}, \frac {1}{n _ {k} + 1}\right)
$$
for all arm  $k\in [K]$
At the  $t$ -th round, we sample a reward  $\widehat{r}_k$  for the  $k$ -th arm from its reward distribution  $\mathcal{D}_k$  for all  $k \in [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.
${}^{1}{T}_{0}$  is included in  $T$  .
We then play the  $a_{t}$ -th arm by distilling  $S_{a_t}$  for  $P$  steps. At any step  $t^\prime \in ((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_{t}^{i}\in \mathbb{R}^{|x|\times d_{t}}$  and  $H_{s}^{i}\in \mathbb{R}^{|x|\times 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:
$$
\mathcal {L} _ {\mathrm {h i d n}} ^ {a _ {t}} \left(\theta_ {s} ^ {(t ^ {\prime})}, \theta_ {t}\right) = \frac {1}{\left| S _ {a _ {t}} \right|} \sum_ {i \in S _ {a _ {t}}} \operatorname {M S E} \left(H _ {t} ^ {i}, H _ {s} ^ {i} W _ {\mathrm {h i d n}} ^ {i}\right), \tag {3}
$$
where  $\mathrm{MSE}(\cdot ,\cdot)$  is the mean-squared error, and  $W_{\mathrm{hidn}}^i\in \mathbb{R}^{d_s\times d_t}$  is a randomly initialized and learnable linear projection that projects  $H_{s}^{i}$  into the same space as  $H_{t}^{i}$ . Similarly, the distillation loss on attention scores is defined as<sup>3</sup>:
$$
\mathcal {L} _ {\text {a t t n}} ^ {a _ {t}} \left(\theta_ {s} ^ {(t ^ {\prime})}, \theta_ {t}\right) = \frac {1}{\left| S _ {a _ {t}} \right|} \sum_ {i \in S _ {a _ {t}}} \operatorname {M S E} \left(A _ {t} ^ {i}, A _ {s} ^ {i}\right), \tag {4}
$$
where  $A_{t}^{i}, A_{s}^{i} \in \mathbb{R}^{|x| \times |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.,
$$
\mathcal {L} _ {\text {t o t a l}} ^ {a t} = \mathcal {L} _ {\text {t r a i n}} + \alpha_ {1} \mathcal {D} _ {\mathrm {K L}} + \alpha_ {2} \mathcal {L} _ {\text {h i d n}} ^ {a t} + \alpha_ {3} \mathcal {L} _ {\text {a t t n}} ^ {a t}, \tag {5}
$$
where  $\mathcal{L}_{\mathrm{train}}$  and  $\mathcal{D}_{\mathrm{KL}}$  are defined in Eq. 1 and  $\alpha_{1},\alpha_{2},\alpha_{3}\geq 0$  are hyper-parameters. Specifically, the student model is updated with an SGD-type algorithm [21]:
$$
\theta_ {s} ^ {(t ^ {\prime} + 1)} \leftarrow \theta_ {s} ^ {(t ^ {\prime})} - \eta \nabla_ {\theta_ {s}} \mathcal {L} _ {\mathrm {t o t a l}} ^ {a _ {t}} (\theta_ {s} ^ {(t ^ {\prime})}, \theta_ {t}),
$$
where  $\eta > 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:
$$
r _ {a _ {t}} = \frac {1}{| U |} \sum_ {\mathcal {L} (\cdot) \in U} \max  \left(0, \frac {\mathcal {L} \left(\theta_ {s} ^ {((t - 1) P)}\right) - \mathcal {L} \left(\theta_ {s} ^ {(t P)}\right)}{\mathcal {L} \left(\theta_ {s} ^ {((t - 1) P)}\right)}\right). \tag {6}
$$
where  $U = \{\mathcal{D}_{\mathrm{KL}}(\cdot),\mathcal{L}_{\mathrm{hidn}}(\cdot),\mathcal{L}_{\mathrm{attn}}(\cdot)\}$ . Note that  $\mathcal{L}_{\mathrm{hidn}}(\cdot)$  and  $\mathcal{L}_{\mathrm{attn}}(\cdot)$  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  $\mathcal{L}(\theta_s^{((t - 1)P)})$  and  $\mathcal{L}(\theta_s^{(tP)})$ . 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  $(\mathcal{D}_{\mathrm{KL}}(\cdot))$  and the layerwise representation distances  $(\mathcal{L}_{\mathrm{hidn}}(\cdot)$  and  $\mathcal{L}_{\mathrm{attn}}(\cdot))$ . 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  $\mathcal{D}_{a_t}$ . Specifically, we update the distribution mean,  $\mu_{a_t}$ , as the exponential moving average of the observed rewards in the past:
$$
\mu_ {a _ {t}} \leftarrow \gamma \mu_ {a _ {t}} + (1 - \gamma) r _ {a _ {t}}, \tag {7}
$$
2We omit the  $\left( t^{\prime} \right)$  superscript on  $\{H_s^i, W_{\mathrm{hldn}}^i\}_{i \in S_{a_t}}$  to simplify the notations. We assume the teacher and student have the same number of layers. The case of having different layers will be elaborated in Section 4.3.
<sup>3</sup>We omit the “ $(t')$ ” superscript on  $\{A_s^i\}_{i \in S_{a_t}}$  to simplify the notations.
where  $\gamma \in (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,  $\mu_{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  $\mu_{a_t}$  as the distribution mean,  $\mathcal{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<sup>4</sup>.
Algorithm 1 OPTIMA: Module Adaptive Distillation
1: Input:  $T$ : the number of total rounds;  $P$ : the number of steps in each round;  $K = 2^c - 1$ : the number of arms;  $\gamma$ : a discounted factor.  $\theta_s^{(0)}$ ,  $\theta_t$ : the initialized student and the teacher models.  
2: Output:  $\theta_s^{(T')}$   
3:  $t' = 0$ ,  $n_k = \mu_k = 0 \forall k = 1, \dots, K$ .  
4: for  $t = 1, \dots, T$  do  
5: for  $k = 1, \dots, K$  do  
6: Sample a reward for arm  $k$ :  $\widehat{r}_k \sim \mathcal{N}\left(\mu_k, \frac{1}{n_k + 1}\right)$ .  
7: end for  
8: Select arm  $a_t \gets \arg \max_k \widehat{r}_k$ .  
9: for  $p = 1, \dots, P$  do  
11: Compute  $\mathcal{L}_{total}^{a_t}(\theta_s^{(t')}, \theta_t)$  following Eq. 5.  
12:  $\theta_s^{(t' + 1)} \gets \theta_s^{(t')} - \eta \nabla_{\theta_s} \mathcal{L}_{total}^{a_t}(\theta_s^{(t')}, \theta_t)$ .  
13:  $t' \gets t' + 1$ .  
14: end for  
15:  
16: Compute  $r_{a_t}$  following Eq. 6.  
17: Update  $\mu_{a_t}$  following Eq. 7.  
18:  $n_k \gets n_k + 1$ .  
19: end for
We verify the effectiveness of OPTIMA on popular multimodal understanding and image captioning benchmarks.
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.
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].
4We assume  $T_0 = 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/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.8M, 12 layers (6/3/3 layers),  $d_s = 768$ ) and CoCa-Tiny $_6$  (54.5M, 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 (293M, 36 layers,  $d = 768$ ).
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.,  $\mathcal{L}_{\mathrm{cap}} = \sum_{\ell=1}^{|x|} \log P_{\theta_t}(y_\ell | y_{<\ell}, x)$ . We follow [48] for all fine-tuning hyper-parameters (See Appendix A.3 for details).