SLA2: Sparse-Linear Attention with Learnable Routing

• SLA2 is a transformer attention mechanism combining sparse and linear branches through dynamic learnable routing and quantization-aware tuning.
• It employs a novel sparse–linear decomposition with direct α-mixing to ensure row normalization and maintain quality with 97% sparsity.
• Empirical results in video diffusion show significant speedups and metric improvements over dense attention, validating its efficiency.

SLA2 refers to “Sparse-Linear Attention with Learnable Routing and QAT”—a transformer attention mechanism that unifies sparse and linear attention branches through dynamic, learned routing and quantization-aware fine-tuning, achieving high compute sparsity and substantial acceleration in diffusion-based video generation without degrading output quality (Zhang et al., 13 Feb 2026).

1. Sparse–Linear Decomposition and Direct α-Mixing

SLA2 is grounded on a direct, mathematically principled decomposition of the attention map. For a query-key-value configuration with $Q,K,V\in\mathbb{R}^{N\times d}$, full softmax attention weights $P = \mathrm{softmax}(QK^T/\sqrt{d})$ are approximated as a convex mixture:

• Sparse branch: $P_s = \mathrm{softmax}(S \odot M)$, where $S=QK^T/\sqrt{d}$ and $M\in\{0,1\}^{N\times N}$ is a learnable mask.
• Linear branch: $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$ with $\phi(\cdot)$ an activation (e.g., softmax), $\mathrm{norm}(\cdot)$ row-normalization.

A learnable, per-query mixing ratio $\alpha\in[0,1]^{N\times 1}$ dictates the combination:

$$P \approx \mathrm{diag}(\alpha)P_s + \mathrm{diag}(1-\alpha)P_l$$

The attended output is then:

$P = \mathrm{softmax}(QK^T/\sqrt{d})$0

Unlike prior SLA approaches that apply a magnitude heuristic to select sparse versus linear computation, SLA2's $P = \mathrm{softmax}(QK^T/\sqrt{d})$1 is optimized end-to-end, eliminating bias and scale mismatches between the two branches and maintaining row normalization without further projection.

2. Learnable Blockwise Router

SLA2’s routing mask $P = \mathrm{softmax}(QK^T/\sqrt{d})$2 is generated by a blockwise compression over $P = \mathrm{softmax}(QK^T/\sqrt{d})$3, followed by differentiable top-$P = \mathrm{softmax}(QK^T/\sqrt{d})$4 selection:

1. Pooling: $P = \mathrm{softmax}(QK^T/\sqrt{d})$5 and $P = \mathrm{softmax}(QK^T/\sqrt{d})$6 are pooled along rows and columns using block sizes $P = \mathrm{softmax}(QK^T/\sqrt{d})$7 to yield $P = \mathrm{softmax}(QK^T/\sqrt{d})$8 and $P = \mathrm{softmax}(QK^T/\sqrt{d})$9.
2. Projected scores: $P_s = \mathrm{softmax}(S \odot M)$0 and $P_s = \mathrm{softmax}(S \odot M)$1 ($P_s = \mathrm{softmax}(S \odot M)$2 learnable), producing $P_s = \mathrm{softmax}(S \odot M)$3.
3. Masking: At inference, hard Top-$P_s = \mathrm{softmax}(S \odot M)$4 per row: $P_s = \mathrm{softmax}(S \odot M)$5 if $P_s = \mathrm{softmax}(S \odot M)$6 is among top $P_s = \mathrm{softmax}(S \odot M)$7, else $P_s = \mathrm{softmax}(S \odot M)$8. Training uses SoftTop-$P_s = \mathrm{softmax}(S \odot M)$9 with a temperature $S=QK^T/\sqrt{d}$0 and Lagrange multiplier $S=QK^T/\sqrt{d}$1 to enforce sparsity constraints, allowing gradient flow.

The mask $S=QK^T/\sqrt{d}$2 is expanded back to $S=QK^T/\sqrt{d}$3 for sparse computation. This mechanism spatially pools attention allocation, scales formally to large $S=QK^T/\sqrt{d}$4, and exposes mask sparsity as a tunable, learnable hyperparameter.

3. Quantization-Aware Sparse Attention (QAT)

To reduce arithmetic intensity and memory overhead, SLA2 incorporates quantization-aware fine-tuning (QAT):

