Paper: arXiv, ICML'26 | Software: GitHub | Keywords: interpretability, long context, attention heads, causal interventions

Abstract: How do language models recover the order of what they read across a long context? Humans and long-context LLMs show the same behavioral signature on this task. In our ICML 2026 paper, we then localize a mechanism that supports this behavior in LLMs. Temporal information is widespread during encoding, but at retrieval a single attention head reinstates a one-dimensional temporal code that is causally important for the judgment. Projecting that direction out of the head's output at the query positions drives order-judgment accuracy toward chance; scaling it up improves accuracy on sequence order recall tasks.

The question

Long-context language models can retrieve content from far back in their context. Recovering when that content appeared is a separate problem. Given two passages from a document the model read tens of thousands of tokens earlier, can it say which one came first, and what in the network computes the answer?

Humans and long-context LLMs share the same behavioral signature here: judging order is hard for passages that occurred close together, and gets easier as they move farther apart. For the models, the mechanistic question is how this behavior is implemented.

This question connects to a line of work relating in-context memory in transformers to episodic memory in humans, which has drawn correspondences between key-value retrieval and self-attention,¹ between in-context learning and computational models of episodic recall,² and between the structure of in-context representations and hippocampal function.³ Temporal order is a concrete point of contact. It is a defining property of episodic memory, and it is directly measurable in a language model.

There are two candidate mechanisms. Order can be reconstructed from content, using causal and narrative structure in the text, without any explicit representation of position. Alternatively, the model can carry a temporal code: a signal monotonic in sequence position, written into token representations during the forward pass and read back at query time.

Our results identify the second as a key mechanism. Temporal information turns out to be widespread across attention heads while the text is read, but at query time a single head reinstates a one-dimensional temporal code that the model relies on for the judgment. Removing the direction from that head’s output, at the query positions only, reduces order-judgment accuracy from 0.68 to 0.52 in Llama 3.1-8B.

Task

The Sequence Order Recall Task (SORT) has two phases within a single context. In the encoding phase the model reads a long text. In the retrieval phase it is shown two non-overlapping 50-word passages taken verbatim from that text, labelled A and B at random, and asked which appeared first. The final prompt then consists of a single long-context prompt in which the full source text appears first (the encoding span), followed by a short task instruction, and the two 50-word segment.

We score by comparing the model’s probabilities for the two labels under both label assignments. There is no external lookup and no iterative search over the source text, so the answer has to come from representations formed while processing the context. The main stimulus is The Murder of Roger Ackroyd, about 93k tokens.

The behavior to be explained

Human readers show a clear distance effect: judgments are near chance for nearby segments and improve as the segments move farther apart. Both Llama models reproduce this profile and reach at least 0.70 accuracy for widely separated pairs.

To test whether this reflects reconstruction from content, we permuted blocks of four sentences throughout the text. Local sentence content survives, while global narrative coherence and causal dependencies are substantially disrupted. Performance is largely preserved under shuffling, and both models retain a strong distance effect. The distance effect therefore does not require intact global narrative coherence or causal structure. Because these shuffled variants introduce new document-level orderings, memorization of the public-domain novel cannot explain their global order.

Where is temporal order represented?

For each head we form one representation per segment in each phase (encoding and retrieval): the mean of the head’s value vectors over the segment’s tokens as the text is read, and the mean of the head’s output over the same tokens when that segment reappears in the query. Stacking \(N\) segments in ground-truth order gives

\[ R_{\mathrm{enc}}^{(\ell,h)},\; R_{\mathrm{ret}}^{(\ell,h)} \in \mathbb{R}^{N \times d_h}. \]

We then look for a direction \(w \in \mathbb{R}^{d_h}\) whose projection scores \(Rw\) sort the segments into their true order, measured by Kendall’s rank correlation

\[ \hat\tau(R; w) = \tau_K\big(\hat\pi(Rw),\, (1, 2, \dots, N)\big), \]

where \(\hat\pi\) is the permutation that sorts the projections. For each phase separately, we take the temporally oriented first principal component as the readout direction \(w\), giving each head an encoding-fitted direction and a retrieval-fitted direction. Temporal orientation here is needed to make sure projected values are expected to be increasing or decreasing because principal components are eigenvectors and thus only defined up to a sign. For the reinstatement analysis below, we then ask whether the retrieval-fitted direction also orders that head’s encoding representations: a genuine reinstatement mechanism should recover temporal order in both phases along the same axis.

That choice is motivated by seriation, the problem of recovering a one-dimensional ordering of items from their pairwise similarities. If a head carries a temporal code, the similarity structure of the rows of \(R\) should admit such an ordering; we use the top principal component as a simple readout of it. Ridge regression probes of segment position produce the same head rankings.

We fit these directions on a 60k-word sequence of i.i.d. English nouns, which carries no narrative or coherent semantic structure, then freeze them and evaluate on five held-out documents: four shuffled novels and a cleaned House of Commons transcript truncated to 100k tokens. This transfer shows that the readout is not tied to the semantics of the training sequence or any one document.

