Distributionally Robust Language Modeling

TL;DR

The paper tackles the problem that language models trained on heterogeneous data can underperform on unseen target distributions. It introduces topic CVaR, a distributionally robust optimization approach that uses topic-level uncertainty sets and a baselined loss to ensure uniform performance across topics, implemented via an online minimax algorithm. Empirical results show substantial perplexity reductions (e.g., 5.5 points on Yelp when tested on reviews) and robustness beyond subpopulation shifts, with strong ablations confirming the importance of topic structure and baselining. This approach enables a single, broadly robust language model that generalizes across unknown distributions without needing target-domain labels.

Abstract

Language models are generally trained on data spanning a wide range of topics (e.g., news, reviews, fiction), but they might be applied to an a priori unknown target distribution (e.g., restaurant reviews). In this paper, we first show that training on text outside the test distribution can degrade test performance when using standard maximum likelihood (MLE) training. To remedy this without the knowledge of the test distribution, we propose an approach which trains a model that performs well over a wide range of potential test distributions. In particular, we derive a new distributionally robust optimization (DRO) procedure which minimizes the loss of the model over the worst-case mixture of topics with sufficient overlap with the training distribution. Our approach, called topic conditional value at risk (topic CVaR), obtains a 5.5 point perplexity reduction over MLE when the language models are trained on a mixture of Yelp reviews and news and tested only on reviews.

Figures (7)

• Figure 1: Illustration of a training corpus as a density (black) with mostly news stories (red) and a small number of restaurant reviews (blue). The standard MLE model (gray) reflects the underlying data and assigns little weight to reviews, and thus performs poorly on reviews. A more robust model should try to equalize the weight across all topics so that it can perform well regardless of which topics appear at test time.
• Figure 2: Toy example of a multinomial distribution over six sentences (top). Different panels illustrate models learned by different training procedures. MLE fits common topics (news) at the expense of rare ones (reviews). Sentence CVaR is too conservative, overemphasizing the ungrammatical sentence. Topic CVaR with log loss overemphasizes difficult topics (news) over easy ones (review). Topic CVaR (with baselining) balances the weight assigned to each topics, as desired.
• Figure 3: Topic CVaR (green) provides substantial improvements in perplexity compared to MLE (black and blue) as the amount of train-test mismatch increases ($\alpha^*\to 0$). This performance is close to the oracle performance which uses ground truth corpus labels and early stopping (orange).
• Figure 4: The robustness improvements from topic CVaR (black vs green and orange) apply even when the test set (TripAdv reviews) is not a subpopulation shift from the training set (Yelp and OneBWord).
• Figure 5: The robustness of topic CVaR degrades when the baseline is removed (teal), but is resistant to being over-conservative in choosing $\alpha$ (yellow).