Journal of Data Science logo


Login Register

  1. Home
  2. To appear
  3. Target-Aware Linear Regression Under Dis ...

Journal of Data Science

Submit your article Information
  • Article info
  • Related articles
  • More
    Article info Related articles

Target-Aware Linear Regression Under Distribution Shift
Zhewen Hou   Tian Zheng  

Authors

 
Placeholder
https://doi.org/10.6339/26-JDS1245
Pub. online: 13 October 2026      Type: Philosophies Of Data Science      Open accessOpen Access

Received
31 August 2026
Accepted
4 September 2026
Published
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 Material
The 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
 
Gloege L, McKinley GA, Landschützer P, Fay AR, Frölicher TL, …, Takano Y (2021). Quantifying errors in observationally based estimates of ocean carbon sink variability. Global Biogeochemical Cycles, 35(4): e2020GB006788.
 
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
 
Kay JE, Deser C, Phillips A, Mai A, Hannay C, …, Vertenstein M (2015). The Community Earth System Model (CESM) large ensemble project: A community resource for studying climate change in the presence of internal climate variability. Bulletin of the American Meteorological Society, 96(8): 1333–1349.
 
Liu Q, Lee J, Jordan M (2016). A kernelized Stein discrepancy for goodness-of-fit tests. In: Balcan MF, Weinberger KQ (eds.), International Conference on Machine Learning 2016, 276–284. PMLR.
 
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
 
Särndal CE (2007). The calibration approach in survey theory and practice. Survey Methodology, 33(2): 99–119.
 
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
 
Stojanov P, Gong M, Carbonell J, Zhang K (2019). Low-dimensional density ratio estimation for covariate shift correction. In: Chaudhuri K, Sugiyama M (eds.), International Conference on Artificial Intelligence and Statistics 2019, 3449–3458. PMLR.
 
Tan Z (2020). Regularized calibrated estimation of propensity scores with model misspecification and high-dimensional data. Biometrika, 107(1): 137–158.MR4064145
 
Tang C, Lenssen N, Wei Y, Zheng T (2023). Wasserstein distributional learning via majorization-minimization. In: Ruiz F, Dy J, van de Meent J-W (eds.), International Conference on Artificial Intelligence and Statistics 2023, 10703–10731. PMLR.
 
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

Related articles PDF XML
Related articles PDF XML

Copyright
2026 The Author(s). Published by the School of Statistics and the Center for Applied Statistics, Renmin University of China.
by logo by logo
Open access article under the CC BY license.

Keywords
calibration distribution shift knowledge informed learning

Metrics
since February 2021
55

Article info
views

21

PDF
downloads

Export citation

Copy and paste formatted citation
Placeholder

Download citation in file


Share


RSS

Journal of data science

  • Online ISSN: 1683-8602
  • Print ISSN: 1680-743X

About

  • About journal
  • Renmin University of China homepage
  • Academic Journal Management
    and Development Center homepage

For contributors

  • Submit
  • OA Policy
  • Become a Peer-reviewer

Contact us

  • JDS@ruc.edu.cn
  • Contact person: Jing Zhou
  • Phone: +86-10-62511318
  • No. 59 Zhongguancun Street, Haidian District Beijing, 100872, P.R. China
Powered by PubliMill  •  Privacy policy