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

Zhangchi Feng; Harry Jiang; Yangyifei Luo; Ray Zhang; Richong Zhang

cs.AIOct 30, 2026

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

Authors: Zhangchi Feng, Harry Jiang, Yangyifei Luo, Ray Zhang, Richong Zhang

Organizations: SpeedAI · CCSE, School of Computer Science and Engineering, Beihang University

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.

Figures & tables

Figure 1: Dual-axis memory initialization. (a) Diverse head-wise timescales broaden memory coverage across evidence–query distances; curves illustrate normalized exponential decay. (b) Schematic uniform-high and uniform-low Random regimes illustrate recent-pair bias and mixed memories, contrasted with DAMI’s low–high–low sensitivity profile. Text excerpts are from RULER; placements and outcomes are schematic.

Figure 2: Overview of DAMI initialization. (a) Head-wise timescales and layer-wise sensitivities define dual-axis targets. (b) Calibration proceeds layer by layer using projected activations without gradients. (c) Each head’s dt_bias is solved before A_log using calibration-set statistics. (d) Schematic parameter layouts illustrate the head and layer axes.

Model	Gate $g_{t}$	A_log	dt_bias	Decayed state	State update
GDN	$-e^{A_{\log}}\phi(z_{t}+\mathrm{dt\_bias})$	Head-level	Head-level	$\widetilde{S}_{t-1}=e^{g_{t}}S_{t-1}$	$S_{t}=\widetilde{S}_{t-1}+\beta_{t}k_{t}(v_{t}^{\top}-k_{t}^{\top}\widetilde{S}_{t-1})$
KDA	$-e^{A_{\log}}\phi(z_{t,k}+\mathrm{dt\_bias}_{k})$	Head-level	Key-channel	$\widetilde{S}_{t-1}=\operatorname{Diag}(e^{g_{t}})S_{t-1}$	$S_{t}=\widetilde{S}_{t-1}+\beta_{t}k_{t}(v_{t}^{\top}-k_{t}^{\top}\widetilde{S}_{t-1})$
GDN2	$-e^{A_{\log}}\phi(z_{t,k}+\mathrm{dt\_bias}_{k})$	Head-level	Key-channel	$\widetilde{S}_{t-1}=\operatorname{Diag}(e^{g_{t}})S_{t-1}$	$S_{t}=\widetilde{S}_{t-1}+k_{t}[(w_{t}\odot v_{t})^{\top}-(b_{t}\odot k_{t})^{\top}\widetilde{S}_{t-1}]$
Mamba-2	$-e^{A_{\log}}\phi(z_{t}+\mathrm{dt\_bias})$	Head-level	Head-level	$\widetilde{S}_{t-1}=e^{g_{t}}S_{t-1}$	$S_{t}=\widetilde{S}_{t-1}+\Delta_{t}B_{t}x_{t}^{\top}$
RWKV-7	$-e^{A_{\log}}\sigma(z_{t,k}+\mathrm{dt\_bias}_{k})$	Head-level (frozen)	Key-channel	$\widetilde{S}_{t-1}=\operatorname{Diag}(e^{g_{t}})S_{t-1}$	$S_{t}=\widetilde{S}_{t-1}+\tilde{k}_{t}v_{t}^{\top}-p_{t}(\hat{\kappa}_{t}^{\top}S_{t-1})$

Table 1: Gate parameterizations and state updates.

\phi=\operatorname{softplus}

\sigma

is sigmoid, and

\odot

denotes element-wise multiplication.

\beta_{t}

is scalar;

b_{t},w_{t}

are independent channel-wise coefficients. For Mamba-2,

\Delta_{t}=\phi(z_{t}+\mathrm{dt\_bias})

, with input vectors

B_{t},x_{t}

. For RWKV-7,

p_{t}=a_{t}\odot\hat{\kappa}_{t}

, with normalized key

\hat{\kappa}_{t}

and channel-wise coefficient

a_{t}

Panel of Figure 3 — Figure 3: Initial–trained Pearson correlations of GDN memory parameters.

Figure 4: Effect of dt_bias on local Softplus sensitivity. Dashed and dotted curves show the asymptotic approximations;

