利用好视觉Attention的“局部性”，清华&字节提出Token Reorder，无损实现5倍稀疏与INT4量化|token|局部性|序列|清华|算法

本文聚焦视觉生成任务中的“局部性”特性，提出通过简单的令牌重排操作，同时提升视觉特征提取的数值局部性与硬件计算的内存局部性，从而兼顾算法性能与硬件效率。PAROAttention虽以推理效率优化为核心，但其局部性利用策略具有更广泛的适用性。多注意力头可在不同维度自适应聚合局部信息，为训练策略、参数化方法、位置编码等方向提供启发，推动具备合理归纳偏置的视觉基座模型发展。

论文标题： PAROAttention: Pattern-Aware ReOrdering for Efficient Sparse and Quantized Attention in Visual Generation Models 论文链接： https://arxiv.org/abs/2506.16054 项目主页： https://a-suozhang.xyz/paroattn.github.io/

PAROAttention：一种面向视觉生成模型的Attention优化方法，通过引入简单且有效的“令牌重排”（Token Reorder）方法，可以显著提升注意力稀疏化与低比特量化的算法性能保持，和硬件效率优化能力。

对视觉生成模型中的Attention稀疏与低比特量化的关键挑战进行了分析，并归纳两者的共同关键挑战为“多样且分散”的视觉注意力模式。
基于视觉特征提取“局部性”的先验，进一步分析了多样且分散的视觉注意力模式的产生原因为Transformer模型的一维“序列化”（Tokenization）过程破坏了三维空间的相邻性，进而归纳出发现“多样且分散”的不同视觉注意力模式，本质上都在描述不同维度上的三维空间局部聚合，因此，可被统一成一种“规则且分块”的注意力模式。
设计了一种简单有效的“令牌重排”方式，有效的将多样的注意力模式转化为统一且硬件友好的分块模式，并针对性的设计了对应的稀疏与量化策略，以获得更优的性能保持与硬件效率提升。经过完整的CUDA Kernel实现，PAROAttention可在主流视频（CogVideo，Wan）与图像生成（Flux）生成模型上，在生成质量与内容不变的前提下，实现20%的稠密度，与INT4量化，以实现5~10x的Attention加速，与2~4x的端到端加速。

一、前言

近年来，随着视觉生成模型的发展，视觉生成任务的输入序列长度逐渐增长（高分辨率生成，视频多帧生成，可达到10K-100K）。与输入序列长度呈平方复杂度的Attention操作，成为主要的性能瓶颈（可占据全模型的60-80%的开销），有明显的效率优化需求。注意力的稀疏化（Sparse Attention）与低比特量化（Attention Quantization）为常用的Attention优化技巧，在许多现有应用中取得优秀的效果。然而，这些方法在视觉生成模型中，在低稠密度（<50%）与低比特（纯INT8/INT4）时面临着显著的性能损失，具有优化的需求。

本文围绕着视觉任务的“局部性”（Locality）特点，首先提出了系统的分析框架，识别出了视觉生成任务Attention优化的关键挑战在于“多样且分散”的注意力模式，并且进一步探索了该模式的产生原因，并揭示了多样且分散的注意力模式，可以被统一为代表“局部聚合”的块状模式。然后，提出了一种简单且硬件友好的离线“令牌重排”方案以实现注意力模式的统一化，并设计了针对性的稀疏与量化方法，配合高效的CUDA系统设计，展现了更优异的算法性能保持与硬件效率提升。最后，我们讨论了该方案更广泛的应用空间，与对视觉生成算法设计的启发。

二、分析框架：关键问题与如何利用局部性（Locality）

如上文所述，一系列现有的注意力稀疏化与低比特量化方案已取得了进展，但是还存在着一定的挑战与改进空间：

对于稀疏化，一系列现有方案（DiTFastAttn，SparseVideoGen，Sparse-vDiT）尝试依据视觉注意力图的独特模式，设计针对性的稀疏掩码（Sparse Mask，如“窗口状的”，“多对角线”，“垂直线的”），并将其进行组合。然而，适配多样且分散的注意力模式，给稀疏掩码的设计与选择机制带来了严峻的挑战。本文尝试采用另一种视角与方法，并不涉及复杂的掩膜选择机制来适配复杂多样的注意力模式，而是设计方案“重整注意力模式”。让多样且复杂的注意力模式，统一为硬件友好的块状注意力模式，让稀疏方案设计更加简单有效。
对于低比特量化，现有方案（SageAttention系列）SageAttentionV2可以将Attention中的QK计算（Query与Key的乘法）降低至INT4，但是PV（AttentionMap与Value的矩阵乘）计算仍然需要保持为较高的FP8。最新版本的SageAttentionV3采用了FP4量化，但仅在最新的B系列Nvidia GPU上有支持。本文尝试分析了更低位宽的定点量化（全流程INT4）的关键问题，并给出了解决方案。

图：本文（PAROAttn）的优化思路：重整注意力图以便稀疏与量化处理

为寻找Attention稀疏与量化的统一解决方案，我们尝试分析Attention效率优化中稀疏与低比特量化的关键问题，来自于视觉注意力图多样且分散的独特数据分布(如下图左侧所示)

稀疏注意力方案设计需要从2方面考虑：保持算法性能，与提升硬件效率
- 从算法性能角度，需要避免在稀疏过程中错误的删除重要值。由于视觉注意力模式存在多样的结构（对角线，纵向，块状等），且这些特征随着不同的时间步，不同的控制信号而动态变化。因此，注意力模式的多样性，与动态变化，导致设计的掩码难以完全涵盖重要值，对算法性能的保持带来困难。
- 从硬件效率角度，需要设计“结构化”的稀疏掩码，以跳过整块计算来获得实际硬件收益（任意不规则稀疏需要引入额外的索引操作，且可能导致计算负载零散，使得加速收益折损）。特别的，由于FlashAttention涉及逐块进行注意力计算，因此稀疏化的过程中，也应考虑如何与其适配。由于视觉注意力图的模式，往往不与FlashAttention中的分块所对应（对角线模式中，每个块中仅有少量较大值）。因此，注意力图模式的分散性，使得结构化稀疏难以取得，难以获得有效的硬件效率提升。
对于低比特量化算法的设计：关键问题为如何尽量减少量化损失
- 现有工作（如ViDiT-Q）已经分析并指出了低比特量化的关键误差来源在于“量化组内的数据分布差异”，对于注意力量化，为适配FlashAttention，需要选择块状的量化分组。然而，“对角线式”的视觉注意力模式，导致块状的量化分组中，对角线上的元素成为离群值，带来了巨大的组内数据差异，而导致了显著的量化损失。因此，视觉注意力模式的数据分布，导致了显著的量化损失。

图：视觉生成稀疏与量化的关键问题来自于多样分散的注意力模式，与本文的解决方案：采用令牌重排以改进注意力图为统一的分块模式

为解决视觉注意力图多样且分散的独特数据分布给注意力稀疏与量化所带来的挑战。本文的技术路线为：对注意力图进行“重整”（Reorganize），以获得更加统一且易处理的注意力模式。受到视觉特征提取具有“局部性”的先验启发（CNN，SwinTransformer的设计理念，与Hubel与Wiesel的生物学实验），我们进一步分析了视觉注意力模式多样性的产生原因，并发现了“多样的视觉注意力模式本质上都在描述空间上的局部聚合”。

如下图所示，在Transformer的处理过程中，原本三维空间（F - 帧数，H，W - 每帧的图像宽高）会被转化为一维的标记序列（Token Sequence），按照默认的[F,H,W]的顺序排列。这会导致在除了内存上连续的最后一维（W）之外维度的三维空间相邻像素，在标记序列中呈现为按照一定的间隔排列。因此，多对角线的注意力模式，本质上是在描述“其他维度上的局部聚合”，并可以通过令牌顺序的重排列，转化为代表局部聚合的块状模式（将局部聚合的维度转化为内存上连续的维度，如[F,H,W] -> [F,W,H]）。

我们进一步验证了，每个不同的注意力头（Head），在不同情况下，呈现出一致的在某个维度上的局部聚合，进而我们可以通过为每个head选取合理的令牌重排（Token Reorder）方案，将多样且分散的注意力模式，转化为统一的，硬件友好的块状模式，以便于Attention的稀疏与量化。该方案利用了算法侧视觉特征提取的局部性（更好的数值Locality），并将其与硬件计算的局部性将对应（更好的内存与计算Locality），从而获得了同时更优的算法性能保持，与硬件效率提升。

图视觉特征提取“局部性”的示意图三、方案设计整体框架

