UDM-GRPO：面向均匀离散扩散模型的稳定高效分组相对策略优化

首次将GRPO引入均匀离散扩散，通过将最终干净样本作为动作并用前向过程重构轨迹，解决训练崩溃。

前置知识

均匀离散扩散模型（UDM）

UDM 是在离散状态空间 $S=[K]^D$ 上定义的生成模型，源分布 $p_0(x)\sim\text{Unif}([K])^D$ 是词表上的均匀分布。它通过 CTMC 概率速度 $u_t$ 迭代精化所有 token，可并行更新所有位置。

UDM 是本文的基座模型，其迭代精化与并行 token 更新机制直接决定了 GRPO 的动作和状态空间设计。

GRPO（分组相对策略优化）

GRPO 是 PPO 简化版：用组内归一化优势 $\hat{A}_i = (R_i-\mu)/\sigma$ 替代价值网络，配合 PPO 裁剪目标 $\min(r_t^i\hat{A}_i,\text{clip}(\cdot)\hat{A}_i)-\beta D_{KL}$。

GRPO 是本文强化学习核心，读者需先理解其优势归一化、PPO 裁剪与 KL 惩罚三要素，才能体会 UDM-GRPO 为何要重写动作和轨迹。

Flow-GRPO

Flow-GRPO 将 Flow Matching 的反向去噪过程构造成多步 MDP，每步去噪后的样本 $x_{t+1}$ 作为动作 $a_t$，首次把 GRPO 扩展到连续扩散模型并取得大幅收益。

Flow-GRPO 是本文的起点与对照基线。它的"中间样本即动作、轨迹即反向过程"设计，正是 UDM-GRPO 颠覆性修改的对象。

前向 vs 反向轨迹

前向轨迹 $X_{\text{pretrain}} = \{x_t\}$ 是把真实 $x_1$ 加噪得到 ($x_t \sim p_t(x\mid x_1)$)；反向轨迹 $X_{\text{backward}}$ 是模型自己从 $x_0$ 去噪到 $\hat{x}_1$，会因误差累积偏离前向分布。

这一差异正是 Insight II 的核心，本文用 FID 量化了反向轨迹的偏移，并据此提出用对 $\hat{x}_1$ 的前向扰动 $X_{\text{forward}}$ 替代反向轨迹。

研究动机

近年文本到图像生成取得了显著进展，但预训练模型往往在需要精确对齐人类偏好的任务（如组合生成、文字渲染）上表现欠佳。强化学习（特别是 GRPO）已经在 LLM 推理上证明有效，Flow-GRPO 也成功地把 GRPO 引入 Flow Matching 并取得大幅收益。然而，当研究者把 Flow-GRPO 的范式直接套用到新兴的均匀离散扩散模型（UDM）时，遇到了严重的训练不稳定问题：作者在 Figure 1 中展示，朴素版基线在 500 步后即发生灾难性崩溃——奖励剧烈震荡、KL 散度爆炸，模型在 GenEval 上的提升很快停滞甚至倒退。具体而言，这种不稳定来自两个本质性的失配：其一，"不准确的中间动作"，即早期去噪步的预测 $x_t^1$ 熵极高、噪声很大，把这些样本当作 PPO 动作等于让模型学习错误信号；其二，"反向轨迹的分布偏移"，即反向过程 $X_{\text{backward}}$ 会因误差累积而偏离预训练前向分布 $X_{\text{pretrain}}$，导致模型暴露在 OOD 状态上，策略更新进一步恶化。

本文的目标是本文的具体目标是提出 UDM-GRPO——第一个将 GRPO 稳定且高效地集成到均匀离散扩散模型中的 T2I 训练框架，并验证该方法能在 GenEval（组合生成）、PickScore（人类偏好）、OCR（文字渲染）三个差异巨大的下游任务上一致地、显著地提升基座模型 URSA 与 FUDOKI 的性能，同时保持低且有界的 KL 散度，避免 Flow-GRPO 在连续域中表现出的训练崩溃问题。具体而言，作者希望用尽量少的代码改动（基于已有 GRPO loss 即可）实现 SOTA，验证"动作换为 $\hat{x}_1$、轨迹换为前向过程"这一对核心洞察的普适性，并在 FUDOKI MLLM 上展示方法的可迁移性。

