RAM：扩散模型强化学习后训练提速50倍的第一性原理

扩散模型的预训练与后训练困境

扩散模型已成为连续表示生成的标准范式——除语言建模外，几乎所有领域都用它来生成连续数据。扩散模型（Diffusion Models）是一类基于马尔可夫链的生成模型，其核心思想是通过逐步向数据添加噪声（前向过程）再学习逆向去噪（反向过程）来捕捉数据分布。

扩散模型的理论根源可追溯至非平衡热力学和随机过程理论。2015年Sohl-Dickstein等人首次将扩散过程引入生成模型，但真正引发工业级应用的是2020年Ho等人提出的DDPM（Denoising Diffusion Probabilistic Models）。此后，Score-based模型（Song等，2021）统一了扩散模型与分数匹配的理论框架，揭示了扩散模型本质上是在学习数据分布的分数函数（score function）——即对数概率密度的梯度∇_x log p(x)。Flow Matching（Lipman等，2022）进一步将离散时间步推广至连续时间流形，以确定性ODE求解器替代随机采样，大幅提升了推理效率，Stable Diffusion 3、Flux等最新文生图模型均采用此框架。

自2020年DDPM提出以来，扩散模型在图像、音频、视频、分子设计等几乎所有连续数据领域成为事实标准。其预训练的可扩展性根源于目标的简洁性：只需最小化模型预测噪声与真实噪声之间的均方误差（score matching目标），这一目标对模型规模和数据规模均表现出良好的扩展律。

TU Munich博士生Andreas Bergmeister在其最新研究中指出了一个核心矛盾：扩散模型的预训练极其简单且可扩展，但基于强化学习的后训练（post-training）却复杂、繁琐且难以扩展。

这项名为 Reinforce Adjoint Matching（RAM） 的工作，试图回到第一性原理，探究究竟是什么让扩散模型预训练如此可扩展，再从中推导出能让强化学习后训练同样简单、可扩展的训练目标。

为什么后训练对扩散模型如此重要？作者给出了一个直观的分层：生成模型有两种学习范式——从样本学习（在数据集上匹配分布，即预训练）和从结果学习（生成后打分并据此调整模型，即后训练）。后训练在大语言模型领域已被RLHF（Reinforcement Learning from Human Feedback）验证为将模型与人类偏好对齐的关键技术，GPT-4、Claude等系统均依赖这一范式。

将类似思路迁移到扩散模型面临独特挑战：语言模型RLHF的成功在一定程度上源于其离散动作空间的天然优势——每步只需从有限词表中选择token，轨迹较短，策略梯度估计方差相对可控。扩散模型的对齐则面临连续高维动作空间、超长去噪轨迹（20-1000步）、梯度在时间维度指数衰减等多重挑战。早期工作如DDPO、DPOK等通过截断反向传播来规避内存问题，但引入了严重偏差。扩散模型的生成过程本质上是随机微分方程（SDE）或常微分方程（ODE）的求解过程，涉及数十至数百个去噪步骤，标准策略梯度方法需要对整条生成轨迹进行反向传播或Monte Carlo估计，计算开销极大。在机器人等领域，扩散模型被用来生成执行策略，而收集模仿学习数据集代价高昂，因此"从结果学习"显得尤为关键。

从最优分布到值函数梯度

后训练的目标可以形式化为：在给定预训练扩散模型和奖励函数的前提下，找到一个能**最大化期望奖励、同时保持接近预训练模型（最小化KL散度）**的分布。KL散度（Kullback-Leibler Divergence）是衡量两个概率分布差异的信息论指标。在后训练的优化目标中，这一KL约束至关重要——它保证模型停留在参考分布的支撑集内，避免漂移到生成不合理样本的区域，也是防止奖励黑客（reward hacking）的主要机制。

这一优化问题的最优解是经典结论：参考分布的指数倾斜（exponential tilting），即最优分布的密度正比于参考模型密度乘以奖励的指数：π*(x) ∝ π_ref(x) · exp(r(x)/β)，其中β是温度参数控制正则化强度。这一结论具有深刻的跨领域普遍性：在物理学中它对应Gibbs-Boltzmann分布，描述热平衡系统中粒子能量的统计分布；在信息论中对应最大熵原理的约束最优解；在金融数学中对应风险中性测度的Esscher变换；在随机最优控制（SOC）中则与著名的Schrödinger Bridge问题等价——最优控制问题等价于最小化从参考过程到控制过程的相对熵。这一结构的普遍性正是RAM能够在扩散模型上建立严格理论保证的根本原因，也解释了为何KL正则化不仅是工程上的防崩措施，更是优化目标的数学核心。

