One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operators
For researchers in probabilistic inference and Bayesian methods, this work provides theoretical foundations for amortized conditioning, potentially enabling foundation models for Bayesian inference.
This paper proposes learning a single neural operator that maps any joint density to its conditional distribution, amortizing over joint-conditional pairs. The authors prove that the conditioning operator can be approximated to arbitrary accuracy by neural operators and demonstrate the approach on Gaussian mixtures.
Probabilistic conditioning is concerned with the identification of a distribution of a random variable $X$ given a random variable $Y$. It is a cornerstone of scientific and engineering applications where modeling uncertainty is key. This problem has traditionally been addressed in machine learning by directly learning the conditional distribution of a fixed joint distribution. This paper introduces a novel perspective: we propose to solve the conditioning problem by identifying a single operator that maps any joint density to its conditional, thus amortizing over joint-conditional pairs. We establish that the conditioning operator can be approximated to arbitrary accuracy by neural operators. Our proof relies on new results establishing continuity of the conditioning operator over suitable classes of densities. Finally, we learn the conditioning map for a class of Gaussian mixtures using neural operators, illustrating the promise of our framework. This work provides the theoretical underpinnings for general-purpose, amortized methods for probabilistic conditioning, such as foundation models for Bayesian inference.