Flow-OPD: On-Policy Distillation for Flow Matching Models

1. 出发点 (Motivation)

T2I 模型 (Stable Diffusion 3.5 等 flow-matching backbone) 想要做"一个模型干很多事" — 同时擅长组合性 (GenEval)、文字渲染 (OCR)、美学 (PickScore)、整体质量 (DeQA)。RLHF 思路移植到这里就是 Flow-GRPO:把 reverse SDE 解释成 Markov 决策过程,reward 上 GRPO。

问题:单 reward 训练可以把一个指标拉满,但**"seesaw effect"** 立刻显现。Tab 1 是论文里那张让人脸疼的数据 — 一步一步叠加 reward,前面拉满的指标就一步一步往下掉:

Table 1: capability degradation under reward mixing
Tab. 1 — 在 SD-3.5-M 上依次叠 GenEval、OCR、PickScore、DeQA 四个 reward。+GenEval 把 GenEval 拉到 0.94;但接着 +OCR,GenEval 掉 5%;+PickScore 再掉 7%;+DeQA 又掉 9%。每加一个新 reward,旧 reward 上的能力就被吃掉。

视觉上更直白 — 单 reward 训出的"专家"在跨任务测试时常常崩坏:

cross-task evaluation of single-reward GRPO
Fig. 2 — 单 reward GRPO 横评。GRPO-OCR 把"philo is a weird"渲染对了,但同一模型在 GenEval task 上画"四个键盘"就崩了;GRPO-GenEval 排列得很整齐但 OCR 渲染又出错。每个专家都被自己的 reward 拽偏。

论文给的因果解释 (Sec. 4.2): 单 reward GRPO 把多维冲突压缩成标量 advantage,模型为了最大化 \(A_1\) 会"吃掉"那些没被监督的参数自由度。一阶 Taylor 近似下,任务 \(T_1\) 对未监督任务 \(T_k\) 的伤害:

当任务梯度 \(\langle \nabla_\theta J_k, \nabla_\theta J_1\rangle \lt 0\) (高维空间里这是常态),没有 \(T_k\) 监督的优化器会"主动"破坏 \(T_k\) 的能力来换 \(T_1\) — 这就是 reward hacking 的根因。标量 reward 的信息密度天然不够,你需要 dense 的 trajectory-level 监督。

LLM 那边怎么解的?On-Policy Distillation (OPD) — 学生自己采样 → 教师在学生的 trajectory 上提供 dense 监督。DeepSeek-V4 / Mimo v2 / GLM-5 全在用。问题是 LLM 是离散 next-token,OPD 怎么"翻译"到连续速度场 flow matching 上?这就是 Flow-OPD 的全部技术贡献。

2. 方法 (Method)

核心思想 (类比)

把 SD-3.5-M 想成一个学画画的学生,要同时学:

单 reward GRPO 像"每天只跟一个老师学,这周构图、下周书法"— 上下周互相覆盖,什么都学不全。

Flow-OPD 的玩法:学生自己画画,每画一笔老师都站旁边看 — 但每张图按内容路由,只让对应的老师评。

关键技术问题是:怎么定量"老师建议学生这一笔往哪挪 X 像素"?在 flow matching 里,每一步的"建议"就是速度场的差。论文的核心数学贡献就是把这一点严格地从 reverse KL 推导出来。

2.1 Flow Matching 速览

Flow Matching (Lipman et al. 2023) 把噪声 \(p_0\) 映射到数据 \(p_{\text{data}}\) 的 ODE:

OT (Optimal Transport) 形式下,直线轨迹 \(x_t = (1-t)x_0 + t x_1\),模型 \(v_\theta\) 学常速度 \((x_1 - x_0)\),训练目标

把这一离散积分过程看成 Markov 步,就能与 RL 接上 — 这是 Flow-GRPO (Liu et al. 2025) 已经做的事,Flow-OPD 直接继承这一视角。

2.2 单 GRPO 为何失败 (sparse reward + 梯度干扰)

论文 Sec. 4 系统讨论了三个问题:

  1. Q1: GRPO 为何起作用? 因为 on-policy 探索打破了 offline SFT 受限于数据集质量的天花板。
  2. Q2: 单 GRPO 为何在多任务上崩? 见前面 Taylor 展开 — 标量 advantage 压缩了多维冲突,模型为提高目标指标会"吃"未监督的自由度。
  3. Q3: 直接混 reward 行不行? 不行 — Tab 1 显示每加一个 reward 旧能力就掉一截。

结论:必须同时满足 (a) on-policy (维持探索) (b) densely uncoupled (每个 task 独立信号,不竞争同一个标量)。这就是 OPD 范式 — 由多个 teacher 在学生的轨迹上分头提供 dense 监督。

2.3 ODE → SDE — 注入探索 (Eq. 5-6)

