Gated Linear Attention (GLA): Mathematical Foundations, Data-Dependent Gating, Chunkwise Parallel Forms, and Sub-Quadratic Sequence Modeling

Standard Transformer architectures rely on softmax self-attention, which exhibits quadratic time and memory complexity $O(N^2)$ with respect to sequence length $N$. While FlashAttention and its successors optimize memory access patterns to execute exact softmax attention efficiently within GPU SRAM, the fundamental quadratic compute barrier remains. Moreover, autoregressive generation requires maintaining a Key-Value (KV) cache that grows linearly with sequence length, creating a memory capacity and memory bandwidth bottleneck during decoding.

Linear attention models replace the exponential softmax kernel with decomposable feature maps, reducing attention complexity to $O(N)$ and admitting a recurrent formulation with constant-size hidden states. However, vanilla linear attention suffers from two major limitations: degraded associative recall capacity compared to softmax attention, and poor hardware utilization during training.

Gated Linear Attention (GLA), introduced by Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim (ICML 2024), resolves these trade-offs. By introducing data-dependent gating into the recurrent state updates and developing an I/O-aware chunkwise parallel training algorithm, GLA achieves competitive language modeling performance against LLaMA-style Transformers while offering linear-time inference and higher training throughput.

The Quadratic Bottleneck and Linear Attention Foundations

Standard causal self-attention computes an output sequence $O \in \mathbb{R}^{N \times d_v}$ from queries $Q \in \mathbb{R}^{N \times d_k}$, keys $K \in \mathbb{R}^{N \times d_k}$, and values $V \in \mathbb{R}^{N \times d_v}$:

$$O_i = \frac{\sum_{j=1}^i \exp(q_i k_j^\top / \sqrt{d_k}) v_j}{\sum_{j=1}^i \exp(q_i k_j^\top / \sqrt{d_k})}$$

Because the softmax function couples $q_i$ and $k_j$ nonlinearly inside the summation, the $N \times N$ attention matrix cannot be factored. In autoregressive decoding, computing token $N+1$ requires computing inner products against all $N$ preceding key vectors stored in high-bandwidth memory (HBM).

Linearized Attention as Associative Memory

Linear attention replaces the softmax kernel with a feature map $\phi(\cdot)$:

$$\operatorname{Sim}(q_i, k_j) = \phi(q_i) \phi(k_j)^\top$$

Applying the associative property of matrix multiplication allows reordering the summation:

$$O_i = \phi(q_i) \sum_{j=1}^i \left( \phi(k_j)^\top v_j \right)$$

Defining the recurrent state $S_i \in \mathbb{R}^{d_k \times d_v}$ as:

$$S_i = S_{i-1} + \phi(k_i)^\top v_i, \quad S_0 = 0$$

The output at step $i$ simplifies to:

$$O_i = \phi(q_i) S_i$$

During inference, updating $S_i$ and computing $O_i$ requires $O(d_k d_v)$ arithmetic operations per step, independent of sequence length $N$. The inference memory footprint remains fixed at $O(d_k d_v)$ elements per head, eliminating KV cache growth.

The Retention and Recall Degradation Problem

Despite $O(N)$ efficiency, unconstrained linear attention struggles on tasks requiring precise key-value retrieval, such as Multi-Query Associative Recall (MQAR). In vanilla linear attention, every historical token contributes additively to $S_i$ without decay. Over long contexts, older key-value pairs accumulate, saturating the bounded state capacity and diluting retrieval signals.

Architectures like RetNet introduced fixed, data-independent exponential decay factors ($\gamma \in (0, 1)$):

$$S_i = \gamma S_{i-1} + k_i^\top v_i$$

While fixed decay prevents gradient explosion and biases attention toward recency, it cannot dynamically retain critical historical tokens while discarding irrelevant noise.

Mathematical Formulation of Gated Linear Attention

