🤖 machine learning

Reinforce Adjoint Matching: Scaling RL Post-Training of Diffusion and Flow-Matching Models

本文介绍了强化伴随匹配（RAM），这是一种用于扩散模型和流匹配模型的可扩展强化学习后训练方法，它通过推导一种简单的一致性损失来修正预训练目标，从而实现对奖励的更优对齐，且无需昂贵的随机微分方程轨迹 rollout、反向伴随扫描或奖励梯度。

原作者： Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, Jakiw Pidstrigach

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你有一位大师级艺术家，他花了数年时间通过临摹成千上万张照片来学习绘画。这位艺术家在重现所见之物方面技艺超群，但他们未必懂得如何遵循具体而棘手的指令，比如“在蓝色卡车旁画一只红色的斑马”，或者“在盾牌上清晰地写下‘HELLO'这个词”。

为了教会这位艺术家这些新技能，我们通常使用一种称为**强化学习（RL）**的方法。这就像一位严厉的艺术评论家，审视画作，给出评分，并告诉艺术家：“下次要做得更好。”

旧方法：昂贵且缓慢的评论家

过去，用评论家来训练这些 AI 艺术家，就像试图通过把某人扔进大海并观察他们挣扎来教他们游泳。

过程：AI 会生成一幅完整的图像（即“游泳”），评论家对其进行评分，然后系统试图弄清楚究竟是哪一笔导致了高分或低分。
问题：为了做到这一点，AI 必须从头到尾模拟整个绘画过程，一遍又一遍，仅仅为了获得一次反馈。这种方法既缓慢，计算成本又高，而且往往不稳定（就像试图在船沉没时计算风速）。

新方法：RAM（强化伴随匹配）

这篇论文的作者是 Andreas Bergmeister 及其团队，他们意识到有一种更聪明的方法。他们发现了一个“捷径”，能够像最初的预训练阶段一样保持训练的快速与简洁。

以下是RAM的工作原理，使用一个简单的类比：

1. “干净终点”技巧
想象绘画过程就像一部倒放的电影。电影从空白画布（噪声）开始，以一幅完成的画作结束。

旧方法：AI 必须观看整部电影，对结局进行评分，然后逐帧倒放电影，以查看哪里出了问题。
RAM 方法：作者们意识到，如果你知道最终的画作（即“干净终点”），你就不需要观看整部电影来理解过程。你只需拿那幅完成的画作，在数学上瞬间对其“添加噪声”，生成一个混乱的版本，就像原始训练所做的那样。

2. “一次完成”的反馈
RAM 不再模拟整个绘画过程来获取反馈，而是这样做：

生成：AI 画出一幅完成的图像（终点）。
评分：评论家给出评分（例如：“文字部分做得很棒！”）。
瞬间倒带：与其模拟倒带，数学方法会瞬间创建该图像的“噪声”版本。
学习：AI 学习如何将那个特定的噪声版本重新转化为高分图像。

神奇的洞察：
论文证明，即使你试图让图像变得更好，“如何给图像添加噪声”的规则也不会改变。唯一改变的是 AI 最初决定绘制哪些图像。由于“噪声规则”保持不变，AI 可以使用与原始训练期间相同的简单数学方法，只是目标不同。

为什么这很重要

速度：论文声称，RAM 比之前的方法快34 到 50 倍。它能在极短的时间内达到相同的技能水平。
简洁性：它不需要复杂且不稳定的计算（如"SDE 展开”或“反向伴随扫描”）。它只是一个简单的回归任务，类似于原始训练。
质量：当他们在Stable Diffusion 3.5M上测试时，AI 在以下方面有了显著提升：
- 组合性：将物体放置在正确的位置（例如，卡车在冰箱的左侧）。
- 文本渲染：在图像中清晰地书写单词。
- 人类偏好：生成人类真正喜欢观看的图像。

总结

将RAM想象成给艺术家一把“魔法橡皮擦”和一张“魔法评分卡”。与其强迫艺术家重画整幅画布以查看哪里出了问题，魔法橡皮擦会瞬间向他们展示他们最佳作品的混乱版本，而评分卡则确切地告诉他们如何修复那个特定的混乱。这使得 AI 能够在没有先前方法那种高昂计算成本的情况下，学习复杂的新技能。

技术摘要：强化伴随匹配（RAM）

问题陈述

扩散模型和流匹配模型之所以能够实现高保真生成，是因为其预训练依赖于监督回归：对干净样本进行解析加噪，并使模型回归至一个闭式目标。然而，通过强化学习（RL）后训练将这些模型与奖励函数（例如用于组合正确性、文本渲染或人类偏好）对齐，会破坏这种高效的结构。

现有的 RL 后训练方法面临显著的 computational 和结构性障碍：

策略梯度方法：将去噪过程视为马尔可夫决策过程，需要随机的随机微分方程（SDE）轨迹展开，并优化变分下界。这引入了高方差，并需要大量的采样步骤。
随机最优控制（伴随）方法：精确针对 KL 正则化的奖励最大化目标，但需要沿存储的 SDE 轨迹进行反向伴随扫描，并将奖励梯度通过采样过程反向传播。这在数值上不稳定且计算昂贵，特别是对于流匹配模型，其噪声调度在接近噪声端时会出现发散。
代理损失：最近的方法用代理（如 ELBO 或对比端点）替换不可处理的似然函数，牺牲了使预训练可扩展的干净回归结构。

核心挑战在于推导一种 RL 后训练目标，既要保留预训练回归框架的简单性和可扩展性，又要能有效优化奖励。

