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

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 33 of its 48 references that are in the map.
- First posted on arXiv 1 Oct 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 33)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Co-teaching: Robust Training of Deep Neural Networks with Extremely Noisy Labels](https://mlsanity.com/paper/1804.06872/) (2018)
- [Modeling Tabular data using Conditional GAN](https://mlsanity.com/paper/1907.00503/) (2019)
- [MentorNet: Learning Data-Driven Curriculum for Very Deep Neural Networks on Corrupted Labels](https://mlsanity.com/paper/1712.05055/) (2017)
- [Deep Neural Networks and Tabular Data: A Survey](https://mlsanity.com/paper/2110.01889/) (2021)
- [TabDDPM: Modelling Tabular Data with Diffusion Models](https://mlsanity.com/paper/2209.15421/) (2022)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [Generating Multi-label Discrete Patient Records using Generative Adversarial Networks](https://mlsanity.com/paper/1703.06490/) (2017)
- [Synthesizing Tabular Data using Generative Adversarial Networks](https://mlsanity.com/paper/1811.11264/) (2018)
- [Model Agnostic Sample Reweighting for Out-of-Distribution Learning](https://mlsanity.com/paper/2301.09819/) (2023)
- [Prototypical Graph Contrastive Learning](https://mlsanity.com/paper/2106.09645/) (2021)
- [PTaRL: Prototype-based Tabular Representation Learning via Space Calibration](https://mlsanity.com/paper/2407.05364/) (2024)
- [Can I trust my fake data -- A comprehensive quality assessment framework for synthetic tabular data in healthcare](https://mlsanity.com/paper/2401.13716/) (2024)
- [A Survey on Deep Learning Approaches for Tabular Data Generation: Utility, Alignment, Fidelity, Privacy, Diversity, and Beyond](https://mlsanity.com/paper/2503.05954/) (2025)
- [CLIMB: Class-imbalanced Learning Benchmark on Tabular Data](https://mlsanity.com/paper/2505.17451/) (2025)
- [Downstream Task-Oriented Generative Model Selections on Synthetic Data Training for Fraud Detection Models](https://mlsanity.com/paper/2401.00974/) (2024)
- [Continuous Optimization for Feature Selection with Permutation-Invariant Embedding and Policy-Guided Search](https://mlsanity.com/paper/2505.11601/) (2025)
- [Drift in a Popular Metal Oxide Sensor Dataset Reveals Limitations for Gas Classification Benchmarks](https://mlsanity.com/paper/2108.08793/) (2021)
- [Interpretable LLMs for Credit Risk: A Systematic Review and Taxonomy](https://mlsanity.com/paper/2506.04290/) (2025)
- [Permutation-Invariant Representation Learning for Robust and Privacy-Preserving Feature Selection](https://mlsanity.com/paper/2510.05535/) (2025)
- [Sim2Act: Robust Simulation-to-Decision Learning via Adversarial Calibration and Group-Relative Perturbation](https://mlsanity.com/paper/2603.09053/) (2026)
- [General regularization in covariate shift adaptation](https://mlsanity.com/paper/2307.11503/) (2023)
- [Mitigating Shortcut Reasoning in Language Models: A Gradient-Aware Training Approach](https://mlsanity.com/paper/2603.20899/) (2026)
- [Hierarchical and Permutation-Invariant Feature Transformation Learning via Policy-Guided Embedding Search](https://mlsanity.com/paper/2609.10225/) (2026)
- [Sensitive Data Detection with High-Throughput Machine Learning Models in Electrical Health Records](https://mlsanity.com/paper/2305.03169/) (2023)

## Built on by (most cited first, 1 of 1)

- [Which and When to Admit: Gradient Admission for Data-Centric Small Language Model Finetuning](https://mlsanity.com/paper/2610.07553/) (2026)

## Abstract (arXiv)

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.