SLA2: Sparse-Linear Attention with Learnable Routing and QAT

SLA2 is a sparse-linear attention method for diffusion models that addresses two limitations of the prior SLA approach: a heuristic split between sparse and linear branches, and a mismatch between SLA's output and the intended sparse-linear decomposition.

SLA2 introduces a learnable router that dynamically selects which attention computations use sparse or linear attention, a learnable ratio to combine the branches, and quantization-aware training (QAT) for low-bit attention to reduce quantization error.

Experiments on video diffusion models (Wan2.1-1.3B and Wan2.1-14B) show that SLA2 achieves 97% attention sparsity and an 18.6× attention speedup while preserving or even improving generation quality compared to baselines at lower sparsity levels.

The method uses a two-stage training strategy: first, the router and ratio are initialized using a differentiable SoftTop-k operator; second, the full diffusion model is fine-tuned end-to-end.

The learnable router uses pooled queries and keys with learnable projections to generate a mask, and the forward pass is implemented efficiently on top of FlashAttention-style kernels.

Ablation studies confirm the benefits of QAT and the learnable router, and show that lower sparsity levels yield better performance.