与已有工作不同的是，与 Flow-GRPO 在连续域中直接使用反向过程-中间样本的设计不同，本文独辟蹊径地提出两条截然不同的设计准则：把动作从"每步中间预测 $x_t^1$"换成"全局唯一的最终干净样本 $\hat{x}_1$"，把训练轨迹从"反向去噪 $X_{\text{backward}}$"换成"对 $\hat{x}_1$ 做前向扰动得到的 $X_{\text{forward}}$"。这一对互换设计既消除了中间高熵信号带来的误导，又通过让轨迹对齐预训练分布，从根本上抑制了 OOD 漂移——这是已有 Flow-GRPO、Mask-GRPO 等工作都没有触及的视角。

核心方法

UDM-GRPO 的整体思路是：先把 GRPO 的多步 MDP 重新表述成"全局-局部"解耦的形式——状态是任意时刻的带噪 token $\hat{x}_t$，动作统一为最终的干净图像 $\hat{x}_1$，策略则是 $p_\theta(\hat{x}_1 \mid \hat{x}_t, c)$；再用前向过程 $p_t(x \mid \hat{x}_1)$ 重新构造训练轨迹，避免反向过程带来的分布漂移。围绕这个核心，本文进一步提出 Reduced-Step 训练（只在前半段高噪声时间步中随机抽 3 个连续步）来缓解多步梯度分散，以及 CFG-Free 训练（在采样和优化中均不依赖 Classifier-Free Guidance）来降低计算开销与条件-无条件双重目标的耦合。直觉上：Flow-GRPO 的做法相当于"让模型在错误的棋盘上下棋"，而 UDM-GRPO 相当于"先在自我对弈中得到最终局面 $\hat{x}_1$，再用前向过程把局面回放到任意中间步，只针对这一条干净的'目标动作'做策略学习"。

本文有两个核心创新点。Insight I：把策略动作从反向过程每步的中间预测 $x_t^1$ 改为对所有 $t$ 一致的最终干净样本 $\hat{x}_1$，即 $a_t \equiv \hat{x}_1$、$\pi(a_t|s_t) = p_\theta(\hat{x}_1|\hat{x}_t, c)$，使奖励与动作完全对齐。Insight II：把训练轨迹从 $X_{back}$ 切换为 $X_{fwd}$——先用反向过程采样出 $\hat{x}_1$，再沿前向过程 $p_t(x|\hat{x}_1)$ 加噪得到 $\hat{x}_t$，FID 证明 $X_{fwd}$ 在所有时间步都更接近 $X_{pretrain}$。两条 Insight 协同把 MDP 改写为 $s_t = (\hat{x}_t, t, c)$，策略比 $r_t^i = p_\theta/p_{old}$，损失形式与 Flow-GRPO 相同但语义根本不同。

方法步骤详情

方法流程如 Algorithm 1 所示。第一步（采样与奖励）：CFG-Free 模式下对每个 prompt $c$ 采样 $G=8$ 条轨迹，每条用 10 步 Euler solver 跑出 $\{\hat{x}_1^i\}$，再用 PickScore/OCR 等奖励器计算奖励并组内归一化得到优势。第二步（Reduced-Step）：从候选步组中随机抽 3 个连续时间步 $(t_1, t_2, t_3)$，约束在前半段，集中梯度在高噪声阶段。第三步（前向扰动）：对每个 $\hat{x}_1^i$ 按 $p_{t_j}(x|\hat{x}_1^i)$ 加噪得到 $\hat{x}_t^i$。第四步（策略更新）：计算策略比 $r_t^i(\theta)$，代入 PPO 裁剪目标 $\min(r_t^i\hat{A}_i, \text{clip}(\cdot)\hat{A}_i) - \beta D_{KL}$ 累加损失，更新参数后令 $\theta_{old} \leftarrow \theta$ 进入下一轮。

