causal linear attention as an RNN

Topics

• linear attention

Causal linear attention can be expressed as an RNN. This applies when a token at position $t$ only attends to tokens at positions $j\le t$.

This formulation, shown by Katharopoulos et al. (2020), defines state at time $t$ using accumulated key-value (context) summary $S_t$ and accumulated key summary $Z_t$:

$$\begin{array}{rl}{S}_t& =\sum _{j=1}^t\varphi (K_j{)}^T{V}_j\in {\mathbb{R}}^{c\times d_v}\\ Z_t& =\sum _{j=1}^t\varphi (K_j{)}^T\in {\mathbb{R}}^c\end{array}$$

For more details on notation, refer linear attention math formulation. Cool observation is that these states update recurrently:

$$\begin{array}{rlr}{S}_t& =S_{t-1}+\varphi (K_t{)}^T{V}_t& (S_0=0)\\ Z_t& =Z_{t-1}+\varphi (K_t{)}^T& (Z_0=0)\end{array}$$

Output at time $t$ computed using current query $\varphi (Q_t)$ and updated state $(S_t,Z_t)$:

$${\text{Output}}_t^{1\times d_v}\approx \frac{\varphi (Q_t)S_t}{\varphi (Q_t)Z_t}$$

Standard causal attention with kv-caching (storing all previous keys/values) needs $O(N)$ memory. Causal linear attention with RNN formulation only needs $O(1)$ memory (specifically $O(c{d}_v+c))$ per step during inference) for fixed-size states $(S_t,Z_t)$, making it efficient for long sequences.

Pros:

• $O(N)$ time complexity and $O(1)$ space complexity per step during inference
• Good for long sequences

Cons:

• Sequential nature of the RNN formulation hinders parallelization during training
• The kernel function is still an approximation, so performance < standard attn

Trends:

• To mitigate training slowdown, use chunkwise processing: input sequence is divided into non-overlapping chunks. Within each chunk, computations can be parallelized. The recurrent state update is performed sequentially between chunks. Nice balance between parallelization and sequential nature required for causal modeling across chunks
• Hardware-aware implementations of linear attention using tools like Triton
• Gaining traction through libraries like flash-linear-attention

Related