DAMI: Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention

AI summary

Overview

Research area: Efficient long-context language modeling, specifically linear attention and gated linear recurrences (GDN, KDA, GDN2, Mamba-2, RWKV-7), with a focus on how the memory/decay parameters are initialized before training.

Technical level: Advanced. The paper builds on the recurrence equations of gated linear attention, Softplus-parameterized decay gates, and approximation theory for power-law mixtures.

Scope: The paper proposes DAMI (Dual-Axis Memory Initialization), a pre-training initialization scheme that jointly sets head-wise decay timescales and layer-wise input sensitivity for recurrent memory models, and evaluates it on five ~100M-parameter architectures trained on 5B tokens.

What This Paper Is About

Linear attention keeps a fixed-size memory instead of a growing cache, so information from far back in a sequence can fade away. How fast it fades is controlled by two learnable parameters per head, A_log and dt_bias, and the paper observes that after training these parameters stay highly correlated with whatever values they started from. The goal is therefore to design a good starting point: DAMI assigns each head a different memory timescale and each layer a different sensitivity to incoming inputs, so that the model can retrieve information at many distances before training even begins.

Key Contributions

1. Evidence that initialization persists. In controlled Gated DeltaNet (GDN) runs, the trained memory parameters remain strongly correlated with their initial values, with initial–trained Pearson correlations of 0.926/0.939 for Random initialization, 0.989/0.949 for PoST-style initialization, and 0.985/0.871 for DAMI (for A_log/dt_bias respectively). This motivates treating initialization as a design choice rather than a random draw.

2. The DAMI method itself. DAMI configures both decay parameters before training along two axes: a geometric spread of per-head timescales (to cover many retrieval distances) and a low–high–low profile of per-layer input sensitivity (early layers low, middle layers high, late layers low). Calibration uses a small sample of real activations, adds no trainable parameters, and leaves the state-update operators unchanged.

3. Cross-architecture extensions. The core scheme is defined for the GDN family (GDN, KDA, GDN2); the paper adds architecture-specific adaptations for Mamba-2 (which uses head-wise Softplus decay) and RWKV-7 (which uses a sigmoid decay gate), all using the same 160 frozen head–layer scales.

4. A five-architecture evaluation. DAMI is compared against Random initialization and a PoST-style spectral baseline across GDN, KDA, GDN2, Mamba-2, and RWKV-7 at approximately 100M parameters and 5B training tokens, on RULER-style retrieval, language modeling, commonsense reasoning, and six real-world retrieval tasks at 2K tokens.

Main Findings

• Retrieval accuracy improves substantially over Random. Average RULER-style NIAH accuracy rises by 13.95, 16.71, 13.44, 14.47, and 5.29 percentage points for GDN, KDA, GDN2, Mamba-2, and RWKV-7 respectively, i.e. a 5.29–16.71 point range overall.

• DAMI also beats the PoST-style spectral baseline. Gains over that control are 10.87, 9.42, 10.24, 10.69, and 8.47 percentage points for the same five architectures. DAMI achieves the best average within all five architectures and the best individual result in 42 of 55 task–length comparisons, including ties.

• Gains extend past the training context. At 8K tokens, twice the 4,096-token training context, DAMI improves NIAH-1 for GDN, KDA, GDN2, and Mamba-2. RWKV-7's gains concentrate on NIAH-1 and MK-NIAH at 1K/2K.

• General capabilities stay broadly comparable. DAMI leads 17 of 30 language-modeling and commonsense comparisons including ties. For Mamba-2 and RWKV-7, WinoGrande improves by 1.82 and 0.79 points over Random, and Social IQa scores are the best within each of those architectures.

• Real-world 2K retrieval is roughly unchanged. Mamba-2's six-task average rises from 3.89 to 4.28 with DAMI (leading SQuAD, FDA, and DROP for that architecture), while RWKV-7 averages span 5.67–6.42 across the three schemes.

• Both axes must be set jointly. In GDN ablations, full DAMI reaches an Overall average of 44.36 versus 30.42 for Random. Restoring only A_log to Random while keeping DAMI's dt_bias drops Overall to 22.80 and NIAH-2 to 2.40; restoring only dt_bias gives 26.33.

• The timescale span matters. Restricting the target timescales to 1–2048 or 1–1024 reduces Overall to 29.07 and 23.13. Linear instead of geometric spacing keeps NIAH-1 close (62.40 versus 64.35) but collapses NIAH-2 to 5.65.

• The low–high–low sensitivity shape matters. Alternative profiles score 32.29 (increasing, L-H), 26.85 (decreasing, H-L), and 32.85 (high–low–high), all below DAMI's 44.36. Shifting the profile by ±0.2 or moving its peak to layer 4 or layer 13 yields 24.31–33.31 Overall; Peak L4 does lead MK-NIAH at 24.13.

• Retrieval gains concentrate at long distances. In a distance-binned GDN case study using 16 shared bins, DAMI keeps higher mean accuracy in distant bins at 1K, 2K, and 4K context lengths.

Methodology in Plain English