技术新颖性

本文的技术新颖性体现在三个层面。第一个层面是 MDP 重新表述：把动作从"时间步相关的中间预测"变为"时间步无关的最终干净样本"，并把轨迹从"模型自身的反向采样"变为"对 $\hat{x}_1$ 的前向扰动"，这种解耦设计在已有 GRPO 工作中前所未有。第二个层面是优化效率：Reduced-Step 通过把优化集中在高噪声时间步显著加速收敛，CFG-Free 通过避免同时优化有条件-无条件双重目标减少约一半的前向计算，并使 KL 散度更低、训练更稳定。第三个层面是理论洞察：作者指出在新形式下，扩散模型的策略损失可视为 LLM 预训练损失的优势加权版本，揭示了 UDM-GRPO 与 LLM-RL 范式的深层联系，理论上更自然，且不存在采样器约束和高显存开销的问题。

**Figure 2**: Illustration of the three trajectories.

**Figure 4**: Entropy and FID visualization across timesteps.

**Figure 10**: Generated samples across successive training iterations.

实验结果

本文在 URSA（1.7B）和 FUDOKI（1.5B）上系统实验，核心发现如下。第一，组合生成大幅跃升：URSA 的 GenEval 整体从 0.69 提升到 0.96，其中 Counting 0.60→0.95、Position 0.28→0.97、Attribute Binding 0.49→0.85，超越 Flow-GRPO+SD3.5-M（0.95）、FLUX.1-Dev（0.66）等更大模型。第二，PickScore 从 20.46 提升到 23.81，超过 SD3.5-L（22.91）。第三，OCR 从 0.08 飙升至 0.57。第四，跨模型泛化性强：FUDOKI 上 PickScore 21.32→23.40、OCR 0.04→0.26、GenEval 0.76→0.86。第五，Table 3 消融验证每项设计：backward+$x_t^1$ 仅 0.84、backward+$\hat{x}_1$ 0.89、forward+$\hat{x}_1$ 0.94、最终 0.96。第六，Figure 6 显示最终方法 KL 散度低且有界，无崩溃现象。

**Table 1**: Comparison result on GenEval.

**Table 2**: Comparison results on GenEval, PickScore, and OCR.

**Table 3**: Ablation study on GenEval, PickScore, and OCR.

**Table 4**: Performance of FUDOKI with UDM-GRPO.

**Figure 5**: Qualitative Comparison on GenEval and PickScore.

**Figure 6**: Performance metrics and KL loss on GenEval, PickScore, OCR.

**Figure 7**: Qualitative comparison of ablation methods.

**Figure 9**: Visualization for different methods.

查看结构化数据

任务	指标	本文	基线	提升
GenEval (Compositional T2I)	Overall Accuracy	0.96	0.69 (URSA)	+0.27（绝对提升）/ +39%（相对提升）
GenEval (Compositional T2I)	Counting 子项	0.95	0.60 (URSA)	+0.35
GenEval (Compositional T2I)	Position 子项	0.97	0.28 (URSA)	+0.69
Human Preference Alignment	PickScore	23.81	20.46 (URSA w/o CFG)	+3.35；并超过 SD3.5-L 22.91 / FLUX.1-Dev 22.84
Visual Text Rendering	OCR Accuracy	0.57	0.08 (URSA)	+0.49（接近 7.1 倍提升）
GenEval (FUDOKI MLLM)	Overall Accuracy	0.86	0.76 (FUDOKI)	+0.10

局限与改进

