用历史轨迹预测未来步骤,让扩散模型生成更快更准。
Autoregressive Distillation of Diffusion Transformers
- 用历史轨迹替代最新样本,缓解误差累积问题。
- 4步生成图像FID达1.84,比基线减少5倍退化。
- 适合需要快速高质生成的文本到图像应用。
基于Transformer架构的扩散模型在生成高保真、高分辨率图像方面表现出色,但其迭代采样过程资源消耗大。已有研究尝试将概率流ODE解压缩为少步学生模型,但受限于依赖最新去噪样本输入,易受暴露偏差影响。为此,我们提出自回归蒸馏(ARD),利用ODE轨迹的历史信息预测未来步骤。该方法通过在教师Transformer中添加逐标记时间嵌入,并采用分块因果注意力掩码进行训练,仅在低层变压器中引入历史输入以提升性能与效率。在ImageNet和T2I生成任务上验证,模型在ImageNet-256上仅需1.1%额外计算量,即实现5倍的FID退化降低;4步生成时FID为1.84,且在提示遵循度上优于公开的1024p文本到图像蒸馏模型,同时保持接近教师模型的FID表现。
原文摘要 · Abstract (English)
Diffusion models with transformer architectures have demonstrated promising capabilities in generating high-fidelity images and scalability for high resolution. However, iterative sampling process required for synthesis is very resource-intensive. A line of work has focused on distilling solutions to probability flow ODEs into few-step student models. Nevertheless, existing methods have been limited by their reliance on the most recent denoised samples as input, rendering them susceptible to exposure bias. To address this limitation, we propose AutoRegressive Distillation (ARD), a novel approach that leverages the historical trajectory of the ODE to predict future steps. ARD offers two key benefits: 1) it mitigates exposure bias by utilizing a predicted historical trajectory that is less susceptible to accumulated errors, and 2) it leverages the previous history of the ODE trajectory as a more effective source of coarse-grained information. ARD modifies the teacher transformer architecture by adding token-wise time embedding to mark each input from the trajectory history and employs a block-wise causal attention mask for training. Furthermore, incorporating historical inputs only in lower transformer layers enhances performance and efficiency. We validate the effectiveness of ARD in a class-conditioned generation on ImageNet and T2I synthesis. Our model achieves a $5\times$ reduction in FID degradation compared to the baseline methods while requiring only 1.1\% extra FLOPs on ImageNet-256. Moreover, ARD reaches FID of 1.84 on ImageNet-256 in merely 4 steps and outperforms the publicly available 1024p text-to-image distilled models in prompt adherence score with a minimal drop in FID compared to the teacher. Project page: https://github.com/alsdudrla10/ARD.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。