💻 computer science

PISA: Piecewise Sparse Attention Is Wiser for Efficient Diffusion Transformers

本文介绍了 PISA，一种无需训练的分段稀疏注意力机制，它通过利用块级泰勒展开而非直接丢弃来近似非关键块，从而在扩散 Transformer 中实现了亚二次复杂度，在保持高生成质量的同时提供了显著的加速。

原作者： Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

发布于 2026-02-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Haopeng Li, Shitong Shao, Wenliang Zhong, Zikai Zhou, Lichen Bai, Hui Xiong, Zeke Xie

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

想象一下，你正试图带领一个艺术家团队绘制一幅规模宏大、细节极其丰富的壁画（就像一段高清视频或图像）。在这个 AI 世界中，这个“团队”就是扩散 Transformer（Diffusion Transformer），而“绘画”的过程涉及观察每一个单笔触（token），以此来决定下一个笔触应该是什么。

问题在于？目前这些 AI 团队的工作方式就像一条严格的规则：“观察整幅壁画中的每一个笔触，以此来决定你的下一步行动。”随着壁画变得越来越大（分辨率更高或视频更长），这变得几乎不可能实现。艺术家们会被淹没，过程会变得极其缓慢，计算机也会耗尽能量。这就是文中提到的“二次方复杂度（quadratic complexity）”瓶板。

为了解决这个问题，之前的方法尝试了一种**“保留或丢弃（Keep or Drop）”**策略。他们说：“好吧，我们忽略 80% 的笔触，只观察最重要的 20%。”

缺陷： 这就像是在画一张脸，却因为眼睛和嘴巴不在你的“前 20%”名单上而忽略了它们。结果往往是模糊、出现故障或缺失关键细节。

新思路：PISA（分段稀疏注意力机制）

该论文的作者提出了一种更聪明的方法，称为 PISA。PISA 不仅仅是忽略那些“不重要”的部分，而是将壁画视为一个分段式的拼图。

以下是它的工作原理，使用了一个简单的类比：

1. “精确 vs. 近似”策略

想象你是一位正在为人群制作一大锅汤的厨师。

旧方法（稠密注意力/Dense Attention）： 你品尝每一勺汤，以获得完美的风味。这很准确，但非常耗时。
“丢弃”法（标准稀疏注意力/Standard Sparse Attention）： 你只品尝前 20% 的汤，然后忽略其余部分。汤的味道会变得很奇怪，因为你错过了后半部分的盐分。
PISA 法：
- 关键区块（“精确”部分）： 你识别出最重要的食材（比如主要的香料或肉类）。你会对这些部分进行精确且细致的品尝。
- 非关键区块（“近似”部分）： 对于剩下的汤（水、高汤、蔬菜），你不需要品尝每一滴。相反，你会使用一个数学捷径（泰勒展开/Taylor expansion），根据该部分的平均味道来估算风味。

2. 为什么这更“聪明”

论文发现了一个迷人的事实：AI 注意力中那些“不重要”的部分（非关键区块）实际上是非常可预测的。它们遵循一种平滑、稳定的模式。

因为它们是可预测的，所以你不需要丢弃它们。你可以通过近似计算来快速处理它们，而不会丢失最终图像的“风味”。
这就像是通过测量房间角落的平均温度，来估算整个房间的温度，而不是测量每一寸空气。

3. 结果：速度与质量兼得

通过结合精确计算重要的部分和对剩余部分的智能近似，PISA 实现了两件事：

速度： 它比目前的最佳方法运行速度快 2 到 2.5 倍。在论文的测试中，生成一段以前需要 26 分钟的视频，现在仅需约 11 分钟。
质量： 生成的图像和视频与那个缓慢的“完美版”一样出色。事实上，在某些测试中，PISA 产生的结果甚至比那些单纯丢弃信息的“快速”方法更好。

背后的“魔法”

论文提到了几个让这一切成为可能且无需重新训练 AI 的技术技巧：

