跳转至

Best-of-\(N\) Guidance for Test-time Diffusion Alignment

会议: NeurIPS 2026
arXiv: 2610.05108
OpenReview: https://openreview.net/forum?id=BvprWJNII9
代码: https://github.com/aailab-kaist/BoNG
领域: 图像生成
关键词: 扩散模型, 测试期对齐, Best-of-N, 引导采样, 人类偏好对齐

一句话总结

提出无需前向/后向展开（Rollout-free）的测试期对齐方法 BoNG，将 Best-of-\(N\) 原理内化到逆向扩散过程，利用粒子群在线推选的高奖赏 Tweedie 估计作为非对称引导信号，兼具卓越的单样本与多样本生成质量且仅增加极小推理延迟。

研究背景与动机

扩散模型在文本到图像生成任务中展现出强大的生成表现，但预训练模型采样生成的图像往往难以充分契合人类偏好与复杂的提示词约束。为了消除这一差距，基于奖赏模型的对齐技术受到了广泛关注。相较于参数微调（如 RL 微调或 DPO）带来的高昂训练开销与灾难性遗忘风险，测试期对齐（Test-time Alignment）以其即插即用、无需重新训练骨干网络的灵活性成为核心研究方向。其中，最经典且实用的基线是事后挑选的 Best-of-\(N\)（BoN）采样：独立生成 \(N\) 个样本并输出奖赏模型打分最高的单张图像。然而，Vanilla BoN 存在根本局限——奖赏信息仅在最终阶段作为筛选器使用，采样轨迹本身没有获得任何引导，导致平均样本质量无法提升，不仅需要巨大的采样数量 \(N\) 才能捕获稀有高奖赏样本，而且完全无法满足现实业务中需要同时输出多张高质量候选图的多输出（Multi-output）需求。

为在扩散轨迹中实时引入奖赏，序贯蒙特卡洛（SMC）通过中间时刻重采样粒子来提升对齐度，但极易引发严重的“粒子坍缩”（Particle Collapse），即多个粒子收敛至同质化的单一轨迹，丧失样本多样性；梯度引导方法（如 UG、DATE）依赖对奖赏与评分网络的反向传播，显存占用暴增近 3 倍且极易陷入奖赏黑客（Reward Hacking）；而免梯度的样本引导方法（如 Tilt、LiDAR）虽然避免了粒子坍缩，却依赖昂贵的多步前向或后向展开（Rollout）来估计预期未来奖赏，引入了 2 到 3 倍的前向推理延迟。如何在无需网络梯度反传、无需额外多步 Rollout 采样且不产生粒子坍缩的前提下，将奖赏引导高效融入逆向扩散动力学，是测试期扩散对齐的核心技术瓶颈。

针对上述矛盾，本文提出了一种全新的视角：Best-of-\(N\) 策略不应只作为事后选择，而可以转化为逆向扩散中的在线动态引导器。核心 idea：利用扩散粒子群自身的单步 Tweedie 清洗估计构建候选池，在线推选出最高奖赏的锚点粒子，并在逆向扩散中形成非对称粒子间正向引导，理论上严格证明其对应奖赏加权倾斜分布在奖赏强度无限大极限下的免展开自正则重要性采样，以极低计算代价同时实现极致的单样本巅峰指标与平均样本质量提升。

方法详解

整体框架

BoNG 针对由 \(N\) 个噪声粒子构成的粒子群开展协同逆向扩散。在指定的时间步窗口内，算法不需要像先验工作那样针对每个粒子跑数十步子扩散去预测最终图像，而是直接复用各粒子当前计算出的评分网络输出，通过解析形式的 Tweedie 公式获得单步清洗样本估计；随后在无梯度的黑盒奖赏模型下对所有粒子的清洗估计打分，选出瞬时最优的 BoN 锚点；最后将该锚点作为引导源，对其余非锚点粒子施加正向吸引引导，而锚点自身保持原本扩散轨迹。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["当前时间步 t 的噪声粒子群 {x_t^i}"] --> Alg1["Tweedie 单步清洗估计<br/>复用已算评分网络预测 x^0_i"]
    Alg1 --> Sel["黑盒奖赏评估与在线 BoN 锚点推选<br/>i* = argmax r(x^0_i, c)"]
    Sel --> Decoup["非对称粒子引导调度<br/>锚点粒子: 维持原生动力学无偏去噪"]
    Sel --> Decoup2["非对称粒子引导调度<br/>非锚点粒子: 受 BoN 评分与位置双重正向牵引"]
    Decoup --> NextStep["调度器迭代推进至时间步 t-1"]
    Decoup2 --> NextStep
    NextStep --> Out["生成图像群：单样本 (BoN) 与多样本 (Mean) 协同优化"]