\phi^{\prime}(-4)\approx 0.018

and

\phi^{\prime}(4)\approx 0.982

Model	1K	2K	4K	8K	1K	2K	4K	8K	1K	2K	4K	Avg.
	Single-Key NIAH-1				Single-Key NIAH-2				Multi-Key MK-NIAH
GDN	92.6	62.0	32.4	13.2	19.8	41.6	13.8	1.6	24.4	20.8	12.4	30.42
+PoST	93.2	64.8	35.6	14.6	47.6	47.6	2.4	9.4	23.6	21.4	8.2	33.49
+DAMI	95.4	84.0	53.8	24.2	89.6	55.6	13.0	12.8	25.0	22.0	12.6	44.36
KDA	63.8	29.4	16.8	5.6	88.8	25.6	1.2	3.8	23.0	18.2	14.0	26.38
+PoST	93.6	72.0	34.2	9.6	79.8	16.8	9.6	5.4	18.2	17.2	14.0	33.67

Table 2: Answer-containment accuracy (%) on the fixed v1 RULER-style NIAH benchmark. Each task–length setting contains the same 500 pregenerated samples for all models. Bold denotes the best result within each architecture; Avg. is the unweighted mean over the 11 settings.

Model	WikiText	WinoG.	ARC-E	ARC-C	OBQA	SIQA
	ppl $\downarrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$	acc $\uparrow$
GDN	44.44	48.93	43.81	25.00	31.60	36.64
+PoST	44.83	48.78	43.90	24.74	30.20	37.15
+DAMI	45.27	52.09	44.07	23.72	31.80	37.00
KDA	42.04	51.30	44.11	24.32	30.60	37.56
+PoST	42.42	51.22	43.52	24.06	31.40	37.10

Table 3: Language modeling and commonsense results. WikiText reports perplexity (lower is better); other tasks report accuracy (%, higher is better). Bold denotes the best result within each architecture. WinoG., ARC-E, ARC-C, OBQA, and SIQA denote WinoGrande, ARC-Easy, ARC-Challenge, OpenBookQA, and Social IQa, respectively.

Model	SWDE	SQuAD	FDA	TQA	NQ	DROP	Avg.
GDN	9.18	18.53	2.00	0.17	0.14	2.16	5.36
+PoST	6.48	18.10	1.54	0.23	0.11	2.49	4.83
+DAMI	7.47	19.81	1.09	0.24	0.14	2.25	5.16
KDA	9.81	21.31	1.18	0.31	0.11	2.33	5.84
+PoST	8.91	19.44	2.18	0.26	0.08	2.30	5.53
+DAMI	8.46	20.27	1.27	0.07	0.03	2.35	5.41

Table 4: Real-world retrieval at 2K tokens (%). TQA denotes TriviaQA; Avg. is the unweighted mean over six tasks. Bold denotes the best result within each architecture.

Variant	NIAH-1	NIAH-2	MK-NIAH	Overall
	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$
Random	50.05	19.20	19.20	30.42
Random A_log	49.15	2.40	14.87	22.80
Random dt_bias	42.55	16.90	17.27	26.33
DAMI	64.35	42.75	19.87	44.36

Table 5: GDN single-parameter replacement ablations (%). The named parameter is restored to Random; the other retains its DAMI value.

Variant	NIAH-1	NIAH-2	MK-NIAH	Overall
	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$
$\tau$ : 1–2048	47.95	17.30	19.60	29.07
$\tau$ : 1–1024	22.05	25.50	21.40	23.13
Linear spacing	62.40	5.65	15.47	28.96
DAMI	64.35	42.75	19.87	44.36

Table 6: Effects of timescale coverage and spacing on GDN retrieval (%).

Figure 5: Normalized target sensitivity profiles across layers. The horizontal axis is normalized depth

r_{\ell}=(\ell-1)/(L-1)

; the vertical axis is

(\mathcal{S}^{\star}_{\ell}-\mathcal{S}_{\min})/(\mathcal{S}_{\max}-\mathcal{S}_{\min})

. Curves indicate initialization targets.