RL 需要 stochastic 行为策略才能做 importance sampling。Flow-OPD 沿用 Flow-GRPO 的做法,把确定性 ODE 转成等价 SDE:

Euler-Maruyama 离散化后,每一步的 transition 分布是各向同性高斯:

\(\sigma_t\) 是注入的噪声幅度 (实现里有 noise_level 超参,Flow-GRPO 默认 0.7),\(\mu_\theta\) 是由学生速度场决定的 Euler 步均值。每个 prompt 采 \(G\) 条 trajectory 得到 on-policy 边缘分布 \(\rho_t^\theta\)。

2.4 KL → L2 的精彩折叠 (Eq. 8-9)

这是论文的数学高潮。要把"OPD 用 reverse-KL 当奖励"翻译到连续域,关键问题是:在 SDE 框架下, \(\mathrm{KL}(\pi_\theta \| \pi_{\text{target}})\) 怎么算?

注意学生策略 \(\pi_\theta = \mathcal{N}(\mu_\theta, \sigma_t^2\Delta t\,\mathbf{I})\) 和教师 \(\pi_{\text{target}} = \mathcal{N}(\mu_{\text{target}}, \sigma_t^2\Delta t\,\mathbf{I})\) 共享同一个协方差矩阵 — 因为 SDE 的 noise 注入是结构性的,不依赖模型预测。两个等方差高斯之间的 KL 有闭式:

代入 \(\Sigma = \sigma_t^2\Delta t\,\mathbf{I}\):

把 \(\mu_\theta\) 按 SDE 离散化展开,常数项消掉,KL 就退化为速度场之间的 L2 距离:

记 \(w(t)\) 为前面的时间相关系数。每一步的 dense reward 就是负的、加权的、速度场 L2 差 (并 detach 掉 \(v_\theta\) 上的梯度):

这就是论文的核心 trick. Reverse KL 本来是 LLM 离散概率分布之间的事,在连续 flow matching 里它等价于速度场 L2 差 — 实现上你只需要每一步比较学生和教师的预测向量,完全不涉及任何概率密度的解析形式。\(\bar v_\theta\) 的 detach 至关重要 (Thinking Machines OPD 的设计哲学): reward 是被监督的目标,不是被反传的对象,梯度只通过 policy ratio 走。

2.5 任务路由 + 多教师 dense 监督 (Eq. 7)

四个 teacher 分别在四个 task 上 GRPO 训到饱和。在线训练时,每个 prompt \(c\) 被一个硬路由 \(R(c)\) 映射到一个 teacher \(k\),只用这个 teacher 的速度场作为目标:

"硬路由"而不是"软混合" — 论文明确说这是为了消除 inter-domain 梯度干扰 — 同一个 prompt 一辈子只跟一个 teacher 学,不会出现"OCR 教师和 GenEval 教师推不同方向"的情况。

Headline performance curves and bars
Fig. 1 — Flow-OPD 训练曲线 (左) 和最终指标 (右):GenEval 训练曲线稳定升到 93,而 vanilla GRPO 提前停在 78;评估指标上 GenEval +19, OCR +11, DeQA 持平 teacher, PickScore +5。

2.6 PPO clip + 冷启动 (Eq. 11)

Dense 高频 reward 会让 policy 跳得太狠,论文借 PPO 的 clipped surrogate 来限制每步的策略漂移。设 policy ratio \(\rho_{t,i,j}(\theta) = \pi_\theta(a_{t,i,j} \mid s_{t,i,j}) / \pi_{\theta_{\text{old}}}(a_{t,i,j} \mid s_{t,i,j})\),对 B 个 prompts × G 条 trajectory × T 个去噪步取平均:

**冷启动 (Sec. 5.1):**如果学生从 base SD-3.5-M 直接进 OPD,初期 trajectory 完全偏离教师 manifold,信号噪声极大。论文给两个变体:

2.7 Manifold Anchor Regularization — 锚定美学 (Eq. 12)

OPD 提供任务对齐的 dense 信号,但有"reward hacking 副作用" — 学生为了 OCR 把整张图画成纯白底黑字,GenEval 拼对了但背景塌成马赛克。MAR 加了一个 任务无关的美学 teacher \(v_{\text{aesthetic}}\) (论文用 DeQA 训出来的 teacher),在所有数据点上提供全场监督:

注意 MAR 是不分 task 全场施加的"美学保安";硬路由的 task-specific teacher 提供主要的能力监督。两套并行。

MAR ablation qualitative
Fig. 5 — MAR 消融。w/o KL Loss 的样本对象单调、背景同质化 (典型的 reward hacking);加上 MAR (w. KL Loss(Ours)) 后橙子有阴影、剪刀有立体感、小鸟羽毛纹理回来了。

2.8 与代码对照

