1. 论文基本信息

1.1. 标题

DiTFastAttn: Attention Compression for Diffusion Transformer Models (DiTFastAttn：扩散 Transformer 模型中的注意力压缩)

1.2. 作者

Zhihang Yuan*1,2 Hanling Zhang*1,2 Pu Lu*1 Xuefei Ning1 Linfeng Linfeng Zhang3 Tianchen Zhao1,2 Shengen Yan2 Guohao Dai3,2 Yu Wang1 *表示贡献同等。 作者来自清华大学 (Tsinghua University)、Infinigence AI 和上海交通大学 (Shanghai Jiao Tong University)。

1.3. 发表期刊/会议

arXiv 预印本。

1.4. 发表年份

2024年6月12日 (UTC)。

1.5. 摘要

扩散 Transformer (DiT) 模型在图像和视频生成方面表现出色，但由于自注意力 (self-attention) 算子的二次复杂度，面临巨大的计算挑战。本文提出了 DiTFastAttn，一种后训练 (post-training) 压缩方法，旨在缓解 DiT 的计算瓶颈。研究识别了 DiT 推理过程中注意力计算的三个关键冗余：(1) 空间冗余 (spatial redundancy)，许多注意力头 (attention heads) 关注局部信息；(2) 时间冗余 (temporal redundancy)，相邻时间步 (timesteps) 的注意力输出高度相似；(3) 条件冗余 (conditional redundancy)，条件和无条件推理表现出显著相似性。为减少这些冗余，本文提出了三项技术：(1) 带有残差共享的窗口注意力 (Window Attention with Residual Sharing, WA-RS) 以减少空间冗余；(2) 跨时间步注意力共享 (Attention Sharing across Timesteps, AST) 以利用时间步之间的相似性；(3) 跨分类器自由引导注意力共享 (Attention Sharing across CFG, ASC) 以跳过条件生成过程中的冗余计算。DiTFastAttn 被应用于 DiT 和 PixArt-Sigma 进行图像生成任务，以及 OpenSora 进行视频生成任务。结果表明，对于图像生成，该方法可将注意力浮点运算次数 (FLOPs) 减少高达 76%，并在高分辨率 (2k x 2k) 生成中实现高达 1.8 倍的端到端加速 (end-to-end speedup)。

1.6. 原文链接

• 原文链接: https://arxiv.org/abs/2406.08552
• PDF 链接: https://arxiv.org/pdf/2406.08552v2.pdf

2. 整体概括

2.1. 研究背景与动机

近年来，扩散 Transformer (Diffusion Transformers, DiT) 模型在图像和视频生成领域取得了显著进展，例如 DiT (Peebles & Xie, 2023)、PixArt-Sigma (Chen et al., 2024) 和 Sora (Brooks et al., 2024)。然而，DiT 模型面临一个主要挑战：其巨大的计算需求，尤其是在生成高分辨率内容时。

核心问题： Transformer 架构中的自注意力 (self-attention) 机制具有对输入词元 (token) 长度 $L$ 的 $\mathcal{O}(L^2)$ 二次复杂度。随着图像和视频分辨率的提高，词元长度 $L$ 急剧增加，导致注意力计算成为推理 (inference) 过程中的主要计算瓶颈 (computational bottleneck)。例如，生成一张 2K x 2K 的图像可能需要处理 16k 个词元，即使在高性能 GPU 上，注意力计算也需数秒。此外，扩散模型的推理过程通常需要大量神经网络推理，这归因于多个去噪步骤和分类器自由引导 (Classifier-Free Guidance, CFG) 技术，进一步加剧了计算负担。

现有研究的挑战或空白： 尽管已有许多工作致力于加速注意力机制，如局部注意力 (Local Attention)、Swin Transformer (Liu et al., 2021) 和组查询注意力 (Group Query Attention, GQA) (Ainslie et al., 2023)，但这些方法主要侧重于设计新的注意力机制或网络架构。这意味着它们需要高昂的重新训练 (retraining) 成本，而训练一个 DiT 模型需要巨大的数据和计算资源。因此，当前领域迫切需要一种后训练 (post-training) 压缩方法，能够在不进行重新训练的情况下加速 DiT 模型的推理。

本文的切入点或创新思路： 本文的切入点是识别并利用 DiT 推理过程中注意力计算固有的冗余。通过深入分析注意力模式，研究发现存在三种主要冗余类型：空间冗余、时间步冗余和条件冗余。基于这些发现，本文提出了一个后训练 (post-training) 压缩框架 DiTFastAttn，旨在针对性地解决这些冗余，从而在保持生成质量的同时大幅提升推理速度。

2.2. 核心贡献/主要发现

本文的核心贡献在于提出了一种名为 DiTFastAttn 的后训练 (post-training) 压缩方法，用于加速扩散 Transformer 模型的推理，主要发现和贡献如下：

• 识别 DiT 注意力计算中的三种关键冗余：

  1. 空间冗余： 许多注意力头主要捕获局部空间信息，对远距离词元 (token) 的注意力值接近于零。
  2. 时间冗余： 在去噪过程中，相同注意力头在相邻时间步的注意力输出表现出高度相似性。
  3. 条件冗余： 在分类器自由引导 (CFG) 中，条件推理和无条件推理的注意力输出在许多层和时间步上呈现显著相似性。

• 提出了三种针对性压缩技术：

  1. 带有残差共享的窗口注意力 (Window Attention with Residual Sharing, WA-RS)： 通过使用窗口注意力减少空间冗余，并通过缓存和重用全注意力与窗口注意力输出之间的残差 (residual) 来弥补远程依赖的损失，从而在不降低性能的情况下实现加速。
  2. 跨时间步注意力共享 (Attention Sharing across Timesteps, AST)： 利用相邻时间步之间注意力输出的相似性，缓存早期时间步的注意力输出并复用于后续时间步，从而跳过重复计算。
  3. 跨分类器自由引导注意力共享 (Attention Sharing across CFG, ASC)： 利用条件和无条件推理之间注意力输出的相似性，重用条件推理的注意力输出用于无条件推理，减少 CFG 带来的双倍计算成本。

• 开发了压缩计划决策方法： 提出了一种简单的贪婪方法 (greedy method) 来为每个层和每个时间步选择最合适的压缩策略组合，以平衡计算效率和生成质量。

• 广泛的实验验证： 在多种 DiT 模型上进行了实验，包括用于图像生成的 DiT-XL 和 PixArt-Sigma，以及用于视频生成的 OpenSora。

  • 结果显示，DiTFastAttn 显著降低了计算成本。对于图像生成，可将注意力浮点运算次数 (FLOPs) 减少高达 76%，并在 2K x 2K 高分辨率生成中实现高达 1.8 倍的端到端加速 (end-to-end speedup)。

  • 随着分辨率的提高，计算节省和延迟 (latency) 降低的效果越显著。

  • 消融实验验证了每种技术以及残差共享机制的有效性。

    这些贡献共同提供了一个有效且实用的后训练 (post-training) 压缩框架，能够显著提高 DiT 模型在实际应用中的效率，特别是在处理高分辨率生成任务时。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 扩散模型 (Diffusion Models)

扩散模型 (Diffusion Models) 是一类新兴的生成模型，通过模拟一个逐步加噪声的扩散过程 (diffusion process) 和一个逐步去噪的逆扩散过程 (reverse diffusion process) 来学习数据的分布。在训练阶段，模型学习如何从带有噪声的数据中恢复原始数据。在推理阶段，模型从一个随机噪声向量开始，通过一系列去噪步骤逐步生成清晰的图像或视频。

• 去噪过程 (Denoising Process): 扩散模型的核心在于一个去噪 U-Net (Denoising U-Net) 或 去噪 Transformer (Denoising Transformer)，它在每个时间步预测添加到数据中的噪声，然后将噪声从当前状态中减去，从而逐步将随机噪声转化为有意义的数据。
• 迭代性质 (Iterative Nature): 扩散模型的生成过程是迭代的，需要多个时间步（通常是几十到几百步）来完成从噪声到数据的转换，这也是其推理速度较慢的一个原因。

3.1.2. 扩散 Transformer (Diffusion Transformers, DiT)

扩散 Transformer (DiT) 是一种特殊类型的扩散模型，它将传统的基于 U-Net 的去噪网络替换为基于 Transformer 架构的模型。

• 优势: Transformer 架构在处理长序列依赖方面表现出色，并具有更好的可扩展性 (scalability)，使得 DiT 能够生成更高质量和更高分辨率的图像和视频。
• 工作原理: 在 DiT 中，图像或视频首先被分解成一系列词元 (token)，然后这些词元被输入到 Transformer 模型中进行处理。Transformer 通过自注意力 (self-attention) 机制捕捉词元之间的关系，并进行特征变换，最终输出预测的噪声。

