Invariant Language Modeling (2022)

Abstract: Modern pretrained language models are critical components of NLP pipelines. Yet, they suffer from spurious correlations, poor out-of-domain generalization, and biases.Inspired by recent progress in causal machine learning, in particular the invariant risk minimization (IRM) paradigm, we propose invariant language modeling, a framework for learning invariant representations that generalize better across multiple environments. In particular, we adapt a game-theoretic implementation of IRM (IRM-games) to language models, where the invariance emerges from a specific training schedule in which all the environments compete to optimize their own environment-specific loss by updating subsets of the model in a round-robin fashion.We focused on controlled experiments to precisely demonstrate the ability of our method to (i) remove structured noise, (ii) ignore specific spurious correlations without affecting global performance, and (iii) achieve better out-of-domain generalization.These benefits come with a negligible computational overhead compared to standard training, do not require changing the local loss, and can be applied to any language model. We believe this framework is promising to help mitigate spurious correlations and biases in language models.

## introduction while modern pretrained transformer models have led to dramatic progress on many nlp tasks, important limitations remain. in particular, pretrained language models suffer from poor generalization, even under small perturbations of the input distribution @xcite . indeed, these models encode @xcite and exploit @xcite spurious correlations, i.e., correlations that do not generalize across data distributions. since language models are trained on large unverified corpora, they also suffer from biases @xcite . here the term "biases" refers to correlations that may or may not be spurious according to the available textual data distributions, but are nevertheless undesired. existing techniques aiming to remove spuriousness or biases involve computationally expensive domain alignment @xcite @xcite , domain transfer @xcite , or the addition of penalty terms to the loss targeted at specific undesired correlations @xcite . alternatively, data preprocessing @xcite or manipulation such as counterfacual data-augmentation @xcite can yield datasets where undesired correlations are less present. pretraining with larger and more diverse datasets can also help @xcite . however, recent works on the theory of causality @xcite argue that removal of spurious correlations requires altogether different learning and training paradigms going beyond purely statistical learning. indeed, generalization, spuriousness, and biases are all better understood in the language of causality @xcite . intuitively, causal relationships are the ones expected to be stable @xcite and generalizable @xcite . when the causal graph underlying the data generation mechanism is known, there exist causal identification algorithms to distinguish desired from undesired correlations @xcite . however, for complex tasks of interest, the underlying causal model is not known. language modeling is one of these tasks, where it is unclear what would even be the relevant random variables constituting the causal model. therefore, causal identification from the causal graph seems out-of-reach for language modeling. similarly, removing undesired correlations one by one is impractical due to the sheer amount of possible correlations to consider. in this work, we propose to benefit from recent progress in causal machine learning to offer a new and more flexible lever for dealing with spuriousness and biases. we take inspiration from the invariance principle, which states that only relationships invariant across training environments should be learned @xcite . under specific assumptions, the invariant representation would then only encode the causal relationships relevant to the task and should thus generalize. environments correspond to different views of the learning task, i.e., different data distributions. the invariance principle is illustrated by fig. 1 with a simplified causal model as an example. e represents environment indices, y is the target variable, x c are the causal features, such that e[y |x c ] is stable across environments (e[y |x c , e] = e[y |x c ]), and x s are the spurious features, not generalizing across environments (e[y |x s , @xmath0, @xmath1. language models trained with standard empirical risk minimization (erm), denoted as elm in this work, exploit all correlations available during training and aim to learn e[y |x c , x s ]. our proposed invariant language models, denoted as ilm, focus on invariant features and aim to learn e[y |x c ]. in practice, since the causal model is unknown, it is the choice of environments that defines what correlations are spurious. invariant learning with appropriate choices of environments is the lever we propose to employ to more flexibly deal with spuriousness and biases. a practical formulation of the invariance principle was proposed by @xcite . they introduced invariant risk minimization (irm), an alternative to erm as a training objective enforcing the learning of invariant representations. @xcite later improved the training procedure to solve the irm objective with a method called irm-games. unlike previous methods for removing biases and spurious correlations, irm-games does not modify the loss with a regularization term and does not compute domain alignment (or matching) statistics. the invariance benefits come from the specific training schedule where environments compete to optimize their own environmentspecific loss by updating subsets of the model in a round-robin fashion. we argue that the irm paradigm, and irmgames specifically, is well-suited to improve nlp systems. textual data naturally comes from different environments, e.g., encyclopedic texts, social media posts, news articles, etc. moreover, not knowing the causal mechanisms behind language generation within these environments is not a blocker, as the relevant variables can now remain latent. by adapting irm-games to language modeling, we introduce invariant language modeling (ilm), where the training of existing pretrained models is continued to enforce invariant representations, using a simple and efficient modification of the training process. we then investigate the ability of ilm to deal with undesired correlations in a series of controlled experiments, answering our core research question: does the invariance principle give rise to a practical strategy to deal with spurious correlations in language models? contributions. (i) we introduce a new training paradigm (ilm) for language models based on the invariance principle (sec. 3). thanks to the use of the irm-games training schedule (see sec. 2), our ilm framework results in negligible computational overhead compared to standard erm training, does not require changing the local loss, and is agnostic to the language model architecture. (ii) in a series of controlled experiments (sec. 4), we demonstrate the ability of ilm to remove structured noise (sec. 4.1), ignore specific spurious correlations without affecting global performance (sec. 4.2), and achieve better out-of-domain generalization (sec. 4.3). (iii) we discuss our contributions in relation to previous work (sec. 5). (iv) finally, we release huggingfacecompatible code for training ilm using existing language model checkpoints @xcite : https://github.com/epfl-dlab/ invariant-language-models 2 background ## model we introduce a way to train language models inspired by the irm-games setup. this involves distinguishing the shared invariant feature extractor ϕ from the environment-specific w e 's. with modern language model architectures, a natural choice emerges: ϕ as the main body of the encoder, and w e as the language modeling head that outputs the logits after the last layer. formally, suppose we have n environments consisting of data {(x e ,y e )} @xmath2,...,n . for a batch (x i , y i ) ∼ p(x i ,y i ) from environment i, the model output is formed using an ensemble of n language modeling heads {w e } @xmath3,...,n on top of the trans- then, a (masked) language modeling loss l is computed on the model output ŷ. note that it is the predictions of the n heads that are averaged not the weights or the gradients. no head gets to predict alone; the n heads always predict together as an ensemble. the heads are subject to competitive gradient updates in a round-robin fashion as described below, which in turn creates the conditions that enforce the invariance. training. the training of ilm follows the pseudocode described in alg. 1, where environments take turns to send a batch of data and update ϕ and their associated head. an illustration is provided in appendix a. each head periodically gets an opportunity to pull the global ensemble classifier w and the feature learner ϕ towards fitting the distribution of its associated environment. intuitively, since each head gets the same amount of updates, the game converges to a global classifier that is simultaneously optimal for each environment, as demonstrated by @xcite . while the v-irm algorithm of @xcite only updates ϕ periodically, we found it more stable to update it together with every head update. algorithm 1 ilm training 1: initialize(ϕ, {w e } @xmath4: for @xmath5, 2, . . . , an advantage of this implementation is that invariance is obtained with few modifications to language models. such simplicity arises from our leveraging of irm-games, where invariance comes from the training schedules and ensembling of classifiers. furthermore, we implement two baselines that appear similar but do not enjoy the same theoretical properties: mtlm and enslm. the multitask baseline @xcite , mtlm, also uses data split into environments with one head per environment and each environment being seen as a different task. the ensemble baseline @xcite , enslm, has a similar architecture as ilm, ensembling n heads for predictions but always updating every head with every batch. the ensemble baseline has the same forward pass as ilm but does not perform the competitive gradient update. these baselines serve as ablations of ilm to demonstrate the importance of splitting the data into environments, ensembling the heads, and using the competitive gradient update. ## experiments invariance training comes with the promise of robustness and generalization @xcite @xcite . in the following series of experiments, we test whether our proposed architecture for language modeling can provide such benefits. since our approach is agnostic to the language model, we focus on two small lms used heavily in practice: distilbert @xcite and roberta @xcite . in this work, we do not aim to engineer the best possible lm but rather precisely test ilm in controlled setups by crafting environments whose difference is known, from which we know the expected behavior. we describe three experiments: robustness to noise, bias removal, and out-of-domain generalization. throughout the experiments, we report estimated uncertainties with 95% confidence intervals. we repeat experiments for varying hyperparameters and different random seeds (see appendix b). ## discussion in this section, we discuss our contributions in the context of previous work. ## limitations in this work, we focus on crafting controlled experiments with easily manageable dataset and language model sizes to carefully test the invariance benefits of ilm. however, it is possible to expect different qualitative behavior for large-scale language models recently deployed due to emergent properties. our implementation could be applied largely to various downstream tasks, other than language modeling measured by perplexity. here, we focus on the language modeling task and perplexity measure because they allow clear and precise experiments measuring the ability of ilm to deal with spurious correlations. the strong positive results observed in this work motivate future work to test ilm in other setups closer to direct practical usecases. it is expected that different choices of environment splits will be useful for different downstream tasks. while this work demonstrates that ilm is useful to remove spurious correlation, it does not say how to choose environments for which tasks. for instance, we observed smaller improvements when using thepile datasets and evaluating on the downstream glue tasks, indicating that thepile environment splits are not optimal for these downstream tasks. we believe that environment design is an important avenue for future research.