用单层Transformer替代28层模型，实现更快更稳的一步生成。
A one-step generation model with a Single-Layer Transformer: Layer number re-distillation of FreeFlow
- 将28层Transformer重构成单层结构，通过深度特征蒸馏逼近原模型。
- 参数量从6.75亿降至430万，每秒可筛查超100个噪声点。
- 适合追求高速高质生成的用户，尤其提升低采样次数下的稳定性。
当前流匹配方法致力于将扩散模型的迭代生成过程压缩至少数甚至一步完成，其中基于常微分方程（ODE）的MeanFlow和FreeFlow是代表性成果。我们观察到FreeFlow的28层Transformer架构可被视作沿深度轴的欧拉离散化方案，层索引即为离散时间步。据此，我们沿用FreeFlow的推导逻辑，对层数量进行再蒸馏，提出SLT（Single-Layer Transformer），使用单一共享DiT模块近似原28层教师模型的深度特征演化。训练中，模型在多个深度片段处匹配教师中间特征，融合片段表示，并同时对齐教师最终速度预测。经蒸馏训练，将教师模型DiT-XL/2的28个独立Transformer块压缩为单个模块，参数量由675M降至4.3M。凭借极小参数与快速采样能力，SLT可在相同时间内筛选更多候选噪声点，从而为教师模型FreeFlow选出更优初始点，显著提升生成图像质量。实验表明，在与两次随机采样相当的时间预算下，本方法可执行逾100次噪声筛查，并利用优选点生成高质量样本。有效规避了有限采样次数下因低质初始噪声导致的质量波动，大幅提高一步生成的稳定性和平均质量。
原文摘要 · Abstract (English)
Currently, Flow matching methods aim to compress the iterative generation process of diffusion models into a few or even a single step, with MeanFlow and FreeFlow being representative achievements of one-step generation based on Ordinary Differential Equations (ODEs). We observe that the 28-layer Transformer architecture of FreeFlow can be characterized as an Euler discretization scheme for an ODE along the depth axis, where the layer index serves as the discrete time step. Therefore, we distill the number of layers of the FreeFlow model, following the same derivation logic as FreeFlow, and propose SLT (Single-Layer Transformer), which uses a single shared DiT block to approximate the depth-wise feature evolution of the 28-layer teacher. During training, it matches the teacher's intermediate features at several depth patches, fuses those patch-level representations, and simultaneously aligns the teacher's final velocity prediction. Through distillation training, we compress the 28 independent Transformer Blocks of the teacher model DiT-XL/2 into a single Transformer Block, reducing the parameter count from 675M to 4.3M. Furthermore, leveraging its minimal parameters and rapid sampling speed, SLT can screen more candidate points in the noise space within the same timeframe, thereby selecting higher-quality initial points for the teacher model FreeFlow and ultimately enhancing the quality of generated images. Experimental results demonstrate that within a time budget comparable to two random samplings of the teacher model, our method performs over 100 noise screenings and produces a high-quality sample through the teacher model using the selected points. Quality fluctuations caused by low-quality initial noise under a limited number of FreeFlow sampling calls are effectively avoided, substantially improving the stability and average generation quality of one-step generation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。