用二次奖励精准调控扩散模型采样,揭示可高效实现的边界
Steering diffusion models with quadratic rewards: a fine-grained analysis
- 基于哈伯-斯特拉顿维奇变换,设计高效采样算法
- 线性奖励可高效采样,低秩正定二次奖励也可行
- 负定二次奖励即使秩为1也难解,理论边界清晰
推理时算法作为新兴范式,将预训练模型作为子程序解决下游任务。但现有方法多为启发式,失败模式多样,且缺乏对改进效率的理解。本文研究从奖励倾斜的扩散模型采样问题:给定奖励函数 $r(x)$ 与预训练扩散模型 $p$,求 $p^{ ext{⋆}}(x) \propto p(x) \exp(r(x))$。针对二次奖励 $r(x) = x^ op A x + b^ op x$,我们给出精细分析。发现线性奖励倾斜始终可高效采样(此前未被注意)。以此为基础,结合新提出的哈伯-斯特拉顿维奇变换,提出低秩正定二次奖励($A$ 正定且秩 $O(1)$)的高效采样算法。而对于负定二次奖励 $r(x) = -x^ op A x$($A$ 正定),即使秩为1,也证明其在计算上不可行(尽管矩阵元素指数级大)。
原文摘要 · Abstract (English)
Inference-time algorithms are an emerging paradigm in which pre-trained models are used as subroutines to solve downstream tasks. Such algorithms have been proposed for tasks ranging from inverse problems and guided image generation to reasoning. However, the methods currently deployed in practice are heuristics with a variety of failure modes -- and we have very little understanding of when these heuristics can be efficiently improved. In this paper, we consider the task of sampling from a reward-tilted diffusion model -- that is, sampling from $p^{\star}(x) \propto p(x) \exp(r(x))$ -- given a reward function $r$ and pre-trained diffusion oracle for $p$. We provide a fine-grained analysis of the computational tractability of this task for quadratic rewards $r(x) = x^\top A x + b^\top x$. We show that linear-reward tilts are always efficiently sampleable -- a simple result that seems to have gone unnoticed in the literature. We use this as a building block, along with a conceptually new ingredient -- the Hubbard-Stratonovich transform -- to provide an efficient algorithm for sampling from low-rank positive-definite quadratic tilts, i.e. $r(x) = x^\top A x$ where $A$ is positive-definite and of rank $O(1)$. For negative-definite tilts, i.e. $r(x) = - x^\top A x$ where $A$ is positive-definite, we prove that the problem is intractable even if $A$ is of rank 1 (albeit with exponentially-large entries).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。