Learning Where to Learn: Training Data Distribution Optimization for Scientific Machine Learning

ICLR 2026 Conference Submission20912 Authors

Keywords: bilevel optimization, distribution shift, experimental design, probability measure, operator learning

TL;DR: This paper tackles the where-to-learn problem in scientific ML, showing how optimizing training distributions via bilevel and alternating algorithms produces more accurate and robust models for function approximation and PDE operator learning

Abstract: 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.

Primary Area: optimization

Submission Number: 20912