💻 computer science

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

CoReDiT 是一种针对扩散变换器的高效令牌剪枝框架，它利用空间相干性来识别并跳过冗余令牌，同时通过邻近令牌重建其输出，从而在视觉质量不受影响的情况下实现显著的计算加速和内存节省。

原作者： Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

发布于 2026-05-15

📖 1 分钟阅读☕ 轻松阅读

原作者： Zhuojin Li, Hsin-Pai Cheng, Hong Cai, Shizhong Han, Fatih Porikli

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图在一面墙上绘制一幅巨大且细节极其丰富的壁画。你有一支艺术家团队（即“令牌”）站在墙前，每一位艺术家都需要与其他每一位艺术家交流，以决定在他们各自的特定位置上涂什么颜色。这确保了画面完美无缺，但过程令人筋疲力尽。如果你有 10,000 位艺术家，他们仅为了商定下一笔笔触就需要进行 1 亿次对话。这就是当前 AI 图像生成器（称为扩散变换器）的工作原理：它们极其出色，但速度慢得惊人，且对计算能力的需求巨大。

本文介绍了一种名为CoReDiT的新方法，旨在加快这一过程，同时不破坏画作。其工作原理可分解为三个简单步骤：

1. “无聊艺术家”检测器（空间一致性）

在一幅典型的画作中，有些区域非常繁忙且细节丰富（如人脸或树木），而另一些区域则只是平淡无奇的背景（如蓝天或平滑的墙壁）。

问题所在： 旧的 AI 对所有艺术家一视同仁，迫使那些绘制无聊蓝天的艺术家与绘制人脸的艺术家进行同样激烈的交流。
CoReDiT 的解决方案： CoReDiT 观察这些艺术家，并问道：“谁站在一个看起来与自己完全相同的人旁边？”如果一位艺术家处于无聊且均匀的区域，那么他的邻居几乎就是克隆体。CoReDiT 会给这些“无聊”的艺术家赋予较高的一致性得分。
采取的行动： 它告诉得分高（即冗余）的艺术家：“你现在不需要和每个人交谈。休息一下。”这跳过了图像中无聊部分的昂贵对话，从而节省了大量时间。

2. “填空”艺术家（重建）

如果你只是让那些“无聊”的艺术家回家并停止工作，你的壁画上就会出现空洞。画面将显得支离破碎、斑驳不全。

问题所在： 简单地移除令牌（艺术家）会造成视觉故障。
CoReDiT 的解决方案： CoReDiT 不会留下空洞，而是利用那些留下的艺术家来填补空白。由于“无聊”的艺术家与他们的邻居如此相似，剩下的艺术家可以轻易地猜出缺失的艺术家原本会画什么。
采取的行动： 这就像一场“传话”游戏，剩下的艺术家将正确的颜色悄悄告诉空位。这确保了最终图像看起来平滑且完整，没有任何可见的缝隙，尽管实际上只有较少的艺术家在承担繁重的工作。

3. 智能管理者（渐进式剪枝）

绘画过程的不同阶段并非完全相同。

问题所在： 在图像生成的最初阶段，画面只是静态噪声（如同电视雪花），因此几乎一切看起来都相同且冗余。但在最后阶段，当细节被锐化时，每一笔笔触都至关重要。如果你过早地削减太多艺术家，AI 就会感到困惑；如果你在最后阶段削减太多，画面就会变得模糊。
CoReDiT 的解决方案： CoReDiT 扮演一位智能管理者的角色，它学会了完美的调度安排。当图像仅仅是噪声时，它首先允许 AI 剔除大量冗余的艺术家。随着图像变得越来越清晰、细节越来越丰富，管理者会慢慢减少剔除艺术家的数量，确保关键细节得到充分关注。
采取的行动： 它自动确定 AI 大脑的哪些部分（模块）以及过程的哪些阶段（时间步）可以承受最大程度的削减，并实时调整计划以保持高质量。

结果

通过使用这种方法，作者发现：

速度： 在强大的云端计算机上，AI 的运行速度提高了1.33 倍；在手机芯片上，速度提高了1.72 倍。
效率： 它将繁重的数学工作（计算量）减少了高达55%。
质量： 生成的图像看起来与慢速版本一样好，没有明显的裂痕或奇怪的伪影。
额外优势： 由于它使用的内存更少，甚至可以在那些通常因内存不足而崩溃的手机上生成更高分辨率的图像。

简而言之，CoReDiT 就像为艺术项目聘请了一位聪明的工头，他确切地知道何时让团队休息，何时让他们努力工作，从而以更快速度完成杰作。

技术摘要：CoReDiT

问题陈述

扩散 Transformer（DiTs）通过注意力机制有效建模长程依赖，已在图像和视频生成领域确立了最先进的性能。然而，其计算成本随令牌（token）数量呈二次方增长，使得高分辨率合成和多帧视频生成对资源受限的移动平台而言代价过高。尽管潜在空间中的大部分令牌对应于视觉上冗余的区域（例如均匀背景），但现有的令牌剪枝方法在应用于扩散推理时面临三大主要挑战：

高效定位：在不产生计算全局注意力分数所带来的二次方开销的情况下，识别低显著性令牌。
语义保持：在跳过令牌时保持视觉连续性并防止伪影（例如纹理断裂、边界畸变），这对于生成任务（不同于判别任务）至关重要。
自适应调度：确定一种剪枝策略，能够适应不同 Transformer 块和去噪时间步中变化的冗余程度。

