Target-Aware Linear Regression Under Distribution Shift

Unknown Date

AI Summary

This study addresses the performance degradation of linear regression caused by distributional shifts in both covariates and responses between training and deployment domains. Under the assumption of conditional mean stability, the authors propose a target-aware regression method that leverages marginal distribution information of covariates and responses from the target domain. Two novel and computationally efficient estimators are introduced: a constrained moment-matching estimator and a two-stage calibration estimator. Theoretical analysis establishes their asymptotic equivalence to an idealized hybrid loss estimator under specific conditions. Through asymptotic mean squared error analysis, nonlinear optimization approximations, and Monte Carlo experiments, the study demonstrates that the two-stage estimator nearly attains the performance of the hybrid benchmark under high signal-to-noise ratios while incurring minimal computational cost, offering a practical and reliable solution for real-world deployment.

Abstract

Distribution shift between training and deployment is a pervasive challenge for modern AI systems. In many cases, the target marginals of covariates and response are known or specified through population-level observations, boundary conditions, properties of simulator configurations, or alignment-time distributional constraints. Such knowledge may provide valuable side information for regression estimation. We study this problem in the multivariate linear regression setting with a stable conditional mean $E[Y\mid X]$ across source and target, and identify the hybrid-loss estimator, which jointly incorporates both target marginals, as a benchmark target-aware estimator. Its direct computation, however, requires solving a coupled nonlinear optimization that is expensive at scale. Our main contribution is to develop and evaluate two computationally tractable alternatives: a constrained moment-matching estimator and a two-stage estimator that augments ordinary least squares with a calibration step. For all three estimators, we derive and compare closed-form asymptotic mean squared errors, yielding conditions under which the tractable alternatives match or closely approximate the hybrid benchmark, and regimes in which they do not. Monte Carlo experiments across three controlled shift regimes validate the theoretical results, investigate the accuracy-runtime tradeoffs among the three estimators, and translate into guidance on estimator choice. In particular, the two-stage estimator nearly matches the hybrid benchmark in the high signal-to-noise regime at essentially no additional cost, providing theoretical grounding for empirical observations in nonlinear settings.

Problem

• distribution shift
• target-aware regression
• linear regression
• marginal distributions
• covariate shift

Innovation

• distribution shift
• target-aware regression
• hybrid-loss estimator
• moment matching
• two-stage estimation