关键设计

1. Tweedie 单步无 Rollout 预测与在线 BoN 锚点推选：摆脱繁重的前向子扩散开销

在先验的样本引导方法中，为了计算含噪样本 \(\mathbf{x}_t\) 的未来预期奖赏，必须从 \(t\) 步展开采样多条完整轨迹直至 \(t=0\)（Rollout），使得单步评分网络前向评估次数成倍增加。BoNG 发现，利用高斯扩散核的 Tweedie 恒等式，无需任何额外网络推理即可得到条件期望估计： $\(\hat{\mathbf{x}}_0^i = \mathbb{E}_{p_\theta(\mathbf{x}_0 \mid \mathbf{x}_t^i, \mathbf{c})}[\mathbf{x}_0] = \mathbf{x}_t^i + \sigma_t^2 \nabla_{\mathbf{x}_t^i} \log p_\theta(\mathbf{x}_t^i \mid \mathbf{c})\)$ 在每一个引导步，仅需调用一次 VAE 解码器和黑盒奖赏函数 \(r(\cdot, \mathbf{c})\)，即可从当前活跃的 \(N\) 个粒子中确定瞬时奖赏最高的锚点索引 \(i^* = \arg\max_{1 \le i \le N} r(\hat{\mathbf{x}}_0^i, \mathbf{c})\)。此举彻底消除了子扩散前向展开，使得额外计算成本严格收敛为 \(N \times K\) 次 VAE 解码与奖赏模型前向推理。

2. 非对称群体引导机制与梯度解耦：兼顾全局探索与局部牵引

确定锚点 \(\hat{\mathbf{x}}_0^{\mathrm{BoN}} = \hat{\mathbf{x}}_0^{i^*}\) 后，BoNG 构造目标分数为匹配该锚点： $\(\nabla_{\mathbf{x}_t^i} \log p_\theta^{\mathrm{BoN}}(\mathbf{x}_t^i \mid \mathbf{c}) = \nabla_{\mathbf{x}_t^i} \log p_\theta(\mathbf{x}_t^i \mid \mathbf{c}) + \frac{s}{\sigma_t^2} \left( \hat{\mathbf{x}}_0^{\mathrm{BoN}} - \hat{\mathbf{x}}_0^i \right)\)$ 对该式进行代数解耦，可以揭示其美妙的物理机制：对于最佳锚点粒子 \(i=i^*\)，引导项由于差值为零而自然消失，保证了最优候选粒子不受任何扰动、沿着原始数据流形稳定去噪；而对于任意非锚点粒子 \(i \neq i^*\)，其更新方向严格等价于两部分贡献的线性和： $\(\nabla_{\mathbf{x}_t^i} \log p_\theta^{\mathrm{BoN}}(\mathbf{x}_t^i \mid \mathbf{c}) = \nabla_{\mathbf{x}_t^{\mathrm{BoN}}} \log p_\theta(\mathbf{x}_t^{\mathrm{BoN}} \mid \mathbf{c}) + \frac{1}{\sigma_t^2}(\mathbf{x}_t^{\mathrm{BoN}} - \mathbf{x}_t^i)\)$ 前一项将非锚点粒子与当前最高奖赏粒子的扩散动力学对齐，后一项则施加朝向高奖赏粒子的正向引力。这一机制避免了 SMC 式的离散粒子复制（防止粒子坍缩），让低分粒子在保留各自差异化语义的同时平滑地被“拉”向高奖赏区域。

3. 自正则重要性采样在无限奖赏极限下的理论统一：建立严格数学支撑

