# Mitigating distribution shift in machine learning-augmented hybrid simulation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 70 references that are in the map.
- First posted on arXiv 17 Jan 2024.

## Builds on (most cited first, 26 of 26)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Lagrangian Neural Networks](https://mlsanity.com/paper/2003.04630/) (2020)
- [A Brief Review of Domain Adaptation](https://mlsanity.com/paper/2010.03978/) (2020)
- [Meta-Learning: A Survey](https://mlsanity.com/paper/1810.03548/) (2018)
- [Deep Potential Molecular Dynamics: a scalable model with the accuracy of quantum mechanics](https://mlsanity.com/paper/1707.09571/) (2017)
- [Failure-informed adaptive sampling for PINNs](https://mlsanity.com/paper/2210.00279/) (2022)
- [Physics-Guided Deep Learning for Dynamical Systems: A Survey](https://mlsanity.com/paper/2107.01272/) (2021)
- [A graph convolutional autoencoder approach to model order reduction for parametrized PDEs](https://mlsanity.com/paper/2305.08573/) (2023)
- [Active Learning of Uniformly Accurate Inter-atomic Potentials for Materials Simulation](https://mlsanity.com/paper/1810.11890/) (2018)
- [OnsagerNet: Learning Stable and Interpretable Dynamics using a Generalized Onsager Principle](https://mlsanity.com/paper/2009.02327/) (2020)
- [Deep Neural Networks for Nonlinear Model Order Reduction of Unsteady Flows](https://mlsanity.com/paper/2007.00936/) (2020)
- [Confidence Regions in Wasserstein Distributionally Robust Estimation](https://mlsanity.com/paper/1906.01614/) (2019)
- [Towards Robust Neural Networks via Close-loop Control](https://mlsanity.com/paper/2102.01862/) (2021)
- [Learning optimal multigrid smoothers via neural networks](https://mlsanity.com/paper/2102.12071/) (2021)
- [A Data Driven Method for Computing Quasipotentials](https://mlsanity.com/paper/2012.09111/) (2020)
- [Coupling parameter and particle dynamics for adaptive sampling in Neural Galerkin schemes](https://mlsanity.com/paper/2306.15630/) (2023)
- [Physics-informed Deep Super-resolution for Spatiotemporal Data](https://mlsanity.com/paper/2208.01462/) (2022)
- [Principled Acceleration of Iterative Numerical Methods Using Machine Learning](https://mlsanity.com/paper/2206.08594/) (2022)
- [Self-Healing Robust Neural Networks via Closed-Loop Control](https://mlsanity.com/paper/2206.12963/) (2022)
- [Why should autoencoders work?](https://mlsanity.com/paper/2310.02250/) (2023)
- [DAGGER: A sequential algorithm for FDR control on DAGs](https://mlsanity.com/paper/1709.10250/) (2017)

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

- [Towards Secure MLOps: Surveying Attacks, Mitigation Strategies, and Research Challenges](https://mlsanity.com/paper/2506.02032/) (2025)
- [Learning task-specific predictive models for scientific computing](https://mlsanity.com/paper/2506.03835/) (2025)

## Abstract (arXiv)

We study the problem of distribution shift generally arising in machine-learning augmented hybrid simulation, where parts of simulation algorithms are replaced by data-driven surrogates. We first establish a mathematical framework to understand the structure of machine-learning augmented hybrid simulation problems, and the cause and effect of the associated distribution shift. We show correlations between distribution shift and simulation error both numerically and theoretically. Then, we propose a simple methodology based on tangent-space regularized estimator to control the distribution shift, thereby improving the long-term accuracy of the simulation results. In the linear dynamics case, we provide a thorough theoretical analysis to quantify the effectiveness of the proposed method. Moreover, we conduct several numerical experiments, including simulating a partially known reaction-diffusion equation and solving Navier-Stokes equations using the projection method with a data-driven pressure solver. In all cases, we observe marked improvements in simulation accuracy under the proposed method, especially for systems with high degrees of distribution shift, such as those with relatively strong non-linear reaction mechanisms, or flows at large Reynolds numbers.