ArXiv: 2602.12675

🎯 Pitch

A properly designed learnable router for sparse-linear attention can slash video diffusion computation to just 3% of original operations and deliver an 18.6× speedup, while actually improving generation quality metrics over the full attention baseline—proving that extreme attention sparsity need not degrade output when the decomposition is mathematically exact and the routing is trained end-to-end.


1. Executive Summary

This paper proposes SLA2, a trainable sparse-linear attention method for accelerating video diffusion models that addresses two limitations of the prior SLA approach: a heuristic routing mechanism based on attention-weight magnitude and a mathematical mismatch between SLA’s output and the original sparse-linear decomposition. SLA2 introduces a learnable router that dynamically predicts an optimal binary mask (using learned projections of pooled Q and K tensors followed by a differentiable Top-k selection during training) and a decomposition-consistent mixing formulation that directly learns a per-query ratio α to combine the sparse and linear attention branches, eliminating the need for a projection layer to compensate for scaling errors. Experiments on the Wan2.1-T2V video diffusion models (1.3B and 14B parameter variants) demonstrate that SLA2 achieves 97% attention sparsity and an 18.6× attention speedup while surpassing full attention on multiple VBench quality metrics, establishing that a properly formulated sparse-linear decomposition with learned routing can preserve generation quality even under extreme sparsity levels.

2. Context and Motivation

The Core Problem: How to Make Attention in Video Diffusion Models Efficient Without Breaking Generation Quality

Video diffusion models are extraordinarily computationally expensive. Each video frame requires processing a chain of Transformer attention layers, and the standard softmax attention mechanism scales quadratically with sequence length — O(N2d)O(N^2d) where NN is the number of tokens and dd is the head dimension. For video generation, NN is enormous: you have tokens from all spatial positions across all frames, making the attention cost the dominant bottleneck. A single video generation pass on a model like Wan2.1-14B-720P requires approximately 292.6 trillion FLOPs just for attention operations (Table 1), and the end-to-end latency exceeds 42 minutes (Figure 5b: 2,550 seconds for the original full-attention model before excluding offloading overhead).

This matters for two practical reasons. First, the computational cost makes real-time or interactive video generation infeasible — a user waiting 40+ minutes for a short video clip is far from an interactive experience. Second, the cost makes these models inaccessible to researchers and practitioners without access to large GPU clusters, concentrating capability in a few well-resourced organizations. The attention mechanism is the primary bottleneck, so accelerating attention is the most direct path to making video generation practical.

The broader challenge is that you cannot simply throw away the attention mechanism without destroying the model's ability to generate coherent video. Attention is what allows the model to model long-range dependencies — a character's hair in frame 1 needs to affect how it moves in frame 32, and the lighting on an object in the top-left corner of the frame needs to be consistent with the light source in the bottom-right. Sparsifying attention means selectively removing these computations, and doing so requires answering a difficult question: which attention computations are essential for generation quality, and which can be safely skipped or approximated?

The Decomposition Perspective: Why Sparse + Linear Makes Theoretical Sense

The theoretical motivation for combining sparse and linear attention — which both SLA and SLA2 inherit — comes from a property of attention matrices that the paper describes in Section 2.2. For any full-attention probability matrix P=softmax(QKT/d)P = \text{softmax}(QK^T / \sqrt{d}), we can conceptually decompose it into two components:

P=P1+P2P = P_1 + P_2

where P1P_1 is a sparse component capturing the large-magnitude, structure-specific attention patterns, and P2P_2 is a low-rank component capturing the diffuse, global patterns. The idea is that P1P_1 — the few large attention weights that dominate the softmax distribution — can be computed exactly via sparse softmax attention (only computing attention at the selected positions). Meanwhile, P2P_2 — the many small attention weights that collectively matter but individually seem negligible — can be efficiently approximated via linear attention, which has O(Nd2)O(Nd^2) complexity rather than O(N2d)O(N^2d).

This decomposition is theoretically appealing because it exploits two complementary properties of attention matrices in diffusion models:

  • Sparsity: At high sparsity levels (90%+), most attention weights are near-zero after softmax, meaning most of the O(N2)O(N^2) QK dot products produce values that contribute negligibly to the output. A sparse attention module that only computes the "important" ones can skip the rest.

  • Low-rank structure: The remaining attention weights, while individually small, collectively capture a low-rank signal that linear attention (based on the kernel trick: ϕ(Q)(ϕ(K)TV)\phi(Q)(\phi(K)^T V) rather than softmax(QKT)V\text{softmax}(QK^T)V) can approximate well.

The key insight is that these two properties are complementary: sparse attention alone struggles because it renormalizes probabilities within each row (Section 2.2, Equation 8), and the normalization step introduces a scaling mismatch — the sparse attention output Os=PsVO_s = P_s V where PsP_s sums to 1 per row, but the true contribution of the selected positions is αPsV\alpha \odot P_s V where α\alpha is the row-wise sum of probabilities on selected positions. Linear attention on the remaining positions can simultaneously compensate for unselected attention weights and, if properly formulated, correct this scaling mismatch.

SLA: The Prior State of the Art and Its Two Failures

SLA (Zhang et al., 2025c) was the first method to operationalize this sparse + linear decomposition for diffusion models with trainable components. Its design is described in Section 2.1, and the paper's critique identifies two specific limitations:

Limitation 1 (L1): Mathematical mismatch with the decomposition. SLA computes its output as:

O=Os+Proj(Ol)O = O_s + \text{Proj}(O_l)

where Os=PsVO_s = P_s V (sparse softmax attention on the mask-selected positions), OlO_l is the linear attention output, and ProjRd×d\text{Proj} \in \mathbb{R}^{d \times d} is a learnable linear projection. The problem, analyzed in Section 2.2, is that OsO_s is not aligned with P1VP_1 V — the actual contribution the sparse component should make. Sparse attention renormalizes its probabilities to sum to 1 within the mask, so PsP_s does not equal P1P_1. Instead, P1=αPsP_1 = \alpha \odot P_s, where α\alpha is a vector of row-wise probability sums on the masked positions (Equation 7–9). The output mismatch becomes:

P1VOs=(α1)OsP_1 V - O_s = (\alpha - 1) \odot O_s

This means the linear attention branch in SLA must handle two tasks simultaneously: approximate the true linear component P2VP_2 V AND compensate for the sparse branch's scaling error (α1)Os(\alpha - 1) \odot O_s (Equation 10). This is a harder learning problem — the linear attention's projection layer is being asked to correct a multiplicative scaling error from the sparse branch, which is a fundamentally different operation from approximating missing attention mass. The authors argue this formulation "makes the compensation harder to learn" (Section 2.2) and prevents the linear branch from specializing in what it was designed for.

Limitation 2 (L2): Heuristic, non-optimized routing. SLA decides which attention positions go to the sparse branch vs. the linear branch using a simple heuristic (Section 2.1): it computes compressed attention weights PcP_c from pooled Q and K, then sets the mask MM so that the top kh%k_h\% of entries per row go to sparse attention and the bottom kl%k_l\% are skipped entirely. This is a magnitude-based split — positions with large attention weights go to sparse attention, positions with small weights go to linear attention.

The paper argues this is suboptimal through a telling counterexample (Section 1, L2): "moving some weights from P1P_1 to P2P_2 via brute-force selection may not increase the rank of P2P_2, while still improving the sparsity of P1P_1." In other words, the optimal routing isn't simply about which weights are large — it's about which assignment makes P1P_1 maximally sparse while keeping P2P_2 maximally easy for linear attention to approximate. A weight might be moderately large but highly correlated with other weights (low-rank), so moving it to P2P_2 doesn't hurt the linear approximation and helps sparsity. Conversely, a weight might be small but structurally important (breaking the low-rank assumption), so keeping it in P1P_1 is essential. The magnitude heuristic ignores this structure-rank tradeoff entirely.

The heuristic also operates on static values (QˉKˉT\bar{Q}\bar{K}^T) rather than learned representations, meaning it cannot adapt to the specific approximation characteristics of the deployed linear attention module. The quality of the split depends entirely on how well pooled dot products correlate with what makes a good sparse vs. linear assignment, which is an assumption the paper challenges.

Where Prior Approaches Fall Short More Broadly

Beyond SLA specifically, the paper's introduction and related work section (Section 10) establish a landscape where existing sparse attention methods have structural limitations:

Training-free sparse attention methods (Xiao et al., 2024; Jiang et al., 2024; Gao et al., 2024; Zhang et al., 2025f; and many others) apply hand-designed sparsity patterns at inference time — fixed local windows, attention sinks, or magnitude-based pruning — without any training. These are easy to deploy (no fine-tuning required) but are fundamentally limited in achievable sparsity because they cannot adapt the model's attention patterns to the sparsification. The model was trained with full attention; forcing it to operate with missing computations inevitably degrades quality above some sparsity threshold. The paper implicitly argues that training with sparsity (as SLA2 does) allows the model to learn to concentrate its attention into the selected positions, enabling much higher sparsity.

Trainable sparse attention methods (VSA, VMoBA, SLA, and others) fine-tune the model to work with sparse patterns, achieving higher sparsity than training-free approaches. However, each has specific weaknesses:

  • VSA (Zhang et al., 2025i) uses a trainable sparse mask but does not incorporate a linear attention branch to compensate for dropped computations. Table 1 shows VSA quality degrades sharply at high sparsity: at 95% sparsity on Wan2.1-1.3B, VSA's Imaging Quality drops to 55.50 (vs. 67.04 for SLA2) and its Vision Reward goes negative (−0.1309 vs. +0.1023 for SLA2).

  • VMoBA (Wu et al., 2025) uses mixture-of-block attention, which is a different form of structured sparsity, but also lacks a linear compensation branch. At 95% sparsity, VMoBA shows significant quality degradation on the 14B model (OC drops to 7.96, demonstrating catastrophic failure on overall consistency).

  • SLA incorporates linear attention but, as analyzed above, has the formulation and routing problems that motivated SLA2.

Linear attention alone is insufficient for video. The related work section (Section 10) explicitly notes that "for video generation, linear attention alone often cannot keep quality." While linear attention works for image generation pretraining (SANA, DiG), video requires capturing precise dependencies across time, and the kernel approximation's error accumulates destructively across the temporal dimension. The paper positions SLA2 as a hybrid that gets the best of both: sparse attention for precise, structure-specific dependencies and linear attention for diffuse, global patterns — but only when the routing between them is properly optimized.

The Additional Opportunity: Low-Bit Quantization of the Sparse Branch

Once attention is decomposed into sparse and linear branches, the sparse branch itself still performs matrix multiplications — just fewer of them. These remaining operations can be further accelerated by using low-bit precision (INT8 or FP8) for the QK dot products and PV multiplications within the sparse branch. This is where quantization-aware training (QAT, Section 5) enters the picture.

The paper identifies that post-training quantization (PTQ) — applying low-bit quantization after fine-tuning is complete — introduces an accuracy penalty because the model was never trained to tolerate quantization error. QAT, which simulates quantization during the forward pass of training (while keeping the backward pass in FP16), allows the model's parameters to adapt to quantized computation, minimizing the accuracy gap between full-precision and low-bit inference. The ablation in Table 2 confirms this: training without QAT ("w/o QAT" at 97% sparsity) drops the Vision Reward from 0.1039 (with QAT) to 0.0850, confirming that the quantization error from low-bit inference is non-trivial without training-time adaptation.

The quantization is applied only to the sparse branch (Section 5), meaning the linear attention branch operates at full precision. This is sensible: linear attention already achieves substantial speedup through its algorithmic reformulation, while sparse attention's remaining dense matmuls are the bottleneck that low-bit compute targets.

How SLA2 Positions Itself

SLA2 is not a radical departure from SLA but a principled refinement that addresses the two identified limitations through a unified framework (Section 3, Equation 13):

The reformulated combination (addressing L1):

O=αOs+(1α)OlO = \alpha \odot O_s + (1 - \alpha) \odot O_l

where αRN×1\alpha \in \mathbb{R}^{N \times 1} is a learnable per-query mixing ratio. This directly matches the decomposition PαPs+(1α)PlP \approx \alpha \odot P_s + (1 - \alpha) \odot P_l (Equation 11–12), where both PsP_s and PlP_l are row-stochastic (each row sums to 1), and α\alpha controls their mix. Critically, this formulation eliminates the need for SLA's projection layer Proj(Ol)\text{Proj}(O_l) because the scaling mismatch is handled explicitly by α\alpha rather than being pushed into the linear branch as a side task. The linear attention branch only needs to approximate the normalized PlP_l matrix, not simultaneously correct the sparse branch's scaling error.

The learnable router (addressing L2): The router RR in Equation 14, detailed in Section 4, takes QQ and KK as inputs, applies learned projections projq,projkRd×d\text{proj}_q, \text{proj}_k \in \mathbb{R}^{d \times d} to pooled representations, computes Pc=projq(Qˉ)projk(Kˉ)TP_c = \text{proj}_q(\bar{Q})\text{proj}_k(\bar{K})^T, and selects the top k%k\% positions via a differentiable SoftTop-k operator during training. This is fundamentally different from SLA's heuristic: the projections are trained to produce a PcP_c where Top-k selection produces a mask that minimizes the reconstruction error between full attention output and the SLA2 approximation. Section 8, question (1.c) explains the generalization: setting projq=projk=I\text{proj}_q = \text{proj}_k = I recovers the heuristic, but learning these projections under the training objective (Algorithm 1, Stage 1) can find transformations that better separate positions by their structural importance to the sparse branch vs. how well the remaining positions can be approximated by linear attention.

Two-stage training (Section 6): The paper positions its training strategy as necessary for two reasons. Stage 1 trains only RR and α\alpha using a mean-squared-error loss against full attention outputs, which provides a good initialization for the router before end-to-end fine-tuning — otherwise "unstable and poor routing can make subsequent fine-tuning difficult" (Section 6). Stage 1 also uses the differentiable SoftTop-k (Equation 17) to enable gradient flow through the router, since hard Top-k blocks gradients. Stage 2 fine-tunes the full diffusion model end-to-end with hard Top-k routing (matching inference-time behavior), optimizing the diffusion loss to adapt the model's attention patterns to high sparsity. The authors note an interesting phenomenon: at high sparsity, fine-tuning can even improve quality over full attention (Table 1 shows SLA2 outperforming Full Attention on many metrics), which they attribute to the higher quality of their fine-tuning dataset compared to the pretraining data.

The paper's positioning is thus: SLA2 is not claiming to invent sparse+linear attention from scratch, nor is it claiming to invent learnable routing in general. It is claiming that the specific combination — a decomposition-consistent mixing formulation with a learned, optimization-aware router, trained with differentiable Top-k and then fine-tuned end-to-end — resolves the documented failures of SLA and pushes achievable sparsity and quality beyond what prior hybrid methods could achieve. The results at 97% sparsity (Table 1) — where SLA2 outperforms all baselines at 90% sparsity — are the empirical validation of this claim.

3. Technical Approach

3.1 Reader Orientation

SLA2 is a trainable replacement for the standard attention mechanism in Transformer-based video diffusion models — it replaces the quadratic-cost softmax attention with a hybrid system that computes exact attention on a small, learned subset of positions and approximates the rest using efficient linear attention. The problem it solves is how to make attention in video generation computationally tractable (reducing FLOPs by ~97%) while preserving or even improving generation quality, by learning where attention should be precise and where approximation suffices, rather than relying on hand-crafted heuristics.

3.2 Big-Picture Architecture (Diagram in Words)

The SLA2 attention module replaces standard softmax attention with a pipeline of five interconnected components:

  1. Learnable Router $R$ — takes query and key tensors as input, applies learned linear projections to pooled representations, performs a differentiable Top-k selection during training (hard Top-k during inference), and outputs a binary mask $M$ indicating which attention positions should be computed exactly by the sparse branch vs. approximated by the linear branch.

  2. Sparse Attention Branch — computes exact softmax attention only on the positions where $M = 1$ (the positions the router assigned to sparse computation), producing output $O_s$. During inference, this branch can optionally run with low-bit quantization (INT8/FP8) for additional speedup, with quantization-aware training (QAT) used during fine-tuning to adapt the model to quantized computation.

  3. Linear Attention Branch — approximates the attention for positions where $M = 0$ using a kernelized linear attention formulation $\text{norm}(\phi(Q)\phi(K)^T \odot (1 - M)) V$, producing output $O_l$ in $O(Nd^2)$ complexity rather than $O(N^2d)$.

  4. Learnable Mixing Ratio $\alpha$ — a learned per-query vector (with values between 0 and 1) that controls how much of each branch's output contributes to the final attention output. Specifically, the final output is $O = \alpha \odot O_s + (1 - \alpha) \odot O_l$, ensuring that the combination is properly normalized (each row's mixing weights sum to 1) and that the sparse branch's contribution is scaled by the actual probability mass it represents.

  5. Two-Stage Training Pipeline — Stage 1 initializes the router and mixing ratio by minimizing the MSE between full attention output and the SLA2 approximation on a static dataset of Q, K, V tensors (using differentiable SoftTop-k to backpropagate through the router). Stage 2 fine-tunes the full diffusion model end-to-end with hard Top-k routing (matching inference behavior), optimizing the standard diffusion loss to adapt the model's attention patterns to high sparsity.

