SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation
Animesh Karnewar, Denis Korzhenkov, Amirhossein Habibian, Mohsen Ghafoorian
cs.CV
2026-08-17
把视频 DiT 的自注意力拆成窗口内+跨窗口两次 softmax,Wan 2.2 5B 上 VBench 83.20 追平原版,注意力 FLOPs 降 67 倍,6 步出片
视频扩散 Transformer(DiT)的算力大头在自注意力,开销随 token 数平方增长。视频的 token 动辄上万:Wan 2.2 5B 在 704p 下处理一支视频要过 18480 个 token,平方项直接吃掉整张卡的预算,分辨率和时长都被它压着。
省钱的路子已有两条,各有代价。线性注意力把 softmax 换成核函数近似,复杂度降到 O(n),但 softmax 的非线性与尖锐的输入相关选择性正是表达力所在,换掉之后质量差距一直补不回来。稀疏注意力保留 softmax,只算选出的 token 对,但要么依赖数据相关的动态选择(VSA),要么要专门的稀疏 kernel 才能把纸面 FLOPs 变成真实时延。混合架构留几层全注意力保底,既复杂又没摆脱对 softmax 的依赖。
SQuad 的出发点是一个已知观察:视频 DiT 的注意力图稀疏且重尾,绝大部分权重集中在少数关键 token 上。既然如此,把通信模式固定成结构化的稀疏模式,softmax 原样保留,也许就够了。
SQuad 注意力把一次全注意力拆成两次普通 softmax 注意力,先后做:
两次 pass 之间只有纯重排(einops 级别的索引变换),没有参数、没有额外算术,不被注意的那个轴直接折叠进注意力头数。整个算子就是两个标准 softmax 注意力的复合,一行定制 kernel 都不用写。
窗口取 w=√n 是全部要点。单层成本是两项之和:局部 hd·nw,全局 hd·n²/w,一项随 w 涨、一项随 w 跌,求导取零点正好落在 w=√n,总代价 2hd·n^1.5,即 O(n√n),名字里的 SQuad 来自 sub-quadratic。n=18480 时窗口取 21×2×4=168 个 token,时间轴整段都在窗口里,时间维的混合是精确的,被切的只有空间维。
拆成两半还能拼回全感受野,附录给了证明:任意源 token 先在自己窗口内走到目标 token 所占的槽位,再跨窗口走到目标窗口,两跳可达任意位置;等效权重是两个 softmax 权重的乘积,严格为正,感受野里没有结构性零。这决定了 SQuad 是全注意力的 drop-in 替换,预训练权重靠蒸馏就能重新落座,滑窗这类纯局部方案做不到这一点。
蒸馏分两段,教师是冻结的原版模型。第一段用原 flow-matching 目标做监督微调,8k 步,把网络在新注意力模式下重新坐稳;第二段上 DMD2 分布匹配蒸馏,15-30k 步,把采样压到 6 次 NFE,classifier-free guidance 也蒸进学生,6 就是字面意义的前向次数。训练数据是 VIPE 1M 的视频配 Qwen3-8B 生成的描述,作者承诺放出这些描述。改动只碰自注意力,cross-attention 与 FFN 原样保留。
Wan 2.2 5B,81×704×1280,n=18480:
| 配置 | VBench 总分 | 注意力 FLOPs/块 | 注意力时延 | NFE |
| 原版 | 83.08 | 4.205 TFLOPs | 47.10ms | 100 |
| SQuad(全 30 块) | 83.20 | 0.063 TFLOPs | 4.27ms | 6 |
注意力 FLOPs 降 66.7 倍,时延降 11 倍,参数零增加。整个 DiT 单次前向在 torch.compile 下从 667ms 降到 314ms(2.1 倍)。对照组里 VSA、Attention Surgery、ReHyAt 各背 72M 到 283M 新参数;后两个在 eager 模式下比 870ms 的原版还慢(1006ms、1757ms),纸面省下的 FLOPs 要靠编译或定制 kernel 兑现。SQuad 的两次 pass 就是普通注意力,一次 torch.compile 直接吃到全部收益。
分辨率拉长,差距继续放大:5B 上注意力 FLOPs 削减从 n=18480 的 67 倍涨到 n=73920 的 129 倍;14B 上从 88 倍涨到 177 倍,块时延 300.29ms 降到 59.24ms(14B 只跑了效率评测,没做蒸馏)。实测 FLOPs 与理论曲线贴合:原版与 4hdn² 的偏差在 0.4% 以内,SQuad 平均只比理想最优高 8%,差额全部来自窗口必须取整。
消融把设计坐实。30 块全换时,只留局部或只留全局,VBench 从 83 掉到 62.6,语义分掉到 18.5;先局部后全局 83.20,反序 82.99。窗口形状上,拉满时间维的 21×2×4 拿 82.99,纯空间窗 82.31。两段训练都必要:SFT 单独 73.03,DMD2 单独 80.91,合起来 82.99。24 人、1179 组对比的用户研究里,拿 SQuad 的 6 步产物对原版 100 步产物,41% 用户更喜欢前者,33% 认为无差别。
对跑 Wan 级视频模型的人,这是能直接抄的提速:质量不掉、6 步出片、时延减半,且不写一行 CUDA,任何能跑 torch.compile 的栈都能用,端侧部署对此最敏感。方法层面它给了一个干净的证据:O(n²) 的全 token 通信未必必要,固定结构化的稀疏模式配上真 softmax 就能顶住,作者也明说这个模式值得作为从零预训练的注意力候选,不只是蒸馏目标。收益随分辨率和时长增长,恰好是视频生成正在去的方向。
作者自列三条:蒸馏与 DMD2 步数蒸馏绑在一起,注意力改动本身无法单独归因;只试了两段复合,更多段可能用一点深度换更便宜的表达;全部验证都在视频上。读下来再补几点。与同类高效注意力的质量对比并不占优:VSA 84.14、Radial Attention 84.56、带训练的 Jenga 84.36 都高于 SQuad 的 83.20,SQuad 追平的是教师,卖点在时延、FLOPs 与零参数,读表格时别混同。Wan 2.1 1.3B 上 SQuad 拿 82.70,低于原版的 83.26,也低于只做 DMD 的 83.04,小模型上质量损失看得见。VBench 上的差距多在 0.1 到 1.4 分之间,接近该基准的噪声水平,用户研究也只有 24 人,结论强度有限。