💻 computer science

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

DiffSparse 提出了一种基于可微分层稀疏度优化的端到端框架，通过结合令牌缓存与两阶段训练策略，在无需完整前向计算的情况下显著降低了扩散 Transformer 模型的推理成本并提升了生成质量。

原作者： Haowei Zhu, Ji Liu, Ziqiong Liu, Dong Li, Junhai Yong, Bin Wang, Emad Barsoum

发布于 2026-04-07

📖 1 分钟阅读☕ 轻松阅读

原作者： Haowei Zhu, Ji Liu, Ziqiong Liu, Dong Li, Junhai Yong, Bin Wang, Emad Barsoum

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一种名为 DiffSparse 的新方法，旨在让现在的 AI 绘画和生成视频模型（比如 Midjourney、Stable Diffusion 的升级版）画得更快，而且画得更好。

为了让你轻松理解，我们可以把 AI 生成图像的过程想象成一位画家在画一幅复杂的油画。

1. 核心问题：画家太累了，画得太慢

现在的 AI 模型（特别是基于 Transformer 架构的 Diffusion 模型）画一幅画，不是一笔就成的。它需要像“去噪”一样，经过很多个步骤（比如 20 步、50 步），一步步把模糊的噪点变成清晰的图像。

现状：每一步，画家都要重新审视画布上的每一个像素点（Token），计算它们该怎么变。这就像画家在画 1000 个像素时，哪怕其中 800 个像素跟上一秒几乎没变，他也要重新算一遍。这非常浪费时间和精力（计算成本极高）。
以前的尝试：以前的加速方法有点像“偷懒”。比如，画家决定：“前 5 步我全画，中间 5 步我直接照搬上一张图，最后 5 步我再全画。”
- 缺点：这种“偷懒”太死板了。有时候该偷懒的时候没偷，有时候该画的时候却偷懒了，导致画出来的东西要么模糊，要么细节丢失。而且，这种“照搬”通常只能用在特定的几步，不能灵活调整。

2. DiffSparse 的解决方案：聪明的“智能助手”

DiffSparse 就像给这位画家配了一个超级聪明的智能助手。这个助手不直接画画，但它负责决定哪些地方需要画家亲自动手，哪些地方可以直接“复制粘贴”上一秒的结果。

核心功能一：动态的“偷懒”策略（可学习的稀疏性）

以前的做法：助手是死板的，它说：“不管画什么，第 3 层和第 5 层必须全画，第 4 层可以偷懒。”这是人工设定的，不够灵活。
DiffSparse 的做法：助手是会学习的。它通过观察，发现：
- 在画画的早期，大部分细节还没出来，所以很多像素点变化不大，可以大胆地“复制粘贴”（节省算力）。
- 在画画的后期，或者在画关键细节（比如眼睛、文字）时，助手会告诉画家：“这里必须亲自动手，不能偷懒！”
- 比喻：这就像你写文章。写大纲时，你可以快速跳过很多词；但写结尾的高潮部分时，你必须字斟句酌。DiffSparse 能自动识别哪里该“快进”，哪里该“精修”。

核心功能二：动态规划（全局最优解）

助手不仅看局部，还看全局。它使用一种叫“动态规划”的数学方法，就像下棋一样，计算每一步怎么分配精力，才能让总耗时最少，同时画出来的画最完美。
它不是随机决定哪里偷懒，而是经过精密计算，确保在节省 50% 工作量的情况下，画出来的画甚至比原来全画还要好。

核心功能三：两阶段训练法（先学后精）

第一阶段：助手先观察画家在“全速运转”（不偷懒）时的表现，学习哪些步骤是必须全画的（比如起稿和收尾），哪些步骤可以完全跳过。
第二阶段：助手把学到的经验整合起来，开始指导画家在“半速”状态下工作。它不再依赖死板的规则，而是根据画作的实时状态，动态调整哪里该快、哪里该慢。
比喻：就像学开车。先跟着教练开（全画），教练告诉你哪里该踩刹车、哪里该加速；然后你自己开（两阶段训练），你学会了根据路况自动调整，既省油又安全。

