改进MeanFlow稳定性与泛化性,实现大模型高效扩散蒸馏
Stabilizing, Scaling & Enhancing MeanFlow for Large-scale Diffusion Distillation

- 用离散解替代微分解缓解训练崩溃,分阶段优化提升稳定性
- 引入轨迹分布对齐,解决极少数步采样下的均值偏差问题
- 在12B到80B参数的大模型上均表现优异,适合工业级扩散模型蒸馏
扩散模型生成能力强大,但推理延迟高限制实际应用。许多研究尝试减少采样步数以加速推理,其中MeanFlow因公式简洁且性能出色备受关注。然而其优化目标不稳及‘均值导向偏差’限制了在大规模工业模型蒸馏中的应用。为稳定大规模模型的MeanFlow蒸馏,我们提出一种热身策略:先用离散解替代原微分解,避免因教师模型未充分训练导致带停止梯度的目标引发训练崩溃;待学生模型初步拟合平均速度场后,再切换回微分解进行精细优化。同时,为缓解极少数步采样下复杂目标分布带来的均值偏差,引入轨迹分布对齐作为辅助目标,促使学生模型轨迹分布更贴近教师模型。该方法在文本到图像模型FLUX.1-dev(最高达120亿参数)上优于现有蒸馏方法;扩展至800亿参数的SOTA模型HunyuanImage 3.0时,仍保持强泛化能力与优异性能。
原文摘要 · Abstract (English)
Diffusion models exhibit remarkable generative capability, but their high latency limits practical deployment. Many studies have attempted to reduce sampling steps to accelerate inference. Among them, MeanFlow has attracted considerable attention due to its concise formulation and remarkable performance. Nevertheless, the instability of its optimization objective and the ''mean-seeking bias'' have limited its applicability to distill large-scale industrial models. To stabilize MeanFlow for distilling large-scale models, we first introduce a warm-up technique, in which the original differential solution of MeanFlow is replaced by a discrete solution. This design avoids training collapse caused by the MeanFlow target containing a stop-gradient term from an undertrained model. Once the model acquires a preliminary ability to fit the average velocity field, we switch the optimization objective back to the differential solution, enabling further refinement. Meanwhile, to alleviate the ''mean-seeking bias'' of MeanFlow under extremely few-step inference with complex target distributions, we incorporate trajectory distribution alignment as an auxiliary objective, encouraging the student model's trajectory distribution to align more closely with that of the teacher model. Our proposed distillation framework achieves superior performance compared to existing distillation approaches when applied to the text-to-image (T2I) model FLUX.1-dev (up to 12B parameters). Furthermore, when extended to the 80B-parameter state-of-the-art (SOTA) T2I model HunyuanImage 3.0, our method continues to demonstrate robust generalization and strong performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。