通过空间一致性剪枝与重建，让扩散Transformer更高效生成图像视频。
CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers
- 用空间一致性评分识别冗余令牌，跳过低效注意力计算。
- 在云GPU上提速1.33倍，移动端达1.72倍，最高省55%算力。
- 适合追求高画质又需轻量化部署的视觉生成任务。
扩散Transformer（DiTs）在图像和视频生成中表现卓越，但计算开销大，限制了可扩展性和设备端部署。本文提出CoReDiT，一种面向多视觉任务的结构化令牌剪枝框架。CoReDiT利用线性时间的空间一致性得分评估潜在令牌网格中的局部冗余，跳过高一致性（冗余）令牌的自注意力计算。为保持表示密度并避免视觉断裂，通过相干性引导的邻域聚合重建被跳过的注意力输出。进一步设计渐进式、块自适应剪枝策略，逐步增加剪枝比例，并将更大剪枝预算分配给冗余度更高的模块与去噪步骤。在PixArt-α和MagicDrive-V2等先进扩散模型上，CoReDiT实现最高55%的自注意力浮点运算量降低，在云GPU上推理加速1.33倍，移动端NPU加速1.72倍，同时保持高视觉质量。值得注意的是，该方法还提升设备端内存空间，支持更高分辨率生成。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) deliver remarkable image and video generation quality but incur high computational cost, limiting scalability and on-device deployment. We introduce CoReDiT, a structured token pruning framework for DiTs across vision tasks. CoReDiT uses a linear-time spatial coherence score to estimate local redundancy in the latent token lattice and skips high coherence (redundant) tokens in self-attention. To maintain a dense representation and avoid visual discontinuities, we reconstruct skipped attention outputs via coherence-guided aggregation of spatially neighboring retained tokens. We further introduce a progressive, block-adaptive pruning schedule that increases pruning gradually and allocates larger budgets to blocks and denoising steps with higher redundancy. Across state-of-the-art diffusion backbones including PixArt-α and MagicDrive-V2, CoReDiT achieves up to 55% self-attention FLOPs reduction and inference speedups of 1.33x on cloud GPUs and 1.72x on mobile NPUs, while maintaining high visual quality. Notably, CoReDiT also increases on-device memory head-room, enabling higher-resolution generation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。