D-OPSD: On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models

1. 出发点 (Motivation)

2026 年的 T2I 主战场是 step-distilled diffusion models — Z-Image-Turbo (6B, 8 步)、FLUX.2-klein (4B, 4 步) 这种"少数几步出图"的型号已经把多步模型(Stable Diffusion 3.5、FLUX-dev)挤下生产线。问题是:怎么继续 fine-tune 这些模型?

三条路全部不通:

论文要的是一个组合方案: (a) 在模型自己 roll-out 的状态上更新 (b) 用 image-text pair 当监督,不要 reward。

LLM 那边的 On-Policy Self-Distillation (OPSD) 正好满足 — DeepSeek-V4、Mimo v2 都在用。同一个模型扮演 student (弱 context) 和 teacher (强 context,比如带 ground-truth answer 的 in-context),在 student 自己采样的输出上做 self-distillation。但 OPSD 怎么"翻译"到扩散模型?

论文的核心观察 (Fig. 1): 现代 T2I 模型 (Z-Image, FLUX.2) 都用 LLM/VLM 作为文本编码器。*这些 VLM 编码器的 in-context 能力被下游 diffusion 模型继承了。*把 prompt 和 target image 一起喂给 VLM,得到的 multimodal feature 当 condition,扩散模型不需要任何额外训练就能生成"保留 target 概念/风格但符合 prompt 的变体"。

In-context capability of LLM/VLM-encoded diffusion models
Fig. 1 — 用 Z-Image-Turbo (8 步) 在不训练的前提下做的实验。左列:target 图 (狗、海港、太阳镜、老虎);中列:只用 text prompt 生图;右列:用 text + target_img 经 VLM 编码后的 multimodal feature 当 condition 生图。后者保留了 target 的概念/风格,但又遵守 prompt。这就是 D-OPSD 的"免费 teacher"。

这条 emergent 性质给了天然的 teacher — 同一个模型,只是 condition 不同。后面的事就顺理成章:让 student (text-only condition) 在它自己的少步采样轨迹上,被 teacher (text+target_img condition) 的速度场预测监督。

2. 方法 (Method)

核心思想 (类比)

把 step-distilled 扩散模型想象成一个学了"4 笔速写"的画家,他原本能用 4 笔出一张高质量的猫。现在你要他学画"你家具体那只英短"。三种教法:

关键 trick:teacher 不是另一个模型,就是同一个模型给了更多 context (target image)。

D-OPSD architecture
Fig. 2 — 总体框架。Encoder 是 VLM (Qwen3-VL for Z-Image),从 text 得到 $c_s$,从 text+img 得到 $c_t$。Student model (条件 $c_s$) 做 on-policy 少步采样得到轨迹 ${x_{t_k}^s}$;Teacher model (条件 $c_t$,与 student 共享 backbone,只是 LoRA 参数不同,EMA 更新) 在同一状态上预测速度场;loss 是两者 MSE。推理时丢掉 teacher 分支,完全保留原少步 T2I pipeline。

2.1 涌现的 in-context 能力 — D-OPSD 的物理基础

论文最重要的实证 claim 是 Fig. 1 揭示的那条 emergent 性质:

当现代 T2I 模型 (LLM/VLM 编码器 + diffusion decoder) 把 prompt 和 target image 一起送入编码器,所得 multimodal feature 当 condition 时,模型直接就能生成"保留 target 概念/风格但符合 prompt"的变体 — 不需要任何额外训练。

这个能力不是 diffusion 学的,而是它从编码器 (Qwen3-VL / FLUX.2 的 LLM encoder) 那里"继承"过来的 — LLM 本来就懂 image+text 的多模态 in-context。下游 diffusion 模型只是把这种 in-context 隐式地翻译成图像条件。

Appendix A 在 FLUX.2-klein 上重复了这个实验,行为一致,说明这是一个 broadly applicable 性质,不限于 Z-Image。没有这条性质,D-OPSD 就不成立 — Fig. 6 也说了,如果 teacher (即"模型 + target image context")不能生成合理图,训练就会失败。

2.2 OPSD 在扩散域的形式化

从 LLM 的 OPSD 出发 (Eq. 1):

\(q\) 是 query,\(r\) 是 stronger context (in-context example / ground-truth answer)。student 用弱 context \(q\) 采样,teacher 用强 context \((q, r)\) 评估。

翻译到扩散: 把 \(q\) 当 text prompt \(y\),\(r\) 当 target image \(x_0\)。两个 condition:

\(f_{\text{text}}\) 只编 text,\(f_{\text{mm}}\) 编 text + image。两者共享同一个 VLM encoder。

Student on-policy 采样轨迹 (Eq. 3):

