论文

CoReDiT:空间相干引导的词元修剪和重建以实现高效扩散 Transformer

CoReDiT: Spatial Coherence-Guided Token Pruning and Reconstruction for Efficient Diffusion Transformers

模型推理推理加速

摘要

Diffusion Transformer (DiTs) 可提供卓越的图像和视频生成质量,但会产生较高的计算成本,限制可扩展性和设备上部署。我们引入了 CoReDiT,这是一种跨视觉任务的 DiT 的结构化词元修剪框架。 CoReDiT 使用线性时间空间相干性分数来估计潜在标记格中的局部冗余,并跳过自注意力中的高相干性(冗余)标记。为了保持密集的表示并避免视觉不连续性,我们通过空间相邻保留标记的连贯性引导聚合来重建跳过的注意力输出。我们进一步引入了一种渐进的、块自适应的剪枝方案,该方案逐渐增加剪枝,并将更大的预算分配给具有更高冗余的块和去噪步骤。在包括 PixArt-α 和 MagicDrive-V2 在内的最先进的扩散主干网络中,CoReDiT 实现了高达 55% 的自注意力 FLOP 减少,云 GPU 上的推理速度提高了 1.33 倍,移动 NPU 上的推理速度提高了 1.72 倍,同时保持了高视觉质量。值得注意的是,CoReDiT 还增加了设备上的内存空间,从而实现更高分辨率的生成。