许多测试期对齐算法以奖赏倾斜分布 \(p_\theta^r(\mathbf{x}_0 \mid \mathbf{c}) \propto p_\theta(\mathbf{x}_0 \mid \mathbf{c}) \exp(\lambda r(\mathbf{x}_0, \mathbf{c}))\) 为目标。本文证明，对于任意提议分布 \(q(\mathbf{x}_0)\)，倾斜后验均值 \(m_\lambda^r(\mathbf{x}_t) = \mathbb{E}_{p_\theta^r(\mathbf{x}_0 \mid \mathbf{x}_t, \mathbf{c})}[\mathbf{x}_0]\) 的自正则重要性采样（SNIS）估计权重为： $\(\alpha_i^{(\lambda)} = \operatorname{softmax}_i \left( \log p_\theta(\hat{\mathbf{x}}_0^i \mid \mathbf{x}_t, \mathbf{c}) - \log q(\hat{\mathbf{x}}_0^i) + \lambda r(\hat{\mathbf{x}}_0^i, \mathbf{c}) \right)\)$ 当对齐奖赏强度趋于极限 \(\lambda \to \infty\) 时，softmax 权重完全退化为 Dirac 测度，即聚焦于单一样本：\(\alpha_{i^*}^{(\lambda)} \to 1\)，其余权重全趋向于 0。此时 SNIS 后验均值严格收敛至 \(\hat{\mathbf{x}}_0^{\mathrm{BoN}}\)，使得基于期望未来奖赏的精确梯度 \(\nabla_{\mathbf{x}_t^i} r_t^\lambda(\mathbf{x}_t^i, \mathbf{c})\) 在大 \(\lambda\) 极限下与 BoNG 的引导项精确吻合。这证明了硬性 BoN 引导并非启发式 trick，而是奖赏倾斜扩散分布在极大化偏好下的严格极限实现。

损失函数 / 训练策略

BoNG 是纯粹的测试期推理算法，无需微调或损失函数更新。在执行推理时，关键策略包括： - 时间步引导窗口 \([t_{\mathrm{start}}, t_{\mathrm{end}}]\)：BoNG 仅在总去噪步数的中间一小段区间激活引导（如 100 步 DDPM 采样中激活 \([40, 44]\)，共 \(K=4\) 步；50 步 DDIM 中仅激活 \([21, 23]\)，共 \(K=2\) 步）。早期时间步噪声过大，Tweedie 估计与最终奖赏的相关性不足 0.6；晚期时间步图像宏观布局已定型，引导无法纠正构图错误；处于二者交界的中间窗口是高奖赏相关性与几何可塑性完美平衡的黄金区间。 - 噪声空间引导尺度 \(s\)：在实际噪声预测空间 \(\bm{\epsilon}^i\) 中，引导项离散化表示为 \(\bm{\epsilon}^i \leftarrow \bm{\epsilon}^i - \sqrt{\bar{\alpha}_t} s (\hat{\mathbf{x}}_0^{\mathrm{BoN}} - \hat{\mathbf{x}}_0^i) / \sigma_t\)，推荐设置 \(s \in [1, 6]\)，以零额外算力平滑控制奖赏与多样性的 Pareto 前沿。 - 正交可组合性：由于 BoNG 作用于分数转移核（Transition Kernel），而 SMC 作用于势函数权重，两者可直接叠加（SMC+BoNG），在不增加评分网络调用的情况下进一步突破对齐上限。

实验关键数据

主实验

在 GenEval 提示词基准上，采用 Stable Diffusion v1.5（0.9B）与 SDXL（2.6B）作为主干网络，以 ImageReward（IR）作为引导奖赏模型，评测单输出（BoN 最高分样本）与多输出（平均得分与生成成本）。

表 1 评估了不同样本预算 \(N \in \{2, 3, 4\}\) 下的单输出（Best-of-\(N\)）表现：