\(\Phi\) 是与推理完全相同的少步 solver(Z-Image 是 8 步,FLUX.2-klein 是 4 步)。从 Gaussian 噪声 \(x_{t_K}^s \sim \mathcal{N}(0, I)\) 出发,student 用自己的 text-only condition 走完整条少步链 — 这就是 on-policy 的物理含义。

同一轨迹上的两套预测 (Eq. 5-6):

\(\bar\theta\) 是 teacher 参数 (EMA 副本)。关键:teacher 评估的是 student 的状态 \(x_{t_k}^s\),不是 teacher 自己采样的状态 — 否则就退化成普通蒸馏。

D-OPSD 损失 (Eq. 7):

\(\mathrm{sg}\) 是 stop-gradient — teacher 不接受梯度,梯度只走 student。直觉:把 student 在自己轨迹上的速度场,拽向 teacher 在同一状态上的速度场;teacher 因 multimodal context 知道更多 target 信息,所以它的速度场指向"target-aware"方向。

为什么 D-OPSD 不破坏少步能力?

  1. 训练状态 = 推理状态 — student 用真正的少步 solver 采轨迹,优化全部发生在这些状态上,没有任何"训练时见、推理时见不到"的中间噪声状态。
  2. 无 ground-truth velocity 监督 — 不像 SFT 把模型往一个外部目标拉,supervision 来自 model 自己 (under stronger context),与原少步动态高度兼容。
  3. 推理时 teacher 分支整个丢弃 — pipeline 与原 step-distilled 模型逐字相同,没有任何 inference overhead。

2.3 完整算法 + 直觉

论文 Algorithm 1 的简化版:

  1. 采一个 batch 的 \((x_0, y)\)。
  2. 编码两个 condition: \(c_s = f_{\text{text}}(y)\), \(c_t = f_{\text{mm}}(y, x_0)\)。
  3. 从噪声 \(x_{t_K}^s \sim \mathcal N(0, I)\) 开始,用 student 跑完整条少步链 (with 与训练同条件 \(c_s\))。
  4. 每一步 \(k = K, \ldots, 1\):
    • \(u_k^s \leftarrow v_\theta(x_{t_k}^s, t_k, c_s)\)
    • \(u_k^t \leftarrow v_{\bar\theta}(x_{t_k}^s, t_k, c_t)\) — 同一 \(x_{t_k}^s\)
    • 累加 \(\frac{1}{K}\|u_k^s - \mathrm{sg}(u_k^t)\|_2^2\)
    • 用 \(u_k^s\) stepping 到 \(x_{t_{k-1}}^s\) (并 detach 避免跨步 BPTT)
  5. 反向传播,更新 student \(\theta\) (用 LoRA 时只更新 student LoRA adapter)。
  6. EMA 更新 teacher: \(\bar\theta \leftarrow 0.9999\,\bar\theta + 0.0001\,\theta\) (Section 3.3 ablation 选的最优值)。

三个会被错过的细节:

  1. 跨步 detach. 代码 (train_dopsd.py:L554) 每步都 latents_student.detach().requires_grad_(True),只让梯度在当前一步内流动。否则会 BPTT 整个少步链,显存爆。
  2. Teacher 即 EMA student. Teacher 不是另一个模型,而是 student 的 EMA 副本 (实现为同一 transformer 上的第二个 LoRA adapter)。EMA decay = 0.9999 (paper §3.3) — 大动量意味着 teacher 比 student 落后约 10000 步,提供平滑稳定的对齐目标。
  3. x_0 损失 vs v 损失 (paper-vs-code 不一致). 代码 train_dopsd.py:L600 用 F.mse_loss(x_0_student, x_0_teacher.detach()) — 对照 \(x_0 = x_t + (1 - t)\, v\),这等价于带时间权重的 velocity MSE(早 t 时刻权重大)。代码注释明说"different as shown in our paper... leads to faster convergence"。一个老实的 admission。

2.4 与代码对照

(a) 双 LoRA adapter 在同一 transformer 上 — Student 训练,Teacher 冻结由 EMA 更新。这是 D-OPSD 最重要的工程优化 — 不复制 6B 模型,只多一份 LoRA。

repo/z-image-turbo_self-distill-vlm/ema_utils.py:L31-L91 — 双 LoRA 初始化

