Rethinking Data Augmentation under Covariate Shift: Invariant-Guided Diffusion and Prototype Reweighting

Hongyu Cao; Xinyuan Wang; Arun Vignesh Malarkkan; Kunpeng Liu; Haifeng Chen; Yanjie Fu

cs.LGOct 1, 2026

Rethinking Data Augmentation under Covariate Shift: Invariant-Guided Diffusion and Prototype Reweighting

Authors: Hongyu Cao, Xinyuan Wang, Arun Vignesh Malarkkan, Kunpeng Liu, Haifeng Chen, Yanjie Fu

Organizations: School of Computing and Augmented Intelligence, Arizona State University, Tempe, AZ 85281 USA

Abstract

In many industrial applications, 1) tabular data is scarce and imbalanced and thus requires synthetic expansion; 2) input distributions drift between training and deployment (covariate shift); 3) validation sets often diverge from unseen test environments; or 4) standard generative models simply mimic outdated source distributions. This learning setting limits the stability of standard augmentation and adaptation pipelines. We generalize the task under such setting as the Augmented and Weighted Learning under Covariate Shift problem (AWL-CS). AWL-CS imposes two critical challenges on existing methods: 1) misleading generative guidance where models optimize for source similarity rather than downstream task relevance, and 2) structural instability of distributional density where reweighting mechanisms overfit to noisy validation signals. To tackle these challenges, we propose IGDPR (Invariant-Guided Diffusion with Prototype Reweighting), a unified framework that synergizes stable synthesis and structural adaptation: i) To achieve task-relevant generation, we steer the diffusion sampling process using invariant potentials to ensure synthetic samples align with stable decision boundaries rather than outdated correlations. ii) To ensure stable adaptation, we develop a prototype-based reweighting strategy that assesses sample reliability through structural clusters instead of isolated points, effectively filtering validation noise. Extensive experiments on real data demonstrate our method improves data quality by augmenting the most beneficial data for robust learning.

Figures & tables

Fig. 1 : Motivation of Augmented and Weighted Learning under Covariate Shift (AWL-CS). Under distribution shift, validation-guided augmentation and reweighting may amplify spurious patterns instead of improving generalization.

Fig. 2 : Overview of the proposed IGDPR (Invariant-Guided Diffusion with Prototype Reweighting) framework. The pipeline consists of four tightly coupled stages. (1) Invariant embedding construction: training data are partitioned into multiple pseudo-environments based on domain features, and a task-invariant latent space is learned via a multi-head TI-VAE with task, density, and invariance supervision. (2) Invariant-guided latent diffusion generation: a latent diffusion model is trained in the invariant space and, during sampling, its denoising trajectories are actively steered by frozen invariant guidance signals toward task-aligned and structurally stable regions. (3) Prototype-based reweighting: real and synthetic samples are clustered into structural prototypes, where reliability scores are estimated at the prototype level to avoid noisy and unstable point-wise reweighting. (4) Weighted downstream training: the downstream predictor is trained on the weighted augmented dataset to improve generalization under covariate shift.

Dataset	Source	Task Type	Samples	Feat Dim.
C-MAPSS	Simulated	Regression	$\sim$ 53k	14
Scania APS	Real	Classification	76k	170
Gas Sensor	Real	Classification	13k	128
MetroPT-3	Real	Regression	1.5M	$\sim$ 5

TABLE I : Summary of Datasets and Industrial Tasks.

Generator	Reweighter	Gas Sensor (Classif.)			NASA C-MAPSS (Reg.)			Scania ComponentX (Classif.)			MetroPT-3 (Reg.)
Generator	Reweighter	Acc $\uparrow$	F1-M $\uparrow$	F1-W $\uparrow$	MSE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$	Acc $\uparrow$	F1-M $\uparrow$	F1-W $\uparrow$	MSE $\downarrow$	RMSE $\downarrow$	$R^{2}$ $\uparrow$
No Aug	Uniform	0.6479	0.6066	0.6069	2069.21	45.49	0.82	0.9083	0.8710	0.9021	0.2605	0.5104	0.71
	Importance	0.6600	0.6225	0.6542	2067.63	45.47	0.83	0.9104	0.8748	0.9046	0.2589	0.5088	0.72
	Approx.	0.4414	0.3921	0.4356	2093.73	45.76	0.78	0.8897	0.8415	0.8823	0.2516	0.5016	0.73
	Prototype	0.6615	0.6284	0.6571	2064.41	45.44	0.83	0.9146	0.8793	0.9085	0.2598	0.5097	0.72