3.1.3. 自注意力机制 (Self-Attention Mechanism)

自注意力机制是 Transformer 架构的核心组件，允许模型在处理序列时，对序列中的每个元素都关注序列中其他所有元素的重要性。

• QKV 模型: 自注意力通过三个线性变换后的向量实现：查询 (Query, $Q$)、键 (Key, $K$) 和值 (Value, $V$)。
  • 查询 (Q): 表示当前词元 (token) 想要关注其他词元的“问题”或“兴趣点”。
  • 键 (K): 表示序列中其他词元可以提供的“信息”或“答案”。
  • 值 (V): 包含序列中其他词元的实际内容信息。
• 计算过程:
  1. 计算 $Q$ 和 $K$ 之间的点积相似度，得到注意力分数 (attention scores)。
  2. 将注意力分数除以 $\sqrt{d_k}$ 进行缩放，其中 $d_k$ 是 $K$ 向量的维度，这有助于稳定梯度。
  3. 对缩放后的分数应用 Softmax 函数，将其转换为注意力权重 (attention weights)，这些权重表示每个词元对其他词元的关注程度，总和为 1。
  4. 将注意力权重与 $V$ 向量加权求和，得到最终的注意力输出。
• 数学公式: 自注意力机制的计算可以表示为： $$\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
  • $Q$ (Query): 查询矩阵，形状为 $(L, d_k)$。
  • $K$ (Key): 键矩阵，形状为 $(L, d_k)$。
  • $V$ (Value): 值矩阵，形状为 $(L, d_v)$。
  • $L$: 输入序列的长度，即词元数量。
  • $d_k$: 键和查询向量的维度。
  • $d_v$: 值向量的维度。
  • $Q K^T$: 计算查询和键之间的相似度。
  • $\sqrt{d_k}$: 缩放因子，防止内积过大导致 Softmax 函数梯度消失。
  • $\mathrm{softmax}(\cdot)$: 将注意力分数转换为概率分布的函数。
  • 最终输出的形状为 $(L, d_v)$。
• 二次复杂度: 由于需要计算每个词元与所有其他词元之间的相似度，因此自注意力机制的计算复杂度为 $\mathcal{O}(L^2)$。这在高分辨率图像和视频生成中，当词元长度 $L$ 很大时，会带来巨大的计算负担。

3.1.4. 分类器自由引导 (Classifier-Free Guidance, CFG)

分类器自由引导 (Classifier-Free Guidance, CFG) (Ho & Salimans, 2022) 是一种在扩散模型中增强条件生成质量的技术。它通过结合条件模型 (conditional model) 和无条件模型 (unconditional model) 的输出来实现。

• 工作原理: 在每个去噪时间步，CFG 需要进行两次神经网络推理：一次是基于给定的条件输入（例如文本提示），另一次是没有条件输入（即生成无条件样本）。然后，通过一个引导尺度参数 (guidance scale) 将这两个预测结果结合起来，以增强条件性，使生成的图像更符合文本提示。
• 计算成本: 由于每个时间步需要进行两次完整的模型推理，CFG 会使计算成本翻倍，从而成为推理过程中的另一个主要瓶颈。

3.1.5. 浮点运算次数 (FLOPs)

浮点运算次数 (FLOPs) 是衡量计算机执行浮点运算数量的指标，常用于评估神经网络模型的计算复杂度。在本文中，Attn FLOPs 特指多头注意力 (multi-head attention) 模块的计算量，通常以与原始模型相比的百分比来表示，用于量化压缩方法带来的计算效率提升。

3.1.6. FID (Frechet Inception Distance)

概念定义： FID (Frechet Inception Distance) 是一种用于评估生成图像质量的指标，尤其是在比较生成模型输出与真实数据分布的相似性方面。它通过计算真实图像和生成图像在 Inception-v3 模型特征空间中的统计距离来衡量两者之间的差异。FID 值越低，表示生成图像的质量越高，且与真实图像的分布越接近。它比 Inception Score (IS) 更能反映生成图像的多样性和模式崩溃问题。

数学公式： 假设真实图像的特征分布为 $\mathcal{N}(\mu_1, \Sigma_1)$，生成图像的特征分布为 $\mathcal{N}(\mu_2, \Sigma_2)$，其中 $\mu$ 是特征均值，$\Sigma$ 是特征协方差矩阵。FID 的计算公式为： $$\text{FID} = ||\mu_1 - \mu_2||_2^2 + \text{Tr}(\Sigma_1 + \Sigma_2 - 2(\Sigma_1 \Sigma_2)^{1/2})$$

符号解释：

• $\mu_1$: 真实图像在 Inception-v3 特征空间中的特征均值向量。
• $\mu_2$: 生成图像在 Inception-v3 特征空间中的特征均值向量。
• $\Sigma_1$: 真实图像在 Inception-v3 特征空间中的特征协方差矩阵。
• $\Sigma_2$: 生成图像在 Inception-v3 特征空间中的特征协方差矩阵。
• $||\cdot||_2^2$: L2 范数的平方，衡量两个均值向量之间的欧氏距离。
• $\text{Tr}(\cdot)$: 矩阵的迹 (trace)，即矩阵对角线元素之和。
• $(\Sigma_1 \Sigma_2)^{1/2}$: 矩阵乘积 $\Sigma_1 \Sigma_2$ 的平方根。

3.1.7. IS (Inception Score)

概念定义： IS (Inception Score) 是另一种用于评估生成对抗网络 (GANs) 或其他生成模型生成图像质量的指标。它基于预训练的 Inception-v3 分类器来评估生成图像的两个方面：(1) 图像质量 (quality)：生成图像是否包含可识别的对象，即分类器对生成图像的预测置信度是否高（熵低）；(2) 图像多样性 (diversity)：生成图像是否包含多种不同的对象，即分类器在所有生成图像上的预测类别分布是否广泛（熵高）。高 IS 值表示生成图像既清晰可辨又具有多样性。

数学公式： 假设 $x$ 是一个生成的图像，$p(y|x)$ 是 Inception-v3 模型预测 $x$ 属于类别 $y$ 的条件概率分布，p(y) 是所有生成图像上类别 $y$ 的边缘概率分布。IS 的计算公式为： $$\text{IS} = \exp\left( \mathbb{E}_{x \sim p_g} [D_{\text{KL}}(p(y|x) || p(y))] \right)$$

符号解释：

• $x$: 从生成模型 $p_g$ 生成的图像样本。
• $p_g$: 生成模型的图像分布。
• $\mathbb{E}_{x \sim p_g}[\cdot]$: 对从生成模型 $p_g$ 中采样的图像 $x$ 求期望。
• $p(y|x)$: Inception-v3 模型在给定图像 $x$ 的情况下，预测其属于类别 $y$ 的概率分布。
• p(y): 所有生成图像在 Inception-v3 模型预测的类别 $y$ 的边缘概率分布。这可以通过对 $p(y|x)$ 在所有生成图像上求平均得到：$p(y) = \mathbb{E}_{x \sim p_g}[p(y|x)]$。
• $D_{\text{KL}}(P || Q)$: 两个概率分布 $P$ 和 $Q$ 之间的 Kullback-Leibler (KL) 散度，衡量它们之间的差异。
• $\exp(\cdot)$: 指数函数。

3.1.8. CLIP Score

概念定义： CLIP Score 是一种用于评估生成图像与给定文本提示之间语义一致性的指标。它利用 OpenAI 提出的 CLIP (Contrastive Language-Image Pre-training) 模型。CLIP 模型在一个巨大的图像-文本对数据集上进行训练，学习将图像和文本嵌入到同一个多模态特征空间中，使得语义相关的图像和文本具有较高的相似度。CLIP Score 通过计算生成图像和其对应的文本提示在 CLIP 嵌入空间中的余弦相似度来量化这种语义一致性。分数越高，表示图像与文本提示的匹配度越好。

数学公式： 假设 $I$ 是一个生成的图像，$T$ 是其对应的文本提示。CLIP Score 的计算公式为： $$\text{CLIP Score} = \text{sim}(\text{CLIP\_ImageEncoder}(I), \text{CLIP\_TextEncoder}(T))$$

符号解释：

• $\text{CLIP\_ImageEncoder}(I)$: CLIP 模型的图像编码器将图像 $I$ 映射到多模态特征空间中的向量。
• $\text{CLIP\_TextEncoder}(T)$: CLIP 模型的文本编码器将文本提示 $T$ 映射到多模态特征空间中的向量。
• $\text{sim}(\cdot, \cdot)$: 两个向量之间的余弦相似度函数，通常定义为 $\frac{A \cdot B}{||A|| \cdot ||B||}$。

3.2. 前人工作

3.2.1. 扩散模型 (Diffusion Models)

