DAdEE: Unsupervised Domain Adaptation in Early Exit PLMs

TL;DR

This paper introduces DADEE, an unsupervised domain adaptation framework for early exit pre-trained language models, improving efficiency and robustness across domains through multi-level adversarial knowledge distillation.

Contribution

DADEE is the first to combine unsupervised domain adaptation with early exit strategies in PLMs using multi-level adversarial training and knowledge distillation.

Findings

DADEE outperforms existing early exit methods in domain shift scenarios.

DADEE reduces inference latency while maintaining high accuracy.

DADEE enhances domain invariance across all model layers.

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.