DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention
Abstract
Dual-Axis Memory Initialization improves long-context retrieval in gated linear recurrent models by jointly configuring decay and sensitivity parameters before training.
Linear attention scales linearly with sequence length, but its finite-state memory can limit long-context retrieval. In gated linear recurrences, memory retention depends on state decay, which is shaped jointly by gate parameters A_log and dt_bias and by input activations. In controlled Gated DeltaNet (GDN) experiments, both parameters remain highly correlated with their initial values after training, highlighting the importance of initialization. Therefore, we introduce Dual-Axis Memory Initialization (DAMI), which jointly configures these parameters before training. DAMI assigns geometric timescales across heads for varied retrieval distances and a low-high-low input-sensitivity profile across layers. A small activation sample calibrates the GDN-family implementation to match both behavioral targets without adding trainable parameters. Across GDN, KDA, GDN2, Mamba-2, and RWKV-7, DAMI improves average RULER-style retrieval accuracy by 5.29 to 16.71 percentage points over the corresponding random-initialization baselines. Performance on language modeling, commonsense reasoning, and 2K real-world retrieval remains broadly comparable to those baselines. Our code is available at https://speedai.com/research/dami.
Community
DAMI 👇 — a fairly systematic study of the initialization of decay-coefficient-related parameters in mainstream linear attention architectures (GDN, GDN2, KDA, Mamba-2, RWKV-7). Comments and discussion are very welcome
Motivation: Linear attention replaces the KV cache with a fixed-size state. Inference gets cheaper, but the old problem of "failing to remember far away" persists, and one of the key control components is the decay coefficient. Through 100M-parameter, 5B-token pretraining experiments, we found that after training, the correlation between the A_log and dt_bias that determine the decay coefficients and their initial values is above 0.9 — the decay-memory structure is essentially fixed at initialization. So how do we set these two parameters correctly before training? We carried out a series of theoretical derivations and controlled experiments, and arrived at DAMI.
DAMI: We spread A_log, which governs the timescale, across heads as a geometric progression (4096 → 1 token), and make dt_bias, which governs input sensitivity, vary along model depth in a "low–high–low" profile (0.25 → 0.75 → 0.25); the two are solved jointly to obtain the initialization values. In practice it requires only a small batch of activations, layer-wise gradient-free forward passes, and a bisection solve — no extra parameters, no operator changes, no impact on inference. The result: across five architectures (GDN / KDA / GDN2 / Mamba-2 / RWKV-7), average RULER retrieval accuracy is 5.29–16.71 higher than random initialization (GDN 30.42 → 44.36, KDA 26.43 → 43.14, GDN2 30.33 → 43.77, Mamba-2 13.7 → 28.2, RWKV-7 23.9 → 29.2), while general capabilities remain essentially unchanged.
Limitations: We only studied the decay coefficients, not the write/update coefficients; larger data scales and model sizes still need validation.
📑 Paper: https://huggingface.co/papers/2610.41994
📎 Blog: https://speedai.com/research/dami
💻 Code: https://github.com/SpeedAI-Lab/DAMI
⚖️ Weights: https://huggingface.co/SpeedAI-Lab/DAMI
Get this paper in your agent:
hf papers read 2610.41994 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 1
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper