Depth Memory

The transformer architecture underpins modern language modeling. Since its introduction, architecture research has mainly seen innovation along two axes: sequence mixing, which includes softmax attention, linear attention, and state-space models, and channel mixing, through works like gated MLPs and MoEs. The standard mechanism for propagating information across depth, however, has remained largely unchanged since the introduction of ResNet: layer outputs are accumulated additively through residual connections. Recent works have established alternatives to this standard mechanism at scale. For example, Hyper-Connections (HC) and Manifold-Constrained Hyper-Connections (mHC) maintain an expanded residual stream, while Attention Residuals (AttnRes) explicitly attend to previous layer outputs. Analogous to sequence and channel mixing operators, we refer to these methods as depth mixing mechanisms.

Our Depth Memory framework describes depth mixing through a set of streams and their read and write operations. For each token, a read forms the next layer's input, and a write incorporates the layer output into the streams. This framework unifies standard pre-layer-norm (PreLN) transformers, AttnRes, and mHC.

The simplest version of a linear associative memory over depth carries two streams: the standard PreLN residual stream \(\mathbf x_l\) and a matrix-valued memory \(\mathbf M_l\). Each layer writes its output \(\mathbf v_l\) into both streams. It adds the output to the residual stream and stores it in the memory under a key \(\mathbf k_l\):

The next layer then reads \(\mathbf x_l+\mathbf M_l\mathbf q_l\) with a query \(\mathbf q_l\). Queries and keys are computed from the residual stream by small MLPs with \(L^2\)-normalized outputs:

Since \(\mathbf M_l=\sum_{i\le l}\mathbf v_i\mathbf k_i^\top\), this read is unnormalized linear attention over depth. Instead of a growing collection of layer outputs, as in AttnRes, the network keeps a fixed-size summary of them. Keeping the two streams separate matters. Supplying only the memory readout to each layer, or removing the residual addition after it, underperforms PreLN in our ablations.

The outer-product write only approximates what the linear depth memory should hold: the linear map that best reconstructs every layer output so far from its key. Akin to test time regression over time, we call this the depth regression objective:

Along the sequence axis, the need for parallelizability of the solve requires approximate solutions to the test time regression objective. Over depth, however, layers are already evaluated one after another. LDM can therefore maintain the exact solution with recursive least squares (RLS). It solves a weighted and regularized version of the objective, with learnable per-layer forgetting factors \(\lambda_l\), dynamic importance weights \(\omega_l\) that the key's MLP also produces, and a ridge penalty scaled by \(\alpha^{-1}\):

Solving it with RLS adds a third stream, the inverse key covariance \(\mathbf P_l\). With the weighted observations \(\bar{\mathbf k}_l=\sqrt{\omega_l}\mathbf k_l\) and \(\bar{\mathbf v}_l=\sqrt{\omega_l}\mathbf v_l\), each layer computes a gain \(\boldsymbol\kappa_l\) and the prediction error \(\mathbf u_l\) of the current memory:

The memory is corrected by what it failed to predict, and the inverse covariance is updated alongside it:

We initialize \(\mathbf M_{-1}=\mathbf 0\) and \(\mathbf P_{-1}=\alpha\mathbf I\), with \(\alpha=1\) in our experiments. Our main models partition the memory into 16 groups. Each group has its own inverse covariance and forgetting factors and is queried independently. The group readouts are concatenated into the full memory readout.