CoReDiT:基于空间一致性引导的令牌剪枝与重建用于高效的扩散变换器
CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers
AI总结:
CoReDiT通过空间一致性引导的令牌剪枝与重建,显著降低扩散变换器的计算量,提升推理速度,同时保持高质量视觉输出。
摘要:
扩散变换器(DiTs)在图像和视频生成方面表现出色,但计算成本高,限制了可扩展性和设备部署。我们引入了CoReDiT,一种适用于视觉任务的结构化令牌剪枝框架。CoReDiT利用线性时间的空间一致性分数来估计潜在令牌晶格中的局部冗余,并在自注意力中跳过高一致性(冗余)令牌。为了保持密集表示并避免视觉断层,我们通过一致性引导的空间邻近保留令牌的聚合来重建跳过的注意力输出。我们进一步引入了一种渐进、块自适应的剪枝计划,逐步增加剪枝程度,并将更大的预算分配给冗余更高的块和去噪步骤。在最先进的扩散骨干网络包括PixArt-α和MagicDrive-V2上,CoReDiT实现了高达55%的自注意力FLOPs减少,并在云GPU上获得1.33倍的推理加速,在移动NPU上获得1.72倍的加速,同时保持高质量的视觉输出。值得注意的是,CoReDiT还增加了设备内存的可用空间,使更高分辨率的生成成为可能。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。