无需重新训练： 因为 PISA 与原始“完美”思维方式的模仿非常接近，你可以直接将其应用于现有的 AI 模型（如 Wan2.1 或 FLUX.1），它就能直接运行。你不需要教 AI 一种新的语言。
“混合”内核（Hybrid Kernel）： 作者构建了一个自定义的计算机程序（内核），可以在“精确品尝”和“估算”之间瞬间切换，这样计算机就不会感到困惑或变慢。

总结

把 PISA 看作是 AI 艺术工作室里的一个聪明经理。

旧经理： “忽略 80% 的工作！”（结果：糟糕的艺术品）。
PISA 经理： “把 100% 的精力集中在重要的细节上，并对背景噪音进行快速、聪明的猜测。”（结果：高质量、高速度的艺术品）。

该论文证明了这种“猜测”并非偷懒的捷径；它是一种在保持杰作完整性的同时，节省时间的、在数学上严谨可行的方法。

技术摘要：PISA —— 用于高效扩散 Transformer 的分段稀疏注意力机制

1. 问题陈述

扩散 Transformer (DiT) 已成为高保真图像和视频生成的标准。然而，随着输入序列长度 ( $L$ ) 的增加（用于更高分辨率或更长视频），其自注意力机制的二次方复杂度 ( $O(L^2)$ ) 成为了扩展性的瓶颈。

现有解决方案面临显著局限性：

稀疏注意力（“保留或丢弃”）： 这些方法通过丢弃非关键的键-值 (KV) 块来加速计算，但由于丢失了上下文，在高度稀疏的情况下会导致性能下降。此外，它们在处理短序列时，选择过程带来的开销往往会抵消其收益。
线性/混合注意力： 改变注意力分布的方法（例如，通过特征映射或将稀疏分支与线性分支进行加性组合）通常会破坏预训练模型的内在归一化。这需要昂贵的重新训练或微调，导致无法直接复用预训练权重。
硬截断 (Hard Truncation)： 直接丢弃块会导致注意力分布发生剧烈偏差，从而造成生成内容的质量损失。

核心挑战在于设计一种具有次二次方复杂度、能加速推理，同时又能保持全量注意力精度，并与预训练权重兼容且无需重新训练的机制。

2. 方法论：分段稀疏注意力 (PISA)

PISA 引入了一种全新的**“精确或近似”**范式，超越了传统的“保留或丢弃”策略。它将注意力视为一个分段过程，以次二次方的复杂度覆盖全量注意力跨度。

核心洞察

作者分析了预训练的 DiT，发现了两个关键特性：

分布稳定性： 非关键块的 pre-softmax 注意力分数呈现出以零或负值为中心的对称钟形分布。这使得它们非常适合通过泰勒展开进行精确近似，而非必须进行精确计算或直接丢弃。
归一化一致性： 为了保持无需训练的兼容性，近似值必须集成到 softmax 的分子和分母内部，而不是简单地将不同分支的输出相加（这会违反加权求和规则）。

核心算法

PISA 将 KV 块划分为两组：

关键块 ( $S_i$ )： 进行精确计算，以保留高显著性信息。
非关键块 ( $U_i$ )： 使用块级泰勒展开进行高效近似。

输出 $o_t$ 通过归一化加权值聚合得出：
$o_t = \frac{N_t}{D_t}$

其中分母 ( $D_t$ ) 和分子 ( $N_t$ ) 由以下部分组成：

精确稀疏项： 为选定的块 $j \in S_i$ 计算。
块级零阶近似： 对于未选中的块，利用块质心 $\bar{k}_j$ 和值的总和进行近似。
全局一阶修正： 为了避免为每个未选中的块都计算独立的一阶项（这会受限于内存带宽），PISA 采用了混合阶近似。它使用一个全局统计量 $\bar{H}$ （所有块中块级一阶矩阵的平均值）来近似整个尾部集合的一阶梯度。

分子 $N_t$ 的公式变为：
$N_t = \sum_{j \in S_i} \sum_{n} e^{q_t k_{j,n}^\top} v_{j,n} + \sum_{j \in U_i} \alpha_{t,j} \left(\sum_{n} v_{j,n}\right) + q_t \bar{H} \sum_{j \in U_i} \alpha_{t,j}$
其中 $\alpha_{t,j} = \exp(q_t \bar{k}_j^\top)$ 。