A gated linear attention layer compresses everything it has read into a fixed-size state. At each token, part of the old state is multiplied by a decay factor between 0 and 1; the rest is updated with new key–value content. That decay factor is not a constant. It is computed as exp(-exp(A_log) · softplus(z_t + dt_bias)), where z_t is a projection of the current input. So A_log sets the overall decay rate, dt_bias sets how responsive the gate is to the input, and the input itself supplies the variation. The baseline initializes both randomly and independently of the data, using the same rule in every layer.

DAMI instead prescribes two targets and then solves for the parameters that hit them:

1. A timescale target per head. The paper treats each head as one exponential decay component and notes that natural language statistics are roughly power-law, which is well approximated by a mixture of exponentials with geometrically spaced rates. With six heads in the GDN-family models, DAMI sets targets of {4096, 776.05, 147.03, 27.86, 5.28, 1} tokens, so some heads remember for a very long time and others turn over quickly.

2. A sensitivity target per layer. The derivative of Softplus is the sigmoid, so dt_bias slides the operating point between a nearly flat region (sensitivity ≈ 0.018 at input −4) and a nearly linear one (≈ 0.982 at input +4). The paper assigns a low–high–low profile across depth: 0.25 at the ends and a central peak of 0.748 in the 18-layer GDN profile.

3. Joint, data-driven calibration. Working layer by layer from shallow to deep, in the spirit of LSUV initialization, DAMI runs a small sample of activations through the model without gradients. For each head it bisects on dt_bias to match the layer's sensitivity target, then solves A_log in closed form from the mean Softplus factor so that the head's mean decay matches its timescale target. Because changing dt_bias shifts the mean Softplus value, A_log must be solved afterwards — this is why the two axes cannot be set independently.

For KDA and GDN2, which have per-channel rather than per-head dt_bias, the within-head means are aligned using centered perturbations with standard deviation 0.03. Mamba-2 keeps the geometric head-wise spectrum and scales its static Softplus biases along depth (resampled to 42 layers, normalized to unit geometric mean). RWKV-7 schedules the mean sigmoid response of its static biases into the same 0.25–0.75–0.25 shape and solves an eight-head timescale ladder from 4,096 to 1, extending its native constant log scale of −0.5. None of these adaptations adds trainable parameters.

Comparisons fix model size, data, optimizer, and training budget; DAMI copies the corresponding Random model's non-memory parameters. The PoST condition applies only spectral initialization with Random's dt_bias and a 0–8.318 spectrum, excluding position-adaptive scaling.

Why This Matters

Impact on research. The paper reframes initialization as a coordinated design choice rather than a nuisance hyperparameter. Its central empirical claim — that memory parameters barely move from where they start — implies that a large part of a recurrent model's long-context ability is decided before the first gradient step. That is a cheap lever compared with changing the recurrence, adding state size, or post-hoc filtering.

Real-world applications:

• Long-document question answering over contracts, filings, or manuals where the answer sits thousands of tokens from the question.
• Retrieval-augmented generation, where many retrieved passages must be held in memory while distractors compete for the same fixed-size state.
• Agent and tool-use transcripts, where earlier instructions must survive long stretches of intervening interaction.
• Codebase-level code modeling, where a relevant definition may be far from the point of use but within one context window.

Industry relevance. DAMI requires no architectural change, no extra parameters, and no change to inference, and it is calibrated from a small activation sample before training. That makes it a low-risk drop-in for teams already training or fine-tuning gated linear attention and hybrid recurrent models, and it targets exactly the retrieval weakness that keeps such models from displacing softmax attention in long-context products.

Future Directions

• Joint decay-side and write-side initialization. The authors explicitly flag this as the main open question: the paper only coordinates decay, holding each architecture's native write parameterization and state-update rule fixed. How writes and decay should be initialized together, to trade off incorporating new information against retaining old information, is left open.

• Scaling behavior is not reported. All experiments use approximately 100M-parameter models trained on 5B tokens with a 4,096-token context, and evaluation goes to 8K. Whether the observed gains persist at larger parameter counts, longer training, or much longer contexts is not established here.

• Why low–high–low works. The profile is motivated by representation stages rather than derived. The ablations show it beats the alternatives tested, but the paper does not provide a mechanistic account of why the middle layers should be most input-sensitive.

• More architectures and hybrid stacks. The extensions to Mamba-2 and RWKV-7 use architecture-specific workarounds and 160 frozen scales. A general recipe that maps the dual-axis targets onto arbitrary decay parameterizations — including hybrid models that mix recurrent and attention layers — remains to be worked out.

Target Audience

Researchers and engineers working on efficient long-context architectures — linear attention, gated DeltaNet variants, state-space models, and RWKV-style recurrences — who want to improve retrieval without changing the model. It is most useful to those who train such models from scratch or adapt them at scale, and to readers interested in initialization and memory dynamics in sequence models. Familiarity with recurrence equations and gate parameterizations is assumed; readers without that background will find the method sections demanding but the empirical results and ablations accessible.

Authors’ abstract

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.