作者沿用Flow Matching框架的思路，为参考分布和最优分布分别定义前向加噪过程，通过score-velocity关系（Score函数即对数概率密度的梯度∇_x log p(x)与速度场之间存在确定的代数关系），将最优过程与参考过程的速度场相减，经代数推导后，两者速度场之差恰好等于值函数的梯度（adjoint，即伴随量）。这意味着：只要能估计值函数的梯度，便可通过回归目标训练模型。

原始Adjoint Matching方法将其转化为on-policy的回归目标，从当前策略估计值函数再回归其梯度。难点在于——如何在奖励不可微的前提下，廉价而稳定地估计这个adjoint。

核心创新：丢弃路径成本

值函数由两部分构成：非可微的奖励项和可微的路径成本项（path cost）。作者利用期望的可加性将两者拆分独立估计。

对奖励项，作者使用REINFORCE/对数导数恒等式——这一经典技巧允许在奖励函数不可微时（例如人类评分、规则判断、黑盒渲染器）仍然通过加权期望形式估计梯度：∇θ E{xp_θ}[f(x)] = E_{xp_θ}[f(x) · ∇_θ log p_θ(x)]。将期望奖励的梯度转化为奖励乘以桥梁分数（bridge score）的期望，计算"给定噪声样本条件下干净样本"的分数时，通过贝叶斯定理交换条件方向，最终可用模型的速度场直接表达。

真正让RAM具备可扩展性的关键决策是：直接丢弃路径成本项。作者坦言项目耗时较长，正是因为团队尝试了各种精确和近似方法估计路径成本，但没有一种能真正扩展。

为什么路径成本如此棘手？这一问题可从Girsanov定理的视角精确量化。Girsanov定理描述了改变漂移项（drift）后SDE的测度变换，其Radon-Nikodym导数恰好包含两个漂移之差的L2范数积分——即路径成本。在d维空间中，即使每个维度的速度场偏差仅为ε，累积路径成本的期望为O(T·d·ε²)，方差为O(T·d·ε⁴)。路径成本（path cost）在随机最优控制理论中衡量从参考过程到当前过程的轨迹偏离代价，数学上等于两个速度场之差的L2范数沿整条轨迹的积分。在文生图等高维场景中，Stable Diffusion 3.5的输出高达约65000维（潜在空间16×64×64≈65536维）。即使模型只有微小偏离，沿20-50步轨迹积分后路径成本也会累积到奖励信号量级的数千倍，产生灾难性的方差爆炸。参考模型与当前模型速度场之差的路径成本极其巨大，这个高方差的巨大标量从奖励中减去后，奖励信号便几乎被淹没——如同获得一张图的奖励是100，下一张却是-200，学习信号完全丧失方向感。这与强化学习中"信用分配问题"和"维度诅咒"有深层联系，也从根本上解释了为何历史上所有尝试精确估计路径成本的方法都在高维场景下折戟沉沙。

值得强调的是，丢弃路径成本并不等于丢弃整个KL正则化。RAM的不动点（fixed point）本身就是KL正则化最优解的推论，只对指数倾斜分布成立。相比之下，Diffusion NFT等基线方法缺少对参考模型的锚定，这也解释了为何RAM中奖励黑客（reward hacking）现象远不如其他基线严重——奖励黑客指模型找到投机取巧的方式来最大化代理奖励而非真正完成任务，例如生成看起来有文字但实际不可读的图像来欺骗OCR奖励模型。

训练流程与50倍加速

RAM的最终损失只依赖三个点：从当前模型采样的干净样本x0、高斯噪声样本x1，以及加噪样本xd——不需要完整的随机rollout。

具体流程为：用ODE（而非随机采样）从当前模型采样x0，评估奖励，然后像预训练一样线性插值构造加噪样本，最后执行回归。ODE采样（确定性求解器）在给定初始噪声的情况下产生完全确定的轨迹，是Flux、SD3等模型推理时的默认选择，也回避了Flow Matching随机调度器在噪声端扩散系数趋于无穷的数值不稳定问题。一旦获得昂贵的端点样本和奖励，便能廉价地构造多个训练目标——每个端点采样一次，可生成多个（作者使用8个）相互独立的加噪样本。

这里的独立性具有深刻的统计学意义：给定端点(x0, x1)，不同插值时刻配合独立高斯噪声向量的训练目标在统计上相互条件独立。从方差分析角度，n个独立样本使方差降低1/n倍，而强相关样本（如同一ODE轨迹上的不同时刻，一个点不过是另一个点的低噪声版本）的有效样本量接近1，无论采集多少步骤方差几乎不降低。RAM的8个独立样本理论上将梯度估计方差降低至1/8，在计算成本几乎不变的前提下提供了接近线性的样本效率提升。相同端点配合独立高斯噪声向量的多次线性插值，从信息论角度等价于对同一观测的多次独立重采样，提供了高质量的去相关学习信号。

