一步生成模型的"既要又要"难题：如何让AI既听话又不丢本事

你有没有想过一个问题：现在很多AI画图模型已经能做到"一步成图"了，输入一句话，瞬间出图，不用像以前那样让模型反复"擦了画、画了擦"几十次。这听起来很爽，但也带来了一个新麻烦。

这些一步生成的模型训练好之后，怎么让它变得更"听话"，比如更懂得数清楚图里有几个苹果、更懂得把物体放在正确的位置上？这件事，远比想象中难。

**核心问题：一步生成模型没有"过程"可抓**

要理解这个难题，得先明白传统AI画图是怎么被"调教"的。

过去的扩散模型，画一张图要经过几十步"去噪"过程，从一团雪花噪点，一步步清晰成一张图。这个过程里有大量中间状态可以观察、可以计算梯度、可以用各种数学技巧去引导模型往"更好"的方向走。研究者可以用强化学习、偏好优化这些成熟工具，因为这些工具都依赖一件事：你得知道模型生成图片的"轨迹"是什么样的。

**去噪轨迹*：扩散模型从随机噪声逐步演变成清晰图像的中间步骤序列，每一步都可以被观察和调整。**

但一步生成模型根本没有这条轨迹。输入一个随机噪声和一句话，模型直接"唰"地吐出一张图，中间发生了什么是个黑箱。更麻烦的是，很多评判图片好坏的"裁判"本身也是黑箱，比如判断"图里是不是真的有三头牛"这种打分器，往往不可导，没法直接用梯度下降去优化。

这就好比你想教一个学生写好作文，但你既看不到他打草稿的过程，又只能拿到一个模糊的分数，连作文哪句话扣分了都不知道。传统的"改这句话、那句话"的精细化辅导方法完全用不上了。

在这篇论文之前，已经有几条路子在尝试解决这个问题。DI++、DI*、DIDR这些方法依赖扩散模型特有的结构，而且需要奖励梯度，也就是说打分器必须可导。FAV方法用了核密度估计和零阶梯度估计，这意味着每一步都要多次调用奖励函数去"猜"梯度方向，代价很大。DrPO则干脆放弃梯度，用排序的方式构造一个"偏好场"，但这种方法在面对复杂的组合性任务时表现有限。

研究团队面对的局面是：市面上的一步生成骨干模型，比如SANA Sprint，在GenEval这个组合性评测基准上只能拿到0.73分（满分1分），而这个分数的提升空间，恰恰卡在了"没有轨迹、没有梯度"这两个硬约束上。

**核心方法：让目标分布自己"学着走"**

这篇论文提出的方法叫Reward-Weighted Transport Distillation，简称RWTD。

**RWTD*：一种只需要生成样本和标量奖励评分、不依赖梯度和去噪轨迹的一步生成模型后训练方法，通过混合当前模型和参考模型的"奖励倾斜分布"构造训练目标。**

要理解RWTD的核心创新，得先知道一个经典概念：奖励倾斜（reward tilting）。

简单说，给定一个基础分布（比如模型原本生成图片的分布），你可以用一个奖励函数给每个样本"打分加权"，分数越高的样本权重越大，这样整个分布就会"倾斜"向高奖励区域。这个操作有个数学上很漂亮的性质：它恰好是在"追求高奖励"和"不要偏离原模型太远"这两个目标之间找到最优平衡点。

传统做法是，只拿参考模型（也就是训练前的原始模型）生成一批样本，做奖励倾斜，当作目标，让当前模型去模仿这个目标。这叫"离策略"（off-policy）倾斜，因为目标始终锚定在固定不变的参考模型上。

这里问题来了：如果参考模型本身在某些高难度任务上表现就很差，比如数数这件事原本就做得不好，那么无论你怎么对它做奖励倾斜，能挖出来的高分样本也是有限的。好比你想从一堆及格线边缘的试卷里挑出满分卷子，挑来挑去也就那么几份，因为大部分试卷本身就没写对。

RWTD的解决思路是：别只盯着参考模型，也看看当前模型自己训练过程中摸索出的高分样本。

具体做法是构造一个混合目标，把"对当前模型做奖励倾斜"和"对参考模型做奖励倾斜"这两部分按比例混在一起，比例由一个叫ρ的超参数控制。当前模型这部分能够捕捉训练过程中新发现的改进方向，参考模型这部分则像一个锚，防止模型跑偏太远，丢掉原本的多样性和稳定性。

