Omitted Variable Bias in Language Models Under Distribution Shift

2026-02-18

AI Summary

Language models are susceptible to omitted variable bias induced by unobserved confounders under distribution shift, leading to unreliable evaluation and optimization. This work presents the first systematic analysis of this issue by decomposing distribution shift into observable and unobservable components, highlighting that conventional approaches only address the former. Integrating concepts from causal inference on omitted variable bias with distributionally robust optimization, we develop a formal framework that maps the strength of unobserved variables to a computable bound on worst-case generalization performance. This bound not only quantifies the impact of such bias but also significantly outperforms standard distribution shift correction methods in both evaluation and optimization, yielding improved out-of-distribution performance. Moreover, when target labels are available, the framework enables backward inference of the unobserved variable strength.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Despite their impressive performance on a wide variety of tasks, modern language models remain susceptible to distribution shifts, exhibiting brittle behavior when evaluated on data that differs in distribution from their training data. In this paper, we describe how distribution shifts in language models can be separated into observable and unobservable components, and we discuss how established approaches for dealing with distribution shift address only the former. Importantly, we identify that the resulting omitted variable bias from unobserved variables can compromise both evaluation and optimization in language models. To address this challenge, we introduce a framework that maps the strength of the omitted variables to bounds on the worst-case generalization performance of language models under distribution shift. In empirical experiments, we show that using these bounds directly in language model evaluation and optimization provides more principled measures of out-of-distribution performance, improves true out-of-distribution performance relative to standard distribution shift adjustment methods, and further enables inference about the strength of the omitted variables when target distribution labels are available.

Problem

• Omitted Variable Bias
• Distribution Shift
• Language Models
• Out-of-Distribution Generalization
• Unobserved Variables

Innovation

• omitted variable bias
• distribution shift
• language models
• generalization bounds
• out-of-distribution evaluation