SD3.5-Flash: Distribution-Guided Distillation of Generative Flows
Hmrishav Bandyopadhyay, Rahim Entezari, Jim Scott, Reshinth Adithyan, Yi-Zhe Song, Varun Jampani
cs.CV, cs.AI
2025-09-26
Stability AI 把 SD3.5 Medium 蒸馏成 2 步和 4 步学生，用 timestep sharing 改写 DMD。8-bit 去掉 T5 后峰值显存 6.61GB，四步比五十步教师快约 18 倍，用户研究画质压过教师。
SD3.5 这类 rectified flow 模型把噪声到图像的路径拉成近似直线，画质上去了，推理还停在机房。论文给的门槛是 25 步以上、16GB 以上显存、单张 30 秒以上。消费级显卡和手机接不住。
少步蒸馏是条出路。DMD（Distribution Matching Distillation）用师生分布 KL 散度的梯度来训练少步学生，不必逐步对齐整条轨迹。它在普通 diffusion 上能用，搬到 flow 上会脏：标准做法从轨迹终点 x0 再加噪声到随机 t，再算 score。rectified flow 的噪声和图像是配对的直线，乱加噪声等于把样本踢出 ODE 路径，速度场预测失真。少步时后面没有余步纠错，质量会塌。步数压下去后模型容量也不够，审美和 prompt 对齐互相抢参数。SD3.5 还挂着 T5-XXL 文本编码器，它自己就占掉大部分峰值显存。
教师是 2.5B 的 SD3.5 Medium。训练图由 8B 的 SD3.5 Large 合成，32 步、CFG 4.0。8 张 H100 上先预训练 2000 iteration（约 17 小时），再进入蒸馏。
先用轨迹引导损失，在学生自己的时间步上拟合教师沿这段轨迹积出来的位移。4 步模型就只有 4 个点。这一步把学生拉到教师 ODE 附近，后面的分布匹配才站得住。
然后上改过的 DMD。改动叫 timestep sharing：算匹配时不再从 x0 随机加噪，直接拿学生轨迹上已经部分去噪的 xt。少步学生在高噪声处估不好 x0，用假 x0 去算 score 会污染梯度；共享时间步等于只在模型真正会走的那几个 t 上对齐。时间步多样性变少，消融里构图和纹理反而更好。
prompt 对齐用 split-timestep fine-tuning 补。把 checkpoint 复制成两支，一支训 t 落在 (0, 500]，一支训 (500, 1000]，EMA 衰减 0.99，收敛后按 3:7（低噪声支:高噪声支）插值合并，比例用 GenEval 选。只在 4 步模型上做，2 步没观察到同样的跳变。
对抗项跟 NitroFusion 同一路：proxy 学生当特征提取器，从 MM-DiT 的第 3、4、5、6、8、10、11 层接 8 层 MLP 判别头，教师合成图当 real。判别头每个 iteration 以 0.005 的概率重置，防止过拟合；判别器和 proxy 每更新 10 次，生成器才更新 1 次。2 步从 4 步继续蒸，额外加师生特征 Gram 矩阵的 MSE。附录里 4 步第二阶段 800 iteration（6 小时），2 步再 1200 iteration（9 小时），split-timestep 再 400 iteration（4 小时）。
部署侧把 T5-XXL 换成空 embedding，SD3.5 预训练做过 encoder dropout，只留 CLIP-L 和 CLIP-G。MM-DiT 从 16-bit 量化到 8-bit，Apple 端再用 CoreML 压到 6-bit，并改写 RMSNorm 以保住 Neural Engine 精度。
COCO-30K 自动指标，RTX 4090：
| 方法 | 步数 | 延迟 | 峰值显存 | ImageReward | GenEval |
| SD3.5M 教师 | 50 | 10.58s | 19.47GB | 0.91 | 0.64 |
| SWD-M | 4 | 0.66s | 17.88GB | 1.12 | 0.72 |
| Flash 16-bit+T5 | 4 | 0.58s | 17.58GB | 1.10 | 0.70 |
| Flash 8-bit 无 T5 | 4 | 0.61s | 6.61GB | 1.08 | 0.68 |
| SANA-Sprint 1.6B | 2 | 0.24s | 10.17GB | 1.01 | 0.73 |
| Flash 16-bit+T5 | 2 | 0.39s | 17.58GB | 1.00 | 0.70 |
相对 50 步教师，4 步 16-bit 大约 18 倍加速（10.58 秒 / 0.58 秒）。ImageReward 从 0.91 到 1.10，GenEval 从 0.64 到 0.70，Aesthetic Score 从 5.99 到 6.38，三项都超过教师。FID 反过来更差：Flash 29.80，教师 20.06，SDXL-DMD2 只有 16.64。论文把原因归到教师家族：SD3.5M 自己的 FID 就比 SDXL 差，同门 SWD 和 Turbo 也偏高。
同门 4 步基线 SWD-M 的 GenEval 是 0.72、IR 是 1.12，两项都略高于 Flash。2 步上 SANA-Sprint 1.6B 延迟更短、GenEval 和 IR 也略高。自动指标上 Sprint 没输。
用户研究找了 124 名有排图经验的标注者，507 条 prompt、4 个 seed，每对图 3 人投票。画质和 prompt 对齐拆开评：评画质时不给 prompt，评对齐时忽略美观。画质上 Flash 压过其他少步模型和 50 步教师；对齐一项所有方法差距不到 ±1.6%，基本打平。各消费级算力档的 ELO 图里 Flash 变体都在最上面，正文没给具体分值。
端侧 4 步、6-bit、无 T5：iPhone A17 在 768px 要 8.32 秒，512px 要 3.25 秒；iPad M4 对应 6.44 秒和 2.62 秒。16-bit 全精度 1024px 在 32GB 的 M3 笔记本上要 18.65 秒。
消融只给了定性图：去掉对抗项训练会散；去掉预训练颜色和构图先崩；不用 timestep sharing 纹理和构图变差；不刷新判别头会过平滑。没有对应的表格数字。
给想在消费卡和手机上跑 SD3.5 画风的人，这是一条能落地的工程路径。4 步、去掉 T5、8-bit，峰值显存从约 18GB 掉到 6.61GB，4090 上 1024px 大约 0.6 秒出图。timestep sharing 是对 DMD 在直线流上的针对性修补，后面蒸 SD3 或 Flux 这类模型可以借。split-timestep 只在训练期扩容，推理不加参数。
别把它读成少步生成的新第一名。同门 SWD-M 的 GenEval 和 IR 略高；2 步自动指标被 SANA-Sprint 压过。卖点是 SD3.5 家族在显存梯度和端侧延迟上的完整套件，不是单点分数。
论文自己写了：蒸馏会牺牲复杂任务上的质量和多样性；去掉 T5 后复杂构图变难。FID 在 COCO 上明显差于 SDXL 系蒸馏。即使用「教师 FID 基因」解释，这套模型也更偏审美分数，不像真实照片分布。
用户研究把画质和对齐拆开，画质赢了、对齐持平。这和「构图理解更强」的定性宣传咬得不紧。消融没有定量。没和 SD3.5 Large、Flux.1-dev 比，理由是消费硬件装不下教师，少步学生之间仍缺这条对照。训练数据来自 8B Large 合成图，教师却是 2.5B Medium，学生有一部分能力可能来自 Large 的数据分布，不完全是 Medium 蒸馏。