协方差感知路由

为了最小化近似误差，PISA 使用了协方差感知 Top-k 策略。块的选择评分不仅包含注意力分数，还结合了块协方差矩阵的范数 ( $M_j = \|H_j - \bar{H}\|_2$ )。具有高语义相关性或高异质性的块被选中进行精确计算，而较平滑的块则被委托给泰勒展开。

硬件实现

PISA 通过一个融合算子 (Fused Kernel) 实现，该算子在在线 softmax 过程中交织进行精确计算和近似扫描（类似于 FlashAttention 风格）。流水线包括：

预计算： 计算块质心、值总和以及全局统计量 $\bar{H}$ 。
阶段 1（精确）： 将选定块加载到 SRAM 中进行精确注意力计算。
阶段 2（近似）： 使用零阶展开对未选中的块聚合向量进行扫描。
阶段 3（修正）： 注入全局一阶修正项。

这种设计确保了方法的高内存效率，并避免了为每个块加载不同矩阵所带来的瓶颈。

3. 核心贡献

新型分段稀疏注意力： 一种无需训练的机制，通过统一精确稀疏计算和块级近似，实现了具有次二次方复杂度的全量注意力跨度覆盖。
混合阶近似方案： 一种结合了块级零阶展开与全局一阶修正的策略。与标准稀疏注意力相比，该策略在几乎不增加计算开销的前提下，显著提升了精度。
协方差感知路由： 基于误差分析的推导，通过结合注意力分数和块协方差范数来选择块，从而最小化近似偏差。
最先进的性能表现： 在多种任务（视频和图像生成）中展示了卓越的效率和质量，且无需重新训练。

4. 实验结果

作者在视频生成模型（Wan2.1-1.3B/14B, Hunyuan-Video）和图像生成模型（SD 3.5, FLUX.1）上评估了 PISA。

视频生成：
- Wan2.1-14B： 在 87.5% 的稀疏度下，相比密集注意力实现了 1.91× 的加速（1468s $\to$ 818s），同时保持或略微提升了质量指标（VBench, PSNR, SSIM）。
- Hunyuan-Video： 实现了 2.57× 的加速（1651s $\to$ 641s），并保持了稳定的质量。
- 对比： 在速度和质量（特别是在保持细粒度细节和时间一致性方面）上均优于现有的稀疏方法，如 SpargeAttn 和 SVG2。
图像生成：
- FLUX.1-dev： 在 85% 的稀疏度下实现了 1.2× 的加速（8.32s $\to$ 6.87s）。值得注意的是，与密集注意力和 SpargeAttn 相比，PISA 在人类偏好指标（ImageReward, HPSv2）和相似度指标（PSNR, SSIM）上取得了更高的分数。
- SD 3.5： 展示了显著的加速效果（例如在 SD 3.5-M 上实现 1.4× 加速），同时保持了与密集注意力相当甚至更好的 FID 分数。
算子效率：
- 在各种序列长度（4K 到 32K）和密度下，PISA 的延迟始终优于 FlashAttention-2/3 和 SpargeAttn。
- 与 SpargeAttn 在短序列上性能下降不同，PISA 在 4K token 时仍保持着速度优势。

5. 重要性与主张

论文声称 PISA 代表了从“保留或丢弃”到“精确或近似”稀疏注意力的范式转变。其主要意义在于：

弥合速度与质量的差距： 它有效地弥合了推理速度与生成质量之间的鸿沟，实现了近乎无损的近似，且开销极小。
无需训练的兼容性： 通过将近似直接集成到 softmax 归一化框架中，PISA 允许直接复用预训练权重，消除了线性或混合注意力方法所需的昂贵重新训练或微调。
可扩展性： 它为高分辨率图像和长视频生成提供了一个统一的解决方案，在不牺牲生成内容保真度的前提下，解决了二次方复杂度的瓶颈。

作者强调，该方法是首个将块级泰勒展开与稀疏注意力原生集成到统一框架中的方法，从而保留了预训练模型的内在分布。