Method
Click any sample to enlarge it, compare it to the corresponding tuned Wan 2.2 5B samples and view its text prompt.
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 $\sim$$2\times$, all while also generating a video in only $6$ Neural Functional Evaluations (NFEs) instead of the default $100$.
This project page is intended to complement the paper by presenting video results that static figures cannot adequately convey. It offers a concise, visual overview of the project rather than an exhaustive account; for the complete methodology, derivations, and quantitative results, we refer the reader to the paper PDF.
Both training stages are necessary. For now. When all 30 blocks of the model are replaced with SQuad-Attention: SFT alone collapses (VBench Tot. 73.03), and DMD2 alone reaches only 80.91, but the two stages together recover the full 82.99 VBench Tot.
Composing both passes matches original's expressivity whereas either one alone is not able to recover the quality. Alternate order swaps local and global block-by-block instead of keeping one order throughout.
SQuad-Attention factorizes Self-Attention into a global and a local branch, where the local branch attends within a window of $T \times H \times W$ tokens. For our $21 \times 22 \times 40$ token grid ($n = 18{,}480$, so $\sqrt{n} \approx 136$), we sweep the window shape while holding its volume close to $\sqrt{n}$, which keeps the $O(n\sqrt{n})$ complexity fixed and isolates the effect of how those tokens are distributed across time v/s space. Drag the slider to sweep from purely spatial windows, through the isotropic ones, to windows spanning the full temporal extent, ending on the $21 \times 2 \times 4$ shape we propose.
@article{karnewar2026squad,
title = {SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation},
author = {Animesh Karnewar and Denis Korzhenkov and Amirhossein Habibian and Mohsen Ghafoorian},
journal = {arXiv preprint arXiv:2608.16585},
year = {2026}
}