SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Abstract: Sparse-Linear Attention (SLA) combines sparse and linear attention to accelerate diffusion models and has shown strong performance in video generation. However, (i) SLA relies on a heuristic split that assigns computations to the sparse or linear branch based on attention-weight magnitude, which can be suboptimal. Additionally, (ii) after formally analyzing the attention error in SLA, we identify a mismatch between SLA and a direct decomposition into sparse and linear attention. We propose SLA2, which introduces (I) a learnable router that dynamically selects whether each attention computation should use sparse or linear attention, (II) a more faithful and direct sparse-linear attention formulation that uses a learnable ratio to combine the sparse and linear attention branches, and (III) a sparse + low-bit attention design, where low-bit attention is introduced via quantization-aware fine-tuning to reduce quantization error. Experiments show that on video diffusion models, SLA2 can achieve 97% attention sparsity and deliver an 18.6x attention speedup while preserving generation quality.

• The paper introduces a learnable router and mixing ratio to precisely decompose the full attention map into high-sparse and low-rank components.
• The paper employs quantization-aware training to achieve efficient low-bit attention processing while significantly reducing inference latency.
• The paper demonstrates superior performance with up to 18.7x kernel speedup and enhanced video quality even at 97% sparsity compared to baselines.

SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Motivation and Technical Innovations

Sparse-Linear Attention (SLA) demonstrated the effectiveness of combining sparse softmax and linear attention for efficient sequence modeling, particularly in diffusion-based video generation. However, SLA suffers from intrinsic limitations: it relies on heuristics for routing computation between sparse and linear branches and introduces a scale mismatch in its decomposition, yielding suboptimal performance and complicating compensation for attention errors. SLA2 systematically addresses these shortcomings through three primary innovations: a learnable router for dynamic computation split, an exact formulation of the sparse-linear decomposition via a learnable mixing ratio, and quantization-aware training (QAT) for low-bit attention acceleration.

The principal insight underpinning SLA2 is the precise decomposition of the full attention map $P$ into high-sparse ($P_1$) and low-rank ($P_2$) components, as $P = P_1 + P_2$, motivating a more faithful reconstruction: $P \approx \alpha \odot P_s + (1-\alpha)\odot P_l$ with learnable $\alpha$. This formulation eliminates the scale mismatch in SLA and avoids the necessity for compensatory projection layers. The learnable router $\mathcal{R}$ eschews the heuristic top-$k$ mechanism, instead learning optimal partitioning based directly on query-key signals, enabled by gradient backpropagation with differentiable Top-$k$ (SoftTop-$k$) during training.

Figure 1: Attention computation pipeline of SLA2.

Methodological Details

SLA2 is instantiated as follows:

• Sparse branch: $O_s = \mathrm{softmax}(QK^\top/\sqrt{d} \odot M) V$
• Linear branch: $$O_l = \mathrm{norm}(\phi(Q) \phi(K)^\top \odot (1-M)) V$$
• Final output: $O = \alpha\odot O_s + (1-\alpha)\odot O_l$

where $M = \mathcal{R}(Q, K)$ is generated by the learnable router, and $\alpha$ is a learnable vector, both trained under a canonical two-stage regime: first initializing $\mathcal{R}$ and $\alpha$ via mean squared error (MSE) minimization between full attention and SLA2 output, then fine-tuning the entire diffusion model including $\alpha$.

The router operates efficiently by pooling adjacent tokens (mean pooling) and projecting pooled $Q, K$ via learnable transformations. Differentiable Top-$k$ (SoftTop-$k$) makes the routing compatible with gradient-based optimization. During inference, hard Top-$k$ is used for routing logic alignment.

SLA2 integrates QAT in its pipeline, quantizing $Q$, $K$, $V$, and $P$ tensors in the sparse branch during forward computation. The backward pass remains in FP16 to avoid accuracy degradation from quantization, ensuring model robustness under low-bit precision.

Empirical Evaluation

Experiments were conducted on Wan2.1-T2V-1.3B-480P and Wan2.1-T2V-14B-720P video diffusion models. SLA2 significantly outperforms baseline methods (VMoBA, VSA, SLA) across all video quality metrics (Imaging Quality, Overall Consistency, Aesthetic Quality, Motion Smoothness, Subject Consistency, Vision Reward) both at 90% and 95% sparsity. Remarkably, SLA2 at 97% sparsity exhibits superior video quality even compared to full attention, attributed to fine-tuning on a higher-quality dataset. These results directly contradict conventional expectations about the trade-off between attention sparsity and fidelity.

Figure 2: Visible examples of SLA2 and baselines on Wan2.1-T2V-1.3B-480P model.

Figure 3: Visible examples of SLA2 and baselines on Wan2.1-T2V-14B-720P model.

Kernel-level and end-to-end efficiency improvements are profound. SLA2 realizes an 18.7x speedup in attention kernel execution relative to FlashAttn2, surpassing VMoBA and VSA by large margins at comparable sparsity levels. End-to-end video generation latency is markedly reduced: 13.9x speedup in attention latency and 2.3x reduction in total inference latency for the 1.3B model, and even greater proportional gains for the 14B model with CPU offloading.

Figure 4: Kernel speed of SLA2 and baselines with different sparsities.

Figure 5: End-to-end generation latency of SLA2 and baselines with different sparsities.

Ablation studies verify the necessity of QAT for maintaining video generation quality under quantization and demonstrate the superiority of the learnable router over the heuristic Top-$k$ mask. Increased sparsity (up to 97%) does not degrade quality, a result not achieved by other methods.

Insights and Theoretical Implications

SLA2's advance lies in formalizing the sparse-linear decomposition, operationalizing this via a learnable mixing ratio, and optimizing the routing split through learnable projections. The primary signal for routing -- the query and key matrices -- is justified both theoretically (attention score construction) and empirically (token distribution smoothness). The block-wise pooling increases computational efficiency without harming partition quality, and learnable projections enable adaptive, model-specific routing, improving upon the fixed heuristic split.

The two-stage training procedure is critical: initialization with SoftTop-$k$ supports stable learning by permitting differentiability, while hard Top-$k$ enforces inference-train consistency. QAT is essential for realizing practical low-bit attention inference without accuracy loss.

Practical Impact and Future Directions

SLA2 provides an effective solution for scaling attention in video diffusion transformers, offering highly sparse (97%) and low-bit acceleration (quantized attention) with no compromise in generation fidelity. The learnable router and mixing ratio formulation are directly applicable to other architectures that require attention partitioning, and the block-wise, pooled design is hardware-friendly.

Looking forward, further investigations can consider adaptive sparsity schedules, hierarchical multi-stage routers, and deeper integration of quantization into both branches. Analysis of generalization performance and transferability of SLA2-trained models to other generative modalities (e.g., text, audio) should be prioritized. Hardware-specific optimizations, including kernel fusion and asynchronous computation, could further extend SLA2's impact on real-world, high-throughput generative systems.

Conclusion

SLA2 systematically addresses the core limitations of SLA by introducing a learnable router, a scale-matched mixing formulation, and quantization-aware training, achieving superior video generation quality and unprecedented efficiency advances in diffusion models. The results validate both the theoretical motivation (exact sparse-linear decomposition) and practical implementation (block-wise learnable routing, low-bit acceleration), establishing SLA2 as an effective mechanism for attention scalability in modern video generative architectures (2602.12675).