SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation
Animesh Karnewar, Denis Korzhenkov, Amirhossein Habibian, Mohsen Ghafoorian
cs.CV
2026-08-17
Two softmax passes, in-window then across-windows, replace full attention at O(n sqrt n). Wan 2.2 5B keeps VBench parity (83.20 vs 83.08) with 67x fewer attention FLOPs.
Self-attention dominates the compute budget of video diffusion transformers, and its cost grows quadratically with token count. Video token counts are huge: Wan 2.2 5B pushes 18,480 latent tokens through attention for a single 704p clip, and the quadratic term caps the resolution and duration a GPU can afford.
Existing fixes each pay a price. Linear and low-rank attention replace softmax with a kernelized surrogate and reach O(n), but the nonlinearity and sharp input-dependent selectivity of softmax are exactly what they give up, and the quality gap never fully closes. Sparse-attention methods keep softmax but compute only selected token pairs, which means either data-dependent token selection (VSA) or hand-tuned sparse kernels to turn paper FLOP savings into wall-clock gains. Hybrid designs keep a few full-attention layers as insurance, staying complex and still dependent on the operation they set out to remove.
SQuad starts from a known observation: attention maps in video DiTs are sparse and heavy-tailed, with most of the mass on a few critical tokens. If that holds, a fixed, structured communication pattern with softmax intact might be enough.
SQuad-Attention composes two ordinary softmax attentions, run one after the other:
Between the passes there is only pure re-indexing (einops-style rearrangement), with no parameters and no arithmetic; the axis not attended over is folded into the head count. The operator is literally two standard softmax attentions composed, which is why it needs no custom kernel.
The window size carries the whole argument. Per-layer cost is the sum of hdnw for the local pass and hdn^2/w for the global pass, one term rising and one falling in w, and the minimum sits exactly at w = sqrt(n), for a total of 2hdn^1.5, that is O(nsqrt(n)), which is where the name SQuad (sub-quadratic) comes from. At n = 18,480 the window is 21x2x4 = 168 tokens, spanning the full temporal extent so temporal mixing stays exact; only space gets cut.
The two passes also recover a full receptive field within one layer, proven in the supplement: any source token reaches any target in two hops, first inside its own window to the slot the target occupies, then across windows at that slot, and the effective weight is a product of two softmax weights, strictly positive with no structural zeros. This makes SQuad a drop-in replacement for full attention rather than a local approximation, and it is why distillation alone can re-seat pretrained weights, something sliding-window schemes cannot do.
Distillation runs in two stages against a frozen teacher. Stage 1 fine-tunes with the original flow-matching objective for 8k iterations to re-seat the network under the new attention. Stage 2 applies DMD2 distribution-matching distillation for 15-30k iterations, compressing sampling to 6 NFEs with classifier-free guidance distilled into the student, so 6 is the literal forward-pass count. Training data is VIPE 1M video with Qwen3-8B-generated captions, which the authors plan to release. Only self-attention is touched; cross-attention and the FFN stay as they are.
Wan 2.2 5B at 81x704x1280, n = 18,480:
| Config | VBench Total | Attn FLOPs/block | Attn latency | NFE |
| Original | 83.08 | 4.205 TFLOPs | 47.10 ms | 100 |
| SQuad (all 30 blocks) | 83.20 | 0.063 TFLOPs | 4.27 ms | 6 |
That is 66.7x fewer attention FLOPs and 11x lower attention latency, with zero added parameters. A full DiT forward drops from 667 ms to 314 ms under torch.compile (2.1x). Among the baselines, VSA, Attention Surgery and ReHyAt carry 72M to 283M extra parameters each, and the latter two are slower than the 870 ms full-attention original in eager mode (1,006 ms and 1,757 ms): their FLOP savings only materialize with compilation or custom kernels. SQuad's two passes are ordinary attentions, so a single torch.compile call captures the whole gain.
The gap widens with scale. On the 5B backbone the attention-FLOP reduction grows from 67x at n = 18,480 to 129x at n = 73,920; on Wan 2.1 14B from 88x to 177x, with block latency falling from 300.29 ms to 59.24 ms (the 14B was efficiency-evaluated only, not distilled). Measured FLOPs track the theory: the original lands within 0.4% of 4hdn^2, and SQuad sits 8% above the ideal optimum on average, with the residual fully explained by windows having to be whole numbers of tokens.
Ablations pin the design down. Replacing all 30 blocks, local-only or global-only collapses VBench from 83 to 62.6, with the semantic score down at 18.5; local-then-global reaches 83.20 against 82.99 for the reverse order. Temporal windows (21x2x4) beat spatial ones, 82.99 against 82.31. Both training stages are load-bearing: SFT alone 73.03, DMD2 alone 80.91, together 82.99. In a 24-person, 1,179-pair user study, 41% preferred SQuad's 6-NFE videos over the original's 100-NFE output, with 33% calling it a tie.
For anyone running Wan-class video models this is a directly usable speedup: teacher-level quality, 6-step sampling, half the latency, and not a line of CUDA, so any stack that can run torch.compile can adopt it, and on-device deployment cares most. Conceptually it is clean evidence that full O(n^2) token-to-token communication may not be necessary: a fixed structured pattern with a true softmax holds up, and the authors argue it deserves a try as a pretraining attention, not only as a distillation target. The gains grow with resolution and duration, which is exactly where video generation is heading.
The authors list three: distillation is coupled with DMD2 step distillation, so the attention change cannot be isolated from the sampling speedup; only two composed passes were tried, and more of them might trade a little depth for a cheaper, more expressive operator; and all validation is on video. A few more caveats from the tables. SQuad does not beat efficient-attention rivals on quality: VSA scores 84.14, Radial Attention 84.56 and trained Jenga 84.36 on VBench against SQuad's 83.20. SQuad matches its teacher; its edge is latency, FLOPs and zero parameters. On Wan 2.1 1.3B it lands at 82.70, below both the 83.26 original and the 83.04 DMD-only baseline, so the quality cost is visible on the smaller model. The VBench margins in play sit between 0.1 and 1.4 points, near the benchmark's noise level, and the user study has only 24 participants.