Bidirectional Sparse Attention for Faster Video Diffusion Training¶

会议: NeurIPS 2026
arXiv: 2509.01085
领域: 视频生成
关键词: 双向稀疏注意力 / 视频扩散模型 / DiT训练加速 / 动态统计阈值 / 查询剪枝

一句话总结¶

针对视频扩散 Transformer（DiT）长序列训练中注意力计算开销占比超 90% 的瓶颈，本文提出首个正交解耦的“查询-键值双向稀疏注意力”（BSA）框架，通过时空局部语义查询剪枝与基于统计分布的动态自适应 KV 块选择，在保证甚至提升生成质量的前提下实现最高 20× 计算量削减与 17.79× 训练注意力加速。

研究背景与动机¶

随着 Sora、Wan2.1、CogVideoX 等视频扩散 Transformer（Video DiT）的发展，模型在生成高分辨率、长时程视频方面展现出强大的物理与视觉连贯性。然而，生成一段仅数秒的高清视频片段，其 3D 时空 Latent 便会展平成数十万个 Token（例如 157 帧 768×1280 分辨率下达 153,600 个 Token）。在全注意力（Full Attention）机制下，计算复杂度随序列长度呈二次方增长（$\mathcal{O}(L^2)$），导致注意力计算占据了视频扩散模型总体训练开销的 90% 以上，极大制约了模型向超长时长与超高分辨率的扩展训练。

现存的高效注意力工作（如 MoBA、VSA、NSA 等）主要借鉴大语言模型（LLM）的思路，仅对 Key-Value（KV）端进行稀疏化建模，假设每个查询只需检索少量全局关键上下文。然而，这种单向剪枝忽略了视频数据内在的“双向冗余”（Bilateral Redundancy）：视频在连续帧之间以及单帧的局部空间区域内存在极高的时空重叠，大量 Query 在注意力特征图上的权重分布高度雷同，对静态背景等冗余区域造成了海量的重复点积计算。此外，现有方法多依赖预设的固定稀疏模式或固定 Top-$k$ 截断，无法适应视频在不同训练步数、不同内容复杂度下的“动态稀疏”（Dynamic Sparsity），极易引发欠选择导致的伪影或过选择引发的计算浪费。

针对上述瓶颈，本文的核心切入点是打破仅剪枝 KV 的传统定势，将稀疏化正交扩展至 Query 侧，并结合数据分布动态调整 KV 预算。核心 idea：提出双向稀疏注意力（BSA）架构，通过 3D 局部块内相似度剪除时空冗余 Query 并用中心输出近似回填，同时基于块间相似度统计分布动态设定阈值筛选全局显著 KV 块，实现长序列视频 DiT 的无损端到端双向训练加速。

方法详解¶

整体框架¶

BSA 针对输入形状为 $(T, H, W)$ 的视频潜变量特征，首先将其划分为 3D 时空长方体块（Cuboid Blocks），以保留局部的物理邻近性并适配现代 GPU（如 FlashAttention 块调度）的硬件执行效率。随后，注意力计算沿着“查询侧”与“键值侧”双向解耦展开：在 Query 侧，算法评估块内所有 Token 与中心代表 Token 的语义余弦相似度，将高相似度的信息冗余 Query 剔除，仅保留低相似度的显著 Token 进入注意力运算，计算完成后再利用中心 Token 输出对被裁剪位置进行特征回填（Token Filling）；在 KV 侧，算法计算查询块与候选 KV 块之间的宏观注意力分数，并依据全局统计分布自适应计算动态保留阈值，贪心累加直至满足重要性预算。两部分正交融合，由定制的 Triton 算子完成前后向高效块级并行计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入视频 Latent<br/>(T, H, W) 展平长序列"] --> B["3D 块划分与索引重排<br/>保持时空局部邻域结构"]
    B --> C["1. 时空查询剪枝与特征回填<br/>中心 Token 相似度评估与去冗余"]
    B --> D["2. 统计动态自适应键值筛选<br/>块间相关性与正态分位数阈值"]
    C --> E["双向稀疏矩阵与 Triton 硬件对齐算子<br/>跳过非活跃块，显存与计算双重加速"]
    D --> E
    E --> F["输出完整时空特征<br/>无损生成高分辨率视频"]

关键设计¶

