IMO: Greedy Layer-Wise Sparse Representation Learning for Out-of-Distribution Text Classification with Pre-trained Models (2024)

Abstract: Machine learning models have made incredible progress, but they still struggle when applied to examples from unseen domains. This study focuses on a specific problem of domain generalization, where a model is trained on one source domain and tested on multiple target domains that are unseen during training. We propose IMO: Invariant features Masks for Out-ofDistribution text classification, to achieve OOD generalization by learning domain-invariant features. During training, IMO employs a greedy algorithm to learn sparse representations for each layer in a top-down manner. It performs better than the opposite direction and learning of sparse representations for all layers simultaneously. Our comprehensive experiments show that IMO substantially outperforms strong baselines such as prompt-based methods and large language models, in terms of various evaluation metrics and settings. 1

heading: 1 Introduction
text: When deploying natural language processing (NLP) models trained on labeled data in the wild, it is well known that their predictive performance declines significantly on samples drawn from distributions that differ from their training data (Wang et al., 2021b). Although various domain adaptation (DA) methods have been proposed (Liu et al., 2022; Saunders, 2022), they assume the availability of labeled or unlabeled data from target domains, along with information about target domains. However, for many real-world applications, especially for early-stage businesses, users may apply their models to arbitrary data so the test data may well be Outof-Distribution (OOD). Hence, target domain information may not be available for DA. In addition, some training datasets are expensive to acquire so
*Corresponding Author. 1Codes are available at https://github.com/
WilliamsToTo/IMO.
they are available only in one domain. Therefore, this work focuses on single-source domain generalization (DG) for text classification, which aims to enable classifiers trained in one source domain to robustly work on the same classification tasks in any unseen OOD data without any model tuning.
Pre-trained large language models (LLMs) have drawn a lot of attentions due to their strong predictive performance across a variety of tasks. Although generative models or classifiers built on top of pre-trained LLMs outperform prior models in multiple domains, their performance is still not robust on tasks when the testing distribution differs substantially from the training distribution (Bang et al., 2023). Recent works (Wang et al., 2021a; Feng et al., 2023; Veitch et al., 2021) show that one of the key reasons is spurious correlations, which refer to the correlations between features and model outputs that are not based on causal relationships.
To take a step towards “train it once, apply it anywhere”, we propose a novel greedy layerwise Invariant Masking technique for OOD text classification, coined IMO, which selects domaininvariant features and key token representations from appropriate layers of a pre-trained deep transformer encoder to mitigate spurious correlations. The resulting hidden representations are sparse from the top layer to a specific layer of the pretrained model. We demonstrate the effectiveness of this technique through theoretical justifications and extensive experiments. Similar to (Zhang et al., 2021) on computer vision tasks, we shed light on how to apply sparsity as an effective inductive bias to deep pre-trained models for OOD text classification. Our contributions are:
• We propose IMO, a novel top-down greedy layer-wise sparse representation learning method for pre-trained text encoders for robust OOD classification by sharply reducing task-specific spurious correlations. In com-
2625
parison with bottom-up layer-wise and simultaneous search across all layers, we discover that the top-down greedy search is decisive for performance improvement.
• We develop a theoretical framework that elucidates the relationship between domaininvariant features and causal features. Additionally, we provide an explanation of how our method learns invariant features.
• Our comprehensive experimental results show that: (i) using IMO with BART (Lewis et al., 2020) significantly outperforms competitive baselines, including CHATGPT, on topic classification and sentiment polarity prediction in most of the target domains. Notably, CHATGPT has 10 times more parameters than BART; (ii) using IMO with CHATYUAN (Clue-AI, 2023) achieves superior performance in Chinese social factor classification compared to strong competitors like CHATGPT; (iii) IMO achieves robust OOD performance w.r.t. varying training data size. The accuracy difference between using 1k and 3.5 million training instances using IMO is less than 6%.

