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

DAMI. Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention. By Zhangchi Feng and four others, from SpeedAI and Beihang University. Published on the thirtieth of October.

Linear attention has an appealing property: its cost grows only linearly with the length of the sequence. The price is paid in memory. Its state is finite, and a finite state can only hold so much, which limits how well a model retrieves something from far back in a long context. In gated linear recurrences, how long a memory survives depends on state decay, and that decay is shaped jointly by two gate parameters and by the input activations flowing through the layer.

The starting observation comes from controlled experiments on Gated DeltaNet, or GDN. After training, both of those gate parameters remain highly correlated with the values they were initialized to. Training does not carry them far from where they began. If the start is largely the finish, then initialization matters a great deal, and that is the lever this work pulls. Dual-Axis Memory Initialization, or DAMI, jointly configures these parameters before training begins.

The two axes are heads and layers. Across heads, DAMI assigns geometric timescales, so that different heads serve different retrieval distances. Across layers, it sets a low, high, low profile of input sensitivity. To land the GDN-family implementation on both of these behavioral targets, a small sample of activations is used for calibration. Nothing trainable is added. The change lives entirely in where the existing parameters start.

The method is tested on five architectures: GDN, KDA, GDN2, Mamba-2, and RWKV-7. On each, DAMI improves average RULER-style retrieval accuracy over the matching random-initialization baseline, by between five point two nine and sixteen point seven one percentage points. Elsewhere, performance holds roughly steady. Language modeling, commonsense reasoning, and the two-K real-world retrieval tests all stay broadly comparable to those baselines.

Very little changes here. There are no new parameters and no new architecture, only a deliberate choice of starting point, motivated by the finding that these gates barely drift from it. For models whose weakness is memory over distance, the retrieval gains come from deciding, before training, how far each head should look back.

DAMI. Dual-Axis Memory Initialization for Enhanced Long-Context Retrieval in Gated Linear Attention. By Zhangchi Feng and four others, from SpeedAI and Beihang University. Published on the thirtieth of October. This is the long version: the whole paper rather than the abstract.

Linear attention trades the growing cache of softmax attention for a fixed-size state. That makes it cheap at long lengths, but a finite memory has to forget something, and long-context retrieval is the task that punishes forgetting the wrong thing. Ask for the number attached to one key, buried thousands of tokens back, and the answer depends on whether that memory survived the decay applied at every step since. The angle taken here is narrow and practical. Leave the recurrence alone, and change only how the decay gate starts life.

In Gated DeltaNet, every head in every layer carries two scalars that shape forgetting. One, called A log, sets a decay multiplier. The other, dt bias, shifts where the incoming activation lands on a Softplus curve. The decay at each step is the multiplier times the Softplus of activation plus bias, and the state is scaled down accordingly. Native initialization samples both at random, with the same rule in every layer, regardless of what the activations look like. The motivating observation is that after training, these parameters stay highly correlated with their starting values. Pearson correlations between initial and trained values run from zero point eight seven one to zero point nine eight nine across the three schemes tested. If training barely moves them, initialization effectively decides them.

DAMI, short for Dual-Axis Memory Initialization, sets two targets on two different axes. The first runs across heads. Each head gets a target timescale: the number of tokens after which a stored contribution has decayed to about one over e. Those timescales are spaced geometrically, from one token up to four thousand ninety-six, the training context. The justification is that dependencies in natural text decay roughly as a power law, and a power law can be approximated by a weighted sum of exponentials. The output projection, which adds up the heads' contributions, plays the part of that weighted sum, and a result from PoST supports geometric spacing as giving an approximation error that shrinks exponentially.

The second axis runs across layers, and concerns sensitivity: the average slope of the Softplus at its operating point, meaning how strongly the gate responds to the content of its input. The profile is low, high, low. About a quarter in early layers, about three quarters in the middle, back to a quarter at the end. Calibration borrows the layer-sequential approach of LSUV. A small calibration sample is pushed through the model without gradients, shallow layers first. For each head, the method collects the activations entering the gate and solves for the bias by bisection, so that the average Softplus slope hits the layer's sensitivity target. Only then does it set A log, choosing the multiplier so that multiplier times mean Softplus value gives one over the head's target timescale. The calibrated layer feeds the next. No parameters are added and the state update is untouched. KDA and GDN2 keep per-channel biases, matched in mean, while Mamba-2 and RWKV-7 get looser, architecture-specific adaptations.

Five architectures, each around a hundred million parameters and trained on five billion tokens with a four-thousand-token context, are compared under three schemes: random initialization, DAMI, and a PoST-style control that sets geometric decay rates but keeps the random bias. Across eleven needle-in-a-haystack settings from RULER, DAMI lifts average accuracy over random by between five point two nine and sixteen point seven one points, and beats the PoST control in every architecture. For GDN, the average goes from thirty point four two to forty-four point three six.

The comparison that tests the thesis is the ablation. Keep DAMI's value for one parameter and restore the other to random. If the axes were independent, each half should keep part of the gain. Instead both halves do worse than fully random: twenty-two point eight with random A log, and twenty-six point three three with a random bias. Timescale ablations agree: linear spacing nearly matches DAMI on the first needle task but collapses on the second, and of four layer profiles, low, high, low wins clearly.

There is a limitations section, and it is brief. The work covers only the decay side, under each architecture's native write rules, and how new content gets written into the state is left for future work. Beyond that, every model is around a hundred million parameters, and each condition is one trained checkpoint with no error bars. For RWKV-7, the DAMI model used a different random seed from its two baselines, and its gains sit at short lengths: on the first needle task at four and eight thousand tokens it falls below random. The multi-key task barely moves for GDN. On six real-world retrieval tasks, absolute scores are tiny, and DAMI's average is slightly below random for GDN, KDA and RWKV-7. GDN perplexity on WikiText also edges up slightly. Broadly comparable is the paper's phrase, and it is a fair one.

The idea to carry away is that in a gated recurrence, forgetting is set by two numbers that only mean something together, and that training largely keeps whatever it is handed. So hand it something deliberate: a ladder of memory lengths across heads, a low, high, low pattern of responsiveness across depth, both calibrated against the activations the model actually produces. No new parameters, no new operator, and at this scale a double-digit retrieval gain for four of the five architectures.