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-of-Distribution text classification, to achieve OOD generalization by learning invariant features. During training, IMO would learn sparse mask layers to remove irrelevant features for prediction, where the remaining features keep invariant. Additionally, IMO has an attention module at the token level to focus on tokens that are useful for prediction. Our comprehensive experiments show that IMO substantially outperforms strong baselines in terms of various evaluation metrics and settings.

## introduction 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 @xcite . although various domain adaptation (da) methods have been proposed @xcite , 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 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 @xcite . recent works @xcite @xcite 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 @xcite 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: 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 @xcite @xcite . this paper focuses on single-source dg, where the model is trained on a single source domain, then evaluated on multiple unseen domains. @xcite proposes a style-complement module to synthesize images with unseen styles, which are out of original distributions. @xcite proposes adversarial domain augmentation to encourage semantic consistency between the augmented and source images in the latent space. @xcite uses a causality-inspired data augmentation approach to encourage network learning domaininvariant features. in terms of text classification, @xcite ; @xcite 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. @xcite 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 @xcite @xcite @xcite . this paper also follows this thread of methods, where we treat features that don't affect prediction as spurious correlated features. ## learning sparse domain-invariant representations 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 @xcite . @xcite 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 @xcite . 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 @xcite 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 p xy on x × y. in the context of a single source dg, we have access to the data of one source domain @xmath0, y s )} drawn from its joint distribution p s xy . the goal is to learn a predictive model f : @xmath1, each of which is associated with a joint distribution p k xy . due to domain drifts, p s xy ̸ = p k xy , ∀@xmath2, ..., k. following @xcite , we make the same assumptions that (i) @xmath3, 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 p s xy ̸ = p (k) xy , ∀@xmath4, ..., k, we show in §3.3 that under all above assumptions, there is a sparse representation h i such that the function @xmath5. 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. ## experiments 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. ## conclusion 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.