Stable Token-by-Token Memory Storage and Retrieval

Why People Wanted It

A long-running assistant needs a way to carry useful information between interactions. Ordinary attention uses the context it receives; an explicit memory state provides another route for retaining and retrieving information across turns.

The problem is not just where to put that information. Updating a numerical memory can destabilize it, and reading memory during generation is a different operation from storing the current interaction. So the real question is: how can the model consult previously accumulated memory while answering, then incorporate the completed interaction without an unstable erase operation?

The timing matters throughout: storage is triggered at the turn level and scans tokens internally; utilization reads memory at the token level during generation. These are separate procedures, not alternating steps in one decoding loop.

The Core Move

MemT-1.0 adds memory blocks to selected Transformer layers: generation reads the turn's accumulated state alongside ordinary attention, while a turn-level storage procedure scans the completed input and response using Stable Gated DeltaNet-2, or SGDN-2.

Let $t$ index turns, let $X_t$ be the current input, and let $Y_t$ be its response. Storage processes their concatenation $S_t=X_t\Vert Y_t$, where $\Vert$ means concatenation. Suppressing layer and channel indices, its state progression is

$$\widetilde M_{t,0}=M_t,\qquad \widetilde M_{t,i}=U_\phi(\widetilde M_{t,i-1},\widetilde h_{t,i}),\qquad M_{t+1}=\widetilde M_{t,L_t}.$$

Here $L_t$ is the number of tokens in the completed interaction, $\widetilde h_{t,i}$ is a token's pre-attention representation in the storage procedure, and $U_\phi$ denotes the SGDN-2 update with learned processing parameters $\phi$. In words: start from the existing state, apply the storage update once per token, and carry the final state into the next turn.

While generating $Y_t$, memory utilization reads $M_t$, not the intermediate storage states $\widetilde M_{t,i}$. The response tokens are part of the interaction that storage subsequently incorporates.

Quiz

Which sequence correctly distinguishes memory utilization from storage in a turn?

Reset memory to zero before every turn, then update the learned projection weights during generation.

Generate while reading the existing memory, then scan the completed input and response to produce the next turn's memory.

Update memory only from the input; the completed response is excluded from storage.

Generate one token, immediately write it into persistent memory, and let the next generated token read that updated state.

$ check

Why the State and the Learned Parameters Are Different Things

Now the version for readers who want more of the machinery: the memory that changes between interactions is a tensor state, not an online update to the learned projection weights.

At selected layer $l$, the state has the form

$$M_t^{(l)}\in\mathbb{R}^{H_{kv}\times d_k\times d_v}.$$

Here $H_{kv}$ counts memory channels, and $d_k$ and $d_v$ are the key and value dimensions. Each channel therefore holds a key-by-value matrix. These states start at zero at the beginning of a new trajectory and evolve across its interactions.

The learned parameters $\phi$ include query, key, and value projections; erase and write gating projections; forget-gate components; and output-gating projections. They are optimized during training and remain fixed during online interaction. A changing matrix state and fixed learned update machinery are different kinds of memory: the former contains online information, while the latter determines how information is processed.

Storage uses the pre-attention representations of the completed $X_t\Vert Y_t$ interaction. For each channel and token, projections form keys and values, and gating machinery supplies erase, write, and forget controls. The token scan begins from that channel's $M_t^{(l)}$ and ends at $M_{t+1}^{(l)}$. Intermediate scan states belong to storage; they are not successive versions of the state consulted while generating the same response.

That defines what changes and when. Why can the erase part of this update need a stability fix?

Why Symmetrizing the Erase Step Matters

Here is the part that gets technical: the stability issue is a mismatch between the direction used to measure an erase and the direction along which its correction is applied.

In the GDN-2 erase analysis, the measurement uses $n=b\odot k$, where $b$ is the erase-gate vector, $k$ is the key vector, and $\odot$ denotes elementwise multiplication. The correction acts along $k$. Those directions need not align, making the erase operator non-symmetric.

The symmetrization instead uses

$$u=\sqrt{b}\odot k,\qquad E=I-uu^\top.$$

The square root is elementwise. Both sides of the rank-one subtraction now use the same vector $u$, so $E$ is symmetric. This is the isolated erase operator, not the complete gated storage update.

To see the consequence, decompose a state vector into a component parallel to $u$ and a perpendicular component:

$$x=x_{\parallel}+x_{\perp},\qquad Ex=(1-\lVert u\rVert_2^2)x_{\parallel}+x_{\perp}.$$

