SQuad: Sub-Quadratic Attention Distillation for Efficient Video Generation

Abstract: Video Diffusion Transformers (DiTs) spend most of their compute inside the Self-Attention operation, whose cost grows quadratically, $\mathcal{O}(n<sup>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<sup>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$.

• The paper introduces SQuad, which combines local and global softmax attention in two re-indexed passes to provide full token-to-token receptive fields at O(n√n) complexity.
• The paper uses flow-matching fine-tuning followed by DMD2 distribution matching to retrofit pretrained video diffusion models, achieving 83.20 VBench versus 83.08 for the teacher at 6 NFEs.
• The paper reports approximately 67× lower attention FLOPs, 2× faster end-to-end inference, zero added parameters, and compatibility with standard PyTorch compilation without custom kernels.

Motivation and problem setting

Video Diffusion Transformers (DiTs) spend most of their compute in softmax Self-Attention, whose cost grows as $\mathcal{O}(n^2)$ in the number of latent tokens $n$. At video-generation resolutions—Wan 2.2 5B produces $n = 18480$ tokens at $81 \times 704 \times 1280$—this quadratic term dominates runtime and memory, capping resolution and clip duration. Existing remedies fall into two camps, each with a drawback: linear and low-rank surrogates ($\mathcal{O}(n)$ or $\mathcal{O}(nk)$) replace the softmax with kernelized approximations that rarely recover full expressivity, while hybrid architectures interleave quadratic and cheap blocks but retain a dependence on full softmax attention and add heterogeneous complexity and extra parameters.

SQuad takes a different position: it retains genuine softmax attention throughout but lowers its complexity to $\mathcal{O}(n\sqrt{n})$, and fits a pretrained quadratic DiT into this modified architecture by distillation rather than training from scratch. The method is evaluated on Wan 2.2 5B and Wan 2.1 1.3B text-to-video models.

The SQuad-Attention operator

SQuad replaces the self-attention sub-layer with the composition of two ordinary softmax attentions, differing only in how the projected $Q$, $K$, $V$ tensors are re-indexed before the head split:

• Local pass: tokens attend within windows of size $n$0, yielding $n$1 parallel attentions of length $n$2.
• Global pass: tokens at the same slot across windows attend over all $n$3 windows, yielding $n$4 parallel attentions of length $n$5.

Both rearrangements are pure, parameter-free, invertible einops-style re-indexings. Crucially, the local output enters the global pass in the value slot—the value stream acts as a shared intermediary, and the passes are neither fused nor pooled. The supplementary proves that this composition yields a full receptive field in a single layer: the effective weight from any source to any target token is a product of two strictly positive softmax weights, $n$6, forming a valid distribution over all $n$7 tokens via a two-hop path. This is what makes SQuad a drop-in replacement rather than a local approximation, and it is the property that permits fitting a pretrained model by distillation alone.

The cost of the two passes is $n$8, which is minimized at $n$9, giving $n = 18480$0 complexity. Both extremes ($n = 18480$1 and $n = 18480$2) recover full quadratic attention through one pass, so the interesting regime is strictly interior.

Two-stage distillation

Since replacing attention changes each block's function, the pretrained weights no longer fit. The authors distill in two stages: Stage 1 is flow-matching supervised fine-tuning (8k iterations on VIPE 1M videos with Qwen3-8B-generated captions), which re-seats the network under the new attention; Stage 2 applies DMD2 distribution-matching distillation (15–30k iterations, captions only), which simultaneously recovers teacher quality and compresses sampling to 6 NFEs, with classifier-free guidance distilled into the student.

The ablation at 30 converted blocks is unambiguous: SFT alone collapses to a VBench total of 73.03, DMD2 alone reaches 80.91, and only the combination recovers 82.99. The authors note that a 20-block ablation makes the stages look interchangeable (all rows within 0.15 points), so the necessity of both stages scales with the depth of the attention replacement—an important caveat for anyone transferring the recipe.

Quality results

On Wan 2.2 5B, the 30-block SQuad model achieves a VBench total of 83.20 versus 83.08 for the original full-attention teacher, with zero added parameters. In a human preference study (1,179 paired comparisons from 24 evaluators over 648 prompts), 35% of users preferred SQuad over Radial Attention and 31% could not distinguish them. The method also transfers to Wan 2.1 1.3B (82.70 total), demonstrating backbone independence.

Among competing efficient-attention methods, Radial Attention posts a higher VBench total (84.56) and trained VSA/Jenga also score well; SQuad's claim rests not on quality superiority but on matching the teacher while achieving the best measured latency and FLOPs without specialized kernels.

Efficiency and hardware results

The efficiency gains are substantial and scale as expected for a sub-quadratic operator:

Setting	Original latency	SQuad latency	Attention FLOPs reduction
Wan 2.2 5B (per block)	62.01 ms	19.04 ms	~67×
Wan 2.1 1.3B (per block)	88.36 ms	15.22 ms	~90×
Wan 2.1 14B (per block)	300.29 ms	59.24 ms	~88×

At the headline setting, end-to-end single-forward DiT latency is 520 ms eager and 314 ms compiled, versus 870/667 ms for the original—2× faster end-to-end, and faster than Jenga (680/427 ms), the closest baseline. Notably, two baselines (Attention Surgery at 1006 ms, ReHyAt at 1757 ms) are slower than the full-attention original in eager mode despite lower FLOP counts, because their savings depend on hand-tuned custom kernels; SQuad's two plain softmax passes separated by pure re-indexing require no custom CUDA and benefit fully from stock torch.compile.

A FLOPs-scaling sweep across ~30× in sequence length confirms the theory quantitatively: measured costs agree with the closed-form prediction to within 0.7% on average, and the attention-FLOPs reduction widens with sequence length (e.g., 67× → 129× on the 5B backbone as $n = 18480$3 grows from 18480 to 73920).

Ablations

Three design choices are characterized. Ordering: local→global is best (83.20 vs. 82.99 for global→local, 82.87 for alternating). Single passes fail at full depth: local-only or global-only attention collapses to ~62.6 VBench total when all 30 blocks are converted, despite being cheaper—confirming that the composition, not either pass alone, carries the receptive field. Window geometry: at a fixed budget near $n = 18480$4, temporal windows ($n = 18480$5) train best (82.99); the authors attribute this tentatively to the asymmetric 1×-temporal/2×-spatial patchification but state that further analysis is needed. A scaling rule that keeps windows spanning the full temporal extent with aspect-ratio-preserving spatial extents independently recovers this trained window at native resolutions.

Limitations and open questions

The paper is candid about several constraints. The distillation strategy is entangled with step reduction: a purely flow-matching variant without DMD2 remains unexplored and would isolate the attention change from the sampling speedup. The measured SQuad FLOPs sit ~8% above the idealized $n = 18480$6 optimum, a gap the authors attribute entirely to window integrality rather than any missing cost term. The preference for temporal windows lacks a mechanistic explanation. Finally, the composition of more than two passes—potentially trading depth for further cost reduction—is hypothesized but untested, and generalization beyond video to other modalities is left open.

Conclusion

SQuad demonstrates that a pretrained quadratic video DiT can be retrofitted to a fixed, structured, $n = 18480$7 softmax attention pattern via two-stage distillation, matching the teacher on VBench (83.20 vs. 83.08) while cutting attention FLOPs by ~67× and end-to-end latency by 2×, at 6 NFEs, with zero added parameters and no custom kernels. Its broader claim—that full quadratic token-to-token communication may be unnecessary, and that softmax-preserving factorizations are an under-explored middle ground between linear surrogates and sparse-mask hybrids—is supported by the receptive-field proof and the quantitative conformance of measured FLOPs to theory, though whether such factorizations suffice as a pretraining architecture, rather than a distillation target, remains an open question.