heading: 2 Related Work
text: Domain Generalization. Numerous DG methods have been proposed in the past decade, and most of them are designed for multi-source DG (Chattopadhyay et al., 2020; Zhao et al., 2020; Ding et al., 2022; Zhang et al., 2022; Lv et al., 2022). Existing DG methods can be roughly classified into two categories: invariant representation learning and data augmentation. The key idea of the former is to reduce the discrepancy between representations of source domains (Muandet et al., 2013; Li et al., 2018a,b; Shao et al., 2019; Arjovsky et al., 2020). The key idea of data augmentation is to generate out-of-distribution samples, which are used to train the neural network with original source samples to improve the generalization ability (Xie et al., 2020; Wei and Zou, 2019; Volpi and Murino, 2019).
This paper focuses on single-source DG, where the model is trained on a single source domain, then evaluated on multiple unseen domains. Wang et al. (2021c) proposes a style-complement module to synthesize images with unseen styles, which are out of original distributions. Qiao et al. (2020) proposes adversarial domain augmentation to encourage semantic consistency between the augmented
and source images in the latent space. Ouyang et al. (2023) uses a causality-inspired data augmentation approach to encourage network learning domaininvariant features. In terms of text classification, Ben-David et al. (2022); Jia and Zhang (2022) apply prompt-based learning methods to generate a prompt for each sample, then use large language models to predict labels. Causal Representation Learning (CRL). CRL addresses OOD generalization by exploring causal features that lead to labels. It is based on the assumption that causal features are stable across different environments or data selections. Since CRL is very ambitious and even infeasible in real application, a more practical method is invariant representation learning. Peters et al. (2016) investigated that invariant features, to some extent, infer the causal structure. Arjovsky et al. (2020) also assumes that prediction conditioned on invariant features is stable under different environments. Following such assumption, a strand of methods tries to learn invariant features by mitigating spurious correlated features, which vary across environments (Muandet et al., 2013; Chattopadhyay et al., 2020; Asgari et al., 2022; Izmailov et al., 2022; Hu et al., 2022b). This paper also follows this thread of methods, where we treat features that don’t affect prediction as spurious correlated features.

heading: 3 Learning Sparse Domain-Invariant Representations
text: LLMs are pre-trained on large-scale corpora so that they can capture rich correlations between tokens across various domains. To enable trained models incorporating LLMs to work across domains, our key idea originates from the Invariance Assumption that the conditional distributions of labels conditioned on invariant features do not change across domains (Peters et al., 2016). Zhang et al. (2021) show that there is a subnetwork inside a full network that can achieve better OOD performance than the full network, if this assumption holds. This hypothesis is also referred to as the functional lottery ticket (Liang et al., 2021). For a specific classification task, such as sentiment polarity analysis, the assumption indicates that there are certain sparse representations that are potential causes of labels (Wang and Jordan, 2022) across domains. Our method IMO realizes this idea by constructing sparse domain-invariant representations from the hidden representations of the selected layers of
pre-trained transformer-based encoders. Let X be the input space and Y be the label space, a domain is characterized by a joint distribution PXY on X × Y . In the context of a single source DG, we have access to the data of one source domain S = {(xs, ys)} drawn from its joint distribution PSXY . The goal is to learn a predictive model f : X → Y using only the data sampled from PSXY to minimize the prediction error on K unseen target domains, each of which is associated with a joint distribution P kXY . Due to domain drifts, PSXY ̸= P kXY ,∀k ∈ 1, ...,K.
Following (Quinzan et al., 2023), we make the same assumptions that (i) Y = f(Pa(Y )) + ϵ, where Pa(Y ) denote the features that directly cause Y , (ii) ϵ is exogenous noise, independent of any features, and (iii) Y has no direct causal effect on any features because classification labels are assigned after observing the corresponding texts. Although PSXY ̸= P (k) XY , ∀k ∈ 1, ...,K, we show in §3.3 that under all above assumptions, there is a sparse representation Hi such that the function Y = f(Hi) + ϵ exists in both source and target domains. We empirically study the presence of invariant representations and influence of spurious correlations in §4.3.
As illustrated in Figure 1, our method constructs sparse domain-invariant representations at both feature and token levels in a top-down manner. At the feature level, given embeddings produced by the transformer block of the top layer, a parametric mask layer identifies invariant features from the embeddings. Then, the mask layer is frozen and the algorithm learns the mask layer for the lower layer. The process is repeated until a pre-specified layer is reached. At the token level, a soft attention mechanism incorporates the selected features from the top layer to identify the tokens strongly correlated with Y and use attention weights to create aggregated sparse representations based on the selected features for binary classification. For multi-class classification tasks, a sparse representation is created for each class so that each of them can focus on class-specific information. The model is regularized during training to increase the divergences of the representations between classes.