• Forward pass: $S=QK^T/\sqrt{d}$5 quantized to INT8/FP8 for all sparse-branch steps; s, the scale, accompanies each tensor.
• Attention computation: $S=QK^T/\sqrt{d}$6, $S=QK^T/\sqrt{d}$7 and $S=QK^T/\sqrt{d}$8 quantized using scale parameters, then dequantized for accumulation.
• Backward pass: Gradients are computed using FP16 (or higher precision), treating quantization as identity for backward propagation.

This design allows model parameters—including the router and mixing ratio $S=QK^T/\sqrt{d}$9—to adapt to quantization artifacts, leading to negligible performance degradation even at low bit-width. Empirical ablations show removing QAT substantially degrades generation metrics.

4. Error Analysis and Computational Complexity

A direct consequence of SLA2’s design is an exact, row-normalized mixture, avoiding the scale-mismatch incurred by prior SLA methods:

• Given the masked “true” sparse map $M\in\{0,1\}^{N\times N}$0 with row sums $M\in\{0,1\}^{N\times N}$1, $M\in\{0,1\}^{N\times N}$2 ($M\in\{0,1\}^{N\times N}$3 is effectively the marginal importance of the sparse branch). SLA2 reconstructs the attention as $M\in\{0,1\}^{N\times N}$4.
• This guarantees the combined output $M\in\{0,1\}^{N\times N}$5 matches the formal mixture of masked softmax and linear attention.

Complexity: For $M\in\{0,1\}^{N\times N}$6 sparsity, sparse branch cost is $M\in\{0,1\}^{N\times N}$7, linear branch is $M\in\{0,1\}^{N\times N}$8, yielding total FLOPs $M\in\{0,1\}^{N\times N}$9. At $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$0, this corresponds to ≈97% reduction in operations compared to dense attention.

5. Empirical Evaluation in Video Diffusion Models

SLA2 has been extensively tested on video diffusion with Wan2.1-T2V-1.3B (480p) and 14B (720p) models. Metrics include VBench (IQ, OC, AQ, MS, SC) and VisionReward, reporting both quality and efficiency.

• At 97% sparsity:
  • IQ: 66.64 (1.3B), 66.93 (14B)
  • VisionReward: 0.1039–0.1149 (matches full attention)
  • Attention FLOPs: 1.82T (1.3B, 97%), 9.26T (14B, 97%) vs 52.75T for dense
  • Kernel speedup: $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$1–$P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$2 over FlashAttn2 @ full attention, $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$3 VMoBA, $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$4 VSA at 95%.
  • End-to-end video latency: 1.3B: 97s$P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$57s (13.9×); 14B: 4.35× speedup, making previously infeasible generations tractable.

Ablations show the necessity of QAT and the learned router; disabling either yields significantly reduced IQ and VR.

Model	Sparsity	IQ	VR	FLOPs (T)	Speedup (FlashAttn2)
Full Attention	0%	63.67	—	52.75	1×
SLA2 (1.3B)	90%	67.70	—	5.51	—
SLA2 (1.3B)	95%	67.04	—	2.87	—
SLA2 (1.3B)	97%	66.64	0.1039	1.82	18.6×
SLA2 (14B)	97%	66.93	0.1149	9.26	18.7×

6. Implementation Notes, Limitations, and Extensions

Implementation uses FlashAttention-style fused kernels for both masked (“sparse”) and blockwise linear branches, avoiding $P_l = \mathrm{norm}(\phi(Q)\phi(K)^T \odot (1-M))$6 operations in practice. Routing masks are computed on compressed blocks, balancing mask granularity and efficiency.

Limitations:

• Performance verified on video diffusion; transferability to other domains (e.g., NLP, general vision) is not demonstrated.
• The blockwise router may be less effective when attention locality is highly irregular.
• Mask granularity is determined by block size, trading off routing flexibility versus computational savings.

Future directions include adapting SLA2’s direct α-mixing and quantized routing infrastructure to multi-modal transformers, NLP, or large-scale foundation models, and exploring alternative sparse-linear fusion strategies to further minimize error and improve generalization.

7. Comparative Positioning and Impact

SLA2 provides a direct, optimization-driven alternative to heuristic sparse/linear splits, achieving extremely high compute sparsity without sacrificing output quality or imposing additional normalization corrections. Empirical results show consistent outperformance over baseline dense and prior sparse attention schemes at a fraction of computational cost. This suggests a general pathway for integrating learnable routing and quantization into mixed-attention architectures for efficient, scalable, and high-fidelity generative modeling (Zhang et al., 13 Feb 2026).