PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.41994v1cs.AIcs.CL

TL;DR

Long-context retrieval remains challenging when linear recurrent models compress history into fixed-size states, and GDN experiments find memory parameters remain correlated with their initial values. DAMI jointly initializes head-wise timescales and layer-wise input sensitivities using actual activations. Across five architectures, it improves average retrieval accuracy over Random while general capabilities remain broadly comparable.

  • Problem

    Linear attention and gated recurrences use fixed-size states, yet long-context retrieval remains challenging despite advances in recurrent architectures.

    Source: §1 · p. 1

  • Method

    DAMI jointly calibrates geometric head-wise timescales and a low–high–low layer-wise input-sensitivity profile from actual activations before training.

    Source: §1 · p. 2

  • Results

    5.29–16.71 percentage points: DAMI improves average RULER-style NIAH accuracy over Random across GDN, KDA, GDN2, Mamba-2, and RWKV-7, with other evaluated capabilities broadly comparable.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings support treating initialization as a coordinated design choice for memory behavior informed by actual input activations.

    Source: §5 · p. 8

    The work focuses on decay-side initialization; jointly studying decay-side and write-side initialization remains future work.

    Source: Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Long-context retrieval remains challenging in fixed-state recurrent models, and controlled GDN experiments motivate structured initialization because trained memory parameters remain correlated with their initial values. DAMI coordinates head-wise timescales and layer-wise input sensitivities, improving retrieval across five architectures while keeping general capabilities broadly comparable.

1 Introduction

  • Controlled GDN runs find that trained A_log and dt_bias remain highly correlated with their initial values, motivating structured initialization.Realized decay also depends on input activations, so the two gate parameters must be coordinated.Source: §1 · pp. 1–2
  • DAMI jointly configures geometric target timescales across heads and a low–high–low input-sensitivity profile across layers.It calibrates both targets from actual activations before training; Mamba-2 and RWKV-7 use architecture-specific extensions.Source: §1 · p. 2
  • 5.29–16.71 percentage points: DAMI improves average RULER-style NIAH accuracy over Random across GDN, KDA, GDN2, Mamba-2, and RWKV-7.Language modeling, commonsense reasoning, and real-world retrieval remain broadly comparable to the baselines.Source: §1 · p. 2
  • Five-architecture experiments report improved retrieval with broadly comparable general capabilities, while GDN ablations establish the benefit of joint configuration.The figure’s placements and outcomes are schematic, not experimental results.Source: Figure 1 · p. 1, §1 · p. 2

2 Related Work

Related work distinguishes efficient recurrent computation from effective recall, and connects multiscale memory to decay spectra and power-law dependencies. DAMI is positioned as an initialization-time approach that calibrates memory parameters using real activations while leaving recurrence and inference unchanged.

Gated Recurrences and Long-Context Retrieval

  • Linear attention compresses history into a fixed-size state, while retrieval evaluations distinguish computational efficiency from effective recall.DAMI operates at initialization and leaves the recurrence and inference procedure unchanged.Source: §2 · p. 2

Decay Spectra and Power-Law Dependencies

  • Approximate power-law dependencies in text motivate exponential mixtures, and PoST applies geometric spectra with position-adaptive scaling to linear recurrences.DAMI builds on this multiscale motivation by jointly realizing target timescales and input responses in GDN gates.Source: §2 · p. 2

Memory Initialization

  • Chrono uses anticipated temporal ranges, LRU combines stable parameterization with initialization and normalization, and LSUV calibrates activation variance using real samples.DAMI uses layer-sequential calibration to jointly determine A_log and dt_bias before training.Source: §2 · p. 2

3 Method

DAMI jointly targets head-wise decay timescales and layer-wise input sensitivities, calibrating both memory parameters before training. The method is motivated by persistent initialization structure and by the dependence of realized decay on parameters and activations.

3 Method

  • DAMI jointly calibrates head-wise timescales and layer-wise Softplus sensitivities by determining A_log and dt_bias before training.It uses a layer-sequential, data-driven procedure.Source: §3 · p. 2, §3.5 · p. 5

3.2 Persistence of Initialization Structure

  • GDN’s A_log and dt_bias remain highly correlated with their initial values: correlations are 0.926/0.939 for RANDOM, 0.989/0.949 for POST, and 0.985/0.871 for DAMI.These correlations indicate that trained memory parameters retain statistical structure associated with initialization.Source: §3.2 · p. 3, Figure 3 · p. 4

3.3 Head-wise Timescale Allocation

  • DAMI assigns geometric target timescales across heads to provide multiscale memory coverage.A head’s e-folding time is the token interval for its retained contribution to reach 1/e; larger values mean slower decay.Source: §3.3 · p. 4

3.4 Layer-wise Input Sensitivity

  • DAMI uses a low–high–low sensitivity profile across layers, with low sensitivity early and late and high sensitivity in the middle.The profile is motivated by balancing token-level stability, content-dependent adaptation, and long-term retention.Source: §3.4 · p. 5

3.5 Joint Data-dependent Initialization

  • For each layer, DAMI collects projected activations without gradients, solves dt_bias by bisection, then matches the target timescale through the mean Softplus factor.The calibrated output feeds the next layer, while the output projection retains its native initialization.Source: §3.5 · p. 5
  • KDA/GDN2, Mamba-2, and RWKV-7 receive architecture-specific adaptations that preserve their state-update operators and add no trainable parameters.The adaptations use channel-wise bias matching, depth-scaled Softplus biases, or scheduled sigmoid response, respectively.Source: §3.5 · p. 5

4 Experiments

