Best-of-N Guidance for Test-time Diffusion Alignment¶
会议: NeurIPS 2026
arXiv: 2610.05108
代码: https://github.com/aailab-kaist/BoNG
领域: 图像生成
关键词: 扩散模型, 测试时对齐, Best-of-N, 样本引导, 人类偏好
一句话总结¶
针对传统 Best-of-N 仅在后处理选优且无法改善平均样本分布、而现有采样引导方法依赖繁重 rollout 的瓶颈,本文提出免展开(rollout-free)的 Best-of-N Guidance (BoNG),在线利用一步去噪样本评选最高奖励粒子作为引导目标,在无需额外神经网络梯度或前向计算的前提下实现了单样本与多样本的高质量对齐。
研究背景与动机¶
扩散模型在文生图等条件生成任务中展现出卓越的表达能力,但预训练生成分布往往无法完全契合由人类偏好奖励模型所定义的细粒度偏好或输入对齐度。为了在推理阶段改善对齐质量,免重训的测试时对齐(test-time alignment)逐渐受到学术界与工业界的青睐。其中,最直接的做法是 Best-of-\(N\)(BoN)采样:独立抽取 \(N\) 个完整样本后取奖励最高者;另一种常见范式是序列蒙特卡洛(SMC)粒子滤波重采样。
然而,现有测试时对齐机制存在明显的效率与分布缺陷。Vanilla BoN 仅将奖励函数作为纯后处理的选优过滤器,整个逆向扩散去噪轨迹完全由基础模型主导,未受奖励信息的任何主动干预,导致样本群体的平均对齐质量毫无改善,且计算成本随候选数量急剧增长,无法直接服务于高质量多输出展示场景。SMC 虽然在去噪中间步引入了重采样交互,但高权重粒子的高频复制会导致严重的“粒子坍缩”(particle collapse),损害生成多样性及最终优选上限。另一方面,基于梯度引导的方法需要穿透奖励模型和扩散网络的昂贵反向传播;而现有的免梯度样本引导方法(如 Tilt、LiDAR)则依赖于多步前向或反向展望展开(rollout),带来 2 到 3 倍以上的采样延迟与额外前向开销。
本文的切入角度是:能否在逆向扩散的多粒子并行生成轨迹中,直接内生化 Best-of-\(N\) 的选优机制,让群体中表现最优的粒子实时“引航”其余粒子,既避免粒子拷贝坍缩,又彻底省去任何展望展开计算?核心 idea:在逆向扩散的关键时间步窗口内,基于各粒子的一步 Tweedie 去噪预测在线评选当前最高奖励粒子,将其构造为非对称引导目标注入逆向扩散得分中,在理论上严格逼近奖励倾斜分布(reward-tilted distribution)的大 \(\lambda\) 极限。
方法详解¶
整体框架¶
BoNG 的整体流程旨在并行演化 \(N\) 个去噪粒子群 \(\{{\mathbf{x}}_t^i\}_{i=1}^N\)。在常规时间步,粒子独立沿基础扩散网络前向反演;而在预设的引导时间步窗口 \([t_{\mathrm{start}}, t_{\mathrm{end}})\) 内,算法在每一步首先通过单步闭式 Tweedie 公式预测每个粒子的干净样本 \(\hat{{\mathbf{x}}}_0^i\),经由 VAE 解码计算标量偏好奖励,在线甄选出当前奖励最高的粒子索引 \(i^*\);随后构建非对称引导项,将处于领先地位的粒子保留原轨迹推进,同时对其余非最优粒子施加指向最高奖励粒子的引力修正,协同将整体粒子群推向高奖励流形。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["多粒子并行输入<br/>当前噪声群体 {x_t^i}"] --> B["一步 Tweedie 预测与在线选优<br/>闭式计算 x̂_0^i 并评选最高奖励粒子 i*"]
B --> C["非对称粒子引导交互<br/>锚定粒子自由演化,其余粒子注入正向引力"]
C --> D["时序引导窗口与步进更新<br/>在预设时序区间缩放得分并推进去噪"]
D --> E["对齐样本输出<br/>兼顾单样本选优与多样本高质多样"]
关键设计¶
1. 一步 Tweedie 预测与在线选优:彻底摆脱展望展开计算
现有样本引导方法为了估算含噪状态 \({\mathbf{x}}_t\) 在最终干净空间的期望奖励,不得不从中间步模拟完整或多步展开轨迹,这构成了计算效率的主要瓶颈。BoNG 的关键洞察在于:利用扩散模型得分函数的 Tweedie 公式,可以在不调用任何额外 UNet 前向计算的前提下,直接由当前步的隐变量与网络预测噪声获得一步去噪期望样本: $$ \hat{{\mathbf{x}}}0^i := \mathbb{E}}({\mathbf{x}}_0 \mid {\mathbf{x}}_t^i, {\mathbf{c}})}[{\mathbf{x}}_0] = {\mathbf{x}t^i + \sigma_t^2 \nabla}t^i} \log p}({\mathbf{x}t^i \mid {\mathbf{c}}) $$ 在粒子群内,这组一步去噪估计 \(\{\hat{{\mathbf{x}}}_0^i\}_{i=1}^N\) 构成了目标流形 \(\mathcal{X}_0\) 上的瞬时候选集。算法仅对这 \(N\) 个估计进行 VAE 解码并输入标量奖励模型 \(r(\cdot, {\mathbf{c}})\),在线甄选出具有全局最优奖励的候选索引: $$ i^* = \arg\max} r(\hat{{\mathbf{x}}0^i, {\mathbf{c}}), \quad \hat{{\mathbf{x}}}_0^{\mathrm{BoN} := \hat{{\mathbf{x}}}_0^{i^}, \quad {\mathbf{x}}_t^{\mathrm{BoN}} := {\mathbf{x}}_t^{i^} $$ 这种做法直接复用了基础扩散采样中本就必须计算的网络得分,将奖励探索完全寄托于现成粒子之间的自然随机差异,彻底消除了传统 rollout 所需的前向求解器开销。
2. 非对称粒子引导交互:引导流形修正而杜绝粒子坍缩
传统 SMC 算法在检测到高奖励粒子后执行重采样,本质上是粒子位置的离散复制与低分粒子的强制淘汰,这在有限粒子数下不可避免地导致所有粒子收敛到相同模式。为了同时保持样本多样性并提升整体验收质量,BoNG 设计了一种连续动力学意义上的非对称引导交互。对于任意粒子 \({\mathbf{x}}_t^i\),定义其对齐得分函数为: $$ \nabla_{{\mathbf{x}}t^i} \log p}^{\mathrm{BoN}}({\mathbf{x}t^i \mid {\mathbf{c}}) := \nabla}t^i} \log p}({\mathbf{x}}_t^i \mid {\mathbf{c}}) + \frac{1}{\sigma_t^2}\left(\hat{{\mathbf{x}}}_0^{\mathrm{BoN}} - \hat{{\mathbf{x}}0^i\right) $$ 展开后可以清晰观察到其非对称性机制:对于获选的最优粒子 \(i^*\),引导项由于 \((\hat{{\mathbf{x}}}_0^{\mathrm{BoN}} - \hat{{\mathbf{x}}}_0^{i^*}) = \mathbf{0}\) 而完全消失,该粒子保持无偏的基础逆向扩散;而对于其余粒子 \(i \ne i^*\),更新方向被精确分解为: $$ \nabla}t^i} \log p}^{\mathrm{BoN}}({\mathbf{x}t^i \mid {\mathbf{c}}) = \nabla}t^{\mathrm{BoN}}} \log p_t^i\right) $$ 前一项赋予非最优粒子向最优粒子的扩散去噪方向看齐的动量,后一项则施加了直接拉向最优粒子空间位置的凸引力。粒子没有被粗暴替换,其固有的轨迹随机性与微观差异得以完整保留,既达成了群体牵引,又避免了模式坍缩。}({\mathbf{x}}_t^{\mathrm{BoN}} \mid {\mathbf{c}}) + \frac{1}{\sigma_t^2}\left({\mathbf{x}}_t^{\mathrm{BoN}} - {\mathbf{x}
3. 时序引导窗口与步进更新:逼近倾斜分布极限与平滑控制
理论上,当把奖励倾斜分布的自归一化重要性采样(SNIS)权重推向温度极限制(\(\lambda \to \infty\))时,权重会完全退化为 Dirac 脉冲并精准聚集于最高奖励候选上,从而在数学上严格证明了 BoNG 正好对应奖励倾斜分布梯度的大 \(\lambda\) 极限。为了在离散采样过程中灵活调节奖励优化强度与多样性边界,BoNG 引入了引导时间步窗口 \([t_{\mathrm{start}}, t_{\mathrm{end}})\) 与缩放系数 \(s\)。在实际调度中,反向扩散得分被缩放修正为: $$ \bm{\epsilon}^i \leftarrow \bm{\epsilon}^i - \sqrt{\bar{\alpha}_t} \cdot s \cdot \frac{\hat{{\mathbf{x}}}_0^{\mathrm{BoN}} - \hat{{\mathbf{x}}}_0^i}{\sigma_t} $$ 作者发现,在扩散早期的超高噪声阶段,Tweedie 估计与最终图像偏好的相关性极低(相关系数不足 0.2),而在晚期细化阶段结构已定型无法改道;因此将窗口精准锁定在相关性超过 0.6 且宏观语义尚未锁死的中间过渡阶段(例如 DDPM 100 步采样的第 40 到 44 步),以极少的时间步激活(\(K \le 4\))即可实现最大化的对齐效能。
损失函数 / 训练策略¶
BoNG 属于纯推理期免微调(test-time alignment)算法,不需要对扩散模型的 UNet/DiT 参数进行重新训练,也不需要针对中间带噪隐变量训练专门的时序奖励模型。在算法落地时,仅需接入预训练的黑盒判别器(如 ImageReward、CLIP、HPSv2 等)对 Tweedie 估计解码图像进行评分即可运作。
实验关键数据¶
主实验¶
论文使用 GenEval 评测提示词库,在 SD v1.5(0.9B)和 SDXL(2.6B)上,全面对比了 Vanilla BoN、SMC 以及 BoNG 在不同粒子预算 \(N \in \{2, 3, 4\}\) 下的最佳单样本(BoN)表现;同时在多输出设定下对比了梯度引导方法(UG、DATE)和最新样本引导 SOTA(LiDAR)。
| 模型与骨干 | 对齐方法 | 样本数 \(N\) | ImageReward (↑) | CLIP Score (↑) | GenEval (↑) | 单次耗时 (s) | 显存 (GiB) |
|---|---|---|---|---|---|---|---|
| SD v1.5 (DDIM 50) | Vanilla | \(N=4\) | 0.682 | 0.283 | 0.520 | 3.54 | 8.90 |
| SD v1.5 (DDIM 50) | SMC | \(N=4\) | 0.636 | 0.282 | 0.509 | 4.12 | 8.90 |
| SD v1.5 (DDIM 50) | BoNG (本文) | \(N=4\) | 0.793 | 0.285 | 0.535 | 4.03 | 8.90 |
| SD v1.5 (DDPM 100) | Vanilla | \(N=4\) | 0.742 | 0.284 | 0.525 | 7.07 | 8.90 |
| SD v1.5 (DDPM 100) | SMC | \(N=4\) | 0.826 | 0.286 | 0.529 | 8.23 | 8.90 |
| SD v1.5 (DDPM 100) | BoNG (本文) | \(N=4\) | 0.894 | 0.286 | 0.555 | 8.06 | 8.90 |
| SDXL (DDIM 50) | Vanilla | \(N=4\) | 1.140 | 0.289 | 0.613 | 11.20 | 18.50 |
| SDXL (DDIM 50) | SMC | \(N=4\) | 1.093 | 0.289 | 0.612 | 12.80 | 18.50 |
| SDXL (DDIM 50) | BoNG (本文) | \(N=4\) | 1.149 | 0.292 | 0.606 | 12.40 | 18.50 |
| 多输出对比 (4图平均) | UG | \(N=4\) | 0.326 | 0.262 | 0.355 | 58.36 | 28.16 |
| 多输出对比 (4图平均) | DATE | \(N=4\) | 0.364 | 0.274 | 0.438 | 32.89 | 24.71 |
| 多输出对比 (4图平均) | LiDAR | \(N=4\) | 0.384 | 0.278 | 0.478 | 13.41 | 8.90 |
| 多输出对比 (4图平均) | BoNG (本文) | \(N=4\) | 0.506 | 0.280 | 0.511 | 8.06 | 8.90 |
在 36 项跨骨干、采样器与候选数对比中,BoNG 在 29 项中位列第一(胜率达 80.56%)。多输出生成均值方面,BoNG 相比前沿方法 LiDAR 取得了约 \(1.3\times\) 的 ImageReward 得分,同时实现 \(1.6\times\) 的端到端推理加速。
消融实验¶
论文针对引导窗口起始点 \(t_{\mathrm{start}}\)、窗口尺寸 \(K\) 以及引导尺度 \(s\) 在 SD v1.5(DDPM 100步, \(N=4\))下展开了严密消融。
| 配置与超参变量 | 调优取值 | 平均 ImageReward (↑) | CLIP Score (↑) | GenEval (↑) | 样本多样性 (↑) |
|---|---|---|---|---|---|
| 窗口区间消融 (\(K=4, s=6\)) | \([20, 24]\)(过早阶段) | 0.753 | 0.282 | 0.523 | - |
| 窗口区间消融 (\(K=4, s=6\)) | \([40, 44]\)(黄金区间) | 0.875 | 0.286 | 0.559 | 0.224 |
| 窗口区间消融 (\(K=4, s=6\)) | \([60, 64]\)(偏后阶段) | 0.849 | 0.287 | 0.526 | - |
| 窗口区间消融 (\(K=4, s=6\)) | \([80, 84]\)(微调阶段) | 0.866 | 0.287 | 0.528 | - |
| 窗口步数 \(K\) 扫描 (\(t=40, s=1\)) | \(K=4\)(标准设置) | 0.019 | 0.272 | - | 0.319 |
| 窗口步数 \(K\) 扫描 (\(t=40, s=1\)) | \(K=15\)(中度引导) | 0.465 | 0.275 | - | 0.262 |
| 窗口步数 \(K\) 扫描 (\(t=40, s=1\)) | \(K=35\)(激进引导) | 0.707 | 0.279 | - | 0.121 |
| 引导尺度 \(s\) 扫描 (\([40,44], K=4\)) | \(s=1\) | 0.019 | 0.272 | - | 0.319 |
| 引导尺度 \(s\) 扫描 (\([40,44], K=4\)) | \(s=3\) | 0.490 | 0.278 | - | 0.255 |
| 引导尺度 \(s\) 扫描 (\([40,44], K=4\)) | \(s=6\)(基准配置) | 0.506 | 0.280 | - | 0.224 |
关键发现¶
- 时序窗口的平衡至关重要:窗口过早(如 \([20, 24]\))由于 Tweedie 图像估计信噪比太低,引导信号失真;窗口过晚(如 \([80, 84]\))虽然图像逼真使得 IR 很高,但全局空间布局与物体对应关系已经固化,GenEval 无法获益。中间段 \([40, 44]\) 是兼顾语义重塑与信噪比的唯一区间。
- 粒子多样性免受坍缩危害:在标准配置(\(K=4\))下,BoNG 生成多样性(0.319)甚至高于 Vanilla 基础采样(0.313),而 SMC 的多样性暴跌至 0.155;即便增加引导步数 \(K\),多样性下降也是平滑单调的,具有更宽广的安全操作窗口。
- 动态锚点漂移机制:监测发现最优粒子锚点在引导过渡步间的切换率(switch rate)均值为 0.2656,证明引导并非由某一特定幸运粒子全局绑架,而是粒子群交互中动态民主甄选的结果。
亮点与洞察¶
- 将 Best-of-N 从终点过滤重构为过程动力学:传统观念视 BoN 为不可导的黑盒后处理机制,本文通过粒子间非对称正向引力与得分协同,巧妙地将离散选优融入扩散逆向 SDE/ODE,展现了优雅的算法重塑能力。
- 数学上严谨对应奖励倾斜分布的极限:借助 Fisher 恒等式与自归一化重要性采样(SNIS),严谨论证了 BoNG 是奖励倾斜分布在 \(\lambda \to \infty\) 时的无展开渐近形态,打通了离散启发式算法与连续测度倾斜理论的桥梁。
- 正交性与即插即用泛化:BoNG 不仅能与流匹配架构(FLUX.1-dev 12B)无缝适配,还能作为提议分布与 SMC 粒子滤波级联组合(SMC + BoNG),双重增益证明了该引导机制的高度通用性。
局限与展望¶
- 受限于奖励模型的保真度:与所有强化学习及测试时对齐算法一致,若黑盒奖励模型存在对抗漏洞或过拟合盲区,激进的引导会放大 reward hacking。
- 单步 Tweedie 近似的流形偏差:在极高噪声尺度下单步高斯反演偏离真实数据流形,导致引导无法在更早的扩散时步生效。
- 软权重平滑松弛的探索空间:本文主实验基于硬截断极限(Hard BoN),附录理论推导的有限 \(\lambda\) 软重要性权重松弛版未来有望在连续奖励控制与多样性保真上提供更精细的调节能力。
相关工作与启发¶
- vs Vanilla BoN: 后者仅在采样完全结束后做单点挑选,群体无信息交互且多样本均值不提升;BoNG 在去噪过程中持续引导,计算成本几乎相同但显著推高了整体验收分布。
- vs SMC (FK-Steering): SMC 通过粒子复制和删除进行重采样,面临严重的粒子坍缩并破坏多样性;BoNG 保持所有粒子生命周期,施加光滑引导力,并可反哺作为 SMC 的转移核。
- vs LiDAR / Tilt: 前者依赖多步展望展开(rollout)来逼近中间步奖励预期,引入数倍于基础采样的计算延迟;BoNG 完全基于当前粒子的一步去噪闭式解,实现了免展开的高效样本引导。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次将 Best-of-N 的核心原则转化为逆向扩散中的无展开粒子间引导机制,理论推导严密。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 SD v1.5、SDXL、FLUX.1-dev 等多种扩散及流匹配模型,消融实验与理论映射分析极为透彻。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,数学符号定义规范,图表信息量大且阐释直击核心。
- 价值: ⭐⭐⭐⭐⭐ 极具落地价值,无需模型微调和反向梯度即可显著拉升生成质量,对文生图工业化生产服务有重大指导意义。