2026-06-08 更新:官方训练代码已释放(CostaliyA/Flow-OPD,"Release Training Code" 已 ✅)。本节已把原先的教学性伪代码全部替换为作者原码的逐行引用(repo/scripts/train_sd3_opd_mix.py 等)。

Flow-OPD 基于 yifan123/flow_grpo 扩展,官方 repo 直接 vendor 了 flow_grpo/ 与 scripts/。核心创新(per-task teacher routing + KL→L2 dense reward + MAR)就落在 train_sd3_opd_mix.py 的训练循环里。下面四段全部 verbatim 引自官方 repo。

(a) SDE 转换 + transition log-prob — 实现 Eq. 5-6。共方差结构在 L44-L50 (std_dev_t),正是后面 KL→L2 折叠之所以成立的根源。

repo/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py:L46-L68 — SDE Euler-Maruyama 步 + Gaussian transition log-prob

    sigma_max = self.sigmas[1].item()
    dt = sigma_prev - sigma

    if sde_type == 'sde':
        std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma)))*noise_level

        # our sde
        prev_sample_mean = sample*(1+std_dev_t**2/(2*sigma)*dt)+model_output*(1+std_dev_t**2*(1-sigma)/(2*sigma))*dt

        if prev_sample is None:
            variance_noise = randn_tensor(
                model_output.shape,
                generator=generator,
                device=model_output.device,
                dtype=model_output.dtype,
            )
            prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1*dt) * variance_noise

        log_prob = (
            -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1*dt))**2))
            - torch.log(std_dev_t * torch.sqrt(-1*dt))
            - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
        )

注意 std_dev_t 完全由 timestep 决定,不依赖 model_output — 这就是为什么学生与教师的 transition 自动同协方差,从而 Eq. 8 中 \(\Sigma^{-1}\) 退化为标量。prev_sample_mean 就是 SDE 步的高斯均值 \(\mu\),KL→L2 折叠的全部魔法都在它身上。

(b) 任务路由(Eq. 7) — 每个 dataset(task)绑定它自己的 teacher LoRA reference:

repo/scripts/train_sd3_opd_mix.py:L1484-L1488 — OPD 模式下按当前 dataset 取该任务的 teacher reference transformer

# In OPD mode, get current dataset's ref transformer for KL reward
if is_opd_mode and is_alternate_mode:
    current_ref_transformer = ref_transformers.get(current_ds_name)
else:
    current_ref_transformer = None

对照 §2.5 的 hard routing \(R(c)\):工程上就是"按 dataset 名查表取 teacher"。ref_transformers 是 {task → teacher LoRA},每个 task 的样本只跟自己的专家算 KL——soft mixing 被显式拒绝,跨任务梯度干扰在 input 端就消除。

(c) KL → 速度场 L2 的 dense reward(Eq. 8-10) — 全文最关键的折叠,这里 verbatim:

repo/scripts/train_sd3_opd_mix.py:L1707-L1728 — reward_mode=="mixed":末步稀疏 task reward + 逐 step 稠密 KL-to-teacher reward

if config.train.get("kl_reward_level") == "step_wise" and config.train.get("kl_scale", 0) != 0:
    if prev_sample_mean_ref_lora is not None:
        kl_reward = ((prev_sample_mean - prev_sample_mean_ref_lora) ** 2).mean(dim=(1, 2, 3), keepdim=True) / (2 * std_dev_t ** 2)
    else:
        kl_reward = torch.zeros_like(prev_sample_mean.mean(dim=(1, 2, 3), keepdim=True))
    kl_reward = kl_reward.squeeze(-1).squeeze(-1)

    if kl_norm == "per_sample":
        kl_reward = (kl_reward - kl_reward.mean()) / (kl_reward.std() + 1e-4)
    elif kl_norm == "per_timestep":
        kl_reward = (kl_reward - kl_reward.mean()) / (kl_reward.std() + 1e-4)
    elif kl_norm == "global":
        kl_reward = (kl_reward - kl_reward.mean()) / (kl_reward.std() + 1e-4)

    advantages = task_advantage + config.train.kl_scale * kl_reward
    advantages = torch.clamp(
        advantages,
        -config.train.adv_clip_max,
        config.train.adv_clip_max,
    )
else:
    advantages = task_advantage

对照 Eq. 8-9:((prev_sample_mean - prev_sample_mean_ref_lora) ** 2) / (2 * std_dev_t ** 2) 就是两个同协方差高斯之间的 KL —— 学生 SDE 均值 \(\mu_\theta\) 与路由到的 teacher 均值 \(\mu_{\phi_k}\) 的逐元素 L2,除以 \(2\sigma_t^2\),正是论文 Eq. 8 的逐字翻译。稀疏的 task_advantage(只在末步给)+ 稠密的 kl_scale * kl_reward(每一步给)= Eq. 10 的混合优势。prev_sample_mean_ref_lora 由 (b) 路由到的 current_ref_transformer 算出。