1. 时空查询剪枝与特征回填：消除跨帧与局部背景的重复计算 在视频序列中，相邻帧之间和单帧背景区域内的 Query 向量高度重复，在全注意力机制下会产生几乎相同的注意力分布图。针对这一单向稀疏无法触及的瓶颈，BSA 在 3D 块内引入基于特征差异度的剪除机制。对于大小为 $B = (C_t, C_h, C_w)$ 的第 $b$ 个 Query 块 $Q_c^{(b)}$，选取其时空中心 Token $q_c^{(b)}$ 作为块锚点（为更精细捕获局部差异，进一步引入小窗口 $(w_t, w_h, w_w)$ 划分）。若块内某 Token $q_i$ 与中心 Token 的余弦相似度极高，则表明该位置承载的信息可被中心 Token 充分代表；反之，与中心 Token 差异较大的特征（即边缘、前景运动等）必须保留。算法按非相似度对块内 Token 降序排列，仅保留前 $r$ 比例的差异化 Token： $$Q^{s}=\bigcup_{b=1}^{N}\left\{q_{i}\in Q_{c}^{(b)}\;\middle|\;\operatorname{rank}_{b}\big(1-\cos(q_{c}^{(b)},q_{i})\big)\leq\lceil r\cdot|Q_{c}^{(b)}|\rceil\right\}$$ 计算完成后，输出序列维度被压缩为 $rL$。为了与后续 DiT 残差连接的原始序列维度严格对齐，BSA 设计了 Token Filling 机制：被剪除的 Token 既然与中心 Token 高度同质，其对应的注意力输出亦可直接由该中心 Token 计算完成的注意力特征直接近似复制填回，从而在 $\mathcal{O}(L) + \mathcal{O}(L \log B)$ 的极低排序开销下实现无损重构。

2. 统计动态自适应键值筛选：打破固定Top-k在时空波动下的刚性失真 不同时空区域的 Query 对上下文信息的依赖范围差异巨大，前景快速移动的主体需要跨越全序列的全局交互，而平缓背景只需要局部联系；若采用固定的剪枝比例或固定 $k$ 个 KV 块，常会导致复杂区域信息欠采样产生瑕疵，或平坦区域冗余保留拖慢计算。BSA 提出基于块间注意力分数统计分布的动态自适应阈值机制。针对各块间的粗粒度注意力得分 $S_b$，模型利用其均值和标准差构建动态门控阈值 $p$： $$p=\operatorname{mean}(S_{b})+\operatorname{std}(S_{b})\cdot U(1-k/n)$$ 其中 $U(\cdot)$ 为分位数函数，$n$ 为总块数，$k$ 为期望的目标候选容量。随后，对于每个具体的 Query 块 $i$，算法以最小索引集合 $S_i$ 贪心纳入相关度最高的 KV 块，直至其注意力权重的归一化累加值达到或超过自适应阈值 $p$： $$\min |S_{i}|\quad\text{s.t.}\quad\sum_{(i,j)\in S_{i}}\frac{\exp(Q_{i}K_{j}^{\top})}{\sum_{j^{\prime}}\exp(Q_{i}K_{j^{\prime}}^{\top})}\geq p$$ 该统计策略使得保留的 KV 块数量随输入内容的实际语义丰富度动态伸缩，从数学上保证了复杂关键内容不失真，平缓低熵内容最大化稀疏。

3. 双向稀疏矩阵与 Triton 硬件对齐算子：将算法稀疏度真实兑现为训练加速 单纯的理论 FLOPs 降低如果缺乏硬件对齐的底层实现，往往会因非连续显存访问和分支发散导致实际运行变慢。BSA 为前后向计算开发了专用的 Triton 块稀疏算子。算子将稀疏矩阵组织为 GPU SM 友好的 Tile 结构：在 Query 侧通过紧凑连续索引映射处理变长块，在 KV 侧利用 $q2k\_num$ 和 $q2k\_index$ 索引张量为每个 Query 块按需指派需加载的 KV 块。在前向传播中，未被选中的块在硬件层面被直接跳过（Skip Tile），彻底规避无用计算与显存加载；在反向传播中，算子沿用完全对称的重排序反向散射逻辑计算梯度，兼顾 FlashAttention 级的吞吐与精确的反向梯度流。

一个完整示例¶

