SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Jintao Zhang; Haoxu Wang; Kai Jiang; Kaiwen Zheng; Youhe Jiang; Ion Stoica; Jianfei Chen; Jun Zhu; Joseph E. Gonzalez

Source-linked AI summary

SLA2: Sparse-Linear Attention with Learnable Routing and QAT

Jintao Zhang, Haoxu Wang, Kai Jiang, Kaiwen Zheng, Youhe Jiang, Ion Stoica, Jianfei Chen, Jun Zhu, Joseph E. Gonzalez

arXiv:2602.12675v1Feb 13, 2026cs.LGcs.AIcs.CV

TL;DR

SLA2 targets SLA’s heuristic branch routing and mismatch with the intended sparse-linear decomposition. It introduces learnable routing, decomposition-consistent mixing, and QAT-based low-bit attention, achieving high sparsity and speedup while preserving video generation quality. The paper also notes that the linear branch in SLA is forced to compensate for sparse-branch scaling error, complicating compensation.

Problem
SLA uses heuristic attention-weight routing and its sparse branch has a scaling mismatch with the direct sparse-linear decomposition.
Method
SLA2 learns branch routing and the combination ratio, then adds low-bit sparse attention with quantization-aware training.
Results
97% attention sparsity and 18.6× attention speedup are reported on video diffusion models while preserving video generation quality.
Takeaways & Limitations
SLA2 offers a trainable sparse-linear attention approach for efficient attention in diffusion models.
Takeaways & Limitations
SLA2’s reported Wan2.1-14B-720P latency excludes sequential CPU offloading overhead during evaluation.

Abstract

from arXiv · show

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.

1. Introduction

SLA2 addresses two limitations of Sparse-Linear Attention: heuristic routing between sparse and linear branches and a scaling mismatch with the intended decomposition. It adds learnable routing and decomposition-consistent mixing, incorporates QAT-based low-bit attention, and reports high sparsity and speedup while maintaining video quality.

Background: SLA combines sparse and linear attention to improve sparsity in diffusion models, including image and video diffusion settings.Its sparse branch is complemented by a linear-attention branch.
Limitations of SLA: SLA’s heuristic magnitude-based routing may produce a suboptimal split between sparse and linear attention.Selecting larger attention weights for the sparse branch does not necessarily yield the best decomposition.
Limitations of SLA: SLA’s sparse attention map differs from the decomposed sparse component by a row-wise scaling factor, requiring an additional projection for compensation.The analysis identifies P1 = αPs, and the projection may not fully address this mismatch.
SLA2: SLA2 learns the ratio α that combines sparse and linear branches, directly matching the intended sparse-linear decomposition.The formulation removes the need for an extra linear-branch projection and preserves row normalization.
SLA2: SLA2 trains a differentiable sparse-attention mask predictor by minimizing the approximation error relative to full attention.Learnable query and key projections allow routing to adapt the partition between branches.
Results: 97% attention sparsity and 18.6× attention runtime speedup are reported on video diffusion models while maintaining video quality.The reported result applies to Wan2.1-1.3B and Wan2.1-14B; 97% sparsity corresponds to about 96.7% computation savings after accounting for the linear branch.

3. SLA2 Design

SLA2 combines sparse and linear attention using a learnable ratio and a router-defined mask, with efficient branch-specific computation that avoids unnecessary full matrix multiplications.

3. SLA2 Design: SLA2 combines sparse and linear attention outputs using a learnable α vector.The formulation uses α ⊙ O_s + (1 − α) ⊙ O_l, with α values between 0 and 1.
3. SLA2 Design: The router R produces a mask M that determines which attention probabilities use the sparse branch.R is a learnable module, while ϕ(·) and normalization define the linear-attention computation.
3. SLA2 Design: SLA2 computes sparse attention only where M = 1, skipping other QK^T and PV operations.The sparse branch is implemented on top of FlashAttention-style computation.
3. SLA2 Design: The linear branch first computes K^T V using positions where M = 0, then multiplies Q by the result.This avoids directly computing the full QK^T matrix for the linear branch.

