全文翻译
摘要
扩散Transformer（DiT）模型凭借其卓越的生成能力和可扩展性，在图像生成领域取得了显著成功。然而，扩散模型（DMs）的迭代特性导致计算复杂度较高，给模型部署带来了挑战。尽管现有的基于缓存的加速方法试图利用时间上的固有相似性来跳过DiT的冗余计算，但缺乏校正可能会导致潜在的质量下降。在本文中，我们提出了增量校准缓存（increment-calibrated caching）方法，这是一种用于DiT加速的无训练方法，其校准参数通过预训练模型自身的低秩近似生成。为解决异常激活可能导致的校正失败问题，我们引入了通道感知奇异值分解（channel-aware Singular Value Decomposition, SVD），进一步增强了校准效果。实验结果表明，在计算资源预算相近的情况下，我们的方法始终比现有的朴素缓存方法表现更优。与35步DDIM相比，我们的方法可减少超过45%的计算量，将Inception Score（IS）提高12，同时FID的增加量小于0.06。代码可在https://github.com/ccccczzy/icc获取。
1. 引言
扩散模型（DMs）[14, 29, 37, 38]已成为最受欢迎的生成模型之一，尤其在图像[31]、视频[36]和音频[6]等连续模态领域展现出卓越的生成质量和多样性。在DMs的反向过程中，期望的内容通过多步细化逐渐合成。尽管最初的DMs基于UNet架构[32]，但最近的研究表明，基于Transformer的DMs（即扩散Transformer，DiTs）可作为一种具有竞争力的替代方案[30]。
DiTs受益于Transformer架构固有的卓越可扩展性[40]，这一点已在图像识别[9]和语言建模[18]等各种任务中得到广泛验证。
然而，DiTs仍然受到DMs固有迭代特性的困扰。这意味着其高保真的合成结果通常以大量延迟为代价，加剧了部署难度，因此亟需提高效率。现有的DiT模型加速策略可主要分为三类：时间步长减少、基于权重的近似[5]和基于缓存的加速[26, 35]。
第一类方法致力于减少所需的反向步数，包括采用快速采样器[23, 24, 37, 44]或将原始轨迹蒸馏为更少的时间步[33, 39]。
基于权重近似的方法灵感来源于降低神经网络复杂度的通用方法。这些工作通常专注于通过剪枝[11]、采用低位量化[5, 12]或使用如奇异值分解（SVD）等低秩近似来消除冗余的权重元素。然而，通过减小DiTs的模型规模来提高生成效率是一项艰巨的任务，尤其是在无法重新训练的情况下。正如图2所示，基于SVD近似的实验结果表明，即使在所有层中移除约33%的秩，计算开销也不会相应减少。此外，生成的图像质量严重下降，并且与原始模型失去一致性。
相反，基于缓存的加速技术旨在利用扩散过程的多步特性，并利用固有的时间相似性来绕过冗余计算[16, 26, 35]。在这种情况下，先前时间步的中间特征被缓存，并直接用于近似后续时间步的输出，而无需任何校正。当前相关工作主要集中在设计最优的缓存机制，而忽略了对缓存值的校准。这种忽视可能导致质量下降，并限制了基于缓存加速的全部潜力。
在本文中，我们提出了增量校准缓存，这是一种新颖的无训练方法，用于加速DiT，它与上述三种方法不同。它可以被视为基于权重近似和基于缓存加速的互补组合，而不是简单地将这两种技术拼凑在一起。与直接使用SVD将原始模型转换为简化版本的方法不同，我们将低秩近似的权重视为校准参数。在反向过程中，这些校准参数通过增量项用于细化从原始模型获得的缓存值，引导其朝向期望的准确结果。这确保了增量校准缓存与传统的朴素缓存机制相比，在效率和性能之间实现了更好的平衡。
此外，我们用通道感知SVD代替了原始的SVD，这是一种改进的变体，考虑了通道敏感性。与原始版本相比，通道感知SVD通过更关注激活敏感的通道来增强校准效果。这一点至关重要，因为正如在各种基于Transformer的模型[43, 45]中所观察到的，DiT中的异常值问题会对推理结果产生不利影响。本文的主要贡献总结如下：
- 我们引入了增量校准缓存，这是一种专为DiT加速设计的无训练方法。据我们所知，这是首次旨在以无训练方式校准DMs缓存的工作。
- 我们提出了通道感知SVD，这是一种改进的SVD变体，可改善DiT的增量校准缓存，并缓解可能破坏校准效果的异常值问题。
- 我们在最先进的DiT模型DiT-XL/2[30]上进行了条件图像合成和PixArt-α[4]文本到图像生成的验证。实验结果表明，在计算预算相近的情况下，我们提出的方法大大优于朴素缓存方法。与35步DDIM相比，我们的方法消除了超过45%的计算量，并将IS提高了12，同时FID的增加量小于0.06。
2. 相关工作
扩散Transformer
早期的扩散模型（DMs）严重依赖U-Net架构，而近期的工作已主要转向基于Transformer的DMs，特别是扩散Transformer（DiTs）。DiT[30]最初应用于类条件图像合成，证明了Transformer架构在扩散框架内的可扩展性和有效性。后续工作如PixArt-α[4]及其后续版本[2,3]通过融入交叉注意力机制和合并自适应层归一化，进一步适配和优化了原始DiT，从而将其能力扩展到文本到图像任务。值得注意的是，最新版本的Stable Diffusion[10]也已转向多模态DiT架构。在视频生成领域，SORA[1]的出现体现了基于DiT构建世界模拟器的潜力，进一步突显了DiT的多功能性和前景。
基于缓存的加速
尽管DMs的迭代特性导致了显著的复杂性，但它也为基于缓存的加速带来了独特机会——利用连续去噪步骤之间的时间相似性。这种相似性允许缓存和重用中间结果，从而绕过冗余计算。一些早期工作致力于利用基于U-Net的DMs的冗余性。例如，DeepCache[27]利用高层特征的时间一致性来简化U-Nets主分支中的计算。类似地，Block Caching[42]引入了一种动态缓存机制，该机制依赖于累积的输入变化。
最近的工作已将重点转向将这些原理应用于基于Transformer架构的DMs。例如，FORA[35]定期存储和检索所有注意力层和MLP层的输出。同时，L2C[26]和HarmoniCa[16]旨在通过基于学习的方法确定最佳缓存策略。Δ-DiT[7]在早期采样阶段加速后部DiT块，在后期阶段加速前部DiT块。最后，TokenCache[22]和ToCa[47]采用基于令牌的缓存方法，而非基于层的缓存，通过区分不同令牌的重要性来实现。
低秩近似
奇异值分解（SVD）及其变体等低秩近似技术已广泛应用于神经网络的压缩和加速[17]。最近，研究人员探索了SVD在大型语言模型（LLMs）中的应用。例如，FWSVD[15]在应用SVD之前结合Fisher信息来权衡参数的重要性。ASVD[45]通过基于激活分布缩放权重矩阵来解决异常值问题。此外，SVD-LLM[41]在奇异值和推理损失之间建立了直接关联。然而，关于SVD在DMs加速中的应用探索有限。我们的实验结果（如图2所示）表明，直接将低秩近似应用于DiT会显著降低生成质量。
3. 方法论
3.1 预备知识
扩散模型的基本框架由两部分组成，即正向过程和反向过程。正向过程描述了在T个时间步内将干净数据$z_0$逐渐转换为高斯噪声$z_T \sim N(0, I)$的过程。在每个时间步$t \in [1, T]$，随机噪声$\epsilon_t$会以噪声调度因子$\beta_t$添加到当前数据$z_{t-1}$中，公式如下：
反向过程从服从标准正态分布的$z_T$开始采样。为了生成服从真实分布的新样本，反向过程使用采样器$\Psi(\cdot)$迭代估计累积噪声并从当前数据中移除。以DDPM[14]为例，其表达式为：
其中，$\alpha_t$、$\overline{\alpha}_t$和$\sigma_t$是依赖于时间步t的常数，$\epsilon \sim N(0, I)$。$\epsilon_\theta(z_t, t)$表示由θ参数化的噪声估计网络，通常采用UNet[14, 29, 31, 37]或Transformer[4, 10, 20, 25, 30, 46]架构。
3.2 扩散Transformer的朴素缓存
DiT模型由堆叠的模块组成，每个模块包含一个多头自注意力（MHSA）层和一个前馈网络（FFN）层，可表示为$F_l(\cdot)_{l=1}^L$，其中L表示网络深度，$F_l$可以是MHSA层或FFN层。由于连续时间步的输出通常具有内在相似性，朴素缓存方法被提出用于重用这两种层的结果并加速反向过程。
如图3(a)所示，朴素缓存机制的关键模块是缓存c、收集操作G和分散操作s。在生成过程中，缓存存储未来可能重用的中间输出，由收集和分散操作共同维护。收集和分散操作均可表示为二进制矩阵，即$G, S \in \{0, 1\}^{T \times L}$，其中元素$G_{t, l}$和$S_{t, l}$分别表示在时间步t是否存储或加载第l层的结果。
给定时间步s的输入$x_{s, l}$，如果收集矩阵对应的位$G_{s, l}$被激活，第l层将被完全计算以更新缓存：$C_l(y) \leftarrow F_l(x_{s, l})$。随后，对于后续时间步m，仅当$S_{m, l} \neq 1$时，第l层需要重新计算；否则，将使用缓存值进行近似以绕过计算：$y_{m, l} \approx C_l(y)$。为了在确保生成质量的同时最小化总计算量，收集和分散矩阵应通过手工设计[35]或基于学习的方法[16, 26]精心设计。
3.3 基于SVD的增量校准缓存
然而，朴素缓存策略通常存在质量下降的风险。解决此问题的一个潜在方法是用增量项校准缓存值，从而引导其朝向准确结果。我们从一个权重为$W \in \mathbb{R}^{C_o \times C_i}$的简单线性层开始讨论，其中$C_i$和$C_o$分别表示输入和输出通道数。该方法可通过校准内部的所有线性操作扩展到MHSA或FFN。
与前一小节建立的符号一致，假设前一步s的输入$x_{s, l}$也与输出$y_{s, l}$同步存储为$C_l(x) \leftarrow x_{s, l}$。因此，时间步m的准确结果$y_{m, l}$可重新表示为缓存输出与增量项之和：
直接计算增量项会导致与完整计算相同的$O(NC_iC_o)$复杂度，其中N是令牌数量。为了在效率和性能之间取得平衡，基于SVD对W进行低秩近似。通过SVD，权重矩阵W可分解为三个矩阵：U、$\sum$和$V^T$，使得$W = U\sum V^T$。对角矩阵$\sum$由按降序排列的正奇异值组成。截断近似保留最大的r个奇异值，以在秩r约束下最小化重建误差。然后，U、$\sum$、$V^T$被替换为$U_r \in \mathbb{R}^{C_o \times r}$、$\sum_r \in \mathbb{R}^{r \times r}$、$V_r^T \in \mathbb{R}^{r \times C_i}$，将W近似为$W_r = U_r\sum_r V_r^T = W_r^a W_r^b$，其中$W_r^a \in \mathbb{R}^{C_o \times r}$和$W_r^b \in \mathbb{R}^{r \times C_i}$作为校准参数保留。
因此，复杂度降低到$O(N(C_i + C_o)r)$，缓存输出可校正为：
3.4 基于通道感知SVD的改进校准
如先前工作[15, 41, 45]所示，直接应用基于SVD的低秩近似可能导致神经网络性能不可接受的下降，尤其是在仅保留有限数量的秩时。原始SVD方法仅关注静态权重，未能考虑输入激活的影响，导致重建误差与实际模型性能不匹配。对于DiT的增量校准缓存，此问题可能更为严重。图4展示了DiT-XL/2[30]单个模块内所有线性层的激活分布，值得注意的是，输入和输出通道维度均存在异常值。这意味着即使是微小的重建误差也可能显著削弱某些通道的校准效果。因此，在校正过程中，必须高度关注这些敏感通道的权重元素。
为此，我们采用通道感知SVD，在近似给定权重矩阵时考虑输入和输出通道的敏感性。假设不同通道的重要性可以由两个对角矩阵$S_i \in \mathbb{R}^{C_i \times C_i}$和$S_o \in \mathbb{R}^{C_o \times C_o}$加权，其中每个元素分别对应一个输入或输出通道。然后，权重矩阵W可缩放为$W’$以考虑通道敏感性：
为了降低其复杂度，对$W’$进行SVD分解，并通过仅保留前r个奇异值将其近似为$U’_r \sum’_r V’^T_r$。然后，我们用矩阵$S_i$和$S_o$的逆对近似的$W’$进行重新缩放，以获得原始矩阵W的近似值。它可以表示为矩阵$\hat{W}_r^a \in \mathbb{R}^{C_o \times r}$和$\hat{W}_r^b \in \mathbb{R}^{r \times C_o}$的乘积，将用作校准参数：
为了获得矩阵$S_o^{-1}$和$S_i^{-1}$，我们提出了两种可行方法：通道激活感知SVD（CA-SVD）和通道差异感知SVD（CD-SVD）。前者是ASVD[45]的扩展版本，其中基于小型校准数据集计算跨通道激活的平均幅度并用作缩放因子。CD-SVD利用时间步间的平均差异来生成缩放矩阵，如算法1所示。
4. 实验
4.1 实验设置
模型、数据集和指标：为了验证所提出方法的有效性，我们基于常用的扩散Transformer模型DiT-XL/2[30]进行类条件图像合成实验，并基于PixArt-α[4]进行文本到图像生成实验。默认采样器分别为DDIM[37]和DPM[23]。从ImageNet[8]的1000个类别中采样50,000张分辨率为256×256的图像，使用Inception Score（IS）[34]、Fréchet Inception Distance（FID）[28]、sFID、Precision[19]和Recall来评估性能。对于文本到图像生成，我们从MSCOCO-2014[21]中随机采样30,000条 caption 来生成图像，分别使用FID-30k和Clip Score[13]来评估图像质量和文本-图像对齐度。计算复杂度通过乘加运算（MAC）的数量来衡量，这是计算密集型矩阵乘法的基本单位。
缓存模式设置：选择FORA[35]的朴素缓存模式作为基线，其中所有MHSA和FFN层的输出根据固定周期p同步缓存。因此，收集矩阵G设置为：$G_{t,l}=\mathbb{1}_{p|t}$，$\forall t \in [1,T]$，$l \in [1,L]$。分散矩阵S可视为G的补集：$S_{t,l}=1-G_{t,l}$。此外，需要注意的是，所提出的校准方法不仅可以应用于这种特定情况，还与3.2小节中定义的其他朴素缓存机制兼容。
校准机制设置：为简化讨论，我们对所有线性层使用统一的秩r来生成校准参数。为了进行CA-SVD或CD-SVD，从训练集中随机选择256张图像作为校准集。
4.2 与朴素缓存的比较
统一缓存模式：将所提出的增量校准缓存与具有相同缓存模式（即FORA[35]）的朴素缓存进行比较。为确保公平性，我们将所有方法的计算开销限制在相同水平，并研究它们在不同采样设置下的性能差异。如表1所示，对于类条件合成，当两者都采用相同的FORA缓存模式时，我们的增量校准缓存在所有性能指标上都优于朴素缓存，尤其是在计算资源受限的情况下。例如，在MAC限制约为2.37T时，所提出的增量校准缓存方法的IS得分比朴素缓存高22分，同时FID和sFID分别降低1.22和3.37。在文本到图像生成中也观察到类似结果，如表2所示。
即使不使用通道感知SVD进行改进，增量校准缓存在FID和Clip Score上也优于朴素缓存。CA-SVD的应用进一步改善了结果，实现了0.51～4.75的FID降低。
不同缓存模式：在这项工作中，我们没有过多关注缓存模式的设计，而是采用了FORA[35]使用的粗粒度方法。另一方面，L2C是一种仅缓存的加速技术，进行细粒度优化，从而进一步最大化缓存的潜力。然而，如表3所示，由于增量校准的校正效果，我们提出的缓存方法与L2C[26]取得了相当的性能，并将计算开销降低了23%。值得注意的是，L2C采用基于学习的方法来优化其缓存模式，需要额外的训练。相反，我们的方法的校准参数是以无训练方式导出的，并且可以在单个NVIDIA 4090D GPU上在几分钟内获得，显著降低了部署的复杂性和工作量。
4.3 与快速采样的比较
对于DMs，减少计算开销的常见做法是以质量损失为代价减少采样步骤。图5和图6展示了增量校准缓存和快速采样器的复杂度-质量权衡曲线。对于类条件合成，当给定固定的计算预算时，我们提出的方法在IS、FID和sFID方面优于DDIM，尤其是在复杂度受限的情况下。这意味着，使用增量校准缓存，有望根据指定指标的下限大幅加速采样。例如，当FID的最大值设置为约2.5时，30步的增量校准缓存以小于0.06的FID增加为代价，IS比35步DDIM提高12。同时，计算量减少了45%，这意味着增量校准缓存相对于DDIM的理论加速比达到1.8。对于文本到图像生成，所提出的方法可用于减少计算开销，同时保持甚至提高FID和Clip Score。例如，当步数达到30时，DPM采样器开始收敛，此时增量校准缓存可以消除45%的MAC操作并将FID降低0.67。
4.4 消融研究
输入和输出通道的敏感性：在通道感知SVD过程中，我们对权重矩阵的输入和输出通道都进行了缩放。为了评估这两个维度的敏感性差异，我们将CD-SVD与两个简化变体进行比较：仅缩放输入通道的CD-SVD_i和仅缩放输出通道的CD-SVD_o。如表4所示，CD-SVD在所有指标上都优于CD-SVD_i和CD-SVD_o。值得注意的是，CD-SVD_i的性能明显优于CD-SVD_o。这表明，尽管输入和输出通道都存在异常激活，但与输入通道相关的权重元素对该问题更为敏感。
校准机制的比较：在这项工作中，增量校准缓存的校准可以通过三种方法实现：原始SVD、CA-SVD和CD-SVD。我们的发现表明，没有一种方法在所有场景下都始终优于其他方法。表5展示了所有三种方法在不同秩值下的结果。当仅保留有限数量的秩作为校准参数时，CA-SVD和CD-SVD表现出优于原始SVD的性能。这种优势源于它们对通道敏感性的考虑，使它们能够过滤出最重要的秩，从而改善FID、sFID和IS指标。然而，随着用于生成校准参数的秩数量增加，CA-SVD或CD-SVD的应用可能会对FID和sFID产生负面影响，尽管它们继续改善IS指标。
不同秩的影响：根据简单的直觉，一个自然的假设是，当校准参数的秩r增加时，图像质量将平滑提高。然而，我们的实验表明，情况并非总是如此，如图7所示。对于文本到图像生成的结果，尽管FID-30k最初随着r的增加而降低，但当r达到256时开始增加。尽管采用通道感知SVD改善了结果，但其曲线仍然不是单调递减的。另一个相关现象是，一些校准设置甚至优于具有相同步数的非缓存采样。这意味着并非所有奇异值都对增量校准的效果有积极影响。因此，所提出的方法有可能在不牺牲生成质量的情况下实现显著的复杂度降低。
结论
在这项工作中，我们提出了增量校准缓存（increment-calibrated caching），这是一种以无训练方式降低DiT计算复杂度的新方法。增量校准缓存结合了低秩近似和基于缓存加速的思想，对前一时间步的缓存激活进行校正，并引导其在后续时间步中朝向所需的准确结果。该过程通过基于SVD的低秩近似从预训练模型自身生成的校准参数来确保。我们进一步引入了通道感知SVD，该方法考虑了通道敏感性并防止了有害的离群值问题。实验结果表明，我们的方法在性能上大幅优于朴素缓存方法。此外，与35步DDIM相比，我们的方法消除了超过45%的计算量，并将IS提高了12，同时FID的增加量小于0.06。