我们的方案流程如下图所示，对Attention计算的主要瓶颈，两个大规模矩阵乘（QK与PV）都进行了稀疏与量化优化，显著减少其硬件开销。我们基于少量矫正数据离线决定了每个注意力头（Head）的令牌重排方案，与对应的稀疏掩码，几乎不在推理时引入额外的开销。在推理时，仅需跳过稀疏掩码所对应的attention分块，并对剩余的部分逐块进行低比特量化。

图 PAROAttention稀疏与量化方案的流程令牌重排方案（PARO：Pattern-Aware Token Reordering）

我们发现每个不同的注意力头（Head），在不同情况下，呈现出一致的在某个维度上的局部聚合。因此，可以离线地对每个注意力头，选择恰当的令牌重排方式，将注意力图转化为展示局部聚合的块状（Block-wise）模式。我们发现了重排列中的一种特殊方式，维度置换（Permutation），就可以取得不错的效果。对于视频生成模型的特征[F,H,W]，我们为每个注意力头6种可能的置换方式，离线选取最优的置换方式，以获得我们需要的数据分布方式。由于对于注意力稀疏与量化，具有不同的数据分布需求。因此，我们针对稀疏和量化分别设计了重排方式的选取指标，并将两者组合作为最终指标。

稀疏角度：为减小结构化稀疏所带来的损失，要求尽量多的分块是完全稀疏的（Block Sparse）
量化角度：为减少块内数据分布差异大而导致的量化损失，要求块内数据分布是尽量均匀的（Block Uniform）

如下图所示，稀疏与量化对注意力图的分布需求不同，需要组合两者需求，才能找到同时适合两者的重排方式。经过合适的重排处理之后，注意力图呈现块状且较为集中的分布，以适配稀疏与量化处理。

图：不同重排方式的注意力图示意稀疏方案：

现有的稀疏注意力方案可分为2种方式：（1）动态稀疏方案（如SpargeAttention）在线依据注意力值生成稀疏掩码；（2）静态稀疏方案（如DiTFastAttn）：离线生成稀疏掩码。两者各有其优劣。尽管本方法设计的令牌重排（PARO）方案能够同时帮助动态与静态方案，本文对两者优劣进行的分析，并最终选取了静态稀疏方案，作为PAROAttention的主要稀疏方案，具体分析如下：

对于动态稀疏（Dynamic Approach）：
- 在性能保持方面，虽然动态的方案能够自然适配动态变化的模式。但是由于需要在线产生稀疏掩码，只能基于Softmax之前（Presoftmax）的注意力值，它们的相对均匀，不呈现明显模式，难以准确的识别出对应模式。
- 在硬件效率方面，动态稀疏方案引入了在线计算出稀疏掩膜的额外开销（overhead），该开销与掩膜预测的准确度互为权衡，若要获得准确的掩膜，则需要引入相对较大的额外计算。该额外预测过程，一般需要精细设计的CUDA Kernel才能够获得较高的效率收益。
- 总结来看，在较低稀疏比下，动态稀疏方式的准确性与效率提升存在瓶颈，因此本文诉诸静态稀疏方案。
对于静态稀疏（Static Approach）：
- 在性能保持方面，由于静态确定的注意力图，难以适配多样且动态变化的注意力模式，因此静态稀疏方案通常会造成相比动态方案更显著的性能损失。然而，PAROAttention的注意力图重整，已将多样动态变化的注意力模式，转化为了规整且统一的模式，解决了这一静态稀疏的关键挑战。因此，通过利用模式更明显的Softmax后注意力图，能够获得比动态方案更优的算法性能保持。
- 在硬件效率方面，虽然避免了在线计算出稀疏掩膜的额外计算开销，但是离线稀疏掩码会带来额外的显存开销。本文针对该问题进行了对应优化（见下文“CUDA系统设计”部分）

经过重排列处理之后，注意力图呈现出统一的集中的分块模式。因此，我们仅需离线统计每块中的attention数据之和，并设计阈值判断当前块是否需要被跳过（该阈值可以用于调节稠密度），就可以离线获取到稀疏掩码，在推理时不引入任何额外开销（overhead）。如下图所示，相比其他现有的静态注意力稀疏方案，由于预先对注意力模式的统一化，PAROAttention避免了复杂且受限的掩膜设计，而能够与原图非常契合的稀疏掩膜。

量化方案：

对于低比特量化，评估量化损失的关键指标是分组内的数据差异，现有文献通常采用不均衡度（Incoherence）进行衡量，被定义为当前数据组中的最大值，除以平均值（x.max() / x.abs().mean()）。经过合适的令牌重排之后，Attention Map块内的显著数据差异得到明显缓解，从而可以支持更低位宽的量化。

