Linear Depth Memory

Why People Wanted It

A transformer can mix earlier layer representations by keeping them available for later layers. The problem is the growing collection. A matrix memory avoids that collection, but simply adding associations does not account for overlap among their keys.

So the real question was never just how to summarize earlier layers: can a fixed-size state maintain a corrected fit between layer keys and layer outputs as depth increases?

Linear Depth Memory (LDM) uses a regression-based associative memory alongside the ordinary PreLN residual stream. The residual stream remains; the memory adds another route for information across layers.

The Core Move

Treat keys as inputs to a linear predictor and layer outputs as its targets, then maintain the regression solution recursively rather than accumulating uncorrected associations.

Stack the keys as rows of $K_{\le l}$ and the corresponding values as rows of $V_{\le l}$. When the key Gram matrix is invertible, the unweighted depth-regression solution is

$$M_l^* = V_{\le l}^{\top}K_{\le l}\left(K_{\le l}^{\top}K_{\le l}\right)^{-1}.$$

The memory maps a key-shaped query to a value-shaped readout. The inverse Gram matrix is the crucial difference from an outer-product accumulator: it accounts for correlations among the keys.

LDM's main instantiation generalizes this fit with learned forgetting, dynamic importance weights, and a weighted ridge penalty. Recursive least squares (RLS) maintains that generalized solution online.

Quiz

What makes the regression-based LDM update different from simply adding key–value outer products?

It replaces the ordinary residual stream entirely.

It accounts for correlations among keys when correcting the memory's fit.

It stores every earlier layer output for exact lookup.

It guarantees error-free retrieval of every stored association.

$ check

Why the Regression Objective Changes the Update

Now the version for readers who want more of the machinery: an outer-product accumulator drops the correction for key correlations that appears in the exact regression solution.

With the row-stacking convention above, the cross-product in that solution is

$$V_{\le l}^{\top}K_{\le l}=\sum_{i=0}^{l}v_i k_i^{\top}.$$

This is precisely the accumulated key–value outer-product memory. Approximating the key Gram matrix by the identity therefore recovers the simple additive memory. The distinction is not whether both approaches store associations; it is whether the predictor corrects for the geometry of the keys used to store them.

The generalized objective changes the fit in three ways. Per-layer forgetting factors change how earlier observations remain influential as depth advances. Dynamic importance weights change how much individual observations count. A weighted ridge penalty regularizes the mapping. These are parts of the regression being solved, not gates attached after an otherwise unchanged solution.

This also separates two meanings of “exact.” RLS maintains the specified regression solution exactly; it does not promise that the compressed predictor reconstructs every stored value without error. What state makes that recursive solution possible?

Why RLS Needs an Inverse-Covariance State

Here is the part that gets technical: the memory matrix alone is insufficient for the covariance-aware recursive correction, so LDM also carries an inverse key-covariance matrix.

The initialization is

$$M_{-1}=0,\qquad P_{-1}=\alpha I.$$

Here $P$ is the inverse-covariance state and $I$ has the key dimension. The observations are scaled by the square root of their importance weight before the RLS calculation. The update computes a prediction error against the previous memory and a gain informed by the inverse covariance, then updates both states through a Sherman–Morrison rank-one correction.

Mechanically, the prediction error says what the current mapping fails to predict. The gain determines how that error changes the mapping given the keys already observed. Updating the inverse covariance carries that information into the next layer's correction. Forgetting and importance weighting modify this recursive fit rather than replacing it with a plain outer-product sum.

This module describes those operations without supplying an explicit weighted gain or covariance recurrence. The code below consequently keeps the numerical RLS solver as an explicit dependency rather than substituting an unweighted recurrence for the generalized update.

Depth mixing makes this solver choice natural: transformer layers already execute sequentially, so the regression update can follow that order. But should every part of the output share the same fit?

Why Groups and Shared Latents Change the Read

Now one layer deeper: grouping partitions the value dimension while giving each partition its own independently queried regression memory and inverse-covariance state.

For model dimension $D$ and $G$ groups, each group produces a readout of dimension $D_v=D/G$. Concatenating those readouts restores the full value dimension. The grouped inverse-covariance state has shape

$$P_l\in\mathbb{R}^{G\times D_k\times D_k}.$$

Forgetting factors and gates are learned independently for each group. This is more than slicing one shared prediction: the groups carry separate covariance information and are queried independently.

The parameterization shares intermediate representations without merging these roles. On the read side, learned affine maps produce a shared latent representation, from which the query and per-group read gate are generated. A learned per-group gain controls the read magnitude. On the write side, another shared latent supplies the key and groupwise importance weights used by the RLS write.

Sharing a latent therefore does not mean identifying a query with a gate, or a key with an importance weight. They are distinct outputs with distinct jobs. The benefit is independently controlled reads and fits; the cost is maintaining their separate state. How much of that flexibility is worth the additional machinery?

Minimal Causal Code

This plain PyTorch skeleton makes state initialization and replacement visible. It is not a standalone RLS implementation: rls_write must implement the generalized weighted, forgetting-aware RLS write. Groups and the surrounding residual operations are omitted.

import torch


def depth_memory_loop(observations, Dv, Dk, alpha, rls_write):
    # Single-group state; observations arrive in layer order.
    M = torch.zeros(Dv, Dk)
    P = alpha * torch.eye(Dk)

    for k, v, forgetting, importance in observations:
        # Both previous states participate in the regression update.
        next_M, next_P = rls_write(
            M, P, k, v, forgetting, importance
        )

        # These, not a cache of all previous pairs, survive to next layer.
        M = next_M
        P = next_P

    return M, P

The important dependency is visible: the next write needs both the previous predictor and its inverse covariance. An outer-product accumulator would not carry that second state.

What You Gain, and What You Lose

LDM replaces a growing collection of layer representations with fixed-size matrix state for a chosen memory configuration. It maintains a specified regression fit instead of treating each association as an independent additive contribution.

The cost is a more involved state and update. RLS carries inverse covariance alongside the memory; grouping maintains separate covariance states; learned forgetting, importance weights, and read gates add control machinery. Exact regression also remains a compressed linear fit, not an exact lookup table for arbitrary earlier representations.

The ordinary residual stream still exists. LDM adds a regression-based information path rather than making the familiar path disappear. It is a different way to retain and retrieve depth information, not simply the same layer cache made cheaper.

Quiz

What additional state does RLS-based LDM maintain beyond the associative memory matrix?

A softmax attention matrix over all previous layers.

Inverse key covariance, maintained separately for each group in the grouped version.

An uncompressed list of every previous key and value.

No additional state; the memory matrix alone determines the RLS correction.

$ check