# Learning where to learn: Training data distribution optimization for scientific machine learning

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 101 references that are in the map.
- First posted on arXiv 27 May 2025.

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

- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Meta-Learning in Neural Networks: A Survey](https://mlsanity.com/paper/2004.05439/) (2020)
- [Generalizing to Unseen Domains: A Survey on Domain Generalization](https://mlsanity.com/paper/2103.03097/) (2021)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks](https://mlsanity.com/paper/2207.10289/) (2022)
- [Sorting out Lipschitz function approximation](https://mlsanity.com/paper/1811.05381/) (2018)
- [Orthogonal Convolutional Neural Networks](https://mlsanity.com/paper/1911.12207/) (2019)
- [Variational Bayesian Optimal Experimental Design](https://mlsanity.com/paper/1903.05480/) (2019)
- [Does Invariant Risk Minimization Capture Invariance?](https://mlsanity.com/paper/2101.01134/) (2021)
- [Less is More: Nystr"om Computational Regularization](https://mlsanity.com/paper/1507.04717/) (2015)
- [Improved architectures and training algorithms for deep operator networks](https://mlsanity.com/paper/2110.01654/) (2021)
- [Operator Learning: Algorithms and Analysis](https://mlsanity.com/paper/2402.15715/) (2024)
- [Distributionally Robust Optimization: A Review](https://mlsanity.com/paper/1908.05659/) (2019)
- [A Mathematical Guide to Operator Learning](https://mlsanity.com/paper/2312.14688/) (2023)
- [DAS-PINNs: A deep adaptive sampling method for solving high-dimensional partial differential equations](https://mlsanity.com/paper/2112.14038/) (2021)
- [Probable Domain Generalization via Quantile Risk Minimization](https://mlsanity.com/paper/2207.09944/) (2022)
- [Neural Inverse Operators for Solving PDE Inverse Problems](https://mlsanity.com/paper/2301.11167/) (2023)
- [Failure-informed adaptive sampling for PINNs](https://mlsanity.com/paper/2210.00279/) (2022)
- [Convergence Rates for Learning Linear Operators from Noisy Data](https://mlsanity.com/paper/2108.12515/) (2021)
- [Optimal experimental design: Formulations and computations](https://mlsanity.com/paper/2407.16212/) (2024)
- [On Fast Leverage Score Sampling and Optimal Learning](https://mlsanity.com/paper/1810.13258/) (2018)
- [Discovering and forecasting extreme events via active learning in neural operators](https://mlsanity.com/paper/2204.02488/) (2022)
- [Randomly pivoted Cholesky: Practical approximation of a kernel matrix with few entry evaluations](https://mlsanity.com/paper/2207.06503/) (2022)
- [Out-of-distributional risk bounds for neural operators with applications to the Helmholtz equation](https://mlsanity.com/paper/2301.11509/) (2023)
- [Iterative Surrogate Model Optimization (ISMO): An active learning algorithm for PDE constrained optimization with deep neural networks](https://mlsanity.com/paper/2008.05730/) (2020)
- [Elliptic PDE learning is provably data-efficient](https://mlsanity.com/paper/2302.12888/) (2023)

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

- [Operator learning for the 2D incompressible Navier-Stokes equations: a conformal prediction approach in the data-scarce regime](https://mlsanity.com/paper/2606.08654/) (2026)
- [Into the danger zone: stable extrapolation in high-dimensional function and operator learning](https://mlsanity.com/paper/2609.36709/) (2026)
- [Stability of Measure-to-Measure Transformers on Sub-Gaussian Data](https://mlsanity.com/paper/2610.07717/) (2026)

## Abstract (arXiv)

In scientific machine learning, models are routinely deployed with parameter values or boundary conditions far from those used in training. This paper studies the learning-where-to-learn problem of designing a training data distribution that minimizes average prediction error across a family of deployment regimes. A theoretical analysis shows how the training distribution shapes deployment accuracy. This motivates two adaptive algorithms based on bilevel or alternating optimization in the space of probability measures. Discretized implementations using parametric distribution classes or nonparametric particle-based gradient flows deliver optimized training distributions that outperform nonadaptive designs. Once trained, the resulting models exhibit improved sample complexity and robustness to distribution shift. This framework unlocks the potential of principled data acquisition for learning functions and solution operators of partial differential equations.