FG$^2$-GDN: Enhancing Long-Context Gated Delta Networks with Doubly Fine-Grained Control

1. Problem Statement

While linear attention mechanisms (e.g., Gated DeltaNet, Kimi Delta Attention) offer $O(L)$ inference complexity compared to the $O(L^2)$ of softmax attention, they still face limitations in long-context modeling and associative recall.

• The Bottleneck: Existing state-of-the-art models like Kimi Delta Attention (KDA) have refined the forgetting mechanism by moving from a scalar decay gate to a channel-wise (dimension-specific) decay vector ($\alpha_t$). However, the learning rate (or writing strength) $\beta_t$ in the delta update rule remains a scalar shared across all feature dimensions.
• The Consequence: This scalar constraint limits the model's ability to perform dimension-specific adaptation. It forces all feature dimensions to update with the same "writing strength," regardless of their individual relevance to the current context, leading to suboptimal associative recall and interference in long sequences.

2. Methodology

The authors propose FG2-GDN (Fine-Grained Gated Delta Network) and its variant FG2-GDN+, which introduce doubly fine-grained control over the delta update rule.

Core Innovation: Per-Coordinate Adaptive Learning Rates

The paper draws an analogy between the delta rule and Stochastic Gradient Descent (SGD). Just as optimizers like AdaGrad and Adam use per-coordinate learning rates to handle varying curvature in parameter space, FG2-GDN replaces the scalar $\beta_t$ with a channel-wise vector $\beta_t \in \mathbb{R}^{d_k}$.

• FG2-GDN Update Rule:
  Instead of a scalar update, the key ($k_t$) and value ($v_t$) vectors are scaled element-wise by the square root of the learning rate vector:
  $$\tilde{k}_t = \sqrt{\beta_t} \odot k_t, \quad \tilde{v}_t = \sqrt{\beta_t} \odot v_t$$
  The state update becomes:
  $$S_t = (I - \tilde{k}_t \tilde{k}_t^\top) \text{Diag}(\alpha_t) S_{t-1} + \tilde{k}_t \tilde{v}_t^\top$$
  • Symmetry Preservation: Crucially, the authors absorb $\beta_t$ symmetrically into both $k$ and $v$. This preserves the Rank-1 Householder structure of the update term ($\tilde{k}_t \tilde{k}_t^\top$), which is essential for maintaining the efficiency of chunkwise parallel training algorithms (WY representation).
  • Online Learning Interpretation: This transforms the update into a gradient step with a diagonal learning rate matrix, allowing dimensions with high gradient variance to adapt faster while others remain stable.

FG2-GDN+ Variant

To further decouple control, FG2-GDN+ introduces separate vectors for keys and values ($\beta^k_t$ and $\beta^v_t$):

• $\beta^k_t$ controls the erasure strength (how aggressively old information is overwritten).
• $\beta^v_t$ controls the write strength (how strongly new information is stored).
  This allows the model to independently manage memory retention and memory injection.

Architecture

The models utilize a Hybrid Architecture (3:1 ratio), interleaving FG2-GDN layers with Multi-head Latent Attention (MLA) layers. This combines the infinite-context efficiency of linear recurrence with the precise pairwise interaction of full attention.

3. Key Contributions

1. Doubly Fine-Grained Control: The first linear attention model to introduce channel-wise adaptive learning rates ($\beta_t$) for the delta update, complementing existing channel-wise forgetting gates ($\alpha_t$).
2. Structural Compatibility: A novel symmetric parameterization that introduces per-coordinate learning rates without breaking the Diagonal-Plus-Low-Rank (DPLR) structure. This ensures the method remains compatible with existing efficient chunkwise parallel implementations (e.g., in KDA) with negligible overhead.
3. Decoupled Erasure and Writing: The FG2-GDN+ variant demonstrates that separating the scaling of keys (erasure) and values (writing) significantly improves the management of competing associations in long contexts.

4. Experimental Results

Experiments were conducted at 340M and 1.3B parameter scales on synthetic and real-world benchmarks.

• Language Modeling (Short Context):

  • FG2-GDN achieved the best average zero-shot accuracy and the lowest LAMBADA perplexity at the 1.3B scale (13.09 vs. 13.80 for KDA).
  • This indicates that fine-grained writing rates significantly benefit short-context understanding and associative recall.

• Long-Context Retrieval (RULER Benchmark):

  • FG2-GDN and FG2-GDN+ outperformed GDN and KDA across sequence lengths of 4k, 8k, and 16k.
  • Key Finding: FG2-GDN showed substantial gains in Single-Key Retrieval (MQ, CWE) by tightening key-value bindings.
  • FG2-GDN+ Superiority: The decoupled variant (FG2-GDN+) led in Multi-Key Retrieval (MK2) and variable tracking, proving that independent control of erasure and writing helps manage interference in complex contexts.
  • Extrapolation: The models maintained high accuracy at 16k (trained on 8k), showing better tolerance to distribution shifts in position.

• LongBench (Real-world Tasks):

  • FG2-GDN+ achieved the highest average score (18.3%) and few-shot score (32.1%) across document QA, summarization, and code completion tasks.

• Efficiency:

  • Throughput: FG2-GDN incurs <5% overhead in prefill throughput compared to KDA.
  • Scaling: It preserves the linear time complexity ($O(L)$) and scales similarly to KDA, offering a 1.7–1.9x speedup over MLA (Multi-head Latent Attention) at long sequence lengths (32k).

5. Significance

• Bridging Optimization and Attention: The work successfully translates the concept of per-coordinate adaptive optimizers (like Adam) from the training phase of neural networks to the inference-time recurrence of linear attention models.
• Efficiency vs. Expressiveness: It demonstrates that increasing the expressiveness of linear attention (via vectorized learning rates) does not require sacrificing the computational efficiency that makes linear attention attractive for long contexts.
• State-of-the-Art Performance: By refining both the forgetting ($\alpha$) and writing ($\beta$) mechanisms, FG2-GDN sets a new benchmark for linear recurrent models, offering a practical path toward scalable, high-performance long-context AI agents without the quadratic cost of standard Transformers.