Dissecting AI
All modules

Depth Memory: Reading and Writing Across Layers

Layer outputs need a route to influence later layers; Depth Memory describes that flow as streams that layers read with queries and update with keys and outputs.

The short version

What it is

  • Depth Memory A way to describe information moving between neural-network layers as one or more states that each layer can read from and update.
  • Explicit and implicit mixing Explicit mixing keeps earlier layer representations as separate entries. Implicit mixing combines new information into a continuing state.

Why it matters

Residual addition, attention over stored layer representations, and expanded residual streams can be compared through the same questions: what state is carried, how does a layer read it, and how does the layer update it?

Before getting into the machinery, here’s the short version: Depth Memory describes what moves between layers as a memory state. Each layer reads from that state, does its computation, and writes its result back. This gives us one vocabulary for the familiar residual connection, attention over earlier layers, and an expanded residual stream.

  1. A common description: Layers need earlier information, but different architectures carry and expose it differently. Depth Memory makes those choices explicit: name the state, then ask how a layer reads and writes it.
  2. Different storage choices: A residual connection keeps one running representation. AttnRes keeps earlier representations available as separate key–value entries. mHC uses an expanded residual state instead. These are not the same memory with different labels; their states and write operations differ.

The canonical concrete example is a residual connection. Imagine a running notebook: each layer reads the current page, computes an update, and adds that update to the page. The notebook analogy helps explain accumulation, though it is not a literal record of separate layer outputs. In Depth Memory terms, the stream is the residual state, the read simply returns it, and the write adds the layer output. So a residual connection is already a memory scheme, even without a cache or attention lookup.

What changes if later layers should be able to retrieve earlier representations separately? AttnRes answers by keeping a growing set of key–value pairs. A layer uses a query to attend over that stored set, and its output is appended as another entry. The useful intuition is a notebook with individually addressable pages rather than one overwritten running page. That preserves explicit entries, but it also means the stored stream grows with depth.

mHC takes the other route: it keeps an expanded residual state of fixed size and mixes information into that state. A layer reads a combination of its state slots, then writes by mixing the prior state and adding a contribution from its output. Think of a fixed number of notebook pages whose contents can be recombined; unlike the AttnRes analogy, old entries are not kept as separate pages. The exact mixing is more structured than ordinary fading or erasing, but the analogy captures the storage tradeoff.

So the real question was never just whether layers use memory. It was what they keep available and how they access it. Depth Memory packages that into streams plus Read and Write operations: a shared interface, not a claim that residual addition, an explicit cache, and implicit state mixing behave identically. What other update rule could fit that interface without changing the basic questions?

Read the deep dive

Why People Wanted It

A transformer layer does not just transform a representation; it also has to decide what information from earlier layers reaches the next one. The familiar residual connection makes that choice in a fixed way: add the new output to the running state. Attention Residuals instead keep earlier representations available for later lookup. The real question is how to describe both choices—and expanded residual streams such as mHC—with one account of what layers read and how they write.

The Core Move

Treat the information carried through depth as one or more memory streams, and define each layer’s input and state update through a read and a write:

δl=Fl(Norm⁡l(Read⁡(Sl−1,ql−1))),Sl=Write⁡(Sl−1,kl,δl).\delta_l = F_l(\operatorname{Norm}_l(\operatorname{Read}(S_{l-1},q_{l-1}))),\qquad S_l=\operatorname{Write}(S_{l-1},k_l,\delta_l).

Here the layer reads with a query, computes its output, then writes that output using a key. The stream and the two operators determine the particular depth-mixing method.

Quiz
In the Depth Memory view, what determines a particular depth-mixing method?
The number of transformer layers, regardless of state representation.
The stream or streams it maintains and the Read and Write operators it uses.
A requirement that every method append every layer output to a cache.
Only the normalization function applied before each layer.

Why a Residual Connection Is Already a Memory

Now the version for readers who want the machinery: the residual stream is a single state, and its read ignores the query while its write adds the layer output. Specifically, set Sl=xlS_l=x_l, with Read⁡(x,q)=x\operatorname{Read}(x, q)=x and Write⁡(x,k,v)=x+v\operatorname{Write}(x,k,v)=x+v. The general framework therefore does not require a separate cache or an attention operation: a plain residual connection is one simple read/write policy.

Why Explicit and Implicit Mixing Differ

Here is the part that gets technical: the distinction is the state representation and what a write does to it. AttnRes uses a stream Ml=(K≤l,V≤l)M_l=(K_{\le l},V_{\le l}). Its read attends to the stored key–value pairs, and its write appends the new pair, so the stream retains previous layer representations explicitly.

For mHC, the stream is an expanded residual state XlX_l. Its read forms a query-weighted combination of the state slots, Read⁡(X,q)=Xq\operatorname{Read}(X,q)=Xq, while its write mixes the previous state and adds an output-weighted write, Xl=Xl−1Hlres(Xl−1)+vlkl⊤X_l=X_{l-1}H_l^{\mathrm{res}}(X_{l-1})+v_lk_l^\top. Unlike the growing AttnRes cache, this stream has fixed size; information is mixed into the state rather than kept as a separate entry for every layer. The framework calls the first style explicit depth mixing and the second implicit. What does this common vocabulary let us compare without claiming their memories behave alike?

Why the Framework Unifies Their Read and Write Roles

Now for the exact constraint that makes the abstraction useful: every instance specifies a tuple of streams and its own Read and Write operators. PreLN and AttnRes each use one stream, as does mHC in this formulation; the framework also permits multiple streams. The commonality is functional, not a demand that states share shape or storage: a read supplies the sublayer input, and a write updates the carried state using the layer output and its key. That is how a fixed residual update and attention over a growing cache fit the same interface. The regression-based linear-memory update is another instantiation, but its derivation is a separate topic. Which read/write design is appropriate depends on what information the next layer should be able to recover.

Minimal Causal Code

This intentionally simple loop makes the state transition visible. The read and write functions stand in for the architecture-specific operators; they are not optimized implementations.

state = initial_state
for layer in layers:
    query = layer.make_query(state)
    layer_input = read(state, query)
    value = layer(norm(layer_input))
    key = layer.make_key(state)
    state = write(state, key, value)

What You Gain, and What You Lose

The framework gives you a shared way to reason about information flow across depth: identify the state, ask what each layer can read, and inspect how its output is written back. That makes residual addition, explicit key–value storage, and fixed-size state mixing comparable at the level of their roles.

The abstraction does not make those mechanisms equivalent. A residual stream adds; AttnRes preserves a growing record and retrieves from it; mHC writes into a fixed-size expanded state. Those choices imply different state structures and access patterns. The framework is a vocabulary for comparing them, not a guarantee that one has the capacity, cost, or behavior of another.

Core Takeaway

Depth Memory asks two questions at every layer: what can you read from the carried state, and how do you write the new output back?