早期的扩散模型如 DDPM (Ho et al., 2020) 和 Latent Diffusion Models (LDM) (Rombach et al., 2022) 通常基于 U-Net 架构。为了实现更好的可扩展性 (scalability)，DiT (Peebles & Xie, 2023) 引入了 Transformer 架构来替代 U-Net。这一演进使得扩散模型在图像和视频生成领域取得了显著突破。例如，PixArt-Sigma (Chen et al., 2024) 展示了 DiT 生成高达 4K 分辨率图像的能力，而 Sora (Brooks et al., 2024) 则展示了 DiT 生成视频的强大能力。

3.2.2. 视觉 Transformer 压缩 (Vision Transformer Compression)

Transformer 的注意力计算开销一直是研究热点。

• FlashAttention (Dao, 2023): 通过将输入词元 (token) 分割成更小的块来优化内存访问，从而提高计算效率和降低延迟 (latency)。
• 词元剪枝 (Token Pruning):
  • DynamicViT (Rao et al., 2021) 使用预测网络动态过滤词元。
  • Adaptive Sparse ViT (Liu et al., 2022) 结合注意力值和特征的 L2 范数来过滤词元。
  • Wu et al. (2023) 发现深层更适合剪枝，浅层更适合合并。
• 词元合并 (Token Merging):
  • Lu et al. (2023) 训练网络利用语义分割标签指导词元合并。
  • Huang et al. (2023) 在下采样词元后进行注意力计算，再上采样恢复空间分辨率。

3.2.3. 局部注意力 (Local Attention)

局部注意力模式通过限制每个词元只关注固定窗口内的邻近词元来减轻计算负担。

• Longformer (Beltagy et al.): 首次引入窗口注意力概念，实现与序列长度线性扩展的注意力机制。
• Bigbird (Zaheer et al., 2020): 结合窗口注意力、随机注意力和全局注意力机制，在保持长距离依赖的同时降低计算成本。
• Swin Transformer (Liu et al., 2021): 在计算机视觉领域，将注意力计算限制在非重叠的局部窗口内，并通过跨层移动窗口来捕捉全局上下文。
• Twins Transformer (Chu et al., 2021), FasterViT (Vasu et al., 2023), Neighborhood Attention Transformer (Hassani et al., 2023): 均采用基于窗口的注意力来提高计算效率，并利用不同的模块设计来利用全局上下文。
• 本文工作的差异： DiTFastAttn 采用固定大小的窗口注意力来加速预训练的 DiT 模型，并引入了新颖的带有残差共享的窗口注意力 (WA-RS) 技术，以在图像词元中保留长距离依赖。

3.2.4. 注意力共享 (Attention Sharing)

利用注意力计算中的相似性来减少冗余。

• GQA (Ainslie et al., 2023): 将查询头 (query heads) 分组，每组共享键 (key) 和值 (value)，以减少内存使用并提高效率。
• PSVIT (Chen et al., 2021): 发现 Vision Transformer (ViT) 不同层之间的注意力图 (attention maps) 具有显著相似性，提出跨层共享注意力图以减少冗余计算。
• Deepcache (Ma et al., 2023): 发现 U-Net 框架扩散模型中的高级特征在不同时间步之间相似，提出重用 U-Net 的高级特征并跳过中间层的计算来加速去噪过程。
• TGATE (Zhang et al., 2024): 发现文本条件扩散模型中的交叉注意力 (cross-attention) 输出在几个去噪时间步后会收敛到固定点，提出一旦收敛就缓存该输出并在剩余去噪步骤中保持固定。
• 本文工作的差异： DiTFastAttn 证明了注意力输出在分类器自由引导 (CFG) 维度和时间步维度上的相似性，并考虑了不同层在不同时间步相似性的差异，从而在这两个维度上共享注意力输出。

3.2.5. 其他扩散模型加速方法

• 网络量化 (Network Quantization): 通过降低权重和激活值的位宽来压缩模型，减少存储和计算开销 (Shang et al., 2023; Zhao et al., 2024b, 2024a)。
• 调度器优化 (Scheduler Optimization): 旨在减少去噪过程中的时间步数量 (Song et al., 2020; Lu et al., 2022; Liu et al., 2023a)。
• 蒸馏 (Distillation): 通过训练一个更小的模型来模仿原始大模型的行为，从而减少去噪所需的时间步 (Salimans & Ho, 2022; Meng et al., 2023; Liu et al., 2023b)。
• 本文工作的互补性： DiTFastAttn 作为一个后训练 (post-training) 压缩方法，与这些技术是互补的，因为它独立于具体的量化位宽、调度器和时间步设置。

3.3. 技术演进

DiTFastAttn 的工作背景建立在扩散模型从最初的 U-Net 架构向更具可扩展性 (scalability) 的 Transformer 架构演进的基础上。传统的 U-Net 在处理高分辨率数据时存在局限性，而 Transformer 凭借其自注意力 (self-attention) 机制处理长序列的能力，成为了 DiT 模型的核心主干网络 (backbone)。然而，Transformer 的二次计算复杂度在高分辨率任务中成为了新的瓶颈。

现有的加速策略，如 FlashAttention，优化了底层计算效率；而词元剪枝和合并等方法则试图减少序列长度。局部注意力模型（如 Swin Transformer）通过限制注意力范围来降低复杂度，但这些方法通常需要在训练阶段进行架构修改。DiTFastAttn 处于这一技术演进的末端，它关注的是如何在模型已预训练完成的情况下，通过识别并利用 DiT 独特的推理冗余来进一步提升效率，而无需重新训练。它借鉴了注意力共享的思想，但将其扩展到扩散模型的特定维度（时间步和 CFG），并提出了新颖的残差共享机制来弥补局部注意力的潜在性能损失。

3.4. 差异化分析

DiTFastAttn 与其他相关工作的主要区别和创新点体现在以下几个方面：

1. 后训练压缩 (Post-Training Compression)： 与许多需要重新训练以修改模型架构（如 Swin Transformer、GQA）或引入新学习模块（如动态词元剪枝）的方法不同，DiTFastAttn 是一种纯粹的后训练 (post-training) 压缩方法。这意味着它可以在不改变现有预训练 DiT 模型参数和不产生额外训练成本的情况下，直接应用于模型进行加速。这对于大型预训练模型（如 DiT、PixArt-Sigma、OpenSora）而言，具有巨大的实用价值。

2. 针对 DiT 特有冗余的识别与利用： 本文深入分析了扩散模型推理过程中的特点，识别出三种 DiT 模型特有的冗余，并为每种冗余设计了特定的压缩技术，这是其与通用 Transformer 压缩方法的显著区别：

   • 空间冗余 (Spatial Redundancy)： 通过 WA-RS 解决，创新点在于引入了残差共享 (Residual Sharing) 机制。与单纯的局部注意力（如 Swin Transformer）相比，WA-RS 在使用窗口注意力的同时，通过缓存和重用全注意力与窗口注意力之间的残差 (residual)，有效保留了长距离依赖，从而避免了严重的性能下降。这使得局部注意力在高分辨率 DiT 中应用成为可能。
   • 时间冗余 (Temporal Redundancy)： 通过 AST 解决。利用扩散模型去噪过程的序列性质，发现相邻时间步的注意力输出相似，从而进行共享。这与 Deepcache 在 U-Net 中共享高层特征、TGATE 共享交叉注意力输出有相似之处，但 AST 专注于 DiT 的自注意力输出，并在不同层和时间步进行差异化共享。
   • 条件冗余 (Conditional Redundancy)： 通过 ASC 解决。专门针对分类器自由引导 (CFG) 机制带来的计算量翻倍问题，通过共享条件和无条件推理的注意力输出来削减成本。这是对 CFG 效率优化的一个新颖切入点。

3. 定制化的压缩计划 (Customized Compression Plan)： 本文提出的贪婪方法 (greedy method) 能够根据不同层和时间步的冗余特性，动态选择最合适的压缩策略组合。这克服了单一压缩策略的局限性，使得压缩效果最大化，同时保持生成质量。

4. 互补性 (Complementarity)： DiTFastAttn 可以与现有的其他加速技术（如量化、调度器优化、蒸馏等）并行使用，共同提升模型推理效率，而非相互替代。

   综上所述，DiTFastAttn 的创新性在于其后训练 (post-training) 特性，对 DiT 模型中独特冗余的精准识别，以及为解决这些冗余而设计的融合了残差共享 (residual sharing) 和多维度注意力共享 (attention sharing) 的策略组合，辅以定制化的压缩计划，使其成为 DiT 模型加速领域的一个重要进展。

4. 方法论

4.1. 方法原理

DiTFastAttn 的核心原理在于识别并利用扩散 Transformer (DiT) 模型推理过程中的三种关键冗余，并针对每种冗余设计相应的后训练 (post-training) 压缩技术。其基本直觉是，在 DiT 的去噪过程中，并非所有的注意力计算都同等重要或必须每次都完整执行。通过有选择地简化或跳过冗余计算，可以在不显著牺牲生成质量的前提下大幅提升推理速度。