本文的方法在 GenEval、PickScore、OCR 等通用基准上取得了 SOTA，但仍存在以下局限。第一，在 OCR 任务上，0.57 的绝对分数虽然相对基线提升巨大，但仍低于 SD3.5-L 的 0.68，说明文字渲染对基座模型容量仍有依赖，1.7B 的小模型仍存在能力天花板。第二，Table 3 显示"forward+$\hat{x}_1$+CFG-Free"的 PickScore（23.81）虽然远高于基线（20.46），但 CFG-Free 在训练初期由于采样质量较差会有一段时间的退化（Figure 6 b1），这给早期阶段的超参选择带来挑战。第三，作者承认训练开销不低——32 张 A100 40GB GPU，且需要分组采样多张图像（每组 8 张、批大小 16 组）才能计算可靠的优势。第四，Reward 仍依赖 Flow-GRPO 使用的 PickScore/OCR 奖励器，方法对奖励黑客（reward hacking）和奖励模型的偏差较为敏感，作者在 Limitations 中未充分讨论。第五，多奖励联合优化的能力未在本文中展开，作者将多奖励任务留作未来工作。

独立分析的弱点

独立分析主要弱点在以下几方面。其一，训练成本与样本效率：每步需对每个 prompt 采 8 张图像且只随机抽 3 个连续时间步，方差较大；策略比仍依赖完整 $\hat{x}_1$ 的 token 级似然，对长序列计算量较大，改进方向可引入 token 子集采样或重要性采样。其二，CFG-Free 早期不稳定：Table 3 显示去掉 CFG 后基线 GenEval 仅 0.36，可先用 SFT/短程 RL 做热身再切换。其三，Reward 单维度限制：每个任务仅训练一个奖励信号，模型可能牺牲多样性。其四，缺少对推理成本（25 步）的优化讨论，部署中需要更少步数时稳定性未验证。其五，超参（学习率 1e-6、batch 16、3 步 Reduced-Step）的鲁棒性缺少 ablation，对不同基座（如 0.5B 或 7B）的迁移性未充分讨论。

未来方向

作者明确指出未来工作将把 UDM-GRPO 扩展到文本到视频生成（Text-to-Video）以及多奖励联合优化任务。从结果可延伸的方向还包括：第一，把 Reduced-Step 思想推广到层次化时间步采样（如基于 reward 难度自适应分配步数），进一步降低训练成本；第二，融合稀疏奖励与稠密奖励，例如引入中间 token 级别的细粒度反馈，弥补仅依赖最终奖励的不足；第三，把 UDM-GRPO 推广到自回归图像生成（VAR / LlamaGen）以及统一多模态理解-生成模型（如 Janus / Show-o）；第四，探索与 DPO / SimPO 等偏好优化方法的结合，避免 on-policy 采样带来的高显存开销；第五，研究 KL 散度约束的理论最优值 $\beta$ 与 Reduced-Step 步数之间的耦合关系，给出自动调度策略。

复现评估

作者在论文末尾提供了代码仓库 https://github.com/Yovecent/UDM-GRPO，复现难度中等偏上。优点：仓库开源、Algorithm 1 详细列出了采样、扰动和策略更新的伪代码，实验设置（AdamW $\beta_1=0.9, \beta_2=0.95$, weight decay=0.01, lr=1e-6, batch=16 groups $\times$ 8 samples, 10/25 推理步）完全披露，并明确说明使用 32 张 A100 40GB GPU，奖励模型和数据集与 Flow-GRPO 完全一致便于对照。不足：未公开具体的 Reduced-Step 时间步区间选取规则、未披露 KL 权重 $\beta$ 和裁剪 $\epsilon$ 的具体取值、未提供不同随机种子下的方差报告，且完整的训练时长未明示（推断至少需要数百 GPU 小时）。下游读者复现时需注意依赖 URSA（1.7B）预训练权重与 PickScore 奖励模型权重。

UDM-GRPO：面向均匀离散扩散模型的稳定高效分组相对策略优化 UDM-GRPO: Stable and Efficient Group Relative Policy Optimization for Uniform Discrete Diffusion Models