🤖 AI

SLA2: Sparse-Linear Attention with Learnable Routing and QAT

This paper proposes SLA2, an enhanced Sparse-Linear Attention framework for video diffusion models that improves upon the original SLA by introducing a learnable routing mechanism, a more faithful attention formulation, and quantization-aware fine-tuning to achieve significant speedups and sparsity while preserving generation quality.

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

Published 2026-02-16

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to direct a massive, chaotic movie scene with thousands of actors (the data) all shouting at once. In the world of AI video generation, this "shouting" is called Attention. The AI needs to figure out which actors are important to look at right now and which ones can be ignored.

The old way of doing this (Full Attention) is like having a director who stops to listen to every single actor in the room before making a decision. It's accurate, but it's incredibly slow and exhausting, especially for long movies (videos).

To fix this, researchers invented SLA (Sparse-Linear Attention). Think of SLA as a smart assistant who splits the room in two:

The "Star" Group: The actors with the loudest voices (high attention). The assistant listens to them closely.
The "Background" Group: The actors whispering. The assistant just glances at them quickly using a shortcut (Linear Attention).

The Problem with the Old Assistant (SLA):
The paper argues that the old assistant was a bit clumsy.

The Heuristic Mistake: It decided who was a "Star" based on a rigid rule (like "if you are louder than 50 decibels, you're a star"). Sometimes, a quiet actor is actually crucial for the plot, but the rule ignored them.
The Math Glitch: When it tried to combine the "Star" list and the "Background" list, the math didn't quite add up. It was like trying to mix a cup of coffee and a cup of milk, but the final drink tasted weird because the ratios were off. The assistant had to add a "fix-it" layer (a projection) to make it taste right, which was inefficient.

Enter SLA2: The Super-Intelligent Director
The authors propose SLA2, a new system that fixes these issues with three main upgrades:

1. The "Smart Router" (Learnable Routing)

Instead of using a rigid rulebook, SLA2 has a Smart Router. Imagine a traffic cop who doesn't just look at speed limits but actually learns the flow of traffic.

How it works: The router looks at the actors (the data) and dynamically decides, "Okay, this actor needs the full spotlight, but that one can just be a background extra."
The Analogy: It's like a DJ who doesn't just play the top 10 hits; they listen to the crowd and mix the perfect song for the moment. It learns to split the work perfectly between the "detailed" part and the "shortcut" part.

2. The "Perfect Mix" (Faithful Formulation)

The authors realized the old way of mixing the "Star" and "Background" groups was mathematically messy.

The Fix: SLA2 uses a Learnable Ratio. Imagine you are making a smoothie. Instead of guessing how much fruit and how much yogurt to add, you have a dial that you can turn to get the exact perfect blend every time.
The Result: The AI no longer needs to guess or add "fix-it" layers. It naturally combines the detailed attention and the shortcut attention in a way that is mathematically perfect, preserving the quality of the video.

3. The "Low-Bit" Shortcut (Quantization-Aware Training)

This is the speed booster.

The Concept: Imagine writing a novel. You can write it with a fancy fountain pen (high precision, slow) or a quick pencil sketch (low precision, fast). Usually, the sketch looks messy.
The Trick: SLA2 uses Quantization-Aware Training (QAT). This is like practicing writing with the pencil while you are learning the story. The AI gets used to the "pencil sketch" style during training so that when it actually writes the final video, the sketch looks just as good as the fountain pen version, but it's done much faster.

The Results: A Movie Magic Trick

The paper tested this on video generation models (like making AI videos from text descriptions).

Speed: SLA2 is 18.6 times faster than the old, slow method. It's like turning a 10-hour movie production into a 30-minute one.
Efficiency: It ignores 97% of the unnecessary data (Sparsity).
Quality: Surprisingly, the videos look better or at least just as good as the slow, perfect method. In fact, at 97% speed, it still beat other "fast" methods that were only 90% fast.

In Summary:
SLA2 is like upgrading a movie director from someone who reads every script line by line to a genius director who knows exactly which lines matter, mixes them perfectly, and can rehearse with a cheap pencil so they can film the final scene at lightning speed without losing any quality.

1. Problem Statement

The paper addresses the computational inefficiency of Diffusion Models, particularly in video generation, where the self-attention mechanism scales quadratically with sequence length ( $O(N^2)$ ). While Sparse-Linear Attention (SLA) was proposed as a solution by combining sparse attention (for high-magnitude weights) and linear attention (for low-magnitude weights), the authors identify two critical limitations in the original SLA formulation:

