Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models

Siddarth Venkatraman, Mohsin Hasan, Minsu Kim, Luca Scimeca, Marcin Sendera, Yoshua Bengio, Glen Berseth, Nikolay Malkin

TL;DR: To sample Bayesian posteriors (for constrained generation, RLHF, and more) under a pretrained prior, we train diffusion models to sample the noise space, applicable to any generative model.

Abstract:

Any well-behaved generative model over a variable $\mathbf{x}$ can be expressed as a deterministic transformation of an exogenous (‘*outsourced'*) Gaussian noise variable $\mathbf{z}$: $\mathbf{x}=f_\theta(\mathbf{z})$.
In such a model (*eg*, a VAE, GAN, or continuous-time flow-based model), sampling of the target variable $\mathbf{x} \sim p_\theta(\mathbf{x})$ is straightforward, but sampling from a posterior distribution of the form $p(\mathbf{x}\mid\mathbf{y}) \propto p_\theta(\mathbf{x})r(\mathbf{x},\mathbf{y})$, where $r$ is a constraint function depending on an auxiliary variable $\mathbf{y}$, is generally intractable.
We propose to amortize the cost of sampling from such posterior distributions with diffusion models that sample a distribution in the noise space ($\mathbf{z}$). These diffusion samplers are trained by reinforcement learning algorithms to enforce that the transformed samples $f_\theta(\mathbf{z})$ are distributed according to the posterior in the data space ($\mathbf{x}$).
For many models and constraints, the posterior in noise space is smoother than in data space, making it more suitable for amortized inference. Our method enables conditional sampling under unconditional GAN, (H)VAE, and flow-based priors, comparing favorably with other inference methods. We demonstrate the proposed ___outsourced diffusion sampling___ in several experiments with large pretrained prior models: conditional image generation, reinforcement learning with human feedback, and protein structure generation.

Lay Summary: The paper introduces a novel approach for steering the output distribution of a generative model toward samples that achieve high reward under a given objective. Instead of modifying the generative model directly, a separate diffusion model is trained using reinforcement learning to sample noise vectors that when passed through the generative model, produce high-reward outputs.

Link To Code: https://github.com/HyperPotatoNeo/Outsourced_Diffusion_Sampling

Primary Area: Probabilistic Methods->Bayesian Models and Methods

Keywords: diffusion, amortized inference, inverse problems, fine-tuning, VAEs, GANs, CNFs, flow matching, generative models

Submission Number: 2615