以一段典型视频特征为例，假设经 VAE 编码后的时空序列尺寸为 $61 \times 448 \times 832$，总序列长度为 $L = 23,296$ 个 Token： 1. 3D 分块：按 $(4, 4, 4)$ 尺寸划分为 3D 块，同时在内部采用 $(2, 2, 2)$ 窗口确定局部中心锚点； 2. Query 剪枝：块内评估余弦距离，设定保留率 $r=0.5$，原本 23,296 个 Query 被实时剔除 50% 的背景冗余特征，仅剩 11,648 个代表性 Query 参与注意力打分； 3. KV 动态选择：计算块间得分分布，统计动态阈值 $p$ 自适应将每个 Query 块关联的平均 KV 块比例收敛至全量的 $14\%$，综合稀疏度提升至 $0.93$； 4. 注意力计算与回填：Triton 算子仅在稀疏激活的掩码块上计算点积，耗时削减至全注意力的 $1/12.85$；计算输出后，11,648 个剪除位置由所属中心 Token 的输出瞬间填充对齐，无缝传递给后续 Feed-Forward 层。

损失函数 / 训练策略¶

BSA 完全保持原有的流匹配（Flow Matching）或扩散均方误差损失（Diffusion MSE Loss），无需引入额外的辅助正则项。为保证模型从头训练（From-Scratch）的数值稳定性，模型设计了退火稀疏度训练调度（Annealed Attention Sparsity Schedule）： - 在初始训练阶段（前若干步）采用 100% 全注意力以稳定初始表征建立； - 随后每隔 30 个 step，稀疏度平滑递增 0.03，直至达到设定的最大稀疏上限（如 0.90 ~ 0.95）； - 在 Sparse KV 中，Top-$k$ 块候选从全量平滑衰减至 $0.1 \times$ 总块数，使动态统计阈值平滑收敛，最终完成 30,000 steps 的稳定收敛。

实验关键数据¶

主实验¶

在 Wan2.1-1.3B 模型上，对不同时空分辨率和 Token 规模展开端到端从头预训练，基于 VBench 基准对比全注意力与 BSA：

序列分辨率与 Token 规模	方法	稀疏度	文本一致性 (TextConsis) ↑	背景一致性 (BGConsis) ↑	图像质量 (ImageQual) ↑	主体一致性 (SubConsist) ↑	FLOPs ↓	注意力训练加速比 ↑
61×448×832 (23,296 tokens)	Full Attention	-	32.71%	95.12%	64.33%	92.34%	$1.51 \times 10^{12}$	1.00× (基准)
61×448×832 (23,296 tokens)	BSA (Ours)	0.93	32.79%	95.22%	64.29%	92.39%	$1.05 \times 10^{11}$	12.85×
157×768×1280 (153,600 tokens)	Full Attention	-	34.76%	93.26%	65.91%	93.79%	$6.99 \times 10^{13}$	1.00× (基准)
157×768×1280 (153,600 tokens)	BSA (Ours)	0.95	34.93%	93.41%	66.03%	94.13%	$3.49 \times 10^{12}$	17.79×

在免训练推理（Training-Free Inference）配置下，BSA 在 23K 与 153K Token 下分别取得 4.6× 和 6.2× 的端到端推理加速，且各项 VBench 指标均持平或略优于 Full Attention 基线。

与相关稀疏注意力基线的横向对比如下：

序列规模	方法	稀疏度	文本一致性 ↑	背景一致性 ↑	图像质量 ↑	主体一致性 ↑	FLOPs ↓	加速比 ↑
23,296 tokens	MoBA	0.80	32.56%	95.14%	64.14%	92.05%	$3.02 \times 10^{11}$	1.2×
23,296 tokens	VSA	0.87	32.65%	95.03%	64.25%	92.21%	$1.96 \times 10^{11}$	4.5×
23,296 tokens	BSA (Ours)	0.93	32.79%	95.22%	64.29%	92.39%	$1.05 \times 10^{11}$	12.85×
153,600 tokens	MoBA	0.80	34.34%	93.05%	65.34%	93.49%	$2.62 \times 10^{12}$	2.3×
153,600 tokens	VSA	0.87	34.72%	93.22%	65.87%	93.72%	$4.54 \times 10^{11}$	6.2×
153,600 tokens	BSA (Ours)	0.95	34.93%	93.41%	66.03%	94.13%	$3.49 \times 10^{12}$	17.79×

