Amortizing intractable inference in diffusion models for vision, language, and control

TL;DR

This paper introduces a theoretically grounded method called relative trajectory balance for amortized inference in diffusion models, enabling efficient posterior sampling in vision, language, and control tasks, with broad practical applications.

Contribution

It proposes a new asymptotically correct learning objective for diffusion-based posterior inference, extending the use of deep reinforcement learning techniques to improve mode coverage.

Findings

Unbiased inference of arbitrary posteriors demonstrated in vision, language, and multimodal tasks.

Achieved state-of-the-art results in offline reinforcement learning using diffusion-based behavior priors.

Theoretical proof of asymptotic correctness for the proposed learning objective.

Abstract

Diffusion models have emerged as effective distribution estimators in vision, language, and reinforcement learning, but their use as priors in downstream tasks poses an intractable posterior inference problem. This paper studies amortized sampling of the posterior over data, $\mathbf{x}\sim p^{\rm post}(\mathbf{x})\propto p(\mathbf{x})r(\mathbf{x})$, in a model that consists of a diffusion generative model prior $p(\mathbf{x})$ and a black-box constraint or likelihood function $r(\mathbf{x})$. We state and prove the asymptotic correctness of a data-free learning objective, relative trajectory balance, for training a diffusion model that samples from this posterior, a problem that existing methods solve only approximately or in restricted cases. Relative trajectory balance arises from the generative flow network perspective on diffusion models, which allows the use of deep reinforcement learning techniques to improve mode coverage. Experiments illustrate the broad potential of unbiased inference of arbitrary posteriors under diffusion priors: in vision (classifier guidance), language (infilling under a discrete diffusion LLM), and multimodal data (text-to-image generation). Beyond generative modeling, we apply relative trajectory balance to the problem of continuous control with a score-based behavior prior, achieving state-of-the-art results on benchmarks in offline reinforcement learning.