The experiments compare Random, PoST-style spectral initialization, and DAMI across five recurrent architectures, evaluating retrieval, general capabilities, ablations, and retrieval by distance. DAMI achieves the best average retrieval accuracy in all five architectures, with gains over Random ranging from 5.29 to 16.71 percentage points.

4.1 Experimental Setup

  • The setup compares Random, PoST-style spectral initialization, and DAMI on five architectures trained for 5B tokens, using fixed RULER-style NIAH tasks and six 2K real-world retrieval tasks.Models are approximately 100M parameters; the primary retrieval evaluation averages answer-containment accuracy over 11 settings.Source: §4.1 · p. 6

4.2 Main Results

  • DAMI improves average retrieval accuracy over Random by 13.95, 16.71, 13.44, 14.47, and 5.29 percentage points in GDN, KDA, GDN2, Mamba-2, and RWKV-7, respectively.It has the best average in all five architectures and the best result in 42 of 55 task–length comparisons, including ties; at 8K, it improves NIAH-1 for GDN, KDA, GDN2, and Mamba-2.Source: §4.2 · pp. 6–7
  • Language-modeling and commonsense performance remain broadly comparable, as do the six 2K real-world retrieval tasks; for Mamba-2, DAMI raises retrieval average from 3.89 to 4.28.DAMI leads 17 of 30 general-capability comparisons, including ties.Source: §4.2 · p. 7

4.3 Ablation Studies

  • Full DAMI reaches 44.36 Overall accuracy, versus 30.42 for Random and 22.80 or 26.33 when A_log or dt_bias alone is randomized.Randomizing A_log yields NIAH-2 accuracy of 2.40, consistent with the reported coupling between the parameters.Source: §4.3 · p. 7, Table 5 · p. 7
  • Restricting timescale coverage to 1–2048 or 1–1024 reduces Overall accuracy to 29.07 or 23.13, while linear spacing gives NIAH-2 accuracy of 5.65.With linear spacing, NIAH-1 accuracy is 62.40 versus 64.35 for DAMI.Source: §4.3 · p. 7, Table 6 · p. 7
  • For GDN retrieval evaluated across four layer-wise sensitivity profiles, DAMI leads all three NIAH families, reaching 44.36 Overall versus 26.85–32.85 for the other tested profiles.Shifting sensitivity by ±0.2 or moving its peak to layer 4 or 13 yields 24.31–33.31 Overall; DAMI leads Overall, while Peak L4 leads MK-NIAH.Source: §4.3 · p. 8, Table 7 · p. 8, Table 8 · p. 8

4.4 Case Study

  • At 1K, 2K, and 4K context lengths, DAMI maintains higher mean accuracy in distant retrieval-distance bins.Figure 6 averages task accuracy in 16 shared bins, with each target–query gap normalized by the background-to-query token span.Source: §4.4 · p. 8, Figure 6 · p. 8

5 Conclusion

DAMI coordinates head-wise timescales and layer-wise input responses through joint calibration across five recurrent architectures. Experiments show improved long-context retrieval with broadly comparable general capabilities, while GDN ablations support jointly configuring the memory parameters.

5 Conclusion

  • DAMI coordinates head-wise timescales and layer-wise input responses through joint calibration for GDN, KDA, and GDN2, with architecture-specific extensions to Mamba-2 and RWKV-7.Across these architectures, experiments show improved long-context retrieval with broadly comparable language modeling, commonsense reasoning, and real-world retrieval performance.Source: §5 · p. 8

Limitations

The paper focuses on decay-side initialization within existing recurrent architectures, leaving coordination with write-side initialization for future study.

Limitations

  • The study addresses decay-side initialization, not joint initialization with the write side that governs how incoming content enters the recurrent state.The authors identify coordinating write-side and decay-side initialization as future work.Source: Limitations · p. 9

A Approximation Basis for Geometric Spectra

The approximation discussion frames geometric spectra as a way to approximate a power law over a finite interval with exponential sums. It connects geometric spacing to asymptotically uniform nodes in the log domain.

A Approximation Basis for Geometric Spectra

B Joint Calibration Derivation

Joint calibration accounts for the dependence between dt_bias and the mean Softplus factor, so A_log must adjust with it. Monotonicity in the bias permits bisection for the calibration.

C Architecture-specific Adaptation Details

Architecture-specific adaptations retain each model's native state-update operator while transferring DAMI's timescale and depth-profile targets. Mamba-2 uses static Softplus biases, whereas RWKV-7 targets its sigmoid response.

C Architecture-specific Adaptation Details

D Detailed Experimental Configurations

The experiments compare five models trained on 5B tokens with a 4,096-token maximum context, using model-specific initialization settings. Distance-based retrieval bins average task accuracies with equal weights.

D Detailed Experimental Configurations

  • All models use hidden size 512 and head dimension 64, train on 5B tokens, and have a maximum context of 4,096 tokens.GDN, KDA, and GDN2 have 18 layers and six heads; Mamba-2 has 42 layers and 16 heads; RWKV-7 has 20 layers and eight heads.Source: D Detailed Experimental Configurations · p. 11
  • GDN-family DAMI uses timescale targets from 4,096 to 1 tokens and an 18-layer sensitivity profile with endpoints 0.25 and central peak 0.748.The settings also specify KDA/GDN2 centered perturbations, Mamba-2 layer-multiplier resampling, and RWKV-7 response and timescale targets.Source: D Detailed Experimental Configurations · p. 11
  • Figure 6 averages task-specific accuracy equally across 16 shared distance bins, using all 500 samples per task–length setting.The analysis uses NIAH-1, NIAH-2, and MK-NIAH at 1K–4K, and NIAH-1 and NIAH-2 at 8K.Source: D Detailed Experimental Configurations · p. 11
Loading 2610.41994v1…