heading: 3.1 Extraction of Invariant Features
text: Given a text input X = [xi]Ti=0, where xi is a token in X , a transformer-based pre-trained language model is employed to convert xi to a continuous token representation. We use hidden states produced
by each transformer layer l as token representations, denoted as H l = [hli] T i=0. h l i embeds both invariant features (useful for prediction in different domains) and spuriously correlated features (irrelevant for prediction) produced by layer l. Based on the Invariance Assumption, the invariant features h∗ ensure pk(Y |h∗) to be the same for each domain k. In a transformer layer l, the spuriously correlated features are filtered out by performing element-wise multiplication between token representation hli and a learnable mask m l.
A parametric filtering vector m = r ⊙ q contains zero and non-zero elements, where we define a trainable weight vector r ∈ Rd and a trainable pruning threshold vector s ∈ Rd. A unit step func-
tion g(t) = { 0 if t < 0 1 if t ≥ 0 is applied to get a binary mask q = g(|r| − s). By applying element-wise multiplication eli = h l i ⊙ml, the zero elements of m remove corresponding features in token embeddings hl, while non-zero elements characterize the importance of corresponding features (Liu et al., 2020).
As the unit step function g is not differentiable, we approximate its derivative by using the derivative estimator proposed in (Xu and Cheung, 2019) such that all parameters of a mask layer are trainable by using back-propagation and the family of stochastic gradient descent algorithms,
d dt g(t) =    2− 4|t|, −0.4 ≤ t ≤ 0.4 0.4, 0.4 ≤ |t| ≤ 1 0, otherwise.
(1)
Following (Xu and Cheung, 2019; Liu et al., 2020), we add a sparse regularization term Lsparse to the training loss to encourage the sparsity of mask layers:
Lsparse = N∑
i=1
exp(−si), s ∈ Rd (2)
where exp(−si) encourages high (but not extremely large) thresholds. A higher threshold leads to removal of more features. During inference, we retain the mask layers to retain invariant features while discarding irrelevant ones.

heading: 3.2 Identification of Invariant Tokens
text: Given a long token sequence, not all information is useful for target tasks. For example, function words, such as ‘the’, or ‘that’, provide little information for predicting sentiment polarity. Thus, we employ a token-level attention mechanism to focus on important tokens. Instead of using all features of a token representation, we compute attention scores by using only the invariant features. The proposed attention mechanism differs slightly between binary and multi-class classification.
Binary Classification. For binary classification, we treat the filtering vector mL from the last layer L as the query vector and compute the attention weight by performing the matrix product between mL and each token embedding from the last layer eLi : ai = m
LeLi . Here, the filtering vector and token embeddings are interpreted as matrices, with mL ∈ R1×d and eLi ∈ Rd×1. For an input token sequence, we aggregate the masked token embeddings to obtain a sequence representation v = ∑T i aie L i , where v ∈ R1×d. Finally, the sequence representation is fed into a fully-connected layer, followed by generating a distribution over the label space as follows: ŷ = softmax(vP).
Multi-class Classification. For the multi-class classification task, we propose using multiple mask layers mLy in the last layer L to capture corresponding features and tokens for labels y. The number of mask layers equals the number of labels. Each label has its own attention weights aLy = m L y e, and its own representation vLy = ∑T i a L yiei. Instead of using a fully-connected layer, we use a learnable weight vector per class to project vL
to a scalar: cL = vLpL, where vL ∈ R1×d and pL ∈ Rd×1. The rationale behind this is that each class should have its own weight vector and hidden representations for encoding class-specific information. Then, we concatenate these scalars to a vector c = [cL], and compute the predictive distribution by ŷ = softmax(c).
To encourage mask layers to extract labelspecific features, we propose the following regularization term to penalize pairwise cosine similarities between the corresponding mask layers (where N is the number of label-specific mask layers):
Ldist = 1 N(N − 1) ∑
i ̸=j cos(mi,mj). (3)
Training Procedure. Rather than training all mask layers simultaneously, we adopt a layer-wise training procedure to train them sequentially from the top layer to the bottom layer. As illustrated in Figure 1, for each layer, a new filtering layer, mL−i, is introduced on the top of the (L − i)th transformer layer, with i ∈ {0, 1, 2, ...L − 1}. Crucially, during this phase, the previously trained mask layers remain frozen to preserve their learned parameters. Upon each layer’s training completion, the model is stored as θL:L−i. This iterative procedure continues until the training of the most bottom filtering vector, m1, is completed. Consequently, a suite of models, ranging from θL to θL:1, is collected. We empirically determine the model’s efficacy by evaluating its performance on the validation set from the source domain. The best-performing model is chosen as the model to test on the target domains.
Objective Function. During training, the overall objective for binary classification is to (1) have good predictive performance on classification tasks and (2) maximize sparsity in mask layers to only keep invariant features. When training mask at layer l, the loss function is:
L = Lce + αLlsparsity (4) where Lce denotes the cross entropy loss and f denotes the predictive model. α, where α > 0, is a hyperparameter that controls the balance between predictive performance and sparsity in mask layers. Llsparsity is the sparse regularization term for mask at layer l.
For multi-class classification, we add a distance regularization term:
L = Lce + αLlsparsity + βLdist (5)
The hyperparameter β serves to calibrate the equilibrium between features specific to individual labels and those shared across multiple labels.

heading: 3.3 Theoretical Analysis
text: Based on our assumptions, Y = f(Hi) + ϵ exists, when Hi are the parent nodes of Y in the underlying causal graph. Because Hi are a subset among all possible hidden representations correlated with Y , there should be a subset of hidden representations serving as parents of Y , otherwise the invariance assumption does not hold. Due to the widely used faithfulness assumption stating that statistical independences imply the corresponding causal structures (Neal, 2020), we aim to find out Hi ⊥̸⊥ Y |Hj , where Hj is any feature set nonoverlapped with Hi.
We start our theoretical analysis by introducing a sparsity regularization term Ω(Y,Hi, ...,Hj), which counts the number of edges between Y and the random variables of features in an underlying causal graph, where Y is the variable for labels and Hk denotes the random variable of the feature hk. Then we introduce a loss function LΩ(Y,Hi, ...,Hj) = Lce+αΩ(Y,Hi, ...,Hj) with α > 0, analogous to Eq. (4).
Considering the simplest case that there is only a causal feature hi and a non-causal feature hj , the corresponding random variables are denoted by Hi and Hj . From any causal graphs in Fig. 2, we conclude that p(Y |Hi, Hj) = p(Y |Hi) so that the cross entropy term in LΩ remains the same when using the term p(Y |Hi), but the loss decreases after removing the non-causal feature from the loss due to the regularization term Ω(Y,Hi, Hj).
The two feature case can be easily extended to the case having more than two features. It is trivial that excluding a non-causal feature from the loss LΩ leads to the decrease of LΩ due to the Markov property of causal graphs (Peters et al., 2017).
Corollary 1. If there is no edge between Y and Hk in a causal graph, then LΩ(Y,Hi, ...,Hj) < LΩ(Y,Hi, ...,Hj , Hk).
During training, we start with a loss LΩ(Y,H1, ...,HN ) with a complete set of features. If a non-causal feature Hk is removed, LΩ(Y,Hi, ...,Hj) decreases according to Corollary 1. In contrast, if a causal feature Hk is removed, the cross entropy term increases because the mutual information I(Y ;Hk|Hi, ...,Hj) > 0. Namely, Hk adds additional information for predicting Y . However, in that case, LΩ(Y,Hi, ...,Hj) may still decrease if the increase of Lce is smaller than the decrease of the regularization term αLΩ(Y,Hi, ...,Hj), where α > 0. The exceptional case can be mitigated if α is sufficiently small. As a result, the loss LΩ provides an effective way to guide the search for the features serving as the causes of the labels, although we cannot recover the underlying true causal graphs. Herein, the loss (4) is a surrogate of LΩ(Y,Hi, ...,Hj) by using a deep neural network.

heading: 4 Experiments
text: We show that our approach significantly outperforms the competitive baselines in almost all settings, followed by empirically verifying that domain-invariant sparse representations indeed exist and spurious features deteriorate model performance in Sec. 4.3, as well as justifying the effectiveness of top-down greedy search strategy and individual modules in the ablation study.

heading: 4.1 Experimental Setup
text: Tasks and Datasets We evaluate our method on binary and multi-class classification tasks. Herein, we adopt accuracy as the metric for binary sentiment polarity classification and macro-F1 for multiclass classification tasks. All models are trained with five different random seeds to assess the statistical significance.
The datasets for binary sentiment analysis include Amazon Review Polarity (Zhang et al., 2015a), Yelp Review Polarity (Zhang et al., 2015a), IMDB (Maas et al., 2011), TweetEval Sentiment (Barbieri et al., 2020) 2 and Yahoo! Answers Sentiment (Li et al., 2019). For multi-class classification, we consider topic classification task in AG News dataset (Gulli, 2005; Del Corso et al., 2005; Zhang et al., 2015b) and social factor prediction task in SocialDial (Zhan et al., 2023, 2024). More details about datasets can be found in Appendix A.2.
2We remove all neutral instances to turn it into a binary classification task.
IMDB→ Amazon→ Yelp→ TweetEval→
Baseline Models. As Gulrajani and Lopez-Paz (2021) showed, simple empirical risk minimization (ERM) outperforms many SOTA domain generalization algorithms. So we finetune BERT (Devlin et al., 2019) and encoder of BART (Lewis et al., 2020) using cross-entropy loss as baselines. For Chinese text classification, we use BERT-zh (Devlin et al., 2019), BART-zh (Shao et al., 2021) and CHATYUAN (Clue-AI, 2023). Domain Generalization Models. PADA (BenDavid et al., 2022) is an example-based autore-
gressive prompt learning algorithm for domain generalization based on the T5 language model (Raffel et al., 2020). PDA (Jia and Zhang, 2022) is a prompt-based learning algorithm for domain generalization. Large Language Models. As CHATGPT shows promising zero-shot ability on various NLP tasks (OpenAI, 2023), we treat CHATGPT (gpt-3.5turbo) as a baseline. ALPACA-7B (Taori et al., 2023) is another baseline, which is finetuned from LLaMA 7B (Touvron et al., 2023) on 52K instruction-following data generated by selfinstruct (Wang et al., 2022). ALPACA-7B-LoRA is a finetuned ALPACA-7B model using low-rank adaptation (Wang, 2023; Hu et al., 2022a). CHATGLM-6B (THUDM, 2023) is an open large language model based on General Language Model (Du et al., 2022), optimized for Chinese questionanswering and dialogue. All LLMs use few-shot in-context learning. The specific query templates used for the LLMs can be found in Appendix A.3. Data Augmentation. Wiles et al. (2022); Gokhale et al. (2022) find data augmentation benefit domain generalization tasks. EDA (Wei and Zou, 2019) uses four operations (i.e., synonym replacement, random insertion, random swap, and random deletion) to augment text data. UDA (Xie et al., 2020) uses back-translation to generate diverse paraphrases while preserving the semantics of the original sentences. PGB (Shiri et al., 2023) generates syntactically and lexically diversified paraphrases using a fine-tuned BART.

heading: 4.2 Domain Generalization Results
text: Binary Classification. Table 1 reports the comparisons between our method and the baselines on sentiment polarity classification. Our method using
BART as backbone (i.e., IMO-BART) achieves superior performance over all baselines in 7 of 12 settings, and outperforms the best baseline CHATGPT by 2.63% on average. Interestingly, CHATGPT stands out as the best model in two out of 12 settings, though it remains unclear whether CHATGPT use those datasets for training. Moreover, it is noteworthy that data augmentation methods (i.e., BERT-EDA, BERT-UDA, BERT-PGB) show slightly inferior performance in comparison to the simple fine-tuning of BERT in terms of average accuracy. This suggests that simply back-translating or paraphrasing instances within source domains does not enhance performance on target domains.
Multi-class Classification. As shown in Table 2 and Table 3, our method outperforms all baselines in terms of average macro-F1 by 3.22% and 5.16% on AG News and SocialDial respectively. Among baselines, CHATGPT exhibits the strongest performance on both datasets and surpasses ALPACA-7B, ALPACA-7B-LoRA, and CHATGLM by a large margin. This superior performance shows that current open-source large language models still have a substantial performance gap with CHATGPT when handling difficult tasks.

heading: 4.3 Analysis of Spurious Features
text: Presence of Invariant Representations. We inspect shared representations at both feature and token levels. Invariant features are expected to have non-zero values across domains. Taking the best performing model IMO-BART in the sentiment analysis as an example, we train the model in each domain respectively and visualize its masks of the top layer in each domains. As depicted in Fig. 3, there are indeed a set of features shared across domains selected by the masks. We further compute Cosine similarities between the filtering vectors m of the top layer trained on different source domains. As shown in Table 8 in Appendix, their similarities range from 0.68 and 0.85. At the token level, we inspect the shared attention weights visualized in Fig.4 (see Appendix A3), which indicate the key-
words shared across domains in sentiment analysis, such as “great” and “slow”.
Impact of Spurious Correlations. To study whether our proposed masking mechanism indeed identifies robust features, we compare the performance of using the selected features with the nonselected ones. Specifically, we run additional experiments by replacing the learned binary masks q with |1− q|, followed by freezing all parameters except the classification head and training a model using those non-selected features. The results in Table 4 show that models using the nonselected features have an approximate 6% accuracy reduction in source domains and perform worse than using all features. In target domains, the corresponding performance drop using non-selected features is significantly higher than that using both our method as well as using all features. Hence, our masks indeed mitigate the use of spurious features.

heading: 4.4 Ablation Study
text: We compare top-down greedy search with alternative methods: bottom-up layer-wise search (B2T), simultaneous search (w/o sq), and only applying a mask on the last layer (last). From Table 1, 2 and 3, we can tell that top-down greedy performs significantly better than the alternative competitors. We conjecture that top-down layer-wise learning serves a regularization method that reduces the risk of loosing crucial features that are well correlated with Y and the corresponding optimization problem is easier to solve than learning all mask layers simultaneously. Representations from higher layers are shown to be more context-specific than lower layer representations (Ethayarajh, 2019). In contrast, the bottom up approach may drop key features in lower layers that significantly contribute to important higher layer features.
We compare variants of IMO by using varying backbone models and removing the corresponding components. For backbones, we compare BART with T5 and BERT, denoted them as IMO-T5, and IMO-BERT. To study the contribution of each component in our approach, we conduct experiments where we exclude the mask layers, attention mechanisms, or both. These models are denoted by w/o m, w/o a, and w/o am, respectively. The corresponding results are reported in Table 5, Table 11 and Table 10 in Appendix. For comparison between backbones, we find that encoderdecoder neural architectures (i.e., BART, T5) con-
IMDB→ Amazon→ Yelp→ TweetEval→
sistently achieve better performance than encoderonly models (i.e., BERT). Compared with variants that remove both the attention module and mask layers, IMO with the attention module or mask module has a significant performance improvement in terms of accuracy or F1 on average, which justifies the usefulness of both modules.
Additionally, we compare IMO with various sparsity methods to implement mask layers, including STR (Kusupati et al., 2020), STE (Bengio
et al., 2013; Liu et al., 2020), and Scalar, which uses a learnable single scalar instead of the threshold vector s. All those alternative methods lead to a significant drop, as seen in Table 5.
To explore the influence of source domain training data size on performance within target domains, we train models based on BART with and without our method on the Amazon review dataset with varying sizes of training data (i.e., 1k, 10k, 100k, 1M, and 3.6M). The results in Table 6 show that our method depends significantly less on training data size, though more training data can improve the performance overall. Notably, 1k training data yields a remarkable decline for the models without using IMO, while the corresponding performance reduction is significantly less by using our method.

heading: 5 Conclusion
text: This paper presents a novel method, coined IMO, which is a greedy layer-wise representation learning method aiming to improve single-source domain generalization on pre-trained deep encoders for text classification tasks. The key idea is to retain invariant features through trainable mask layers
and incorporate a token-level attention module to focus on the tokens that directly lead to the prediction of labels. Through extensive experiments, we demonstrate that IMO achieves superior OOD performance over competitive baselines on multiple datasets. The visualization of masks and attention weights empirically justifies the effectiveness of identified invariant sparse representations.

heading: Limitations
text: Our work focuses on the text classification task, intending to investigate how to learn invariant features to improve out-of-domain generalization. However, the proposed method has promising potential for domain generalization in various NLP tasks, such as question answering and text generation tasks. Future work may consider more tasks beyond text classification.
It is worth noting that IMO needs to be trained in a large source domain. The size of the source domain should ideally exceed 10,000 samples to achieve consistently good performance. However, this requirement may pose challenges in lowresource learning scenarios.

heading: Ethics Statement
text: This research is dedicated to augmenting the reliability and safety of text classification models, particularly in the context of domain shifts, as highlighted by Ribeiro et al. (2020). By focusing on the learning of invariant features across diverse domains, our approach aims to provide tangible benefits to applications that serve a wide array of user groups. From a user-centric perspective, the implementation of our methodology is expected to bolster the trustworthiness and diminish potential biases in language models.
It is pertinent to note that our study does not involve human subjects, nor does it contravene any legal or ethical standards. We foresee no detrimental impacts arising from our research endeavors. The experimental work underpinning this study was exclusively conducted using datasets that are publicly accessible. Our overarching goal is to foster enhanced academic and societal consciousness regarding the challenges of domain generalization in the field of natural language processing.

heading: Acknowledgement
text: This material is based on research sponsored by DARPA under agreement number
HR001122C0029. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon.

heading: A Appendix

heading: A.1 Background of Causal Representation Learning
text: Causal representation learning aims to learn and leverage the causal relations within data to enhance model generalization and robustness against distribution shifts in the data generation process (Shen et al., 2021). This approach differs from traditional machine learning methods that predominantly focus on correlational patterns without distinguishing causation and correlation. Causation refers to the underlying mechanisms that connect variables, implying that alterations in a causal variable will consequentially affect the associated effect variable, a
process known as an intervention. In contrast, correlation does not necessarily indicate a direct mechanistic link. For instance, a model might infer ’it is raining’ upon observing people with open umbrellas, which recognizes a correlation between these events. However, the act of closing umbrellas does not influence the weather. This example shows the difference between correlation and causation. Predictions relying on correlation may yield erroneous outcomes when the environment (i.e., data distributions) change. For example, if umbrellas are opened due to sunlight rather than rain, a model trained on the correlation might inaccurately predict rain. This indicates the importance of making predictions based on causation rather than correlation. A causal model should predict the weather based on temperature, humidity, air pressure, etc. Prediction based on causation can enhance out-ofdistribution performance, which is supported by the assumption that causal relations remain constant across diverse environments (Bühlmann, 2018).
However, learning the complete causal structure is ambitious and may not be realized in practice. A more feasible approach involves identifying invariant features that reliably predict target variables across varying environments. A series of methods (Muandet et al., 2013; Chattopadhyay et al., 2020; Asgari et al., 2022; Izmailov et al., 2022; Hu et al., 2022b) have been proposed by leveraging the invariance between environments. They leverage the fact that when conditioning all direct causes of a target variable, the conditional distribution of the target will not change when interventions are applied to all other variables in the model except the target itself. Building upon this foundational idea, our work seeks to identify and utilize these direct causes (i.e., invariant features across environments) for the accurate prediction of target variables in the out-of-distribution setting.

heading: A.2 Experiment Datasets
text: AG News (Gulli, 2005; Del Corso et al., 2005; Zhang et al., 2015b) is a collection of news articles used for topic classification, which contains news titles, and news descriptions assigned to four topic classes. Titles and descriptions are employed as different domains. For social factor prediction, we use SocialDial (Zhan et al., 2023), which is a Chinese socially-aware dialogue corpus consisting of synthetic conversations generated by CHATGPT and human-written conversations. Both are annotated with social factors such as location, social distance,
and social relation. Synthetic conversations and human-written conversations are considered as different domains. The statistics of datasets are listed in Table 7.

heading: Binary Classification Dataset Domain #Train #Dev #Test

heading: Multi-class Classification Dataset Domain #Train #Dev #Test

heading: A.3 Training details
text: We use the encoder of BART (Lewis et al., 2020) as the default pre-trained language model. All models are trained up to 100 epochs with a minibatch size of 32 in the source domain. We use the Adam (Kingma and Ba, 2015) optimizer with hyperparameters tuned on the validation sets. As a result, we run Adam with β1 = 0.9 and β2 = 0.999. The learning rate is 5× 10−5. We use a linear learning rate scheduler that dynamically decreases the learning rate after a warm-up period. All experiments are conducted on NVIDIA A40 GPU.
The process of model selection in domain generalization is inherently a learning problem. In this approach, we employ training-domain validation, which is one of the three selection methods introduced by Gulrajani and Lopez-Paz (2021). We divide each training domain into separate training and validation sets. Models are trained on the training set, and the model that achieves the highest accuracy on the validation set is chosen as the selected model.
When using large language models to predict target classification labels, the query template for sentiment analysis is: “There are some examples about sentiment analysis: {examples}. Given text: {sentence}, what is the sentiment conveyed? Please select the answer from ‘positive’ or ‘negative’.”. The query template for AG News topic classification is "There are some examples for topic classification: {examples}. Given text: {sentence}, what is the topic of this text? Please select the answer from ‘Business’, ‘Sci/Tech’, ‘World’ or ‘Sports’.” The query templates for SocialDial are “There are some examples for classification: {exam-
ples}. Given conversation: {conversation}, what’s the location/social distance/social relation of this conversation? Please select the answer from {labels}”3 (Min et al., 2022; Wang et al., 2023; Yang et al., 2023).
A.4 Visual Explanation To intuitively show how the attention module and mask module work in models, we visualize attention weights on tokens and mask vectors in Figure 4 and 3, respectively. We also demonstrate cosine similarities between mask vectors m trained on different source domains and Jaccard similarities between binary vectors q trained on different source domains on Table 8 and Table 9, respectively.
From Figure 4, we can find that our model primarily focuses its attention on sentiment-indicative tokens. Notably, positive reviews exhibit high attention weights for tokens like ‘good,’ ‘great,’ and ‘nice,’ indicating their significance. Conversely, negative reviews assign high attention weights to tokens such as ‘horrible’ and ‘slow,’ highlighting their importance in expressing negativity.
In Figure 3, we visualize mask vectors m and binary vectors q trained on different source domains across dimensions. It can be observed that certain dimensions are consistently assigned zero (or nonzero) values across different training domains, indicating our mask layers can capture some features that are irrelevant (or invariant) across domains. We quantify invariant features across domains by computing vector similarity. We calculate cosine similarities between different mask vectors m. The results are shown in Table 8. We can find that most mask vector pairs have over 0.75 similarity, except the Yelp-TweetEval pair, which is probably because of a larger divergence between Yelp and TweetEval domains. Table 9 shows Jaccard similarities between binary vectors q. Most binary vector pairs have similarities of over 0.5, except the Yelp-TweetEval pair, with a similarity of 0.45.

heading: A.5 Additional Experimental Results
text: 3Since SocialDial is a dataset in Chinese, we provided queries translated from Chinese to English for use with CHATGPT.

heading: SocialDial

heading: Future Work
text: Future work may consider more tasks beyond text classification. It is worth noting that IMO needs to be trained in a large source domain. The size of the source domain should ideally exceed 10,000 samples to achieve consistently good performance. However, this requirement may pose challenges in lowresource learning scenarios