该方法通过以下三项技术来实现：

1. Window Attention with Residual Sharing (WA-RS)： 解决注意力机制中的空间冗余 (spatial redundancy)。直观地说，许多注意力头主要关注图像的局部区域，对远距离区域的关注度很低。WA-RS 用计算量更小的窗口注意力 (window attention) 取代全注意力 (full attention)，并通过缓存和重用全注意力与窗口注意力之间的残差 (residual) 来弥补局部注意力可能丢失的全局信息，从而保持性能。

2. Attention Sharing across Timesteps (AST)： 解决时间冗余 (temporal redundancy)。观察到扩散模型在连续的去噪时间步 (timesteps) 中，其注意力输出常常高度相似。AST 利用这一特性，在相似的时间步之间共享注意力输出，避免重复计算。

3. Attention Sharing across CFG (ASC)： 解决条件冗余 (conditional redundancy)。在分类器自由引导 (CFG) 机制下，模型需要进行条件和无条件两次推理。研究发现这两次推理的注意力输出在许多情况下非常相似。ASC 通过重用条件推理的注意力输出，跳过无条件推理中的冗余计算。

   此外，由于不同层和不同时间步的冗余程度不同，DiTFastAttn 还引入了一种贪婪方法 (greedy method) 来动态决定为每个层和时间步应用哪种压缩策略，以实现计算效率和生成质量的最佳平衡。

4.2. 核心方法详解 (逐层深入)

4.2.1. 带有残差共享的窗口注意力 (Window Attention with Residual Sharing, WA-RS)

冗余类型： 空间冗余。 方法原理： 在许多预训练的 DiT 模型中，注意力机制表现出显著的空间局部性 (spatial locality)。这意味着注意力值主要集中在一个固定大小的窗口内，而对远距离词元 (token) 的注意力值往往接近于零。因此，用计算成本更低的固定大小窗口注意力 (window attention) 替代全注意力 (full attention) 可以大幅减少计算量。

然而，简单地丢弃窗口外的所有注意力计算会导致性能下降，因为一些词元仍然需要关注少量空间上较远的词元（即长距离依赖）。为了在高压缩率下保持性能，本文提出了缓存并重用窗口注意力残差 (Cache and Reuse the Residual for Window Attention) 的策略。

如下图 (原文 Figure 3) 所示，全注意力输出与窗口注意力输出之间的残差 (residual) 在不同时间步之间变化很小。这启发了作者缓存一个时间步的残差，并在后续的多个时间步中重复使用它。

该图像是示意图，展示了滑动窗口注意力机制的计算过程，左侧为不同时间步的注意图，右侧为全注意力与窗口注意力的比较及残差计算公式，其中 $R^t = O^t - W^t$ 表示残差计算。

图示：带有残差共享的窗口注意力。(a) 左图：注意力图显示窗口模式示例。右图：前一步和当前步窗口注意力输出的均方误差（黄色线）与前一步和当前步窗口注意力与全注意力输出残差的均方误差（蓝色线）对比。输出残差在各步骤中变化极小。(b) 带有残差共享的窗口注意力计算流程。变化显著的窗口注意力会被重新计算。变化极小的残差则被缓存并在后续步骤中重复使用。

计算流程： 设 $\mathbf{K}$ 为共享残差值 $\mathbf{R}_r$ 的时间步集合。对于集合中的第一个时间步 r = \min(\mathbf{K})，WA-RS 的计算如下： $$\begin{array} { r l } & { \mathbf { O } _ { r } = \operatorname { A t t e n t i o n } ( \mathbf { Q } _ { r } , \mathbf { K } _ { r } , \mathbf { V } _ { r } ) , } \\ & { \mathbf { W } _ { r } = \operatorname { W i n d o w A t t e n t i o n } ( \mathbf { Q } _ { r } , \mathbf { K } _ { r } , \mathbf { V } _ { r } ) , } \\ & { \mathbf { R } _ { r } = \mathbf { O } _ { r } - \mathbf { W } _ { r } . } \end{array}$$ 符号解释：

• $\mathbf{O}_r$: 在时间步 $r$ 时，通过全注意力 (Full Attention) 计算得到的输出。

• $\operatorname{Attention}(\mathbf{Q}_r, \mathbf{K}_r, \mathbf{V}_r)$: 表示在时间步 $r$ 使用查询 (Query) $\mathbf{Q}_r$、键 (Key) $\mathbf{K}_r$ 和值 (Value) $\mathbf{V}_r$ 进行全注意力计算。

• $\mathbf{W}_r$: 在时间步 $r$ 时，通过窗口注意力 (Window Attention) 计算得到的输出。

• $\operatorname{WindowAttention}(\mathbf{Q}_r, \mathbf{K}_r, \mathbf{V}_r)$: 表示在时间步 $r$ 使用查询 (Query) $\mathbf{Q}_r$、键 (Key) $\mathbf{K}_r$ 和值 (Value) $\mathbf{V}_r$ 进行窗口注意力计算。窗口注意力只考虑每个词元固定邻域内的其他词元。

• $\mathbf{R}_r$: 在时间步 $r$ 计算得到的残差 (residual)，它是全注意力输出 $\mathbf{O}_r$ 与窗口注意力输出 $\mathbf{W}_r$ 之间的差值。这个残差捕获了窗口注意力丢失的远程依赖信息。

  对于集合 $\mathbf{K}$ 中的后续时间步 $k \in \mathbf{K}$（其中 $k > r$），WA-RS 的计算如下： $$\begin{array} { r l } & { \mathbf { W } _ { k } = { \mathrm { W i n d o w A t t e n t i o n } } ( \mathbf { Q } _ { k } , \mathbf { K } _ { k } , \mathbf { V } _ { k } ) , } \\ & { \mathbf { O } _ { k } = \mathbf { W } _ { k } + \mathbf { R } _ { r } . } \end{array}$$ 符号解释：

• $\mathbf{W}_k$: 在后续时间步 $k$ 时，通过窗口注意力 (Window Attention) 计算得到的输出。

• $\mathbf{O}_k$: 在后续时间步 $k$ 时，通过 WA-RS 得到的最终注意力输出。它由当前时间步的窗口注意力输出 $\mathbf{W}_k$ 和之前缓存的残差 $\mathbf{R}_r$ 相加得到。这意味着在这些后续时间步中，无需重新计算全注意力，只需计算窗口注意力并加上缓存的残差即可，从而大大减少了计算量。

4.2.2. 跨时间步注意力共享 (Attention Sharing across Timesteps, AST)

冗余类型： 时间冗余。 方法原理： 扩散模型去噪过程的顺序性 (sequential nature) 带来了推理速度的瓶颈。研究发现，在去噪过程中，对于某些层，相同注意力头在相邻时间步 (timesteps) 的注意力输出表现出显著的相似性。

如下图 (原文 Figure 4a) 所示，不同时间步的注意力输出之间存在明显的时间相似性 (temporal similarity)。这种相似性在不同的时间步和层之间有所不同。

该图像是图表，显示了不同层次下的注意力输出的相似性。左侧为在不同时间步刻度下，第5层和第25层的注意力输出的余弦相似性；右侧为条件生成和无条件生成之间的注意力输出相似性的热图。

图示：DiT 中注意力输出在时间步和 CFG 维度上的相似性。(a) 不同层中注意力输出在时间步维度上的相似性。(b) 不同时间步下，不同层中条件和无条件注意力输出之间的相似性。

计算流程： 为了利用这种相似性来降低计算成本，AST 技术被提出。具体来说，对于一组注意力输出彼此相似的时间步，本文缓存最早时间步 (earliest step) 的注意力输出 $\mathbf{O}$，并在后续时间步中重复使用它，从而跳过这些后续时间步的注意力计算。

4.2.3. 跨分类器自由引导注意力共享 (Attention Sharing across CFG, ASC)

冗余类型： 条件冗余。 方法原理： 分类器自由引导 (Classifier-Free Guidance, CFG) (Ho & Salimans, 2022) 广泛应用于条件生成中，但其机制要求在每个推理步骤中进行两次神经网络评估：一次带条件输入，一次不带条件输入。这使得计算成本相比无条件生成翻倍。

研究观察到，如下图 (原文 Figure 4b) 所示，对于许多层和时间步，条件神经网络评估和无条件神经网络评估的注意力输出之间存在高度相似性 (例如，结构相似性指数 SSIM $\ge 0.95$)。

计算流程： 基于这一观察，ASC 技术被提出。该技术重用条件神经网络评估 (conditional neural network evaluation) 的注意力输出，用于无条件神经网络评估 (unconditional neural network evaluation)。这意味着在进行无条件推理时，可以直接使用已计算出的条件推理的注意力输出，从而跳过无条件推理中注意力模块的重复计算，有效地将 CFG 的注意力计算成本减半。

