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.
翻译:扩散模型已成为视觉、语言和强化学习领域有效的分布估计器,但其在下游任务中作为先验使用时,会引出一个难解的后验推理问题。本文研究数据后验 $\mathbf{x}\sim p^{\rm post}(\mathbf{x})\propto p(\mathbf{x})r(\mathbf{x})$ 的摊销采样问题,该模型由扩散生成模型先验 $p(\mathbf{x})$ 与黑盒约束或似然函数 $r(\mathbf{x})$ 构成。我们提出并证明了一种无数据学习目标——相对轨迹平衡的渐近正确性,该目标用于训练能够从该后验中采样的扩散模型,而现有方法仅能在近似或受限情况下解决此问题。相对轨迹平衡源于扩散模型的生成流网络视角,该视角允许利用深度强化学习技术来提升模态覆盖度。实验展示了在扩散先验下对任意后验进行无偏推理的广泛潜力:在视觉领域(分类器引导)、语言领域(基于离散扩散大语言模型的文本填充)以及多模态数据领域(文本到图像生成)。除生成建模外,我们将相对轨迹平衡应用于基于分数行为先验的连续控制问题,在离线强化学习基准测试中取得了最先进的结果。