SenseFlow: Scaling Distribution Matching for Flow-based Text-to-Image Distillation

📄 arXiv: 2506.00523v1 📥 PDF

作者: Xingtong Ge, Xin Zhang, Tongda Xu, Yi Zhang, Xinjie Zhang, Yan Wang, Jun Zhang

分类: cs.CV

发布日期: 2025-05-31

备注: under review

🔗 代码/项目: GITHUB

💡 一句话要点

SenseFlow：通过缩放分布匹配实现Flow模型文本到图像的蒸馏

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 文本到图像生成 蒸馏训练 Flow模型 分布匹配 隐式分布对齐 段内指导 模型压缩

📋 核心要点

  1. 现有DMD方法在应用于大规模Flow模型（如SD 3.5和FLUX）时，面临收敛困难，限制了其在大规模模型上的应用。
  2. 论文提出隐式分布对齐(IDA)和段内指导(ISG)两种方法，分别用于规范生成器分布和重定位时间步重要性，以解决收敛问题。
  3. 实验结果表明，SenseFlow在SDXL、SD 3.5 Large和FLUX等模型上实现了卓越的蒸馏性能，验证了所提出方法的有效性。

📝 摘要（中文）

分布匹配蒸馏(DMD)已成功应用于Stable Diffusion (SD) 1.5等文本到图像扩散模型。然而，原始DMD在大规模基于Flow的文本到图像模型（如SD 3.5和FLUX）上存在收敛困难。本文首先分析了将原始DMD应用于大规模模型时的问题。然后，为了克服可扩展性挑战，我们提出了隐式分布对齐(IDA)来规范生成器和伪分布之间的距离。此外，我们提出了段内指导(ISG)来重新定位教师模型的时间步重要性分布。仅使用IDA，DMD即可收敛于SD 3.5；同时使用IDA和ISG，DMD可收敛于SD 3.5和FLUX.1 dev。结合其他改进，如扩大判别器模型，我们最终的模型，称为 extbf{SenseFlow}，在基于扩散的文本到图像模型（如SDXL）和Flow匹配模型（如SD 3.5 Large和FLUX）的蒸馏中都实现了卓越的性能。源代码将在https://github.com/XingtongGe/SenseFlow上提供。

🔬 方法详解

问题定义：论文旨在解决将分布匹配蒸馏(DMD)应用于大规模Flow-based文本到图像模型时遇到的收敛困难问题。现有DMD方法在处理如SD 3.5和FLUX等大型模型时，由于模型复杂性和训练数据规模的增加，容易出现训练不稳定和收敛速度慢的问题，导致蒸馏效果不佳。

核心思路：论文的核心思路是通过引入隐式分布对齐(IDA)和段内指导(ISG)来稳定训练过程并提高蒸馏效果。IDA通过规范生成器和伪分布之间的距离，使得生成器能够更好地学习教师模型的分布。ISG则通过重新定位时间步重要性分布，使得学生模型能够更好地关注重要的时间步，从而提高生成质量。

技术框架：SenseFlow的整体框架基于DMD，并在此基础上添加了IDA和ISG模块。首先，使用教师模型生成伪数据。然后，利用IDA来约束生成器和伪数据分布之间的距离。同时，利用ISG来调整时间步重要性分布。最后，通过判别器来区分真实数据和生成数据，从而训练生成器。

关键创新：论文的关键创新在于提出了IDA和ISG两种方法。IDA通过隐式地对齐生成器和伪数据分布，避免了直接计算分布距离的困难，从而提高了训练的稳定性。ISG则通过重新定位时间步重要性分布，使得学生模型能够更好地学习教师模型的生成过程。

关键设计：IDA通过最小化生成器输出和伪数据之间的距离来实现，可以使用各种距离度量，如L1或L2距离。ISG通过调整时间步的权重来实现，可以使用各种权重调整策略，如基于教师模型输出的梯度或方差来调整权重。此外，论文还使用了更大的判别器模型来提高判别能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SenseFlow在SDXL、SD 3.5 Large和FLUX等模型上进行了实验，结果表明，该方法能够显著提高蒸馏模型的性能。例如，在SD 3.5 Large上，SenseFlow能够实现与教师模型相媲美的生成质量，同时显著减少了模型大小和计算量。实验结果验证了IDA和ISG的有效性，并表明SenseFlow是一种有效的文本到图像蒸馏方法。

🎯 应用场景

SenseFlow的研究成果可应用于各种文本到图像生成任务，尤其是在资源受限的场景下，可以通过蒸馏将大型模型压缩为小型模型，从而在移动设备或边缘设备上部署高质量的图像生成应用。此外，该方法还可以用于加速新模型的训练过程，提高开发效率。

📄 摘要（原文）

The Distribution Matching Distillation (DMD) has been successfully applied to text-to-image diffusion models such as Stable Diffusion (SD) 1.5. However, vanilla DMD suffers from convergence difficulties on large-scale flow-based text-to-image models, such as SD 3.5 and FLUX. In this paper, we first analyze the issues when applying vanilla DMD on large-scale models. Then, to overcome the scalability challenge, we propose implicit distribution alignment (IDA) to regularize the distance between the generator and fake distribution. Furthermore, we propose intra-segment guidance (ISG) to relocate the timestep importance distribution from the teacher model. With IDA alone, DMD converges for SD 3.5; employing both IDA and ISG, DMD converges for SD 3.5 and FLUX.1 dev. Along with other improvements such as scaled up discriminator models, our final model, dubbed \textbf{SenseFlow}, achieves superior performance in distillation for both diffusion based text-to-image models such as SDXL, and flow-matching models such as SD 3.5 Large and FLUX. The source code will be avaliable at https://github.com/XingtongGe/SenseFlow.