def init_dual_lora_transformer(
    transformer,
    lora_rank=16,
    lora_alpha=16,
    target_modules=None,
    current_adapter_name="current",
    old_adapter_name="old",
    old_init_from_current=True,
):
    """
    Attach two LoRA adapters to the same transformer:
      - current: trainable
      - old: frozen, updated via EMA
    """
    # Freeze the base model
    for p in transformer.parameters():
        p.requires_grad = False

    lora_config = LoraConfig(r=lora_rank, lora_alpha=lora_alpha,
                             init_lora_weights="gaussian",
                             target_modules=target_modules)

    transformer = get_peft_model(transformer, lora_config, adapter_name=current_adapter_name)
    transformer.add_adapter(old_adapter_name, lora_config)
    transformer.set_adapter(current_adapter_name)

    set_adapter_trainable(transformer, current_adapter_name, True)
    set_adapter_trainable(transformer, old_adapter_name, False)

    if old_init_from_current:
        copy_lora_adapter_weights(transformer,
                                   src_adapter=current_adapter_name,
                                   dst_adapter=old_adapter_name)
    return transformer

训练时通过 gen_model.set_adapter("student" / "teacher") 切换。两份 LoRA 在显存里共存,但只前向其中一个 — 这就是为什么论文 §4 说 "roughly 4× FLOPs and 2× wall time" 而不是 2 倍显存。

(b) D-OPSD 主训练循环 — 实现 Eq. 5-7 + Algorithm 1。这一段是论文心脏,逐字对照 paper 公式。

repo/z-image-turbo_self-distill-vlm/train_dopsd.py:L541-L611 — D-OPSD 训练主循环

for back_step in range(len(timesteps)):
    t = timesteps[back_step].expand(bsz) / 1000
    if back_step < len(timesteps) - 1:
        next_t = timesteps[back_step + 1].expand(bsz) / 1000
    else:
        next_t = torch.ones_like(t)
    dt = next_t - t

    # detach current state to avoid cross-timestep BPTT
    latents_student = latents_student.detach().requires_grad_(True)
    latents_teacher = latents_teacher.detach()

    # === teacher branch (Eq. 6): v_θ̄(x_t^s, t, c_t) ===
    with torch.no_grad():
        gen_model.set_adapter("teacher")
        v_pred_teacher = gen_model(
            latents_student_list, t, prompt_embeds_list_vl, return_dict=False)[0]
        v_pred_teacher = torch.stack(v_pred_teacher, dim=0).squeeze(2)
        latents_teacher_cur = latents_student            # ← teacher 评估 STUDENT 的 state
        x_0_teacher = latents_teacher_cur + (1 - t.reshape(bsz,1,1,1)) * v_pred_teacher
        latents_teacher = latents_teacher_cur + v_pred_teacher * dt.reshape(bsz,1,1,1)

    # === student branch (Eq. 5): v_θ(x_t^s, t, c_s) ===
    gen_model.set_adapter("student")
    v_pred_student = gen_model(
        latents_student_list, t, prompt_embeds_list, return_dict=False)[0]
    v_pred_student = torch.stack(v_pred_student, dim=0).squeeze(2)
    latents_student_cur = latents_student
    x_0_student = latents_student_cur + (1 - t.reshape(bsz,1,1,1)) * v_pred_student
    latents_student = latents_student_cur + v_pred_student * dt.reshape(bsz,1,1,1)

    # === Eq. 7 — but in x_0 space, weighted v loss ===
    #  we use x_0 loss here, which is different as shown in our paper, It can be regarded
    #  as a weighted sum of v loss and t (with a greater weight in the early steps).
    #  We found that this leads to faster convergence.
    loss_dopsd = F.mse_loss(x_0_student, x_0_teacher.detach(), reduction="mean")
    total_loss = total_loss + loss_dopsd

三个对照点:

(c) EMA 更新 teacher LoRA — 每个 sync step 后从 student LoRA 蒸馏到 teacher LoRA。

repo/z-image-turbo_self-distill-vlm/ema_utils.py:L93-L114 + train_dopsd.py:L633-L638 — EMA

# ema_utils.py
@torch.no_grad()
def ema_update_lora_adapter(model, src_adapter, dst_adapter, ema_decay=0.999):
    """dst = ema_decay * dst + (1 - ema_decay) * src"""
    named_params = dict(model.named_parameters())
    for name, src_param in named_params.items():
        if src_adapter not in name:
            continue
        dst_name = name.replace(src_adapter, dst_adapter)
        if dst_name not in named_params:
            continue
        dst_param = named_params[dst_name]
        dst_param.data.mul_(ema_decay).add_(src_param.data, alpha=1.0 - ema_decay)

# train_dopsd.py — call site after each optimizer step
ema_update_lora_adapter(gen_model, src_adapter="student",
                        dst_adapter="teacher", ema_decay=args.ema_decay)

论文 §3.3 ablation 显示 ema_decay=0.9999 最优 (而不是 default 的 0.999) — 直接用 student 副本 (decay=0) 会训练崩坏,需要"极平滑"才能让 alignment target 稳定。

(d) VLM multimodal feature 提取 — Eq. 4 的 \(f_{\text{mm}}\) 实现。用 Qwen3-VL 的倒数第 2 层 hidden states 当 condition。