4. Learnable Router

SLA2 uses a learnable router to predict a sparse-attention mask from Q and K, reducing router cost through pooling and enabling differentiable training through SoftTop-k.

4. Learnable Router: The router R takes Q and K to dynamically decide which probabilities use sparse attention, independently of V.Adjacent-token similarity motivates mean pooling over consecutive query and key blocks before routing.
4. Learnable Router: Two learnable projections transform pooled Q and K before the router constructs the mask M.The projections are proj_q and proj_k, each mapping from d dimensions to d dimensions.
4. Learnable Router: Top-k sets the highest-scoring k% positions in each row to 1 and the remainder to 0.The compressed mask can be expanded to an N × N mask, while GPU kernels operate on the compressed block mask.
4. Learnable Router: Training replaces non-differentiable Top-k with SoftTop-k so gradients can propagate through routing.SoftTop-k preserves the row-wise sum constraint associated with Top-k.

5. Quantization-aware Training

SLA2 accelerates its sparse branch with low-bit attention while using quantization-aware training to adapt the model and preserve inference accuracy.

5. Quantization-aware Training: Quantization-aware training incorporates quantization effects during training, unlike post-training quantization applied after training.QAT lets model parameters adapt to quantization error and improves low-bit accuracy at inference.
5. Quantization-aware Training: SLA2 uses low-bit attention in the forward pass while retaining fully FP16 computation in the backward pass.This design provides low-bit speedups while minimizing end-to-end accuracy loss.
5. Quantization-aware Training: The low-bit forward pass quantizes Q and K before computing attention scores.The resulting tensors and scales are later used in the quantized attention computation.
5. Quantization-aware Training: SLA2 subsequently quantizes P and V, then dequantizes their product to produce O_s.The output is rescaled back to FP16 after quantized multiplication.
5. Quantization-aware Training: The quantization scheme maps FP16 tensors to low-bit tensors such as INT8 or FP8 with scales, then rescales results during dequantization.The implementation follows SageAttention2++.

6. Training with SLA2

SLA2 training initializes the router and combination ratio against full attention, then fine-tunes the diffusion model end to end using differentiable sparse routing and the branch-combination formulation.

6. Training with SLA2: SLA2 uses two training stages: initialize R and α, then fine-tune diffusion parameters end to end.The second stage optimizes the diffusion loss over Θ and α without R.
6. Training with SLA2: During stage one, Q, K, and V from attention layers across diffusion timesteps train R and α for each sparsity setting.The reported settings use k% values of 5%, 4%, and 3%.
6. Training with SLA2: SoftTop-k uses a sigmoid, temperature τ, and per-row λ_i solved by binary search to enforce the target row sum.Each row sums to k% × N/b_k.
6. Training with SLA2: Algorithm 2 partitions Q, K, and V into query and key-value blocks before constructing the compressed routing mask.It also applies the stated key smoothing and feature transformations before block processing.
6. Training with SLA2: For M = 1 blocks, SLA2 computes sparse attention scores and values; for M = 0 blocks, it accumulates linear-attention statistics.The outputs are combined as O = α ⊙ O_s + (1 − α) ⊙ O_l.
6. Training with SLA2: SoftTop-k retains the row-wise Top-k sum constraint while enabling gradient backpropagation during training.The complete training procedure uses SLA2 as the attention operator.

7. Inference with SLA2

At inference, SLA2 replaces the diffusion model’s attention modules and uses hard Top-k routing in its forward pass.

SLA2 inference replaces the diffusion model’s attention modules with the SLA2 forward pass.
Hard Top-k, rather than SoftTop-k, is used for routing during inference.

8. Insights

SLA2’s routing design uses Q and K to learn a task-adaptive sparse–linear decomposition efficiently, while two-stage training supports stable optimization and inference consistency.