这正是加速的根本来源。作者将RAM与代表策略梯度方法的 Flow-GRPO 进行对比：

• 策略梯度方法必须进行随机采样，而flow matching模型的随机调度器在噪声端会发散、系数趋于无穷，导致模拟困难且端点模糊。更关键的是，训练步只能在采样轨迹的固定点上计算，这些点高度相关（一个点不过是另一个点的低噪声版本），学习信号十分微弱。这与REINFORCE方法的高方差本质问题相通——DeepSeek-R1和GRPO通过组内相对奖励作为baseline大幅降低了语言模型场景的方差，但扩散模型连续动作空间中的去相关问题比离散语言模型场景更为棘手，因为连续空间中相邻时刻的状态之间存在确定性的代数关系，无法通过简单的基线减法消除相关性。
• RAM先用ODE采样端点，再向后跳跃构造任意数量的训练步。由于加噪时采用独立的高斯噪声向量，各样本相互独立、相关性低，提供了更丰富的学习信号。

实验结果令人瞩目：在GenEval（复杂组合提示，如"红色斑马"）、视觉文本渲染、人类偏好模型等基准上，RAM达到相同乃至更高奖励仅需1/50的GPU小时，且每个训练步比Flow-GRPO还略便宜。RAM在美学和人类偏好指标上同样表现更优，有效避免了"文字渲染正确但图像失真"的问题。

工程实现要点

在答疑环节，作者澄清了实现细节。算法实际需要三个模型：

1. 当前模型（正在优化的对象）
2. 参考模型（预训练模型，提供KL锚定）
3. 指数移动平均（EMA）模型（用于采样，提升训练稳定性）

EMA（指数移动平均）模型是一种常见的训练稳定化技术，用历史权重的指数加权平均替代当前瞬时权重进行推理，类似于优化中的momentum机制，可有效抑制训练过程中的参数震荡，在扩散模型训练中已被广泛验证有效。此外还有第四个EMA模型用于评估，但非必需。

由于后训练采用LoRA微调，保存多个模型副本成本很低。LoRA（Low-Rank Adaptation）通过在原始权重矩阵旁并联两个低秩矩阵（W = W_0 + BA，其中秩r远小于原始维度）来大幅减少可训练参数量——对于数十亿参数的扩散模型，全量微调需要数百GB显存，而LoRA只需几十MB额外参数，使得单机维护多套模型参数成为可能：只需一个基础模型加不同的LoRA权重。LoRA的低秩假设背后有理论支撑：语言模型和扩散模型在微调时的有效参数更新通常具有低内在维度，即改变模型行为所需的自由度远小于参数总量，这与神经网络损失景观的局部低维结构研究一致。作者还指出，虽然采样是off-policy的，但未使用重要性采样比率，因为对扩散模型而言这些比率难以可靠估计。

采样设置为：classifier-free guidance、20步采样、每个提示24个rollout，并将平均奖励作为baseline减去以降低方差。

对于路径成本的未来方向，作者认为与其纠结如何更稳定地估计它，不如从RAM不动点的其他视角出发。他提到可以用EMA或旧版当前模型替代参考模型并降低正则化强度，这与近期"tilt matching"工作中将奖励缩放、分多步倾斜以减小累积近似误差的思路相通，也让人联想到信赖域（trust region）方法的设计哲学——信赖域方法（如TRPO和PPO）通过在每次更新时约束策略变化幅度来保证学习稳定性，与RAM的KL正则化框架在本质上殊途同归：两者都是在"充分利用当前信息"与"保持对已知良好区域的信任"之间寻求动态平衡，只是前者通过显式约束实现，后者通过正则化项软约束实现。

小结

RAM的价值在于回归第一性原理：识别出高维空间中路径成本才是扩散模型强化学习难以扩展的根本原因——这一结论有Girsanov定理的严格数学支撑，不是经验性的启发式发现。在保留KL正则化本质（通过不动点条件而非显式惩罚项）的前提下大胆丢弃路径成本，配合ODE确定性采样消除数值不稳定性、去相关训练目标构造提升样本效率，最终实现了数量级的效率提升。这一看似激进的简化背后，是对理论结构的深刻理解：当路径成本的方差在高维空间中必然淹没奖励信号时，精确估计它反而是有害的，而丢弃它反而能让不动点条件精确成立。对于机器人控制、文生图等依赖"从结果学习"的领域，RAM提供了一条更简洁、更可扩展的扩散模型对齐路径，也为连接随机最优控制理论与实用深度学习后训练技术之间的鸿沟架设了一座新的桥梁。

核心要点

核心要点