MLP GAN	Uniform	0.6410	0.6032	0.6354	2117.05	46.01	0.75	0.9121	0.8769	0.9068	0.2521	0.5021	0.73

TABLE II : Main Results on Industrial Domain Shift Benchmarks. We compare different generative augmentation methods combined with various reweighting strategies. We report Accuracy , F1-Macro (F1-M), and F1-Weighted (F1-W) for classification tasks, and MSE , RMSE , and

R^{2}

for regression tasks. The Guided Latent Diffusion generator denotes our invariant-guided latent diffusion module; paired with the Prototype reweighter (bold row) it constitutes our full method IGDPR. Bold indicates best performance.

Hyper-parameter	Scania	C-MAPSS	Gas Sensor	MetroPT-3
Batch Size	128	256	64	128
Learning Rate	1e-3	5e-4	1e-3	1e-3
Latent Dim ( $z$ )	16	16	16	16
Prototypes ( $K$ )	20	20	20	10
Reweight $\beta$	1.0	0.5	1.0	1.0
Reweight $\gamma$	1.0	0.5	0.5	1.0

TABLE III : Per-dataset hyperparameter configuration.

Fig. 3 : Ablation Study on Backbone, Reweighting Strategy, and Prototype Count. The bar charts (bottom axis) compare different generative backbones and reweighting strategies. The hatched bars denote our full method IGDPR (Guided Latent Diffusion + Prototype). The red lines (top axis) indicate the sensitivity to the number of prototypes

K

on a logarithmic scale. Metrics are Accuracy for the classification datasets (Gas Sensor, Scania), RMSE for C-MAPSS, and MSE for MetroPT-3.

Method	Gas Sensor			C-MAPSS			Scania			MetroPT-3
	Acc	F1-M	F1-W	MSE	RMSE	$R^{2}$	Acc	F1-M	F1-W	MSE	RMSE	$R^{2}$
No Aug + Uniform	0.006	0.008	0.007	42.1	0.46	0.012	0.004	0.006	0.005	0.006	0.005	0.015
TabDDPM + Prototype	0.005	0.006	0.006	38.5	0.41	0.010	0.003	0.005	0.004	0.005	0.004	0.012
Guided Latent Diffusion + Prototype (IGDPR)	0.004	0.005	0.005	35.2	0.38	0.009	0.003	0.004	0.004	0.004	0.003	0.010

TABLE IV : Standard deviation across five random seeds , corresponding to the entries in Table II . Gains in the main table consistently exceed seed-level variance.

Fig. 4 : Sensitivity Analysis of Hyperparameters $\beta$ and $\gamma$ . We report Accuracy for Gas Sensor, F1-Macro for Scania, and MSE for C-MAPSS and MetroPT-3. Brighter colors (yellow) indicate better performance. For MSE, the color scale is inverted so that lower error corresponds to brighter colors. The optimal performance is consistently observed near $\beta=1.0$ and $\gamma\in[0.5,1.0]$ .

Panel of Figure 9 — Fig. 5 : t-SNE visualization of latent representations with generation and reweighting. Colors denote class labels, marker size indicates sample importance, and cross-shaped markers represent generated samples. Generated points that align with the original data manifold receive higher weights and help fill sparse regions, while off-manifold samples are assigned low importance.

Fig. 6 : Information Plane Analysis. Trade-off between Domain Information

I(Z;E)

and Task Information

I(Z;Y)

across datasets. “Ours” (marked IGD in the legend) denotes IGDPR, which consistently moves representations toward the ideal region (top-left), achieving lower domain dependence and higher task relevance compared to baselines.

$\lambda_{\text{inv}}$	0.0	0.1	0.5	1.0	2.0	10.0
RMSE	111.48	110.27	108.82	107.60	108.99	108.85
$\Delta$ vs $\lambda{=}0$	—	$-1.21$	$-2.66$	$\mathbf{-3.88}$	$-2.49$	$-2.63$

TABLE V :

\lambda_{\text{inv}}

sweep on C-MAPSS (RMSE

\downarrow

, single seed). The 5-seed validation at the optimum

\lambda_{\text{inv}}{=}1.0

yields

108.74\pm 0.96

