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

2026-05-07

AI Summary

This work addresses the limitation of traditional methods that learn conditional distributions separately for fixed joint distributions, thereby struggling to generalize to unseen distribution pairs. It reframes conditional probability modeling as a universal operator learning problem across distributions and proposes a single neural operator that amortizes the approximation of conditional densities by mapping arbitrary joint densities to their corresponding conditionals. Leveraging neural operator theory and continuity analysis in spaces of probability densities, the authors establish that this operator is continuous over suitable classes of densities and can be arbitrarily well approximated by neural networks. Empirical validation on Gaussian mixture distributions demonstrates the framework’s effectiveness, showcasing strong generalization capabilities and promising applicability in tasks such as Bayesian inference.

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.

Problem

• probabilistic conditioning
• neural operators
• amortized inference
• conditional distribution
• joint density

Innovation

• neural operators
• amortized inference
• probabilistic conditioning
• conditional density estimation
• foundation models