repo/z-image-turbo_self-distill-vlm/vlm_utils.py:L38-L158 — get_qwen3vl_zimage_prompt_embeds

@torch.no_grad()
def get_qwen3vl_zimage_prompt_embeds(
    vl_model, processor,
    prompts, images=None,
    hidden_state_layer: int = -2,
    ...
):
    for text, image in zip(prompts, images):
        user_content = []
        if image is not None:
            user_content.append({"type": "image", "image": image})  # ← 把 target image 当 in-context
        user_content.append({"type": "text", "text": text})

        conv = [{"role": "user", "content": user_content}]
        text_input = processor.apply_chat_template(conv, tokenize=False,
                                                    add_generation_prompt=add_generation_prompt)

        model_inputs = processor(text=[text_input], images=[image] if image else None,
                                  padding=True, truncation=True, return_tensors="pt")

        outputs = vl_model(**model_inputs, output_hidden_states=True, return_dict=True,
                           use_cache=False)
        prompt_embeds_full = outputs.hidden_states[hidden_state_layer]   # ← 倒数第 2 层
        sample_hidden = _extract_masked_hidden(prompt_embeds_full, attention_mask)
        split_hidden_states.append(sample_hidden[0][:max_sequence_length])
    return split_hidden_states

\(c_t\) 不是显式的 "image embedding + text embedding 拼起来" — 而是把 image 当作 chat 消息里的一个 user_content item,让 VLM 自己跑完 multimodal attention,取倒数第 2 层 hidden states。这种做法的关键是 VLM 已经训练过处理图文混合输入,扩散模型继承的就是这一能力 — 这也是 §2.1 强调的 emergent property 的来源。

3. 结论 (Key Findings)

LoRA 自定义训练 (4 张图 + 一点风格图,DreamBooth 数据集):

Table 1 LoRA main results
Tab. 1 — Z-Image-Turbo & FLUX.2-klein 上 LoRA 训练对比。D-OPSD 在 CLIP-S (0.3664)、Quality-S (3.7965)、Aesthetic-S (3.1710) 全部最高。DINO-D / LPIPS-D 上 PSO 略优是因为 PSO 过拟合训练图;但 PSO 的 CLIP-S 反而比 base 还低 (0.2893 vs 0.3043) — 它学了 target 但失去了 prompt 跟随。D-OPSD 是唯一在所有"学到了"和"还能用"指标上同时领先的方法。

全量 fine-tune (anime 大数据集):

Table 2 full-finetune results
Tab. 2 — 全量 fine-tune 把模型推向 anime 域。D-OPSD 在 FID 40.49 / 38.29 (vs base 48.69 / 45.23) 最低,Quality-S 3.84 / 3.83 (vs SFT 2.61 / 2.50) 几乎不掉,GenEval 0.7170 / 0.7298 (vs base 0.7543 / 0.8155) 只小幅下降,DPG 84.11 / 83.89 (vs base 84.76 / 85.56) 几乎持平。Vanilla SFT 在所有保留指标上全面塌方 (Quality 2.61, GenEval 0.16)。

LoRA 定性 (Fig. 3):

LoRA qualitative
Fig. 3 — 4 张训练图 → 4 类新场景生成。SFT 普遍模糊 (少步能力坏掉),PSO 过拟合 (鸭子原样复制,加不上场景),D-OPSD 既学到对象又能放进新场景。

全量 fine-tune 定性 (Fig. 4):

Full-finetune qualitative
Fig. 4 — 左半为目标域 (anime),右半为原域。SFT/PSO 训练后两边都坏 (颜色失真 + 模糊);D-OPSD 学到 anime 偏好的同时原域质量保留。

Ablation 验证 on-policy + self-distillation 必须同时存在:

Ablation curves
Fig. 5 — (a) 训练策略消融:SFT on target images (蓝) Quality-S 急跌;SFT on teacher samples (橙) 持平;off-policy distillation (紫) DINO 上得慢;on-policy distillation = D-OPSD (绿) DINO 收敛最快且 Quality 最高。(b) Teacher 构造消融:student 副本 (蓝,decay=0) DINO 不动,训练崩;ema=0.9 (橙) 慢;ema=0.9999 (绿) 最优。

关键量化对比 (Z-Image-Turbo,Tab 1+2 摘要):

4. 实现细节 (Implementation Notes)

训练代码已公开 — repo vvvvvjdy/D-OPSD 提供 Z-Image-Turbo 的 self-distill-vlm LoRA 训练完整代码 (commit 4a4cd7c),FLUX.2-klein 的 editing-branch 版本"稍后释出"。

5. 批判性总结 (Critical Assessment)

5.1 优点

5.2 不足 / 疑点

5.3 适用 vs 不适用

5.4 进一步阅读

讨论 / Comments

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