SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation

📄 arXiv: 2608.16585v1 📥 PDF

作者: Animesh Karnewar, Denis Korzhenkov, Amirhossein Habibian, Mohsen Ghafoorian

分类: cs.CV

发布日期: 2026-08-17

💡 一句话要点

提出SQuad以解决视频生成中的自注意力计算效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 视频生成 自注意力 蒸馏训练 计算效率 多模态学习

📋 核心要点

  1. 现有的视频生成方法在自注意力计算中面临平方复杂度的问题，导致资源消耗过大，限制了生成质量和时长。
  2. SQuad通过引入子平方注意力蒸馏框架，优化了自注意力的计算复杂度，平衡了效率与表达能力。
  3. 在Wan~2.2 5B文本到视频模型上，SQuad在保持生成质量的同时，将计算复杂度降低了约67倍，显著提升了生成效率。

📝 摘要（中文）

视频扩散变换器（DiTs）在自注意力操作中消耗了大部分计算资源，其成本随着潜在标记数量的增加而呈现平方增长，导致生成视频的分辨率和时长受到限制。为了解决这一问题，本文提出了SQuad，一个子平方注意力蒸馏框架，能够在保持表达能力的同时，实现复杂度为$ ext{O}(n ext{sqrt}(n))$的高效注意力计算。通过对预训练的全softmax自注意力DiT进行两阶段蒸馏，SQuad在多个性能指标上超越了传统方法，显著提高了计算效率。

🔬 方法详解

问题定义：本文旨在解决视频生成中自注意力计算的高复杂度问题，现有方法在处理大量潜在标记时，计算和内存消耗过大，限制了生成能力。

核心思路：提出SQuad框架，通过蒸馏预训练的全softmax自注意力DiT，采用两阶段的流匹配监督微调和改进的分布匹配蒸馏，达到更高的计算效率。

技术框架：SQuad的整体架构包括两个主要阶段：首先进行流匹配监督微调（SFT），然后进行改进的分布匹配蒸馏（DMD2），以提高采样效率。

关键创新：SQuad的核心创新在于实现了$ ext{O}(n ext{sqrt}(n))$的复杂度，显著降低了计算资源的消耗，同时保持了生成质量，解决了传统方法的表达能力不足问题。

关键设计：在设计中，采用了特定的损失函数以优化蒸馏过程，并调整了网络结构以适应新的注意力机制，确保在减少计算的同时不损失生成效果。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果表明，SQuad在VBench上与传统的二次教师模型相比，生成质量相当（$83.20$对$83.08$），同时将每步每块注意力的FLOPs减少约67倍，注意力延迟减少约11倍，整体DiT延迟降低2倍，且仅需6次神经功能评估生成视频，显著提升了效率。

🎯 应用场景

该研究的潜在应用领域包括高效的视频生成、实时视频处理和多媒体内容创作等。通过提高生成效率，SQuad能够在资源受限的环境中实现高质量的视频生成，推动相关领域的发展，具有重要的实际价值和未来影响。

📄 摘要（原文）

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$.