**奖励倾斜分布*：给定一个基础概率分布，用奖励函数的指数加权对其重新归一化后得到的新分布，奖励越高的样本在新分布里占比越大。**

这个设计让我想起一个很具体的场景。假设你在学一门新技能，比如学做菜。如果你只盯着一本老菜谱反复练习（纯参考模型），进步会很慢，因为菜谱本身写的方法可能就不是最优的。但如果你完全只相信自己做菜摸索出的新手法（纯当前模型），很容易越练越偏，最后做出来的菜连基本的咸淡都把握不住，因为你没有一个靠谱的参照系去校验自己是不是走偏了。真正靠谱的做法是：一边参考经典菜谱打基础，一边把自己摸索出的新技巧融进去，两者混合着来。如果完全抛弃菜谱，你可能会在追求某个单一风味的路上越走越极端，最后菜难吃到没法吃，这正是论文里纯"在策略"（ρ=0）训练出现的情况，多样性会严重崩塌。

论文给出了具体数据支撑这个判断。当ρ=0，也就是完全依赖当前模型自己摸索时，样本多样性指标（用LPIPS衡量的平均两两感知距离）从基线的0.652骤降到0.414，掉了超过三分之一。而加入哪怕15%的参考模型混合（ρ=0.15），这个数字立刻回升到0.580，稳定在正常区间内。

**特征空间最优传输：把抽象目标变成具体的"搬运指令"**

光有一个混合目标还不够，接下来的问题是：怎么让当前生成的样本，真的朝这个目标分布靠近？

这里论文引入了最优传输（Optimal Transport，简称OT）的概念。

**最优传输*：一种数学框架，用于计算把一堆"货物"（概率质量）从一个分布搬运到另一个分布所需的最小代价，广泛用于衡量和实现分布之间的转化。**

直观理解最优传输，可以想象你有一堆沙子堆成某种形状（当前模型生成的样本分布），需要把它重新堆成另一种形状（目标分布）。最优传输要解决的问题是：怎么搬运这堆沙子，让总的搬运距离最短？每一铲沙子该从哪里挖、挖多少、搬到哪里去，都是最优传输要给出的答案。

论文用熵正则化的最优传输（通过经典的Sinkhorn算法求解）来计算一个"传输方案"，告诉每个当前模型生成的样本，应该朝目标分布里的哪些样本靠拢，靠拢多少比例。

但这里还有个细节问题：直接在像素空间里算距离，意义不大。两张图片哪怕内容相似，像素层面的数值差异可能很大；反过来两张像素很接近的图可能内容完全不同。这就像用两份菜的"重量差"去判断它们味道像不像，显然不靠谱。所以论文选择在一个冻结的视觉编码器（比如DINOv2或MAE）提取出的特征空间里计算距离，这样算出的"相似度"更贴近人类对图像内容和风格的感知。

**冻结视觉编码器*：一个参数不再更新、专门用来把图像转换成特征向量的神经网络，在RWTD中用于在语义层面衡量图像之间的相似度，而非直接比较像素。**

算出传输方案之后，每个当前样本会得到一个"目标落点"，这个落点是加权平均出来的重心（barycentric destination）。但论文没有让样本直接跳到这个重心位置，而是引入了一个步长参数η，让样本只朝目标方向走一部分距离，而不是一步到位。

这个设计的用意也很实际。因为实际训练中用的是小批量数据，计算出来的传输方案本身带有噪声，如果每次都让样本一步跳到位，训练过程会很不稳定，容易抖动甚至崩溃。保守地走一部分，相当于给训练过程加了个"安全带"。

**部分传输蒸馏*：RWTD中将源特征以步长η的比例朝目标重心移动，而不是一次性完全移动到目标位置的训练策略，用于降低噪声敏感性、稳定训练过程。**

最后一步是把这些传输目标"蒸馏"进模型里：让当前模型生成图片，提取特征，然后用回归损失（MSE）把特征拉向刚才算出的目标，整个过程中目标值是被detach的（停止梯度），只有生成器本身在被更新。这个循环会不断重复，每一轮都基于最新的当前模型重新计算目标，这就是论文所说的"定点回归"。

**理论分析：这个混合目标最终会稳定在哪里**

论文花了相当篇幅做理论分析，想搞清楚这个循环迭代的过程最终会收敛到什么样的分布。

分析结果显示，当ρ=1（纯粹依赖参考模型）时，整个迭代过程会收敛到标准的奖励倾斜参考分布，这跟传统方法是一致的，相当于理论上验证了"老方法是这个新框架的一个特例"。

当ρ=0（纯粹依赖当前模型）时，论文证明这种情况下根本没有唯一的不动点，分布会随着迭代不断朝奖励最大化的区域收缩集中，这也从理论上解释了为什么实验里ρ=0的多样性会崩掉。

另外论文还证明了一个挺有意思的结论：部分传输的步长η不会改变最终的收敛点集合，换句话说，无论你是保守地小步走还是激进地一步到位，理论上最终都会走到同一个地方，只是走的过程稳不稳的问题。这个结论给了保守步长策略一个理论背书：你可以放心地用小步长换稳定性，不用担心小步长会让模型最终收敛到一个"打了折扣"的次优结果。

**实验结果：数字说话**

理论讲完了，关键看实际效果怎么样。

在GenEval这个专门评测组合性理解能力（比如数数、颜色搭配、物体位置关系）的基准上，SANA Sprint 1.6B这个一步生成模型原始分数是0.73，经过RWTD训练后提升到了0.80，这个成绩甚至超过了很多需要50步生成的更重型多步模型，比如FLUX.1 Dev（0.66）和SD3.5-L（0.71）。相比之下，DrPO只能把分数推到0.75，FAV几乎没有提升，还是0.73。

提升最明显的是位置（Position）和颜色属性（Color attribution）这两个出了名难优化的子任务，Position这一项从0.54直接跳到0.73，提升了近20个百分点。这说明RWTD在处理那些需要精细空间推理和属性绑定的复杂指令时特别有效。

在另一个实验里，研究团队用HPSv2（一个基于人类偏好训练出的打分模型）作为训练奖励，然后在完全没有见过的PartiPrompts数据集上测试泛化能力。

**跨奖励泛化*：模型只针对某一种奖励信号做训练，但在评测时表现出对其他未训练过的评价指标同样有提升，这说明优化没有"投机取巧"，而是真正学到了更普适的生成能力。**

结果很能说明问题：RWTD是唯一一个同时提升了全部五项偏好指标（PickScore、HPSv2、CLIP、美学分数、ImageReward）以及域外GenEval分数的方法。而使用梯度优化的DRaFT虽然把HPSv2刷到了37.02的高分，但CLIP分数却从0.2753掉到了0.2573，域外GenEval分数更是从0.73骤降到0.62，掉了11个百分点。

这组对比数据揭示了一个值得警惕的现象：梯度方法能更狠地优化目标奖励本身，但代价是在其他没被直接监督的维度上出现明显的"跑偏"。论文配的图片对比里能看出来，DRaFT和FAV生成的图片出现了明显的风格化过拟合现象，比如把写实照片画成过度艺术化的插画风格，而RWTD生成的图片保持了更自然的写实感。论文里举了个很具体的例子，让模型生成"雪山和倒映在湖面的树"，DRaFT和FAV生成的倒影里会凭空多出一些不存在的树，几何上根本说不通，而RWTD生成的倒影是几何一致的。

这背后的原因值得琢磨一下。梯度方法是直接朝着让打分器给高分的方向硬推，打分器觉得"色彩浓郁、细节丰富"的图片分高，模型就会不断朝这个方向使劲，哪怕代价是牺牲现实感和逻辑一致性，这就是典型的奖励过优化（reward overoptimization）现象。而RWTD通过混合参考分布的方式，天然地给这种"钻空子"行为设了一道刹车。

**计算成本：好处不是白来的吗**

一个自然的疑问是：RWTD听起来需要额外算传输方案、算特征距离，会不会特别费算力？

论文给出了实打实的对比数据。在GenEval任务上，RWTD训练一次需要197个GPU小时，DrPO需要194个GPU小时，基本持平；而FAV需要431个GPU小时，是RWTD的2.2倍还多，这是因为FAV用的零阶梯度估计需要反复多次调用奖励函数去估算梯度方向。在HPSv2任务上，RWTD是121小时，DrPO是119小时，FAV是123小时，DRaFT因为要对奖励模型做梯度反传，开销最大，达到315小时。

这说明RWTD虽然引入了最优传输这个看起来"重"的数学工具，但由于视觉编码器是冻结的、不参与训练反传，Sinkhorn算法本身的计算开销在实际训练批量下并不算大，整体成本和现有方法处在同一量级。

**消融实验：每个超参数都在干什么**

论文还做了细致的消融实验，搞清楚每个关键超参数到底该怎么调。

