arXiv:2604.03674cs.CV2026-04中稿 · ICLR被引 2

通过可学习的稀疏性优化,显著加速扩散Transformer生成速度。

DiffSparse: Accelerating Diffusion Transformers with Learned Token Sparsity

  • 设计可微分的逐层稀疏优化框架,动态分配计算资源。
  • 在20步采样下降低54%计算量,且生成质量更优。
  • 适合追求高效生成的图像模型部署与推理优化场景。

扩散模型在图像生成中表现优异,但其多步推理机制带来巨大计算开销。现有方法通过层或令牌缓存技术降低计算成本,但在少步扩散Transformer模型中因特征缓存效率低、手动稀疏分配及保留完整前向计算而效果受限。为此,本文提出一种可微分的逐层稀疏优化框架,利用令牌缓存减少令牌计算开销并提升加速效果。通过可学习网络结合动态规划求解器,端到端优化逐层稀疏分配策略。此外,提出的两阶段训练策略无需全步处理,进一步提升效率。在DiT-XL/2、PixArt-α、FLUX和Wan2.1等扩散变压器模型上广泛实验表明,本方法在不降低生成质量的前提下持续提升效率。例如,在PixArt-α使用20个采样步骤时,计算成本降低54%,生成指标超越原模型,显著优于现有方法。结果证明该方法在大幅提升效率的同时,常能改善生成质量。

原文摘要 · Abstract (English)

Diffusion models demonstrate outstanding performance in image generation, but their multi-step inference mechanism requires immense computational cost. Previous works accelerate inference by leveraging layer or token cache techniques to reduce computational cost. However, these methods fail to achieve superior acceleration performance in few-step diffusion transformer models due to inefficient feature caching strategies, manually designed sparsity allocation, and the practice of retaining complete forward computations in several steps in these token cache methods. To tackle these challenges, we propose a differentiable layer-wise sparsity optimization framework for diffusion transformer models, leveraging token caching to reduce token computation costs and enhance acceleration. Our method optimizes layer-wise sparsity allocation in an end-to-end manner through a learnable network combined with a dynamic programming solver. Additionally, our proposed two-stage training strategy eliminates the need for full-step processing in existing methods, further improving efficiency. We conducted extensive experiments on a range of diffusion-transformer models, including DiT-XL/2, PixArt-$α$, FLUX, and Wan2.1. Across these architectures, our method consistently improves efficiency without degrading sample quality. For example, on PixArt-$α$ with 20 sampling steps, we reduce computational cost by $54\%$ while achieving generation metrics that surpass those of the original model, substantially outperforming prior approaches. These results demonstrate that our method delivers large efficiency gains while often improving generation quality.

扩散模型稀疏性优化加速推理Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。