Reconstructing and Batching Depth-Memory Updates

Why People Wanted It

A recurrent memory changes as the network moves through its layers. The straightforward implementation keeps each intermediate state so backpropagation can revisit it, and reads or writes that state at every layer. That is convenient, but state-sized tensors are costly to save and move. The real question is: can we recover the states when they are needed, and avoid writing every intermediate version to memory?

The Core Move

Keep compact update contexts instead of every memory state, then reconstruct states backward and defer forward writes across short windows. For the memory update

$M_l=M_{l-1}+u_l\kappa_l^\top,$

where $u_l$ is the write's residual value and $\kappa_l$ its gain, the previous state is simply the current state minus that update.

Why the State Can Be Unwound

Here is the part that gets technical: reversibility applies to the memory state because the update's ingredients are saved, not because the whole computation is magically invertible. During the forward pass, retain the pseudo-value $u_l$, gain $\kappa_l$, key $k_l$, gated query $\tilde q_l$, forgetting factor $\lambda_l$, and the prior inverse covariance $P_{l-1}$. Then, starting with the final state, recover each preceding state:

$M_{l-1}=M_l-u_l\kappa_l^\top.$

This subtraction restores the exact previous memory if the saved update context is available. The context also supports the gradients: the read gradient needs $M_l$, while the write computation needs $M_{l-1}$. The inverse covariance is saved rather than unwound, and its symmetry means only its nonredundant triangular entries need be kept. So the backward pass walks the recurrence in reverse, reconstructing a state just before it is needed rather than retaining all states from the forward pass.

Why Windowed Writes Avoid State Traffic

Now the version for readers who want more of the machinery: the same additive update lets a short run of writes be represented as a sum of rank-one changes to the state at the window's start. For a window from $b+1$ through $i$,

$M_i=M_b+\sum_{b<m\le i}u_m\kappa_m^\top.$

In words, the state at any point inside the window equals its starting state plus the updates accumulated so far. A product of that state with a vector can therefore be evaluated from $M_b$ and the saved update vectors, without materializing each intermediate matrix. Reads and the key products required to form writes use this representation; only the window's final state is written out.

Backward follows the same structure: undo the window's accumulated updates once to recover its starting state, and compute the needed products and gradients from that state plus the window contexts. Intermediate state gradients inside the window need not be formed. The window length trades fewer state-sized transfers for more accesses to vector-sized update values; the reported main experiments use a window length of four. How large a window is worthwhile depends on that traffic tradeoff.

Minimal Causal Code

This loop shows the state arithmetic and the reverse reconstruction. It is intentionally simple: it materializes each state and does not implement window fusion or all gradients.

import torch

# M is the memory state; contexts hold each additive update.
M = torch.zeros(value_dim, key_dim, device=device)
contexts = []

for k, v, gain in steps:
    u = v - M @ k
    contexts.append((u, gain, k))
    M = M + u[:, None] * gain[None, :]

# The forward pass retained only the final M plus update contexts.
for u, gain, k in reversed(contexts):
    M_prev = M - u[:, None] * gain[None, :]
    # M_prev is now the exact state before this write.
    M = M_prev

What You Gain, and What You Lose

You avoid retaining a full memory state for every layer, and windowing reduces state-sized writes and reads by keeping intermediate states implicit. The price is extra bookkeeping: compact contexts must be saved, and reconstructing products from a window requires accessing its update vectors. Window size is therefore a hardware-aware balance, not a free reduction in work. This is an implementation strategy for the memory recurrence, not a claim that every operation in the model becomes cheaper.

Quiz

What makes it possible to reconstruct a prior memory state during the backward pass?

The inverse covariance alone contains a copy of every prior memory state.

The memory state is unchanged by writes.

The backward pass reruns the entire forward model with different weights.

The update is additive, and its pseudo-value and gain are saved as context.

$ check

Quiz

What is the main effect of deferring writes within a window?

It avoids materializing every intermediate memory state, reducing state-sized traffic at the cost of accessing update vectors.

It removes the need to compute reads inside the window.

It makes the memory update non-recurrent.

It reduces all memory and vector traffic without a tradeoff.

$ check