消融实验¶

在 23,296 tokens 序列长度下对 Query 剪枝与 KV 统计筛选各组件进行系统消融：

组件配置	设定分支	稀疏度	验证损失 (Val Loss) ↓	文本一致性 ↑	背景一致性 ↑	图像质量 ↑	主体一致性 ↑	FLOPs ↓	加速比 ↑
Sparse Query 单独	基础块剪枝 (Original)	0.50	0.211	32.83%	95.25%	64.34%	92.44%	$7.5 \times 10^{11}$	1.96×
Sparse Query 单独	带小窗口机制 (w/ Window)	0.50	0.208	32.85%	95.29%	64.36%	92.44%	$7.5 \times 10^{11}$	1.98×
Sparse KV 单独	固定阈值 (Original)	0.86	0.210	32.84%	95.24%	64.30%	92.41%	$2.1 \times 10^{11}$	6.05×
Sparse KV 单独	统计动态阈值 (w/ Statistic)	0.89	0.209	32.82%	95.25%	64.28%	92.42%	$1.67 \times 10^{11}$	6.12×
Full Attention	-	0.00	0.213	32.71%	95.12%	64.33%	92.34%	$1.51 \times 10^{12}$	1.00×
BSA 完整方案	Sparse Query + KV	0.93	0.212	32.79%	95.22%	64.29%	92.39%	$1.73 \times 10^{11}$	12.85×

关键发现¶

双向正交性倍增收益：单做 Query 剪枝带来 1.98× 加速，单做 KV 稀疏带来 6.12× 加速；二者结合时稀疏度相乘达到 0.93，训练注意力加速直接跃升至 12.85×，且验证损失从全注意力的 0.213 降至 0.212，充分证明两端冗余的正交独立性。
序列越长，加速与质量优势越显著：在 153K tokens 序列下，BSA 实现高达 17.79× 的加速与 20× 的 FLOPs 削减，且生成的多项一致性指标（文本 34.93% vs 34.76%、背景 93.41% vs 93.26%）反超全注意力。这表明剪除高度重复的静态 Query 能迫使模型注意力更精准聚焦于动态前景与实质语义变化。
稀疏度临界拐点：实验显示当稀疏度从 0 提升至 0.93 时验证损失稳定在 0.212 左右；然而一旦稀疏度突破 0.95 极限，损失剧烈飙升，生成质量出现阶跃式退化，表明 0.93~0.95 是视频 DiT 动态信息的保真临界域。

亮点与洞察¶

首次打破长序列视频注意力的单向 KV 剪枝惯性，揭示并系统论证了 Query 侧时空分布冗余现象，通过中心 Token 余弦距离筛选与特征回填构造出闭环无损计算流。
引入正态分位数统计动态阈值替代静态超参 Top-$k$，让每个 Query 块自适应决定所关注的上下文容量，兼顾了平缓背景的高压缩与运动前景的高精度。
为现代 GPU 架构量身定制了 Triton 块对齐核心，规避了非规则稀疏带来的内存碎片，将超过 90% 的理论稀疏度转化为实打实的 17.79× 硬件级训练吞吐提升。

局限与展望¶

时空块尺寸依赖先验经验：当前 3D 划分尺寸采用固定的 $(4,4,4)$ 块及 $(2,2,2)$ 窗口，对于极端高帧率或非均匀长宽比视频，最优块几何构型仍需自适应搜索。
大模型长文本交互影响待进一步拓展：论文主体实验在 1.3B 模型上完整从头训练并在 14B 模型上完成微调验证，未来需进一步在大规模 30B+ 模型及具备长文本交互的复杂影视级管线上展开全面验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出面向视频扩散训练的双向正交注意力稀疏框架与统计自适应阈值机制。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖从头训练、免训练推理、多尺度长序列、全维度消融与 14B 用户评测，数据扎实详尽。
写作质量: ⭐⭐⭐⭐⭐ 动机清晰逻辑严密，图表对比鲜明，算法与底层工程实现交待透彻。
价值: ⭐⭐⭐⭐⭐ 显著攻克高分辨率超长视频 DiT 训练最大的显存与算力墙，工业界落地实用性极强。