arXiv:2506.20701cs.LGcs.AI2025-06NeurIPS被引 34

让扩散模型在推理时更准更快,靠树搜索复用历史计算结果。

Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models

  • 用树搜索思路重用之前生成结果,迭代优化采样方向。
  • 在MNIST/CIFAR-10上用1/10算力达到最佳基线的生成质量。
  • 适合需要高效高质生成的文本到图像与语言补全任务。

将预训练扩散模型适配新目标的推理时对齐仍是生成建模中的开放问题。现有引导方法在高噪声水平下价值估计不准确,导致引导偏差;且过往生成信息未被复用,造成计算效率低下。受蒙特卡洛树搜索成功启发,我们将推理时对齐视为可复用历史计算的搜索问题。提出基于树的采样方法,通过反向传播终端奖励并迭代修正价值估计,实现从目标分布中渐进精确采样。所提方法Diffusion Tree Sampling(DTS)在无限回滚下收敛至目标分布精确样本,其贪心变体Diffusion Tree Search(DTS$^ ext{star}$)可全局搜索高奖励样本。在MNIST与CIFAR-10类别条件生成任务中,DTS以最高10倍更少算力达到最优基线的FID表现;在文本到图像生成与语言补全任务中,DTS$^ ext{star}$以最高5倍更少算力实现与best-of-N相当的高质量样本。通过复用先前生成信息,该方法成为任意时间算法,额外算力持续提升样本质量,为扩散模型推理时对齐提供可扩展方案。

原文摘要 · Abstract (English)

Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling. Existing steering methods suffer from inaccurate value estimation, especially at high noise levels, which biases guidance. Moreover, information from past runs is not reused to improve sample quality, resulting in inefficient use of compute. Inspired by the success of Monte Carlo Tree Search, we address these limitations by casting inference-time alignment as a search problem that reuses past computations. We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain and iteratively refining value estimates with each additional generation. Our proposed method, Diffusion Tree Sampling (DTS), produces asymptotically exact samples from the target distribution in the limit of infinite rollouts, and its greedy variant, Diffusion Tree Search (DTS$^\star$), performs a global search for high reward samples. On MNIST and CIFAR-10 class-conditional generation, DTS matches the FID of the best-performing baseline with up to $10\times$ less compute. In text-to-image generation and language completion tasks, DTS$^\star$ effectively searches for high reward samples that match best-of-N with up to $5\times$ less compute. By reusing information from previous generations, we get an anytime algorithm that turns additional compute into steadily better samples, providing a scalable approach for inference-time alignment of diffusion models.

扩散模型推理优化树搜索生成质量

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。