全文翻译
摘要
扩散Transformer（DiTs）已实现了最先进（SOTA）的图像生成质量，但存在延迟高和内存效率低的问题，使其难以在资源受限的设备上部署。一个主要的效率瓶颈是，现有的DiTs在图像的所有区域上应用了同等的计算量。然而，并非所有图像令牌都同等重要，某些局部区域（如物体）需要更多计算。为解决这一问题，我们提出了DiffCR，这是一种具有可微分压缩比的动态DiT推理框架，它能自动学习为每个图像令牌跨层和跨时间步动态分配计算资源，从而实现高效的DiTs。具体而言，DiffCR集成了三个特性：（1）令牌级路由机制，其中每个DiT层包含一个路由器，该路由器与模型权重联合微调以预测令牌重要性分数。通过这种方式，不重要的令牌可以绕过整个层的计算；（2）层级可微分比率机制，不同的DiT层从零初始化开始自动学习不同的压缩比，使得冗余层的压缩比较大，而其他层的压缩比较小甚至不压缩；（3）时间步级可微分比率机制，每个去噪时间步学习其自身的压缩比。由此产生的模式显示，在噪声较大的时间步压缩比较高，而随着图像变得更清晰，压缩比逐渐降低。在文本到图像和图像修复任务上的大量实验表明，DiffCR有效地捕捉了令牌、层和时间步三个维度的动态性，与先前的工作相比，在生成质量和效率之间取得了更优的权衡。
1. 引言
扩散模型最近在图像生成方面展现出了卓越的性能，其架构从U-Net[13, 30, 34, 41]发展到了Transformer[1, 3, 4, 29]。其中，扩散Transformer（DiTs）[1, 29]因其出色的可扩展性而脱颖而出。然而，扩散模型，特别是DiTs，受到大量计算和内存需求的限制，这降低了它们在生成和部署方面的效率。例如，在单个A100 GPU上以完整上下文生成一张1024像素的图像可能需要19.48秒，并且需要超过40GB的GPU内存[28, 30]。大多数DiTs的一个主要效率瓶颈源于在所有图像区域上均匀地应用计算，尽管不同区域的复杂程度存在差异[28, 32, 57]。这种效率瓶颈表明，理想的DiT推理框架可以在DiTs的三个关键维度（令牌、层和时间步）上进行自适应和动态的计算。
已经提出了多种技术来解决上述三个关键维度上的效率瓶颈：（1）令牌合并[2]、剪枝[45]和下采样[39]；（2）层[17]或通道[9]剪枝；（3）少步蒸馏和采样[23, 35, 46, 56]。尽管这些技术很有前景，但（1-2）中的技术在很大程度上依赖于启发式方法，例如用于令牌重要性的启发式规则以及通道和层剪枝规则。此外，压缩比在各层之间通常是统一的，或者是根据先前的经验凭经验调整的。另外，大多数方法只关注单一的效率维度，忽略了在所有三个维度上结合优化所产生的复合效应。
为了实现一个统一的、可学习的动态DiT推理框架，使其在层和时间步上具有可微分的压缩比，必须解决三个关键挑战：（1）令牌视角：开发一种可学习的、适应视觉内容的令牌重要性度量，因为并非所有令牌都同等重要。（2）层视角：设计能够自主学习每个层的自适应压缩比的机制，以优化处理效率，因为并非所有层的贡献都同等重要。（3）时间步视角：开发能够跨时间步有效学习和应用压缩比模式的方法，因为并非所有时间步都同等重要。为了应对这些挑战，我们做出了以下贡献：
- 我们提出了一种具有可微分压缩比的动态DiT推理框架，称为DiffCR，它能自动为每个图像令牌学习跨层和跨时间步的最优动态计算，从而为内容生成任务得到高效的DiT模型。
支撑技术1：我们采用了一种受深度混合（MoD）[33]启发的令牌级路由机制，该机制能自动学习令牌重要性分数。每个DiT层都包含一个轻量级路由器，该路由器与模型权重联合微调。根据压缩比，不太重要的令牌会绕过整个层的计算。据我们所知，我们是第一个将MoD应用于视觉领域的。我们的路由分析表明，令牌重要性在不同层和时间步上存在差异。
支撑技术2：基于我们的分析，我们引入了一种新颖的DiffCR模块，该模块使令牌路由机制相对于压缩比具有可微分性，从而允许模型从零初始化开始学习每个层的自适应压缩比。冗余层会学习到更高的压缩比，而重要层则保持较低的压缩率甚至完全不压缩。
支撑技术3：我们进一步提出了一种时间步级可微分比率机制，使每个层和去噪时间步都能学习自身的压缩比。这形成了一种模式：噪声较大的时间步采用较高的压缩比，而图像较清晰的阶段则保持较低的压缩比。
我们在图像修复和文本到图像（T2I）任务上进行的大量实验一致表明，与最具竞争力的基线相比，DiffCR在生成质量和效率之间取得了更优的权衡，在保持相当的延迟和内存使用的同时，FID平均降低了8.51。
2. 相关工作
2.1. 扩散模型
扩散模型[13, 40]在图像合成任务中展现出了优于以往最先进的生成对抗网络（GANs）的性能[7]。早期的扩散模型主要采用U-Net架构。后续研究带来了多项改进，例如先进的采样方法[16, 23, 41]和无分类器引导[12]。尽管这些模型效果显著，但由于直接在像素空间中进行处理，它们存在生成延迟高的问题，从而限制了其实际应用。潜在扩散模型（LDMs）[34]的引入标志着一项重大进步，它通过训练变分自编码器（VAE）将像素空间编码到更紧凑的潜在空间中。这降低了扩散过程的计算成本，为诸如稳定扩散模型（SDMs）[30]等广泛使用的模型铺平了道路。最近，研究人员探索了将Transformer[43]架构用于扩散模型，进而开发出了扩散Transformer（DiTs）[1, 29]，该模型采用纯Transformer骨干网络，并表现出更好的可扩展性。我们提出的DiffCR是一种新颖的、具有可微分压缩比的动态DiT推理框架，并且与所有最新的DiT模型兼容。
2.2. 高效扩散模型和DiT模型
由于采用Transformer架构，DiTs[29]对资源的需求很大，其注意力模块的复杂度与令牌数量呈二次关系。先前的工作主要集中在从三个维度优化DiTs的部署效率：令牌、层和时间步。在令牌方面，研究人员引入了诸如令牌合并[2]（合并相似令牌）、令牌剪枝[45]或图像分辨率下采样[39]（移除冗余令牌）以及LazyDiffusion[28]（专门用于修复任务，绕过背景令牌的生成）等技术。在层方面，已经提出了诸如层剪枝[17]、通道剪枝[9]以及中间特征缓存[22, 25, 54]等方法来跳过冗余计算。在时间步方面，策略包括蒸馏以减少所需的时间步数量（这一方法已在UNets上进行了探索[15, 23, 35, 36, 55, 56]，尽管没有理由认为这些技术不能应用于Transformer），以及非对称采样（已应用于Transformer架构，为采样不足的阶段分配更多样本，为已经收敛的阶段分配更少样本[31, 46]）。此外，为了加速扩散文本到图像模型，还引入了更多专门的技术[3, 4]。相比之下，我们提出的DiffCR是一种可学习的、统一的动态DiT推理框架，在层和时间步上具有可微分的压缩比，探索了在所有三个维度上进行压缩所产生的复合效应。然而，在本文中我们没有探索少步蒸馏（例如[55]），因为它是一种与我们的方法互补的正交加速方法。
2.3. 动态推理
模型压缩[6]提供了一种静态方法来提高推理效率，而动态推理[33, 48, 49, 52, 58]则允许根据输入、层或其他条件进行自适应压缩。例如，早期退出方法[14, 24, 42]预测中间层内的最佳提前终止点，使模型能够在完成所有计算之前退出。动态层跳过方法[48, 49, 52]为每个输入选择性地执行层的子集，通常利用门控网络实时做出决策。在更精细的粒度上，研究人员还探索了通道跳过[9, 26]和深度混合（MoD）方法[33]，这些方法为单个令牌选择特定的层子集，而不是对整个输入进行统一处理。相比之下，我们的DiffCR是第一个引入统一的动态DiT推理框架的方法，该框架在三个维度（令牌、层和时间步）上进行优化。它还支持与网络联合微调的可微分压缩比，从而提高了适应性和效率。
3. 所提出的DiffCR框架
在本节中，我们将介绍所提出的DiffCR框架。首先，我们会对该方法进行概述。然后，详细阐述三个支撑技术：（1）第3.2节中的令牌级路由机制；（2）第3.3节中的层级可微分压缩比机制；（3）第3.4节中的时间步级可微分压缩比机制。
3.1. DiffCR概述
出于在DiT推理过程中实现统一且动态压缩的需求，DiffCR引入了一种令牌级路由机制，以动态学习每个令牌的重要性。如图1（a）所示，与先前用于自然语言处理任务的深度混合（MoD）工作[33]类似，每个DiT层都包含一个轻量级路由器，该路由器使用单个线性层，基于输入图像/噪声和文本嵌入来预测每个令牌的重要性。这使我们能够在每个层中绕过不太重要的令牌的计算，并将它们的激活直接传递到层的输出。因此，每个令牌仅由一部分选定的层进行处理。对这些路由器预测结果的可视化显示，不同的层或时间步倾向于不同的压缩比——例如，某些层优先生成物体，而其他层则专注于背景——这凸显了跨层和跨时间步进行自适应压缩的必要性。为了实现这种动态压缩，DiffCR集成了一种可微分压缩比机制，如图1（b）所示。该机制包括一个表示连续压缩比的可学习标量参数，以及作为代理比率的预定义离散比率区间。该标量查询这些区间以确定较低和较高的区间比率，从而创建具有不同压缩比的两条独立路径。最终输出是这两条路径的线性组合，其权重由学习到的比率与每个区间的接近程度决定。我们应用均方误差（MSE）损失，以确保跨层或跨时间步的平均学习比率收敛到目标比率。通过这种方式，DiffCR以可微分的方式学习自适应压缩比，从而得到高效且动态的深度混合DiT。
3.2. 支撑技术1：令牌级路由机制
动机：不同令牌的计算需求存在差异，许多令牌只需较少的层即可进行高效处理，这是我们开展此项研究的动机。我们的研究始于与MoD[33]相同的令牌级路由机制。我们移除了MoD中两个专门针对非因果自然语言处理任务的特性：具体而言，我们移除了其论文3.5节中的辅助损失和辅助MLP预测器。据我们所知，本文是首次将MoD应用于视觉领域，因此我们接下来将回顾该路由机制，进行一些可视化，并报告在视觉任务中的见解。
令牌级路由：DiT将噪声和条件文本嵌入作为输入，旨在以端到端的方式对图像进行去噪和生成。为了预测令牌的重要性，我们采用了来自MoD[33]的一种简单但有效的令牌级路由机制。如图1（a）所示，每个DiT层都包含一个轻量级路由器，该路由器由单个带有sigmoid激活函数的线性层组成，可在0到1的范围内预测每个令牌的重要性。经过路由器后，我们选择前k个最重要的令牌进行该层的处理，而其他令牌的激活则被缓存并与层输出连接，绕过包括注意力和MLP在内的整个层计算。为了在与预训练DiT模型联合微调期间使梯度能够流向路由器的权重，与MoD[33]相同，我们通过将前k个令牌的输出激活乘以路由器的预测结果来对其进行重新缩放。这种重新缩放确保了在反向传播过程中梯度能够有效地传播到路由器。k的值根据压缩比确定；我们通过实验发现，20%或30%的压缩率在延迟/内存效率和生成质量的最小下降之间实现了最佳权衡。与先前的令牌合并技术[2]不同，如图3所示，由于额外开销的存在，其计算节省与合并比率并不直接成比例，而我们方法中的MoD由于开销可忽略不计，实际延迟和内存使用量得到了更大幅度的减少。
3.3. 支撑技术2：层级可微分比率
动机：鉴于不同的层优先处理不同的物体或背景元素，因此它们受益于不同的压缩比，我们提出了一种新颖的层级可微分压缩比机制。这种方法以可微分的方式从零初始化开始自动学习每个层的压缩比，以适应不同层的冗余程度差异。
设计选择：在设计DiffCR之前，我们需要解决一个关键选择：是采用离散代理还是连续比率表示。先前的工作[5]使用具有多个压缩比候选和可学习概率的离散代理，但这种方法对于MoD存在三个挑战：（1）它缺乏有效的初始化，因为最终比率依赖于候选和概率的乘积，这使得将所有比率初始化为零变得困难；（2）MoD需要一个可微分且可学习的路由器，与需要多组前k个令牌的离散代理不兼容；（3）它引入了大量可学习参数，使训练和可解释性变得复杂，并导致难以解释的比率分布。相比之下，我们使用单个连续标量参数来表示每个层的压缩比，将参数数量减少到DiT的28层[3]所需的28个，并直接表示压缩比。
层级DiffCR：如图1（b）所示，我们为每个层分配一个可学习参数，并引入间隔为10%、范围从0%到100%的离散MoD压缩比区间。在训练期间，可学习的MoD比率查询最近的两个离散区间，以获取较低和较高的区间比率。例如，22%的可学习比率将对应20%的较低区间和30%的较高区间。然后，我们使用每个区间比率通过DiT层（带有MoD路由器）进行前向传播，产生两个输出分支。最终输出是这些分支的加权线性组合，其中权重由可学习比率与每个区间的接近程度决定——例如，对于12%的比率，输出将80%权重分配给10%的分支，20%权重分配给20%的分支。尽管这种方法在训练期间使前向传播的成本增加了一倍，但在推理期间我们只需选择最近的区间作为最终压缩比，从而消除了任何开销。为了确保比率收敛到我们的目标值，我们在批次中所有层的当前学习平均比率与作为超参数的目标比率之间引入了一个额外的MSE损失。
3.4. 支撑技术3：时间步级可微分比率
动机：除了层级比率差异外，我们还观察到模型在不同时间步上表现出不同程度的冗余。这促使我们也探索时间步级压缩比的方法。
时间步级可微分比率：在层级DiffCR的基础上，我们引入了特定于不同时间步区域的可学习参数。对于图像修复任务，遵循先前的最先进方法LazyDiffusion[28]，我们使用1000个训练时间步和100个采样时间步，并将它们均匀地划分为10个区域。我们相应地为每个层分配10个可学习比率，总共得到280个可学习参数。同样，对于文本到图像任务，遵循PixArt-Σ[4]的20个时间步，我们将它们划分为4个区域，为每个层分配4个可学习比率，得到112个可学习参数。与之前一样，我们在批次内的平均学习比率与目标比率之间应用MSE损失，以确保收敛。
4. 实验
4.1. 实验设置
任务、数据集和模型。任务与数据集：我们在两个具有代表性的图像生成任务上，使用相应的基准数据集对DiffCR进行评估：（1）基于包含2.2亿张高质量图像的内部数据集的图像修复任务，掩码和文本提示的生成遵循[28,53]；（2）基于LAION-5B数据集[38]的文本到图像（T2I）任务，该数据集仅限于具有高美学分数、英文文本且文本相似度分数不低于0.24的图像样本。模型：我们将提出的DiffCR与最先进的模型相结合。对于图像修复任务，我们使用Lazy Diffusion（一种经过适配的PixArt-α模型，带有额外的ViT编码器）生成1024×1024分辨率的图像。对于文本到图像任务，我们使用PixArt-Σ生成512×512分辨率的图像。
训练和采样设置。对于图像修复任务，我们使用AdamW优化器[20]对模型参数进行微调直至收敛，学习率为10⁻⁴，权重衰减为3×10⁻²。在采样时，使用IDDPM[27]生成图像，采用100个时间步和4.5的CFG因子。对于文本到图像任务，我们使用秩为32的LoRA适配器对模型进行微调，直至训练损失和验证损失均收敛，且DiffCR模型的当前压缩比与目标压缩比之间的MSE损失降至接近零。在训练过程中，我们使用IDDPM[27]在1000个时间步上计算扩散损失。在采样时，使用DPM-solver[21]生成图像，采用20个时间步和4.5的CFG因子。所有训练均在8×A100-80GB GPU集群上进行。
基线和评估指标。基线：在文本到图像和图像修复任务中，我们将提出的DiffCR与最先进的基线进行比较，包括ToMe[2]、AT-EDM[45]以及我们适配的具有统一MoD压缩比的模型。对于图像修复任务，我们还与RegenerateCrop和RegenerateImage进行了比较，其中RegenerateCrop围绕掩码区域生成一个紧凑的方形裁剪区域，与主流软件框架[44,50]类似，RegenerateImage生成整个图像，这在相关文献中较为常见[30,34,47,53]。评估指标：我们使用FID分数[11]评估生成图像的质量，使用CLIP分数[10]评估文本-图像对齐程度，并在A100 GPU上通过推理FLOPs、延迟和内存使用量评估效率。对于图像修复模型和文本到图像模型，我们分别在LAION-400M[37]和LAION-5B[38]的10K张图像上进行评估，均排除了训练样本。
4.2. DiffCR与最先进基线的比较
文本到图像。为了评估我们提出的DiffCR的有效性，我们将其应用于一般的文本到图像任务，并与先前的令牌合并[2]和剪枝[45]基线进行比较。具体而言，我们在PixArt-Σ上应用这些压缩方法，PixArt-Σ是一种最先进的公开可访问的文本到图像模型，以其高分辨率图像生成质量和效率权衡而闻名。如表1所示，配备DiffCR的PixArt-Σ显著提高了生成质量，与ToMe[2]和AT-EDM[45]相比，FID分别降低了57.83和241.11，同时延迟（降低8.59%~20.15%）和内存使用量（降低-2.71%~0.72%）相当甚至更低。此外，在与具有20%压缩比的ToMe[2]延迟相近的情况下，DiffCR实现了335.23的FID降低。此外，配备DiffCR的PixArt-Σ还实现了与未压缩的PixArt-Σ相当的图像生成质量，同时实现了20.68%和8.33%的延迟和内存节省。请注意，为了进行公平比较，我们与在LAION数据集上微调的PixArt-Σ进行了比较。这组实验证明了DiffCR在一般文本到图像任务中的有效性。
图像修复。我们进一步将DiffCR扩展到图像修复任务。具体而言，我们将其应用于最先进的Lazy Diffusion（LD）[28]之上，Lazy Diffusion使用DiT解码器仅生成掩码区域而非整个图像，并利用单独的ViT编码器捕获输入掩码图像的全局上下文。我们将我们的DiffCR方法与两类基线进行比较：（1）RegenerateImage和RegenerateCrop；（2）采用先前的令牌合并[2]或剪枝[45]技术的LD。如表2所示，我们的DiffCR在精度-效率权衡方面始终优于所有基线。例如，配备DiffCR的LD与配备ToME[2]或AT-EDM[45]的LD相比，FID分别降低了47.35和189.93，同时实现了相当或高达23.61%和13.63%的延迟和内存节省。此外，在与具有30%压缩比的ToMe[2]内存使用量相近的情况下，配备DiffCR的LD实现了265.94的FID降低，同时实现了高达21.54%的延迟节省。此外，与RegenerateImage相比，我们的方法在修复1024²图像中的256掩码大小时，实现了73.51%/60.26%的FLOPs和延迟节省。值得注意的是，与Lazy Diffusion一样，我们的方法的复杂度随掩码大小而变化，而RegenerateImage基于完整图像分辨率进行生成，这使其在较小的掩码大小下效率较低。与RegenerateCrop相比，我们的方法实现了显著更高的图像生成质量（+41.01 FID），同时实现了22.96%的内存节省。请注意，所有内存测量均在批处理大小为128的情况下进行。这组实验验证了我们的DiffCR在应用于图像修复任务时的有效性。
4.3. DiffCR的消融研究
我们对DiffCR进行了消融研究，分析了第3节中描述的三个支撑技术的贡献。如表格1和表格2所示，我们分别报告了在文本到图像和图像修复任务中，配备MoD路由器（第3.2节）、DiffCR-L（第3.3节）、DiffCR-LT（第3.4节）的LD或PixArt-Σ的性能。结果一致表明，我们的DiffCR的所有组件都对最终性能有贡献。具体而言，仅MoD就比ToMe[2]和AT-EDM[45]在文本到图像任务上平均降低了323.13和229.01的FID，同时延迟和内存使用量相当甚至更低。DiffCR-L和DiffCR-LT进一步提高了生成质量，在图像修复任务上额外降低了4.81/4.92的FID，在文本到图像任务上额外降低了10.5/12.1的FID。
我们的DiffCR的一个关键优势是，在微调过程中，所有层的平均压缩比逐渐收敛到目标比，从而产生一系列具有不同压缩比的“副产品”模型。如图7所示，我们可视化了图像修复和文本到图像任务的模型轨迹以及相应的FID分数和压缩比。观察结果表明，对于文本到图像任务，FID随着压缩比的增加而逐渐增加，达到了预期结果。相比之下，对于图像修复任务，FID逐渐降低。这种差异的产生是因为与文本到图像任务相比，图像修复任务和LD模型对剪枝更敏感，需要更长时间的微调来提高生成质量。这也反映在表格2和表格1中，图像修复任务在剪枝后FID增加，而文本到图像任务的FID甚至降低。
4.4. 定性视觉示例
视觉示例。我们选择具有挑战性的输入提示来评估我们提出的DiffCR的定性结果。如图6所示，这些示例表明，在图像修复和文本到图像任务中，DiffCR实现了与RegenerateCrop基线甚至未压缩的LD或PixArt-Σ相当甚至更优的生成质量。请注意，此处省略了ToMe和AT-EDM，因为即使在仅10%的压缩比下，它们在应用于DiTs时生成质量也很差。
人类偏好分数。我们使用计算机视觉模型来估计可能的人类偏好，并评估模型生成高质量、上下文相关图像的能力。具体而言，我们为文本到图像任务生成了2K个样本，并使用HPSv2[51]评估不同方法生成的图像的人类偏好。如表5所示，对于文本到图像任务，我们将所有压缩方法应用于PixArt-Σ[4]。与先前的压缩方法ToMe[2]和 vanilla MoD[33]相比，DiffCR实现了更高的人类偏好分数，分别为4.685/0.847。
5. 结论
在这项工作中，我们提出了DiffCR，这是一种具有可微分压缩比的动态DiT推理框架，它能跨令牌、层和时间步自适应地分配计算资源，从而得到高效的DiT模型。具体而言，DiffCR整合了一种基于深度混合（MoD）的令牌级路由机制，该机制能动态学习每个令牌的重要性分数，同时还包含一个新颖的模块，该模块使MoD相对于压缩比具有可微分性，从而使模型能够学习每个层和时间步的自适应压缩比。冗余层和时间步会学习到更高的压缩比，而关键层和时间步则保持较低的压缩率甚至完全不压缩。在图像修复和文本到图像（T2I）任务上的大量实验一致表明，与其他压缩方法相比，DiffCR在图像生成质量和效率之间取得了更优的权衡。
层级和时间步自适应可微分令牌压缩比用于高效扩散Transformer补充材料
A. 令牌路由器的更多可视化
在3.2节中，我们提供了路由器预测的示例可视化，以评估我们的DiffCR路由器的有效性。在这里，我们在图9中展示了更多的可视化示例，以进一步验证我们的发现。我们的观察结果一致表明：（1）路由器能有效捕捉语义信息，清晰地勾勒出物体形状，并在显著降低计算成本的同时达到类似注意力的效果；（2）预测的令牌重要性在不同层和时间步之间存在差异。例如，一些层专注于物体生成，而另一些层则强调背景区域。此外，随着时间步的推进，路由器越来越多地捕捉到物体的语义轮廓，这凸显了动态令牌重要性估计的重要性；（3）最佳压缩比在不同层和时间步之间存在差异。例如，一些层将高重要性分配给所有令牌，表明冗余度极小，而另一些层则有选择地修剪来自具有不同形状的物体或背景的令牌，这需要不同的压缩比。这种差异在不同时间步之间也有体现。在之前的MoD[33]方法中，固定的全局压缩率被均匀地应用于所有层和时间步，而忽略了它们各自的重要性。这种均匀修剪可能会过度修剪关键层或时间步，同时对冗余层或时间步的压缩不足。这一观察结果强调了需要为层和时间步量身定制自适应和动态的压缩比。
B. 文本到图像任务的比率轨迹分析
在3.3节中，我们可视化了使用我们提出的层级DiffCR训练的修复任务的比率轨迹。在这里，我们还提供了在文本到图像任务上对PixArt-Σ模型进行微调期间所有层的压缩比训练轨迹，如图8（a-c）所示。可视化结果一致表明：（1）每个层都学习到其独特的压缩比，冗余层实现了更高的压缩率，而关键层则保持较低的压缩率或完全不压缩；（2）各层的平均比率逐渐收敛到目标比率。在这个例子中，目标比率为20%，最终实现的平均比率约为19%，差距很小。值得注意的是，收敛速度和生成质量之间存在权衡：比率的MSE损失系数越高，收敛速度越快，但可能因压缩过快而降低质量；而系数越小，收敛越平缓，质量得以保持，但训练速度较慢。在实践中，我们将初始系数设置为0.3，并在训练过程中动态调整，以有效平衡速度和质量；（3）中间层的冗余度较大，而后层的冗余度通常较低，往往无法被压缩。早期层的冗余度水平则各不相同。
需要注意的是，为了防止模型学习到所有层的压缩比都为0%，我们使用一个系数来平衡扩散损失（倾向于较低的比率以获得更高的质量）和MSE损失（驱动目标平均比率），而无需额外的正则化或惩罚。较高的系数会加快收敛速度，但可能会影响质量；而较小的系数则能确保收敛平缓并保持质量。一些层自然会学习到0%的比率，这凸显了它们的重要性。
C. 学习到的压缩比与路由器预测之间的相关性
我们选择了三个具有高、中、低学习压缩比的代表性层，以可视化DiffCR路由器的相应预测，并分析可能存在的相关性。如图10所示（其中“C.R.”表示压缩比），我们观察到学习到的比率与路由器的预测之间存在很强的相关性。对于具有高压缩比的层，例如修复任务中的层1或文本到图像任务中的层9，路由器始终对许多语义区域预测较低的重要性分数，采用极其“懒惰的行为”来节省计算。相反，对于具有低压缩比的层，路由器为大多数区域分配较高的重要性分数。这种可视化验证了我们的令牌级路由器与可微分比率之间的联合学习效果。
D. 时间步区域选择的权衡
在3.4节中，我们介绍了时间步级DiffCR，其中对于具有100个采样时间步的修复任务，时间步区域被均匀地划分为10个区域；对于具有20个时间步的文本到图像任务，时间步区域被划分为4个区域。在这里，我们提供了关于选择时间步区域数量及相关权衡的额外指导。更多的时间步区域可以在所有时间步上学习更精细、更精确的压缩比。
然而，过多的区域会使训练不稳定且具有挑战性。为了降低训练复杂性并提高稳定性，我们选择较少的区域，例如文本到图像任务中选择4个区域。相反，使用过少的区域可能会使方法过于简化，使其沦为类似SpeeD[46]的启发式方法，后者手动定义了三个时间步区域。在实践中，我们选择4到10个时间步区域，以平衡粒度和稳定性。虽然我们的方法与SpeeD的一般见解一致，但它更具系统性和适应性。与手动探索大量设计空间不同，我们的方法以一种原则性的方式有效地处理数量多得多的区域，平衡了粒度和训练稳定性。
E. 整体比较图
在4.2节中，我们全面比较了我们的DiffCR方法与修复和文本到图像任务的基线方法。在这里，我们提供了整体比较图，以更好地说明在FID和延迟减少方面取得的改进。如图11所示，我们的DiffCR在FID和延迟之间始终实现了更优的权衡，与最具竞争力的基线相比，在相当的GPU延迟下，文本到图像和修复任务的FID分别降低了12.10和4.92。
F. DiffCR的模型轨迹
在4.2节中，我们可视化了DiffCR-L在文本到图像和修复任务训练期间的模型轨迹。这揭示了一个关键优势：在微调过程中，所有层的平均压缩比逐渐收敛到目标比，从而产生一系列具有不同压缩比的“副产品”模型。在这里，我们还提供了DiffCR-LT（“-LT”表示层级和时间步级DiffCR）的模型轨迹。如图13所示，我们可视化了DiffCR-LT微调期间的FID分数和相应的压缩比。观察结果一致验证了这种方法的优势，表明它能够生成一系列具有不同压缩比的模型。此外，我们观察到与文本到图像任务相比，修复任务和潜在扩散（LD）模型[28]对修剪更敏感，需要更长时间的微调才能有效提高生成质量。此外，对于文本到图像任务，与DiffCR-L相比，DiffCR-LT在模型轨迹方面表现出略高的稳定性。
G. 视觉示例的更多可视化
在4.4节中，我们选择了具有挑战性的输入提示来评估我们提出的DiffCR的定性性能。
在这里，我们提供了额外的视觉示例，如图12所示。这些示例一致表明，在修复和文本到图像任务中，DiffCR实现了与RegenerateCrop基线甚至未压缩的LD或PixArt-Σ相当甚至更优的生成质量。请注意，此处省略了ToMe[2]和AT-EDM[45]，因为即使在10%的适度压缩比下，它们在应用于DiTs时生成质量也很差。
H. 与基于缓存的基线的比较
我们在表4中总结了我们的方法和基于缓存的基线的特点。DeepCache[54]和CMYC[25]是为基于U-Net的模型设计的，因此难以直接比较，而L2C[22]和TGATE[19]通过缓存层特征来减少未来时间步的重新计算，从而针对DiTs。与这些方法不同，我们的方法专注于具有可学习的层级和时间步相关压缩比的令牌修剪，虽然它不采用时间缓存，但与这些技术兼容。为了进行直接比较，我们在MS-COCO-30K数据集（文本到图像任务）上使用PixArt-Σ评估了所有方法，在约25%的延迟节省下，L2C的FID为28.6（我们训练的路由器重现了与报告相似的缓存模式），TGATE的FID为43.6，而我们的DiffCR的FID为28.6。这些结果表明，我们的方法与基于缓存的基线相当或更好，并且可以进一步与它们结合使用，以实现额外15-30%的延迟减少。
I. 修复任务的人类偏好分数
在4.4节中，我们使用计算机视觉模型来估计可能的人类偏好，并评估模型生成高质量、上下文相关的文本到图像任务图像的能力。在这里，我们还提供了修复任务的评估。具体而言，我们为修复任务生成了2K个样本，并使用HPSv2[51]评估不同方法生成的图像的人类偏好。如表5所示，对于修复任务，我们将所有压缩方法应用于Lazy Diffusion（LD）[28]。与先前的压缩方法ToMe[2]和vanilla MoD[33]相比，DiffCR实现了更高的人类偏好分数，分别为2.181/0.263。
J. 压缩比的消融分析
在这项工作中，我们以降低延迟为目标，作为向边缘部署迈出的一步。为了分析不同压缩比的影响，我们在MS-COCO-30K数据集[18]上使用PixArt-Σ模型进行了消融研究。值得注意的是，1/3的时间步被分配给全模型推理以保持准确性。下表中的结果表明，我们的方法能有效地扩展到更大的压缩比，FID仅略有增加（<1）。先前选择30%的压缩比用于具有挑战性的生成任务，以在现有最先进的高效方法的基础上保持准确性。
K. 仅MSE损失是否足够？
我们发现，仅使用MSE损失就能有效地引导比率向目标值靠近，无需额外的正则化，因此我们将其固定为MSE损失，但其他损失函数也可能同样有效。此外，尽管我们没有强制进行二元预测，但在某些层中，路由器倾向于学习极化分布，将重要令牌与不重要令牌区分开来，学习到的比率也相应地对齐，如图10所示。