---
title: 'SLA2: Sparse-Linear Attention with Learnable Routing'
url: https://www.emergentmind.com/topics/sla2
type: topic
---

# SLA2: Sparse-Linear Attention with Learnable Routing

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 [2602.12675].

## 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:

$$
O = \alpha \odot (P_s V) + (1-\alpha) \odot (P_l V)
$$

Unlike prior SLA approaches that apply a magnitude heuristic to select sparse versus linear computation, SLA2's $\alpha$ 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 $M$ is generated by a blockwise compression over $Q,K$, followed by differentiable top-$k$ selection:

1. **Pooling**: $Q$ and $K$ are pooled along rows and columns using block sizes $b_q, b_k$ to yield $\bar{Q}\in\mathbb{R}^{N/b_q\times d}$ and $\bar{K}\in\mathbb{R}^{N/b_k\times d}$.
2. **Projected scores**: $\bar{Q}W_q$ and $\bar{K}W_k$ ($W_{q,k}\in\mathbb{R}^{d\times d}$ learnable), producing $P_c = (\bar{Q} W_q) (\bar{K} W_k)^T / \sqrt{d}$.
3. **Masking**: At inference, hard Top-$k$ per row: $M_c[i,j]=1$ if $P_c[i,j]$ is among top $k\%$, else $0$. Training uses SoftTop-$k$ with a temperature $\tau$ and Lagrange multiplier $\lambda_i$ to enforce sparsity constraints, allowing gradient flow.

The mask $M$ is expanded back to $\{0,1\}^{N\times N}$ for sparse computation. This mechanism spatially pools attention allocation, scales formally to large $N$, 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**: $Q,K,P,V$ quantized to INT8/FP8 for all sparse-branch steps; s, the scale, accompanies each tensor.
- **Attention computation**: $S_q = \hat{Q}\hat{K}^T/\sqrt{d}$, $P_s$ and $O_s$ 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 $\alpha$—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 $P_1 = P\odot M$ with row sums $\alpha$, $P_s = P_1/\alpha$ ($\alpha$ is effectively the marginal importance of the sparse branch). SLA2 reconstructs the attention as $P = \alpha \odot P_s + (1-\alpha)\odot P_l$.
- This guarantees the combined output $O$ matches the formal mixture of masked softmax and linear attention.

**Complexity**: For $k\%$ sparsity, sparse branch cost is $k\% N^2 d$, linear branch is $N d^2$, yielding total FLOPs $k\% N^2 d + N d^2$. At $k=3\%$, 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**: $18.6\times$–$18.7\times$ over FlashAttn2 @ full attention, $11.7\times$ VMoBA, $2.6\times$ VSA at 95%.
  - **End-to-end video latency**: 1.3B: 97s$\to$7s (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 $O(N^2)$ 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 [2602.12675].

Source: https://www.emergentmind.com/topics/sla2