关于奖励温度β，数值越大，奖励倾斜越"狠"，对HPSv2这类目标奖励的优化效果越强，但CLIP分数会逐渐下滑，PickScore在β=5时达到峰值后开始回落。这是一个典型的"优化力度"与"语义保真"之间的权衡，β太小，优化不够狠；太大，又容易把其他能力带偏。

关于参考质量比例ρ，前面已经提到ρ=0会导致多样性崩溃，论文进一步发现ρ=0.15能在PickScore、美学分数、ImageReward三项上拿到最佳表现，是个很平衡的选择点。

关于传输步长η，论文发现模型对这个参数相对不敏感，PickScore、CLIP、ImageReward在η=0.2时达到最优，而HPSv2和美学分数则在更大的η值时略好，但整体差异不大，这也呼应了前面理论分析里"η不改变收敛点"的结论。

论文还单独验证了最优传输这个设计本身是否真的必要，设置了三种变体对比：默认的重心法（Barycentric，用加权平均目标）、采样法（Sampled，从传输方案里按概率采样一个具体目标点）、以及完全去掉传输几何结构的奖励加权回归（RWR，独立采样目标粒子，只保留边际分布一致，但丢掉样本对样本的最优分配关系）。

结果显示，重心法和采样法表现非常接近，这符合直觉，因为重心本来就是采样分布的期望值。但RWR在美学分数上表现更好，却在CLIP和PickScore上明显落后于前两者。这说明RWTD真正起作用的关键不是"用加权平均这个技巧"，而是最优传输算出的那套样本对样本的精准分配关系，这套关系能让每个生成样本找到和它语义最接近的目标，而不是随机匹配到一个同样高分但内容风马牛不相及的目标。

写在后面

读完这篇论文，有一个细节让我反复琢磨：为什么"把当前模型也纳入目标分布"这件事，之前没人系统地做过。

仔细想想会发现，这个想法其实有点反常识。如果一个模型还没训练好，用它自己生成的样本去指导自己训练，听起来有点像"自己骗自己"的风险，万一当前模型本身就跑偏了呢？但论文的理论分析恰恰说明了，这种担心是可以被参考分布的锚定作用化解的，关键不在于要不要用当前模型，而在于怎么配比。这提醒我一个更普遍的道理：很多看似矛盾的设计选择，答案往往不是"二选一"，而是"按比例混合"，而这个比例本身是可以被严格建模和求解的，不是拍脑袋定的。

另一个让我意外的地方是梯度方法的"过拟合"现象表现得如此直接和可视化。通常我们讨论模型过拟合都是在看数字，但这篇论文给出的那几张倒影图片对比特别直观，湖面倒影里凭空多出树木这种几何错误，一眼就能看出模型在某种程度上已经"学歪了"，这比任何一个分数下降都更有说服力。

这篇论文也留了一些没彻底解决的问题。比如混合比例ρ目前还是靠人工网格搜索调出来的固定超参数，如果能让ρ本身随训练动态调整，是不是能进一步提升效果和效率？另外论文的理论分析都是在"理想化"的连续分布假设下做的，实际训练用的是有限的小批量样本，这中间的理论与实践落差有多大，论文自己也坦承没有给出严格保证。这些问题留给后来者去接棒了。

Q&A

Q1：RWTD是什么，和传统的奖励对齐方法有什么不同？

A：RWTD全称Reward-Weighted Transport Distillation，是一种只需要生成样本和标量奖励评分的一步生成模型后训练方法。和传统方法只对固定参考模型做奖励倾斜不同，RWTD会把当前模型和参考模型分别做奖励倾斜后按比例混合，作为训练目标，既能加速优化又能保留原有能力。

Q2：RWTD在实际测试中效果如何？

A：在GenEval组合性评测基准上，RWTD把SANA Sprint 1.6B这个一步生成模型的分数从0.73提升到0.80，超过了多个更重型的多步生成模型。在HPSv2偏好对齐实验中，RWTD是唯一一个同时提升全部五项偏好指标以及域外泛化分数的方法。

Q3：为什么不直接用梯度方法优化奖励，而要用最优传输这种复杂方式？

A：很多实际应用中的奖励函数（比如组合性评测打分器）本身不可导，没法直接用梯度下降优化。而且实验显示梯度方法容易出现奖励过优化，导致图片风格失真、域外泛化能力下降，RWTD通过最优传输和混合目标的方式能更好地平衡奖励提升和原有能力保留。