One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operators

TL;DR

This paper proposes a neural operator approach to approximate the conditioning operator for any joint density, enabling amortized probabilistic inference with theoretical guarantees and practical demonstrations on Gaussian mixtures.

Contribution

It introduces a novel neural operator framework for universal approximation of the conditioning operator, supported by new continuity results and empirical learning examples.

Findings

Neural operators can approximate the conditioning operator arbitrarily well.

The framework is applicable to a class of Gaussian mixtures.

The approach provides a foundation for general-purpose Bayesian inference methods.

Abstract

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.