方法论：强化伴随匹配（RAM）

作者提出了强化伴随匹配（RAM），该方法将预训练回归结构扩展至 RL 后训练，且无需 SDE 轨迹展开、反向伴随扫描或奖励梯度。

理论基础

该方法基于表述为随机最优控制问题的KL 正则化奖励最大化。

最优过程结构：在 KL 正则化下，最优生成过程将干净端点（ $X_0$ ）的分布向高奖励样本倾斜（ $p^*_0 \propto p_{ref}^0 \exp(r)$ ），但保持给定干净端点的噪声状态的条件律不变。具体而言，从固定 $X_0$ 到噪声状态 $X_t$ 的轨迹仍保持预训练中使用的相同解析高斯核。
伴随匹配条件：最优控制 $u_t$ 满足一个不动点条件，即控制等于价值函数（伴随）的负缩放梯度： $u_t = -\sigma_t \nabla_x V_t(x)$ 。
REINFORCE 分解：为了避免计算 $\nabla r$ （其可能不可微），作者将价值函数应用对数导数（REINFORCE）恒等式。这将伴随分解为加权了反向桥接分数的奖励项，以及路径成本修正项。

RAM 目标

RAM 通过忽略路径成本修正（基于低方差和大规模稳定性的合理性）并利用加噪核保持解析这一结构洞察，来近似伴随。

端点采样：RAM 不使用受控 SDE 进行模拟，而是利用任何现成的 ODE 采样器，从当前模型的分布中抽取干净端点 $X_0$ 。
解析加噪：对采样的 $X_0$ 进行解析加噪以生成训练状态 $X_t = (1-t)X_0 + t\epsilon$ ，这与预训练完全一致。
闭式目标：通过将奖励项和贝叶斯桥接分数（对于线性加噪具有闭式表达式）代入伴随条件，推导出速度场 $v_\theta$ 的回归目标。

由此产生的损失函数为：
$L_{RAM}(\theta) = \mathbb{E}_t \left[ \left\| v_\theta^t(X_t) - \text{sg}\left( v_{ref}^t(X_t) + r(X_0) \left( (\epsilon - X_0) - v_\theta^t(X_t) \right) \right) \right\|^2 \right]$
其中 $X_0 \sim p_\theta^0$ ， $\epsilon \sim \mathcal{N}(0, I)$ ，且 $\text{sg}(\cdot)$ 表示停止梯度算子。

关键效率机制

无需 SDE 轨迹展开：训练状态通过对单个端点 $X_0$ 使用独立噪声向量加噪 $K$ 次来生成。这为每次模型前向传播提供了 $K$ 个条件独立的训练样本，相比依赖单条 SDE 轨迹的方法，每一步提供了更多的梯度信号。
无需奖励梯度：该方法适用于黑盒奖励。
无需反向扫描：加噪核的解析性质消除了对反向伴随积分的需求。

主要贡献

结构扩展：论文证明，在 KL 正则化下，预训练的回归结构可扩展至 RL 后训练。最优过程仅改变端点分布，而不改变条件加噪律。
RAM 算法：一种新颖的一致性损失，利用闭式贝叶斯桥接分数用奖励信息修正预训练目标，避免了现有方法的计算瓶颈。
理论依据：推导表明 RAM 的不动点对应 KL 正则化最优解，对于线性奖励下的高斯参考保证了精确性，并在一般情况下具有一阶一致性。
方差分析：对路径成本估计器的分析表明，忽略路径成本修正（如 RAM 所做）相比全视界或局部伴随估计器显著降低了方差，使其成为图像级后训练中唯一可行的估计器。

实验结果

作者在Stable Diffusion 3.5 Medium (SD3.5M) 上评估了 RAM，涵盖三项任务：

组合图像生成（GenEval）：衡量物体、属性和空间关系的正确性。
视觉文本渲染（OCR）：衡量生成图像中文本的可读性。
人类偏好对齐（PickScore）：与人类偏好对齐。

性能表现：

奖励最大化：RAM 在所有三项任务中均取得了最高的训练奖励，优于包括 Flow-GRPO、DiffusionNFT 和 AWM 在内的基线方法。
训练效率：RAM 达到 Flow-GRPO 的峰值奖励所需的步数更少：GenEval 任务少50 倍，OCR 任务少48 倍，PickScore 任务少34 倍。
图像质量：与一些因“奖励黑客”（为任务指标而降低美学质量）而受损的基线不同，RAM 在最大化任务奖励的同时，保持或提升了图像质量指标（Aesthetic、DeQA、ImageReward、HPSv2）。
CFG 兼容性：RAM 在后训练后仍与 Classifier-Free Guidance (CFG) 兼容，而 DiffusionNFT 和 AWM 等基线在应用 CFG 时任务奖励显著下降。

意义与主张

论文声称，RAM 代表了扩散和流匹配模型 RL 后训练的范式转变。通过利用最优控制问题的特定结构属性，RAM 恢复了使预训练可扩展的“简单回归”范式。

作者断言：

RAM 消除了对昂贵的 SDE 轨迹展开和反向伴随扫描的需求，这些一直是将 RL 扩展至高分辨率图像生成的主要瓶颈。
该方法“简单且可扩展”，以更少的训练步数实现了最先进（SOTA）的结果，且每步计算成本与现有方法相当。
该方法成功地将模型与复杂奖励（组合性、文本、偏好）对齐，且未出现可见的质量下降或奖励黑客现象，证明了在满足正确理论条件时，预训练结构足以稳健地处理后训练对齐。