FREE : Fast and Robust Vision Language Models with Early Exits (2025)

Abstract: In recent years, Vision-Language Models (VLMs) have shown remarkable performance improvements in Vision-Language tasks. However, their large size poses challenges for real-world applications where inference latency is a concern. To tackle this issue, we propose employing Early Exit (EE) strategies in VLMs. However, training exit classifiers in VLMs is challenging, particularly with limited labeled training data. To address this, we introduce FREE, an adversarial training approach within a GAN-based framework. Here, each exit consists of a transformer layer and a classifier. The transformer layer is adversarially trained to produce feature representations similar to the final layer, while a feature classifier serves as the discriminator. Our method focuses on performing input-adaptive inference that increases inference speed with minimal drop in performance. Experimental results demonstrate the effectiveness of our approach in enhancing accuracy and model robustness by mitigating overthinking and the phenomenon of mid-crisis that we highlight. We experimentally validate that our method speeds up the inference process by more than 1.51\times while retaining comparable performance. The anonymized source code is available at https://github.com/Div290/BLIPEE.

## introduction vision-language pre-training (vlp) has evolved significantly with the emergence of sophisticated pre-trained vision language models (vlms). these models have consistently pushed the performance boundaries across various vision-language tasks. however, their demanding computational requirements and inference latency pose challenges for real-time applications. several models, such as blip-2 @xcite , @xcite etc., leverage off-the-shelf large-scale pre-trained models as building components with their parameters frozen. this reduces vlms training parameters but makes them extremely slow during inference, as highlighted in (bajpai and hanawal, 2024a) leading to higher inference time. the use of the language model (lm) component with frozen parameters not only makes vlm susceptible to overthinking but also to another phenomenon that we term mid-crisis @xcite . this phenomenon occurs when intermediate layers suffer performance drops due to the search for irrelevant features. while initial layers capture shallow representations and syntactic information, and deep layers learn semantic-fusion relations @xcite , intermediary layers tend to capture dataset patterns that degrade their performance, even losing the information learned by initial layers, and the model has to regain the lost information again in deeper layers. we illustrate this phenomenon in the left figure of fig. 3 and sec. 3.1, showcasing accuracy on the vqav2 (visual question answering) @xcite dataset across different exits when trained with dedicated classifiers as in vanilla early exit setup for the intermediate layers. this raises the question: how can we mitigate mid-crisis and overthinking to enhance the accuracy and efficiency of vlms? we address this using early exit (ee) techniques @xcite @xcite , an input-adaptive method that reduces computational costs by bypassing certain layers while preserving the knowledge encoded in largescale models. since real-world datasets contain both "easy" and "hard" samples, ee ensures adaptive computation per sample. however, applying ee to vlms presents challenges: (1) exit classifiers introduce significant overhead-e.g., a single exit for opt 2.7b @xcite decoder adds approximately 130m parameters; (2) training these exits requires sub-how many dogs are there? looks like 2. ## related works we discuss the vlps with lm components and ee strategies related to our work below. vision-language pre-training: different model architectures have been proposed for specific downstream tasks in vlps, including dual-encoder architectures @xcite , encoder-decoder architectures @xcite @xcite . various pre-training objectives have also been introduced, such as image-text contrastive learning @xcite @xcite , the image-text matching @xcite @xcite , and masked language modeling @xcite @xcite . however, these end-to-end models are inflexible to leverage readily available pre-trained models, such as llms @xcite @xcite . recent approaches in vision-language pretraining have adopted the strategy of utilizing offthe-shelf pre-trained models and keeping them frozen during training. some methods freeze only the image encoder @xcite @xcite , and recent lit @xcite which uses a frozen pre-trained encoder for clip @xcite pre-training, while others freeze the language model to leverage knowledge from language-only pre-trained models for vision-to-language generation tasks @xcite @xcite @xcite . the primary challenge lies in aligning visual features with text space. to address this challenge, techniques like frozen @xcite finetune image encoders or insert new cross-attention layers into language models to incorporate visual features. blip-2 @xcite employs both frozen image encoders and language models for visionlanguage tasks, achieving strong performance. early exits: to minimize inference latency in deep neural networks, branchynet @xcite introduced attaching exits classifiers at intermediary layers. this concept was extended by shallow-deep @xcite , which effectively determines when to exit based on confidence distribution at each exit classifier. approaches like @xcite @xcite improve ees for image tasks by dynamically choosing different depths for different regions of the image. approaches like @xcite have employed knowledge distillation for image classification. for nlp tasks, several early exit frameworks have emerged @xcite @xcite @xcite @xcite @xcite @xcite , primarily built on the bert backbone. deecap @xcite introduces ee for image captioning tasks, employing an imitation network to replicate outputs from transformer layers. mue @xcite applies ee to ofa @xcite , a vlm tailored for multi-modal applications. capeen (bajpai and hanawal, 2024a) makes the ees robust to noise by adapting to the exiting threshold. the key differences in our work are: 1) we employ adversarial training for efficient learning of ee models. 2) our method can work under both supervised and unsupervised setups by utilizing the zero-shot capabilities of the vlms, while previous methods require a good amount of highquality labeled training data. ## methodology we begin with a pre-trained vision-language model (vlm) comprising three key components: a visual encoder, a projection layer, and a language model (lm). the visual encoder processes the input image, which may be a standalone image encoder @xcite or include an imagegrounded text encoder (e.g., q-former in blip-2). the projection layer then transforms the visual features to align with the lm's input space which is passed as an input to the decoder. as autoregressive decoding requires multiple forward passes of lm, our focus will primarily be on optimizing the lm component for efficient inference. we assume that lm consists of n layers, and we attach exits to the k chosen layers. each exit consists of one exit transformer (et) layer with the same configuration as one lm layer and an exit classifier (ec). the exit layer is such the parameters of the et are trainable, and those of ec are frozen. before indulging into architectural details, we motivate our method by highlighting the major concerns on the well-known blip-2 model. ## conclusion in this study, we introduced a novel inference technique free, which leverages adversarial training of exits alongside the zero-shot capabilities of the vlms. by employing free, we effectively reduce the dependency on a vast amount of labeled training data typically required for exit training. our approach involves adversarially training exits to generate representations similar to those of the final layer, thereby minimizing the need for extensive labeled data. moreover, our exit design reduces the number of trainable parameters, resulting in lower computational costs. experimental results demonstrate that our method significantly enhances inference speed while yielding high-quality outputs.