Target-Aware Linear Regression Under Distribution Shift

When training data and deployment data differ, practitioners often have access to "side information" about the target population, such as known marginal distributions of covariates or responses. This paper investigates how to formally incorporate this target-marginal information into multivariate linear regression to improve predictive performance beyond standard ordinary least squares (OLS).

The authors analyze three target-aware estimators in a linear regression setting where the conditional mean is stable across distributions:

1. Hybrid-loss estimator: A benchmark method that jointly penalizes the empirical risk and the Wasserstein distance between the model-implied and target marginals. While theoretically optimal, it requires solving a computationally intensive nonlinear optimization problem.
2. Constrained moment-matching estimator: A method that enforces target mean and variance constraints exactly, formulated as a constrained optimization problem.
3. Two-stage estimator: A computationally efficient approach that performs standard OLS followed by a post-hoc calibration step to rescale predictions to match target moments.

The authors derive closed-form asymptotic mean squared errors (MSE) for all three estimators, providing a rigorous framework to evaluate accuracy-complexity tradeoffs.

The study reveals that the benefit of using target marginals is highly dependent on the signal-to-noise ratio. In high signal-to-noise regimes—where base models are already accurate—the simple two-stage estimator performs nearly as well as the complex hybrid benchmark at a fraction of the computational cost. Conversely, when the noise level is high, the target-aware estimators may offer little improvement or even perform worse than OLS. The theoretical results are validated through extensive Monte Carlo experiments, which demonstrate that the two-stage and moment-matching methods provide a practical, low-cost alternative to iterative optimization.