用新方法让视频生成快11倍,还保持原质量
SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation

- 设计了复杂度仅为O(n√n)的注意力蒸馏框架,比传统方法更高效
- 在50亿参数模型上实现与原版几乎相同的质量(83.20 vs 83.08)
- 只需6次神经函数评估即可生成视频,大幅降低计算成本
视频扩散变换器(DiT)的自注意力计算开销随潜在令牌数n呈二次增长,$ ext{O}(n^2)$,导致视频生成时长和分辨率受限。现有线性或低秩替代方案虽降低复杂度,但表达能力下降明显。本文提出SQuad,一种子二次注意力蒸馏框架,将复杂度降至$ ext{O}(n ext{√}n)$,平衡效率与表达力。不从头训练模型,而是通过两阶段蒸馏:流匹配监督微调(SFT)和改进的分布匹配蒸馏(DMD2),使采样更高效。在Wan 2.2 5B文生视频模型上,SQuad在VBench得分达83.20(原模型为83.08),每步每块注意力计算量减少约67倍,注意力延迟降低约11倍,端到端延迟减半,且仅需6次神经函数评估(NFE)即可生成视频,远低于默认的100次。
原文摘要 · Abstract (English)
Video Diffusion Transformers (DiTs) spend most of their compute inside the Self-Attention operation, whose cost grows quadratically, $\mathcal{O}(n^2)$, with the number of latent tokens $n$. For the task of video generation, the token count is large, so this term dominates runtime and memory, and thereby caps the resolution and duration we can generate. Linear $\mathcal{O}(n)$ and low-rank $\mathcal{O}(nk)$ surrogates of Self-Attention trade the full softmax $QK^T$ for cheaper kernels, but rarely recover the original's expressivity, leaving a stubborn quality gap. Motivated by this, we propose SQuad, a Sub-Quadratic Attention Distillation framework that achieves a complexity of $\mathcal{O}(n\sqrt{n})$ in the resulting distilled Attention, naturally balancing the efficiency v/s expressivity trade-off. Instead of training our own Video DiT from scratch, which is prohibitively expensive, we fit a pretrained full softmax Self-Attention DiT into our proposed SQuad-Attention one by distilling the former in two stages: Flow-Matching Supervised Fine-Tuning (SFT), followed by improved Distribution Matching Distillation (DMD2) which additionally makes the sampling more efficient. On the Wan~2.2 5B text-to-video model, SQuAD matches the quadratic teacher on VBench ($83.20$ v/s $83.08$) while cutting the per-step per-block attention FLOPs by $\sim$$67\times$ and attention latency by $\sim$$11\times$, and end-to-end DiT latency by 2$\times$, all while also generating a video in only $6$ Neural Functional Evaluations (NFEs) instead of the default $100$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。