用线性注意力改造扩散Transformer,高效生成高质量图像。
LiT: Delving into a Simple Linear Diffusion Transformer for Image Generation
- 采用轻量线性注意力与深度卷积结合,简化结构提升效率。
- 仅需20%~33%训练步数即可达到DiT同等图像生成效果。
- 适用于图像和文生图任务,可快速迁移至PixArt-$Σ$等模型。
本文探索如何将预训练的扩散Transformer(DiT)转化为纯线性注意力的LiT,以提升图像生成的简洁性、并行性和效率。通过系统研究,提出五条实用指导:1)在简单线性注意力中引入深度卷积足以支持图像生成;2)减少注意力头数可免费提升性能且不增加延迟;3)继承已完全收敛的预训练DiT权重;4)加载除线性注意力外的所有参数;5)采用混合知识蒸馏,用预训练教师模型同时监督噪声预测与反向扩散过程的方差。基于这些准则,提出线性扩散Transformer(LiT),作为纯线性注意力的高效基准。在256×256与512×512图像生成任务中,LiT分别仅需20%和33%的训练步数即可媲美原DiT性能,且优于基于Mamba或门控线性注意力的方法。该方法还可推广至文本到图像生成:从PixArt-Σ快速转换为LiT,生成高质量图像并保持相近GenEval得分。
原文摘要 · Abstract (English)
In this paper, we investigate how to convert a pre-trained Diffusion Transformer (DiT) into a linear DiT, as its simplicity, parallelism, and efficiency for image generation. Through detailed exploration, we offer a suite of ready-to-use solutions, ranging from linear attention design to optimization strategies. Our core contributions include 5 practical guidelines: 1) Applying depth-wise convolution within simple linear attention is sufficient for image generation. 2) Using fewer heads in linear attention provides a free-lunch performance boost without increasing latency. 3) Inheriting weights from a fully converged, pre-trained DiT. 4) Loading all parameters except those related to linear attention. 5) Hybrid knowledge distillation: using a pre-trained teacher DiT to help the training of the student linear DiT, supervising not only the predicted noise but also the variance of the reverse diffusion process. These guidelines lead to our proposed \underline{L}inear D\underline{i}ffusion \underline{T}ransformer (LiT), which serves as a safe and efficient alternative baseline for DiT with pure linear attention. In class-conditional 256$\times$256 and 512$\times$512 ImageNet generation, LiT can be quickly adapted from DiT using only $20\%$ and $33\%$ of DiT's training steps, respectively, while achieving comparable performance. LiT also rivals methods based on Mamba or Gated Linear Attention. Moreover, the same guidelines generalize to text-to-image generation: LiT can be swiftly converted from PixArt-$Σ$ to generate high-quality images, maintaining comparable GenEval scores.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。