Recurrent Memory › Transformer-XL: Segment-Level Recurrence

Now that we understand the problem, let's explore how Transformer-XL solves it. The solution is remarkably simple: instead of discarding the previous segment entirely, cache its hidden states and make them available during attention computation for the current segment.

Think of it like this: when you read a new paragraph, you don't forget the previous one. You carry forward a mental summary of what came before. That summary isn't the original words themselves, but your processed understanding of them. Transformer-XL does exactly this, but with hidden states instead of mental summaries. The cached states are not raw token embeddings. They are processed representations that already encode relationships between tokens, contextual information from within that segment, and whatever higher-level abstractions the network has learned to extract.

The key insight is that we do not need to preserve the original tokens from previous segments. We need to preserve the information they contained in a form that is useful for processing the current segment. Hidden states serve this purpose because they encode meaning rather than surface form alone. A token's hidden state at layer 3 is less like "the word 'cat'" and more like "a specific cat being discussed in a context involving a mat and some sitting." That richer representation is what gets cached and carried forward.

This approach also fits naturally with the autoregressive nature of language modeling. When generating text one token at a time, we process segments sequentially anyway. Caching the hidden states of the just-processed segment and feeding them to the next segment adds minimal overhead to the already-sequential processing. There is no need to maintain an exponentially growing memory buffer; we only ever need the most recent segment's hidden states.

The recurrence operates at each layer independently. When processing layer $n$ of segment $\tau$, the model:

1. Retrieves the cached hidden states $\mathbf{h}_{\tau-1}^{n-1}$ from processing the previous segment at layer $n-1$
2. Concatenates these cached states with the current segment's hidden states
3. Computes attention where queries come only from the current segment, but keys and values span both the cached and current states
4. Caches the current segment's output hidden states for use when processing the next segment

This creates an asymmetric attention pattern: current tokens can "look back" at cached tokens, but we never recompute outputs for the cached tokens. They're frozen representations from the previous forward pass. This asymmetry is not a limitation but a design choice. Recomputing the outputs of cached tokens would require backpropagating gradients through an unbounded history, which is exactly the computational cost we are trying to avoid. By treating cached states as read-only context, Transformer-XL achieves information flow across segments at constant computational cost per segment.

Another subtle but important point is that the cached states are from layer $n-1$, not layer $n$. When processing layer $n$, we concatenate the previous segment's layer $n-1$ output with the current segment's layer $n-1$ output and compute attention. This is not an implementation detail but a consequence of how information flows through a multi-layer network. At layer $n$, the current segment's representation already incorporates information from lower layers. The cached representation we need to extend must be at the same "processing depth" to combine coherently with the current representation.

Let's formalize this mechanism. Suppose we're processing segment $\tau$, which contains $L$ tokens. The previous segment's hidden states, also of length $L$ (or some memory length $M$), have been cached. At layer $n$, we want to compute attention over an extended context that includes both segments.

The extended context for attention at layer $n$ concatenates the stop-gradient of the previous segment's hidden states with the current segment's hidden states. Concretely:

$$\tilde{\mathbf{h}}_\tau^{n-1} = \text{StopGrad}(\mathbf{h}_{\tau-1}^{n-1}) \circ \mathbf{h}_\tau^{n-1}$$

where:

• $\tilde{\mathbf{h}}_\tau^{n-1}$: the extended hidden states combining previous and current segments, used as keys and values
• $\mathbf{h}_{\tau-1}^{n-1}$: cached hidden states from the previous segment at layer $n-1$
• $\mathbf{h}_\tau^{n-1}$: current segment's hidden states at layer $n-1$
• $\text{StopGrad}(\cdot)$: stops gradient flow to prevent backpropagating through the cached states
• $\circ$: concatenation along the sequence dimension, resulting in a tensor of length $M + L$

Why does StopGrad matter? If we allowed gradients to flow back through the cached states, we would need to store the computation graph for the previous segment's forward pass, then the segment before that, and so on, going back all the way to the beginning of the document. This would make the memory cost proportional to the total document length, defeating the purpose of segmentation. By stopping gradients at the cached states, each segment's backward pass touches only the current segment's computation, keeping memory usage constant regardless of document length.

The attention computation then becomes:

$$\mathbf{q}_\tau^n = \mathbf{h}_\tau^{n-1} W_q^n, \quad \mathbf{k}_\tau^n = \tilde{\mathbf{h}}_\tau^{n-1} W_k^n, \quad \mathbf{v}_\tau^n = \tilde{\mathbf{h}}_\tau^{n-1} W_v^n$$ $$\mathbf{h}_\tau^n = \text{Attention}(\mathbf{q}_\tau^n, \mathbf{k}_\tau^n, \mathbf{v}_\tau^n)$$

where:

• $\mathbf{q}_\tau^n \in \mathbb{R}^{L \times d}$: queries from the current segment only
• $\mathbf{k}_\tau^n, \mathbf{v}_\tau^n \in \mathbb{R}^{(L + M) \times d}$: keys and values from the extended context (current + cached)
• $W_q^n, W_k^n, W_v^n$: learnable projection matrices for layer $n$
• $M$: the length of the cached memory (typically equal to segment length $L$)

The critical detail is that queries come only from the current segment while keys and values include the cached previous segment. This asymmetry is intentional: we want current tokens to attend to past context, but we don't want to regenerate outputs for past tokens. The cached states are read-only: they provide context but don't receive updates.

Why does this formula make sense? Notice that the query dimension remains $L$ (only current tokens ask questions) while the key and value dimensions are $L + M$ (both cached and current tokens can be referenced). The attention matrix has shape $(L, L+M)$, which means each current token computes a score against every cached token and every preceding current token. The resulting weighted sum draws from both the cached and current portions of the extended context, exactly what we want.

Let's translate this into code. The core operation is straightforward: concatenate cached and current hidden states, project to queries/keys/values, and compute attention with appropriate masking.