主干模型 & 采样器 方法 \(N=2\) IR \(N=2\) GenEval \(N=3\) IR \(N=3\) GenEval \(N=4\) IR \(N=4\) GenEval
SD v1.5 (DDIM 50步) Vanilla BoN 0.317 0.473 0.545 0.496 0.682 0.520
SD v1.5 (DDIM 50步) SMC 0.317 0.473 0.532 0.512 0.636 0.509
SD v1.5 (DDIM 50步) BoNG (本文) 0.407 0.480 0.659 0.507 0.793 0.535
SD v1.5 (DDPM 100步) Vanilla BoN 0.396 0.481 0.605 0.513 0.742 0.525
SD v1.5 (DDPM 100步) SMC 0.396 0.481 0.675 0.511 0.826 0.529
SD v1.5 (DDPM 100步) BoNG (本文) 0.511 0.491 0.749 0.521 0.894 0.555
SDXL (DDIM 50步) Vanilla BoN 0.899 0.578 1.060 0.600 1.140 0.613
SDXL (DDIM 50步) SMC 0.899 0.578 1.011 0.596 1.093 0.612
SDXL (DDIM 50步) BoNG (本文) 0.911 0.573 1.073 0.594 1.149 0.606
SDXL (DDPM 100步) Vanilla BoN 1.037 0.600 1.162 0.617 1.228 0.634
SDXL (DDPM 100步) SMC 1.037 0.600 1.191 0.623 1.264 0.631
SDXL (DDPM 100步) BoNG (本文) 1.059 0.597 1.194 0.618 1.259 0.648

表 2 评估了多输出场景下的平均指标与单张 A100 GPU 上的推理开销（生成 4 张图）：

方法 平均 IR (↑) 平均 CLIP (↑) 平均 GenEval (↑) 耗时 Time (s) (↓) 显存 Mem (GB) (↓)
Vanilla -0.001 0.271 0.426 7.07 8.90
Vanilla (250 steps) 0.003 0.272 0.430 17.42 8.90
Vanilla Best-4-of-5 0.180 0.275 0.453 8.83 9.97
UG (Bansal et al., 2024) 0.326 0.262 0.355 58.36 28.16
DATE (Na et al., 2025) 0.364 0.274 0.438 32.89 24.71
LiDAR (Kim et al., 2026) 0.384 0.278 0.478 13.41 8.90
BoNG (本文) 0.506 0.280 0.511 8.06 8.90

消融实验

表 3 验证了时间步引导窗口起始位置 \(t_{\mathrm{start}}\) 的影响（SD v1.5, DDPM 100步, \(N=4, K=4, s=6\)）：

窗口区间 \([t_{\mathrm{start}}, t_{\mathrm{end}}]\) BoN IR (↑) BoN CLIP (↑) BoN GenEval (↑) 说明
Vanilla (无引导) 0.713 0.284 0.532 基线基准
\([20, 24]\) (过早) 0.753 0.282 0.523 Tweedie 与最终奖赏相关性过低导致引导噪声大
\([40, 44]\) (推荐黄金窗口) 0.875 0.286 0.559 预测可靠且几何宏观布局尚未固化，双向增益最强
\([60, 64]\) (中偏晚) 0.849 0.287 0.526 宏观对象构成已定型，GenEval 指标回落
\([80, 84]\) (过晚) 0.866 0.287 0.528 仅能微调细微纹理，对结构对齐无明显改善

表 4 扫描了引导尺度 \(s\) 与窗口步数 \(K\) 对奖赏与多样性（CLIP 特征距离）的权衡：

引导步数 \(K\) / 尺度 \(s\) 平均 IR (↑) 多样性 Diversity (↑) 说明
Vanilla -0.001 0.313 无引导 baseline
\(K=4, s=1\) 0.019 0.319 轻度介入，多样性不降反升
\(K=4, s=2\) 0.224 0.315 奖赏大幅增长，多样性完全保持
\(K=4, s=6\) (主实验) 0.506 0.224 优异的多输出均值，未出现退化坍缩
\(K=20, s=1\) 0.576 0.216 延长窗口带来强奖赏但增加推理延迟
\(K=35, s=1\) 0.707 0.121 接近极端奖赏倾斜，多样性出现显著压缩