RMSE, indicating that the sweep range (

\Delta_{\text{RMSE}}{=}3.88

) is approximately

4\times

the seed-level noise.

$\lambda_{\text{inv}}$	0.0	0.1	0.5	1.0	2.0	10.0
RMSE	0.583	0.509	0.486	0.578	0.476	0.543
$\Delta$ vs $\lambda{=}0$	—	$-0.074$	$-0.097$	$-0.005$	$-0.107$	$-0.040$

TABLE VI :

\lambda_{\text{inv}}

sweep on MetroPT-3 (RMSE

\downarrow

, single seed). The 5-seed validation at

\lambda_{\text{inv}}{=}0.1

yields

0.510\pm 0.029

RMSE; seed noise is comparable to inter-

\lambda

gaps, so single-seed rankings within the

[0.1,2.0]

plateau should be interpreted as a low-error region rather than a strict ordering.

Dataset	Env. construction	$I(E;X)_{\max}$	$I(E;X)_{\Sigma}$	$I(E;Y)$	$I(E;Y)-I(E;\hat{Y})$
Scania	residual $K$ -means ( $K{=}3$ )	0.396	4.757	0.000	$-0.002$
C-MAPSS	explicit op. condition	0.501	0.954	0.007	$-0.000$
Gas Sensor	explicit batch id	0.591	7.790	0.132	$-0.003$
MetroPT-3	$K$ -means ( $K{=}3$ ) on PCA(5) features	0.626	1.146	0.469	$-0.161$

TABLE VII : MI probe of per-dataset environment constructions. All four clear the

I(E;X)_{\max}\geq 0.05

nats identifiability floor by at least an order of magnitude; the non-positive residual

I(E;Y)-I(E;\hat{Y})

is consistent with env–label dependence being mediated by

X

Explore similar work

SGN: A Similarity-based Generative Network for Data Generation under Distribution Shift
Jul 20, 2026Jiaqi Zhu, Xincheng Chen, Yuncheng Wu +2Data Augmentation Generative Architectures
National University of Singapore · Renmin University of China · Beijing Institute of Technology +1
Active Tabular Augmentation via Policy-Guided Diffusion Inpainting
May 11, 2026Zheyu Zhang, Shuo Yang, Bardh Prenkaj +1Augmentation Generative Architectures
Technical University of Munich · Munich Center for Machine Learning (MCML)
Preserve the Hard, Regenerate the Rest: Uncertainty-Guided Synthetic Training Data Augmentation with Diffusion Models
Jun 30, 2026Nikolai Röhrich, Julian Gleißner, Ahmed H. A. Ibrahim +2Data Augmentation Augmentation
1XITASO GmbH · 2Zuse School of Excellence in Reliable AI · 3Technische Hochschule Ingolstadt +1

Jul 20, 2026, cs.LGJ/K move · Enter open · S save

SGN: A Similarity-based Generative Network for Data Generation under Distribution Shift

Jiaqi Zhu, Xincheng Chen, Yuncheng Wu, Zhaojing Luo+1

National University of Singapore · Renmin University of China · Beijing Institute of Technology+1

Data Augmentation Generative Architectures Distribution Shifts Data Generation Latent Space Similarity and Metrics

Open paper AskPDF arXiv

Generative models trained on a source domain often produce samples that are poorly aligned with shifted target domains, limiting their effectiveness for target-domain data augmentation. Although target-specific adaptation can reduce this mismatch, it typically requires additional optimization and domain-specific parameters. We propose a Similarity-based Generative Network (SGN), a reusable framework that is trained once on labeled source data and applied to new target domains without parameter updates. SGN learns a latent space structured by label-induced pairwise similarities while preserving reconstructive information through an encoder-decoder architecture. At generation time, a small labeled representative set from the target domain is encoded and combined in the learned latent space, allowing the generated samples to inherit target-specific characteristics while maintaining class consistency. We further analyze the realizability and dimensionality requirements of the proposed similarity structure. Experiments on image and tabular datasets demonstrate the effectiveness of SGN for target-guided data augmentation under source-to-target distribution shifts.

Abstract

Explore similar work

SGN: A Similarity-based Generative Network for Data Generation under Distribution Shift

Active Tabular Augmentation via Policy-Guided Diffusion Inpainting

Preserve the Hard, Regenerate the Rest: Uncertainty-Guided Synthetic Training Data Augmentation with Diffusion Models