DiTFastAttn: Attention Compression for Diffusion Transformer Models
针对扩散 Transformer 的注意力压缩后训练方法，通过三种冗余消除技术实现高达 76% FLOPs 减少
DiTFastAttn: Attention Compression for Diffusion Transformer Models
一、论文概述
| 项目 | 内容 |
|---|---|
| 标题 | DiTFastAttn: Attention Compression for Diffusion Transformer Models |
| 作者 | Zhihang Yuan, Hanling Zhang, Pu Lu, Xuefei Ning, Linfeng Zhang, Tianchen Zhao, Shengen Yan, Guohao Dai, Yu Wang |
| 机构 | Tsinghua University, Infinigence AI, Shanghai Jiao Tong University |
| 论文 | https://arxiv.org/abs/2406.08552 |
| 代码 | 未公开 |
| 发布 | 2024-06-12 (v1), 2024-10-18 (v2) |
| 许可 | 未明确 |
二、核心思想
问题定义
扩散 Transformer (DiT) 在图像和视频生成方面表现出色，但由于自注意力算子的二次复杂度面临计算挑战。在高分辨率生成时，注意力计算成为主要瓶颈。
三种冗余识别
DiTFastAttn 识别出 DiT 推理过程中注意力计算的三种关键冗余：
- 空间冗余 (Spatial Redundancy)：许多注意力头专注于局部信息，对远距离 token 的注意力值接近零
- 时间冗余 (Temporal Redundancy)：相邻步骤之间的注意力输出具有高相似度
- 条件冗余 (Conditional Redundancy)：条件推理和无条件推理之间表现出显著相似度（SSIM ≥ 0.95）
解决方案概述
DiTFastAttn 是一种后训练压缩方法，提出三种互补技术来消除上述冗余：
- Window Attention with Residual Sharing (WA-RS)：用窗口注意力替代全注意力处理空间冗余
- Attention Sharing across Timesteps (AST)：跨时间步共享注意力输出利用时间冗余
- Attention Sharing across CFG (ASC)：跨 CFG 条件/无条件推理共享注意力利用条件冗余
三、技术架构
核心方法
1. Window Attention with Residual Sharing (WA-RS)
针对空间冗余：许多注意力头只关注局部区域，对远距离 token 的注意力接近零。
- 窗口注意力 (Window Attention)：将全注意力替换为窗口注意力，只计算局部窗口内的注意力
- 残差共享 (Residual Sharing)：直接丢弃窗口外的注意力计算会导致信息损失。WA-RS 通过缓存并共享窗口外的残差部分来弥补：
- 计算全注意力与窗口注意力的差值作为残差
- 在后续步骤中复用该残差，避免重复计算
2. Attention Sharing across Timesteps (AST)
针对时间冗余：相邻去噪步骤的注意力输出高度相似。
- 观察：相邻步骤的注意力输出 SSIM 值很高
- 方法：缓存前一步的注意力输出，当下一步的相似度超过阈值时直接复用
- 适用层：主要用于中间层和后期层
3. Attention Sharing across CFG (ASC)
针对条件冗余：Classifier-Free Guidance (CFG) 中条件和无条件推理的注意力输出高度相似。
- 观察：某些头和时间步的条件/无条件注意力输出 SSIM ≥ 0.95
- 方法：先执行条件推理，缓存注意力输出；无条件推理时直接复用
- 效果：减少 50% 的注意力计算（CFG 需要两次推理）
压缩计划搜索
自适应压缩策略：
不同模型和不同层的冗余模式不同，不存在通用的压缩策略。DiTFastAttn 通过搜索最优压缩计划：
- 采样校准：使用少量样本评估不同压缩配置
- 逐层决策：对每一层独立选择最优压缩技术（WA-RS/AST/ASC 或不压缩）
- 阈值优化：通过控制相似度阈值调节压缩比例
评估指标：使用注意力输出的相对误差作为压缩质量指标（而非 LPIPS 或 SSIM，因为计算速度更快）
四、核心创新
| 创新点 | 说明 | 理论/实验依据 |
|---|---|---|
| 三种冗余识别 | 系统性识别空间、时间、条件三种冗余 | SSIM ≥ 0.95 的条件冗余，局部注意力模式的空间冗余 |
| WA-RS 技术 | 窗口注意力+残差共享，避免信息损失 | 残差共享比直接窗口注意力质量更高 |
| ASC 技术 | 利用 CFG 的条件/无条件相似度 | 50% FLOPs 减少（CFG 需要两次推理） |
| 自适应压缩计划 | 逐层搜索最优压缩配置 | 不同模型冗余分布不同，需要定制化策略 |
| 后训练方法 | 无需重新训练或微调 | 即插即用，与量化等其他方法正交 |
五、实验结果
图像生成性能
DiT-XL-2 (512×512)：
| 配置 | IS↑ | FID↓ | Attn FLOPs |
|---|---|---|---|
| Raw | 408.16 | 25.43 | 100% |
| D1 | 412.24 | 25.32 | 85% |
| D2 | 412.18 | 24.67 | 69% |
| D3 | 411.74 | 23.76 | 59% |
| D4 | 391.80 | 21.52 | 49% |
| D5 | 370.07 | 19.32 | 41% |
PixArt-Sigma-XL (1024×1024)：
| 配置 | IS↑ | FID↓ | CLIP↑ | Attn FLOPs |
|---|---|---|---|---|
| Raw | 24.33 | 55.65 | 31.27 | 100% |
| D1 | 24.27 | 55.73 | 31.27 | 90% |
| D2 | 24.25 | 55.69 | 31.26 | 74% |
| D3 | 24.16 | 55.61 | 31.25 | 63% |
| D4 | 24.07 | 55.32 | 31.24 | 52% |
PixArt-Sigma-XL (2048×2048)：
| 配置 | IS↑ | FID↓ | CLIP↑ | Attn FLOPs |
|---|---|---|---|---|
| Raw | 23.67 | 51.89 | 31.47 | 100% |
| D3 | 22.96 | 52.54 | 31.36 | 46% |
| D4 | 22.95 | 51.74 | 31.39 | 36% |
关键发现：
- D1-D3 配置几乎匹配原始模型性能
- 随着分辨率增加，DiTFastAttn 不仅实现更大压缩，还更好地保持生成性能
- 高分辨率（2K）下可减少高达 76% 的注意力 FLOPs
注意力压缩效果
各技术独立效果：
| 模型 | 序列长度 | 技术 | Attn FLOPs | Attn Latency |
|---|---|---|---|---|
| DiT-XL-2 (512) | 1024 | ASC | 50% | 59% |
| DiT-XL-2 (512) | 1024 | WA-RS | 77% | 85% |
| DiT-XL-2 (512) | 1024 | WA-RS+ASC | 38% | 51% |
| PixArt (1024) | 4096 | ASC | 50% | 54% |
| PixArt (1024) | 4096 | WA-RS | 51% | 54% |
| PixArt (1024) | 4096 | WA-RS+ASC | 26% | 31% |
| PixArt (2048) | 16384 | ASC | 50% | 52% |
| PixArt (2048) | 16384 | WA-RS | 33% | 35% |
| PixArt (2048) | 16384 | WA-RS+ASC | 16% | 19% |
关键观察：
- ASC 固定减少 50% FLOPs（CFG 需要两次推理）
- WA-RS 的压缩效果随分辨率增加而提升（77% → 33%）
- 两种技术正交，可同时使用
端到端加速
- 高分辨率 (2K×2K)：实现 1.8× 端到端加速
- 注意力延迟：WA-RS+ASC 在 2K 分辨率下将注意力延迟降至 19%
视频生成
应用于 OpenSora V1.1 (240p, 16 frames)：
- D1-D6 配置的计算减少：7.63%, 19.50%, 30.16%, 37.66%, 40.52%
- 视觉质量在低压缩配置下保持良好
消融实验
- DiTFastAttn 组合方法优于任何单一方法
- 不同时间步需要不同的压缩策略
- 残差共享 (RS) 比直接窗口注意力质量更高
- 压缩计划搜索对不同模型需要定制化
六、相关工作
| 方法类别 | 代表方法 | 与 DiTFastAttn 对比 |
|---|---|---|
| 量化 | Q-DiT, PTQD | 正交，可同时使用 |
| 调度器优化 | DPM-Solver, DDIM | 正交，不影响 |
| 特征缓存 | DeepCache | DiTFastAttn 关注注意力层 |
| Token 合并 | ToMe | 不同维度的优化 |
| 稀疏注意力 | 各种稀疏注意力方法 | DiTFastAttn 是后训练方法 |
七、总结
核心贡献
- 系统性冗余识别：识别 DiT 推理中的三种关键冗余（空间/时间/条件）
- 三种互补技术：WA-RS、AST、ASC 分别消除三种冗余
- 自适应压缩计划：逐层搜索最优压缩配置，适应不同模型
- 显著效果：减少高达 76% 注意力 FLOPs，高分辨率下 1.8× 端到端加速
- 后训练方法：无需重新训练，即插即用
技术影响
- 实用性：后训练方法可直接应用于已部署的 DiT 模型
- 正交性：与量化、调度器优化等方法正交，可组合使用
- 可扩展性：在更高分辨率下效果更好
局限性
- 代码未公开
- 压缩计划搜索需要校准时间和计算
- 极高压缩比下质量有一定下降
- 主要在图像生成上验证，视频生成效果待进一步探索
八、参考资源
- 论文: https://arxiv.org/abs/2406.08552
- 相关工作:
- DiT: Scalable Diffusion Models with Transformers
- PixArt-Sigma: 弱到强扩散 Transformer
- OpenSora: 开源视频生成
- FlashAttention-2: 高效注意力实现