Source-linked AI summary
Distributionally Robust Language Modeling
Yonatan Oren, Shiori Sagawa, Tatsunori B. Hashimoto, Percy Liang
TL;DR
Broad language-model training can degrade performance on an unknown target distribution because MLE favors common topics. The paper introduces topic CVaR, a topic-based DRO procedure that optimizes worst-case mixtures with sufficient training-distribution overlap. It reduces Yelp-review perplexity by 5.5 points versus MLE in a Yelp–news mixture setting.
Problem
Training on heterogeneous corpora can degrade performance on target distributions because MLE emphasizes common sentences and topics over infrequent ones.
Method
Topic CVaR minimizes risk over worst-case meaningful topic subpopulations, using a baselined loss to account for topic entropy.
Results
5.5-point perplexity reduction on Yelp reviews versus MLE is obtained when training uses a Yelp–news mixture.
Takeaways & Limitations
Topic CVaR improves robustness to subpopulation shifts and similar shifts, with topics encoding prior information about potential test distributions.
Takeaways & Limitations
The approach targets subpopulation shifts, assuming sufficient distributional overlap between training and potential test distributions rather than shared examples.
Abstract
from arXiv · showhide
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.
1 Introduction
Heterogeneous training data can hurt performance on unknown target distributions because MLE emphasizes common topics. The paper proposes topic CVaR, a DRO approach designed to improve robustness across possible topic shifts.
- Language models trained on broad corpora are applied to tasks with specific, potentially unknown test distributions.
- Perplexity rises from 32 to 43 when training changes from Yelp reviews alone to a mixture of 10% Yelp and 90% newswire sentences.The passage attributes this degradation to MLE emphasizing common sentences and topics.
- Fine-tuning and domain adaptation require knowing the test distribution and training a separate model for each target.
- The paper instead seeks a single model that performs well across many unknown test distributions, particularly subpopulation shifts.The subpopulation assumption concerns distributional overlap rather than shared individual examples.
- Topic CVaR minimizes worst-case risk over meaningful topic-defined subpopulations and corrects topic difficulty using a baselined loss.The method addresses arbitrary-subpopulation conservatism and high-entropy bias in naive CVaR.
- 5.5-point perplexity reduction is achieved on Yelp reviews versus MLE when training uses a Yelp–One Billion Word Benchmark mixture.
2 Problem Statement
Language models trained with MLE can perform poorly on unknown test distributions when those distributions emphasize topics that are rare in training. DRO addresses this mismatch by optimizing for the worst-case distribution in an uncertainty set, whose objective does not require knowing the test distribution.
- The target distribution may be unknown, so the goal is good performance across a large set of potential test distributions.The setting assumes the test distribution is a subpopulation of the training distribution.
- MLE can perform poorly under train-test mismatch when the test set emphasizes topics that are infrequent during training.Such topics receive low probabilities, producing high losses even with abundant data.
- DRO minimizes model risk under the worst-case distribution in an uncertainty set of potential test distributions.The framework optimizes a specified loss over the distribution considered most adverse within that set.
- The worst-case objective does not depend on the unknown test distribution and upper-bounds test risk whenever that distribution belongs to the uncertainty set.
- The choice of uncertainty set and loss function is critical to the guarantees and behavior of DRO.The paper notes that these choices are discussed as central design decisions.
3 Robust Language Modeling
Standard DRO can be too conservative for language modeling because arbitrary sentence groups and absolute log loss do not match meaningful topic shifts or topic difficulty. Topic CVaR addresses these issues by defining robustness over latent-topic mixtures and using a baselined loss.
- Motivation: Existing DRO uncertainty sets are too conservative, and log loss overemphasizes topics with inherently high entropy.These two drawbacks motivate topic CVaR.
- Robustness to arbitrary subpopulations: Sentence CVaR protects against arbitrary sentence groups, including groups composed solely of invalid sentences.This can encourage modeling ungrammatical sentences at the expense of valid ones.
- Robustness over Topics: Topic CVaR defines the uncertainty set over latent topics, targeting meaningful shifts in topic proportions while keeping each topic’s conditional sentence distribution fixed.The objective upweights high-loss topics and downweights low-loss topics.
- Baselined Loss Function: The baselined loss measures model performance relative to the best possible model for each topic, encouraging uniform relative performance.It uses log p_x|z(x | z) − log p_θ(x) as the topic-relative loss.
- Baselined Loss Function: Topic CVaR minimizes per-topic KL divergences, aligning robust training with matching the test distribution rather than minimizing ordinary log loss.This interpretation distinguishes topic CVaR from MLE.
- Results: In the toy example, topic CVaR produces a model performing relatively well on both news and reviews, with equal topic weights.The resulting model is a mixture of the news and review distributions.
4 Algorithm
The algorithm operationalizes topic CVaR by selecting latent topics, estimating topic entropies with baseline models, and optimizing a two-player minimax objective online. Historical topic losses guide the adversarial topic distribution while stochastic gradient updates train the model.
- Topic Selection: The algorithm chooses latent topics with LDA, assigning each sentence to the topic with highest total posterior probability.The topics are intended to capture realistic potential test-distribution shifts.
- Baseline Estimation: Topic CVaR uses KL divergence as its loss, but topic conditional entropy is difficult to estimate directly from samples.The entropy term is H(X | Z = z) := E_p(x|z)[−log p_x|z(x | z)].
- Baseline Estimation: A separate baseline model for each topic estimates conditional entropy; in practice, a scalable bigram model is used.The estimated entropy is H_β(X | Z = z) := E_p(x|z)[−log p_β(x | z)].
- Online Optimization: No scalable online algorithm previously handled topic CVaR’s nested maximization-expectation structure, motivating a stochastic-gradient-compatible procedure.Batch dual methods require exact expectations, while existing online DRO methods do not handle this structure.
- Online Optimization: Topic CVaR is optimized as a two-player minimax game between model parameters θ and an adversarial topic distribution p_z.The topic player maximizes the robust objective while θ minimizes it.
- Online Optimization: At each round, the algorithm updates p_z using historical per-topic losses and updates θ through interleaved minimax optimization.It tracks empirical topic probabilities and historical average losses for the observed topics.
5 Experiments
The experiments evaluate topic CVaR against MLE across synthetic and real train-test shifts, showing improved robustness to nuisance data, imperfect topic overlap, and hyperparameter misspecification.
- Experimental setup: Models use YELP, ONEBWORD, and TRIPADV corpora with a shared fixed vocabulary and topic clusters for evaluation.Transformer language models compare MLE and topic CVaR under matched preprocessing and training settings.
- Subpopulation shift: Topic CVaR substantially improves robustness as YELP’s fraction in YELP–ONEBWORD training mixtures decreases.MLE incurs up to 10 additional perplexity points below α∗=0.4, while topic CVaR retains near-baseline performance and approaches oracle performance.
- Beyond subpopulation shift: Topic CVaR also improves performance on TRIPADV reviews, whose distribution is similar to but not strictly a subpopulation of YELP and newswire.The MLE degradation between α∗=0.7 and 0.1 is substantially reduced by topic CVaR.
- Ablations: Removing topic structure causes perplexity to exceed 80 at α=0.2, while removing the baseline reduces robustness as α∗ decreases.Without topics, worst-case groups can contain disfluent sentences; without the baseline, optimization focuses heavily on higher-entropy ONEBWORD data.
- Hyperparameter robustness: Selecting α imperfectly degrades perplexity by 2–3 points, but α=0.2 still outperforms MLE for small α∗ and incurs only 2 points of loss at α∗=1.0.At α∗=0.1, every tested α outperforms MLE, with an approximately linear penalty for incorrect choices.
- Analysis: Topic CVaR assigns higher probabilities to Yelp and infrequent sentences, whereas MLE performs better on frequent ONEBWORD sentences and some location- or accident-related Yelp examples.The sentence-level and example-level analyses align with topic CVaR’s lower losses on typical Yelp content.
6 Related Work
Related work covers domain adaptation, latent-domain discovery, topic-based adaptation, and distributionally robust optimization. The paper differs by targeting unknown language-model test shifts without test-time inputs or domain annotations.
- Domain adaptation: Domain adaptation methods learn robust models or domain-invariant features when source and target domains are known, but require accurate domain membership annotations.Prior multi-source approaches instead use clustering to identify candidate domains when annotations are absent.
- Topic and domain adaptation: Topic-based domain-adaptation methods for conditional tasks use topics at test time by inferring domains from input variables.Language modeling lacks such inputs, motivating robustness to unknown domain shifts at test time.
- Distributional robustness: Prior DRO work optimizes over distributional uncertainty sets, including subpopulations and class-structured groups.This paper derives an efficient procedure for topic-based uncertainty sets and examines why naive DRO application is insufficient.
7 Discussion
The paper finds that heterogeneous training data can hurt language-model performance under train-test mismatch and proposes topic CVaR as a DRO-based remedy. Its experiments support robustness across subpopulation and similar shifts.
- Discussion: Language-model performance degrades as text outside the test distribution increases, consistent with MLE emphasizing common corpus sentences.The paper proposes distributionally robust optimization to address this mismatch.
- Discussion: DRO-based topic CVaR is more robust than MLE to subpopulation shifts and similar shifts.The paper argues that topics encode prior information about potential test distributions and baselines normalize difficulty across topics.