Recurrent Memory › Worked Example: Tracing Information Through Segments

To make the segment-level recurrence concrete, let's trace how information about a specific token propagates across three segments of a simple model. Imagine a 2-layer model with segment length 4 processing the sentence: "The professor lectured. Students took notes. They asked questions."

We have three segments: ["The", "professor", "lectured", "."], ["Students", "took", "notes", "."], and ["They", "asked", "questions", "."]. The token "They" in segment 3 is a pronoun that refers to "Students" in segment 2. Let's trace how the model can resolve this reference.

Segment 1 processing. No cached memory exists. Layer 1 processes ["The", "professor", "lectured", "."] with standard causal attention. The hidden states at layer 1 encode things like "lecturer introducing a topic." Layer 2 receives these layer-1 states as keys and values and produces layer-2 hidden states. These layer-2 states are cached as $\mathbf{h}_1^2$ (segment 1, layer 2 output) and $\mathbf{h}_1^1$ (segment 1, layer 1 output).

Segment 2 processing. Layer 1 of segment 2 receives the layer-1 cached states from segment 1: $\mathbf{h}_1^0$ (the embedding layer outputs). The cache at layer $n$ uses the previous segment's layer $n-1$ output. Thus, layer 1 of segment 2 receives the embedding-level outputs of segment 1 as its extended key-value context. Layer 1 computes attention over ["Students" attended to "The", "professor", "lectured", "."] plus itself. The hidden state for "Students" at layer 1 now encodes "a plural subject, appearing after a lecturing context." Layer 2 of segment 2 receives layer-1 outputs of segment 1 as its extended context. The hidden state for "Students" at layer 2 now carries information about both the current segment and the previous one.

Segment 3 processing. Layer 1 of segment 3 receives the layer-1 outputs of segment 2 as its extended context. "They" at layer 1 can attend to the layer-1 representation of "Students" from segment 2. The layer-1 hidden state for "They" begins to encode a connection to the plural subject from the previous segment. Layer 2 of segment 3 receives the layer-2 outputs of segment 2 as its extended context. "They" at layer 2 now attends to the richer layer-2 representation of "Students," which already encoded the connection between "Students" and the lecturing context from segment 1. The final hidden state of "They" captures, through this chain of attention, a connection back to "Students" and even further to the "professor" from segment 1.

This trace illustrates several important properties. First, the coreference link between "They" and "Students" is resolved by a single attention operation at layer 1 of segment 3, which can directly attend to "Students" in the cached memory. Second, deeper connections (like the link between the students and the professor) are accessible through the layer-2 representations, which already aggregated information from segment 1 when processing segment 2. Third, the information path follows the recurrence chain: segment 3 layer 1 attends to segment 2 layer 0 cached states, and segment 3 layer 2 attends to segment 2 layer 1 cached states, which themselves incorporated segment 1 information.

The numerical details would require running an actual trained model, but the logical structure of the information flow is clear: information about tokens from prior segments is compressed into the hidden states and carried forward one segment at a time, with each layer having access to the corresponding layer's cached representation from the previous segment.