C ap EEN : Image Captioning with Early Exits and Knowledge Distillation (2024)

Abstract: Deep neural networks (DNNs) have made significant progress in recognizing visual elements and generating descriptive text in image-captioning tasks. However, their improved performance comes from increased computational burden and inference latency. Early Exit (EE) strategies can be used to enhance their efficiency, but their adaptation presents challenges in image captioning as it requires varying levels of semantic information for accurate predictions. To overcome this, we introduce CapEEN to improve the performance of EE strategies using knowledge distillation. Inference in CapEEN is completed at intermediary layers if prediction confidence exceeds a predefined value learned from the training data. To account for real-world deployments, where target distributions could drift from that of training samples, we introduce a variant A-CapEEN to adapt the thresholds on the fly using Multi-armed bandits framework. Experiments on the MS COCO and Flickr30k datasets show that CapEEN gains speedup of 1.77\times while maintaining competitive performance compared to the final layer, and A-CapEEN additionally offers robustness against distortions. The source code is available at https://github.com/Div290/CapEEN.

## introduction image captioning, a multifaceted challenge at the intersection of computer vision and natural language processing, has reaped the benefits of deep neural networks (dnns), characterized by their increased scale and complexity. this task entails not only the identification of visual elements within an image but also the intricate interpretation of their relationships. notably, the encoderdecoder framework has made significant strides in sentence generation by anticipating the next word in a sequence @xcite ; chen encoder 2 a big cat is chasing a mouse on a wooden floor encoder 1 encoder n decoder n decoder 2 decoder 1 exit-1 exit-2 figure 1 : the encoder-decoder framework with attached exits. the figure states that low-level features could be extracted from early classifiers and inferred there, while high-level features are inferred at deeper classifiers. the color of the text in the caption is the same as the color of the classifier after that layer. @xcite @xcite @xcite @xcite @xcite @xcite . this predictive modeling considers both the image's content and the preceding partial sentence, resulting in substantial progress in the field @xcite . however, their large size restricts deployment in resource-constrained scenarios requiring fast inference. early exit (ee) strategies have emerged as a strategic solution to overcome this challenge. in ee strategies, classifiers are attached to the intermediary layers. each sample can exit from one of them without requiring to pass through all layers (see fig. 1 ). this brings down computational requirements and improves latency @xcite . however, this generic approach of attaching exits to the pre-trained backbone may not be suitable for image captioning @xcite -in the layered hierarchy of representation in transformer-based models, the initial layers focus on extracting lowlevel features, while deeper layers delve into the complexities of semantic fusion relations @xcite . consequently, even 'easy samples' require a certain level of high-level information present at deeper layers. moreover, the decisions of early exit are based on the confidence at the intermediary layers being above a predefined threshold. the threshold used to compare the confidence levels significantly impacts the amount of latency and accuracy. these thresholds are learned during training and serve as a crucial reference point during inference. post-deployment, it may be possible that the distribution of the target sample could drift away from that of training samples. this is often encountered in real-world scenarios, e.g., blurred images due to the camera being out of focus @xcite during inference. such drifts could affect the threshold choice and significantly lower the dnns performance (see figure 2 ) prompting the question: how to adjust the threshold of deployed pre-trained models when the latent distribution of target samples differs from the training samples due to variation in distortion levels? also, this adaptation has to be unsupervised, as the ground truth labels may not be available during inference. this motivates a method that 1) gives initial layers high-level information during training and 2) adjusts early exit thresholds during inference for efficiency and robustness against distortions. we introduce a new approach that extends the idea of knowledge distillation in early exits @xcite to image captioning tasks named image captioning with early exits and knowledge distillation (capeen) to improve the efficiency of ees in image captioning tasks. by distilling the knowledge residing in deeper layers (teacher), capeen empowers initial classifiers (students) to utilize the richness of deeper representations, which enhances both the performance and speed of the ee model. to circumvent the issue of threshold choice under distribution change due to distortion in incoming samples during inference, we propose a novel online learning algorithm a-capeen based on the multi-armed bandits (mab) framework @xcite to learn the optimal threshold as per latent distributions of input samples. a-capeen is activated during the inference phase and adapts to various levels of distortions in test samples with minimal computational require- our experiments on the ms coco @xcite and @xcite dataset demonstrate that capeen significantly increases speedup while maintaining competitive accuracy performance (1.77x speedup in runtime as compared to the final layer). a-capeen further enhances efficiency by dynamically adjusting exit thresholds based on the latent distribution of the test dataset to make it robust to noise present in the datasets during inference. in summary, our contributions are as follows: 1) we present a novel self-distillation framework named capeen tailored for early exiting in image captioning. 2) we introduce an online learning algorithm, a-capeen, designed to dynamically choose optimal thresholds as per the latent data distribution utilizing the mab framework. the algorithm uses the confidence scores to learn the optimal threshold. it makes the method robust to different levels of distortion in the test samples. 3) comprehensive experiments conducted on the ms-coco and flickr30k dataset, reveal the superior performance of capeen across all key metrics compared to previous methodologies. additionally, we validate the effectiveness of a-capeen in consistently determining optimal threshold values for various data distributions. ## related works image captioning: recent research has witnessed a surge in exploring efficient text description generation for input images @xcite . encoder-decoder frameworks have gained prominence for their exceptional performance in text generation tasks, leveraging contextual infor-mation @xcite @xcite @xcite @xcite @xcite @xcite . with the advent of the attention mechanism @xcite , the focus has shifted towards employing multiple layers of transformers as both encoders and decoders. early-exits. in recent years, the issue of inference latency has gained substantial attention @xcite . to address this issue, dnns are implemented with internal classifiers in the intermediate layers. @xcite explores early classification at intermediate layers for images, while spinn @xcite uses a mobile cloud setup to split the dnn. see @xcite performs the early exiting in a service outage scenario. multiple early exiting frameworks have also been proposed such as @xcite @xcite for improving early exits for image tasks dynamically choosing the depth of network for different regions for an image. other works like @xcite have utilized knowledge distillation in early exit framework but not for image captioning. also, it performs joint training of teacher and student that deteriorates the optimality of the backbone. numerous early exiting frameworks have been devised for natural language processing tasks @xcite @xcite @xcite @xcite @xcite , primarily based on the bert backbone. deecap @xcite introduces early exiting to image captioning, employing an imitation network to replicate outputs from computationally intensive transformer layers within an encoder-decoder architecture. similarly, mue @xcite applies early exits to ofa @xcite , a unified vision language model designed for multi-modal applications. multi-armed bandits in early exits. several works utilize the mab framework to adapt to different scenarios. notable methods like lee @xcite , dee @xcite aim to learn the optimal exit points in scenarios like mobile devices with restricted computational resources. additionally, epnet @xcite adopts an offline approach to learning when to exit based on considerations of computational overhead and accuracy. on the other hand, ueeucb @xcite employs a multi-armed bandit (mab) framework to dynamically learn the optimal exit strategy in an online and unsupervised manner. ueeucb relies on the assumption of strong dominance in neural networks, wherein accuracy increases with the layer number in the neural network. the key differences are: 1) to the best of our knowledge, improving the performance of early exits using knowledge distillation has not been studied for image captioning. 2) our online algorithm based on the mab setup overcomes the challenge of choosing the optimal threshold by adapting to the underlying latent distributions of the noise present in the test dataset. note: we are the first to apply knowledge distillation for image captioning. previous methods have applied knowledge distillation for early exits in text and image classification which are much simpler tasks than image captioning and they simultaneously perform distillation in a singlestage training. this cannot be extended to image captioning as it requires high-quality knowledge transfer else it can lead to incorrect learning paths. given the complexity of image captioning, we perform a two-stage training that not only transfers high-quality knowledge but also maintains the optimality of the backbone. this approach provides us with a state-of-the-art backbone and serves as a testbed for a-capeen. ## methodology in this section, we discuss our method of fusing knowledge distillation with early exit layers by treating the final layer classifier as the teacher and the early exit classifiers as the students. ## learning of thresholds recall the discussion from the introduction that a threshold set using the validation set may not result in better performance when test data distribution drifts from that of the train data distributions, especially when the images in the test data are distorted. as the threshold used in early exit decisions significantly impacts both computational requirements and accuracy, setting it appropriately as per the test data distribution is crucial for optimal performance. we address this challenge by adjusting the threshold as per input data distribution in an online fashion using the mab framework @xcite . in the mab setup, a decision-maker repeatedly selects from a set of arms (or actions) while adapting to the unknown environment. each arm corresponds to a specific choice or decision. the challenge lies in learning which arms yield the most favourable outcomes (highest reward) over time. this adaptation and learning process, central to mab setups, aligns with the dynamic nature of online learning problems, where decisions are made sequentially based on incoming data. by leveraging the principles of exploration and exploitation, mab frameworks facilitate learning the optimal action for the latent distribution. in this context, we define the action set as k thresholds for each exit as @xmath0, α 2 , . . . , α k }. for exit i, we define the latency factor as the cost of processing the sample from the 1st exit to the ith exit and denote it as o i . let [n ] denote the set {1, 2, . . . , n }. a token will exit the backbone only if the confidence is above the chosen threshold. for a given threshold α where α ∈ a, suppose that c @xmath1, and the reward is defined as: the reward could be interpreted as follows: if a sample exits at layer i, then the gain is the confidence gain in the inference at layer i compared to the inference at the first layer, and the cost incurred in processing the token from the first exit to the ith exit. the reward is the net gain, expressed as the difference between gain and cost. the objective is to maximize the expected reward as : pass s j till layer i 14: apply threshold β t and observe c i 15: return c j , j ← j + 1 34: end for we introduce an algorithm called adaptive-capeen (a-capeen) to adaptively choose the threshold values, and its pseudocode is outlined in algorithm 2. the input for this algorithm includes latency factors o i for each exit @xmath2, the exploration parameter γ. the expectation above is with respect to the randomness induced by latent sample distribution in the selection of actions. a policy denoted as π * is characterized as sub-linear when the average cumulative regret diminishes, that is, r(π * , t )/@xmath3. our primary goal is to devise a learning algorithm that has a sub-linear regret guarantee. the initialization of the algorithm involves ensuring that each action is played at least once. in the subsequent rounds, the algorithm selects the arm with the highest ucb index, denoted as β t (line 5). ucb indices comprise weighted averages of rewards q t (α) and incorporate confidence bonuses with γ as the exploration parameter (weight). this confidence threshold is associated with each exit i.e., β t is the threshold chosen for all exits for the tth token. the token is then processed until the confidence is above β t (line 10) and if the token never gains sufficient confidence, it is inferred at the final layer (line 16). after the token exits from one of the layers, the average reward of the played arm is updated (line 14). the caption starts with @xmath4. note that there are two counters, one is for the image, once @xmath5. hence, the algorithm learns faster as the thresholds (arms) are learnt on the number of tokens passed instead of the number of images. drawing from the analysis of ucb1 @xcite , the regret of a-capeen can be shown to be of the order o α∈a\α * log(n) ∆α , where ∆ α = r(α * )r(α) denotes the suboptimality gap. for completeness, we provide proof in the appendix a.1. training: the encoder-decoder backbone is initially fine-tuned for 10 epochs with a starting learning rate of 1e-5, which decays by 0.5 every 3 epochs. subsequently, self-critical training is employed for 5 epochs with an initial learning rate of 7e-6, also decaying by 0.5 every 3 epochs. the backbone weights are frozen post-fine-tuning, and exits are added to the decoder, whose weights are further trained for 5 epochs. the adam optimizer and a batch size of 8 are chosen, with a threshold of 0.6 chosen based on the accuracy-efficiency trade-off on the validation split. inference is conducted on the karpathy test split of the ms-coco dataset with a batch size of 1, using nvidia rtx 2070 gpus. results are presented in we set µ = 1/n . the latency cost o i could be understood as the cost of processing the samples from 1st exit to ith exit. hence we set the latency cost as o @xmath7. since the value of threshold α is adaptively chosen in this case, λ models the trade-off between accuracy and efficiency. a higher value of λ will provide higher accuracy while a lower value will give high efficiency. we use the value of λ = 1. a detailed ablation study of this hyperparameter can be found in the appendix b.4. ## conclusion we introduced a new encoder-decoder backbone for image captioning with early exits and knowledge distillation named capeen. using the student-teacher model, we trained early exit classifiers using knowledge distillation to capture highlevel semantic representations available at the deeper layers. we demonstrated that capeen offers a significant increase in speedup while maintaining a competitive performance guarantee. further, we introduced a-capeen, where the threshold used for early exit decisions can be adaptively learned for distributions that differ from the training data due to changes in the distortion levels.