# Target-Aware Linear Regression Under Distribution Shift

arXiv: https://arxiv.org/abs/2606.22775 · HTML: https://mlsanity.com/paper/2606.22775/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 29 references that are in the map.
- First posted on arXiv 22 Jun 2026.

## Builds on (most cited first, 10 of 10)

- [A Kernelized Stein Discrepancy for Goodness-of-fit Tests and Model Evaluation](https://mlsanity.com/paper/1602.03253/) (2016)
- [Augmented Minimax Linear Estimation](https://mlsanity.com/paper/1712.00038/) (2017)
- [User-defined Event Sampling and Uncertainty Quantification in Diffusion Models for Physical Dynamical Systems](https://mlsanity.com/paper/2306.07526/) (2023)
- [Efficient and Adaptive Linear Regression in Semi-Supervised Settings](https://mlsanity.com/paper/1701.04889/) (2017)
- [PI-VAE: Physics-Informed Variational Auto-Encoder for stochastic differential equations](https://mlsanity.com/paper/2203.11363/) (2022)
- [Augmented balancing weights as linear regression](https://mlsanity.com/paper/2304.14545/) (2023)
- [Regularized calibrated estimation of propensity scores with model misspecification and high-dimensional data](https://mlsanity.com/paper/1710.08074/) (2017)
- [Adaptive learning of density ratios in RKHS](https://mlsanity.com/paper/2307.16164/) (2023)
- [Nearest Neighbor Sampling for Covariate Shift Adaptation](https://mlsanity.com/paper/2312.09969/) (2023)
- [Calibrating Geophysical Predictions under Constrained Probabilistic Distributions](https://mlsanity.com/paper/2512.03081/) (2025)

## Abstract (arXiv)

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[Ymid 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.