SLA2: Sparse-Linear Attention v2 for Diffusion Transformers

algorithm 2602.12675
sparse-attentionlinear-attentiondiffusion-transformervideo-generationrouting-fix
← Synthesis ← algorithm Survey Source ↗

§1 TL;DR #

SLA2 修正 SLA 的 renormalization mismatch（用可学习 α 替代 projection）并引入 learnable router + QAT，在视频扩散模型上实现 97% attention sparsity / 18.6× attention speedup 且质量超越 full attention。

§2 痛点 / 方法 / 结果 #

Q1 痛点 #

SLA (Sparse-Linear Attention) 将注意力分解为 sparse + linear 两支，但存在两个系统性缺陷：

  1. Renormalization mismatch: sparse softmax 对 masked positions 重新归一化产生行级缩放因子 $\alpha$，使得 $P_1 = \alpha \cdot P_s$ 而非 $P_1 = P_s$。SLA 用一个额外 linear projection 来 jointly 补偿 (a) linear component 和 (b) 缩放误差，学习负担过重。
    1. Heuristic routing: 基于 attention weight magnitude 的 Top-k 选择不保证最优 sparse/low-rank 分解——可能把能增加 $P_2$ rank 的 entry 错误分配给 sparse branch。

      2. Q2 方法 #

      组件Before (SLA)After (SLA2)
      Sparse-linear mixing$O = O_s + \mathrm{Proj}(O_l)$ — projection 要同时补偿 scaling 和 linear$O = \alpha \odot O_s + (1-\alpha) \odot O_l$ — convex combination, 无 projection
      RoutingHeuristic: $\mathrm{softmax}(\bar{Q}\bar{K}^\top)$ → Top-kLearnable: $\mathrm{proj}_q(\bar{Q}) \cdot \mathrm{proj}_k(\bar{K})^\top$ → Top-k
      Acceleration—Sparse branch: QAT (INT8/FP8 forward, FP16 backward)
      Training单阶段 fine-tune两阶段: Stage 1 init router+α on QKV data; Stage 2 end-to-end diffusion loss

      核心技术壁垒: $\alpha = P_1 \mathbf{1}$（行概率质量）的精确公式化使得 sparse 和 linear 分支可用 convex combination 直接重组——消除 projection 后线性 branch 只需补偿 $(1-\alpha)$ 比例的 output，学习难度大幅降低。

      Q3 结果 #

      • 97% sparsity: 18.6× attention speedup, 2.30× (1.3B) / 4.35× (14B) end-to-end speedup
      • 质量超越 Full Attention: Wan2.1-1.3B IQ 66.64 vs 63.67 (+4.7%), VR 0.1039 vs 0.1084
      • 97% SLA2 outperforms all baselines at 90% sparsity
      • VMoBA 在 14B/95% 时 catastrophic failure (IQ 21.27)，SLA2 保持稳定

      §3 架构 / 方法图 #

      flowchart TB subgraph Router["Learnable Router R(Q,K)"] Q_pool["pool(Q) → Q̄ ∈ R^{N/b_q × d}"] K_pool["pool(K) → K̄ ∈ R^{N/b_k × d}"] Q_pool --> Pq["proj_q(Q̄)"] K_pool --> Pk["proj_k(K̄)"] Pq --> Pc["P_c = softmax(proj_q · proj_k^T / √d)"] Pk --> Pc Pc --> TK["Top-k(k%, P_c) → M_c"] TK --> M["Expand M_c → M ∈ {0,1}^{N×N}"] end subgraph Sparse["Sparse Branch (M=1 positions)"] QKV1["Q, K, V"] --> S_attn["softmax(QK^T/√d ⊙ M)V"] S_attn --> Quant["QAT: INT8 forward, FP16 backward"] Quant --> Os["O_s"] end subgraph Linear["Linear Branch (M=0 positions)"] QKV2["Q, K, V"] --> Lin["norm(φ(Q)φ(K)^T ⊙ (1-M))V"] Lin --> Ol["O_l"] end M --> Sparse M --> Linear Os --> Mix["O = α ⊙ O_s + (1-α) ⊙ O_l"] Ol --> Mix Alpha["Learnable α ∈ [0,1]^{N/b_q}"] --> Mix

      §4 作者证明 #

      符号表 #

      符号含义维度
      $Q, K, V$Query, Key, Value$\mathbb{R}^{N \times d}$
      $S = QK^\top/\sqrt{d}$Attention scores$\mathbb{R}^{N \times N}$
      $P = \mathrm{softmax}(S)$Full attention weights$\mathbb{R}^{N \times N}$
      $M$Binary sparse mask$\{0,1\}^{N \times N}$
      $P_1 = P \odot M$Mask-selected part$\mathbb{R}^{N \times N}$
      $P_s = P_1 / \alpha$Sparse attention (renormalized)$\mathbb{R}^{N \times N}$
      $\alpha = P_1 \mathbf{1}$Row-wise probability mass$\mathbb{R}^{N \times 1}$
      $\phi(\cdot)$Linear attention activation (softmax)—
      $\mathrm{proj}_q, \mathrm{proj}_k$Router projections$\mathbb{R}^{d \times d}$
      $b_q, b_k$Block sizes128, 64

      方程物理意义 #

      Eq. 9: $P_1 V = \alpha \odot O_s$ — sparse softmax 的重新归一化使输出多了一个行级 scale $\alpha$。这是 SLA 中 projection 必须补偿的根本原因。

      Eq. 11: $P \approx \alpha \odot P_s + (1-\alpha) \odot P_l$ — full attention 可精确表达为 sparse 和 linear 的 convex combination（$\alpha + (1-\alpha) = 1$），行归一化自动保持。

      Eq. 17 SoftTop-k: $\sigma((P_c)_{ij}/\tau + \lambda_i)$ — sigmoid 做 differentiable relaxation，$\lambda_i$ 通过 binary search 保证每行 sum = k%，使 gradient 可流过 router。

      6 项检查 #

      1. Decomposition 正确性: Eq. 5-9 链式推导，$P = P_1 + P_2$, $P_1 = \alpha P_s$ — 代数验证无误
      2. Convexity: $\alpha \in [0,1]$ 保证 mixing 是 valid convex combination
      3. SoftTop-k 约束: binary search 精确求解 $\lambda_i$ 使 row sum = k%
      4. QAT 正确性: forward 量化 + backward FP16 = standard STE (Straight-Through Estimator)
      5. $\alpha$ 初始化: Stage 1 用 MSE loss 直接优化 → 收敛到接近 true $P_1 \mathbf{1}$
      6. 无收敛定理: heuristic two-stage training, empirically validated

        7. §5 实验与数据 #

        训练配置 #

        StagePurposeDataStepsModel
        Stage 1Init router + αQ,K,V tensors from all layers/timesteps—Wan2.1-1.3B/14B
        Stage 2E2E diffusion fine-tune3000 private videos (~5s each), Qwen3-VL-Flash captions500 stepsWan2.1-1.3B (bs=64) / 14B (bs=15)

        关键实验 (Table 1) #

        Wan2.1-1.3B-480P @ 97% sparsity:

        MethodIQ↑OC↑AQ↑VR↑FLOPs
        Full Attention63.6720.2764.410.108452.75T
        VMoBA @90%65.3120.8264.140.09365.28T
        SLA @90%63.1020.8864.340.08725.40T
        SLA2 @97%66.6421.4264.620.10391.82T

        Wan2.1-14B-720P @ 95% sparsity:

        MethodIQ↑VR↑FLOPs
        Full Attention68.010.1238292.6T
        VMoBA @95%21.27-0.096514.63T
        SLA @95%64.430.107814.87T
        SLA2 @95%69.020.112515.11T

        Ablation (Table 2, 1.3B @ 97%) #

        VariantVR↑Δ vs SLA2
        w/o QAT0.0850-18%
        Topk-router (heuristic)0.0876-16%
        SLA2 (full)0.1039baseline

        Efficiency (Fig. 4-5) #

        • Kernel speedup @ 97%: 18.7× over FlashAttn2 (RTX5090)
        • End-to-end Wan-1.3B: attention 97s → 7s (13.9×), overall 2.30×
        • End-to-end Wan-14B: 4.35× overall

        §6 论证链 #

        Step论据证据结论
        1SLA 的 sparse branch 存在 renormalization mismatch$P_s = P_1/\alpha$ (Eq. 7-8), projection 需同时补偿 scaling + linearSLA formulation 非最优
        2Convex combination $\alpha \odot O_s + (1-\alpha) \odot O_l$ 消除 mismatchEq. 9-12 推导; 无需 extra projection更 faithful 的分解
        3Heuristic routing 不保证最优 sparse/low-rank splitAblation: Topk-router VR=0.0876 vs learned router 0.1039Learnable projections 显著改善分配质量
        4QAT 进一步加速 sparse branchAblation: w/o QAT VR drops to 0.0850; QAT 提供 ~1.3× kernel speedupQAT 同时改善质量(通过 fine-tune)和速度
        5三组件协同: 97% sparsity 质量仍超 baselines @90%Table 1 across both modelsSLA2 是 sparse attention 在 video diffusion 的 Pareto-optimal

        §7 实现 cross-reference #

        代码: 论文使用 SLA 的 official open-source implementation 作为 baseline; SLA2 自身代码状态未明确声明。

        关键实现细节:

        1. K smoothing: $K = K - \mathrm{colmean}(K)$ 在 forward pass 开头执行（Algorithm 2 line 1），stabilize linear attention 的 kernel feature map。
        2. Block-wise IO pattern: sparse blocks 走 FlashAttention-style tiled computation (仅计算 $M_c[i,j]=1$ 的 block pairs); linear blocks 预计算 $h_j = (K_j^\phi)^\top V_j$ 然后 accumulate — 内存访问模式对 GPU 友好。

          3. Backward pass: 手动推导 gradient（Appendix A, Algorithm 3），linear branch 的 $dH_i, dZ_i$ 预计算后主循环仅需一次 matrix addition。

          [实现未公开]