(d) PPO clip(Eq. 11)+ MAR 美学锚定(Eq. 12):

repo/scripts/train_sd3_opd_mix.py:L1753-L1767 — clipped surrogate;beta>0 时加 MAR(到 task-agnostic teacher 的 KL)

else:
    ratio = torch.exp(log_prob - sample["log_probs"][:, j])
    unclipped_loss = -advantages * ratio
    clipped_loss = -advantages * torch.clamp(
        ratio,
        1.0 - config.train.clip_range,
        1.0 + config.train.clip_range,
    )
    policy_loss = torch.mean(torch.maximum(unclipped_loss, clipped_loss))
    if config.train.beta > 0:
        kl_loss = ((prev_sample_mean - prev_sample_mean_ref_base) ** 2).mean(dim=(1,2,3), keepdim=True) / (2 * std_dev_t ** 2)
        kl_loss = torch.mean(kl_loss)
        loss = policy_loss + config.train.beta * kl_loss
    else:
        loss = policy_loss

prev_sample_mean_ref_base 由 mar_transformer(单独的 task-agnostic 美学 teacher,加载时 requires_grad_(False) + eval())算出;loss = policy_loss + beta * kl_loss 就是 Eq. 12 的 MAR 锚定项 —— 用一个"全数据美学 teacher"把生成拉回高质量 manifold,缓解纯 RL 的美学退化。

三个关键观察(基于真实代码,非伪代码):

3. 结论 (Key Findings)

在 SD-3.5-M 上的四个 benchmark:

Main results table
Tab. 2 — Flow-OPD (Merge cold-start) 在 GenEval (0.92)、OCR (0.94)、DeQA (4.35)、PickScore (23.08) 上拿 best;平均 0.9044,比最强基线 GRPO-Mix (0.8165) 高 8.8 pt,比 base 高 18.8 pt。GRPO-GenEval 单 teacher 把 GenEval 拉到 0.94 但 PickScore 掉到 21.53;Flow-OPD 是同一模型在四个指标上都接近或超过专业 teacher。

Teacher-Surpassing 现象 (paper Sec. 6.2): 在 OCR 上学生 0.94 超过 OCR teacher 的 0.92,DeQA 上学生 4.35 超过 DeQA teacher 的 4.23 — 单一 teacher 在自己专攻领域可能 cherry-pick 过头,而多 teacher dense 监督迫使学生学到更"全息"的速度场,在某些边界 prompt 上比任何单 teacher 都好。论文把这归因到"潜在 flow manifold 内的知识交叉传授"。

Cold-start ablation
Fig. 4 — 冷启动消融。从 Base (浅灰) → GRPO-Mix → OPD w/o cold start → Ours(SFT) → Ours(Merge)。重要观察:即使不冷启动,纯 OPD 也能从 63 直接跳到 81 (GenEval),证明 dense 多教师监督本身的有效性;冷启动 (尤其是 Merge) 再加 10 pt 到 92。

定性比较 (Fig. 3):

Qualitative comparison
Fig. 3 — 8 个模型在 6 类 prompt 上的对比。最右两列 (Merge / Ours = Flow-OPD with Merge cold-start) 在 "philo is a weird" 海报、橙色剪刀、蘑菇猫、八辆车、火山实验海报、人 + 红绿灯 上同时做到指令遵循 + 美学。其他基线总在某一个维度崩盘。

OOD 验证 (T2I-CompBench++): 在 paper 没训过的 7 维组合性 benchmark (Color, Shape, Texture, Complex, 3D-Spatial, Numeracy, Non-Spatial) 上 Flow-OPD 也全部最强 — 表明 dense 多教师监督对未见 prompt 也泛化。

MAR 的量化收益 (Tab 4): 在 ImageReward (1.02→1.36)、Aesthetic (5.87→6.23)、UnifiedReward (3.339→3.659)、HPS-v2.1 (0.2982→0.3302)、QwenVL Score (3.45→4.05) 上全面提升 — w/o MAR 这些指标都明显回退。

4. 实现细节 (Implementation Notes)

**代码状态:**官方 repo CostaliyA/Flow-OPD 目前只发布:README + 项目页 + arxiv PDF + HF model checkpoint。TODO 列表第一条"Release full training code"标为 in progress。训练循环的真实代码不可访问。下方实现细节来自 paper §6.1, Appendix, 加上致谢的基础 codebase yifan123/flow_grpo。

5. 批判性总结 (Critical Assessment)

5.1 优点

5.2 不足 / 疑点

5.3 适用 vs 不适用

5.4 进一步阅读

讨论 / Comments

评论托管在本仓库的 GitHub Discussions, 需 GitHub 账号。