4.2.4. 压缩计划决策方法 (Method for Deciding the Compression Plan)

问题： 前述的 WA-RS、AST 和 ASC 技术都能有效降低计算成本并保持性能。然而，正如 Figure 3 和 Figure 4 所示，不同层在不同时间步的冗余类型和程度是不同的。因此，为每个层和每个时间步正确选择合适的压缩技术组合（即决定压缩计划 (compression plan)）至关重要。

方法： 本文开发了一种简单的贪婪方法 (greedy method) 来选择合适的策略。策略列表 $S$ 包含了四种候选项：[AST, WA-RS + ASC, WA-RS, ASC]。算法会根据压缩比 (compression ratio) 升序排列这些策略，优先尝试压缩率高的策略。

算法 1 (Algorithm 1): 压缩计划决策方法 输入：

• Transformer 模型 $M$
• 总时间步 $T$
• 压缩策略列表 $S$
• 阈值 $\delta$

输出：

• dict：存储所选压缩技术的字典（键为 (step, layer_index)，值为选定的策略）。

初始化：

• 初始化 dict 为空。

流程：

1. 对于 $T$ 中的每个时间步 $t$：
   1. 计算未压缩 (uncompressed) 模型 $M$ 在当前时间步的输出 $O$。
2. 对于 $M$ 中的每个 Transformer 层 $i$：
   1. 对于 $S$ 中的每个压缩策略 $m$ (按压缩比升序排序)：
      1. 使用压缩策略 $m$ 压缩时间步 $t$ 的层 $i$。
      2. 计算应用压缩后模型 $M$ 的输出 $O'$。
      3. 如果 $L(O, O')^* < \frac{i}{|M|}\delta$：
         • 更新 dict，将 $m$ 作为层 $i$ 和时间步 $t$ 的选择策略。
         • 跳出当前策略循环，进入下一个层或时间步。
      4. 否则： 继续尝试 $S$ 中的下一个策略。
   2. 如果 没有策略满足阈值条件：

      • 则层 $i$ 在时间步 $t$ 不应用任何压缩。

        返回： dict。

说明：

• $|M|$ 是模型中的 Transformer 层数。
• $\frac{i}{|M|}\delta$ 是一个动态阈值，它允许深层（$i$ 越大）有更大的误差容忍度，这通常是因为深层对模型输出的影响更小。
• L(O, O') 是用于衡量压缩前后模型输出差异的损失函数。

损失函数 L(O, O') 详解 (Appendix A.2): 本文使用均相对绝对误差 (mean relative absolute error) 作为损失函数 L(O, O')，用于评估模型输出 $O'$ 与原始输出 $O$ 之间的平均相对偏差。 L ( O , O ^ { \prime } ) = \frac { 1 } { | O | _ { 1 } } \sum _ { i } \mathrm { c l i p } \left( \frac { | O _ { i } - O _ _ { i } ^ { \prime } | } { \operatorname* { m a x } ( | O _ { i } | , | O _ { i } ^ { \prime } | ) + \epsilon } , 0 , 1 0 \right) 符号解释：

• $|O|_1$: 原始输出向量 $O$ 中的元素数量。

• $i$: 遍历向量 $O$ 和 $O'$ 中每个元素的索引。

• $O_i$: 原始输出向量 $O$ 的第 $i$ 个元素。

• O'_i: 压缩后输出向量 $O'$ 的第 $i$ 个元素。

• $|O_i - O'_i|$: 原始输出与压缩后输出之间绝对差值。

• $\operatorname{max}(|O_i|, |O'_i|) + \epsilon$: 作为归一化因子，使误差相对于输出值的幅度。为了避免数值不稳定性，当 $O_i$ 和 O'_i 都非常小或为零时，会添加一个小的正常量 $\epsilon$（在实验中设置为 $10^{-6}$）。

• $\mathrm{clip}(\cdot, 0, 10)$: 裁剪函数 (clip function)，确保相对误差值被限制在 [0, 10] 的范围内，防止极端值主导总误差。

• $\sum_i$: 对所有元素的裁剪相对误差求和。

• $\frac{1}{|O|_1}$: 将求和结果除以元素总数，得到平均相对绝对误差。

  这个指标提供了一个归一化的平均相对偏差度量，其值范围从 0 到 10（最大允许误差）。通过这个损失函数，算法能够权衡压缩带来的计算效益与对模型输出质量的影响。

5. 实验设置

5.1. 数据集

• 图像生成任务：
  • DiT (Diffusion Transformer) 模型： 使用 ImageNet 数据集进行评估。ImageNet 是一个包含数百万张带标签图像的大型数据集，常用于评估图像生成模型的性能。
  • PixArt-Sigma 模型： 使用 MS-COCO 数据集进行评估。MS-COCO 2014 Caption (Microsoft Common Objects in Context) 数据集被用作文本提示 (text prompt) 来指导图像生成。MS-COCO 是一个广泛用于目标检测、图像分割和图像描述的数据集，包含大量带有详细文本描述的图像。
• 视频生成任务：
  • OpenSora 模型： 具体的视频数据集在原文中未直接提及，但 OpenSora 项目本身旨在复现 OpenAI Sora 的能力，通常会使用大规模视频数据集进行训练和评估。本文在 OpenSora 模型上进行验证。

数据集中的具体样本示例： 原文未提供数据集中的具体样本示例，但根据描述，ImageNet 包含各种类别的图像，MS-COCO 包含带有复杂场景和对象描述的图像。例如，对于 MS-COCO，文本提示可能是“A group of people standing next to a large body of water under a blue sky.”（一群人站在一大片水域旁边，蓝天下）。

5.2. 评估指标

论文中使用了多种标准指标来评估生成图像和视频的质量以及模型的计算效率。

5.2.1. 图像生成评估指标

• FID (Frechet Inception Distance)：

  • 概念定义 (Conceptual Definition)： FID 衡量真实图像和生成图像特征分布之间的距离，越小越好。它通过比较两组图像在 Inception-v3 模型特征空间中的均值和协方差来量化它们的相似性。FID 值能够捕捉生成图像的质量和多样性，是当前评估生成模型性能最广泛和权威的指标之一。
  • 数学公式 (Mathematical Formula)： $$\text{FID} = ||\mu_1 - \mu_2||_2^2 + \text{Tr}(\Sigma_1 + \Sigma_2 - 2(\Sigma_1 \Sigma_2)^{1/2})$$
  • 符号解释 (Symbol Explanation)：
    • $\mu_1$: 真实图像在 Inception-v3 特征空间中的特征均值向量。
    • $\mu_2$: 生成图像在 Inception-v3 特征空间中的特征均值向量。
    • $\Sigma_1$: 真实图像在 Inception-v3 特征空间中的特征协方差矩阵。
    • $\Sigma_2$: 生成图像在 Inception-v3 特征空间中的特征协方差矩阵。
    • $||\cdot||_2^2$: L2 范数的平方，衡量两个均值向量之间的欧氏距离。
    • $\text{Tr}(\cdot)$: 矩阵的迹，即矩阵对角线元素之和。
    • $(\Sigma_1 \Sigma_2)^{1/2}$: 矩阵乘积 $\Sigma_1 \Sigma_2$ 的平方根。

• IS (Inception Score)：

  • 概念定义 (Conceptual Definition)： IS 评估生成图像的质量和多样性，越大越好。它利用预训练的 Inception-v3 分类器来判断生成图像是否包含可识别的对象（质量高）以及是否覆盖了多种不同的对象类别（多样性高）。
  • 数学公式 (Mathematical Formula)： $$\text{IS} = \exp\left( \mathbb{E}_{x \sim p_g} [D_{\text{KL}}(p(y|x) || p(y))] \right)$$
  • 符号解释 (Symbol Explanation)：
    • $x$: 从生成模型 $p_g$ 生成的图像样本。
    • $p_g$: 生成模型的图像分布。
    • $\mathbb{E}_{x \sim p_g}[\cdot]$: 对从生成模型 $p_g$ 中采样的图像 $x$ 求期望。
    • $p(y|x)$: Inception-v3 模型在给定图像 $x$ 的情况下，预测其属于类别 $y$ 的概率分布。
    • p(y): 所有生成图像在 Inception-v3 模型预测的类别 $y$ 的边缘概率分布。这可以通过对 $p(y|x)$ 在所有生成图像上求平均得到：$p(y) = \mathbb{E}_{x \sim p_g}[p(y|x)]$。
    • $D_{\text{KL}}(P || Q)$: 两个概率分布 $P$ 和 $Q$ 之间的 Kullback-Leibler (KL) 散度，衡量它们之间的差异。
    • $\exp(\cdot)$: 指数函数。

• CLIP Score：

  • 概念定义 (Conceptual Definition)： CLIP Score 衡量生成图像与文本提示之间的语义一致性，越高越好。它使用 CLIP (Contrastive Language-Image Pre-training) 模型将图像和文本嵌入到同一特征空间，然后计算它们嵌入向量的余弦相似度。
  • 数学公式 (Mathematical Formula)： $$\text{CLIP Score} = \text{sim}(\text{CLIP\_ImageEncoder}(I), \text{CLIP\_TextEncoder}(T))$$
  • 符号解释 (Symbol Explanation)：
    • $I$: 生成的图像。
    • $T$: 对应的文本提示。
    • $\text{CLIP\_ImageEncoder}(I)$: CLIP 模型的图像编码器将图像 $I$ 映射到多模态特征空间中的向量。
    • $\text{CLIP\_TextEncoder}(T)$: CLIP 模型的文本编码器将文本提示 $T$ 映射到多模态特征空间中的向量。
    • $\text{sim}(\cdot, \cdot)$: 两个向量之间的余弦相似度函数，通常定义为 $\frac{A \cdot B}{||A|| \cdot ||B||}$。

5.2.2. 计算效率评估指标

• Attn FLOPs (Attention Floating Point Operations)： 表示在多头注意力 (multi-head attention) 模块中进行的浮点运算次数。通常以相对于原始模型的百分比表示，用于量化压缩方法带来的计算量减少。
• Latency (延迟)： 衡量模型完成一次推理所需的时间，通常以秒 (s) 为单位。分为注意力延迟 (Attn Latency)（仅注意力模块的延迟）和端到端延迟 (End-to-End Latency)（整个生成过程的延迟）。

5.3. 对比基线

本文主要将 DiTFastAttn 方法的性能与原始的、未压缩的 DiT 模型进行比较。这些原始模型作为基线，用于展示 DiTFastAttn 在计算成本降低和生成质量保持方面的效果。

实验中使用的具体基线模型包括：

• DiT-XL-2-512： 用于图像生成，生成 512x512 尺寸的图像。

• PixArt-Sigma-XL-1024： 用于图像生成，生成 1024x1024 尺寸的图像。

• PixArt-Sigma-XL-2K： 用于图像生成，生成 2048x2048 尺寸的图像。

• OpenSora (OpenSora, 2024)： 用于视频生成任务。

  这些基线模型都是当前最先进的扩散 Transformer 模型，因此与它们的比较能够充分验证 DiTFastAttn 的有效性。论文旨在展示其后训练 (post-training) 压缩能力，因此未直接与其他需要重新训练的注意力加速方法（如 Swin Transformer）进行并列性能对比。

5.4. 其他实验设置

• 采样方法：
  • DiT 和 PixArt-Sigma 模型：使用 50 步的 DPM-Solver (Lu et al., 2022) 作为快速采样器。
  • OpenSora 模型：使用 200 步的 IDDPM (Improved Denoising Diffusion Probabilistic Models) (Nichol & Dhariwal, 2021)。
• 硬件： 所有延迟 (latency) 测量均在单个 Nvidia A100 GPU 上进行。
• 评估样本数量：
  • DiT 模型：生成 50k 张图像用于计算质量指标。
  • PixArt-Sigma 模型：生成 30k 张图像用于计算质量指标。
• 损失函数和阈值：
  • 用于决定压缩计划的损失函数 L(O, O') 采用均相对绝对误差 (mean relative absolute error)。
  • 阈值 $\delta$ 在 0.025 到 0.15 之间以 0.025 的间隔进行实验。这些阈值设置被标记为 D1 ($\delta=0.025$)、D2 ($\delta=0.05$)、...、D6 ($\delta=0.15$)。
• WA-RS 窗口大小： 带有残差共享的窗口注意力 (WA-RS) 的窗口大小设置为词元 (token) 大小的 1/8。
• 实现基础： DiTFastAttn 基于 FlashAttention-2 (Dao, 2023) 实现，以确保底层注意力计算的高效率。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 图像生成性能和注意力 FLOPs

以下是原文 Table 1 的结果：

• 对 DiT-XL-2-512 和 PixArt-Sigma-1024 的分析：

  • 在 D1、D2 和 D3 配置下，DiTFastAttn 在 IS 和 FID 指标上与原始模型表现几乎一致，同时显著降低了注意力 FLOPs。例如，DiT-XL-2-512 在 D3 配置下注意力 FLOPs 降至 59%，性能却能很好地保持。PixArt-Sigma-1024 在 D3 配置下注意力 FLOPs 降至 63%，IS 和 FID 几乎无损。
  • 随着压缩等级的提高（D4-D6），注意力 FLOPs 进一步降低（例如 DiT-XL-2-512 在 D6 降至 34%），但 IS 和 FID 指标开始出现一定程度的下降。不过，这种下降仍在可接受范围内，尤其是在需要更高效率的场景下。

• 对 PixArt-Sigma-XL-2K 的分析：

  • 对于更高分辨率的 PixArt-Sigma-XL-2K 模型，DiTFastAttn 展现出更优异的性能保持能力和压缩效果。
  • 在 D1、D2、D3 甚至 D4 配置下，PixArt-Sigma-XL-2K 的 IS、FID 和 CLIP Score 仍然非常接近原始模型。
  • 特别是在 D6 配置下，PixArt-Sigma-XL-2K 的注意力 FLOPs 仅为原始模型的 24%，即减少了 76% 的计算量，而其 IS、FID 和 CLIP Score 仅有轻微下降。
  • 这一结果表明，分辨率越高，DiTFastAttn 能够实现的压缩率越高，同时对生成性能的保持越好。这验证了高分辨率场景下，注意力计算中的冗余越大，且 DiTFastAttn 能够更有效地利用这些冗余。

6.1.2. 压缩计划分布

下图 (原文 Figure 5) 展示了 DiT-XL-512、PixArt-Sigma-XL-1024 和 PixArt-Sigma-XL-2K 在 D6 配置下，DPM-Solver 步数为 50 时的压缩计划。

该图像是示意图，展示了DiT-XL-2、PixArt-Sigma-XL-1024和PixArt-Sigma-XL-2K在D6配置下的压缩计划，阈值设定为0.15。图中展示了不同层与时间步的注意力分布，包括完全注意力（Full Attn）、窗口注意力与残差共享（WA-RS）、注意力共享（ASC）、窗口注意力与残差共享及注意力共享（WA-RS+ASC）和自适应共享技术（AST）的效果。

图示：DiT-XL-512、PixArt-Sigma-XL-1024 和 PixArt-Sigma-XL-2K 在 D6 配置下，DPM-Solver 步数设为 50 时的压缩计划。

• 对 DiT 模型的分析：
  • AST (跨时间步注意力共享) 和 ASC (跨 CFG 注意力共享) 主要在早期时间步 (early timesteps) 被广泛使用。
  • 全注意力 (Full Attention) 主要出现在初始的注意力层 (initial attention layers)。这可能表明早期时间步和初始层对细节和全局上下文的捕捉更为关键，需要更精确的计算。
• 对 PixArt-Sigma 模型的分析：
  • AST 在前两层和中间时间步的中间注意力层中零星出现。
  • WA-RS + ASC (窗口注意力与残差共享 + 跨 CFG 注意力共享) 的组合在最终时间步 (final timesteps) 占据主导地位。这可能反映了在去噪过程的后期，图像结构已经基本确定，此时可以更激进地采用局部注意力和共享策略。
• 结论： 压缩计划在不同模型之间存在显著差异，且在时间步和层维度上分布不均。这强调了没有通用压缩策略 (no universal compression strategy)，必须通过定制化的搜索来决定每个层和时间步的最佳压缩方案。

6.1.3. 可视化生成结果

下图 (原文 Figure 6) 展示了 DiTFastAttn 在不同分辨率和压缩比下的图像生成样本。

该图像是图像生成示例，展示了在不同压缩比和分辨率下生成的图像，共有三种分辨率，分别为512x512、1024x1024和2048x2048。图中呈现了不同样本的细节和清晰度。

图示：在不同图像分辨率和压缩比下的图像生成样本。

• DiT-XL-2-512 和 PixArt-Sigma-1024： 在 D1、D2 和 D3 配置下生成的图像，其视觉质量与原始模型相当。图像细节丰富，整体构图和颜色保持良好。
• DiT-XL-2-512 和 PixArt-Sigma-1024 的 D4、D5、D6 配置： 实现了更大的压缩，图像在细节上可能存在轻微的差异，但整体仍能生成可接受质量的图像。
• PixArt-Sigma-2K 模型： 即使在 D4 配置下也能保持与原始模型相似的图像质量，而在 D5 和 D6 配置下也生成了高质量的输出。
• 结论： DiTFastAttn 能够有效地在大幅减少注意力计算（例如在高分辨率下减少超过 50% 并压缩至 33%）的同时，保持较高的生成质量。

6.1.4. 视频生成结果

下图 (原文 Figure 7) 展示了使用 OpenSora V1.1 在 240p 分辨率下生成 16 帧视频的对比。

该图像是一个视频生成对比图，展示了在240p分辨率下使用OpenSora V1.1生成的16帧视频。图像的左侧展示了热气球的场景，中间展示了海洋中的海龟，右侧展示了夜晚城市的景象，体现了不同场景下的视频生成效果。

图示：使用 OpenSora V1.1 在 240p 分辨率下生成 16 帧视频的对比。

原文 Appendix A.3 (Figure 10) 提供了更详细的视频生成结果，这里截取了一部分进行描述。

• 在 D1 到 D4 配置下，DiTFastAttn 在视频生成任务中表现出有效性，能够在平衡计算效率的同时，保持生成视频的视觉质量。这些配置下生成的视频流畅，帧间过渡自然，关键细节得到保留。
• 在 D5 和 D6 配置下，虽然应用了更激进的压缩，视频仍保持连贯和流畅，基本能传达预期的叙事或提示。然而，与原始视频相比，视觉质量出现明显下降。
• 结论： 视频生成结果强调了在应用 DiTFastAttn 时，需要在计算效率和生成质量之间找到最佳平衡点。对于视频生成，更高的压缩率（D5、D6）可能导致感知质量下降，但在计算资源受限的场景下仍有其价值。

6.1.5. FLOPs 减少和加速

以下是原文 Table 2 的结果：

• 单一技术效果：

  • ASC (跨 CFG 注意力共享) 技术：将注意力计算量减少 50%，因为分类器自由引导 (CFG) 需要进行两次推理，ASC 通过重用其中一次的注意力输出，直接将计算量减半。延迟 (latency) 降低约 46%~48% (从 100% 降至 54%~52%)。
  • WA-RS (带有残差共享的窗口注意力) 技术：对注意力计算的减少量与序列长度 (Sequence Length)（即图像分辨率）呈负相关。对于 DiT-XL-2 512x512 (Seqlen 1024)，WA-RS 减少了 77% 的注意力 FLOPs；对于 PixArt-Sigma-XL 2048x2048 (Seqlen 16384)，它减少了 33% 的注意力 FLOPs。这是因为在序列长度较长时，即使使用窗口注意力，窗口内的词元数量也相对较多，但其相对节省依然显著。延迟降低约 15% (512x512) 到 65% (2048x2048)。
  • AST (跨时间步注意力共享) 技术：在 FLOPs 计数上显示为 0% 减少，这可能是因为它通过跳过完整模块 (full module) 的计算来节省实际计算时间，但其 FLOPs 计数方式可能与 ASC 和 WA-RS 的“内部模块计算减少”不同。但它带来了 1% 到 4% 的注意力延迟降低，这表明它确实减少了实际的计算开销。

• 组合技术效果：

  • WA-RS + ASC 组合技术：由于 WA-RS 和 ASC 是正交的（解决不同类型的冗余），它们可以同时应用而不会带来额外的开销。这种组合带来了最大的 FLOPs 减少和延迟降低。对于 PixArt-Sigma-XL 2048x2048，FLOPs 降至 16%，延迟降至 19%，这意味着高达 84% 的 FLOPs 减少和 81% 的延迟降低。

• 结论： 随着序列长度和分辨率的增加，DiTFastAttn 带来的注意力 FLOPs 减少和延迟 (latency) 降低效果越显著。

  下图 (原文 Figure 8) 展示了不同分辨率图像生成在不同压缩比下的延迟 (latency) 表现。

  该图像是一个图表，展示了在不同分辨率下的图像生成延迟与注意力计算 TFLOPs 的关系。左侧是 512x512 尺寸的 DiT-XL，中间为 1024x1024 尺寸的 PixArt-Sigma-XL，右侧则为 2048x2048 尺寸的 PixArt-Sigma-XL。蓝线表示图像生成延迟，橙线表示多头注意力模块的延迟。

图示：不同分辨率图像生成在不同压缩比下的延迟。DiT 以 batch size 8 运行，PixArt-Sigma 模型以 batch size 1 运行。蓝线表示端到端图像生成 (end-to-end image generation) 延迟，橙线表示多头注意力模块 (multi-head attention module) 的延迟。

• 对所有模型的分析：
  • DiTFastAttn 在所有模型和所有压缩比设置下，都实现了端到端延迟 (end-to-end latency) 的降低。
  • 随着注意力计算的减少，图像生成总延迟和注意力模块延迟均随之下降。
• 高分辨率模型的优势：
  • PixArt-Sigma-2K 表现出最佳的加速性能。在 D6 配置下，其总生成延迟降至原始模型的 56%，注意力总延迟降至原始模型的 37%。
  • 这进一步印证了随着分辨率的增加，DiTFastAttn 在降低整体注意力和图像生成延迟方面取得了更好的效果。

6.1.6. 消融实验

下图 (原文 Figure 9) 展示了在 DiT-XL-2-512 模型上进行的消融实验结果，旨在评估不同技术的影响。

该图像是图表，展示了DiTFastAttn与其他方法（如ASC、WA-RS和AST）在不同Attention TFLOPs下的Inception Score表现。左侧图表比较了不同方法的得分与TFLOPs的关系，中间图表分析了不同时间步对得分的影响，右侧图表则对WA-RS和WA进行了直接比较。

图示：在 DiT-XL-2-512 上的消融实验。考察方法影响（左）、时间步变异性（中）和残差共享（右）。“WA”表示没有残差共享（RS）的窗口注意力。

• DiTFastAttn 优于单一方法 (左图)：

  • 图表比较了 DiTFastAttn 与 ASC、WA-RS 和 AST 等单一技术在相同计算预算（Attention TFLOPs）下，Inception Score (IS) 表现。
  • DiTFastAttn 始终保持了更高的生成质量 (IS)，尤其是在低 TFLOPs（高压缩比）区域。
  • 在单一技术中，AST 显示出最好的生成质量，但当计算量降至 2.2 FLOPs 以下时，其性能会显著下降，导致搜索算法终止。这表明 AST 有其极限。
  • DiTFastAttn 能够支持更激进的压缩，同时保持更好的质量，这归因于其结合多种技术并进行智能决策的能力。

• 更多时间步提升 DiTFastAttn 性能 (中图)：

  • 该图比较了 DiTFastAttn 在不同时间步数量（50、100、150、200）下的性能。
  • 结果显示，随着时间步数量的增加，DiTFastAttn 能够在保持甚至提高生成质量 (IS) 的同时，实现更大的计算压缩。
  • 这可能是因为更多的去噪时间步提供了更多的冗余利用机会，或者更长的去噪过程本身对少量误差的容忍度更高。

• 残差缓存对于保持性能至关重要 (右图)：

  • 该图对比了带有残差共享的窗口注意力 (WA-RS) 与纯窗口注意力 (Window Attention, WA)（不带残差共享）在相同压缩比下的生成性能。
  • WA-RS 显著优于 WA。不使用残差共享的窗口注意力会导致生成性能的显著下降。
  • 这验证了残差 (residual) 缓存技术对于 WA-RS 至关重要，它有效地弥补了局部注意力可能丢失的远程依赖信息，从而在保持计算效率的同时，维持了高质量的生成。

6.2. 数据呈现 (表格)

本节已在 6.1.1. 图像生成性能和注意力 FLOPs 和 6.1.5. FLOPs 减少和加速 中完整转录了原文 Table 1 和 Table 2。原文 Appendix 中还提供了更详细的延迟数据，这里也进行转录。

以下是原文 Table 3 的结果：

以下是原文 Table 4 的结果：

以下是原文 Table 5 的结果：

这些表格提供了在不同配置（D1-D6）下，不同 DiT 模型在图像生成任务中的延迟 (latency)、注意力延迟 (Attn Latency) 以及质量指标（FID、IS、CLIP Score）的详细数据。

• 延迟与压缩关系： 可以看出，随着压缩等级的增加（D1到D6），总延迟 (Latency) 和注意力延迟 (Attn Latency) 均显著下降。例如，PixArt-Sigma-XL 2048x2048 在 D6 配置下，总延迟从 39.86s 降至 22.27s，注意力延迟从 27.57s 降至 10.13s，实现了大幅加速。
• 质量权衡： 在较低压缩等级（D1-D3），模型能很好地保持甚至略微提升质量指标。在较高压缩等级（D4-D6），质量指标可能会有一定程度的下降，这体现了压缩效率和生成质量之间的权衡。
• 高分辨率优势： 如前所述，高分辨率模型在压缩方面具有更大潜力，能在更激进的压缩下保持相对较高的生成质量。

6.3. 压缩计划搜索时间

以下是原文 Table 6 的结果：

• 搜索时间与阈值关系： 压缩计划的搜索时间与阈值 $\delta$ 成反比，即随着阈值的增加（更激进的压缩），搜索时间会减少。这是因为更宽松的阈值允许算法更早地找到满足条件的压缩策略，从而更快地终止内部循环。
• 可接受的开销： 尽管搜索时间从几分钟到一小时不等，但考虑到一次搜索可以为模型提供持续的推理加速，这种一次性开销是合理的。例如，生成一张 512x512 图像需要约 2 秒，而搜索一个压缩计划大约需要 $2 \text{s} \times 28 \text{ (Transformer 层数)} \times 4 \text{ (方法候选)}$，约为 224 秒（约 4 分钟），这相对于整体推理时间来说是可接受的。

6.4. 负面条件 (Negative Conditioning)

下图 (原文 Figure 14) 展示了使用 PixArt-Sigma-XL-1024 在有/无负提示（如“Low quality”）下，在不同阈值（D2、D4、D6）生成的图像。

该图像是展示了使用 PixArt-Sigma-XL-1024 在不同阈值下生成的图像，包括有无负提示的比较。上部分为车辆图像，底部为蓝色房间图像。图中分别标注了 D2、D4 和 D6。

图示：PixArt-Sigma-XL-1024 在不同阈值下有/无负提示生成的图像。

• 实验表明，DiTFastAttn 在应用通用负提示（例如“Low quality”）时，能够很好地保留负提示对生成图像的影响。这意味着即使在压缩模式下，模型仍能理解并排除负面描述的内容，从而保持生成图像的质量和控制力。

7. 总结与思考

7.1. 结论总结

本文提出了 DiTFastAttn，一个针对扩散 Transformer (DiT) 模型的后训练 (post-training) 注意力压缩方法。该方法的核心在于识别了 DiT 推理过程中注意力计算的三种关键冗余：空间冗余、时间冗余和条件冗余。

为解决这些冗余，DiTFastAttn 提出了三项创新技术：

1. 带有残差共享的窗口注意力 (WA-RS)：通过局部注意力减少空间冗余，并通过缓存和重用残差 (residual) 来有效保留长距离依赖，避免性能下降。

2. 跨时间步注意力共享 (AST)：利用去噪时间步之间的注意力输出相似性，通过共享减少重复计算。

3. 跨分类器自由引导注意力共享 (ASC)：针对分类器自由引导 (CFG) 机制，通过重用条件推理的注意力输出，大幅降低计算成本。

   此外，研究还开发了一种贪婪方法 (greedy method) 来动态决定每个层和时间步的最佳压缩策略。

实验结果在 DiT、PixArt-Sigma（图像生成）和 OpenSora（视频生成）模型上验证了 DiTFastAttn 的有效性。该方法显著减少了注意力浮点运算次数 (FLOPs)（高达 76%），并在高分辨率图像生成中实现了高达 1.8 倍的端到端加速 (end-to-end speedup)，同时在保持生成质量方面表现出色，尤其在高分辨率任务中优势更明显。这使得 DiT 模型在实际部署中更具效率和可行性。

7.2. 局限性与未来工作

论文作者指出了 DiTFastAttn 的以下局限性：

1. 后训练方法 (Post-Training Limitation)： 作为一个后训练 (post-training) 压缩技术，DiTFastAttn 无法利用训练过程来优化性能或避免潜在的性能下降。这意味着它必须在不改变模型学习到的参数的情况下工作，这限制了其可能达到的最佳压缩效果。

2. 侧重推理加速，非显存减少 (VRAM Reduction)： 该方法主要关注推理速度的提升，而非显存 (VRAM) 的减少。特别是当应用 AST (跨时间步注意力共享) 时，为了共享注意力隐状态 (hidden states)，需要存储前一个时间步的注意力隐状态，这反而可能带来额外的显存使用。

3. 贪婪搜索的次优性 (Suboptimal Compression Plan)： 提出的简单贪婪方法 (greedy method) 可能无法找到全局最优的压缩计划。更复杂的优化算法可能会发现更好的策略组合。

4. 仅优化注意力模块 (Attention Module Only)： DiTFastAttn 仅减少了注意力模块的计算成本。DiT 模型中还包含其他计算密集型模块（如前馈网络，Feed-Forward Networks），这些模块的成本并未被优化。

   基于这些局限性，未来可能的研究方向包括：

• 探索将 DiTFastAttn 与训练时优化 (training-time optimization) 方法相结合，以实现更大的性能提升和压缩率。
• 研究能够同时减少计算和显存 (VRAM) 的方法，特别是在处理 AST 引入的额外显存开销方面。
• 开发更先进的压缩计划搜索 (compression plan search) 算法（例如，基于强化学习或更复杂的优化方法），以找到更优的策略组合。
• 将压缩方法扩展到 Transformer 的其他模块，如前馈网络 (Feed-Forward Networks)，以实现更全面的模型加速。

7.3. 个人启发与批判

7.3.1. 个人启发

这篇论文提供了一个非常实用的视角来解决大型预训练模型（尤其是 DiT 这种高计算需求模型）的部署效率问题。

1. 冗余利用的普适性： 识别并利用模型推理过程中的冗余是一个非常通用且强大的优化思路。DiTFastAttn 在扩散模型的特定上下文（空间、时间步、CFG）中，精妙地发现了这些冗余，并设计了相应的解决方案。这种“对症下药”的方法比通用压缩技术更有效。
2. 后训练压缩的巨大价值： 对于动辄数百万甚至数十亿参数的预训练模型，重新训练的成本几乎是不可承受的。DiTFastAttn 证明了后训练 (post-training) 压缩在不牺牲过多性能的情况下，可以带来巨大的实际效益。这对于推动这些先进模型走向更广泛的应用场景（例如边缘设备、资源有限的云环境）具有重要意义。
3. 残差共享的巧妙设计： WA-RS 中的残差共享 (residual sharing) 机制非常巧妙。它解决了局部注意力固有的问题（丢失全局信息），通过简单的缓存和重用，以极低的成本恢复了大部分性能，这比重新设计复杂的全局-局部注意力机制要高效得多。这种“小修小补”却能带来大收益的思想值得借鉴。
4. 高分辨率场景的显著收益： 论文清晰地展示了在高分辨率生成任务中，DiTFastAttn 带来了更大的 FLOPs 减少和延迟降低。这直接解决了 DiT 模型在实际应用中的一个主要痛点，使得生成 2K 甚至 4K 图像和视频变得更加高效。

7.3.2. 批判性思考

尽管 DiTFastAttn 提供了强大的解决方案，但仍有一些值得批判性思考和改进的地方：

1. AST 的 FLOPs 计数解释： 在 Table 2 中，AST 的 Attn FLOPs 减少量显示为 0%。这在直观上似乎有些矛盾，因为 AST 确实通过跳过计算来节省时间（Attn Latency 有所减少）。这可能意味着 FLOPs 的计算方式并未完全捕捉到 AST 带来的“跳过整个模块计算”的效益，或者 AST 更多是一种时间优化而非底层算术运算的减少。如果能更清晰地解释 AST 的 FLOPs 计数逻辑，将有助于初学者理解。

2. 贪婪搜索的局限性： 论文承认其贪婪方法 (greedy method) 可能无法找到最优的压缩计划。虽然它在实际中可能足够好，但探索更先进的优化技术（如使用遗传算法 (genetic algorithms)、贝叶斯优化 (Bayesian optimization) 或强化学习 (reinforcement learning)）来搜索压缩计划，可能会发现全局最优或更好的权衡点。然而，这些方法本身的计算成本也会更高，需要权衡。

3. 显存 (VRAM) 影响的更深入分析： 论文提及 AST 可能增加显存 (VRAM) 使用。对于资源受限的部署场景，显存占用与计算速度同等重要。如果能够提供 DiTFastAttn 在不同配置下的详细显存使用数据，并与原始模型进行对比，将更有助于评估其在实际系统中的适用性。

4. 负面条件 (Negative Conditioning) 的定量分析： 论文在附录中展示了负面条件下生成图像的视觉效果。如果能提供更定量的指标（例如，使用 CLIP Score 评估负面提示的抑制效果，或使用其他感知指标），将使这一分析更具说服力。

5. 与其他后训练压缩方法的直接对比： 尽管论文强调了其方法的创新点，但如果能与一些已有的、同样是后训练 (post-training) 的注意力稀疏化或量化方法进行直接的性能对比，将能更全面地评估 DiTFastAttn 的相对优势。

   总的来说，DiTFastAttn 是一项扎实且具有高度实用价值的工作，它巧妙地利用了扩散 Transformer 模型的内在冗余，为高分辨率生成任务提供了有效的后训练 (post-training) 加速方案。其创新点和实验结果都令人印象深刻，对于推动扩散模型在真实世界的应用具有重要意义。