3. 效果如何？

论文在多种模型（如 PixArt-α, FLUX, Wan2.1）上做了测试，效果非常惊人：

速度提升：在 PixArt-α 模型上，他们把计算量减少了 54%（相当于画一幅画的时间缩短了一半多），速度提升了 1.91 倍。
质量提升：最不可思议的是，画得更快，画得更好了！ 以前加速往往会导致画质下降（比如模糊、变形），但 DiffSparse 因为更聪明地分配了精力，反而让生成的图片质量比原来的全速模型还要高（FID 分数更低，代表更逼真）。
通用性：这个方法不仅适用于画 256x256 的小图，甚至直接套用到 512x512 的大图上，不需要重新训练，依然有效。

总结

DiffSparse 就像是给 AI 画家装上了一个拥有“上帝视角”的导航系统。

以前：画家不管路多堵，都按固定路线开，或者盲目地抄近道，容易迷路或撞车（画质差）。
现在：导航系统（DiffSparse）实时分析路况，告诉画家：“前面这段路车少，我们可以开快点（复用特征）；前面那个路口很复杂，必须慢慢开（重新计算）。”

结果是：路走得更顺了（速度更快），而且准时甚至提前到达了目的地（画质更好）。 这让 AI 绘画和生成视频变得更加高效，未来在手机上运行大型 AI 模型也将成为可能。

这篇论文提出了一种名为 DiffSparse 的新框架，旨在通过可学习的层间 Token 稀疏性优化来加速扩散 Transformer（Diffusion Transformers, DiT）模型的推理过程。该方法在保持甚至提升生成质量的同时，显著降低了计算成本。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

背景：扩散模型（特别是基于 Transformer 的 DiT 架构，如 PixArt-α, FLUX, Wan2.1）在图像和视频生成方面表现卓越，但其多步去噪推理机制带来了巨大的计算开销，限制了实际部署。
现有方法的局限性：
- 特征缓存（Feature Caching）：现有的缓存方法（如 Token Cache）利用相邻时间步特征的冗余性来跳过计算。然而，这些方法通常存在以下问题：
  1. 人工设计：需要手动设定稀疏率分配策略和调度表，缺乏自适应能力。
  2. 全步计算依赖：为了维持生成质量，现有方法通常强制在去噪过程的某些步骤（如前几步或每隔几步）进行完整的正向传播（Full Forward Pass），这限制了加速潜力。
  3. 效率瓶颈：由于保留了部分完整计算步骤，无法在 Token 级别实现最大化的加速。

2. 核心方法论 (Methodology)

DiffSparse 将 Token 缓存优化建模为一个基于动态规划的层间稀疏性分配问题，并通过端到端的可学习框架解决。

2.1 核心组件

可学习的稀疏性成本预测器 (Learnable Sparsity Cost Predictor)：
- 这是一个轻量级的可学习网络，用于预测在不同去噪时间步（ $t$ ）和不同网络层（ $l$ ）应用不同稀疏率（ $s \in S$ ）时的“成本”。
- 成本反映了在该配置下对生成质量（如感知相似度）的潜在损害。
- 该预测器的规模仅取决于时间步数 $T$ 、层数 $L$ 和候选稀疏率集合大小 $|S|$ ，与 Token 序列长度无关，因此具有良好的可扩展性。
动态规划求解器 (Dynamic Programming Solver)：
- 在给定全局稀疏率约束（即总计算量减少比例 $R$ ）下，利用动态规划算法在预测的成本矩阵中寻找最优的稀疏性配置方案。
- 该算法能够自动决定在每一层、每一个时间步保留多少 Token 进行计算，从而在满足总稀疏率的前提下最小化累积成本。
Token 选择器 (Token Selector)：
- 根据学习到的稀疏率，动态选择具体的 Token 进行重用（从缓存读取）或重新计算。
- 选择依据包括注意力分数、交叉注意力焦点、缓存重用频率以及空间分布均匀性（通过复合评分函数 $S(\hat{x}_i)$ 实现）。