CUDA系统设计

最小化额外开销：PAROAttention所引入的额外开销主要有以下两方面，我们在系统层面进行了针对性优化以最小化额外开销。
- 在线的令牌重排开销：虽然重排方式离线确定，但是令牌重排的过程（维度置换）需要在线进行。为了避免一次显示的从GPU Global Memory到Shared Memory的内存搬移开销，我们进行了算子融合（Layer Fusion）的操作，仅修改重排前算子写入地址的顺序，所引入的额外开销可忽略。
- 静态稀疏掩码的显存开销：由于注意力图体量较大，离线决定的稀疏掩码，可能会占用GB级别的GPU显存。为减少该开销，本文采用了预取（Prefetch）策略，通过新建一个CUDA Stream，在每次运算时，只读取当前层的稀疏掩码，可以将额外的显存开销降低到若干MB级别。
兼容性：PAROAttention的稀疏与量化方案都逐块处理，可直接与兼容FlashAttention。由于重排与稀疏掩码均离线完成，无需精细的CUDA Kernel优化，仅需基于FlashAttention进行跳过整块计算的支持，能够广泛适配各种场景。

四、软硬件实验结果算法性能保持效果

本文在主流视频（CogVideo）与图片生成模型（Flux）上测试了多方面指标，包括了：

视频质量指标：CLIPSIM衡量语义一致性；VQA衡量视频质量；FlowScore衡量时间一致性；
与浮点生成差异：如FVD-FP16衡量特征空间差异，PSNR/CosSim衡量像素空间差异，SSIM衡量结构相似性

典型的实验结论概括如下：

其他基线的稀疏方案在相对较高稀疏比（50%）时，仍会造成可观的质量损失，包括内容变化，图像模糊等；而PAROAttention的稀疏化方案，可以在20%的较高稀疏比情况下，依然生成和浮点结果非常相似的结果，获得比基线方案50%更好的多方面指标。
令牌重排方案PARO，并不局限于静态稀疏方案。其与动态稀疏方案SpargeAttention能够直接适配，并提升生成效果。将30%稠密度的SpargeAttention组合PARO，可以获得与50%稠密度SpargeAttention同等的生成质量。将加速比从1.67x提升至2.22x。
相比于SageAttentionV2（QK INT4，PV FP8），PAROAttention的量化方案可以在无精度损失的情况下，进一步将PV量化到INT4。
PAROAttention的稀疏与量化方案可以并行使用，最激进的优化方案（50%+INT4）相比浮点能取得近10倍的Attention部分延迟优化，同时获得与仅能取得2x左右延迟优化的基线方法类似的算法性能保持。

本文进一步对系统层面优化技巧进行了分析，关键实验结论如下：

PAROAttention的稀疏方案，同时取得了更优的算法性能保持与效率提升。以50%稠密度为例，PAROAttention取得了1.73x的attention加速，超过同等情况下的SpargeAttention（1.67x）与SparseVideoGen（1.42x），由于静态稀疏方案几乎不会引入额外开销，而基线方案的在线稀疏掩膜生成/选择会造成6%到10%左右的额外开销，该开销在更低稠密度下显得给更为明显。
PAROAttention的加速比与理论上限较为接近（50%稠密度，理论2倍，实际1.73倍），凸显了方案的硬件友好性。
PAROAttention的各方面额外开销overhead得到了有效减少，控制在整体的1%之内

五、总结与未来指引

总结来看，本文关注了视觉生成任务的“局部性”特性。通过一个简单且有效的令牌重排操作，可以同时实现算法侧视觉特征提取的局部性（更好的数值Locality），并将其与硬件计算的局部性相对应（更好的内存与计算Locality），从而获得了同时更优的算法性能保持，与硬件效率提升。PAROAttention的方案主要围绕推理效率优化设计，但是采用令牌重排来更好利用特征提取局部性的思想并不局限于推理优化中。不同的注意力头自主的学习到在不同维度上的局部聚合，可以启发优化训练方法，与图像的参数化方式，三维空间的位置编码设计，并进一步推动具有合理归纳偏置（Inductive Bias）的视觉基座模型的构建。

llustration From IconScout By Scout Stores

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线600+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

melodybai@thejiangmen.com

或添加工作人员微信（yellowsubbj）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。