Prompt-based Distribution Alignment for Domain Generalization in Text Classification (2022)

Abstract: Prompt-based learning (a.k.a. prompting) achieves high performance by bridging the gap between the objectives of language modeling and downstream tasks. Domain generalization ability can be improved by prompting since classification across different domains can be unified into the prediction of the same set of label words. The remaining challenge for domain generalization by prompting comes from discrepancies between the data distribution of different domains. To improve domain generalization with prompting, we learn distributional invariance across source domains via two alignment regularization loss functions. The first is vocabulary distribution alignment, which uses a Kullback-Leibler divergence regularization on source-domain vocabulary distributions. The second is feature distribution alignment, which uses a novel adversarial training strategy to learn domain invariant representation across source domains. Experiments on sentiment analysis and natural language inference show the effectiveness of our method and achieve state-of-the-art results on six datasets.

## introduction pretrained language models (plms) have achieved promising results on a range of natural language processing (nlp) tasks @xcite @xcite . the framework of tuning a plm by task data has achieved competitive performance @xcite . however, it suffers significant performance degradation when the tuned plm is directly applied to out-of-domain examples @xcite . to tackle the problem of domain shift where the training set and the test set come from different data distributions, unsupervised domain adaptation @xcite uses unlabeled target data cooperated with the labeled source data for training. however, in many real-world systems, access to unlabeled data in the target domain is also impossible. this paper focuses on domain generalization (dg), the practical and challenge setting, where a model trained on multiple source domains can be directly generalized to a target domain without any labeled or unlabeled data from the target domain @xcite . prompt-based learning (a.k.a. prompting) @xcite @xcite makes better use of pretrained knowledge by bridging the gap between objectives of language modeling and downstream task. in particular, prompt-based text classification uses an identical projection from the probability distribution on label worlds to the probability distribution on classification classes and has achieved state-of-the-art results @xcite @xcite . intuitively, dg benefits from prompting in that different domains can share a unified set of label words, and thus no additional parameters such as output layers can carry domain-specific information to hinder the generalization performance for an unseen domain. however, there remains a crucial challenge for directly using prompting for dg. as shown in figure 1 (middle right), different domains have intrinsically different feature distributions, and instances from different domains have different predicted vocabulary distributions (top left). for example, the preferable positive set of label words for book reviews can consist of "helpful" and "well-written" more frequently, but for film reviews, it can consist of "amazing" and "real", etc more frequently. to tackle the above challenge, we propose two regularization loss functions to better align different domains, so that the gap between source domains and a new target domain can be reduced. the first alignment loss is vocabulary distribution alignment (vda), which is used to reduce the divergence of predicted vocabulary distributions of different domains. to achieve this goal, we use kullback-leibler divergence regularization between predicted vocabulary distributions of different source domains in the same category. as a result, the optimized distribution alignment for the predicted vocabulary can improve the generalization performance for unseen domains, as shown in figure 1 (top). the second alignment loss is feature distribution alignment (fda), which is used to learn a generalized feature space corresponding to the predicted token for different domains. to achieve this goal, we use an adversarial training strategy @xcite to reduce the domain discrepancy @xcite . different from traditional domain adversarial training for da @xcite , we learn the domain-invariant representation for each category across source domains. as a result, the optimized feature distribution alignment across source domains can improve dg performance, as shown in figure 1 (middle). we conduct experiments on three sentiment analysis datasets and three natural language inference datasets. results show that our method can effectively learn distribution alignment and achieve the best results under both leave-one-domain-out evaluation and cross-dataset evaluation settings. to our knowledge, we are the first to use regularization to improve dg for prompting and the first to simultaneously learn domain invariance over representation and predicted probability in dg. the code will be released at https://github.com/ jiachenwestlake/pda . ## related work prompt-based learning. adapting the plms for downstream tasks via fine-tuning has become a dominant framework for nlp in recent years @xcite @xcite . prompt-based learning applies a fixed function to condition the model, so that the language model gets additional instructions to perform the downstream task. prompt tuning the plms with manually designed prompts has achieved promising results on few-shot classification tasks such as sentiment analysis and natural language inference @xcite . however, designing prompting function is challenging and requires heuristics. to this end, recent work propose to apply prompts as learnable parameters, such as soft prompts @xcite @xcite , p-tuning v2 @xcite and prefix tuning @xcite . prompts capture task-specific knowledge with much smaller additional parameters than its competitors, such as adapter @xcite and lora @xcite . however, relatively little work has considered that the domain discrepancy can affect the performance of prompt-based learning. this paper focuses on improving the domain generalization ability for prompt-based learning via a domain invariance learning strategy. domain generalization. different from domain adaptation (da) @xcite , dg improves out-of-domain (ood) robustness with no need to explicitly know the data distribution in target domain @xcite . while a plethora of dg methods have been proposed for objective recognition during the last decade @xcite , invariance learning has shown high success and has become a prevalent approach for dg. such algorithms include domain adversarial training @xcite @xcite ; deep coral @xcite , which optimizes the second-order statistics over feature space and irm @xcite , which tackles domain shift with intrinsic relationship between feature representation and labeling prediction. in contrast to these work, we tackle dg via simultaneously learning the domain invariance for both feature representation and predicted probability using novel regularization methods with prompting. besides, recent work on da with plms use domain-invariant feature regularization based on adversarial fine-tuning @xcite @xcite , which are difficult to directly apply for dg. in terms of prompting for classification on ood instances, a recent work pada @xcite first generates a prompt for each instance and then applies the example-specific prompt to a t5 model for classification. orthogonal to this example-based method, we aim to tackle a more general problem setting, generalization for any data distribution. ## prompt-based distribution alignment the overall framework of our method is shown in figure 3 . the proposed method consists of three parts: (i) prompt-tuning for text classification ( §3.1). (ii) vocabulary distribution alignment (vda) via kullback-leibler divergence regularization ( §4.1) and (iii) feature distribution alignment (fda) via labeling awareness domain adversarial training ( §4.2). ## experiments we evaluate our method on both sentiment classification and natural language inference. ## conclusion we investigated how to improve out-of-domain (ood) robustness for prompt-based learning, by learning domain invariance via vocabulary distribution alignment and feature distribution alignment with prompting. experiments show that we achieve the best results on six datasets of sentiment analysis and natural language inference under both the standard leave-one-domain-out evaluation setting and a novel cross-dataset evaluation setting compared with a range of strong baselines. moreover, the proposed distribution alignment method can be seen as a general regularization technique for domain generalization beyond the text classification task.