Our results indicate that temporal information is widespread in encoding-phase value representations, and that retrieval-phase temporal structure is more concentrated. Interestingly, the stronger localization result comes from asking which heads preserve the same temporal axis across encoding and retrieval.

Is the head retrieving time, or inheriting it?

A head’s retrieval representations can be temporally ordered without the head retrieving anything, if an upstream layer writes position information into the residual stream at the query tokens and the head inherits it. To separate the two cases, we require the same fixed direction to order both \(R_{\mathrm{enc}}\) and \(R_{\mathrm{ret}}\), and define a temporal reinstatement score

\[ \mathrm{TRS}^{(\ell,h)} = \mathrm{sign}\Big(\min\big(\bar\tau_{\mathrm{enc}},\, \bar\tau_{\mathrm{ret}}\big)\Big) \cdot \big|\bar\tau_{\mathrm{enc}}\, \bar\tau_{\mathrm{ret}}\big|, \]

with both rank correlations computed using the retrieval-fitted direction and averaged over held-out documents. The sign term assigns a negative score to heads whose ordering flips between phases.

The distribution is striking. One head dominates this score in each model. In Llama 3.1-8B it is L8H11, among 1,024 heads, and in Llama 3.1-70B it is L14H21, among 5,120. Running the same localization on two models from other families gives the same picture: the top reinstatement head is L7H18 in Mistral-7B-v0.2 and L14H0 in Qwen2.5-7B.

Both Llama heads act as retrieval heads for the queried segments, placing 64% and 39% of their query-time attention mass respectively on the segments’ original encoding spans.

Causal interventions

To show that the heads are uniquely causally implicated in the model’s task performance, we perform causal interventions that remove or amplify the time-coding directions in the retrieval-phase representations of each head.

Let \(\hat w\) be the temporal direction of the top reinstatement head and \(\mathcal{T}^{\mathrm{ret}}_{\mathrm{seg}}\) the token positions of the two segments in the query. We replace the head’s output \(o_t\) for \(t \in \mathcal{T}^{\mathrm{ret}}_{\mathrm{seg}}\) with

\[ \Pi_{\perp \hat w}(o_t) = o_t - \langle o_t, \hat w\rangle\, \hat w, \]

leaving all other positions, heads and directions unchanged. The encoding pass is unaffected, so the intervention targets reinstatement rather than storage.

Accuracy falls by 0.15 in the 8B and 0.11 in the 70B. Sweeping every head in the 8B, L8H11 is the only one whose temporal-direction ablation substantially reduces accuracy. Across the 180 heads tested in the 70B, L14H21 is the only one whose removal reduces accuracy; removing the temporal direction of other heads slightly increases it.

We also scale the component along a direction rather than removing it, replacing \(o_t\) with

\[ \tilde o_t = o_t + (\alpha - 1)\langle o_t, w\rangle\, w. \]

At \(\alpha = 1.75\) accuracy in the 8B increases by up to 0.09 over the unmodified model. Random directions have no comparable effect. Together with the ablation result, this shows that the temporal component is causally important for the judgment and that, at baseline, Llama 3.1-8B under-weights it for SORT.

Conclusion

Temporal information is distributed across many attention heads while the text is read, but at query time a single head reinstates it. The reinstated direction is causally important for SORT order judgments, and amplifying it can improve accuracy.

Several questions follow from this. The first is what these heads do outside our task. We isolated them with an explicit ordering question, but a code that tracks position in the context could be recruited wherever recency, distance or event structure matters, including in behaviors that are never framed as ordering.

The second is how the code is formed. It is not simply inherited from RoPE: temporal information in these heads survives, and in fact strengthens, when RoPE is disabled, while a randomly initialized model shows little temporal structure. The contrast with random initialization indicates that the code is learned during training rather than inherited directly from RoPE; what shapes it remains open.

The third is how this relates to computational models of episodic memory, in which retrieval works by reinstating the context state that was active during encoding.⁴ A mechanism of that form arising in a system trained only to predict text is worth understanding on both sides of the comparison.

Finally, the attention scores of these heads offer a natural candidate readout. Because one head dominates temporal reinstatement, where it attends during generation may indicate which part of the context the model is drawing on, giving a compact signal of which episode is currently being retrieved.

Resources

• Paper: arXiv (ICML 2026), by Mathis Pink, Vy Ai Vo, Qinyuan Wu, Jianing Mu, Javier Turek, Uri Hasson, Kenneth A. Norman, Sebastian Michelmann, Alexander Huth and Mariya Toneva
• Code: github.com/mathispink/temporal-context-reinstatement
• Toolkit: ICMemory-Kit, for long-context memory experiments and mechanistic interpretability
• Human data: 990 order judgments from 97 readers of a full-length novel

Funded in part by the Deutsche Forschungsgemeinschaft (DFG), GRK 2853/1 “Neuroexplicit Models of Language, Vision, and Action”, project number 471607914.

References