DiffSparse：利用学习词元稀疏性加速扩散 Transformer

扩散模型在图像生成方面表现出出色的性能，但其多步推理机制需要巨大的计算成本。之前的工作通过利用层或词元缓存技术来降低计算成本来加速推理。然而，由于这些词元缓存方法中低效的特征缓存策略、手动设计的稀疏分配以及在几个步骤中保留完整的前向计算的做法，这些方法无法在少步扩散Transformer模型中实现优异的加速性能。为了应对这些挑战，我们提出了一种用于扩散 Transformer 模型的可微分层稀疏优化框架，利用词元缓存来降低词元计算成本并增强加速。我们的方法通过可学习网络与动态规划求解器相结合，以端到端的方式优化分层稀疏分配。此外，我们提出的两阶段训练策略消除了现有方法中全步骤处理的需要，进一步提高了效率。我们对一系列扩散-Transformer 模型进行了广泛的实验，包括 DiT-XL/2、PixArt-$α$、FLUX 和 Wan2.1。在这些架构中，我们的方法在不降低样本质量的情况下持续提高效率。例如，在具有 20 个采样步骤的 PixArt-$α$ 上，我们将计算成本降低了 54\%$，同时实现了超越原始模型的生成指标，大大优于先前的方法。这些结果表明，我们的方法可带来巨大的效率增益，同时经常提高生成质量。