将扩散Transformer的注意力机制改为线性复杂度,大幅加速高分辨率图像生成。
CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up
- 提出类似卷积的局部注意力策略,实现线性计算复杂度。
- 仅用10K样本微调10K步,生成效果接近原模型,推理提速6.3倍。
- 适合需要高效生成高分辨率图像的研究者与开发者。
扩散Transformer(DiT)已成为图像生成的主流架构。然而,其注意力机制带来的二次方复杂度在生成高分辨率图像时导致显著延迟。本文旨在实现预训练DiT的线性注意力机制,将复杂度降至线性。通过系统分析现有高效注意力方法,识别出四个关键因素:局部性、公式一致性、高秩注意力图和特征完整性。基于此,提出名为CLEAR的卷积类局部注意力策略,限制每个查询令牌的特征交互范围在局部窗口内,从而实现线性复杂度。实验表明,仅需对注意力层在10,000个自生成样本上微调10,000次迭代,即可有效将知识从教师模型迁移到线性复杂度的学生模型,生成效果与教师模型相当。同时,注意力计算量减少99.5%,8K分辨率图像生成速度提升6.3倍。此外,蒸馏后的注意力层展现出零样本跨模型泛化能力及多GPU并行推理支持等优势。代码与模型已开源:https://github.com/Huage001/CLEAR。
原文摘要 · Abstract (English)
Diffusion Transformers (DiT) have become a leading architecture in image generation. However, the quadratic complexity of attention mechanisms, which are responsible for modeling token-wise relationships, results in significant latency when generating high-resolution images. To address this issue, we aim at a linear attention mechanism in this paper that reduces the complexity of pre-trained DiTs to linear. We begin our exploration with a comprehensive summary of existing efficient attention mechanisms and identify four key factors crucial for successful linearization of pre-trained DiTs: locality, formulation consistency, high-rank attention maps, and feature integrity. Based on these insights, we introduce a convolution-like local attention strategy termed CLEAR, which limits feature interactions to a local window around each query token, and thus achieves linear complexity. Our experiments indicate that, by fine-tuning the attention layer on merely 10K self-generated samples for 10K iterations, we can effectively transfer knowledge from a pre-trained DiT to a student model with linear complexity, yielding results comparable to the teacher model. Simultaneously, it reduces attention computations by 99.5% and accelerates generation by 6.3 times for generating 8K-resolution images. Furthermore, we investigate favorable properties in the distilled attention layers, such as zero-shot generalization cross various models and plugins, and improved support for multi-GPU parallel inference. Models and codes are available here: https://github.com/Huage001/CLEAR.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。