关键发现

  • 显著超越 SMC 且无粒子坍缩：在 36 项单样本跨模型/采样器对比中，BoNG 在 29 项（80.56%）中拔得头筹。SMC 在多输出下由于频繁复制粒子导致多样性骤降至 0.155，而 BoNG 依靠平滑引导保持了 0.224 的高多样性。
  • 极致的生成效率：BoNG 生成 4 张图仅耗时 8.06s，相比 Vanilla 仅增加 0.99s 开销（+14%），比当前 SOTA 样本引导方法 LiDAR 快 1.66 倍，比梯度反传方法 DATE 快 4 倍，同时峰值显存与 Vanilla 严格一致（8.90 GB），无计算图缓存开销。
  • 跨骨干与奖赏通用性：BoNG 在 12B 参数的最新 Flow Matching 架构 FLUX.1-dev 上无缝适配（IR 从 1.353 提升至 1.370），且可在 HPSv2、HPSv3、CLIP 及多目标线性加权奖赏下自适应工作。

亮点与洞察

  • 将离散筛选转化为连续流体牵引：传统 BoN 的最大缺陷在于把挑选留在最后一步，导致无效探索。BoNG 巧妙地将粒子群自身作为动态参考系，让当轮最好的粒子成为引力源，在不进行粒子强行复制（Resampling）的前提下，把整个粒子群向高奖赏盆地拉动。
  • 免子扩散展开（Rollout-free）的数学降维：前人为了估计未来奖赏梯度往往需要做子采样甚至蒸馏，而 BoNG 直接借用 Tweedie 解析期望作为一阶替代，并从理论上证明其等于大 \(\lambda\) 极限下的重要性采样，将复杂的动态规划降维为一次矩阵差值引导。
  • 模块正交性带来的即插即用：由于该引导机制修改的是分数项而非权重分配，它可以与现有的各类采样器（DDIM、DDPM、Euler）及重采样技术（SMC）无缝嵌套，无需修改底模权重即可带来免费提升。

局限与展望

  • 依赖预训练奖赏模型的保真度：若使用的奖赏模型本身存在偏差或对抗性脆弱点，过强的引导尺度 \(s\) 容易放大奖赏黑客效应，导致图像出现饱和过曝或非真实纹理。
  • 软性加权松弛机制（Finite-\(\lambda\) Relaxation）的计算拓展：论文在附录探讨了非极大化的 Soft-weighted 变体，但仍依赖前向核近似；未来可进一步探索在非无限大 \(\lambda\) 下的高效粒子概率加权。
  • 动态自适应窗口选择：当前引导窗口依靠预设步数区间，若能根据实时各粒子的奖赏方差或 Tweedie 置信度动态开启与终止引导，有望进一步压缩延迟并提升对齐精度。

相关工作与启发

  • vs Sequential Monte Carlo (SMC / FK-Steering)：SMC 通过势函数对粒子做多项式重采样，会导致同一母粒子克隆多次而杀死其他探索分支，造成多样性崩溃；BoNG 保持所有粒子生命周期，通过软性梯度向量牵引，在多样性与高奖赏间取得显著更优的 Pareto 前沿。
  • vs Gradient Guidance (Universal Guidance / DATE)：UG 和 DATE 需要穿透 UNet/DiT 做全网络反向传播以求导奖赏，耗时是原采样的 4 到 8 倍且显存爆炸；BoNG 纯属前向样本引导，免反传且显存恒定。
  • vs Sample-based Guidance (Tilt / LiDAR)：Tilt 依赖每个时间步做完整的多步后向展开，LiDAR 则需要独立前向生成 \(n\) 个前瞻样本并消耗额外评分网络评估；BoNG 原地复用正在采样的 \(N\) 个并行粒子，做到了完全零 Rollout。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [将经典 Best-of-N 从事后选择创新性重构为逆向扩散阶段的在线粒子间正向引导，理论证明严谨]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 SD1.5、SDXL 及 FLUX.1-dev，涵盖单输出、多输出、消融、多样性及多奖赏模型测试]
  • 写作质量: ⭐⭐⭐⭐⭐ [数学推导流畅，物理机制解释清晰，算法与图表结构严谨]
  • 价值: ⭐⭐⭐⭐⭐ [极具工程实用价值，为高并发在线图像生成服务提供了几乎零显存增加、极低耗时的测试期对齐新范式]

相关论文