方法论：CoReDiT

CoReDiT 是一个结构化的令牌剪枝框架，旨在加速预训练的 DiT，同时保持输出保真度。它通过三个核心组件运行：

1. 基于空间相干的令牌选择

CoReDiT 不依赖全局注意力分数（其复杂度为 $O(N^2)$ ），而是利用**空间相干性（SC）**分数以线性时间（ $O(N)$ ）估计局部冗余。

机制：将令牌晶格划分为不重叠的网格。对于每个令牌，SC 分数计算为该令牌的归一化嵌入与其所在网格的平均归一化嵌入之间的平均余弦相似度。
逻辑：具有高 SC 分数的令牌被视为冗余，因为它们能被其空间邻居很好地解释。
选择：选择 SC 分数最高的前 $K$ 个令牌进行跳过（ $X_S$ ），而剩余的令牌（ $X_R$ ）继续进入多头自注意力（MHSA）层。

2. 相干性引导的令牌重建

为了防止因跳过令牌而造成的空间不连续性，CoReDiT 利用保留邻居的输出（ $Y_R$ ）重建被跳过令牌的注意力输出（ $Y_S$ ）。

近似：该方法不使用成对相似度计算进行重建（这会带来额外开销），而是利用邻居预计算的 SC 分数来近似相似度权重。
局部子网格：为了避免过度平滑，重建在更小的子网格内进行。
微设计：
- 交替网格尺寸：网格维度在连续块之间交替，以转移边界伪影并允许信息跨越网格边界传播。
- m-步长保留：一种确定性的步长模式确保每个子网格中至少保留一个令牌，为重建提供必要的锚点。

3. 渐进式、块自适应剪枝调度

鉴于冗余度在不同 Transformer 块和去噪步骤中各不相同，CoReDiT 在轻量级微调期间采用渐进式剪枝策略。

块自适应：将剪枝预算分配给测量冗余度最高的块（通过候选令牌的 SC 分数之和估算）。
时间步衰减：在早期去噪步骤（噪声占主导且结构粗糙）中剪枝更为激进，而在后期步骤（精细细节被优化）中则减少剪枝。采用两阶段衰减调度，以在效率与管理特定步骤计算图的复杂性之间取得平衡。
训练：模型使用蒸馏损失进行微调，该损失结合了原始模型输出损失和剪枝块损失，以恢复性能。

主要贡献

线性时间冗余估计：引入空间相干性分数，以 $O(N)$ 的时间识别冗余令牌，避免了全局注意力分析的计算瓶颈。
致密晶格保持：一种重建机制，从邻居合成被跳过的注意力输出，确保令牌晶格对后续层保持致密，并防止视觉伪影。
自适应剪枝策略：一种渐进式、块自适应的调度，根据深度和时间上的经验冗余统计动态分配剪枝预算，促进稳定的适应。

实验结果

CoReDiT 在最先进的 DiT 骨干网络上进行了评估，包括 PixArt-α、PixArt-Σ 和 MagicDrive-V2。

效率提升：
- 实现了高达 55% 的自注意力 FLOPs 减少。
- 推理加速：在云 GPU（Nvidia H100）上加速 1.33 倍，在移动 NPU（Qualcomm Snapdragon 8 Elite）上加速 1.72 倍。
- 内存空间：在移动设备上实现了高达 1920 分辨率的生成，而基线模型在 1600 分辨率时即遇到内存溢出（OOM）错误。
质量保持：
- 在 PixArt-α-1024 上，40% 的剪枝率下，CoReDiT 保持了 28.7 的 FID（基线为 27.3）以及具有竞争力的 CLIP/IS 分数。经过蒸馏后，FID 差距缩小至 27.4。
- 在 PixArt-Σ-2048 上，23% 的平均剪枝率将自注意力延迟降低了 22%，而质量指标仅略有下降。
- 在 MagicDrive-V2（视频）上，自注意力 FLOPs 减少 39%，保持了逐帧质量（PSNR、LPIPS、SSIM）和条件对齐（mAP、mIoU），尽管观察到时间一致性（FVD）略有下降。
消融研究：
- 随机令牌选择导致灾难性的质量下降（FID 644.5），验证了基于相干性选择的必要性。
- 禁用重建导致显著的质量损失（FID 30.4 对比 28.5），证实了重建步骤的重要性。
- 均匀剪枝后接微调的表现不如所提出的渐进式自适应调度。

意义与主张

该论文声称，CoReDiT 通过解决生成性令牌剪枝的具体挑战，为在边缘设备上部署高保真扩散模型提供了一种实用解决方案。其意义在于：

可扩展性：通过减少计算和内存占用，使移动硬件能够支持更高分辨率的合成和视频生成。
架构兼容性：与现有 DiT 架构保持兼容，无需结构更改或复杂的路由模块。
效率 - 质量权衡：证明了结构化、相干性引导的跳过可以大幅降低推理成本（自注意力 FLOPs 减少高达 55%），同时保持与完整模型相当的感知质量。

作者指出，虽然当前方法侧重于空间冗余，但未来的工作旨在将效率扩展至 FFN/MLP 组件，并纳入跨帧相干性目标，以进一步改善视频的时间一致性。