The erase scales the parallel component and leaves the perpendicular component unchanged. With the forget operator set to identity, no write term, and the stated assumption

$$0<\lVert u\rVert_2^2<1,$$

this map contracts along $u$ without expanding perpendicular directions. It is non-expansive in Euclidean norm under those conditions.

Contrast this with the original erase geometry. For a unit vector $e$ orthogonal to $k$, a mismatched correction can produce

$$e'=e-\gamma k.$$

Orthogonality removes the cross term, giving

$$\lVert e'\rVert_2^2=1+\gamma^2\lVert k\rVert_2^2.$$

A nonzero correction along a nonzero key increases the norm even though the operation is meant to erase. Symmetrization removes this particular coupling between perpendicular directions. It does not establish that arbitrary writes, gates, or complete trajectories can never produce large states.

Once storage has produced a state, how does that state enter the next response?

Why Memory Joins Attention Rather Than Replacing It

Now the deeper architectural point: a read of persistent memory is an additional attention path, not a replacement for the backbone's ordinary attention or an instruction to update memory during decoding.

At a selected layer, both paths receive the pre-attention hidden representation. Ordinary attention runs through the backbone path. The memory path projects a query, reads the turn's current state $M_t^{(l)}$, and processes the result with learned output gating. Their outputs are combined before subsequent computation.

The backbone's pre-attention, attention, and post-attention operations are preserved. The additional MemoryAttention path makes previously accumulated information available during autoregressive generation.

There are two distinct uses of “parallel” here. Architecturally, memory attention runs alongside ordinary attention. Procedurally, memory utilization and storage have separate jobs. Neither means that a generated token must immediately write a new state that the next generated token reads. Within turn $t$, utilization queries $M_t$; storage scans the completed interaction to obtain $M_{t+1}$.

This leaves a practical question: how do you keep those two clocks explicit in an implementation?

Minimal Causal Code

This plain PyTorch control-flow sketch shows one selected layer's memory. The supplied callbacks stand for the backbone operations, gated memory read, and SGDN-2 arithmetic. In particular, sgdn2_step includes the projections and erase, write, and forget operations; this is a timing skeleton, not a numerical implementation of that update.

import torch

@torch.no_grad()
def run_turn(prompt, memory, max_new_tokens, *,
             decode_hidden, backbone_attention, memory_attention,
             combine, post_attention_and_sample, is_end,
             storage_hidden_states, sgdn2_step):
    # memory is M_t. Do not modify it during generation.
    response = []
    for _ in range(max_new_tokens):
        prefix = prompt + response
        hidden = decode_hidden(prefix)
        ordinary = backbone_attention(hidden)
        retrieved = memory_attention(hidden, memory)
        token = post_attention_and_sample(combine(ordinary, retrieved))
        response.append(token)
        if is_end(token):
            break

    # Storage is triggered for the completed interaction X_t || Y_t.
    # These are storage representations, not a per-token decode write.
    complete_turn = prompt + response
    storage_hidden = storage_hidden_states(complete_turn)
    working = memory.clone()  # M_t -> initial storage scan state
    for hidden in storage_hidden:
        working = sgdn2_step(working, hidden)

    next_memory = working  # Commit M_{t+1} after the full scan.
    return response, next_memory

# Initialize once per new trajectory, not once per turn.
# H_kv, d_k, and d_v are the configured memory dimensions.
memory = torch.zeros(H_kv, d_k, d_v)
# For each turn: response, memory = run_turn(prompt, memory, ...)

For multiple selected layers, carry a separate state for each layer. The important invariant is unchanged: decoding reads the existing turn state without mutating it; the subsequent storage scan produces the state carried to the next turn. The sketch omits batching and execution optimizations so that this boundary stays visible.

What You Gain, and What You Lose

You gain an explicit state that evolves across interactions, a learned way to store the completed interaction, and a dedicated read path alongside ordinary attention. SGDN-2 addresses the specific non-symmetric erase mechanism described above.

You also take on projections, gates, per-layer state management, and a storage pass over the completed input and response. Token-granular storage does not mean newly generated content is immediately written into persistent memory during that response. Current-turn information remains available through the backbone's context processing; its incorporation into the persistent state serves later turns.

Finally, the erase-only non-expansion argument is narrower than a guarantee about the full memory system. The mechanism is neither an exact transcript store nor simply ordinary attention made faster: it is an additional, learned route for carrying information across turn boundaries.