SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation

2026-08-17

AI Summary

This study addresses the quadratic complexity bottleneck of self-attention in video DiTs and the quality limitations of existing linear approximations by proposing SQuad. This framework introduces a novel O(n√n) SQuad-Attention mechanism, combined with flow-matching fine-tuning and an improved two-stage DMD2 distillation strategy, to compress self-attention to sub-quadratic complexity and accelerate sampling without retraining. Applied to the Wan2.2 5B model, SQuad generates high-quality videos in merely six steps, achieving a VBench score of 83.20. Notably, it reduces attention FLOPs by 67× and inference latency by 11×, yielding a 2× end-to-end speedup. These advancements enable, for the first time, efficient and high-fidelity generation of high-resolution long-form videos.

Abstract

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

Problem

• Video Diffusion Transformers
• Self-Attention
• Computational Complexity
• Expressivity Gap
• Video Generation

Innovation

• Sub-Quadratic Attention
• Attention Distillation
• Video Diffusion Transformers
• Distribution Matching Distillation
• Flow-Matching SFT