Gated Linear Attention augments the linear recurrence with a data-dependent, fine-grained gating mechanism. Instead of a scalar constant, GLA modulates each hidden feature dimension dynamically based on the input representation.

Data-Dependent Forget Gate Parametrization

Given input projection $x_t \in \mathbb{R}^{d_{in}}$, GLA computes queries $q_t$, keys $k_t$, values $v_t$, and gate logits $\alpha_t \in \mathbb{R}^{d_k}$:

$$q_t = x_t W_q, \quad k_t = x_t W_k, \quad v_t = x_t W_v$$

$$\alpha_t = \sigma\left( \frac{1}{\tau} x_t W_{\alpha 1} W_{\alpha 2} \right)$$

Here, $\sigma(\cdot)$ denotes the sigmoid activation function, $W_{\alpha 1} \in \mathbb{R}^{d_{in} \times d_{gate}}$, $W_{\alpha 2} \in \mathbb{R}^{d_{gate} \times d_k}$ define a low-rank projection bottleneck, and $\tau > 0$ is a learnable or fixed temperature parameter that sharpens gate values.

Because $\alpha_t \in (0, 1)^{d_k}$, each key dimension receives an independent decay rate between 0 and 1 at time step $t$.

The Gated Recurrent Update

The hidden state $S_t \in \mathbb{R}^{d_k \times d_v}$ evolves according to:

$$S_t = \operatorname{diag}(\alpha_t) S_{t-1} + k_t^\top v_t$$

$$o_t = q_t S_t$$

Expanding this recurrence unrolls the state into a cumulative product of historical decay vectors:

$$S_t = \sum_{j=1}^t \left( \prod_{m=j+1}^t \operatorname{diag}(\alpha_m) \right) k_j^\top v_j$$

Letting $\Lambda_{j, t} = \prod_{m=j+1}^t \alpha_m \in \mathbb{R}^{d_k}$ denote the cumulative decay vector from step $j$ to step $t$, the output at position $t$ is expressed in dual attention form:

$$o_t = \sum_{j=1}^t \left( (q_t \odot k_j)^\top \Lambda_{j, t} \right) v_j$$

This formulation reveals that GLA functions as a relative-position-weighted attention mechanism where the decay rate across token intervals is dynamic, input-conditioned, and channel-specific.

Hardware-Efficient Training: The Chunkwise Parallel Form

Evaluating the sequential recurrence step-by-step during training fails to utilize GPU Tensor Cores, which require dense matrix multiplications (GEMMs). Conversely, naive parallel materialization of the full $N \times N$ attention matrix incurs quadratic memory overhead and memory bandwidth saturation.

To achieve maximum compute density and I/O efficiency, GLA employs a hybrid chunkwise parallel formulation implemented in Triton via the FlashLinearAttention library.

Sequence Partitioning and State Factorization

The input sequence of length $N$ is divided into non-overlapping chunks of length $C$ (typically $C = 64$ or $C = 128$), yielding $U = N / C$ chunks. For chunk index $u \in \{1, \dots, U\}$:

• $Q_{[u]}, K_{[u]} \in \mathbb{R}^{C \times d_k}$
• $V_{[u]} \in \mathbb{R}^{C \times d_v}$
• $A_{[u]} \in \mathbb{R}^{C \times d_k}$, where $A_{[u], i} = \log \alpha_{(u-1)C + i}$

The chunk-level hidden state entering chunk $u$ is $S_{[u-1]} \in \mathbb{R}^{d_k \times d_v}$.

Within chunk $u$, token representations decompose into two distinct components:

1. Intra-chunk local interaction: Interactions between queries and keys located within the same chunk $u$.
2. Inter-chunk historical context: Interactions between queries in chunk $u$ and the accumulated state $S_{[u-1]}$ from all preceding chunks.

Mathematical Formulation of Chunk Computations

Let $i, j \in \{1, \dots, C\}$ index positions within chunk $u$:

Cumulative Decay Representation

The cumulative log-gate prefix sum within chunk $u$ is:

