Multimodal Scientific Learning Beyond Diffusions and Flows

TL;DR

The paper argues that uncertainty in Scientific Machine Learning is often multimodal due to ill-posed inverses, chaotic dynamics, and multistability, challenging unimodal regression. It champions Mixture Density Networks as explicit density estimators that model p(y|x) with a finite Gaussian mixture, providing data efficiency and interpretability compared to diffusion/flow-based models. A unified probabilistic framework positions MDNs between point-predictors and implicit samplers, and empirical results across inverse, dynamical, and chaotic tasks show competitive or superior performance in low-data regimes, with direct insight from mixture weights into physical regimes. The work also introduces JaxMix and discusses limitations and future directions toward broader adoption and domain-aligned backbones.

Abstract

Scientific machine learning (SciML) increasingly requires models that capture multimodal conditional uncertainty arising from ill-posed inverse problems, multistability, and chaotic dynamics. While recent work has favored highly expressive implicit generative models such as diffusion and flow-based methods, these approaches are often data-hungry, computationally costly, and misaligned with the structured solution spaces frequently found in scientific problems. We demonstrate that Mixture Density Networks (MDNs) provide a principled yet largely overlooked alternative for multimodal uncertainty quantification in SciML. As explicit parametric density estimators, MDNs impose an inductive bias tailored to low-dimensional, multimodal physics, enabling direct global allocation of probability mass across distinct solution branches. This structure delivers strong data efficiency, allowing reliable recovery of separated modes in regimes where scientific data is scarce. We formalize these insights through a unified probabilistic framework contrasting explicit and implicit distribution networks, and demonstrate empirically that MDNs achieve superior generalization, interpretability, and sample efficiency across a range of inverse, multistable, and chaotic scientific regression tasks.

Figures (20)

• Figure 1: Illustration of increasing UQ complexity. Point estimates provide no uncertainty (left), while arbitrary generative models are powerful but data/compute intensive (right). Many scientific problems are well approximated by simpler multimodal mixture distributions.
• Figure 2: True conditional density of $x$ given $y$ in the inverse sine problem. (left) Conditional NLL. (right) Conditional likelihood.
• Figure 3: Comparison of MDN and CFM on the sinusoidal inverse problem. Top: test NLL versus training set size, with bands showing 1 Std. Dev. over 12 seeds. Bottom: averaged NLL surfaces.
• Figure 4: Top three mixture elements (by marginal probability) for the three scenarios in the gravitational attractor example. The left columns show mixture element samples, while the right show the weight of that mixture element given the input. Sample opacities are weighted by mixture probability. (left) Case 1. (middle) Case 2. (right) Case 3.
• Figure 5: Data efficiency comparison in the bifurcation system. Shaded regions indicate $\pm 1$ Std. Dev. of 12 ensemble members.