通过奖励倾斜分布匹配增强少步生成器

TL;DR RTDMD 将奖励倾斜分布匹配分解为分布匹配与奖励最大化，通过两阶段训练让少步生成模型在仅 4 步推理下对齐人类偏好，在 SD3/FLUX 等模型上取得新 SOTA。

问题

问题背景

当前文生图领域，扩散模型蒸馏技术（如 DMD、LCM）已能在 4 步内生成高保真图像，但如何让这些极低步数的模型同时具备与人类审美、复杂文本组合性对齐的能力，成为继速度优化之后的关键需求。

现有方法局限

• 蒸馏方法与对齐目标分离：传统蒸馏（如 DMD）仅最小化与教师分布的距离，忽略奖励信号，导致生成结果虽快但未必符合人类偏好；而直接嫁接 RLHF 至扩散模型（如 DPO、D3PO）通常依赖全步采样，计算开销巨大，无法适配少步场景。
• 策略梯度估计的高方差：少步扩散模型（如 4 步）的中间状态随机性大，直接使用 REINFORCE 式的梯度估计会导致训练不稳定；若只对最终确定步回传奖励，又无法充分优化中间步的策略。
• 分布坍塌风险：单独最大化奖励会使生成偏离教师分布覆盖的高质量区域，产生“奖励黑客”现象，如图像细节破坏或模式坍塌。

为什么这个问题难/重要

• 技术挑战：蒸馏压缩轨迹后，模型可调整的自由度极低；奖励信号需在稀疏的几步中有效分配，且必须与分布保持平衡。该问题本质上是动态分布匹配与高方差策略优化的耦合，需要精巧的分解与正则化设计。
• 业界关注度：在实时交互（如 AIGC 工具、虚拟试穿）中，4 步推理是工业落地的性能门槛；同时，美学与组合性评估（如 GenEval、HPSv2）直接影响产品体验，解决该问题可大幅提升商业模型的竞争力。

行业类比

如同在机器人操控中，让一只仅能执行 4 个动作的机械臂，不仅快速完成任务，还要使动作姿态严格符合人类审美标准，每一次动作分配都必须精准且优雅。

核心洞察

• 通过最小化与奖励倾斜教师分布的 KL 散度，RTDMD 自然地将分布匹配蒸馏与偏好强化学习统一为单一目标。这种分解避免了先蒸馏后微调的分离流程中常见的分布偏移与次优权衡，为少步生成器对齐提供了更原则性的优化路径。
• 混合策略梯度利用确定性最后一步进行直接奖励反向传播，显著降低信用分配难度和方差；中间随机步骤则采用 GRPO 估计，并引入 SubGRPO 进一步稳化训练。这一设计精确适配了少步生成器中每步决策的关键性，使 4 步推理即可达到 SOTA 偏好与美学指标。

方法

输入与整体框架

RTDMD 的输入为文本提示 $c$ 与初始噪声 $x_T$，通过一个预训练的少步流生成器（few‑step flow generator）生成图像。该方法以两阶段训练为核心：第一阶段借助分布匹配蒸馏建立快速生成能力，第二阶段通过奖励引导的强化学习将生成分布向人类偏好方向对齐。

第一阶段：AC‑DMD 分布匹配蒸馏

第一阶段提出 Ambient‑Consistent Distribution Matching Distillation (AC‑DMD)，解决标准分布匹配蒸馏（DMD）在少步生成下的两个缺陷——伪造分数模型（fake score model） 难以跟上生成器分布变化，以及跨步去噪不一致。

• 子区间匹配：将扩散轨迹划分为多个子区间，每个子区间内用单步生成器近似教师模型的多步去噪结果，从而降低单步映射的学习难度。
• 伪造分数模型与一致性正则化：训练一个伪造分数模型去逼近真实分数，同时加入一致性正则化损失，约束相邻时间步的分数预测满足一致性关系，使其在有限更新步数下仍能准确跟踪生成器分布的变化。

该阶段输出一个具有快速生成能力的生成器，但尚未针对人类偏好优化。

第二阶段：奖励倾斜分布匹配与混合策略梯度

第二阶段将训练目标形式化为最小化生成器分布与奖励倾斜教师分布（reward‑tilted teacher distribution）的 KL 散度。该 KL 散度可分解为分布匹配项和奖励最大化项之和，实现分布对齐与奖励优化的统一。

• 分布匹配项沿用第一阶段 AC‑DMD 的目标，保持生成质量。
• 奖励最大化项采用创新的混合策略梯度：
  1. 中间随机步骤：对随机采样的中间状态，使用 GRPO（Group Relative Policy Gradient，基于组内相对优势的梯度估计）计算梯度，有效抑制高方差。
  2. 最终确定性步骤：最后的去噪操作是确定性的，因此直接对奖励模型做直接反向传播，获得低方差梯度。
  3. SubGRPO 方差削减：进一步提出 step‑subset GRPO，仅在部分关键步骤子集上应用 GRPO，避免对所有步骤采样带来的额外噪声，显著降低梯度估计方差。

输出与效果

经过两阶段训练，生成器仅需 4 步推理即可输出图像，在偏好得分、美学指标、组合性指标上均达到新 SOTA，超越以往少步文生图方法。

与同类方法的差异：不同于将分布匹配与强化学习孤立使用的方案，RTDMD 从奖励倾斜分布理论推导出统一目标，并针对少步生成器的多步随机性特点设计混合策略梯度，首次在高效率推理条件下实现稳定的强化微调。

实验

实验设计

