Gradient-Free Noise Optimization for Reward Alignment in Generative Models

技术摘要：ZeNO（零阶噪声优化）

问题陈述
现有的生成模型（如扩散模型和流模型）奖励对齐方法通常依赖于优化多步随机轨迹。虽然这些方法在上述场景中效果显著，但难以扩展到确定性或单步生成器（例如一致性模型、蒸馏扩散模型），因为这些生成器并未自然暴露相同的轨迹结构。一个自然的替代方案是直接优化输入噪声。然而，标准的基于梯度的噪声优化需要对生成器和奖励函数进行反向传播。这限制了其在可微分管道中的应用，并且当梯度反映的是虚假的局部敏感性而非有意义的语义改进时，该方法可能不可靠。此外，许多实际奖励函数（例如蛋白质可设计性、组合式文本 - 图像对齐）是不可微分的或黑盒的，这使得基于梯度方法变得不可行。

方法论：ZeNO
作者提出了ZeNO（零阶噪声优化），这是一个无梯度框架，将噪声优化表述为路径积分控制（PIC）问题。其核心思想是将输入噪声的演化视为一个可控随机过程，仅利用零阶奖励评估对其进行优化。

1. 路径积分控制公式化：
   该方法将噪声优化建模为受控随机微分方程（SDE）。其目标是将噪声轨迹引导至能产生高奖励终端状态的区域，同时对控制努力进行正则化。最优控制 $u^*(z, t)$ 以路径积分形式推导得出：
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   关键在于，该估计器仅依赖于终端成本评估（奖励），无需对生成器 $G_\theta$ 或奖励函数 $r$ 进行微分。

2. 参考动力学（Ornstein–Uhlenbeck）：
   为了确保更新后的噪声与预训练生成器兼容，ZeNO 采用 Ornstein–Uhlenbeck（OU）过程作为参考动力学。与方差爆炸的 SDE 不同，OU 过程保持了高斯平稳分布 $N(0, I)$，这与大多数预训练生成器的先验相匹配。这一选择将噪声更新与隐式针对奖励倾斜分布 $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$ 的朗之万动力学联系起来。

3. 线性化估计器与方差缩减：
   标准的路径积分估计器由于指数加权而遭受高方差困扰。作者通过减去平均奖励 $\bar{r}$ 作为基线，并在小奖励方差假设下线性化指数项，引入了线性化近似（命题 1）：
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   该估计器与自然进化策略中使用的得分函数估计器一致，并且在经验上优于指数变体。

4. 单步滚动时域：
   为了避免长时域 rollout 的计算成本，ZeNO 采用单步滚动时域近似（$H=1$）。在每次迭代中，从参考动力学中采样 $N$ 个扰动噪声候选，由奖励模型进行评估，并利用线性化控制估计来更新当前噪声状态。这允许通过增加粒子数（$N$）或迭代次数（$M$）来实现推理时的扩展，而无需修改生成器。

主要贡献

• 无梯度对齐： ZeNO 使得针对确定性和单步生成器的奖励对齐成为可能，无需对生成器或奖励模型进行反向传播。
• 黑盒兼容性： 该方法适用于不可微分和黑盒奖励函数，而在这些场景下，现有的基于梯度的噪声优化方法（如 ORIGEN、ReNO）会失效。
• 理论联系： 该框架建立了零阶噪声优化与针对奖励倾斜分布的朗之万采样之间的理论联系，为更新规则提供了原则性依据。
• 推理时扩展： ZeNO 引入了一种灵活的机制，通过权衡计算预算（通过 $N$ 和 $M$）来换取奖励性能的提升。

实验结果
作者在多种生成器（SDXL-Turbo、DMD2、LCM、DFGAN）和奖励函数（美学评分、PickScore、GenEval 和蛋白质可设计性）上评估了 ZeNO。

• 图像生成： ZeNO 在所有测试的生成器和奖励类型上均一致地提高了奖励分数。与基于梯度的基线（ORIGEN、$\Psi$-Sampler）和 Best-of-N 采样相比，它实现了具有竞争力或更优越的性能。值得注意的是，在 LCM 生成器上，基于梯度的方法性能下降至低于未控制基线，而 ZeNO 保持了提升，突显了零阶方法对局部梯度敏感性的鲁棒性。
• 不可微分奖励： 在 GenEval 任务（对象计数和定位）中，由于奖励不可微分，ZeNO 在计数任务上优于 Best-of-N，证明了迭代噪声优化比独立采样能从相同的计算预算中提取更多价值。
• 蛋白质结构生成： 应用于单步黎曼 MeanFlow（RMF）蛋白质骨架生成器时，ZeNO 显著提高了可设计性（通过自一致性 RMSD 衡量）。它仅使用单次流评估就达到了与 5 步 RMF 基线相当的可设计性水平，展示了其通过噪声优化补偿减少采样步骤的能力。
• 多样性： 使用 Vendi 分数的分析表明，ZeNO 在提高奖励的同时保持了样本多样性，避免了基于梯度方法中常见的坍缩到单一模式的现象。

意义与主张
本文主张 ZeNO 为将日益增长的高效、确定性和单步生成模型类与奖励函数对齐提供了一个实用且具有理论依据的解决方案。其主要意义在于能够在梯度不可用或不可靠的场景下运行，例如黑盒奖励模型或复杂的生物设计任务。通过将噪声优化表述为路径积分控制问题，ZeNO 提供了一个可扩展的、与生成器无关的推理时对齐框架，无需重新训练或修改底层生成模型。作者指出一个局限性：该方法在每个更新步骤需要多次奖励评估，如果奖励模型本身计算昂贵，这可能会带来计算负担。