Router motivation: SLA2 aims to route a low-rank portion of attention to the linear branch while making the sparse branch as sparse as possible without harming quality.
Router inputs: Q and K are sufficient routing signals because they determine attention weights, whereas V does not affect attention-map formation.
Efficient routing: Mean pooling adjacent Q and K tokens reduces the router’s cost because a full QK^T computation would have O(N^2) complexity.
Learnable routing: Learnable projections make Top-k selection task-adaptive, improving alignment with the desired sparse and low-rank decomposition over the identity-projection heuristic.
Training: Two-stage training first initializes the router and ratio, then fine-tunes the diffusion model with inference-aligned hard Top-k routing.

9. Experiments

Experiments evaluate SLA2 on Wan2.1 video diffusion models using video-quality, computational-cost, kernel-efficiency, and end-to-end latency metrics. SLA2 maintains strong generation quality while achieving high sparsity and substantial speedups.

Setup: SLA2 is fine-tuned on Wan2.1-1.3B-480P and Wan2.1-14B-720P using a private dataset of 3,000 videos.
Metrics: Video quality is evaluated with VBench dimensions and Vision Reward, while efficiency is measured using FLOPs, C/t, and end-to-end latency.
Effectiveness: At 90% and 95% sparsity, SLA2 consistently outperforms all baselines across every reported video-quality metric on both models.
Effectiveness: At 97% sparsity, SLA2 surpasses baselines at 90% sparsity and achieves 29× speedup over Full Attention.
Efficiency: 18.7× kernel speedup over FlashAttn2 is achieved at 97% sparsity, while end-to-end latency falls 2.30× on Wan-1.3B-480P and 4.35× on Wan-14B-720P.The 14B evaluation uses sequential CPU offloading, whose overhead is excluded from the reported latency.
Quantization: Quantized inference without QAT reduces video quality, while low-bit quantization provides approximately 1.3x kernel speedup.
Ablation study: The learnable router significantly outperforms SLA’s Top-k router, and lower sparsity consistently produces better performance across the tested range.

10. Related Work

Related work covers sparse and linear attention methods for accelerating Transformer-based models, including training-free and trainable sparse-attention approaches.

Sparse attention methods are categorized by whether they require training, alongside linear attention as another major acceleration approach.

11. Conclusion

SLA2 addresses SLA’s heuristic routing and decomposition mismatch with learnable routing, decomposition-consistent mixing, and quantization-aware sparse plus low-bit attention. It reaches high sparsity and speedup while preserving video generation quality.

11. Conclusion: SLA2 introduces a learnable router and decomposition-consistent mixing formulation to address SLA’s routing and decomposition limitations.The method replaces heuristic routing and directly learns the ratio combining sparse and linear attention branches.
11. Conclusion: SLA2 adds sparse plus low-bit attention through quantization-aware fine-tuning for further acceleration.Quantization-aware training helps the model adapt to low-bit quantization and improve inference accuracy.
11. Conclusion: 97% attention sparsity and an 18.6× attention speedup are achieved while preserving video generation quality.These are reported experimental outcomes for diffusion-model attention.

A. Backward Pass of SLA2

The SLA2 backward pass manually derives selected gradients while using automatic differentiation for the remaining ones. Precomputed intermediate gradients reduce the main procedure to a single matrix addition, improving computational efficiency.

A. Backward Pass of SLA2: SLA2 manually derives gradients with respect to Q, K, V, Qϕ, and Kϕ, while PyTorch automatic differentiation computes the remaining gradients.This follows the backward-pass approach used by SLA.
A. Backward Pass of SLA2: A single matrix addition forms the main backward procedure because dHi and dZi are precomputed.The algorithm accumulates dH and dZ for unselected blocks before completing the remaining gradient calculations.
A. Backward Pass of SLA2: The backward algorithm takes forward-pass quantities and output gradients as inputs, then iterates over query and key blocks to accumulate derivatives.The block computation branches according to the coarse mask Mc and returns gradients for Q, K, V, Qϕ, and Kϕ.