Target-Aware Linear Regression Under Distribution Shift
Pub. online: 13 October 2026
Type: Philosophies Of Data Science
Open Access
Received
31 August 2026
31 August 2026
Accepted
4 September 2026
4 September 2026
Published
13 October 2026
13 October 2026
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 $\mathbb{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.
Supplementary material
Supplementary MaterialThe supplementary material for this article is available on the official website of the Journal of Data Science at https://jds-online.org/. It may include additional results, appendices, tables, figures, code, and data.
References
Ben-David S, Blitzer J, Crammer K, Kulesza A, Pereira F, Vaughan JW (2010). A theory of learning from different domains. Machine Learning, 79(1–2): 151–175.MR3108150
Bruns-Smith D, Dukes O, Feller A, Ogburn EL (2026). Augmented balancing weights as linear regression. Journal of the Royal Statistical Society Series B: Statistical Methodology, 88(3), 699–723. https://doi.org/10.1093/jrsssb/qkaf019.MR5102315
Chakrabortty A, Cai T (2018). Efficient and adaptive linear regression in semi-supervised settings. The Annals of Statistics, 46(4): 1541–1572. https://doi.org/10.1214/17-AOS1594.MR3819109
Chen X, Hong H, Tarozzi A (2008). Semiparametric efficiency in GMM models with auxiliary data. The Annals of Statistics, 36(2): 808–843.MR2396816
Devaud D, Tillé Y (2019). Deville and Särndal’s calibration: revisiting a 25-years-old successful optimization problem. Test, 28(4): 1033–1065.MR4036996
Deville JC, Särndal CE (1992). Calibration estimators in survey sampling. Journal of the American Statistical Association, 87(418): 376–382.MR1173804
Efron B, Morris C (1973). Stein’s estimation rule and its competitors—an empirical Bayes approach. Journal of the American Statistical Association, 68(341): 117–130.MR0388597
Finzi MA, Boral A, Wilson AG, Sha F, Zepeda-Núñez L (2023). In: Krause A, Brunskill E, Cho K, Engelhardt B, Sabato S, Scarlett J (eds.), User-defined event sampling and uncertainty quantification in diffusion models for physical dynamical systems. In: International Conference on Machine Learning 2023, 10136–10152. PMLR.
Gao F, Chan K (2023). Noniterative adjustment to regression estimators with population-based auxiliary information for semiparametric models. Biometrics, 79(1): 140–150.MR4572511
Han P, Lawless JF (2019). Empirical likelihood estimation using auxiliary summary information with different covariate distributions. Statistica Sinica, 29(3): 1321–1342.MR3932520
Han P, Li H, Park SK, Mukherjee B, Taylor JM (2024). Improving prediction of linear regression models by integrating external information from heterogeneous populations: James–Stein estimators. Biometrics, 80(3): ujae072.MR4896878
Hansen LP (1982). Large sample properties of generalized method of moments estimators. Econometrica, 50(4): 1029–1054.MR0666123
Hirshberg DA, Wager S (2021). Augmented minimax linear estimation. The Annals of Statistics, 49(6): 3206–3227.MR4352528
Hou Z, Sun J, Venkatasubramanian S, Jin P, Li S, Zheng T (2025). Calibrating geophysical predictions under constrained probabilistic distributions. arXiv preprint: arXiv:2512.03081.
Imbens GW (1997). One-step estimators for over-identified generalized method of moments models. The Review of Economic Studies, 64(3): 359–383.MR1456135
Judge GG, Griffiths WE, Hill RC, Lee TC (1980). The Theory and Practice of Econometrics. Wiley.MR0594435
Portier F, Truquet L, Yamane I (2024). Nearest neighbor sampling for covariate shift adaptation. Journal of Machine Learning Research, 25(410): 1–42.MR4869955
Qin J, Lawless J (1994). Empirical likelihood and general estimating equations. The Annals of Statistics, 22(1): 300–325.MR1272085
Shimodaira H (2000). Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of Statistical Planning and Inference, 90(2): 227–244.MR1795598
Stein C (1956). Inadmissibility of the usual estimator for the mean of a multivariate normal distribution. In: Neyman J (ed.), Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, Volume 1: Contributions to the Theory of Statistics, 197–206. University of California Press.MR0084922
Tan Z (2020). Regularized calibrated estimation of propensity scores with model misspecification and high-dimensional data. Biometrika, 107(1): 137–158.MR4064145
Zellinger W, Kindermann S, Pereverzyev SV (2023). Adaptive learning of density ratios in RKHS. Journal of Machine Learning Research, 24(395): 1–28.MR4720851
Zhong W, Meidani H (2023). PI-VAE: Physics-informed variational auto-encoder for stochastic differential equations. Computer Methods in Applied Mechanics and Engineering, 403: 115664.MR4500668