$$\Gamma_{[u], i} = \sum_{m=1}^i A_{[u], m}$$

The decay from token $j$ to token $i$ ($i \ge j$) is:

$$\lambda_{[u], i, j} = \exp\left( \Gamma_{[u], i} - \Gamma_{[u], j} \right) \in \mathbb{R}^{d_k}$$

Intra-Chunk Matrix Multiplications

The intra-chunk output is computed in parallel across the chunk using Tensor Cores:

$$M_{[u]} = \left( Q_{[u]} K_{[u]}^\top \right) \odot D_{[u]}$$

Where $D_{[u]} \in \mathbb{R}^{C \times C}$ represents the causal masked decay tensor. The local attention contribution is then:

$$O_{[u]}^{\text{intra}} = M_{[u]} V_{[u]}$$

Inter-Chunk Recurrence and Output Projection

The historical hidden state transferred across chunk boundaries updates recursively:

$$S_{[u]} = \operatorname{diag}\left( \exp(\Gamma_{[u], C}) \right) S_{[u-1]} + \left( K_{[u]} \odot \exp(\Gamma_{[u], C} - \Gamma_{[u]}) \right)^\top V_{[u]}$$

The contribution of historical chunks to the current chunk outputs is:

$$O_{[u]}^{\text{inter}} = \left( Q_{[u]} \odot \exp(\Gamma_{[u]}) \right) S_{[u-1]}$$

The complete output for chunk $u$ is the sum of both terms:

$$O_{[u]} = O_{[u]}^{\text{intra}} + O_{[u]}^{\text{inter}}$$

I/O-Aware Kernel Design in FlashLinearAttention

The FlashLinearAttention Triton kernel exploits GPU memory hierarchies by mapping chunk operations to SRAM:

• SRAM Residency: The chunk hidden state $S_{[u]}$ is maintained in on-chip SRAM across the chunk-level sequence iteration, avoiding intermediate writes and reads to global High-Bandwidth Memory (HBM).
• GEMM Acceleration: Intra-chunk operations are structured as standard matrix multiplications ($Q K^\top$ and $M V$), fully utilizing hardware Tensor Cores.
• Recomputation in Backward Pass: Rather than caching all intermediate states $S_{[u]}$ for backpropagation, the backward kernel recomputes states on-the-fly during gradient propagation, reducing peak training memory to $O(N d)$.

Structural Comparison Across Sub-Quadratic Architectures

GLA sits within a broader landscape of linear-time sequence models and State Space Models (SSMs). The key distinction lies in the state rank and gating topology.

• Softmax Attention (Transformer / FlashAttention): State size is $O(N d)$, compute per token during inference is $O(N d)$, training complexity is $O(N^2 d)$, associative recall is exact.
• Linear Attention (Katharopoulos et al.): State size is $O(d^2)$ (matrix state), compute per token during inference is $O(d^2)$, training complexity is $O(N d^2)$, gating is none (additive accumulation), associative recall is weak.
• RetNet (Retention): State size is $O(d^2)$ (matrix state), compute per token during inference is $O(d^2)$, training complexity is $O(N d^2)$, gating is data-independent scalar decay ($\gamma$), associative recall is moderate.
• Mamba (Selective SSM): State size is $O(d \cdot d_{state})$ (vector state per channel), compute per token during inference is $O(d \cdot d_{state})$, training complexity is $O(N d \cdot d_{state})$, gating is data-dependent continuous discretization ($\Delta, A$), associative recall is strong.
• Mamba-2 (State Space Duality / SSD): State size is $O(d \cdot d_{state})$ (1D scalar semi-separable), compute per token during inference is $O(d \cdot d_{state})$, training complexity is $O(N d)$, gating is 1-scalar data-dependent decay per block, associative recall is strong.
• Gated Linear Attention (GLA): State size is $O(d_k d_v)$ (2D full matrix state), compute per token during inference is $O(d_k d_v)$, training complexity is $O(N d_k d_v)$, gating is data-dependent channel-wise vector gate ($\alpha_t \in \mathbb{R}^{d_k}$), associative recall is high.