Shape	NIAH-1	NIAH-2	MK-NIAH	Overall
	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$
L-H	54.70	19.55	19.40	32.29
H-L	50.55	9.15	18.87	26.85
L-H-L (DAMI)	64.35	42.75	19.87	44.36
H-L-H	60.05	16.60	18.27	32.85

Table 7: GDN retrieval with four layer-wise sensitivity profiles (%).

Variant	NIAH-1	NIAH-2	MK-NIAH	Overall
	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$	Avg. $\uparrow$
Shift $+0.2$	45.25	6.10	20.67	24.31
Shift $-0.2$	57.40	14.75	23.13	32.55
Peak L4	51.25	22.25	24.13	33.31
Peak L13	52.60	20.80	18.60	31.76
DAMI	64.35	42.75	19.87	44.36

Table 8: GDN sensitivity-level and peak-location perturbations (%). Shift denotes profile translation; Peak denotes the peak layer.

Figure 6: GDN retrieval accuracy by relative distance at four context lengths. Points show task-averaged accuracy in 16 shared bins.

Long-Context Retrieval

Long-context retrieval finds and uses relevant evidence across lengthy documents, conversations, and model contexts.

See all 49 papers

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention
May 21, 2026Ali Hatamizadeh, Yejin Choi, Jan KautzLong-Context Language Modeling Recurrent Transformers
NVIDIA
Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity
Jul 8, 2026Loïc Cabannes, Pierre-Emmanuel Mazaré, Gergely Szilvasy +6Long-Context Retrieval Memory-Augmented Language Models
Meta FAIR · Inria Paris & ENS-PSL University · University of Tübingen
HLA: Expressive Hybrid Linear Attention via Chunk-Wise Dynamic Mixing
Oct 5, 2026Zhuokun Chen, Xi Lin, Xiyu Wu +3Long-Context Language Modeling Linear Attention
Monash University · Zhejiang University

May 21, 2026J/K move · Enter open · S save

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

Ali Hatamizadeh, Yejin Choi, Jan Kautz

NVIDIA

Long-Context Language Modeling Recurrent Transformers Linear Attention Delta-Rule Linear Attention Gated Delta Networks

Open paper AskPDF arXiv

Linear attention replaces the unbounded cache of softmax attention with a fixed-size recurrent state, reducing sequence mixing to linear time and decoding to constant memory. The hard part is not just what to forget, but how to edit this compressed memory without scrambling existing associations. Delta-rule models subtract the current read before writing a new value, and Kimi Delta Attention (KDA) sharpens forgetting with channel-wise decay. But the active edit still uses a single scalar gate to control two different things: how much old content to erase on the key side and how much new content to commit on the value side. We introduce Gated DeltaNet-2, which generalizes both Gated DeltaNet and KDA by inheriting adaptive forgetting and channel-wise decay while addressing their shared limitation, the scalar tie between erasing and writing. Gated Delta Rule-2 separates these roles with a channel-wise erase gate b_t and a channel-wise write gate w_t, reducing to KDA when both gates collapse to the same scalar and to Gated DeltaNet when the decay also collapses. We derive a fast-weight update view, a chunkwise WY algorithm with channel-wise decay absorbed into asymmetric erase factors, and a gate-aware backward pass that preserves efficient parallel training. At 1.3B parameters trained on 100B FineWeb-Edu tokens, Gated DeltaNet-2 achieves the strongest overall results among Mamba-2, Gated DeltaNet, KDA, and Mamba-3 variants across language modeling, commonsense reasoning, and retrieval. Its advantage is most pronounced on long-context RULER needle-in-a-haystack benchmarks, where it improves the evaluated multi-key retrieval setting and remains strong in both recurrent and hybrid settings. Code is available at https://github.com/NVlabs/GatedDeltaNet-2.

Abstract

Long-Context Retrieval

Explore similar work

Gated DeltaNet-2: Decoupling Erase and Write in Linear Attention

Sparse Delta Memory: Scaling the State of Linear RNNs through Sparsity

HLA: Expressive Hybrid Linear Attention via Chunk-Wise Dynamic Mixing