Amortized Probabilistic Conditioning for Optimization, Simulation and Inference

2024-10-20

arXiv.org

AI Summary

Existing amortized meta-learning approaches struggle to flexibly condition on and extract probabilistic latent variables at inference time, limiting their applicability in Bayesian inference, optimization, and simulation. To address this, we propose the Amortized Conditioning Engine (ACE), a Transformer-based, conditional neural process model that enables *bidirectional explicit manipulation* of probabilistic latent variables for the first time—supporting dynamic injection of observed data and interpretable latents, prior embedding, and joint generation of discrete/continuous data and predictive latent distributions. ACE unifies variational inference, amortized inference, and conditional generative modeling to handle supervised learning, Bayesian optimization, and simulation-based inference within a single framework. Experiments demonstrate that ACE significantly outperforms existing neural process methods on image completion and classification, black-box optimization, and simulation-based inference—achieving superior flexibility, probabilistic calibration, and cross-task generalization.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Amortized meta-learning methods based on pre-training have propelled fields like natural language processing and vision. Transformer-based neural processes and their variants are leading models for probabilistic meta-learning with a tractable objective. Often trained on synthetic data, these models implicitly capture essential latent information in the data-generation process. However, existing methods do not allow users to flexibly inject (condition on) and extract (predict) this probabilistic latent information at runtime, which is key to many tasks. We introduce the Amortized Conditioning Engine (ACE), a new transformer-based meta-learning model that explicitly represents latent variables of interest. ACE affords conditioning on both observed data and interpretable latent variables, the inclusion of priors at runtime, and outputs predictive distributions for discrete and continuous data and latents. We show ACE's modeling flexibility and performance in diverse tasks such as image completion and classification, Bayesian optimization, and simulation-based inference.

Problem

• Flexible runtime conditioning on probabilistic latent information
• Explicit representation of interpretable latent variables
• Improved performance in optimization, simulation, and inference tasks

Innovation

• Transformer-based meta-learning model
• Explicit representation of latent variables
• Conditioning on observed and latent data