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

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.

• The paper introduces an operator-theoretic formulation of conditioning that maps joint densities to conditional densities using neural operators.
• It establishes universal approximation theorems and stability estimates, demonstrating that structured neural networks can uniformly approximate these conditioning operators.
• Empirical tests with Fourier and Transformer Neural Operators show reduced L1 errors in kernel and in-context settings, enhancing Bayesian inference efficiency.

Amortized Conditioning Operators via Neural Operator Learning

Problem Formulation and Operator-Theoretic Perspective

The paper "One Operator for Many Densities: Amortized Approximation of Conditioning by Neural Operators" (2605.06873) presents a novel formulation for probabilistic conditioning. Instead of learning conditional distributions for each fixed joint density, the authors recast conditioning as a nonlinear operator $\Psi^\star$ mapping any joint density $\rho$ and query $y$ to the conditional density $\rho(x\,|\,y)$. Two operator formulations are studied: the in-context conditioning operator, which returns a conditional density given $(\rho, y)$, and the kernel conditioning operator, which maps $\rho$ to a kernel $\kappa_\rho(x,y)$ representing the family of conditionals.

This operator-theoretic formalism leads to a function space perspective: conditioning becomes a mapping in infinite-dimensional Banach spaces (e.g., spaces of continuous probability densities or $L^1$ spaces). Crucially, results are established regarding the continuity and stability of these operators under appropriate topologies, notably the supremum and $L^1$ norms, with quantitative Hölder and Lipschitz bounds for stability under perturbation of the input density.

Universal Approximation with Neural Operators

Building on the continuity properties, the paper proves that both conditioning operators admit uniform approximation over compact sets of densities by neural operators (NOs)—structured neural networks designed to learn mappings between function spaces. The authors provide formal definitions for lifting, hidden, and projection layers relevant to neural operator architectures, including augmented NOs (AugNOs) operating on function-vector inputs.

Two main universal approximation theorems are established:

• For any compact set of admissible densities $K$ and $\rho$0, there exists a neural operator $\rho$1 such that the kernel conditioning operator is approximated to within $\rho$2 uniformly on $\rho$3.
• Similarly, for the in-context operator, an AugNO achieves uniform accuracy over $\rho$4.

These results extend to densities supported on unbounded domains (e.g., $\rho$5), using full space neural operators (FullNOs) involving truncation and extension layers. Approximation guarantees are given for compact sets in weighted Sobolev spaces, notably including Gaussian mixtures with bounded means/covariances.

Stability and Extensions Beyond Lower-Bounded Densities

Local and global stability results, with explicit dependence on uniform lower bounds and supremum norms, are provided for the kernel and in-context operators. The paper demonstrates continuity and derives quantitative estimates over balls in Hölder spaces and more general subsets. Extensions beyond the requirement of strictly positive marginals are constructed, giving limits of convolutions with mollifiers as continuous extensions.

Examples show that the operator-theoretic conditioning cannot be continuously extended to all probability measures with respect to weaker (e.g., weak or narrow) convergence topologies, illuminating fundamental limitations.

Numerical Demonstrations

Proof-of-concept experiments support the theoretical claims. Fourier Neural Operators (FNOs) and Transformer Neural Operators (TNOs) are trained to approximate the conditioning operator over classes of Gaussian mixture joints. Results show superior median relative $\rho$6 error compared to plug-in estimators, particularly notable in kernel density estimation (KDE) settings for $\rho$7. The architectures also scale to $\rho$8 mixtures, though maximum errors indicate challenges in more heterogeneous cases, motivating further methodological advances for high-dimensional settings.

Figures illustrate per-sample error comparisons and model predictions, validating the practical implications of operator learning for probabilistic conditioning.

Theoretical and Practical Implications

The central implication is that conditioning, understood as an operator on function spaces, can be learned and amortized by a generic neural operator. This enables foundation models for Bayesian inference—neural architectures trained on extensive datasets of joint and conditional densities, capable of fast, general-purpose inference on unseen input distributions. This paradigm removes the necessity of re-training for each new scenario, which is particularly relevant in sequential inference, nonlinear filtering, and design applications.

From a theoretical viewpoint, the paper unifies operator learning methods (DeepONet, FNO, TNO) with probabilistic conditioning and draws connections to conditional mean embedding, measure transport, and modern distribution-valued learning. The universal approximation theorems facilitate rigorous integration of neural operator design into statistical inference pipelines.

Future Research Directions

Several avenues are identified:

• High-dimensional scalability: Efficient density representations for neural operators, beyond coarse grid discretization, are essential for practical deployment.
• Sample-based input: Combining neural operator conditioning with generative or density estimation models to enable conditioning from empirical samples, rather than analytic densities.
• Generalized measure conditioning: Extending the operator learning approach and continuity theory to classes of probability measures lacking absolutely continuous densities or including singular measures, aiming for broader universal approximation results.

Conclusion

This work rigorously formalizes probabilistic conditioning as a continuous operator in function spaces and demonstrates universal approximation by neural operators. The mathematical and computational results represent a foundational step toward building amortized, general-purpose models for Bayesian inference, opening up new directions in uncertainty quantification and probabilistic machine learning. The theoretical guarantees provided support the design of scalable, efficient inference architectures that operate across diverse classes of densities.