DA d EE : Unsupervised Domain Adaptation in Early Exit PLM s (2024)

Abstract: Pre-trained Language Models (PLMs) exhibit good accuracy and generalization ability across various tasks using self-supervision, but their large size results in high inference latency. Early Exit (EE) strategies handle the issue by allowing the samples to exit from classifiers attached to the intermediary layers, but they do not generalize well, as exit classifiers can be sensitive to domain changes. To address this, we propose Unsupervised Domain Adaptation in EE framework (DAdEE) that employs multi-level adaptation using knowledge distillation. DAdEE utilizes GAN-based adversarial adaptation at each layer to achieve domain-invariant representations, reducing the domain gap between the source and target domain across all layers. The attached exits not only speed up inference but also enhance domain adaptation by reducing catastrophic forgetting and mode collapse, making it more suitable for real-world scenarios. Experiments on tasks such as sentiment analysis, entailment classification, and natural language inference demonstrate that DAdEE consistently outperforms not only early exit methods but also various domain adaptation methods under domain shift scenarios. The anonymized source code is available at https://github.com/Div290/DAdEE.

## introduction pre-trained language models (plms) such as bert @xcite , gpt @xcite , xlnet @xcite , and roberta @xcite have experienced substantial growth in size to attain state-of-the-art performances across a wide range of natural language processing tasks. despite their remarkable efficacy, plms suffer from inference latencies, limiting their utility in industrial applications requiring faster inference. prior research @xcite has also highlighted overthinking issues in plms. more specifically, shallow representations in the initial layers may suffice for correct inference of 'easy' samples, while final layer representations may overfit or be influenced by irrelevant features, resulting in poor generalization and computational inefficiency. to circumvent this, several methods such as pruning @xcite , quantization @xcite and knowledge distillation @xcite have been proposed. among them, early exit methods @xcite have gained significant attention, where inference can be made at classifiers attached at intermediary layer based on the hardness of the input samples. they perform adaptive inference strategies to address both efficiency and overthinking concerns. nevertheless, all these methods do not generalize well to new domains. also, as ee introduces more parameters in the exit classifers, it requires high-quality labeled training data to learn the exit weights. the cost of creating labeled training data is often prohibitively expensive. the question then arises: how can we adapt early exit plms to diverse domains in an unsupervised setup? domain adaptation is a vital technique in machine learning to ensure models perform well on data from a target domain, even if trained on a different source domain. unsupervised domain adaptation methods primarily focus on aligning source and target data in a shared feature space. this alignment is achieved by optimizing the model's representation to minimize domain discrepancies, such as maximum mean discrepancy @xcite or correlation instances @xcite . previous approaches, like adversarial training methods @xcite @xcite , have utilized adversarial objectives to bridge domain gaps by producing domain-invariant features at the final layer. however, since the early exit plms (eeplms) have added exits, domain adaptation methods cannot be directly used as they only focus on adapting to the final layer performance. this necessitates domain adaptation at every layer such that domain invariant features are not only available at the final layer but across all the layers of plm. to tackle this problem, we present a novel strategy that integrates adversarial domain adaptation techniques in early exits that adapt eeplms to diverse domains. our approach, named unsupervised domain adaptation in ee plms (dadee), emphasizes achieving domain invariant representations across all the layers such that exit classifiers trained on a source domain can be directly utilized for the target domain. this allows for the exit classifiers trained on the source domain to be directly utilized in the target domain without requiring labels. our method not only speeds up inference but also allows for fast and robust bridging of domain gaps compared to traditional methods that rely solely on final layer representations in adversarial setups, which is insufficient for such large and complex models. during adversarial training, the exits aid the adaptive process by mitigating the risk of catastrophic forgetting and mode collapse using knowledge distillation between source and target representations across all the layers. by leveraging multi-level domain feature adaptation, dadee enhances overall effectiveness across real-world applications, addressing the challenges of inference speed and domain adaptation in plms. also, in scenarios where the size of the source dataset is limited, early exit models can adapt more readily as these models do not depend on the prediction of just the final layer but multiple layers, improving generalization. our method achieves improved performance metrics on sentiment analysis, entailment classification, and natural language inference (nli) tasks. our main contributions are as follows: our method gets the best of both early exit and domain adaptation methods to simultaneously increase both performance and speed. ## related works in this section, we discuss studies relevant to our work in domain adaptation and early exiting. ## methodology in this section, we detail our method. we start with a plm such as bert/roberta with l layers. we attach exit classifiers to each layer and train them using the source dataset. we then perform multilevel adversarial training on target data. ## experiments in this section, we elaborate on all the experimental details of our work. ## conclusion we present a new method dadee for multi-level domain adaptation in plms using the early exit approach and adversarial training. the exits attached lower the chances of catastrophic forgetting by incorporating knowledge distillation across all layers. also, the attached exits decrease the inference time by inferring the easier samples early. t-sne plots and a-distance demonstrate the effectiveness of our method in domain adaptation. extensive experiments demonstrate that our method not only bridges the domain gap in a cross-domain setup but also provides faster inference making it suitable for real-world applications.