arXiv:2412.06028cs.CV2024-12NeurIPS被引 9

通过稀疏化令牌提升扩散Transformer效率,生成质量不降反升。

SparseDiT: Token Sparsification for Efficient Diffusion Transformer

  • 分层稀疏设计:底层池化、中层混合、顶层密集,动态适配特征需求。
  • 节省55%计算量,图像生成速度提升175%,视频生成效率提升56%。
  • 适合追求高效高质生成的科研与工程人员,兼容现有优化技术。

扩散Transformer(DiT)虽生成效果出色,但受限于自注意力的二次复杂度和大量采样步骤,计算开销巨大。现有研究多聚焦加速采样过程,却忽视了架构本身的低效问题。本文提出SparseDiT,通过在空间和时间维度实施令牌稀疏化,显著提升效率同时保持生成质量。空间上采用三段式结构:底层使用Poolingformer高效提取全局特征,中层引入稀疏-密集令牌模块(SDTM)平衡全局上下文与局部细节,顶层保留密集令牌以精炼高频信息。时间上动态调节各去噪阶段的令牌密度,随时间步推进逐步增加令牌数量,以捕捉精细结构。该空间自适应架构与时间剪枝策略协同作用,实现生成全过程的效率与保真度统一。实验表明,SparseDiT在DiT-XL上实现55%的FLOPs减少和175%的推理加速,图像生成FID分数相当;在视频生成数据集上降低56%的计算量;在PixArt-α上实现69%的推理加速,且文本到图像任务的FID下降0.24。SparseDiT为高质量扩散生成提供可扩展的高效解决方案,兼容现有采样优化方法。

原文摘要 · Abstract (English)

Diffusion Transformers (DiT) are renowned for their impressive generative performance; however, they are significantly constrained by considerable computational costs due to the quadratic complexity in self-attention and the extensive sampling steps required. While advancements have been made in expediting the sampling process, the underlying architectural inefficiencies within DiT remain underexplored. We introduce SparseDiT, a novel framework that implements token sparsification across spatial and temporal dimensions to enhance computational efficiency while preserving generative quality. Spatially, SparseDiT employs a tri-segment architecture that allocates token density based on feature requirements at each layer: Poolingformer in the bottom layers for efficient global feature extraction, Sparse-Dense Token Modules (SDTM) in the middle layers to balance global context with local detail, and dense tokens in the top layers to refine high-frequency details. Temporally, SparseDiT dynamically modulates token density across denoising stages, progressively increasing token count as finer details emerge in later timesteps. This synergy between SparseDiT spatially adaptive architecture and its temporal pruning strategy enables a unified framework that balances efficiency and fidelity throughout the generation process. Our experiments demonstrate SparseDiT effectiveness, achieving a 55% reduction in FLOPs and a 175% improvement in inference speed on DiT-XL with similar FID score on 512x512 ImageNet, a 56% reduction in FLOPs across video generation datasets, and a 69% improvement in inference speed on PixArt-$α$ on text-to-image generation task with a 0.24 FID score decrease. SparseDiT provides a scalable solution for high-quality diffusion-based generation compatible with sampling optimization techniques.

扩散模型效率优化稀疏化图像生成

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