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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Venkatraman, Siddarth, Jain, Moksh, Scimeca, Luca, Kim, Minsu, Sendera, Marcin, Hasan, Mohsin, Rowe, Luke, Mittal, Sarthak, Lemos, Pablo, Bengio, Emmanuel, Adam, Alexandre, Rector-Brooks, Jarrid, Bengio, Yoshua, Berseth, Glen, Malkin, Nikolay
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912184777310208
author Venkatraman, Siddarth
Jain, Moksh
Scimeca, Luca
Kim, Minsu
Sendera, Marcin
Hasan, Mohsin
Rowe, Luke
Mittal, Sarthak
Lemos, Pablo
Bengio, Emmanuel
Adam, Alexandre
Rector-Brooks, Jarrid
Bengio, Yoshua
Berseth, Glen
Malkin, Nikolay
author_facet Venkatraman, Siddarth
Jain, Moksh
Scimeca, Luca
Kim, Minsu
Sendera, Marcin
Hasan, Mohsin
Rowe, Luke
Mittal, Sarthak
Lemos, Pablo
Bengio, Emmanuel
Adam, Alexandre
Rector-Brooks, Jarrid
Bengio, Yoshua
Berseth, Glen
Malkin, Nikolay
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2405_20971
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Amortizing intractable inference in diffusion models for vision, language, and control
Venkatraman, Siddarth
Jain, Moksh
Scimeca, Luca
Kim, Minsu
Sendera, Marcin
Hasan, Mohsin
Rowe, Luke
Mittal, Sarthak
Lemos, Pablo
Bengio, Emmanuel
Adam, Alexandre
Rector-Brooks, Jarrid
Bengio, Yoshua
Berseth, Glen
Malkin, Nikolay
Machine Learning
Computer Vision and Pattern Recognition
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.
title Amortizing intractable inference in diffusion models for vision, language, and control
topic Machine Learning
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2405.20971