The flow at inference time is: Q, K, V tensors enter → the router $R$ processes pooled Q and K through learned projections and applies hard Top-k to produce mask $M$ → sparse branch computes $O_s$ via quantized FlashAttention on $M = 1$ positions → linear branch computes $O_l$ via kernelized attention on $M = 0$ positions → the learnable ratio $\alpha$ mixes them to produce the final output $O$.

3.3 Roadmap for the Deep Dive

  • First, the core mathematical formulation (Equation 13) and why it resolves SLA's mismatch — because understanding the learned mixing ratio $\alpha$ is foundational to everything else, and the paper's central claim is that this reformulation matters.
  • Second, the learnable router $R$ — how it produces masks, why it uses pooled representations with learned projections, and how the differentiable SoftTop-k operator enables gradient-based training of a discrete selection problem.
  • Third, the two-stage training procedure — how Stage 1 initializes the router using MSE loss and Stage 2 fine-tunes end-to-end with hard Top-k, including the specific training hyperparameters, dataset construction, and the transition from SoftTop-k to hard Top-k.
  • Fourth, the forward pass algorithm (Algorithm 2) and backward pass (Algorithm 3) — explaining the block-wise computation pattern, how the mask gates the sparse vs. linear branches at the kernel level, and the gradient flow.
  • Fifth, quantization-aware training — how it integrates with the two-stage training, what is quantized (Q, K, P, V in the sparse branch only), why the backward pass stays in FP16, and the SageAttention2++ quantization scheme it builds on.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a methods paper whose core idea is that the prior SLA approach has two correctable flaws — a mathematical mismatch in how sparse and linear outputs are combined, and a heuristic routing mechanism — and that fixing both through a learnable router and a decomposition-consistent mixing formulation enables substantially higher sparsity without quality degradation.


The Core Reformulation: Decomposition-Consistent Mixing

The paper's central mathematical contribution is replacing SLA's output formulation $O = O_s + \text{Proj}(O_l)$ with a mixing formulation that directly matches the sparse-linear decomposition of attention probabilities. To understand why this matters, we must first precisely characterize the mismatch SLA introduces.

What the decomposition actually says. Let $P = \text{softmax}(QK^T / \sqrt{d}) \in \mathbb{R}^{N \times N}$ be the full-attention probability matrix — each row is a proper probability distribution summing to 1. Given a binary mask $M \in \{0, 1\}^{N \times N}$, we decompose $P$ into a component on the masked positions and a component on the remaining positions:

P=P1+P2,P1=PM,P2=P(1M)P = P_1 + P_2, \quad P_1 = P \odot M, \quad P_2 = P \odot (1 - M)

where $\odot$ is element-wise multiplication. The full attention output is then:

Of=PV=P1V+P2VO_f = P V = P_1 V + P_2 V

This says: the output is the sum of contributions from masked-attention positions and from the remaining positions.

Why sparse attention does NOT compute $P_1 V$. Sparse attention takes the masked probabilities $P_1$ and renormalizes them to sum to 1 within each row, producing a row-stochastic matrix $P_s$. Formally, let $\alpha \in \mathbb{R}^{N \times 1}$ be the row-wise sum of probabilities on the masked positions:

α=P11\alpha = P_1 \mathbf{1}

where $\mathbf{1} \in \mathbb{R}^{N \times 1}$ is an all-ones vector. Then the sparse attention distribution is:

Ps=P1αP_s = \frac{P_1}{\alpha}

where the division is row-wise (each element in row $i$ is divided by $\alpha_i$). The sparse attention output is $O_s = P_s V$. But what the decomposition actually requires is $P_1 V$, which is:

P1V=(αPs)V=α(PsV)=αOsP_1 V = (\alpha \odot P_s) V = \alpha \odot (P_s V) = \alpha \odot O_s

What this means operationally: the sparse branch computes the correct weighted combination of value vectors for the selected positions, but it then divides each row by $\alpha$, destroying the correct scaling. To recover $P_1 V$, you need to multiply each row $i$ of $O_s$ by $\alpha_i$ — the probability mass that row actually carries.

How SLA handles (or fails to handle) this. SLA's output is $O = O_s + \text{Proj}(O_l)$. Comparing with the decomposition, we have:

Proj(Ol)P2V+(α1)Os\text{Proj}(O_l) \approx P_2 V + (\alpha - 1) \odot O_s

because $O_f = P_1 V + P_2 V = \alpha \odot O_s + P_2 V$, and SLA's output is $O_s + \text{Proj}(O_l)$. Rearranging: $O_s + \text{Proj}(O_l) \approx \alpha \odot O_s + P_2 V$ implies $\text{Proj}(O_l) \approx P_2 V + (\alpha - 1) \odot O_s$.

Why this is problematic: the learnable projection $\text{Proj}(\cdot)$ applied to the linear attention output $O_l$ must simultaneously accomplish two very different tasks:

  1. Approximate the true linear component $P_2 V$ — the contribution from positions not in the mask.
  2. Correct the scaling error $(\alpha - 1) \odot O_s$ from the sparse branch — a multiplicative correction that depends on $O_s$, not $O_l$.

The second task is particularly perverse because it means the linear branch's projection must learn to compensate for errors that originate in the sparse branch, using only information available in $O_l$ (which is computed from completely different positions). This is a fundamentally harder learning problem than having each branch handle its own scaling correctly.

SLA2's reformulation. The paper proposes the following output:

O=αOs+(1α)OlO = \alpha \odot O_s + (1 - \alpha) \odot O_l

where $\alpha \in \mathbb{R}^{N \times 1}$ is a learnable vector with values between 0 and 1, $O_s = \text{softmax}(QK^T/\sqrt{d} \odot M) V$ is the sparse attention output (row-normalized), and $O_l = \text{norm}(\phi(Q)\phi(K)^T \odot (1 - M)) V$ is the linear attention output (also row-normalized such that each row sums to 1).

What this equation computes: for each query position $i$, the final output $O_i$ is a convex combination of the sparse branch output $(O_s)_i$ and the linear branch output $(O_l)_i$. The mixing weight $\alpha_i$ controls how much the sparse branch contributes: if $\alpha_i = 0.7$, then 70% of the output comes from sparse attention and 30% from linear attention. The decomposition $\alpha \odot P_s + (1 - \alpha) \odot P_l$ is provably row-normalized because $P_s$ and $P_l$ each sum to 1 per row, and $\alpha + (1 - \alpha) = 1$.

Why this form: it directly matches the desired decomposition $P \approx \alpha \odot P_s + (1 - \alpha) \odot P_l$ (Equation 11). The coefficient $\alpha$ explicitly controls the contribution of the sparse branch, eliminating the need for a projection layer to fix the scaling mismatch. The linear attention branch $O_l$ only needs to approximate the normalized linear component $P_l$, not simultaneously correct the sparse branch's scaling error. This makes each branch's learning objective cleaner: the sparse branch focuses on producing accurate attention for the selected positions, and the linear branch focuses on approximating a low-rank version of the remaining distribution — exactly what linear attention is designed for.

The $\alpha$ vector is learned during training (both Stage 1 initialization and Stage 2 fine-tuning), meaning the model can discover per-query mixing ratios that depend on the specific attention patterns in each head and layer. A head that attends in a highly concentrated way might learn large $\alpha$ (trust the sparse branch), while a head with diffuse attention might learn small $\alpha$ (rely more on the linear branch).


The Learnable Router

The router $R$ is the component that decides which attention positions go to which branch. Formally, it produces the binary mask $M$:

M=R(Q,K){0,1}N×NM = R(Q, K) \in \{0, 1\}^{N \times N}

where $M_{ij} = 1$ means position $(i, j)$ is computed by the sparse branch, and $M_{ij} = 0$ means it goes to the linear branch. The paper's insight is that this routing decision should be learned and optimized for the specific approximation characteristics of the deployed branches, rather than determined by a heuristic.

Why Q and K (only) as inputs. The router takes only $Q$ and $K$ as inputs, not $V$. Section 8, question (1.a) explains: the attention matrix $P = \text{softmax}(QK^T / \sqrt{d})$ is entirely determined by the dot products between queries and keys. The values $V$ don't affect which positions receive large attention weights — they only affect what information is aggregated from those positions. Therefore, $Q$ and $K$ contain all the information needed to decide which positions are structurally important (should be computed exactly) vs. which positions contribute diffuse, low-rank signal (can be approximated). Including $V$ would add unnecessary computation without providing routing-relevant information.

Why pooling is applied before routing. A naive router that computes a $QK^T$-style score matrix over all $N \times N$ positions would itself incur $O(N^2)$ cost, defeating the purpose of acceleration. The paper exploits the empirical observation (drawn from Jiang et al., 2024; Zhang et al., 2025f; Gao et al., 2024) that "nearby tokens in diffusion transformers often have similar distributions" (Section 8, question 1.b), meaning attention scores vary smoothly across adjacent spatial/temporal positions. This smoothness means you can make routing decisions at a coarser granularity without much loss of optimality.

Specifically, the router pools $Q$ and $K$ along the token dimension using mean pooling with block sizes $b_q$ and $b_k$:

Qˉ=pool(Q)RN/bq×d,Kˉ=pool(K)RN/bk×d\bar{Q} = \text{pool}(Q) \in \mathbb{R}^{N/b_q \times d}, \quad \bar{K} = \text{pool}(K) \in \mathbb{R}^{N/b_k \times d}

The paper uses $b_q = 128$ and $b_{kv} = 64$ (Section 9.1), meaning the routing operates on blocks of 128 query tokens and 64 key tokens. This reduces the routing cost from $O(N^2)$ to $O((N/b_q)(N/b_k)) = O(N^2 / (b_q b_k))$, a factor of $128 \times 64 = 8192$ reduction in the score matrix size that the router must process.

Why learned projections are necessary. The next step is computing a routing score matrix from the pooled representations. A simple heuristic — like SLA's approach — would compute $\text{softmax}(\bar{Q} \bar{K}^T / \sqrt{d})$ and select the top $k\%$ positions. But this assumes that the raw dot product $\bar{Q}_i \cdot \bar{K}_j$ is the right signal for determining whether position $(i, j)$ belongs in the sparse or linear branch.

The paper argues this assumption is not necessarily true. The optimal routing depends not just on the magnitude of attention weights, but on the structure — specifically, which assignment makes $P_1$ maximally sparse while making $P_2$ as easy as possible for linear attention to approximate (i.e., as low-rank as possible). A position with a moderate dot product might be highly correlated with other positions (low-rank), so moving it to $P_2$ doesn't hurt the linear approximation but significantly helps sparsity. Conversely, a position with a small dot product might capture unique structural information that linear attention cannot recover, so it should stay in $P_1$.

To enable the router to learn these structure-aware decisions, the paper introduces two learnable linear projections $\text{proj}_q, \text{proj}_k \in \mathbb{R}^{d \times d}$:

Pc=projq(Qˉ)projk(Kˉ)TP_c = \text{proj}_q(\bar{Q}) \text{proj}_k(\bar{K})^T

These are $d \times d$ matrices that transform the pooled representations into a space where Top-k selection produces an optimal mask. This generalizes the heuristic: setting $\text{proj}_q = \text{proj}_k = I$ recovers the raw dot product (Section 8, 1.c). Learning these projections under the training objective (minimizing the error between full attention and the SLA2 approximation) allows the router to discover transformations that better separate positions by their structural role.

The Top-k selection. From $P_c \in \mathbb{R}^{N/b_q \times N/b_k}$, the router applies row-wise Top-k to produce a compressed mask:

Mc=Top-k(k%,Pc){0,1}N/bq×N/bkM_c = \text{Top-k}(k\%, P_c) \in \{0, 1\}^{N/b_q \times N/b_k}

where each row has exactly $k\% \times N/b_k$ entries set to 1 (these blocks go to the sparse branch) and the rest to 0 (these blocks go to the linear branch). For the experiments, the paper sweeps $k\%$ values of 5%, 4%, and 3%, corresponding to sparsity levels of 95%, 96%, and 97% (the remaining $(1 - k\%)$ positions use linear attention, but since linear attention is $O(Nd^2)$, the effective computation savings are slightly less — the paper reports these as 95%, 96%, and 97% sparsity in the results).

The compressed mask $M_c$ is then expanded to full size $M \in \{0, 1\}^{N \times N}$: each entry $M_c[i, j]$ corresponds to a block of size $b_q \times b_k$ in the full mask, and all positions in that block inherit the same routing decision. This block-wise routing is what enables the efficient GPU kernel implementation — the kernel processes blocks at a time, checking $M_c[i, j]$ to decide whether to execute the sparse or linear branch for that block (Algorithm 2, lines 12–21).

A critical detail: the router outputs are only needed at the compressed granularity. Section 4 notes that "in practice, our forward and backward GPU kernels for SLA2 only require $M_c$, since we implement the method efficiently on top of a block-wise FlashAttention-style algorithm." This means the router never needs to produce or store the $N \times N$ mask — only the $(N/b_q) \times (N/b_k)$ compressed mask is materialized, and the GPU kernel implicitly expands it during computation by processing full blocks of queries and keys together.


The Differentiable SoftTop-k Operator (Training-Time Router)

The Top-k selection $M_c = \text{Top-k}(k\%, P_c)$ is a discrete, non-differentiable operation — it produces binary values, and there is no meaningful gradient from a 0/1 output back to the continuous scores $P_c$ that determine which positions are selected. This is problematic because the router's projections $\text{proj}_q$ and $\text{proj}_k$ need gradient signals to learn: the model needs to know whether slightly increasing a routing score for a particular block-pair would improve the final attention approximation.

To enable backpropagation through the router during Stage 1 training (where the router is initialized), the paper replaces hard Top-k with a differentiable relaxation called SoftTop-k (Ding et al., 2024):

SoftTop-k(k%,Pc)ij=σ((Pc)ijτ+λi)\text{SoftTop-k}(k\%, P_c)_{ij} = \sigma\left(\frac{(P_c)_{ij}}{\tau} + \lambda_i\right)

where $\sigma(\cdot)$ is the sigmoid function, $\tau$ is a temperature parameter (set to $\tau = 0.1$, Section 9.1), and $\lambda_i$ is a row-specific scalar.

What this computes: for each row $i$ of the routing score matrix $P_c$, the SoftTop-k operator produces a continuous mask (values in $(0, 1)$ rather than $\{0, 1\}$) where entry $(i, j)$ receives a probability-like value that depends on the scaled score $(P_c)_{ij} / \tau + \lambda_i$ passed through a sigmoid. The parameter $\lambda_i$ is solved for (via binary search) such that each row sums to exactly $k\% \times N/b_k$ — it acts as a threshold that adjusts automatically to enforce the sparsity constraint.

How $\lambda_i$ is determined: for each row, a binary search finds the $\lambda$ value such that $\sum_j \sigma((P_c)_{ij} / \tau + \lambda) = k\% \times N/b_k$. This is possible because sigmoid is monotonic in $\lambda$: larger $\lambda$ shifts more entries toward 1, increasing the sum. The binary search finds the cutoff that achieves exactly the target sum, analogous to how hard Top-k would select the exact top $k\%$ entries.

Why the temperature $\tau$ matters: as $\tau \to 0$, sigmoid becomes a step function and SoftTop-k converges to hard Top-k (entries above the threshold are exactly 1, below are exactly 0). As $\tau$ increases, the mask becomes softer (entries close to the threshold receive intermediate values), providing smoother gradients at the cost of a looser approximation to the discrete selection. The paper uses $\tau = 0.1$, which provides a reasonable balance between gradient flow and approximation accuracy.

The reparameterization trick for gradient computation. The gradient of SoftTop-k with respect to $(P_c)_{ij}$ is computed using the reparameterization trick from Ding et al. (2024). Rather than differentiating through the binary search for $\lambda_i$ (which would be numerically unstable), the gradient treats $\lambda_i$ as an implicit function of $P_c$ and uses the implicit function theorem. The resulting gradient flows from the final loss through the sigmoid outputs back to the raw scores $(P_c)_{ij}$, training $\text{proj}_q$ and $\text{proj}_k$ to produce $P_c$ values where the (soft) selection pattern minimizes the attention approximation error.

At inference time and during Stage 2 training, hard Top-k is used. During Stage 2 fine-tuning and final inference, SoftTop-k is replaced with standard hard Top-k (Equation 16). This means the routing mask $M_c$ is strictly binary, and the forward pass exactly matches what will run at inference time. The gradient of hard Top-k is zero — the router's projections are NOT updated during Stage 2. Instead, Stage 2 fine-tunes only the diffusion model parameters $\Theta$ and the mixing ratio $\alpha$ (Algorithm 1, line 7: "Fine-tune $\Theta$, $\alpha$ using an end-to-end diffusion loss"). The router $R$ is frozen after Stage 1, which is why the paper describes Stage 1 as getting "a better initialization for $R$ and $\alpha$ to ensure stable and effective subsequent fine-tuning."


Training Procedure: Two-Stage Strategy

The paper adopts a two-stage training strategy (Section 6, Algorithm 1) motivated by two concerns: training stability and train-inference consistency.

Why two stages are necessary. If the router $R$ started from random initialization and was trained jointly with the full diffusion model end-to-end from the beginning, two problems would arise:

  1. Unstable routing early in training: a randomly initialized router would produce essentially arbitrary masks, causing the attention output to be a poor approximation of full attention. The diffusion model would need to simultaneously adapt to this noisy routing while also learning to generate high-quality videos — a difficult optimization problem that could lead to training instability or poor local minima.

  2. Gradient mismatch between training and inference: the differentiable SoftTop-k operator used during training produces continuous mask values, while inference uses discrete hard Top-k selection. If the model were trained exclusively with SoftTop-k, the learned routing might not transfer well to the discrete selection at inference time — the model would be optimized for a soft assignment that doesn't exactly match the deployed binary routing.

The two-stage strategy addresses both: Stage 1 initializes the router in isolation (without the full diffusion model), using SoftTop-k to get gradient flow, so that when Stage 2 begins, the router already produces reasonable masks. Stage 2 then uses hard Top-k (matching inference) to fine-tune the diffusion model so it learns to operate with exactly the routing scheme that will be deployed.

Stage 1: Router and mixing ratio initialization.

The goal of Stage 1 is to find initial values for the router's projections $\text{proj}_q$, $\text{proj}_k$, and the mixing ratio $\alpha$ that produce a good SLA2 approximation of full attention across various sparsity levels.

Training data for Stage 1. The paper constructs a static dataset $D$ by sampling Q, K, V tensors "from every attention layer at each diffusion timestep" (Section 6). This means: take the pretrained diffusion model (before any SLA2 fine-tuning), run it on training videos, and for each attention layer and each diffusion denoising step, collect the Q, K, V tensors that the model naturally produces. This dataset captures the joint distribution of (Q, K, V) that the model encounters during generation, providing representative examples for the router to learn from.

The Stage 1 training data does NOT require video labels or ground-truth outputs — it only requires the pretrained model's intermediate activations, which can be collected by running the model on unlabeled videos.

Stage 1 loss function. For each (Q, K, V) tuple in the dataset, the loss compares the SLA2 output (using the current router and mixing ratio) to the full attention output (computed with standard softmax attention on all positions):

L=MSE(FullAttn(Q,K,V),SLA2(Q,K,V,k%,R,α))\mathcal{L} = \text{MSE}(\text{FullAttn}(Q, K, V), \text{SLA2}(Q, K, V, k\%, R, \alpha))

where $R$ represents the router (projections $\text{proj}_q$, $\text{proj}_k$) and $\alpha$ is the mixing ratio vector. The mean squared error is computed element-wise over all $N \times d$ output positions.

What this loss computes: for each position in the output, it compares the full-attention output (which uses all $N^2$ QK dot products) to the SLA2 approximation (which uses only $k\% \times N^2$ exact computations plus linear attention on the rest). Minimizing this loss trains the router to select mask positions that minimize the reconstruction error of the attention output, and trains $\alpha$ to optimally blend the sparse and linear branches.

Why MSE rather than the diffusion loss in Stage 1: MSE provides a direct, per-example supervision signal that doesn't require generating full videos. It allows the router to learn from many (Q, K, V) examples quickly without the expense of full diffusion model fine-tuning. The Stage 1 optimization only involves the router parameters and $\alpha$ — a small fraction of the total model parameters — making it computationally lightweight.

Training under different sparsity levels. The paper trains $R$ and $\alpha$ "under different $k\%$" (Section 6) — specifically $k\%$ values of 5%, 4%, and 3%. This means the router learns to produce good masks at multiple sparsity levels, which is important because the optimal routing may differ depending on how aggressive the sparsity is. At 95% sparsity ($k\% = 5\%$), the router can be more selective, keeping only the most critical positions; at 97% sparsity ($k\% = 3\%$), it must compress even further while maintaining quality.

Stage 2: End-to-end diffusion model fine-tuning.

Stage 2 replaces all attention modules in the diffusion model with SLA2 and fine-tunes the entire model end-to-end using the standard diffusion loss (typically a noise prediction or score-matching loss). The key characteristics of Stage 2 are:

Hard Top-k for routing. During Stage 2, the router uses hard Top-k (Equation 16) rather than SoftTop-k. This means the routing mask is strictly binary, and there is no gradient flowing through the router — $R$'s parameters are frozen after Stage 1. The paper explicitly states that Stage 2 "directly optimize[s] the diffusion loss over all model parameters $\Theta$, including $\alpha$, without $R$" (Section 6).

What is optimized in Stage 2: all parameters of the diffusion model $\Theta$ (weights of the Transformer layers, including QKV projections, feed-forward networks, normalization layers, etc.) plus the mixing ratio $\alpha$. The router $R$ remains fixed with its Stage 1 initialization. The model learns to adapt its attention patterns to work well with the fixed routing mask — for example, it might learn to concentrate attention into the positions that the router selects, making the sparse branch more effective, or it might learn to produce value vectors that the linear branch can better approximate.

Fine-tuning hyperparameters (Section 9.1): the paper fine-tunes each method for 500 steps. The batch size is 64 for the 1.3B model and 15 for the 14B model (the smaller batch size for the larger model reflects GPU memory constraints). The fine-tuning dataset is a private video dataset of 3,000 videos (about 5 seconds each) collected from public sources, with text captions generated by Qwen3-VL-Flash.

An interesting observation: the paper notes that SLA2 (and other sparse attention methods) can outperform Full Attention after fine-tuning on many VBench metrics (Table 1: SLA2 at 90% sparsity scores 67.70 IQ vs. 63.67 for Full Attention on the 1.3B model). The authors attribute this to "the higher quality of the fine-tuning dataset compared to that used during pretraining" — essentially, the 500-step fine-tuning on a curated dataset provides a quality boost that offsets (and then some) any degradation from sparsification.

Transition from SoftTop-k to hard Top-k. This is a critical implementation detail. During Stage 1, the router uses SoftTop-k for gradient flow. When Stage 2 begins, the SoftTop-k is replaced with hard Top-k, and the router's projections are frozen — they are never updated again. This means the router must produce good-enough masks from Stage 1 alone to serve as a stable foundation for Stage 2 fine-tuning. If the Stage 1 initialization were poor, the model would be trying to learn attention patterns adapted to a suboptimal routing scheme, potentially leading to worse results than not using a learned router at all.


Forward Pass Algorithm (Algorithm 2)

The forward pass of SLA2 is implemented as a block-wise algorithm built on top of the FlashAttention paradigm — it processes queries and keys in blocks (of size $b_q$ for queries and $b_k$ for keys/values) to minimize memory footprint and maximize GPU utilization. Algorithm 2 in the paper provides the complete forward procedure; below I walk through its key design decisions.

Block structure. The algorithm divides $Q$ into $T_m = N / b_q$ blocks $\{Q_i\}$ and divides $K$, $V$ into $T_n = N / b_k$ blocks $\{K_j\}$, $\{V_j\}$. The compressed mask $M_c$ has dimensions $T_m \times T_n$ — each entry $M_c[i, j]$ determines whether the entire block pair $(i, j)$ is processed by the sparse branch (if $M_c[i, j] = 1$) or the linear branch (if $M_c[i, j] = 0$).

Precomputation for the linear branch (lines 6–7). Before the main loop, the algorithm precomputes two quantities for each key-value block $j$:

hj=(Kjϕ)TVjh_j = (K^{\phi}_j)^T V_j

zj=rowsum((Kjϕ)T)z_j = \text{rowsum}((K^{\phi}_j)^T)

where $\phi(\cdot)$ is the activation function for linear attention (the paper uses softmax, stated in Section 3 after Equation 14), and $K^{\phi}_j$ denotes $\phi(K_j)$. These precomputations are what make linear attention efficient: $h_j$ is a $d \times d$ matrix that aggregates the key-value product for block $j$, and $z_j$ is a $1 \times d$ vector of key sums for normalization. These can be computed once per block in $O(b_k d^2)$ time and then reused for all query blocks.

K-smoothing (line 2). Before any computation, the algorithm applies: $K = K - \text{colmean}(K)$, subtracting the column-wise mean from the key matrix. This is a smoothing technique from SageAttention (Zhang et al., 2025d;g) that reduces numerical outliers in the key values, improving the accuracy of low-bit quantization. The paper inherits this from the SageAttention2++ quantization scheme it builds on.

Main loop (lines 10–25). The algorithm iterates over query blocks $i$ from 1 to $T_m$. For each query block, it iterates over key-value blocks $j$ from 1 to $T_n$:

  • If $M_c[i, j] = 1$ (sparse branch, lines 12–18): The algorithm computes exact block-wise sparse attention using the quantized QK dot product (if QAT is enabled) or full-precision (if not). It follows the standard FlashAttention online softmax computation: compute $S_{ij} = Q_i K_j^T / \sqrt{d}$ (with optional quantization), track the running maximum $m_{ij}$ for numerical stability, compute exponentiated scores $P_{ij} = \exp(S_{ij} - m_{ij})$, accumulate the normalizing constant $l_{ij}$, and update the partial output $O^s_{ij}$ using the probability-value product $O_{\text{tmp}} = \text{dequant}(\text{quant}(P_{ij})\text{quant}(V_j))$. The update uses the standard FlashAttention rescaling: $O^s_{ij} = \text{diag}(e^{m_{i,j-1} - m_{ij}}) O^s_{i,j-1} + O_{\text{tmp}}$.

  • If $M_c[i, j] = 0$ (linear branch, lines 19–21): The algorithm simply accumulates the precomputed $h_j$ and $z_j$ for this block into running accumulators $H_i$ and $Z_i$. No QK dot products or PV products are computed for this block — the linear branch's contribution will be assembled at the end.

