Target-Aware Linear Regression Under Distribution Shift

• URL: http://arxiv.org/abs/2606.22775v2
• Date: Tue, 23 Jun 2026 16:56:48 GMT
• Title: Target-Aware Linear Regression Under Distribution Shift
• Abstract summary: Distribution shift between training and deployment is a pervasive challenge for modern AI systems.<n>We study this problem with a stable conditional mean $E[Ymid X]$ across source and target.<n>We identify the hybrid-loss estimator, which jointly incorporates both target marginals, as a benchmark target-aware estimator.
• Score: 2.705806052277982
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• On the Optimal Construction of Unbiased Gradient Estimators for Zeroth-Order Optimization [57.179679246370114]
  A potential limitation of existing methods is the bias inherent in most perturbation estimators unless a stepsize is proposed.<n>We propose a novel family of unbiased gradient scaling estimators that eliminate bias while maintaining favorable construction.
  arXiv  Detail & Related papers  (2025-10-22T18:25:43Z)
• Doubly Calibrated Estimator for Recommendation on Data Missing Not At Random [20.889464448762176]
  We argue that existing estimators rely on miscalibrated imputed errors and propensity scores. We propose a Doubly Calibrated Estimator that involves the calibration of both the imputation and propensity models.
  arXiv  Detail & Related papers  (2024-02-26T05:08:52Z)
• Nearest Neighbor Sampling for Covariate Shift Adaptation [7.940293148084844]
  We propose a new covariate shift adaptation method which avoids estimating the weights. The basic idea is to directly work on unlabeled target data, labeled according to the $k$-nearest neighbors in the source dataset. Our experiments show that it achieves drastic reduction in the running time with remarkable accuracy.
  arXiv  Detail & Related papers  (2023-12-15T17:28:09Z)
• Leveraging Variational Autoencoders for Parameterized MMSE Estimation [10.141454378473972]
  We propose a variational autoencoder-based framework for parameterizing a conditional linear minimum mean squared error estimator. The derived estimator is shown to approximate the minimum mean squared error estimator by utilizing the variational autoencoder as a generative prior for the estimation problem. We conduct a rigorous analysis by bounding the difference between the proposed and the minimum mean squared error estimator.
  arXiv  Detail & Related papers  (2023-07-11T15:41:34Z)
• Multi-Fidelity Covariance Estimation in the Log-Euclidean Geometry [0.0]
  We introduce a multi-fidelity estimator of covariance matrices that employs the log-Euclidean geometry of the symmetric positive-definite manifold. We develop an optimal sample allocation scheme that minimizes the mean-squared error of the estimator given a fixed budget. Evaluations of our approach using data from physical applications demonstrate more accurate metric learning and speedups of more than one order of magnitude compared to benchmarks.
  arXiv  Detail & Related papers  (2023-01-31T16:33:46Z)
• Pessimistic Minimax Value Iteration: Provably Efficient Equilibrium Learning from Offline Datasets [101.5329678997916]
  We study episodic two-player zero-sum Markov games (MGs) in the offline setting. The goal is to find an approximate Nash equilibrium (NE) policy pair based on a dataset collected a priori.
  arXiv  Detail & Related papers  (2022-02-15T15:39:30Z)
• Learning to Estimate Without Bias [57.82628598276623]
  Gauss theorem states that the weighted least squares estimator is a linear minimum variance unbiased estimation (MVUE) in linear models. In this paper, we take a first step towards extending this result to non linear settings via deep learning with bias constraints. A second motivation to BCE is in applications where multiple estimates of the same unknown are averaged for improved performance.
  arXiv  Detail & Related papers  (2021-10-24T10:23:51Z)
• Near-optimal inference in adaptive linear regression [60.08422051718195]
  Even simple methods like least squares can exhibit non-normal behavior when data is collected in an adaptive manner. We propose a family of online debiasing estimators to correct these distributional anomalies in at least squares estimation. We demonstrate the usefulness of our theory via applications to multi-armed bandit, autoregressive time series estimation, and active learning with exploration.
  arXiv  Detail & Related papers  (2021-07-05T21:05:11Z)
• Distribution-free binary classification: prediction sets, confidence intervals and calibration [106.50279469344937]
  We study three notions of uncertainty quantification -- calibration, confidence intervals and prediction sets -- for binary classification in the distribution-free setting. We derive confidence intervals for binned probabilities for both fixed-width and uniform-mass binning. As a consequence of our 'tripod' theorems, these confidence intervals for binned probabilities lead to distribution-free calibration.
  arXiv  Detail & Related papers  (2020-06-18T14:17:29Z)
• GenDICE: Generalized Offline Estimation of Stationary Values [108.17309783125398]
  We show that effective estimation can still be achieved in important applications. Our approach is based on estimating a ratio that corrects for the discrepancy between the stationary and empirical distributions. The resulting algorithm, GenDICE, is straightforward and effective.
  arXiv  Detail & Related papers  (2020-02-21T00:27:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.