Invariant Language Modeling

TL;DR

This paper introduces invariant language modeling, a new framework inspired by causal learning principles, to improve out-of-domain generalization and reduce biases in large pretrained language models through invariant representations.

Contribution

It adapts a game-theoretic IRM approach to language models, enabling invariant representations that enhance robustness and mitigate spurious correlations.

Findings

Removes structured noise effectively

Ignores specific spurious correlations without harming performance

Achieves better out-of-domain generalization

Abstract

Large pretrained language models are critical components of modern 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 formulation 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 focus 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.