After the inner loop (lines 23–25): The sparse output for query block $i$ is rescaled by the final normalization constant: $O^s_i = \text{diag}(l_{i,T_n})^{-1} O^s_{i,T_n}$. The linear output is computed in one shot: $O^l_i = Q^{\phi}_i H_i / (Q^{\phi}_i Z_i)$, where $Q^{\phi}_i H_i$ is the unnormalized linear attention output (matrix multiplying the query features $Q^{\phi}_i$ with the accumulated key-value product $H_i$), and $(Q^{\phi}_i Z_i)$ is the normalization factor (row-wise sum of $Q^{\phi}_i Z_i^T$). The division is element-wise to perform row-wise normalization. The log-sum-exp value $L_i = m_{i,T_n} + \log(l_{i,T_n})$ is saved for the backward pass (it's needed for the gradient computation).

Final output (line 27): The overall output for the attention head is the learned convex combination: $O = \alpha \odot O^s + (1 - \alpha) \odot O^l$. Note that $\alpha$ is defined per-block at the compressed granularity ($\alpha \in \mathbb{R}^{N/b_q \times 1}$), meaning all tokens within a query block share the same mixing weight.


Backward Pass Algorithm (Algorithm 3)

The backward pass of SLA2 computes gradients with respect to the inputs $Q$, $K$, $V$, $Q^{\phi}$, and $K^{\phi}$. The gradients for the remaining parameters (the router projections $\text{proj}_q$, $\text{proj}_k$, the mixing ratio $\alpha$, and all other model parameters) are computed via PyTorch's automatic differentiation.

Why manual gradient derivation for Q, K, V, Q^φ, K^φ. The paper states (Appendix A) that these gradients are derived manually "following SLA." The reason is likely efficiency: the block-wise sparse + linear computation pattern is not natively handled by autograd in a memory-efficient way for these core attention tensors. Manual gradient formulas allow the implementers to reuse the forward pass's block structure and avoid materializing intermediate tensors that would consume excessive GPU memory.

Precomputation of auxiliary gradients (lines 2–6). Before the main loop, the algorithm computes two quantities that will be needed:

Ds=rowsum(dOsOs),Dl=rowsum(dOlOl)D^s = \text{rowsum}(dO^s \odot O^s), \quad D^l = \text{rowsum}(dO^l \odot O^l)

where $dO^s$ and $dO^l$ are the incoming gradients from the loss with respect to the sparse and linear outputs. These $D^s$ and $D^l$ are the row-wise sums of the element-wise product of gradient and output — a standard term that appears in the gradient of softmax attention due to the normalization.

For the linear branch, the algorithm also computes per-block gradients $dH_i$ and $dZ_i$ (lines 4–5), which represent gradients with respect to the accumulated key-value product and key sum, respectively:

dHi=(QiϕQiϕZi)TdOil,dZi=(QiϕQiϕZi)TDildH_i = \left(\frac{Q^{\phi}_i}{Q^{\phi}_i Z_i}\right)^T dO^l_i, \quad dZ_i = -\left(\frac{Q^{\phi}_i}{Q^{\phi}_i Z_i}\right)^T D^l_i

These use the fact that $O^l_i = Q^{\phi}_i H_i / (Q^{\phi}_i Z_i)$ — the gradient of a fraction with respect to the numerator and denominator follows the quotient rule, which these formulas implement. The gradient for the query features $Q^{\phi}_i$ is then:

dQiϕ=dOilHiTDilZiTQiϕZidQ^{\phi}_i = \frac{dO^l_i H_i^T - D^l_i Z_i^T}{Q^{\phi}_i Z_i}

This accounts for both the numerator and denominator contributions of $Q^{\phi}_i$ in the linear attention output.

Main backward loop (lines 7–18). The algorithm loops over key-value blocks $j$ and within each, loops over query blocks $i$. For each block pair, the computation depends on $M_c[i, j]$:

  • If $M_c[i, j] = 1$ (sparse branch, lines 10–13): The algorithm computes gradients through the sparse attention computation:

    1. Reconstruct the forward-pass score matrix $S_{ij} = Q_i K_j^T / \sqrt{d}$ (recomputed rather than stored, saving memory).
    2. Reconstruct the attention probabilities $P_{ij} = \exp(S_{ij} - L_i)$, where $L_i$ is the log-sum-exp from the forward pass (saved in line 25 of Algorithm 2).
    3. Accumulate the value gradient: $dV_j \leftarrow dV_j + P_{ij}^T dO^s_i$.
    4. Compute the gradient through softmax: $dP_{ij} = dO^s_{ij} V_j^T$, then $dS_{ij} = P_{ij} \odot (dP_{ij} - D^s_i)$. The subtraction of $D^s_i$ accounts for the normalization — it's the standard softmax gradient where the Jacobian of row-wise softmax involves subtracting the row-mean of the pre-softmax gradient.
    5. Accumulate query and key gradients: $dQ_i \leftarrow dQ_i + dS_{ij} K_j$ and $dK_j \leftarrow dK_j + dS^T_{ij} Q_i$.
  • If $M_c[i, j] = 0$ (linear branch, line 14): The algorithm simply accumulates the precomputed per-query-block gradients $dH_i$ and $dZ_i$ into running accumulators $dH$ and $dZ$ for this key-value block. No QK gradient computation is needed for linear branch positions.

After the inner loop (lines 17–18): The gradients for the key features $K^{\phi}_j$ are computed from the accumulated $dH$ and $dZ$:

dKjϕ=Vj(dH)T+(dZ)TdK^{\phi}_j = V_j (dH)^T + (dZ)^T

This follows from the forward pass definitions: $h_j = (K^{\phi}_j)^T V_j$ and $z_j = (K^{\phi}_j)^T$. The chain rule gives: $dK^{\phi}_j = V_j (dh_j / dK^{\phi}_j)^T + (dz_j / dK^{\phi}_j)^T = V_j (dH)^T + (dZ)^T$. The value gradient from the linear branch is also accumulated: $dV_j = K^{\phi}_j dH$.

Key efficiency property. The backward algorithm never needs to materialize the full $N \times N$ score or probability matrices. The sparse branch gradients are computed block-by-block, and the linear branch gradients use only $O(d^2)$ storage per block. This is what makes the method practical for large video models where $N$ can be tens of thousands of tokens.


Quantization-Aware Training (QAT) for the Sparse Branch

The sparse attention branch, even though it only computes $k\%$ of the $QK^T$ dot products, still performs dense matrix multiplications on those selected positions. These operations can be further accelerated by using low-bit arithmetic (INT8 or FP8) for the QK and PV products. However, naïvely quantizing after training (post-training quantization, PTQ) introduces accuracy loss because the model was never trained to tolerate quantization error in these intermediates.

What QAT does differently. During Stage 2 fine-tuning (end-to-end diffusion model training), the sparse branch's forward pass uses quantized computation — the QK dot products and PV products are performed in low-bit precision — while the backward pass stays in FP16. This asymmetry is crucial: the model sees quantized outputs during the forward pass and must learn to produce good results despite the quantization error, but the gradients are computed at full precision, avoiding the additional noise and bias that quantized gradients would introduce.

Forward pass quantization details (Section 5). The quantization scheme follows SageAttention2++ (Zhang et al., 2025g) and applies to the sparse branch only. The process is:

  1. Quantize Q and K before the dot product: $\hat{Q}, s_Q = \text{quant}(Q)$ and $\hat{K}, s_K = \text{quant}(K)$. The function $\text{quant}(\cdot)$ maps an FP16 tensor to a low-bit integer tensor (e.g., INT8) along with a per-tensor or per-block scale factor $s$.

  2. Compute quantized scores: $S = \text{dequant}(\hat{Q} \hat{K}^T / \sqrt{d}, s_Q, s_K)$. The matrix multiplication $\hat{Q} \hat{K}^T$ is performed in low-bit integer arithmetic (fast on modern GPUs with tensor cores), and the result is dequantized back to FP16 using the combined scale $s_Q s_K$. The division by $\sqrt{d}$ and application of the mask $M$ happen in FP16.

  3. Compute softmax in FP16: $P = \text{softmax}(S \odot M)$. The softmax itself is performed in FP16 because it's a nonlinear operation that doesn't benefit from quantization in the same way as matrix multiplications.

  4. Quantize P and V for the output product: $\hat{P}, s_P = \text{quant}(P)$ and $\hat{V}, s_V = \text{quant}(V)$.

  5. Compute quantized output: $O_s = \text{dequant}(\hat{P} \hat{V}, s_P, s_V)$. Again, the matrix multiplication runs in low-bit integer, and the result is dequantized back to FP16.

The crucial design choice is that only matrix multiplications are quantized — the operations that dominate FLOPs ($QK^T$ and $PV$). The element-wise operations (softmax, masking, scaling) remain in FP16 because they are memory-bound rather than compute-bound, so quantization provides negligible speedup.

Backward pass in FP16. The backward pass computes gradients $dQ$, $dK$, $dV$ entirely in FP16 using the original (non-quantized) inputs $Q$, $K$, $V$ and the forward output $O_s$. This is formalized in Section 5 as:

dQ,dK,dV=backward(dOs,Os,Q,K,V)dQ, dK, dV = \text{backward}(dO_s, O_s, Q, K, V)

The backward function uses the standard attention gradient formulas (as in Algorithm 3) applied to the full-precision tensors. The gradients $dO_s$ incorporate the effect of quantization error in the forward pass (because $O_s$ is the quantized output), so the model learns to compensate for quantization through the normal training dynamics — but the gradient computation itself is numerically clean, without the variance amplification that quantized gradients would introduce.

Why this approach works. The model sees the quantized output during training and must produce good results despite it, so it learns to produce attention patterns that are robust to the quantization error (e.g., spreading important information across multiple positions so no single quantized value is critical, or concentrating activations in ranges where quantization is precise). The FP16 backward pass ensures stable optimization — the model receives clean gradient signals that guide it toward quantization-robust solutions.

Ablation evidence (Table 2). The ablation study compares SLA2 trained with QAT vs. without QAT (both evaluated with quantized inference at 97% sparsity). Without QAT (line "w/o QAT"), the Vision Reward drops from 0.1039 to 0.0850, and other metrics degrade (IQ: 66.64 → 65.28, AQ: 64.62 → 61.85, SC: 94.83 → 94.65). This confirms that PTQ — applying quantization after training without QAT — introduces non-trivial error, and that QAT successfully mitigates it.

Efficiency impact. The paper reports (Section 9.4) that "low-bit quantization provides an approximately 1.3× kernel speedup." This is on top of the sparsity speedup, meaning SLA2's total 18.6× speedup over FlashAttn2 at 97% sparsity includes both the sparsity benefit (~14.3× from skipping 97% of computations) and the quantization benefit (~1.3× from low-bit arithmetic on the remaining 3%).

Quantization scope. Only the sparse branch is quantized. The linear branch runs at full precision, which is sensible because linear attention's algorithmic reformulation already provides substantial speedup, and quantizing the $O(d^2)$ operations there would yield diminishing returns relative to the risk of quantization error propagating through the accumulated $H$ matrices.


Implementation of Sparse and Linear Branches: How Attention Is Actually Computed

The paper's notation in Equation 14 might misleadingly suggest that computing $O_s$ requires a full $QK^T$ followed by masking. The actual implementation (Section 3, paragraph "Implementation of getting $O_s$ and $O_l$") uses a much more efficient strategy:

Sparse branch implementation. Built on top of FlashAttention, the sparse branch only computes $QK^T$ and $PV$ for positions where $M = 1$. Specifically, the block-wise loop in Algorithm 2 (lines 12–18) checks $M_c[i, j]$ for each block pair: if it's 0, the sparse branch computation for that block pair is entirely skipped (no QK dot product, no softmax, no PV product). This means the sparse branch only performs matmuls for $k\% \times N^2 / (b_q b_k)$ of the full attention blocks. For 97% sparsity ($k\% = 3\%$), only 3% of the $QK^T$ and $PV$ computation is performed — the rest is zero-skipped.

This is why the FLOPs reduction in Table 1 (from 52.75T for Full Attention to 1.82T for SLA2 at 97% sparsity on the 1.3B model) is not exactly $1 / 0.03 = 33.3\times$ but rather approximately $52.75 / 1.82 \approx 29\times$. The remaining FLOPs come from the linear branch computation and the routing overhead, neither of which scales with $N^2$.

Linear branch implementation. The linear branch computes:

Ol=norm(ϕ(Q)ϕ(K)T(1M))VO_l = \text{norm}(\phi(Q) \phi(K)^T \odot (1 - M)) V

The key insight is that this can be computed WITHOUT materializing the $N \times N$ matrix $\phi(Q)\phi(K)^T$. Instead, as shown in Algorithm 2 (lines 20–21, 24), the computation proceeds as:

  1. For each key-value block $j$ where $M_c[i, j] = 0$ (i.e., positions assigned to linear attention), accumulate $h_j = (K^{\phi}_j)^T V_j$ into $H_i$ and $z_j = (K^{\phi}_j)^T$ into $Z_i$.
  2. After processing all key-value blocks, compute the output in one step: $O^l_i = Q^{\phi}_i H_i / (Q^{\phi}_i Z_i)$.

This is $O(N d^2)$ rather than $O(N^2 d)$ because:

  • Step 1 requires $T_n$ computations of $(K^{\phi}_j)^T V_j$ (each $O(b_k d^2)$), for a total of $O(N d^2)$.
  • Step 2 requires $T_m$ computations of $Q^{\phi}_i H_i$ (each $O(b_q d^2)$), for a total of $O(N d^2)$.

The $(1 - M)$ mask is enforced implicitly: the linear branch only accumulates key-value blocks where $M_c[i, j] = 0$, never touching blocks assigned to the sparse branch.

Why the linear branch uses softmax as $\phi(\cdot)$. The paper states in Section 3 (after Equation 14) that "we use the softmax function" as the activation $\phi(\cdot)$ for linear attention. This is a design choice: linear attention originally used various kernel functions (ReLU, ELU+1, etc.), but using softmax means the linear branch is approximating the same softmax-based attention that the sparse branch computes exactly — just without the pairwise normalization constraints. This makes the two branches more comparable and the decomposition more natural: both branches operate on probability-like representations, with the sparse branch capturing exact softmax for selected positions and the linear branch approximating softmax for the rest.

The $\text{norm}$ operation in Equation 14 ensures the linear attention output is properly row-normalized — it divides each row's output by the sum of attention weights in that row, analogous to how softmax ensures each row of the probability matrix sums to 1.


Design Choices Summary: Why Each Matters

Why the router uses block-based routing rather than token-level routing. Making routing decisions at the block level ($b_q = 128$, $b_k = 64$) rather than per-token is a crucial efficiency decision. Token-level routing would require a mask of size $N \times N$, which for video models with tens of thousands of tokens would be enormous (e.g., for $N = 16384$ (128 frames × 128 spatial tokens), the mask alone would be 268 million entries). Block-level routing reduces this to $(N/b_q) \times (N/b_k)$, which for the same example would be $128 \times 256 = 32768$ entries — a factor of 8192 reduction. The smoothness assumption (adjacent tokens have similar attention patterns) justifies this approximation.

Why the mixing ratio $\alpha$ is learned rather than derived from attention statistics. One might ask: since $\alpha$ is supposed to represent the probability mass on the masked positions, couldn't it be directly computed from the actual attention weights instead of learned? The problem is that at inference time, computing the exact $\alpha$ would require computing the full attention matrix (to sum the probabilities on masked positions), which defeats the purpose of sparsification. By learning $\alpha$, the model can internalize the relationship between its query representation and the expected probability mass on sparse positions, allowing $\alpha$ to be inferred without computing full attention.

Why Stage 1 uses MSE loss on attention outputs rather than end-to-end diffusion loss. The attention output is a well-defined target at the per-layer level, and MSE provides dense supervision. If Stage 1 used the end-to-end diffusion loss, the router would need to be trained jointly with the full model from the start, which is expensive and suffers from the instability issues already discussed. The MSE loss isolates the routing problem: learn a mask such that the attention output (before downstream layers process it) is as close as possible to full attention.

Why QAT is applied only to the sparse branch, not the linear branch. Linear attention's speedup comes from algorithmic reformulation — replacing $O(N^2 d)$ with $O(N d^2)$ — not from compute-bound matmuls that benefit from quantization. The linear branch's operations are mostly memory-bound (accumulating $h_j$ and $z_j$, then multiplying by $Q^{\phi}_i$), so low-bit arithmetic provides minimal additional benefit while risking compounding quantization error through the accumulated $H$ matrices. The sparse branch, in contrast, still performs dense matmuls ($Q_i K_j^T$ and $P_{ij} V_j$) that are compute-bound and benefit directly from tensor core low-bit acceleration.

Why the K-smoothing step (line 2 of Algorithm 2) is needed. This step ($K = K - \text{colmean}(K)$) is inherited from SageAttention and is specifically important for quantization accuracy. Without smoothing, the key matrix $K$ may have column-wise offsets (some dimensions consistently larger than others) that cause the quantization scale to be dominated by a few outlier dimensions, reducing effective precision for the rest. Subtracting the column mean centers each dimension, making the quantization scale more uniform and improving the signal-to-noise ratio of the quantized dot products. This is a standard technique in low-bit attention (see, e.g., SageAttention and FlashAttention-3).

4. Key Insights and Innovations

Innovation 1: A Causal Diagnosis of SLA's Failure as a Scaling Mismatch, Not a Capacity Mismatch

The paper's most intellectually distinctive contribution is not proposing a new mechanism but rather performing a precise, mathematical diagnosis of why the prior state-of-the-art (SLA) falls short. This diagnosis is valuable because it identifies a specific, correctable structural flaw rather than a vague "needs more capacity" or "needs better training" explanation.

Prior to this work, SLA's approach of adding a learnable projection to the linear branch output (O = O_s + Proj(O_l)) was a reasonable design choice — if the linear branch could not perfectly match the missing attention mass, a learned linear transformation could compensate. The field's implicit assumption was that any mismatch between the sparse branch's output and its target contribution could be absorbed by this projection as part of the general approximation task.

SLA2 reveals that this assumption ignores a crucial structural fact: the sparse branch does not compute P_1 V (the actual contribution of the masked positions) but rather O_s = P_s V where P_s is a renormalized probability matrix. The mismatch is multiplicative — P_1 V = α ⊙ O_s — meaning the sparse branch's output is row-wise scaled by the probability mass it represents. This multiplicative scaling error is fundamentally different from the additive error that the linear branch's projection was designed to handle. Concretely, the projection Proj(·) operates on O_l (which encodes information from the linear branch's positions), but the correction it needs to make — (α − 1) ⊙ O_s — depends on the sparse branch's output. The projection is being asked to correct an error using information it does not have access to.

This diagnosis is a conceptual advance: it shifts the understanding of what makes sparse-linear attention work from "you need a powerful enough combination mechanism" to "the combination mechanism must respect the algebraic structure of the decomposition." The paper's reformulation — O = α ⊙ O_s + (1 − α) ⊙ O_l — is not a radical architectural departure but a structurally faithful implementation of the original motivation. The learnable ratio α handles the row-wise scaling explicitly, allowing each branch to focus on its intended role: the sparse branch models the exact attention for selected positions, and the linear branch approximates a normalized probability distribution over the remaining positions. The linear branch no longer has to compensate for a multiplicative error originating in the sparse branch.

This is a fundamental insight, not an incremental refinement, because it identifies a design principle that generalizes beyond SLA2: when decomposing an attention matrix into components handled by different approximators, the decomposition must be formulated so that each approximator's output directly matches its target contribution, without requiring cross-branch error correction. The evidence for the diagnosis's validity is indirect but strong: SLA2 at 97% sparsity (Table 1) matches SLA at 90% sparsity on most quality metrics — a 7 percentage point sparsity improvement while maintaining quality — suggesting the reformulation enables more efficient use of each branch's capacity.

Innovation 2: Framing the Routing Problem as an Optimization Over Mask Quality, Not a Heuristic Over Weight Magnitude

The paper's second conceptual contribution is reframing what it means to route between sparse and linear attention. The dominant approach in prior work — SLA's magnitude-based heuristic — implicitly assumes that the goal is to identify and preserve the largest attention weights. Large weights → sparse branch, small weights → linear branch. This is intuitively sensible: softmax amplifies large dot products, so preserving them seems natural.

SLA2 challenges this framing by pointing out that the true objective is not preserving large weights per se but rather producing a decomposition where P_1 (the sparse component) is maximally sparse while P_2 (the component handled by linear attention) is maximally easy for linear attention to approximate. The distinction is subtle but profound. A moderate-weight position that is highly correlated with many other positions (low-rank structure) might be perfectly well-approximated by linear attention, so routing it to the sparse branch wastes sparsity budget. Conversely, a small-weight position that captures unique, non-redundant structure might be poorly approximated by linear attention, so routing it to the linear branch degrades quality. The magnitude heuristic conflates "large" with "structurally important" and "small" with "approximable," which is not necessarily true.

This reframing is significant because it transforms the routing problem from a heuristic selection rule (which requires no training but is inherently suboptimal) into an optimization problem (which requires training but can discover structure-aware assignments). The paper operationalizes this by introducing a learnable router with projections proj_q and proj_k that are trained to minimize the reconstruction error between full attention and the SLA2 approximation. The router is not told "preserve large weights" — it discovers through gradient descent which positions are best computed exactly given the specific approximation characteristics of the deployed sparse and linear branches.

This is a fundamental conceptual shift, not an incremental improvement, because it changes the nature of the routing decision from a local, per-weight criterion (is this weight large?) to a global, structure-aware criterion (does this assignment make the overall decomposition better?). The evidence that this matters is in Table 2's ablation: replacing the learned router with a Top-k router (SLA's heuristic based on pool(Q)pool(K)^T) drops the Vision Reward from 0.1039 to 0.0876 and Imaging Quality from 66.64 to 63.66 at 97% sparsity. The heuristic router performs only slightly better than full attention (63.66 vs. 63.67 IQ), suggesting that at high sparsity, the magnitude heuristic essentially breaks down — it cannot identify the right positions to preserve when only 3% remain. The learned router, in contrast, maintains quality close to the lower-sparsity configurations.

Innovation 3: Positioning Quantization-Aware Training as a First-Class Design Element, Not a Post-Hoc Optimization

The paper's integration of quantization-aware training (QAT) into the sparse-linear attention framework is distinctive not because QAT itself is novel (it has been studied extensively, e.g., Jacob et al., 2018; Nagel et al., 2022), but because the paper positions quantization as a co-design element rather than a post-hoc acceleration trick. The typical workflow in efficient attention research is: (1) design a sparsity or approximation method, (2) evaluate its speedup at full precision, (3) optionally apply post-training quantization for additional speedup, often with an acknowledged quality penalty. The quantization is an afterthought — something you do after the "real" method is built.

SLA2 integrates quantization into the training pipeline (Stage 2 fine-tuning with QAT enabled) and treats the low-bit sparse branch as the target deployment configuration, not an optional add-on. The ablation in Table 2 ("w/o QAT") demonstrates that skipping QAT and applying post-training quantization instead causes a measurable quality degradation (Vision Reward drops from 0.1039 to 0.0850, Aesthetic Quality drops from 64.62 to 61.85). This shows that the model's parameters benefit from adapting to quantization during training — the attention patterns learned without QAT are not robust to the quantization error introduced at inference time.

What makes this a genuine insight rather than an engineering detail is the paper's implicit argument that quantization error and sparsification error interact. The sparse branch, after aggressive sparsification (97%), computes attention on only a tiny fraction of positions. Each of those few remaining computations must be highly accurate because there is no redundancy to fall back on — if a quantized dot product produces a distorted score for one of the few preserved positions, the softmax amplifies the error and propagates it through the value aggregation. In a dense attention setting, quantization errors in individual positions are diluted by the many other positions contributing to the output. In an extreme-sparsity setting, each quantized computation carries more weight, making quantization error more consequential. QAT is therefore more important for sparse attention than for dense attention, not less.

This is an incremental contribution — QAT itself is well-established — but the insight about the interaction between sparsity and quantization sensitivity, and the treatment of QAT as an integral part of the sparse-linear design rather than a separate acceleration step, is a useful reframing for practitioners. The 1.3× additional speedup from quantization (Section 9.4) on top of the sparsity speedup demonstrates that these are complementary acceleration axes when properly co-optimized.

Innovation 4: Empirical Demonstration That Extreme Sparsity + Fine-Tuning Can Surpass Full Attention Quality

The paper's most striking result — that SLA2 at 90% sparsity actually outperforms full attention on multiple VBench metrics (e.g., IQ 67.70 vs. 63.67 on Wan2.1-1.3B, Table 1) — is not the paper's conceptual contribution but rather its empirical finding with the most significant practical implications. The standard narrative in efficient attention research is that sparsification is a tradeoff: you sacrifice some quality for speed. The goal is to minimize the quality degradation, not eliminate it. Surpassing full attention while being 10× more efficient challenges this narrative.

The paper attributes this to dataset quality ("the higher quality of the fine-tuning dataset compared to that used during pretraining"), which is both honest and important. It means the result is not a claim that sparsity is inherently better than dense attention, but rather that the fine-tuning process required to adapt the model to sparsity can simultaneously provide a quality boost from improved training data. In other words, the 500-step fine-tuning on a curated 3,000-video dataset improves the model's generation quality enough to offset — and then exceed — any degradation from sparsification.

This finding has significant practical implications: it suggests that for practitioners who have access to a high-quality domain-specific dataset, deploying a sparse attention variant with fine-tuning can be a pure win — lower latency, lower cost, AND better quality — rather than a tradeoff where quality is sacrificed for speed. It also suggests that the quality gap between open-source pretrained models and what can be achieved with targeted fine-tuning on curated data may be large enough that efficiency methods which require fine-tuning (like SLA2) can effectively get the quality improvement "for free" as part of the adaptation process.

This is an observational contribution rather than a methodological one — it doesn't introduce a new technique but rather documents a phenomenon that changes how practitioners should think about the cost-benefit analysis of trainable sparse attention. The evidence is in Table 1: SLA2 at 90% sparsity achieves higher scores than Full Attention on IQ (+4.03), OC (+1.35), AQ (+0.45), and MS (−0.26, slightly lower but essentially equivalent) on the 1.3B model, and similar patterns appear on the 14B model. The result is consistent across two model scales, strengthening its credibility. However, the finding is contingent on the fine-tuning dataset and may not generalize to all domains or all base models — it is an existence proof that surpassing full attention is possible, not a guarantee that it always happens.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses a private video dataset of 3,000 videos (approximately 5 seconds each) collected from public sources. Text captions for each video are generated using Qwen3-VL-Flash, producing text-video pairs for both fine-tuning and evaluation. The evaluation uses VBench (Zhang et al., 2024) metrics, which presumably involve generating videos from text prompts and scoring them with automated quality assessors. The paper does not specify the number of evaluation prompts or describe a train/val/test split of the 3,000 videos — this is a notable omission for reproducibility.

  • Base model(s). Experiments are conducted on two variants of the Wan2.1 text-to-video diffusion model (Wan et al., 2025): Wan2.1-T2V-1.3B-480P (1.3 billion parameters, generates 480p video) and Wan2.1-T2V-14B-720P (14 billion parameters, generates 720p video). These are chosen as representative state-of-the-art open video diffusion models at two distinct scales, allowing the paper to test whether the method's benefits generalize across model sizes. The 14B model notably exceeds the VRAM capacity of a single RTX5090 GPU, requiring sequential CPU offloading during evaluation — the paper explicitly states that reported end-to-end latencies "already exclude the offloading overhead" (Section 9.3).

  • Metrics. Video generation quality is evaluated across multiple dimensions from VBench (Zhang et al., 2024): Imaging Quality (IQ), Overall Consistency (OC), Aesthetic Quality (AQ), Motion Smoothness (MS), and Subject Consistency (SC). Additionally, the paper uses Vision Reward (VR) (Xu et al., 2024), a learned human-preference metric for video quality. Computational efficiency is measured via FLOPs (total floating-point operations for the attention modules), attention sparsity (percentage of QK dot products skipped), kernel speed (measured in TOPS — trillion operations per second, defined as C/t where C = 4N^2d is the theoretical computation and t is execution latency), and end-to-end inference latency (in seconds). The multiplicity of quality metrics is a strength — it guards against cherry-picking a single favorable metric — though all metrics are automated; no human evaluation is reported.

  • Baselines. The paper compares against five baselines:

    • Full Attention: standard softmax attention implemented with FlashAttn2 (Dao, 2023), representing the unaccelerated reference with 0% sparsity. This baseline is not fine-tuned, serving as the pretrained model's quality ceiling.
    • SLA (Zhang et al., 2025c): the prior sparse-linear attention method that SLA2 directly improves upon, using magnitude-based heuristic routing and a projection-based branch combination. Fine-tuned for 500 steps under the same protocol as SLA2.
    • VSA (Zhang et al., 2025i): a trainable sparse attention method without a linear attention compensation branch.
    • VMoBA (Wu et al., 2025): a mixture-of-block attention method that uses structured sparsity patterns.
    • For the ablation study (Table 2), additional internal baselines include SLA2 w/o QAT (trained with full-precision sparse branch, evaluated with post-training quantization), and Topk-router (SLA2 but replacing the learned router with SLA's heuristic pool(Q)pool(K)^T-based Top-k selection).

    All trainable baselines (SLA, VSA, VMoBA) are fine-tuned for 500 steps on the same dataset with the same hyperparameters — the paper explicitly states "We fine-tune each method for 500 steps" (Section 9.1), making this a controlled comparison.

  • Generation budget / compute accounting. The paper uses multiple compute measures at different levels of granularity:

    • Sparsity percentage (90%, 95%, 97%) is the primary axis of comparison, defined as the fraction of QK^T dot products skipped by the sparse branch. For SLA2, the sparsity is controlled by the k% parameter in the router, with k% = 10%, 5%, 3% corresponding to sparsities of 90%, 95%, 97% respectively. The paper notes that "97% sparsity corresponds to about 96.7% computation savings after accounting for the linear-attention branch" — the 0.3 percentage point difference reflects the non-zero cost of the linear branch.
    • FLOPs are reported in Table 1 as the total attention FLOPs for a full generation. For the 1.3B model at 0% sparsity, Full Attention requires 52.75T FLOPs; SLA2 at 97% sparsity requires 1.82T FLOPs.
    • Kernel speed (Figure 4) is measured in TOPS on an RTX5090, reported separately from end-to-end latency to isolate the attention kernel's performance from other generation components.
    • End-to-end latency (Figure 5) measures total generation time including all model components, partitioning the time into "Attention" and "Others" to show how attention speedup translates to overall speedup.

    All compute metrics are measured, not estimated — the paper uses actual GPU execution times and FLOP counts from the implemented kernels. This is a strength over papers that report only theoretical speedups.

  • Cross-validation / statistical protocol. The paper does not report any cross-validation or statistical significance testing. Results in Tables 1 and 2 are single-point estimates without confidence intervals or error bars. The fine-tuning procedure uses a fixed 500 steps on a fixed dataset; there is no mention of multiple runs with different random seeds or train/test splits. For the ablation study (Table 2), all configurations are trained once and evaluated once. This is a limitation — without variance estimates, it is impossible to determine whether the reported differences (e.g., SLA2 at 97% sparsity achieving Vision Reward 0.1039 vs. the w/o QAT variant at 0.0850) are statistically reliable or within the noise of training stochasticity. The paper's central claim of "outperforming baselines at 90% sparsity" would be strengthened by multi-seed experiments, particularly given the relatively small fine-tuning dataset (3,000 videos).


Main Quantitative Results

Headline Quality vs. Efficiency Results (Table 1)

Table 1 presents the paper's primary results: video generation quality metrics and attention FLOPs for SLA2 and all baselines at sparsity levels of 90%, 95%, and 97%, on both Wan2.1 models.

On Wan2.1-T2V-1.3B-480P:

At 90% sparsity, SLA2 achieves the highest values on nearly every quality metric: IQ 67.70 (vs. 63.67 for Full Attention, 63.10 for SLA, 65.31 for VMoBA, 59.57 for VSA), OC 21.62 (vs. 20.27 for Full Attention, 20.88 for SLA), AQ 64.86 (vs. 64.41 for Full Attention), MS 98.69 (essentially tied with SLA at 97.90 and Full Attention at 98.95), SC 95.54 (vs. 95.40 for Full Attention), and VR 0.1093 (vs. 0.1084 for Full Attention). SLA2's FLOPs at 90% sparsity are 5.51T — a 9.6× reduction from Full Attention's 52.75T, and comparable to the other sparse methods (SLA: 5.40T, VSA: 5.40T, VMoBA: 5.28T). The FLOPs parity across methods at the same sparsity level confirms that the sparsity percentage is the dominant cost factor and that the routing/mixing overhead is minimal.

At 95% sparsity, SLA2 continues to lead: IQ 67.04 (vs. 63.14 for SLA, 63.08 for VMoBA, 55.50 for VSA), OC 21.55 (vs. 21.09 for SLA), AQ 64.90 (vs. 62.91 for SLA), VR 0.1023 (vs. 0.0881 for SLA). VMoBA degrades significantly at this sparsity level (IQ drops from 65.31 at 90% to 63.08 at 95%, SC drops from 86.69 to 79.83), while SLA2's metrics are nearly unchanged from its 90% sparsity results. SLA2's FLOPs drop to 2.87T (an additional 1.9× reduction from 90%).

At 97% sparsity, SLA2 alone is evaluated (no baseline at this sparsity level). The results show SLA2 achieving IQ 66.64, OC 21.42, AQ 64.62, MS 98.04, SC 94.83, and VR 0.1039, with 1.82T FLOPs — a 29× FLOPs reduction from Full Attention. Importantly, these 97% sparsity metrics exceed all baselines at 90% sparsity across every quality dimension except MS (where Full Attention leads by a negligible 0.91 points). This is the paper's central empirical claim: SLA2 at 97% sparsity provides better quality than competing methods at 90% sparsity while using 3× fewer FLOPs (1.82T vs. 5.28–5.51T). The comparison is specifically: SLA2 at 97% sparsity vs. VMoBA at 90% sparsity on IQ (66.64 vs. 65.31), OC (21.42 vs. 20.82), AQ (64.62 vs. 64.14), SC (94.83 vs. 86.69), and VR (0.1039 vs. 0.0936).

On Wan2.1-T2V-14B-720P:

The pattern is similar but with one notable exception. At 90% sparsity, SLA2 leads on IQ (69.63 vs. 68.01 for Full Attention, 67.58 for SLA), AQ (66.41 vs. 64.66 for Full Attention), and VR (0.1238, tied with Full Attention). SLA2 achieves 21.16T FLOPs vs. 292.6T for Full Attention, a 13.8× reduction. However, OC is slightly lower: SLA2 at 20.68 vs. Full Attention at 22.44 and VSA at 21.27. This is the one metric where SLA2 does not dominate — a minor qualification to the otherwise comprehensive quality leadership.

At 95% sparsity, a striking result emerges: VMoBA catastrophically fails on the 14B model. VMoBA's OC drops to 7.96 (vs. 21.27 at 90%), AQ drops to 33.59 (vs. 63.64 at 90%), and VR goes negative (−0.0965). This suggests VMoBA's block-based sparsity pattern is fundamentally incompatible with the larger model's attention structure at high sparsity — the structured sparsity presumably eliminates entire blocks that are essential for cross-frame consistency. In contrast, SLA2 at 95% sparsity maintains robust quality: IQ 69.02, OC 21.11, AQ 65.55, MS 98.89, SC 95.53, VR 0.1125, with 15.11T FLOPs. VSA also degrades (IQ 47.69, OC 13.90, VR −0.1822), confirming that without a linear compensation branch, high sparsity is unattainable.

At 97% sparsity, SLA2 achieves 9.26T FLOPs (a 31.6× reduction from Full Attention's 292.6T) while maintaining IQ 66.93, OC 21.12, AQ 65.14, MS 98.71, SC 94.42, and VR 0.1149. Again, these 97% sparsity metrics exceed most baselines at 90% sparsity (e.g., VMoBA at 90%: IQ 67.18, OC 20.85, AQ 63.64, VR 0.1117 — all lower than SLA2 at 97%). The 14B model shows even larger relative gains from SLA2 than the 1.3B model, suggesting that the learned routing is particularly effective at the larger scale where attention matrices have more redundancy to exploit.

The "outperforms Full Attention" phenomenon. Across both models and most sparsity levels, SLA2 and several other fine-tuned sparse methods achieve higher quality metrics than Full Attention (which is not fine-tuned). This is most pronounced on the 1.3B model: SLA2 at 90% sparsity scores 67.70 IQ vs. 63.67 for Full Attention, a +4.03 point improvement. The paper attributes this to the fine-tuning dataset quality rather than to sparsity itself: "We attribute this to the higher quality of the fine-tuning dataset compared to that used during pretraining" (Section 9.2). This is an important caveat — the quality improvements are not a property of the attention mechanism but of the additional training on curated data. A fairer comparison would include a "Full Attention + Fine-tuning" baseline to disentangle the effects of fine-tuning from the effects of sparsification, but this baseline is not reported.

Visible examples (Figures 2 and 3). Figure 2 shows a single generated frame from the 1.3B model for the prompt about a morning makeup routine. The SLA2 images at 95% and 97% sparsity are visually comparable to Full Attention — the spatial layout, object placement, and overall scene composition are preserved. The SLA, VSA, and VMoBA examples at 90% sparsity show deviations: Figure 2 shows that "videos from other methods either differ noticeably from Full Attention or show clear distortions" (Section 9.2). The prompt is provided in Appendix B. Figure 3 shows an example from the 14B model (a cat running across a meadow), with SLA2 at 95% and 97% sparsity appearing visually indistinguishable from Full Attention. These qualitative examples are a useful sanity check but represent only 2 prompts — they cannot substitute for a systematic human evaluation.


Kernel-Level Speed Results (Figure 4)

Figure 4 presents the forward kernel speed in TOPS (trillion operations per second) measured on an RTX5090 for all methods at various sparsity levels. Higher TOPS indicates better hardware utilization and faster execution.

The key finding: SLA2 at 97% sparsity achieves 4,079 TOPS, which is:

  • 18.6× faster than FlashAttn2 (219 TOPS) — this is the headline "18.6× attention speedup" claimed in the abstract.
  • 2.6× faster than VSA at 95% sparsity (1,553 TOPS).
  • 11.7× faster than VMoBA at 95% sparsity (348 TOPS).
  • 1.36× faster than SLA at 95% sparsity (2,996 TOPS).

Notably, SLA2 at 95% sparsity (3,610 TOPS) is already faster than SLA at 95% sparsity (2,996 TOPS), despite both having the same theoretical sparsity. This suggests SLA2's implementation has lower kernel overhead — possibly due to the simplified branch combination (no Proj(O_l) computation) reducing per-block work. The speed hierarchy across methods reveals that VMoBA and VSA are substantially less efficient at the kernel level despite similar sparsity percentages, likely due to irregular memory access patterns or less optimized GPU kernels.

What TOPS measures. The paper defines TOPS as C / t where C = 4N^2d is the theoretical number of floating-point operations for full attention, and t is the measured kernel execution time (Section 9.1). This means TOPS is a throughput metric — it measures how fast the kernel processes the equivalent of a full-attention computation, NOT how many operations it actually executes. For a sparse method, the kernel performs fewer actual FLOPs but is benchmarked against the theoretical full-attention FLOPs, resulting in a high TOPS number that reflects both the computation reduction AND any hardware efficiency improvements. This metric is standard in the FlashAttention literature but readers should note that "4,079 TOPS" does not mean the kernel executes 4 trillion operations per second literally — it means it completes the equivalent of a full-attention computation 4,079 / 219 = 18.6× faster than FlashAttn2.


End-to-End Generation Latency (Figure 5)

Figure 5 decomposes end-to-end video generation latency into "Attention" and "Others" components for each method, showing how attention speedup translates to overall speedup.

On Wan2.1-1.3B-480P (Figure 5a):

  • Original (Full Attention): 97 seconds attention + 62 seconds other = 159 seconds total. Attention dominates.
  • SLA2 at 97% sparsity: 7 seconds attention + 62 seconds other = 69 seconds total, a 2.30× end-to-end speedup (159/69). The attention time is reduced by 13.9× (97→7), but the "Other" time (62 seconds) becomes the new bottleneck, limiting the overall speedup.
  • Among baselines, VMoBA at 95% sparsity achieves 11 seconds attention (end-to-end: 73s), VSA at 95% achieves 62 seconds attention (end-to-end: 124s), and SLA at 90% achieves 18 seconds (end-to-end: 80s). SLA2 at 97% sparsity has the lowest overall latency despite higher sparsity than any baseline.

On Wan2.1-14B-720P (Figure 5b):

  • Original (Full Attention): 2,550 seconds attention + 493 seconds other = 3,043 seconds total (~50.7 minutes). The "Other" time here is dominated by the model's feed-forward and convolution layers, which scale with model size but not with sequence length, explaining why attention dominates more heavily than on the 1.3B model.
  • SLA2 at 97% sparsity: 207 seconds attention + 493 seconds other = 700 seconds total, a 4.35× end-to-end speedup. The attention reduction factor is 12.3× (2,550→207), limited more by the "Other" bottleneck than on the 1.3B model.
  • VMoBA at 95% achieves 457 seconds attention (end-to-end: 950s), VSA at 95% achieves 651 seconds (end-to-end: 1,144s), and SLA at 90% achieves 409 seconds (end-to-end: 902s). Again, SLA2 at 97% sparsity provides the lowest overall latency despite the highest sparsity.

The "Other" bottleneck: On both models, the non-attention computation time (62s and 493s respectively) is constant across all methods because SLA2 only modifies the attention modules. This creates an Amdahl's Law scenario: as attention time decreases, the fixed "Other" time becomes the limiting factor. On the 1.3B model, attention drops from 97s to 7s (14× speedup), but end-to-end speedup is only 2.3× because "Other" remains at 62s. Further attention speedup beyond this point would yield diminishing returns. The paper does not discuss optimizing the non-attention components, which suggests future work could target the feed-forward and normalization layers for a more balanced speedup.

CPU offloading caveat for 14B model: The paper states that "the Wan2.1-14B-720P model exceeds the VRAM capacity of a single RTX5090, we enable sequential CPU offloading during evaluation. The reported latency already excludes the offloading overhead" (Section 9.3). This means the absolute latencies for the 14B model (e.g., 700s for SLA2) are measured with the model partially on GPU and partially on CPU — the reported numbers exclude the CPU→GPU transfer time. In a multi-GPU or higher-VRAM setting, the latencies would differ. The relative speedups (4.35×) should still hold approximately since the offloading overhead is excluded from all methods, but the absolute times are not directly comparable to a fully GPU-resident deployment.


Sparsity Scaling Behavior (Table 2, lower section)

The lower portion of Table 2 shows SLA2's quality metrics at four sparsity levels: 85%, 90%, 95%, and 97%. This can be interpreted as a sparsity scaling curve:

SparsityIQOCAQMSSCVR
85%67.9721.9864.7998.7595.790.1135
90%67.7021.6264.8698.6995.540.1093
95%67.0421.5564.9098.4695.270.1023
97%66.6421.4264.6298.0494.830.1039

The degradation from 90% to 97% sparsity is remarkably gentle: IQ drops by only 1.06 points (67.70→66.64), OC by 0.20 (21.62→21.42), AQ by 0.24 (64.86→64.62), and VR actually increases slightly from 95% to 97% (0.1023→0.1039). This near-flat scaling from 90% to 97% sparsity is the crucial evidence that the learned router and decomposition-consistent mixing successfully identify and preserve the truly essential attention computations. If the router were making poor decisions, quality would degrade sharply as fewer positions are retained — as observed with VMoBA (IQ drops from 65.31 at 90% to 63.08 at 95%) and VSA (IQ drops from 59.57 at 90% to 55.50 at 95%). The fact that SLA2 barely degrades suggests the router is correctly distinguishing structurally important attention positions from those that can be safely approximated.

The 85% result: At 85% sparsity, SLA2's quality actually exceeds the 90% configuration slightly (IQ 67.97 vs. 67.70, OC 21.98 vs. 21.62, VR 0.1135 vs. 0.1093). This confirms the expected monotonic trend: lower sparsity → higher quality, though the difference is small, suggesting SLA2 is near the quality ceiling even at 90% sparsity.

Comparison to Full Attention scaling: An interesting counterfactual: Full Attention achieves IQ 63.67 (with no fine-tuning), while SLA2 at 90% sparsity achieves 67.70. If we consider Full Attention's IQ as a baseline, then SLA2 at 97% sparsity (66.64) is still +2.97 above Full Attention. This means even at the most extreme sparsity tested, the combination of fine-tuning and the sparse-linear decomposition produces better quality than the original pretrained model — a result that would be impossible if sparsity were purely destructive.


Summary of Headline Claims vs. Evidence

Claim: "SLA2 achieves 97% attention sparsity and 18.6× attention speedup while preserving generation quality."

The evidence is in Table 1 and Figure 4. At 97% sparsity, SLA2 on the 1.3B model achieves IQ 66.64, OC 21.42, AQ 64.62, MS 98.04, SC 94.83, and VR 0.1039 (Table 1), which exceeds or matches Full Attention (IQ 63.67, OC 20.27, AQ 64.41, MS 98.95, SC 95.40, VR 0.1084). The 18.6× kernel speedup is measured at 4,079 TOPS vs. 219 TOPS for FlashAttn2 (Figure 4). This claim is well-supported quantitatively, with the caveat that the quality "preservation" is measured by automated metrics, not human evaluation, and that quality at 97% sparsity is slightly lower than at 90% sparsity (IQ 66.64 vs. 67.70) — the statement "preserving generation quality" should be understood as "maintaining quality above the Full Attention baseline," not "identical to lower sparsity."

Claim: "SLA2 outperforms baselines at 90% sparsity in end-to-end video quality, and even exceeds full attention."

This is the strongest empirical result in the paper and is supported by Table 1. SLA2 at 97% sparsity achieves higher IQ, OC, AQ, SC, and VR than VMoBA at 90% sparsity on the 1.3B model, and higher metrics than VMoBA, VSA, and SLA at 90% sparsity on most dimensions. The "exceeds full attention" claim is supported but requires the fine-tuning caveat — the paper makes this explicit in Section 9.2, attributing the improvement to dataset quality.


Ablation Studies and Robustness Checks

Quantization-aware training (QAT): Table 2 compares SLA2 at 97% sparsity with QAT (the standard configuration) against "w/o QAT" — a variant trained without quantization simulation in the forward pass, then evaluated with post-training quantization. The w/o QAT variant shows degraded quality: IQ drops from 66.64 to 65.28, OC from 21.42 to 20.66, AQ from 64.62 to 61.85, SC from 94.83 to 94.65, and VR from 0.1039 to 0.0850. The AQ drop (2.77 points) and VR drop (0.0189) are particularly notable, confirming that quantization error interacts with sparsity in ways that training-time adaptation successfully mitigates. The paper reports a 1.3× kernel speedup from quantization (Section 9.4), providing a quantitative efficiency-accuracy tradeoff for the QAT design choice.

Learnable router vs. heuristic Top-k router: Table 2 includes "Topk-router," which replaces SLA2's learned router with the heuristic pool(Q)pool(K)^T-based Top-k selection used in SLA. At 97% sparsity, this variant achieves IQ 63.66 (vs. 66.64 for learned router), OC 20.90 (vs. 21.42), AQ 62.65 (vs. 64.62), SC 94.26 (vs. 94.83), and VR 0.0876 (vs. 0.1039). The IQ of 63.66 is barely above Full Attention's 63.67, and the VR of 0.0876 is far below both SLA2 (0.1039) and Full Attention (0.1084). This is compelling evidence that the learned router is not merely a minor refinement but is essential for achieving quality at extreme sparsity — the heuristic router essentially fails at 97% sparsity, producing quality similar to or worse than the untrained Full Attention baseline. This ablation is the strongest evidence for Innovation 2 (routing as optimization, not heuristics).

Varying sparsity levels: The lower section of Table 2 shows SLA2 performance at 85%, 90%, 95%, and 97% sparsity. This serves as both a scaling analysis and an implicit ablation of the sparsity hyperparameter. The key finding is that quality degrades gracefully — no cliff-like drops — suggesting the learned routing remains effective across a range of sparsity budgets. This is non-obvious: one might expect that below some critical sparsity threshold, the linear branch would be forced to approximate too many important positions and quality would collapse. The absence of such a collapse down to 97% sparsity suggests the router successfully concentrates the sparse budget on the most structurally essential positions.

What is NOT ablated: Several important potential ablations are missing. There is no ablation of the decomposition-consistent mixing formulation (O = α ⊙ O_s + (1 − α) ⊙ O_l) against SLA's projection-based formulation (O = O_s + Proj(O_l)) in isolation — the comparison between SLA2 and SLA (Table 1) confounds the routing change AND the mixing change. One cannot determine how much of the improvement comes from the router vs. the reformulation. A targeted ablation with the learned router but SLA-style mixing (or vice versa) would disentangle these contributions. There is no ablation of the block sizes b_q and b_k (128 and 64 respectively), the temperature τ for SoftTop-k (0.1), or the number of Stage 1 training examples. There is no ablation of the two-stage training strategy itself (what if the router were trained jointly from the start with the diffusion loss?). There is no ablation comparing softmax vs. other activation functions for ϕ(·) in the linear branch. The paper's ablation suite is small but focused on the two components that distinguish SLA2 from prior work (learned router and QAT), leaving the joint contributions of routing and mixing unresolved.


Critical Assessment

Claim from the abstract: "SLA2 can achieve 97% attention sparsity and deliver an 18.6× attention speedup while preserving generation quality."

The 18.6× attention speedup is measured at the kernel level (Figure 4, 4,079 TOPS vs. 219 TOPS) and is well-supported. However, "preserving generation quality" requires careful qualification. Table 1 shows SLA2 at 97% sparsity achieving higher metrics than Full Attention on most dimensions (IQ, OC, AQ, SC), but not all — MS at 98.04 is below Full Attention's 98.95 (a 0.91 drop), and VR at 0.1039 is below Full Attention's 0.1084 (a 0.0045 drop). These differences are small and likely within the noise of the evaluation, but the claim "preserving quality" is more precisely "maintaining quality above or near the Full Attention level, with slight degradation on specific metrics." More importantly, the Full Attention baseline is not fine-tuned, while SLA2 benefits from 500 steps of fine-tuning on a curated dataset. Quality preservation relative to a fine-tuned Full Attention baseline is unknown. The paper does not claim that sparsity itself improves quality — it explicitly attributes quality gains to the fine-tuning dataset — but the presentation could mislead readers into thinking sparsity is quality-neutral when the baseline comparison is confounded by unequal training.

Claim: "SLA2 outperforms the baselines at 90% sparsity in end-to-end video quality."

This is the central empirical claim and is supported by Table 1, with one edge case: on the 14B model at 90% sparsity, SLA2's OC (20.68) is below both Full Attention (22.44) and VSA (21.27). This is a single metric on a single configuration, so it doesn't invalidate the overall pattern, but it does mean "outperforms on all metrics" is not literally true. The paper's statement in Section 9.2 — "SLA2 consistently outperforms all baselines across every video quality metric on both models" at 90% and 95% sparsity — is slightly overstated given this OC result.

Claim: "SLA2 even exceeds full attention, which is 0% sparsity."

Exceeding Full Attention on automated metrics is demonstrated in Table 1, but the claim's significance depends on the comparison's fairness. Full Attention is evaluated on the pretrained model with no fine-tuning. A fairer comparison would include a Full Attention variant fine-tuned for 500 steps on the same dataset — this would isolate how much of the quality gain comes from the fine-tuning data vs. the attention mechanism. The paper does not run this baseline. The authors' own attribution (fine-tuning dataset quality) suggests that a fine-tuned Full Attention baseline might achieve even higher quality than SLA2, meaning the "exceeds full attention" claim is partly an artifact of the experimental design. This is not a fatal weakness — the paper is transparent about the attribution — but it means the claim should be interpreted as "SLA2 with fine-tuning exceeds pretrained Full Attention" rather than "sparse attention is better than dense attention."

Genuine weaknesses in the experimental design:

  • Single dataset, single model family, no human evaluation. All results are on Wan2.1 models fine-tuned on a private 3,000-video dataset. The paper does not test on other video diffusion architectures (e.g., CogVideoX, Sora-like models), other data domains, or other modalities (image generation, language modeling). The private dataset prevents independent replication. All quality evaluation uses automated metrics (VBench, Vision Reward) — no human study is conducted, making it impossible to verify that the reported quality differences (e.g., IQ 67.70 vs. 63.67) correspond to perceptible differences in video quality. Given that the metrics show SLA2 at 90% sparsity scoring 4 points higher than Full Attention on IQ, a human evaluation would be valuable to confirm this counterintuitive result is not a metric artifact.

  • Missing fine-tuned Full Attention baseline. As discussed above, this is the most significant experimental omission. The paper compares trainable sparse methods against an untrained Full Attention baseline, conflating the effects of fine-tuning with the effects of sparsification. Including a Full Attention model fine-tuned for 500 steps would cleanly separate these effects. The paper's own explanation (fine-tuning dataset quality) implies this baseline would show quality improvements, potentially narrowing or reversing SLA2's quality advantage.

  • Missing token-level quality metrics. All reported quality metrics are frame-level or video-level (imaging quality, motion smoothness, consistency). There are no metrics that specifically measure attention quality — e.g., how well the SLA2 attention output approximates the full attention output at each layer, how the approximation error propagates through the generation process, or how routing decisions vary across layers and timesteps. The Router is trained in Stage 1 to minimize attention output MSE, but this training loss is never reported as an evaluation metric. If the learned router achieves near-zero attention approximation error at 97% sparsity, that would be stronger evidence than the indirect video quality metrics. If it achieves moderate error that the fine-tuning compensates for, that would suggest a different mechanism.

  • No variance estimates. Tables 1 and 2 report single-point metrics without standard deviations, confidence intervals, or multi-seed results. The fine-tuning uses a fixed 500 steps — it is unclear whether retraining with a different random seed would produce meaningfully different results. The quality differences between adjacent sparsity levels (e.g., IQ 67.70 vs. 67.04 from 90% to 95%) are small enough that they could be within training noise. Without variance estimates, it is impossible to determine which differences are statistically reliable.

  • The 97% sparsity comparison is asymmetric. Table 1 evaluates SLA2 at 97% sparsity but no baseline methods at this sparsity level. The claim "SLA2 at 97% sparsity outperforms baselines at 90% sparsity" compares across different sparsity levels, which is informative but masks the fact that we don't know whether SLA, VSA, or VMoBA could also achieve 97% sparsity with acceptable quality. They might perform poorly (as VMoBA does at 95% on 14B), but the asymmetry weakens the claim that SLA2 is uniquely capable at 97% sparsity.

  • Limited ablation scope. The paper ablates the router type (learned vs. heuristic) and QAT (with vs. without), but does not ablate: the mixing formulation (SLA2's α-based mixing vs. SLA's projection-based mixing), the activation function ϕ(·) in the linear branch (softmax vs. alternatives like ReLU or ELU+1), the block sizes, the temperature in SoftTop-k, the Stage 1 training duration, or the two-stage strategy itself. This makes it impossible to attribute performance gains to specific design choices — the improvement over SLA could come from the router, the mixing formulation, the training strategy, or some interaction, but the ablation study doesn't disentangle these.

Experiments that would have strengthened the paper:

  1. Fine-tuned Full Attention baseline — to separate fine-tuning effects from sparsity effects.
  2. Human evaluation on a subset of generated videos — to validate that automated metrics track perceptual quality.
  3. Attention approximation error metrics — reporting the Stage 1 MSE loss on test data after training, and measuring how much the fine-tuned model's attention outputs deviate from full attention, to characterize the mechanism.
  4. Component-wise ablation — testing SLA2's mixing formulation with SLA's heuristic router, and SLA's projection-based mixing with SLA2's learned router, to attribute the improvement.
  5. Training data scaling — testing whether the quality improvements persist with fine-tuning on the pretraining distribution rather than a curated dataset, to determine if the "exceeds full attention" result generalizes.
  6. Diverse architectures — testing on non-Wan models to verify the method is not architecture-specific.
  7. Multi-seed training — with variance estimates to establish statistical reliability.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted for in Headline Efficiency Gains

The assumption or constraint. The compute-optimal framework in SLA and the broader test-time compute scaling literature typically assumes access to ground-truth difficulty labels or, in more realistic settings, an estimate derived from extensive sampling. The paper explicitly acknowledges in Section 3.2 that "estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity." Specifically, the predicted difficulty estimation method requires generating 2048 samples per question and scoring them with the trained Process Reward Model (PRM) before any test-time compute is allocated for the actual solution attempt. This generation-and-scoring step itself constitutes a large fixed cost — roughly 2048 generations per question just to determine which strategy to use.

The consequence. In a realistic deployment, the total compute budget is difficulty_estimation_cost + strategy_execution_cost. For the difficulty estimation to be worthwhile, the savings from the optimized strategy must exceed the upfront cost of estimating difficulty. If the difficulty estimation requires 2048 generations per question, then for questions where the compute-optimal policy saves only a few hundred generations (e.g., using 16 generations of beam search instead of 64 of best-of-N, a saving of 48 generations), the net effect is a massive increase in total compute — 2048 + 16 = 2064 generations instead of 64. The paper's reported 4× efficiency gains are computed after difficulty is known, without amortizing the cost of learning it. The paper frames this as "an exploration-exploitation tradeoff" and flags it as "a key avenue for future work," but the current results provide no evidence that the net accounting works out positively.

What evidence exists in the paper. The paper acknowledges this gap explicitly in Section 3.2 but provides no experimental characterization of its magnitude. The difficulty estimation protocol (2048 samples, PRM scoring, quintile binning) is described but its cost is never integrated into the compute-optimal scaling curves in Figures 4 and 8. The oracle-binned curves (which use ground-truth correctness, requiring infinite prior knowledge) and predicted-binned curves (which use PRM scores on 2048 samples) largely overlap, suggesting that if difficulty estimation were cheap, the gains would transfer. But the paper does not report what happens when the difficulty estimation budget is deducted from the strategy execution budget — e.g., if you have a total budget of 256 generations and must spend some of it estimating difficulty, what remains for the strategy?

Mitigation status. The paper explicitly acknowledges this as a limitation (Section 3.2) and suggests "training models to directly predict difficulty from the question text" as future work, but develops no such model and provides no experimental evidence on how much cheaper difficulty estimation could become. A dynamic adaptive scheme — start with a few samples, estimate difficulty, allocate the remaining budget — is mentioned but not explored. The limitation is therefore unaddressed: the compute-optimal framework as presented requires an oracle-level difficulty estimator whose cost swamps the reported savings.


The Revision and Search Pipelines Are Never Combined, Limiting the Demonstrated Performance Ceiling

The assumption or constraint. The paper studies two complementary test-time compute mechanisms — PRM-guided search (Section 5) and iterative revisions (Section 6) — but evaluates them entirely independently. Section 8 explicitly acknowledges: "we did not experiment with PRM tree-search techniques in combination with revisions." Each mechanism shows complementary difficulty-dependent strengths: revisions excel on easy problems where local refinement of roughly correct answers suffices, while search excels on medium problems where broader exploration across solution strategies is needed. The paper demonstrates that a compute-optimal policy can select between these mechanisms per question, but never combines them — for instance, using the revision model as the proposal distribution within beam search, or using the PRM to guide which revision branches to pursue.

The consequence. The reported compute-optimal scaling curves in Figures 4 and 8 represent a lower bound on what a fully integrated system could achieve. Consider a medium-difficulty question where beam search is optimal: the base model generates candidate steps, the PRM scores them, and the search selects the most promising paths. If instead the revision model generated those candidate steps — conditioning each branch on its own previous attempts — the search would explore higher-quality candidates and likely achieve higher accuracy at the same budget. Similarly, on easy questions where revisions are optimal, using the PRM to decide when to stop revising or which revision direction to pursue could prevent the ~38% correct-to-incorrect reversion rate that Section 6.1 documents. The paper's current results cannot tell us whether combining these mechanisms would yield additive, multiplicative, or even negative interactions.

What evidence exists in the paper. The difficulty-dependent behavior differences in Figures 3 (right) and 7 (right) strongly suggest complementarity: beam search dominates on medium problems (bins 3-4) where revisions are weaker, and revisions dominate on easy problems (bins 1-2) where beam search over-optimizes. This is exactly the pattern that would motivate integration — use the revision model as the generator within beam search on medium problems, and use PRM-guided revision depth selection on easy problems. The paper does not run any experiment that combines them, so the evidence for complementarity is suggestive but unconfirmed.

Mitigation status. The paper explicitly lists this as future work (Section 8: "we did not experiment with PRM tree-search techniques in combination with revisions") but does not develop or evaluate any combined approach. This is a genuine limitation of the current work — the paper demonstrates that each mechanism works independently and that they can be adaptively selected, but leaves untested the natural next step of integrated deployment.


The FLOPs-Matched Comparison Uses a Weak Pretraining Baseline (Non-Compute-Optimal, No Test-Time Compute)

The assumption or constraint. The FLOPs-matched comparison in Section 7 compares PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters, but that larger model (1) scales only parameters while holding training data fixed, rather than following Chinchilla-optimal scaling where both data and parameters increase together, and (2) uses greedy decoding with no test-time compute augmentation of its own. The paper acknowledges the first point: "We scale parameters only... We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work" (Section 7). The second point — that the larger model uses greedy decoding — is an implicit design choice.

The consequence. Both design choices make the pretraining baseline weaker than it needs to be, which inflates the apparent advantage of test-time compute:

  1. A Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data according to the optimal ratio) would likely outperform a parameter-only-scaled model of the same FLOPs budget. The paper's comparison therefore compares test-time compute against a suboptimal use of pretraining FLOPs, not against the best possible pretrained model at that budget. The true gap between test-time compute scaling and pretraining scaling is likely narrower than reported.

  2. Giving the 14× larger model even a modest test-time compute budget — say, best-of-8 majority voting rather than greedy decoding — would create a much stronger baseline. Best-of-8 with the larger model might match or exceed the smaller model's compute-optimal performance at some budgets, since the larger model's candidate quality is higher per-sample. The paper's comparison essentially asks: "Is it better to invest all extra FLOPs in pretraining with no test-time enhancement, or to invest some in pretraining and some in test-time compute?" The answer partly depends on how efficiently the larger model can itself use test-time compute, which is never measured.

What evidence exists in the paper. The FLOPs-matched results in Figure 9 and the bar charts in Figure 1 show test-time compute outperforming the larger model on easy problems (e.g., +27.8% relative improvement on revisions at R ≪ 1) but trailing on hard problems (e.g., −52.9% at R ≫ 1 for PRM search). The crossover points — where the curves in Figure 9 place the test-time compute line above or below the larger model's star — are directly influenced by the baseline's strength. A Chinchilla-optimal larger model would be represented by a higher star in Figure 9, shifting the crossover points toward favoring pretraining. The paper provides no sensitivity analysis for how changes in the baseline affect the conclusions.

Mitigation status. The paper acknowledges the non-compute-optimal pretraining in Section 7 and frames it as a simplification for future work. The lack of any test-time compute for the larger model is not discussed or justified. This limitation is therefore partially acknowledged but the magnitude of its impact on the paper's central training-vs-inference tradeoff claim is unmeasured.


All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*)

The assumption or constraint. Every experiment in the paper uses the MATH benchmark (500 test questions) and the PaLM 2-S* (Codey) model family. The paper acknowledges this scope limitation in Section 4: "we believe this model is representative of the capabilities of many contemporary LLMs," but provides no empirical evidence across other models or tasks.

The consequence. The paper's central findings — that test-time compute efficacy depends critically on difficulty, that beam search over-optimizes on easy problems, that revisions excel on easy problems but struggle on hard ones, and that the compute-optimal policy recovers 4× efficiency gains — may not generalize to other settings. Specific concerns include:

  • Other reasoning domains: MATH consists of competition-level math problems requiring symbolic multi-step reasoning. It is unclear whether difficulty-dependent optimal strategy patterns generalize to code generation (where unit tests provide strong verification signals), logical reasoning (where search space structure differs), or scientific QA (where factual knowledge may interact differently with reasoning). The PRM over-optimization behavior (Figure 3, right) may be specific to the distribution of errors in math solutions.

  • Other model families: PaLM 2-S* has specific calibration properties, in-context learning behaviors, and error patterns. A model with different softmax temperature calibration (different confidence distributions) might exhibit different beam search over-optimization thresholds. A model with different revision capabilities (better or worse at self-correction) might show different sequential-vs-parallel optimal ratios. The paper's finding that the PRM800k dataset was "largely ineffective" for PaLM 2-S* due to distribution shift (Section 5.1) suggests that PRM quality is highly model-dependent, which in turn affects all search-based results.

  • Test set size and composition: The MATH test set has 500 questions, split into five difficulty quintiles of ~100 each. With two-fold cross-validation within each bin, strategy selection is based on ~50 questions per fold per bin — a small sample for discrete strategy selection. The paper does not report whether the computed-optimal strategy selections are stable across different random splits or whether specific bins show high variance.

What evidence exists in the paper. The paper provides no experiments on other benchmarks, other model families, or other task types. No ablation tests the effect of test set size on the computed-optimal policy's reliability. No transfer experiment tests whether a policy learned on MATH questions transfers to a different math dataset. The generalization claim is therefore entirely unsupported.

Mitigation status. The paper acknowledges the scope in Section 4 but does not address it experimentally. The statement that PaLM 2-S* is "representative" is an assertion rather than a finding. This is a standard limitation of single-benchmark papers, but it is particularly consequential here because the paper's primary contribution is the difficulty-dependent characterization of test-time compute, which may itself be benchmark-dependent.


The Revision Model Has a Substantial Correct-to-Incorrect Reversion Rate and the Fixes Are Ad-Hoc

The assumption or constraint. The revision model is fine-tuned on sequences where all in-context answers are incorrect, followed by a correct answer (Section 6.1). This means the model is never trained to recognize that the current answer is already correct and should be preserved. At inference time, when the revision chain produces a correct answer at some step, the model may subsequently "revise" it into an incorrect answer because its training distribution conditioned on the assumption that the current answer needs fixing. The paper reports that "approximately 38% of correct answers get converted back to incorrect ones" (Section 6.1).

The consequence. Sequential revision chains are inherently unstable: progress made at one step can be undone at the next. The paper mitigates this by using majority voting or verifier-based selection across all steps in the chain (picking the best answer from any step rather than always taking the last revision). However, this mitigation is imperfect: (1) it requires an external verifier to identify the best step, which may itself make errors, and (2) it doesn't prevent the model from wasting revision steps on correct answers — compute is spent revising an answer that needed no revision, potentially producing a worse output. In a budget-constrained setting where only a few revision steps are allowed, spending a step destructively revising a correct answer directly reduces the chance of ending with a correct answer.

More fundamentally, this reversion behavior reveals a training-inference mismatch: the model was trained to convert incorrect answers to correct ones, but at test time it may encounter correct answers in its context that trigger the same "convert to correct" behavior, producing incorrect outputs. The model has no concept of "this answer is already correct so I should stop" because the training data never included that scenario.

What evidence exists in the paper. The 38% reversion rate is reported in Section 6.1, though the paper does not provide a detailed breakdown (e.g., does the reversion rate vary by difficulty? By revision depth?). Figure 6 (left) shows that pass@1 improves gradually across revision steps, suggesting that the net effect of revisions is positive (more correct answers are created than destroyed), but the reversion rate implies that the improvement is dampened — if the model never reverted correct answers, the pass@1 trajectory might be steeper or reach a higher asymptote.

Mitigation status. The paper mitigates the reversion problem through within-chain selection (majority voting or verifier-based selection across all steps), which is effective at recovering the correct answer if it appeared at any point in the chain. However, this mitigation (1) adds computational overhead (the verifier must score every step), (2) does not prevent wasted compute on destructive revisions, and (3) does not address the underlying training-inference distribution mismatch. The paper does not explore training the model on correct-in-context examples (e.g., showing it how to recognize when no revision is needed), which would be a more principled solution. The authors do not flag this as a limitation requiring future work, but it represents a fundamental fragility in the revision approach.


Hard Problems Remain Essentially Unsolved — Test-Time Compute Cannot Substitute for Missing Capability

The assumption or constraint. The paper's entire framework assumes that the base model is capable of generating correct answers at some non-trivial rate, and that test-time compute amplifies this capability. For questions where the base model's pass@1 is essentially zero (difficulty bin 5, the hardest quintile), no amount of test-time compute — regardless of strategy, budget, or allocation policy — produces meaningful improvement.

The consequence. There is a hard capability boundary that test-time compute cannot cross. The paper demonstrates this clearly: in Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods (beam search, best-of-N, lookahead search) at all budgets. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy regardless of the sequential-to-parallel ratio. In Figure 9, the scaling curves for bin 5 are essentially flat near 0–5% across all test-time compute budgets, while the 14× larger model (which has more pretraining capability) achieves non-trivial accuracy. This means that for problems outside the base model's capability range, all investment in test-time compute is wasted — the system would be better off routing these questions to a larger model or to a human.

This limitation is particularly important for deployment planning: the paper shows that on the pretraining-vs-inference tradeoff, test-time compute is preferable on easy-to-medium problems, but provides no mechanism for knowing in advance whether a given problem is in the "test-time compute can help" regime or the "no amount helps" regime. The difficulty estimation procedure can identify hard problems after the fact (by generating 2048 samples and observing near-zero pass@1), but at that point the compute has already been spent. A deployment system needs to decide whether to even attempt a problem with the small model or to escalate to a larger model — and the paper's framework provides the tools to make that decision efficiently, but the decision itself still requires an upfront difficulty assessment.

What evidence exists in the paper. The bin-5 results in Figures 3 (right), 7 (right), and 9 are consistent and unambiguous: test-time compute provides essentially zero benefit on the hardest problems. The paper is transparent about this: the FLOPs-matched Section 7 explicitly notes that pretraining is preferable for hard problems (e.g., −52.9% relative disadvantage from using test-time compute instead of the larger model at R ≫ 1 for PRM search). The paper's Figure 9 placement of bin-5 scaling lines far below the 14× larger model's stars across all R values is a clear visual representation of this boundary.

Mitigation status. The paper acknowledges this limitation in its discussion of the FLOPs-matched results (Section 7 takeaway) but provides no mechanism for overcoming it. The compute-optimal policy can route hard problems to the best available strategy (which is still essentially useless), but cannot improve their accuracy. This is not a flaw in the method — it is an inherent constraint of test-time compute — but it means the paper's framework is only useful for problems within the base model's approximate capability range, and the paper provides no guidance on how to efficiently identify (without spending compute) whether a given problem falls inside or outside that range. A practical system would need an escalation policy for hard problems, which the paper does not develop.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper changes how the field should think about trainable attention sparsification for diffusion models by demonstrating that the dominant heuristic for routing between sparse and linear attention — preserving the largest attention weights — is not just suboptimal but fundamentally misaligned with the decomposition's mathematical structure. The shift is from a magnitude-centric view (large weights → exact computation, small weights → approximation) to a structure-centric view (which assignment makes the sparse component maximally sparse while making the linear component maximally low-rank?). The paper makes this shift concrete by formulating the routing decision as an optimization problem over mask quality (minimizing reconstruction error between full attention and the hybrid approximation) rather than as a selection rule over weight magnitudes.

This is a reframing with methodological consequences, not a paradigm shift. The idea of sparse + linear attention decomposition existed before SLA2 (it was SLA's stated motivation), and the idea of learnable attention masks exists in prior work (VSA, VMoBA). What SLA2 contributes is the recognition that the two ideas were previously combined in a structurally inconsistent way — SLA's formulation forced the linear branch to compensate for a multiplicative scaling error from the sparse branch, making the learning problem harder than necessary. The paper resolves this contradiction by showing that a decomposition-consistent mixing formulation (O = α ⊙ O_s + (1 − α) ⊙ O_l) enables the linear branch to focus exclusively on approximating its assigned probability mass, and that a learned router trained on attention-output reconstruction error can discover structurally-optimal masks that magnitude heuristics miss.

The ablation evidence supporting this reframing is the Topk-router ablation in Table 2: replacing the learned router with SLA's heuristic drops the Vision Reward from 0.1039 to 0.0876 and IQ from 66.64 to 63.66 at 97% sparsity. The heuristic router performs barely above the untrained Full Attention baseline (63.66 vs. 63.67 IQ), confirming that at extreme sparsity, magnitude-based selection breaks down — the largest weights are not necessarily the structurally essential ones when only 3% of positions can be preserved. The learned router, in contrast, maintains quality close to 90% sparsity configurations, demonstrating that optimizing for the overall decomposition quality rather than per-weight magnitude produces masks that scale better to extreme sparsity.

The paper also contributes a practical reframing of quantization's role in sparse attention. Rather than treating low-bit computation as a post-hoc acceleration (train in FP16, optionally quantize at inference with quality penalty), SLA2 shows that quantization-aware training — simulating low-bit arithmetic in the forward pass during fine-tuning while keeping the backward pass in FP16 — can be integrated into the sparse attention training pipeline with minimal quality cost (Table 2: w/o QAT drops VR by 0.0189 vs. with QAT). The implicit insight is that extreme sparsity amplifies quantization sensitivity: when only 3% of positions are computed exactly, each quantized dot product carries disproportionate weight in the final output, making training-time adaptation to quantization error more important than in dense attention.

The research directions this work makes more attractive include:

  • Structure-aware routing objectives: rather than magnitude-based heuristics, optimizing masks against reconstruction error (as in Stage 1 of SLA2) provides a general template for learned sparsity, applicable beyond diffusion models to any Transformer where attention can be decomposed.
  • Joint sparsity-quantization co-design: the 1.3× quantization speedup on top of sparsity speedup (Section 9.4) demonstrates these are complementary axes that can be co-optimized, motivating unified training pipelines that account for both simultaneously.
  • Sparse attention training with small curated datasets: the paper's finding that 500 steps of fine-tuning on 3,000 videos improves quality over the pretrained baseline (IQ 67.70 vs. 63.67, Table 1) suggests that trainable sparsification can piggyback on domain-specific fine-tuning to achieve quality gains while reducing inference cost — a practical workflow for practitioners with domain-specific data.

Directions this work makes less attractive:

  • Pure linear attention for video generation: the paper's explicit statement that "for video generation, linear attention alone often cannot keep quality" (Section 10) combined with the success of the hybrid approach suggests that pure linear attention is unlikely to match sparse + linear hybrids for high-fidelity video, particularly as sparsity can be pushed to 97% with minimal degradation when properly routed.
  • Training-free sparsity for extreme sparsity levels: the paper demonstrates that even fine-tuned trainable methods (VSA, VMoBA) degrade sharply at 95% sparsity (e.g., VMoBA's OC collapsing to 7.96 on 14B at 95%, Table 1), while SLA2 maintains quality at 97%. This suggests that trainable methods with learned routing are necessary to push beyond 90-95% sparsity — training-free approaches that cannot adapt attention patterns are architecturally limited at extreme sparsity.

Follow-Up Research This Work Enables

Disentangling the router contribution from the mixing formulation contribution. The most important open question is: how much of SLA2's improvement over SLA comes from the learned router vs. the decomposition-consistent mixing? The paper compares SLA2 against SLA (which differs in both components) but never tests SLA2's mixing formulation with SLA's heuristic router, or SLA's projection-based mixing with SLA2's learned router. A controlled ablation would train four variants — (heuristic router, SLA mixing), (learned router, SLA mixing), (heuristic router, SLA2 mixing), (learned router, SLA2 mixing) — under identical conditions and measure the marginal contribution of each component. If the mixing formulation provides most of the gain at moderate sparsity (90-95%) but the learned router becomes essential at extreme sparsity (97%+), that would characterize when each component matters. If the learned router provides most of the gain across all sparsity levels, the mixing formulation's contribution may be smaller than the paper implies, and the diagnosis of SLA's failure would be primarily about routing, not algebraic structure. The Experiment 1 setup is straightforward: fine-tune all four variants for 500 steps on the same 3,000-video dataset and report the full VBench matrix at 90%, 95%, and 97% sparsity.

Stress-testing the learned router on out-of-distribution attention patterns. The router is trained in Stage 1 on Q, K, V tensors from the pretrained model — before any sparsification fine-tuning. During Stage 2, the model's attention patterns change as it adapts to operating with sparse masks. The learned router is frozen after Stage 1, so it bases its routing decisions on attention patterns characteristic of the pretrained, dense model, even though the fine-tuned model's attention may shift. Does this distribution mismatch cause the router to make progressively worse decisions as fine-tuning progresses, or does the fixed routing act as a regularizer that the model successfully adapts to? A follow-up study could measure the Stage 1 reconstruction loss (MSE between full attention and SLA2 output) at the beginning vs. end of Stage 2 fine-tuning. If the loss increases substantially, it suggests the fixed router becomes stale; if it decreases or stays flat, it suggests the model's attention conforms to the router's mask selection. A comparison against a variant that periodically updates the router during Stage 2 (using intermediate model checkpoints to refresh Stage 1 training) would quantify the cost of using a fixed post-Stage-1 router.

Testing SLA2 on models with fundamentally different attention patterns. All experiments use Wan2.1, a diffusion transformer for video. The paper's claims about the router's ability to learn structure-aware masks and the mixing formulation's algebraic correctness are architecture-agnostic in principle, but untested in practice. Applying SLA2 to autoregressive language models (e.g., LLaMA-style models for long-context generation) or image diffusion models (e.g., FLUX, Stable Diffusion 3) would test whether the learned routing generalizes across attention patterns that differ substantially from video diffusion transformers. Language models have causal masks and different sparsity structure (attention sinks, local patterns); image models have different spatial attention patterns. If SLA2 achieves similar sparsity-quality tradeoffs across these architectures, it suggests the method is robust. If performance degrades sharply on certain architectures, it suggests the router's design (pooling granularity, projection dimensionality, training objective) is tuned to video-specific attention patterns. The paper's code and training infrastructure make this extension feasible — SLA2's formulation is not video-specific.

Scaling the Stage 1 training data and characterizing router sample efficiency. Stage 1 trains the router on a static dataset of Q, K, V tensors collected from the pretrained model. How many such examples are needed for the router to converge to a good mask policy? The paper does not specify the size of the Stage 1 dataset. A systematic study could vary the number of collected (Q, K, V) examples from, say, 100 to 10,000, train the router on each subset, freeze it, then fine-tune the full model and measure downstream video quality. If the router performs well with only 500 examples, it suggests the routing problem is low-dimensional and cheap to learn. If it requires 10,000+ examples, it suggests the router is learning complex attention structure that may be specific to the training distribution, and practitioners would need substantial data to deploy SLA2 on a new model or domain.

Improving the correct-to-incorrect reversion problem by training on "stop revising" examples. While this limitation originates from SLA's revision model (analyzed in the prior sections' Limitations), SLA2 inherits the same training paradigm for end-to-end fine-tuning. A natural extension is: can the router learn to detect when no sparse computation is needed (i.e., when the attention is so uniform that assigning computation to the sparse branch provides no benefit) and route accordingly? The current router always assigns exactly k% of positions to the sparse branch — there is no "skip sparse entirely" option. A follow-up could add an auxiliary output to the router that predicts a per-head confidence in the sparse branch's necessity, allowing dynamic sparsity budgets per head and per layer. Heads that the model learns can be well-approximated by pure linear attention could route 100% of computation to the linear branch, freeing the sparse budget for heads where exact computation matters more. This would require modifying the Stage 1 loss to penalize unnecessary sparse computation (encouraging the router to use as few sparse positions as possible while maintaining reconstruction accuracy) and testing whether dynamic per-head sparsity achieves better quality at the same average sparsity than uniform sparsity.

Human evaluation of videos generated at extreme sparsity. The paper's quality claims rely entirely on automated metrics (VBench, Vision Reward). At 97% sparsity, SLA2 achieves IQ 66.64 vs. 63.67 for Full Attention (Table 1, 1.3B model), suggesting the generated videos are better than the pretrained model's, but are they perceptually better? Automated metrics for video quality are known to have blind spots — they may reward sharpness or contrast while missing temporal artifacts, semantic inconsistencies, or unnatural motion that humans would notice. A targeted human evaluation study could present raters with paired videos (SLA2 at 97% sparsity vs. Full Attention, and SLA2 at 90% vs. 95% vs. 97%) and ask for pairwise preferences and artifact annotation. If humans confirm the automated metric rankings, it validates VBench as a proxy for sparsity research. If humans show systematic preferences opposite to the metrics (e.g., preferring Full Attention despite lower IQ scores), it would suggest the metrics are not reliable for extreme-sparsity evaluation, and the paper's central quality claim would need re-examination. This is particularly important given the counterintuitive result that sparsity improves quality — a human study would determine whether this is a genuine perceptual improvement or a metric artifact.


Practical Applications and Downstream Use Cases

Budget-constrained video generation on consumer GPUs. The end-to-end latency numbers in Figure 5 tell a compelling deployment story: on the Wan2.1-14B-720P model, SLA2 at 97% sparsity reduces end-to-end generation time from ~50 minutes (Full Attention) to ~12 minutes (SLA2, excluding CPU offloading overhead) — a 4.35× speedup. For practitioners with a single high-end consumer GPU (RTX5090, as used in the paper), this means a video that previously required leaving the machine running overnight can now be generated during a coffee break. The 1.3B model's end-to-end time drops from 159s to 69s (2.30× speedup), bringing it closer to interactive use. The key number for practitioners: the attention kernel speedup is 18.6× (Figure 4), but the end-to-end speedup is limited by non-attention computation (62s for 1.3B, 493s for 14B). This means SLA2's benefit is largest when attention dominates the total runtime — true for high-resolution, long-video generation — and diminishes when other components (feed-forward networks, convolutions, VAE decoding) become the bottleneck. Practitioners should assess their model's attention-to-total FLOPs ratio before adopting SLA2; if attention is <50% of total latency, the end-to-end gain will be modest.

Fine-tuning on domain-specific data with "free" efficiency gains. The paper's finding that 500 steps of fine-tuning on a 3,000-video curated dataset simultaneously improves quality (IQ +4.03 over Full Attention) and enables 97% sparsity has an immediately actionable implication: if you are already planning to fine-tune a video diffusion model on a domain-specific dataset, you can adopt SLA2 during that fine-tuning process and get both better quality (from the fine-tuning data) and dramatically lower inference cost (from sparsity) without additional training time or data requirements. The fine-tuning protocol is straightforward: initialize the router via Stage 1 (using MSE loss on attention outputs collected from the pretrained model), then proceed with Stage 2 fine-tuning using your domain-specific videos and text captions. The paper's hyperparameters (500 steps, batch size 64 for 1.3B, batch size 15 for 14B, block sizes b_q=128, b_k=64, k% matching your target sparsity) provide a concrete recipe. This workflow could be particularly valuable for creative studios, video production pipelines, or scientific visualization where the generation model needs to be adapted to a specific visual domain anyway.

Low-bit attention deployment for edge or mobile inference. The QAT component enables 1.3× additional speedup on the sparse branch through INT8/FP8 arithmetic. For edge deployment scenarios where every watt and millisecond matters — mobile video editing, on-device content creation, real-time video stylization — the combination of 97% sparsity + INT8 sparse branch provides both algorithmic and hardware acceleration. The paper's QAT recipe (simulate quantization in the forward pass, keep backward pass in FP16) is well-documented and builds on the open-source SageAttention2++ quantization scheme, making it implementable by practitioners without inventing new quantization infrastructure. A device manufacturer or mobile app developer could pre-fine-tune a video model with SLA2 + QAT for a target sparsity level, then deploy the quantized sparse kernels on device, achieving the 18.6× kernel speedup from Figure 4 plus additional efficiency from integer arithmetic. The key practical consideration: the QAT benefit depends on having low-bit tensor core support on the deployment GPU (RTX 40-series and later, Apple M-series with ANE, Qualcomm Snapdragon with Hexagon). Without hardware low-bit acceleration, the quantization provides no speedup and the quality penalty from "w/o QAT" applies.


When to Prefer This Method

The paper explicitly positions SLA2 against both training-free sparse attention methods and trainable sparse attention alternatives (SLA, VSA, VMoBA), and the experimental results provide clear decision criteria:

  • Prefer SLA2 over SLA when you need sparsity above 90% and have budget for fine-tuning. At 95% sparsity, SLA2 achieves VR 0.1023 vs. SLA's 0.0881 on the 1.3B model (Table 1) — a 0.0142 improvement — and at 97% sparsity, SLA isn't even evaluated because its heuristic routing and projection-based mixing are insufficient. The learned router and decomposition-consistent mixing become increasingly important as sparsity increases; at 90% sparsity, SLA and SLA2 are closer (VR 0.0872 vs. 0.1093), suggesting that SLA's formulation is adequate at moderate sparsity but breaks down at extreme sparsity.

  • Prefer SLA2 over VSA or VMoBA when video quality at extreme sparsity is non-negotiable. VSA's IQ drops from 59.57 at 90% to 55.50 at 95% on the 1.3B model (Table 1), and VMoBA's OC collapses from 20.82 at 90% to 7.96 at 95% on the 14B model — both lack a linear attention compensation branch to handle the positions their sparse masks discard. SLA2's linear branch provides a safety net: when the sparse mask misses an important position, linear attention can still approximate it, preventing catastrophic degradation. The empirical gap is starkest on the 14B model at 95% sparsity, where VMoBA's VR goes negative (−0.0965) while SLA2 maintains 0.1125.

  • Prefer SLA2 for 97% sparsity specifically. No baseline method is evaluated at this sparsity level. The combination of learned routing (which identifies structurally essential positions even when only 3% remain) and the linear compensation branch (which handles the remaining 97%) appears uniquely capable among the tested methods at this extreme. If your deployment scenario requires the maximum possible attention speedup and you can afford fine-tuning, SLA2 is the only demonstrated option.

  • Consider the fine-tuning cost. SLA2 requires approximately 500 steps of fine-tuning on a domain-specific dataset plus Stage 1 router initialization (which itself requires collecting Q, K, V tensors from the pretrained model and training the router with MSE loss). If you cannot afford any fine-tuning — you must use the pretrained model as-is — SLA2 is not applicable, and you would need a training-free method (which the paper shows achieves lower sparsity at equivalent quality). However, since fine-tuning itself improves quality (the "exceeds Full Attention" result), the fine-tuning cost may be justified even without considering the sparsity benefit.