D-OPSD: On-Policy Self-Distillation for Continuously Tuning Step-Distilled Diffusion Models¶
会议: NeurIPS 2026
arXiv: 2605.05204
论文: https://openreview.net/forum?id=HM5bMcRfJ3
领域: 图像生成
关键词: 步数蒸馏扩散模型, 同策自蒸馏, 持续微调, 流匹配, 图像生成
一句话总结¶
本文发现采用大语言模型或视觉语言模型作为编码器的现代扩散模型天然继承了编码器的上下文学习能力,进而提出同策自蒸馏框架 D-OPSD,让同一模型分别在纯文本条件与图文多模态条件下扮演学生与教师,在学生自身的少步采样轨迹上进行速度场对齐,从而在无需外部奖励模型的前提下实现步数蒸馏扩散模型的高质量概念定制与持续微调。
研究背景与动机¶
近年来文本到图像(T2I)生成领域正迅速从高延迟的多步扩散模型向高效的少步蒸馏模型(如 Z-Image-Turbo 和 FLUX.2-klein)演进,这类少步模型不仅将函数评估次数(NFE)压缩至 4 到 8 步,甚至在生成质量上也开始超越多步前辈。然而,如何在下游成对图文数据上对这些步数蒸馏模型进行持续微调(如 LoRA 概念定制或特定领域全量微调)却是一个棘手难题。直接套用常规多步扩散模型的标准监督微调(Vanilla SFT,即对目标图像加噪并回归流匹配真值速度),是在外部真实图像诱导的离线状态分布上定义训练信号,而非模型自身少步采样器实际访问的状态轨迹。由于步数蒸馏模型的生成质量高度依赖少数几步精细校准的去噪更新,这种训练与测试状态失配会严重扰乱原有的少步动力学,导致微调后少步生成质量急剧崩塌、图像模糊。而另一类基于配对样本优化的离线偏好方法(如 PSO)虽然限制在少步时间步上训练,但输入状态和监督信号依然绑定在离线目标图像上,在小样本下极易过拟合训练集而丧失对新场景提示词的泛化能力。
相比之下,在线强化学习(如 ReFL、Flow-GRPO)在模型自己采样的少步轨迹上进行优化,能够有效保持训练与推理一致性,从而不破坏少步生成能力。但在线强化学习依赖精心设计的奖励函数或外部人类偏好模型,这对于绝大多数仅拥有少量图文对(例如几张特定物体或风格示例图)的社区开发者而言难以获取。因此,理想的步数蒸馏模型持续微调方案应当兼具两者的核心优势:既要在模型自身的同策(on-policy)采样轨迹上更新参数以维持少步推理能力,又要能直接从配对的目标图文数据中提取监督信号而无需外部奖励模型。在自回归大语言模型(LLM)中,同策自蒸馏(OPSD)通过把参考答案拼入输入序列构造强上下文教师来解决这一问题;但对于扩散模型,若直接将目标图像塞进加噪去噪轨迹中,就会改变采样轨迹本身并退化回破坏少步能力的离线 SFT。
本文的切入角度是:当前最先进的少步扩散模型普遍采用 LLM 或 VLM(如 Qwen3)作为条件编码器,作者通过实验发现后续的扩散主干网络无需额外训练即可继承编码器的多模态上下文能力——当把条件从纯文本特征替换为“文本提示 + 目标图像”的多模态特征时,模型直接少步采样就能生成保持目标概念或风格的多样化图像。核心 idea:利用现代扩散模型继承自多模态编码器的上下文能力,将目标图像作为编码端的特权上下文而非去噪端的加噪目标,让同一模型在纯文本条件下作为学生展开少步同策轨迹,并在相同轨迹状态上用图文多模态条件下的教师速度场进行自蒸馏监督。
方法详解¶
整体框架¶
D-OPSD 是一个单阶段、无需奖励模型的同策自蒸馏微调框架。给定训练图文对 \((x_0, y)\),框架首先通过编码器分别提取仅含文本提示的学生条件特征 \(c_s\),以及融合了文本提示与目标图像的教师多模态条件特征 \(c_t\)。随后,学生模型在纯文本条件 \(c_s\) 下从随机高斯噪声出发,使用与测试时完全一致的少步求解器执行前向去噪,生成一条完整的同策少步状态轨迹。接着,学生分支和基于指数移动平均(EMA)更新的教师分支在这条学生轨迹所经过的每一个中间潜空间状态上,分别根据各自的条件 \(c_s\) 与 \(c_t\) 预测流匹配速度场,并通过最小化两者速度预测的均方误差来更新学生模型。训练结束后,多模态教师分支直接丢弃,推理阶段完全保持原少步模型标准的纯文本到图像生成管线。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入:训练图文对 (x0, y)<br/>与初始高斯噪声"] --> B["多模态上下文条件构建<br/>纯文本学生条件 vs 图文教师条件"]
B --> C["学生同策少步轨迹采样<br/>沿纯文本条件展开少步ODE轨迹"]
C --> D["同轨迹速度场自蒸馏对齐<br/>同状态双分支速度场MSE与EMA教师"]
D --> E["输出:微调后的少步扩散模型<br/>推理仅保留纯文本学生分支"]
关键设计¶
1. 多模态上下文条件构建:利用编码器涌现的上下文能力无损注入目标图像监督
在扩散模型中实现同策自蒸馏的最大障碍,是如何在不触碰潜空间去噪状态的前提下把目标图像 \(x_0\) 的监督信息告诉模型。本文发现,以 Qwen3-4B 等大语言模型作为文本编码器的现代扩散模型(如 Z-Image-Turbo、FLUX.2-klein),其去噪主干对编码器特征空间具有极强的上下文适应性。基于此,本文在条件编码阶段构造非对称的双路上下文:学生分支仅使用原文本编码器处理文本提示得到纯文本条件 \(c_s = f_{\mathrm{text}}(y)\),确保其推理路径与原始文生图任务完全一致;教师分支则引入配对的视觉语言编码器(如 Qwen3-VL-4B)同时编码文本提示和目标图像,得到富含目标概念、外观或风格先验的多模态条件 \(c_t = f_{\mathrm{mm}}(y, x_0)\)。为了消除 VLM 继续训练带来的特征分布偏移(直接替换会导致生成图像出现高频伪影和过度锐化),作者在工程实现上保留了 Qwen3-VL-4B 的视觉编码器(ViT)与连接器(Connector)权重,但将其内部的 LLM 参数替换为原始文本编码器 Qwen3-4B 的权重。这一巧妙的权重回退操作在激活多模态上下文理解能力的同时,保证教师条件特征分布与扩散主干预训练时的文本特征空间高度对齐。
2. 学生同策少步轨迹采样:在模型自身真实推理状态上更新以消除训练-测试失配
标准 SFT 之所以会摧毁步数蒸馏模型的少步采样能力,根本原因在于多步模型可以通过数十步迭代逐步修正单步误差,而少步模型(仅 4 或 8 步)对每一步潜变量状态所在流形极其敏感,在真实图像加噪状态上强行拟合真值速度会直接扭曲蒸馏好的少步跳步轨迹。为此,D-OPSD 严格遵循同策(on-policy)学习原则,在每次训练迭代中,先固定测试时的 \(K\) 步时间步调度 \(1 = t_K > t_{K-1} > \cdots > t_0 = 0\),从标准高斯噪声 \(x_{t_K}^s \sim \mathcal{N}(0, I)\) 出发,仅以纯文本特征 \(c_s\) 为条件,调用原模型的少步 ODE 求解器 \(\Phi\) 逐步推演生成学生的少步去噪轨迹:
其中 \(\mathrm{sg}(\cdot)\) 表示停止梯度操作,即采样轨迹本身仅作为收集同策访问状态的 rollout 过程,不跨时间步反向传播梯度。通过这种方式,模型参数更新所面对的输入潜变量状态 \(\{x_{t_k}^s\}_{k=1}^K\) 百分之百来自当前学生策略在测试时的真实少步分布,从源头上杜绝了离线加噪状态带来的分布偏移。
3. 同轨迹速度场自蒸馏对齐:在同策轨迹状态上对齐双分支速度预测并稳定自蒸馏
在收集到学生的少步同策状态序列 \(\{x_{t_k}^s\}_{k=1}^K\) 后,需要将多模态教师分支的知识蒸馏给学生。不同于自回归语言模型在每个时间步直接输出离散词表上的概率分布(可直接计算 KL 散度),基于流匹配(Flow Matching)的扩散模型通过连续条件速度场来参数化去噪动力学。因此,D-OPSD 让参数为 \(\theta\) 的学生网络和参数为 \(\bar{\theta}\) 的教师网络在同一个同策状态 \(x_{t_k}^s\) 和时间步 \(t_k\) 上分别在各自条件下预测速度向量 \(u_k^s = v_\theta(x_{t_k}^s, t_k, c_s)\) 与 \(u_k^t = v_{\bar{\theta}}(x_{t_k}^s, t_k, c_t)\),并最小化整条轨迹上两者的均方误差损失:
这里的核心机制在于,不仅优化状态来自学生自身的少步采样,连监督信号 \(u_k^t\) 也是在同一少步动力学框架下针对同一状态实时评估得到的,从而将学生条件下的去噪轨迹自然拉向富含目标图像语义的教师轨迹分布。此外,如果直接拷贝当前学生权重作为教师,自指涉的高方差目标会导致训练迅速坍缩;本文采用极大动量系数(0.9999)的指数移动平均(EMA)来缓慢更新教师参数 \(\bar{\theta}\),在平滑高方差对齐目标、保证训练稳定性的同时,让教师能够持续追踪学生的学习进展。
损失函数 / 训练策略¶
- 总体训练目标:全程仅使用式(6)定义的单阶段同策速度场自蒸馏损失 \(\mathcal{L}_{\mathrm{D\text{-}OPSD}}\),无需任何外部奖励模型、判别器或先微调多步模型再二次蒸馏(如 DMD)的多阶段管线。
- 教师模型构建与显存优化:学生与教师均初始化自预训练步数蒸馏基座模型 \(\phi\),每步迭代后通过 \(\bar{\theta} \leftarrow \mu \bar{\theta} + (1-\mu)\theta\) 更新教师,默认 EMA 衰减率 \(\mu = 0.9999\)。在小样本 LoRA 微调场景下(rank 64,alpha 128),主干权重共享,仅对轻量级的 LoRA 权重维护 EMA 副本,大幅降低了显存开销。
- 训练超参数配置:小样本概念/风格定制(DreamBooth 等 30 个概念类 + 10 个风格类,每类约 4 张图)在单张 H800 GPU 上训练 1K 步,batch size 为 4,Z-Image-Turbo(8 步采样,6B)学习率设为 \(4\times 10^{-5}\),FLUX.2-klein(4 步采样,4B)学习率设为 \(1\times 10^{-5}\);大规模领域微调(25K 高质量动漫数据集全量微调)在 32 张 H800 GPU 上训练 10K 步,总 batch size 为 256,两者学习率分别设为 \(3\times 10^{-5}\) 和 \(8\times 10^{-6}\)。
实验关键数据¶
主实验¶
表 1:小样本定制化 LoRA 微调系统级对比(涵盖 30 个概念类与 10 个风格类的平均结果)
| 基座模型与方法 | DINO-D \(\downarrow\) | LPIPS-D \(\downarrow\) | VLM-J \(\uparrow\) | CLIP-S \(\uparrow\) | Quality-S \(\uparrow\) | Aesthetic-S \(\uparrow\) |
|---|---|---|---|---|---|---|
| Z-Image-Turbo (8 NFEs) | ||||||
| Base-Model | 0.2373 | 0.7520 | 1.0000 | 0.3043 | 3.5088 | 2.9667 |
| Vanilla SFT | 0.2212 | 0.6501 | 1.3293 | 0.3095 | 2.4236 | 2.3582 |
| SFT + LoRA on distilled | 0.1588 | 0.7243 | 1.8525 | 0.3201 | 3.6081 | 3.0027 |
| DreamBooth | 0.0902 | 0.6424 | 3.0625 | 0.3328 | 2.5582 | 2.3755 |
| PSO | 0.0570 | 0.4974 | 3.3333 | 0.2893 | 3.3422 | 3.0820 |
| D-OPSD (本文) | 0.0823 | 0.5803 | 3.3333 | 0.3664 | 3.7965 | 3.1710 |
| FLUX.2-klein (4 NFEs) | ||||||
| Base-Model | 0.2386 | 0.7518 | 1.0000 | 0.3020 | 3.5076 | 2.9665 |
| Vanilla SFT | 0.2212 | 0.6500 | 1.2083 | 0.3002 | 2.4119 | 2.3595 |
| SFT + LoRA on distilled | 0.1564 | 0.7288 | 1.0000 | 0.3151 | 3.6079 | 3.0021 |
| DreamBooth | 0.1053 | 0.6824 | 3.0283 | 0.3266 | 2.5492 | 2.3683 |
| PSO | 0.0593 | 0.5120 | 3.3015 | 0.2801 | 3.3258 | 3.0566 |
| D-OPSD (本文) | 0.0822 | 0.5648 | 3.1255 | 0.3689 | 3.7957 | 3.1521 |
表 2:大规模数据集(25K 动漫域)全参数微调系统级对比
| 基座模型与方法 | FID \(\downarrow\) | DINO-D \(\downarrow\) | LPIPS-D \(\downarrow\) | Quality-S \(\uparrow\) | Aesthetic-S \(\uparrow\) | GenEval \(\uparrow\) | DPG \(\uparrow\) |
|---|---|---|---|---|---|---|---|
| Z-Image-Turbo (8 NFEs) | |||||||
| Base-Model | 48.6858 | 0.1274 | 0.7036 | 3.7624 | 3.5218 | 0.7543 | 84.7645 |
| Vanilla SFT | 82.2036 | 0.1896 | 0.6787 | 2.6121 | 2.4852 | 0.1588 | 69.9746 |
| PSO | 88.4343 | 0.1716 | 0.6530 | 2.8653 | 2.6424 | 0.2475 | 72.7363 |
| D-OPSD (本文) | 40.4938 | 0.1088 | 0.6419 | 3.8438 | 3.6195 | 0.7170 | 84.1116 |
| FLUX.2-klein (4 NFEs) | |||||||
| Base-Model | 45.2335 | 0.1210 | 0.6894 | 3.7626 | 3.4465 | 0.8155 | 85.5624 |
| Vanilla SFT | 98.7893 | 0.2565 | 0.6567 | 2.5038 | 2.4121 | 0.1113 | 56.2138 |
| PSO | 82.1438 | 0.2185 | 0.6511 | 2.8066 | 2.5952 | 0.2406 | 63.4174 |
| D-OPSD (本文) | 38.2932 | 0.0902 | 0.6476 | 3.8328 | 3.5511 | 0.7298 | 83.8927 |
消融实验¶
表 3:单阶段 D-OPSD 与两阶段微调+重蒸馏管线(SFT on Multi-step + DMD)的精度与成本对比(基于 Z-Image / Z-Image-Turbo)
| 训练设置与对比方案 | DINO-D \(\downarrow\) | Quality-S \(\uparrow\) | GenEval \(\uparrow\) | 推理步数 (NFE) \(\downarrow\) | 训练耗时 (H800 GPU hours) \(\downarrow\) |
|---|---|---|---|---|---|
| LoRA 定制设置 | |||||
| SFT on Z-Image-Turbo (单阶段直接SFT) | 0.2212 | 2.4236 | - | 8 | 5 |
| SFT on Z-Image (多步模型直接SFT) | 0.0547 | 3.3488 | - | 100 | 5 |
| SFT on Z-Image + DMD (两阶段微调+蒸馏) | 0.1024 | 3.3372 | - | 8 | 5 + 882 (共 887) |
| D-OPSD (本文,单阶段少步微调) | 0.0823 | 3.7965 | - | 8 | 9 |
| 全量微调设置 | |||||
| SFT on Z-Image-Turbo (单阶段直接SFT) | 0.2565 | 2.5038 | 0.1588 | 8 | 1066 |
| SFT on Z-Image (多步模型直接SFT) | 0.0843 | 3.3801 | 0.6578 | 100 | 1066 |
| SFT on Z-Image + DMD (两阶段微调+蒸馏) | 0.1094 | 3.3652 | 0.6084 | 8 | 1066 + 2054 (共 3120) |
| D-OPSD (本文,单阶段少步微调) | 0.1088 | 3.8438 | 0.7170 | 8 | 1918 |
关键发现¶
- 消除概念学习与少步生成质量之间的冲突:在表 1 的 LoRA 微调中,Vanilla SFT 和 DreamBooth 虽然一定程度降低了 DINO-D,但少步生成质量(Quality-S)从基座的 3.5088 暴跌至 2.4236 和 2.5582;PSO 虽在训练集近义改写提示词上取得最低的 DINO-D(0.0570),但其在新场景泛化提示词上的 CLIP-S 跌至 0.2893(甚至低于基座的 0.3043),说明严重过拟合了参考图背景。D-OPSD 则在保持最高主体一致性评分(VLM-J 达 3.3333)的同时,将新场景泛化 CLIP-S 提升至 0.3664,且 Quality-S 反超基座达到 3.7965。
- 大规模持续微调不发生灾难性遗忘:在表 2 的 25K 动漫图像全参数微调中,Vanilla SFT 与 PSO 在通用指令遵循基准 GenEval 上分别从 0.7543/0.8155 崩溃至 0.1588/0.1113 和 0.2475/0.2406,几乎丧失了原有的通用生成与少步采样能力;而 D-OPSD 不仅将 FID 大幅降低至 40.4938(Z-Image-Turbo)和 38.2932(FLUX.2-klein),还将 GenEval 维持在 0.7170 和 0.7298 的高位,证明同策自蒸馏能在向目标域偏移的同时完好保留先验知识。
- 训练策略与教师构建消融结论:在图 6 的四组训练范式消融中,随着训练步数增加,Vanilla SFT 的图像质量持续下滑;若仅用多模态教师生成图像做离线 SFT 或在固定数据集上做离线蒸馏(Off-policy distillation),收敛速度明显变慢;唯独同策自蒸馏(D-OPSD)同时实现了最快的 DINO 相似度上升和最高的生成质量。在教师构建方式上,直接用当前学生权重作教师会导致训练崩溃,使用冻结基座或高动量 EMA(0.9999)教师最为稳定,其中 EMA 0.9999 因能渐进追踪学生能力而效果最优。
- 相比“多步微调 + 重新蒸馏”两阶段管线极具性价比:如表 3 所示,虽然单步迭代中 D-OPSD 因需 rollout 和教师前向计算导致单步耗时约为直接 SFT 的 2 倍(LoRA 下 9 GPU 小时 vs 5 GPU 小时),但直接 SFT 完全不可用;若采用先微调多步模型再用开源 DMD 重新蒸馏的两阶段方案,LoRA 定制总耗时高达 887 GPU 小时(慢近 98 倍),且最终 Quality-S(3.3372 vs 3.7965)和全量微调 GenEval(0.6084 vs 0.7170)均明显劣于单阶段 D-OPSD。
亮点与洞察¶
- 发掘并巧妙复用现代扩散模型编码器的零样本多模态上下文涌现能力:本文最核心的“啊哈”时刻在于发现采用 LLM/VLM 作编码器的扩散模型(如 Z-Image-Turbo、FLUX.2-klein),无需训练去噪主干即可在传入“文本+参考图”多模态特征时生成保持目标身份与风格的变体。作者顺势将这一零样本特性转化为同策自蒸馏中的特权教师信号,彻底绕开了扩散模型无法像语言模型那样直接把目标图像拼入生成轨迹的死结。
- 回退 VLM 语言骨干权重的特征空间对齐技巧(Weight Reversion Trick):在将纯文本编码器 Qwen3-4B 升级为多模态编码器 Qwen3-VL-4B 以构建教师时,直接替换会因 VLM 后训练带来的分布漂移而在生成图像中引入严重高频伪影。作者仅保留 Qwen3-VL 的 ViT 与 Connector,而把其内部 LLM 权重换回原始 Qwen3-4B,这一零成本工程技巧非常适合迁移到任何基于预训练 LLM 扩展多模态条件控制的扩散模型任务中。
- 统一在线强化学习的轨迹一致性与监督学习的简易性:D-OPSD 在概念上打通了扩散模型后训练的四大维度——同策采样(On-Policy)、无需奖励模型(Reward-Free)、训练推理状态一致(Train-Inference Match)以及利用成对图文监督。这为所有少步生成模型(包括一致性模型、少步视频扩散模型)的后训练定制提供了一套通用范式。
局限与展望¶
- 单步训练计算开销增加:由于每次迭代需要执行学生模型的少步同策 rollout(如 4 或 8 步前向)以及教师分支在轨迹状态上的前向推理,D-OPSD 每步迭代的 FLOPs 约为普通 SFT 的 4 倍,实际挂钟训练时间约为普通 SFT 的 2 倍(尽管相比重新做步数蒸馏的数百 GPU 小时仍然极为廉价)。
- 强依赖基座编码器与教师的多模态上下文能力上限:D-OPSD 的成败完全取决于多模态教师分支能否在给定“文本+目标图像”时提供准确的语义与风格引导。如图 7 所示,当面对某些极其复杂或生僻的特定主体 ID,若教师在多模态条件下也无法生成与目标概念一致的隐式表征,自蒸馏就会因缺乏有效监督信号而训练失败;此外,这也要求扩散基座必须采用具备对应多模态版本(或原生多模态)的 LLM/VLM 编码器,难以直接用于早期仅采用纯 CLIP/T5 编码器的旧架构。
- 未来改进方向:一是引入图像编辑模型(如 FLUX.1 Kontext)或视频生成模型构建更强的教师端多模态条件上下文;二是将自蒸馏目标与中间层表征自对齐(如 SRA)等内部正则约束结合;三是探索多专家同策蒸馏(Multi-expert OPD),先分别训练多个垂直领域的专家模型,再通过 D-OPSD 框架将多专家能力无损蒸馏融合回单一少步基座模型中。
相关工作与启发¶
- vs Vanilla SFT / DreamBooth:传统微调方法在目标图像加噪得到的离线状态上回归真值速度场,严重破坏步数蒸馏模型的少步 ODE 动力学,导致少步推理图像模糊崩塌;D-OPSD 在学生自己的少步采样轨迹上优化并由同状态下的多模态教师速度场提供监督,完美保持了少步生成质量。
- vs PSO (Pairwise Sample Optimization):PSO 针对步数蒸馏模型引入类似 Diffusion-DPO 的成对样本优化,虽然只在少步采样时间步上训练,但其输入状态与真值速度信号依然来自离线目标图像,导致小样本下严重过拟合参考图而丧失对新场景提示词的泛化能力,在大规模全量微调下同样发生质量退化;D-OPSD 则是纯正的同策轨迹优化,兼具概念捕获与跨场景提示词泛化能力。
- vs 在线强化学习微调(如 Flow-GRPO / ReFL / DMD-RL):在线 RL 同样在同策少步轨迹上更新因而能保持少步能力,但必须依赖外部奖励模型或可微评分函数,无法直接从几张示例图文对中学习特定新概念;D-OPSD 用多模态上下文自蒸馏替代外部奖励,仅凭普通图文对即可完成定制与微调。
- vs 大语言模型同策自蒸馏(LLM OPSD,如 Self-Distilled Reasoner):LLM 中的 OPSD 通过在输入序列中直接拼接参考答案构造教师并最小化词表 KL 散度;D-OPSD 首次将该思想推广至连续流匹配扩散模型,通过编码器端的多模态特征注入解决“图像无法直接拼入去噪轨迹”的难题,并用同策状态上的速度场均方误差替代离散分布散度。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次发现并利用现代扩散模型编码器的零样本多模态上下文继承特性,巧妙解决了同策自蒸馏向步数蒸馏扩散模型迁移时的核心矛盾。
- 实验充分度: ⭐⭐⭐⭐⭐ 在 Z-Image-Turbo 和 FLUX.2-klein 两个主流少步模型上同时验证了小样本 LoRA 定制与 25K 大规模全参数微调,包含两阶段重蒸馏成本对比、四类范式消融与百组人工评测。
- 写作质量: ⭐⭐⭐⭐⭐ 痛点分析直击步数蒸馏模型微调的社区真问题,从 LLM OPSD 到扩散速度场对齐的数学映射清晰透彻,附录对编码器权重回退等工程细节坦诚详尽。
- 价值: ⭐⭐⭐⭐⭐ 极具实用价值,以单卡数小时的轻量成本解决了少步蒸馏模型微调必崩的行业痛点,无需奖励模型或二次蒸馏即可直接服务于开源社区的 LoRA 定制与领域适配。