• 基础模型：基于 SD3、SD3.5 和 FLUX.2 的 flow matching 生成器进行蒸馏。
• 两阶段流程：
  1. AC-DMD 第一阶段：通过子区间分布匹配和一致性正则化训练假评分模型，快速将多步教师模型压缩为少步学生模型。
  2. RTDMD 第二阶段：在蒸馏损失基础上添加奖励最大化的策略梯度项，联合优化分布匹配和人类偏好对齐。
• RL 细节：采用混合策略梯度，对随机中间步使用 GRPO 风格估计，对确定性最后一步直接回传奖励；引入SubGRPO 仅对某几步子集估计优势，降低方差。
• 评估基准：覆盖偏好（HPSv2、ImageReward）、美学（LAION-Aesthetics）、组合性（T2I-CompBench、GenEval）多个维度。

关键发现

• 4 步推理即 SOTA：在 SD3、SD3.5、FLUX.2 上，RTDMD 仅需 4 个 NFE 便在偏好、美学、组合性指标上全面超越之前少步蒸馏方法（如 DMD2、LCM 等）。
• 分布匹配与 RL 的协同：消融验证移除第二阶段的奖励最大化或第一阶段的 AC-DMD 均导致性能下降，表明奖励倾斜分布 分解为 DM + RL 是有效的，且 AC-DMD 提供的高质量初始策略对 RL 训练至关重要。
• SubGRPO 降方差：相较于全步 GRPO，SubGRPO 在保持改进幅度的同时显著降低训练波动，提升样本效率。

与基线对比

• 对比 DMD2：纯分布匹配方法未显式优化人类偏好，RTDMD 通过显式最大化奖励信号大幅提升偏好/美学得分，且不牺牲组合性。
• 对比直接 DPO/RLOO：在扩散时序上直接应用策略梯度存在方差高、效率低的问题；RTDMD 的混合梯度设计与 SubGRPO 针对性解决了少步生成器中间步的信用分配，性能更优且稳定。
• 该框架为后续少步生成模型对齐提供了一种模块化、理论驱动的路径：AC-DMD 负责高效压缩，RTDMD 负责偏好注入。

行业影响

落地场景

RTDMD 将少步扩散模型的对齐能力提升到新水平，尤其适合需要低延迟、高吞吐、高审美的图像生成服务。典型落地场景包括：

• 内容创作平台：辅助创作者在 4 步内生成符合复杂文本描述的高质量图像，用于社交媒体配图、数字艺术等；
• 电商产品图生成：根据商品文本描述快速合成多角度、高细节的产品展示图，支持实时编辑与批量生成；
• 游戏与影视概念设计：为美术团队提供迭代速度极快的视觉探索工具，降低预可视化成本。

商业价值

• 降本：将老师模型（如 SD3、FLUX）数十步的推理压缩至 4 步，同等硬件下吞吐量提升 5-10 倍，大幅降低云部署的 GPU 小时成本；
• 增收：对齐人类偏好后，用户对生成图像的接受率更高，可直接提升 AIGC 产品的付费转化与留存；
• 体验提升：AC-DMD 阶段保证保真度，RL 阶段强化审美与指令遵循，让最终输出在 ImageReward、PickScore 等偏好指标上达到 SOTA，减少用户反复抽卡的等待时间。

与现有工作流的接口

RTDMD 可直接嵌入主流扩散模型的蒸馏 + 微调流水线：

• 第一阶段（AC-DMD）：在原有 DMD 框架上引入子区间分布匹配与一致性正则，仅需少量更新步数即可得到稳定的少步生成器，可与现有 SD3/FLUX 蒸馏代码库无缝衔接；
• 第二阶段（RL 微调）：利用预训练奖励模型（如 ImageReward）计算奖励，通过 混合策略梯度（GRPO + 确定性最终步反向传播）进行在线或离线微调，对模型架构无侵入，可直接替换推理用 checkpoint；
• 推理侧零改动：蒸馏后模型保持原始架构，服务框架无需适配，只需将训练好的权重部署至 vLLM、TGI 等推理引擎，或直接集成于 ComfyUI / Diffusers 管线。

具体落地案例

1. 全球电商平台：为卖家提供“一句话生成商品主图”功能，4 步推理可在 1 秒内返回高分辨率结果，允许买家实时调整风格（如场景、色调），提升成交率。RTDMD 的高偏好得分确保生成图符合平台质量规范，避免人工审核退回。
2. 社交媒体内容平台：在创作者工具中内嵌基于 FLUX 的快速生图功能，每天支撑百万次请求的生成服务；通过混合策略梯度持续优化生成结果与用户点赞/分享指标的相关性，用 RL 主动适应社区审美变迁，保持功能活跃度。

局限

• **两阶段训练流程复杂度高**：RTDMD 需先完成 AC-DMD 的第一阶段蒸馏，再执行第二阶段的联合优化，涉及多个损失项和超参数，训练开销远大于单一蒸馏或 RL 方法。第二阶段中分布匹配损失与奖励最大化目标的平衡系数需要精细调整，若设置不当可能导致生成质量与偏好对齐之间的折衷不佳。
• **奖励模型局限性**：方法依赖预训练的多个奖励模型（如 PickScore、HPSv2、Aesthetic Predictor）来引导生成，但这些模型本身可能存在偏差，无法全面刻画人类偏好，且容易受到奖励黑客攻击。当奖励函数与真实偏好不一致时，第二阶段 RL 微调可能使生成分布偏离真实数据分布，引入伪影或模式坍塌。
• **实验范围有限**：论文仅在 4 步推理设置下验证了 SD3、SD3.5 和 FLUX.2 模型，尚未探索更极端的少步场景（如 2 步）或更多步数下的性能趋势。同时，方法对教师模型质量的依赖较强，在教师性能较差时的鲁棒性未被评估，且未在其他非扩散架构上验证通用性。