Gated Linear Attention

Summary

• Flash Linear Attention (FLA): hardware/IO-aware linear attention kernel.
  • Chunk-wise block-parallel form interpolates between the parallel and recurrent form.
  • $O\left(\frac{L}{C}\left(C^{2}d+C{d}^2\right)\right)=O\left(LCd+L{d}^2\right)<O(L^{2}d)\text{when}L>d$
• Data-dependent gating mechanism for linear attention that still allows hardware-efficient chunkwise form.

Motivation

• Current implementations of linear attention lack I/O-awareness and are thus slower than highly optimized implementations of softmax attention.
• Forget gates are shown crucial in RNNs.
  • RetNet: non-data-dependent decay factor.
  • Mamba cannot parallelize training and is not hardware-efficient.
  • Mamba2 uses a more restricted gating mechanism: ${\mathbf{G}}_t={\gamma }_{t}11^{\top },{\gamma }_t\in (0,1)$ is a scalar.

Methodology

Chunk-Wise Block-Parallel

• Softmax attention can be parallelized during training at the cost of quadratic complexity.

GLA

Recurrent Form

A general data-dependent forgetting gate can be formulated as

$${\mathbf{S}}_t={\mathbf{G}}_t\odot {\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top },{\mathbf{G}}_t\in (0,1{)}^{d_k\times d_k}$$

• A naive mapping $x_t\to {\mathbf{G}}_t$ would require matrix of size $d\times d_k\times d_v$, which would be parameter-inefficient.
• Outer-product-based low-rank parameterization: ${\mathbf{G}}_t={\alpha }_t{\beta }_t^{\top }$.
• Mamba2 uses ${\mathbf{G}}_t={\gamma }_{t}11^{\top }$, where ${\gamma }_t\in (0,1)$ is a data-dependent scalar.

GLA adopts a middle ground:

$$\begin{array}{c}{\mathbf{G}}_t={\mathbf{\alpha }}_t{1}^{\top }\\ {\mathbf{S}}_t=\left({\mathbf{\alpha }}_t{1}^{\top }\right)\odot {\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top }=\text{Diag}({\mathbf{\alpha }}_{\mathbf{t}}){\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top }\\ {\alpha }_t=\sigma \left({\mathbf{W}}_{\alpha ,\text{up}}\cdot {\mathbf{W}}_{\alpha ,\text{down}}\cdot {\mathbf{x}}_t\right)\end{array}$$

flowchart LR

subgraph Recurrent
    q_proj("Q")
    k_proj("K")
    v_proj("V")
    alpha_down[/"`$$W_{\alpha^-}$$`"\]
    alpha_up[\"`$$W_{\alpha^+}$$`"/]
    sigmoid(("`$$\sigma$$`"))
    matmul(("`$$\times$$`"))
    add(("`$$+$$`"))
    matmul2(("`$$\times$$`"))
    diag("`$$\text{Diag}(\cdot)$$`")
    
    outer_product(("`$$\times$$`"))
    v_proj --> outer_product
    k_proj --> outer_product
    outer_product --> add
    alpha_down --> alpha_up
    alpha_up --> sigmoid --"`$$\alpha_t$$`"--> diag --"`$$G_t$$`"--> matmul --> add
    add --> matmul2
    q_proj --> matmul2
end

state(("`$$S_{t - 1}$$`"))
input(("`$$x_t$$`"))
output(("`$$o_t$$`"))
state_next(("`$$S_t$$`"))

input --> q_proj
input --> k_proj
input --> v_proj 
input --> alpha_down
state --> matmul
matmul2 --> output
add --> state_next