2.2 两阶段训练策略 (Two-Stage Training Strategy)

为了解决现有方法依赖“全步计算”的问题，DiffSparse 提出了一种创新的两阶段训练策略：

第一阶段：独立优化“全步位置”的成本矩阵和“层稀疏性”成本矩阵。利用动态规划确定哪些时间步必须保留为全步计算（Full Steps），并以此作为先验。
第二阶段：将第一阶段的先验信息（通过常数 $\delta$ 调整成本）整合到层稀疏性优化中，进行微调。
优势：这种策略允许模型逐渐减少对全步计算的依赖，动态地在所有步骤中重新分配计算资源，从而完全释放 Token 级缓存的加速潜力。

2.3 训练目标

使用 LPIPS (Learned Perceptual Image Patch Similarity) 作为感知蒸馏损失。
将未剪枝的原始模型作为教师网络，剪枝后的模型作为学生网络。
通过最小化 LPIPS 损失，确保在大幅减少计算量的同时，生成的图像在感知质量上与原始模型保持一致。
使用 Straight-Through Estimator (STE) 处理离散掩码带来的不可导问题，实现端到端优化。

3. 主要贡献 (Key Contributions)

DiffSparse 框架：提出了一种可微分的、端到端的层间 Token 稀疏性优化方法，无需人工启发式规则即可自动分配稀疏性。
两阶段训练策略：消除了现有方法对预定义全步计算的依赖，通过动态优化稀疏模式，进一步提升了加速效率。
动态规划求解：将稀疏性分配问题形式化为动态规划问题，确保了在满足全局约束下的最优解。
广泛的实验验证：在多个主流 DiT 模型（PixArt-α, FLUX, Wan2.1, DiT-XL/2）和任务（文生图、文生视频）上进行了验证。

4. 实验结果 (Results)

实验在多个基准数据集（COCO, ImageNet, PartiPrompts, VBench）上进行，结果显示 DiffSparse 在加速比和生成质量上均优于现有 SOTA 方法：

PixArt-α (文生图)：
- 在 20 步采样下，减少 54% 的计算量（加速 1.91 倍），FID 达到 27.79。
- 相比之下，ToCa (1.75×) 的 FID 为 28.35，TaylorSeer (1.83×) 的 FID 为 29.08。
- 甚至在 1.74 倍加速下，FID (26.91) 优于原始全量模型 (28.20)，表明稀疏性分配实际上改善了收敛性。
DiT-XL/2 (类条件生成)：
- 在 2.07 倍加速下，FID 从 ToCa 的 3.05 提升至 2.81。
Wan2.1 (文生视频)：
- 在 2.05 倍加速下，VBench 得分达到 43.83，优于 DuCa 等方法的 43.56，且计算成本更低。
泛化性：在 256x256 分辨率训练的稀疏性预测器，直接应用于 512x512 分辨率时，仍能保持优异性能，无需重新训练。

5. 意义与影响 (Significance)

效率与质量的平衡：DiffSparse 证明了通过智能的 Token 级稀疏性分配，可以在大幅降低计算成本（FLOPs）的同时，不仅不降低甚至能提升生成质量。
自动化与适应性：摆脱了人工设计稀疏调度表的繁琐过程，使得加速策略能够自适应不同模型架构和任务需求。
工业部署潜力：显著降低了扩散模型的推理门槛，使其更易于在资源受限的设备（如移动端、边缘计算）上部署，同时兼容现有的 DiT 架构。
方法论创新：将动态规划引入扩散模型的稀疏性优化，为后续研究提供了一种新的优化范式，即通过可学习的成本预测结合全局优化求解器来替代启发式规则。

综上所述，DiffSparse 通过结合可学习的成本预测、动态规划求解和两阶段训练策略，成功解决了扩散 Transformer 推理加速中的关键瓶颈，为高效生成式 AI 的部署提供了强有力的技术支持。