Empirical Performance and Benchmarks

Experimental evaluations across language modeling benchmarks demonstrate the empirical impact of data-dependent matrix gating.

Language Modeling Perplexity

On standard pre-training runs evaluated at equal parameter counts and token budgets:

• 340M Parameter Models (15B Tokens Pre-training):
• Transformer++ baseline: 28.39 perplexity, 41.2 zero-shot accuracy
• RetNet: 32.33 perplexity, 41.0 zero-shot accuracy
• Mamba: 28.39 perplexity, 41.8 zero-shot accuracy
• GLA Transformer: 28.65 perplexity, 41.5 zero-shot accuracy
• 1.3B Parameter Models (100B Tokens Pre-training):
• Transformer++ baseline: 16.85 perplexity, 50.9 zero-shot accuracy
• RetNet: 18.64 perplexity, 48.9 zero-shot accuracy
• Mamba: 17.06 perplexity, 50.0 zero-shot accuracy
• GLA Transformer: 17.22 perplexity, 51.0 zero-shot accuracy

Associative Recall and Synthetic Needle Retrieval

On synthetic Multi-Query Associative Recall (MQAR) and selective retrieval tasks, data-independent models (such as RetNet) degrade sharply as sequence length scales due to uniform state decay. GLA maintains high recall accuracy by selectively assigning near-zero decay to key memory slots containing factual associations, matching or exceeding Mamba on synthetic memory benchmarks.

Context Length Extrapolation

Because GLA maintains a 2D matrix recurrent state modulated by normalized bounded gates ($\alpha_t \in (0, 1)$), the spectral radius of the state transition operator is strictly bounded by 1. Consequently, models trained on 2K context windows extrapolate to sequences exceeding 20K tokens without experiencing numerical overflow or perplexity explosion.

Training and Inference Throughput

Benchmarked against FlashAttention-2 and Mamba:

• Training Speed: On sequence lengths of 4K to 16K, FlashLinearAttention achieves 1.5x to 2.5x higher training throughput than pure recurrent implementations and matches or exceeds FlashAttention-2 throughput due to chunkwise tiling.
• Inference Decoding Speed: For single-token autoregressive generation, GLA achieves constant $O(1)$ latency per step with zero KV cache allocation, delivering up to 3x higher decoding throughput at sequence lengths above 8K compared to standard Transformer baselines.

Architectural Trade-Offs and Limitations

While Gated Linear Attention narrows the expressivity gap between linear-time architectures and softmax attention, several engineering considerations remain:

1. State Capacity Limits: A matrix state $S \in \mathbb{R}^{d_k \times d_v}$ has a fixed numerical capacity. For complex in-context retrieval tasks requiring simultaneous retrieval of hundreds of arbitrary key-value bindings, softmax attention still maintains an expressivity advantage by storing exact historical states.
2. Chunk Size Tuning: Chunkwise training performance depends on selecting optimal chunk dimensions ($C$) matched to GPU shared memory limits. Extremely small chunks increase inter-chunk state passing overhead, while oversized chunks increase intra-chunk quadratic compute costs.
3. Head Configuration: Similar to Multi-Head Attention, GLA benefits from Multi-Head Gated Linear Attention configurations where independent gate projections govern separate representation subspaces.

Sources

• Gated Linear Attention Transformers with Hardware-Efficient Training (Yang et al., ICML 2024 / arXiv:2312.06635)
• FlashLinearAttention Library GitHub Repository
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention (Katharopoulos et al., ICML 2020)
• Retentive Network: A Successor to Transformer for Large Language Models (Sun et al., 2023 / arXiv:2307.08621)
• Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Gu and Dao, 2023 / arXiv:2312.00752)
• Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Dao and Gu, ICML 2024 / arXiv:2405.21060)