Formulation Mismatch: The original SLA output does not mathematically align with the theoretical decomposition of attention into sparse ( $P_1$ ) and linear ( $P_2$ ) components. Specifically, the sparse branch renormalizes probabilities, introducing a row-wise scaling factor ( $\alpha$ ) that creates a mismatch. SLA attempts to compensate for this via a learnable projection, which is suboptimal and difficult to train.
Heuristic Routing: The original SLA uses a heuristic "Top-k" split based on the magnitude of attention weights derived from pooled queries and keys. This static, non-learnable split is not optimal; it may assign computations to the wrong branch (e.g., moving weights that could be low-rank to the sparse branch), failing to maximize sparsity without degrading quality.
Quantization Limits: Existing methods do not fully leverage low-bit quantization for further speedups without introducing significant accuracy degradation.

2. Methodology: SLA2

The authors propose SLA2, a refined sparse-linear attention architecture that introduces three core innovations:

A. Faithful Sparse-Linear Formulation

Instead of the original SLA's additive projection, SLA2 reformulates the attention output to directly match the theoretical decomposition $P \approx P_1 + P_2$ .

Learnable Ratio ( $\alpha$ ): The model learns a vector $\alpha$ (values between 0 and 1) to dynamically weight the contributions of the sparse and linear branches.
Equation: The output is computed as:
$O = \alpha \odot O_s + (1 - \alpha) \odot O_l$
Where $O_s$ is the sparse attention output and $O_l$ is the linear attention output. This formulation eliminates the need for a compensatory projection layer and ensures the output remains row-normalized, preventing magnitude drift.

B. Learnable Router ( $R$ )

To replace the heuristic Top-k split, SLA2 introduces a Learnable Router that dynamically decides which tokens should be processed by the sparse branch.

Inputs: The router takes pooled Query ( $\bar{Q}$ ) and Key ( $\bar{K}$ ) matrices as input. It does not use Value ( $V$ ) because attention weights are determined solely by $Q$ and $K$ .
Mechanism: It applies learnable linear projections ( $proj_q, proj_k$ ) to the pooled inputs, computes a compressed attention score, and uses a SoftTop-k operator during training to generate a differentiable mask $M$ .
Efficiency: By pooling adjacent tokens and using a compressed mask, the router avoids $O(N^2)$ complexity, making it scalable for long video sequences.

C. Quantization-Aware Training (QAT)

To achieve further acceleration, SLA2 integrates low-bit attention (e.g., INT8/FP8) specifically into the sparse branch.

Strategy: During the forward pass, quantization is applied to $Q, K, P,$ and $V$ . However, the backward pass remains in FP16 using the original full-precision inputs and outputs.
Benefit: This allows the model parameters to adapt to quantization noise during training, minimizing accuracy loss while enabling significant inference speedups.

D. Two-Stage Training Strategy

To ensure stability and consistency between training and inference:

Stage 1 (Initialization): The router ( $R$ ) and mixing ratio ( $\alpha$ ) are trained using a Mean Squared Error (MSE) loss against Full Attention outputs. A differentiable SoftTop-k is used here to allow gradient flow.
Stage 2 (Fine-tuning): The entire diffusion model is fine-tuned end-to-end using the standard diffusion loss. The router switches to hard Top-k (matching inference logic) while $\alpha$ continues to be optimized.

3. Key Contributions

Theoretical Correction: Identified and fixed the mathematical mismatch in original SLA, proposing a decomposition-consistent formulation that removes the need for error-compensating projections.
Learnable Routing: Replaced heuristic weight-based splitting with a trainable router that optimizes the sparse/linear split based on an explicit error minimization objective.
QAT Integration: Successfully integrated low-bit quantization into sparse attention via QAT, achieving high sparsity without quality degradation.
Efficiency: Demonstrated that combining learnable routing, faithful formulation, and quantization yields massive speedups while maintaining or improving generation quality.

4. Experimental Results

The method was evaluated on Wan2.1 video diffusion models (1.3B and 14B parameters) using a private video dataset.

Sparsity & Speed:
- Achieved 97% attention sparsity.
- Delivered an 18.6× speedup in attention runtime compared to Full Attention (FlashAttention).
- On the 14B model, end-to-end video generation latency was reduced by 4.35×.
Quality:
- At 97% sparsity, SLA2 outperformed baseline methods (including Full Attention) on multiple quality metrics (Imaging Quality, Aesthetic Quality, Motion Smoothness) on the 1.3B model.
- On the 14B model, SLA2 maintained quality comparable to Full Attention, significantly outperforming other sparse methods (SLA, VSA, VMoBA) at lower sparsity levels (90-95%).
Ablation Studies:
- Removing QAT resulted in a noticeable drop in video quality, confirming its necessity.
- The Learnable Router significantly outperformed the static Top-k router used in the original SLA.

5. Significance

SLA2 represents a significant advancement in efficient video generation. By moving from heuristic, static attention splitting to a learnable, mathematically grounded, and quantization-aware approach, it bridges the gap between theoretical efficiency and practical performance. It demonstrates that diffusion models can operate at extreme sparsity levels (97%) with massive computational savings (18.6×) without sacrificing the high fidelity required for professional video generation. This makes high-resolution video generation more accessible on consumer-grade hardware.