ArXiv: 2407.09450

🎯 Pitch

LLMs can now process 10 million tokens—a scale where full-context models collapse—by simply segmenting text using the model’s own moment-to-moment surprise, mimicking how the human brain chunks experiences. This zero-fine-tuning method outperforms both retrieval-augmented generation and full-context models on most long-range tasks, while its event boundaries align with human-perceived boundaries.

1. Executive Summary

EM-LLM introduces a training-free architecture that integrates key aspects of human episodic memory and event cognition into Transformer-based LLMs, enabling them to handle practically infinite context lengths by organizing token sequences into coherent episodic events using Bayesian surprise (the model's negative log-likelihood of observing the current token given prior tokens, thresholded to detect novelty-driven event boundaries) with graph-theoretic boundary refinement (optimizing modularity or conductance on attention-key similarity matrices to maximize within-event cohesion and cross-event separation) and a two-stage memory retrieval process (similarity-based k-NN retrieval combined with a temporal contiguity buffer that enqueues neighboring events to replicate human-like recall patterns). Evaluated on LongBench and ∞-Bench across five base LLMs, EM-LLM outperforms the state-of-the-art InfLLM retrieval model on 80% of task groups and surpasses both RAG methods (by 30.5% on LongBench with NV-Embed-v2) and full-context models on most tasks, while successfully performing passkey retrieval across 10 million tokens — a scale computationally infeasible for full-context models — establishing that surprise-based dynamic event segmentation and contiguity-aware retrieval can substitute for brute-force attention only when the underlying LLM's surprise signal correlates with human-perceived event boundaries.

2. Context and Motivation

The Core Problem: LLMs Cannot Process Truly Long Contexts

The fundamental problem this paper addresses is straightforward in its statement but pervasive in its consequences: contemporary LLMs, despite ever-expanding context windows, perform markedly worse on long sequences than on short ones, and the gap widens dramatically as context length increases beyond the model's training regime. This is not merely an inconvenience — it is a structural limitation that constrains how LLMs can be deployed in any setting requiring sustained, coherent interaction with large bodies of information.

The paper frames this through two distinct but compounding failure modes (Section 1):

Failure Mode 1: Attention Dilution. Softmax attention, the core mechanism in Transformers, aggregates information from all previous tokens into a single weighted-sum embedding at each generation step. As the sequence grows, this aggregated representation must encode information from an ever-larger set of tokens, becoming increasingly noisy and losing the distinctiveness needed to discriminate between relevant and irrelevant context. The paper cites Tworkowski et al. (2023) to support this: the weighted sums of value vectors "risk becoming excessively noisy and losing their distinctiveness." This is not a theoretical concern — it manifests as the well-documented "lost in the middle" phenomenon (Liu et al., 2024a), where LLMs can access information at the beginning or end of a context window but largely fail to retrieve facts buried in the middle.

Failure Mode 2: Positional Encoding Extrapolation. Transformer models are trained with a finite context window, and their positional encodings — the mechanism that tells the model where tokens appear in the sequence — are poorly suited to positions outside that training regime. The paper specifically calls out Rotary Position Embeddings (RoPE; Su et al., 2024), the dominant positional encoding scheme in modern open-weight LLMs (LLaMA, Mistral, Phi families). Kazemnejad et al. (2024) demonstrated that Transformers "struggle with extrapolating to contexts longer than their training window size," and the positional encoding is a primary culprit. When a model trained on 4K-token sequences encounters position 20,000, the rotational angles in RoPE have never been observed during training, and the resulting attention patterns degrade.

These two failure modes are orthogonal and compounding: even if you solve the positional encoding problem (which several approaches attempt, as discussed below), you still face attention dilution. Even if you solve attention dilution, you still face positional encoding failures. A comprehensive solution must address both.

Why This Problem Is Important (and Becoming More So)

The paper's introduction identifies three intersecting trends that make long-context processing an urgent priority:

Trend 1: The context window is the primary mechanism for incorporating external knowledge. Unlike training-time knowledge, which is frozen at the model's creation, the context window is where domain-specific information, private data, up-to-date facts, and user-provided documents enter the model's reasoning process. As the paper states, "the context window serves as the primary mechanism to incorporate domain-specific, private, or common up-to-date information" (Section 1). This makes context window capacity a ceiling on knowledge integration: a model with a 4K context window cannot reason over a 50-page legal document, a multi-hour meeting transcript, or a complete codebase. Expanding this ceiling directly expands what LLMs can do.

Trend 2: Real-world use cases demand increasingly long contexts. The paper's evaluation benchmarks reflect this: LongBench includes tasks like multi-document QA (HotpotQA, 2WikiMQA, Musique) where the model must synthesize facts across multiple documents, summarization of entire reports and meetings (GovReport, QMSum), and passage retrieval requiring location of specific information in long texts. ∞-Bench (Zhang et al., 2024) pushes this further with tasks spanning 100K+ tokens. These are not artificial stress tests — they model genuine applications in legal analysis, scientific research, customer support history, and codebase understanding.

Trend 3: The training-inference length gap is growing. While pretraining context windows have expanded (from 512 tokens in early GPT models to 128K in recent LLaMA and Mistral variants), the demand for inference-time context vastly exceeds what can be economically trained. Training with 1M-token context windows is computationally prohibitive because the quadratic complexity of attention makes pretraining costs explode. This creates an asymmetry: we can train models with moderately long contexts, but we need them to handle extremely long contexts at inference time. The paper explicitly targets this gap: EM-LLM requires no fine-tuning and works with existing pretrained models, addressing the inference-side problem without touching the training pipeline.

A fourth, unstated motivation runs beneath the surface: as LLMs are deployed in continual, interactive settings (long-running conversations, agent loops, personal assistants that accumulate history), the context grows unboundedly. Without mechanisms for organizing and selectively retrieving from this growing history, the model's usable memory effectively shrinks over time as the past dilutes the present. This is precisely the problem that biological episodic memory solves — and the paper's explicit inspiration.

Prior Approaches and Where They Fall Short

The paper identifies four broad categories of prior work, each with specific limitations that EM-LLM is designed to overcome. Understanding these limitations is essential to appreciating why the paper's approach is non-obvious.

1. Positional Encoding Extrapolation Methods

These methods modify the positional encoding to handle positions beyond the training length without retraining the entire model:

  • Positional Interpolation (Chen et al., 2023; Press et al., 2021): Scale down all position indices by a constant factor so that extended positions map back into the training range. For example, with a scaling factor of 8, position 32,000 maps to position 4,000 in the original encoding. The problem: this compresses the effective resolution of positional information, potentially making nearby tokens indistinguishable.

  • NTK-Aware Scaling (Xiong et al., 2023; Liu et al., 2024b): Rather than scaling positions uniformly, modify the base constant in RoPE to stretch the high-frequency components less than the low-frequency ones. This preserves local positional discrimination better than naive interpolation. The problem: it only addresses the extrapolation issue, not the attention dilution or computational cost issues. The model can technically attend to position 100,000, but the resulting attention weights may be diffuse and the computational cost is still quadratic.

  • YaRN (Peng et al., 2024): Combines NTK-aware scaling with a temperature adjustment to the softmax attention, effectively sharpening the attention distribution to compensate for the expanded positional range. The problem: same as above — it addresses extrapolation but not the fundamental attention quality and cost issues.

  • LongRoPE (Ding et al., 2024): Extends RoPE to 2 million tokens through a combination of progressive extension during fine-tuning and search over the optimal rescaling factors. The problem: requires some fine-tuning, still faces quadratic attention cost, and as Figure 1 (bottom) shows, performance on passkey retrieval still degrades at extreme lengths.

The fundamental limitation: All these methods address only Failure Mode 2 (positional encoding) while leaving Failure Mode 1 (attention dilution) and the computational cost problem untouched. They make it possible to attend to long sequences but do not make it efficient or effective.

2. Efficient Attention Mechanisms

These methods reduce the computational cost of attention, typically from O(n²) to O(n log n) or O(n):

  • Linear Attention (Katharopoulos et al., 2020): Replace the softmax kernel with a linear kernel, allowing the attention computation to be reordered from (Q × K^T) × V to Q × (K^T × V), changing the complexity from O(n²) to O(n). The problem: linear attention has consistently underperformed softmax attention in practice on complex reasoning tasks, because the softmax nonlinearity provides important selectivity that the linear kernel lacks.

  • FlashAttention (Dao, 2024): Not a complexity reduction but an IO-aware implementation that makes softmax attention dramatically faster in practice by minimizing GPU memory reads/writes. The problem: it makes long-context attention feasible on modern hardware but does not address the representational issue — attention is still diluted across the full sequence.

  • RingAttention (Liu et al., 2024c): Distributes the attention computation across multiple devices in a ring topology, scaling context length with device count. The problem: requires multiple GPUs and still computes full attention, which becomes increasingly diluted.

The fundamental limitation: These methods make long attention faster or more memory-efficient but do not make it smarter. The model still attends to everything, and the aggregated representation still becomes noisy as context grows.

3. Compression and KV-Cache Eviction Methods

These methods selectively discard or compress past information to keep the active context window manageable:

  • H2O (Zhang et al., 2023): Identifies "heavy hitter" tokens that consistently receive high attention scores and evicts the rest, keeping only a small fraction of the KV cache. The problem: this is a lossy compression — once tokens are evicted, they cannot be recalled. Information that was not identified as a heavy hitter at the time of encoding is permanently lost.

  • Dynamic Memory Compression (Nawrot et al., 2024): Compresses the KV cache by merging nearby key-value pairs, reducing the effective sequence length. The problem: merging blurs the distinction between tokens, potentially losing fine-grained information.

The fundamental limitation: These methods are irreversible — once information is compressed or evicted, it cannot be recovered. This is fundamentally different from how human memory works: we may not hold all past experience in active working memory, but we can retrieve it when cued. The paper explicitly positions EM-LLM against this limitation by storing all past information and retrieving it on demand.

4. Retrieval-Based Methods (Closest Competitors)

These methods, including the paper's primary baseline, selectively retrieve relevant past information into the active context window:

KV-Cache k-NN Retrieval:

  • Memorizing Transformers (Wu et al., 2022): Store all past key-value pairs in a database and use k-nearest neighbor search at each attention head to retrieve the most relevant ones. This was the first demonstration that k-NN retrieval over the KV cache could approximate full attention. The problem: per-token retrieval is expensive, and individual token retrieval fragments the context, losing the sequential coherence that Transformers rely on.

  • Unlimiformer (Bertsch et al., 2023): Similar k-NN retrieval approach but retrieves tokens rather than KV pairs, and encodes the full context with a separate encoder to generate query vectors. The problem: requires an additional encoding pass and retrieves individual tokens, not coherent segments.

  • InfLLM (Xiao et al., 2024a): The state-of-the-art baseline and the most direct predecessor to EM-LLM. InfLLM segments the entire context into fixed-size memory units (e.g., 512 tokens each), selects representative tokens per unit (those with the highest accumulated attention scores), and retrieves the most relevant units using k-NN search. This addressed a key weakness of per-token retrieval: by retrieving coherent blocks of tokens, the model gets sequential context within each block, which is essential for the Transformer's ability to understand local relationships.

InfLLM's critical limitation, which motivates EM-LLM: The fixed-size segmentation is arbitrary with respect to the content. A block boundary might fall in the middle of a sentence, a paragraph, a logical argument, or a multi-step reasoning chain. This has two consequences:

  1. Fragmented semantics: Information that belongs together is split across blocks, meaning the model might retrieve part of a relevant passage but miss the continuation in the next block. Or it might waste retrieval budget on a block that only partially contains relevant content.

  2. Inefficient retrieval: Because the blocks are content-agnostic, the retrieval signal (based on the representative tokens within a block) is noisier than it would be if blocks corresponded to coherent semantic units. A block that happens to contain the queried keyword in a tangential context might score higher than a block containing deeply relevant information expressed in different vocabulary.

The paper's core insight is that the segmentation itself should be adaptive to the content, and that a principled way to achieve this is to mimic how the human brain segments continuous experience into episodic events — using prediction error (surprise) as the boundary signal and optimizing for within-event coherence.

Retrieval-Augmented Generation (RAG):

  • Standard RAG (Lewis et al., 2020; Gao et al., 2024): Chunk the context into fixed-size segments, embed each chunk with a separate retriever model, index them in a vector database, and retrieve the top-k most similar chunks at query time. The problem: as the paper's results show (Table 9, Appendix A.2), even with a state-of-the-art retriever (NV-Embed-v2; Lee et al., 2024), RAG underperforms both full-context models and EM-LLM on most tasks. The paper identifies several reasons:
    • Single retrieval step: RAG retrieves once before generation, whereas EM-LLM retrieves at every layer, allowing different attention heads to focus on different parts of the context. The paper quantifies this in Appendix Figure 5: the ratio of blocks retrieved uniquely by a single layer (i.e., not retrieved by any other layer) is substantial, meaning layer-wise retrieval provides genuinely complementary information.
    • Retriever bottleneck: The retriever model (typically much smaller than the LLM) may not capture the nuances of relevance that the LLM itself would. In EM-LLM, the retrieval is done using the LLM's own query and key representations, which are naturally aligned with its processing.
    • Fixed chunking: Same problem as InfLLM — arbitrary boundaries fragment semantics.

How EM-LLM Positions Itself

EM-LLM does not propose an entirely new paradigm. Instead, it extends the group-based KV cache retrieval paradigm pioneered by InfLLM with three specific innovations drawn from cognitive science:

  1. Dynamic, surprise-based segmentation replaces fixed-size chunking. The paper draws directly on neuroscientific evidence: event boundaries in human perception correspond to moments of high prediction error — when the brain's internal model of the world is violated by incoming sensory input (Zacks et al., 2007; 2011; Sinclair et al., 2021). The paper adapts this to LLMs by using the model's own next-token prediction surprise (negative log-likelihood) as the boundary signal. This is computationally elegant because surprise is already computed during inference — the model produces token probabilities as part of generation — so segmentation incurs zero additional forward-pass cost.

  2. Graph-theoretic boundary refinement optimizes within-event coherence. Recognizing that surprise alone provides an imperfect segmentation, the paper introduces a second pass that treats the attention key similarity matrix as an adjacency matrix and adjusts boundaries to maximize modularity (or minimize conductance) — standard community detection metrics from network science. This is the first application of graph clustering to in-context segmentation in LLMs.

  3. Temporal contiguity buffer adds a human-like retrieval dynamic. Drawing on the free recall literature in cognitive psychology (Howard and Kahana, 2002) and the recent finding that Transformer attention heads exhibit the same contiguity and asymmetry effects as human memory (Ji-An et al., 2024), the paper introduces a second retrieval stage that prioritizes events adjacent to those retrieved by similarity. This is completely absent from InfLLM, which retrieves solely by similarity.

The paper explicitly frames itself as training-free — it works with any pretrained LLM without modification — which distinguishes it from methods that require long-context fine-tuning (LongLoRA, PoSE, etc.) and positions it as a drop-in solution for existing models. It also explicitly positions itself as a bridge between cognitive science and LLM engineering, arguing that the parallels between human event segmentation and LLM surprise (Section 4.2, Figure 4) suggest a deeper alignment worth exploring both for engineering applications and as a computational model of human memory.

The Paper's Central Hypothesis (and Its Risk)

The paper's approach rests on a hypothesis that is both elegant and potentially fragile: that the information structure relevant to an LLM's attention mechanism can be discovered through the same surprise-based segmentation that humans use for perceptual experience. If the LLM's next-token prediction errors cluster around genuine semantic or structural shifts in the content, then surprise-based segmentation will produce memory units that are coherent and easily retrievable. If not — if the LLM's surprise is driven by superficial features (rare words, formatting changes, tokenization artifacts) rather than meaningful event boundaries — then the segmentation will be noisy and the retrieval will be no better (or worse) than fixed-size chunking.

The paper addresses this risk in two ways: first, by empirically validating that LLM surprise correlates with human-perceived events (Section 4.2, using human-annotated podcast data from Kumar et al., 2023), and second, by introducing the boundary refinement step that can correct surprise-based boundaries when they fail to produce cohesive segments (Section 3.3). The refinement acts as a safety net: even if surprise identifies suboptimal boundaries, the modularity optimization can shift them to more natural breakpoints.

A Subtle but Critical Design Principle

Before diving into the technical details (Section 3), it is worth noting a design principle that permeates EM-LLM and distinguishes it from many retrieval-based approaches: EM-LLM never discards information. Unlike KV-cache eviction methods (H2O, Dynamic Memory Compression) that permanently delete tokens, EM-LLM stores all past tokens in an organized structure and retrieves subsets on demand. This means that even if the retrieval is imperfect for a particular query, the information is still available for future queries. The event segmentation and retrieval serve to prioritize what enters the active context window, not to limit what can ever be accessed. This aligns with the paper's stated goal of "practically infinite context lengths" — the memory store grows with the sequence, but the active processing window remains constant.

This design principle also explains why EM-LLM can succeed on passkey retrieval at 10M tokens (Figure 1, bottom): the passkey is stored as part of an event somewhere in the sequence, and when the query asks for it, similarity-based retrieval locates it regardless of its absolute position. A compression or eviction method would have almost certainly discarded it.

3. Technical Approach

3.1 Reader Orientation

EM-LLM is a training-free inference-time architecture that augments any pretrained Transformer-based LLM with an external episodic memory system, enabling the model to process sequences far longer than its original training context window without any fine-tuning. The system solves the long-context degradation problem — where LLMs lose coherence and retrieval accuracy as sequences grow beyond their training length — by mimicking how the human brain segments continuous experience into discrete events, stores them in organized memory, and retrieves them through a combination of similarity matching and temporal context.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components that operate in a processing loop as the LLM ingests a long sequence chunk by chunk:

  1. Base LLM — any pretrained Transformer model (e.g., LLaMA-3.1-8B, Mistral-7B, Phi-3.5). This is the "brain" that generates text and, critically, produces the next-token probabilities used for surprise computation. It is never modified or fine-tuned.

  2. Event Segmenter — processes each new chunk of incoming tokens and identifies where meaningful event boundaries occur. It does this in two stages: (a) surprise-based boundary detection using the LLM's own next-token prediction log-probabilities to find points where the model is "surprised" by the token it just saw, and (b) graph-theoretic boundary refinement that treats the attention key similarity matrix as an adjacency matrix and adjusts boundaries to maximize within-event cohesion and cross-event separation.

  3. Episodic Memory Store — a key-value (KV) cache database that stores all previously processed tokens, organized into variable-length segments (events) rather than fixed-size blocks. Each event stores: the KV pairs for all tokens within it, representative token indices (those with highest accumulated attention scores, used for retrieval), and its temporal position in the original sequence. The memory store grows unboundedly with sequence length; only CPU/disk memory limits its capacity.

  4. Two-Stage Memory Retriever — selects which stored events to load into the LLM's active context window when generating each new chunk of output. Stage one (similarity-based): uses k-nearest neighbors (k-NN) search to find events whose representative keys have highest dot-product similarity with the current query. Stage two (contiguity-based): enqueues events temporally adjacent to those retrieved in stage one into a decaying buffer that maintains temporal context. Both stages operate independently at every Transformer layer.

  5. Context Window Manager — assembles the final context for each attention computation by combining four groups of tokens: (a) initial tokens (first 128 tokens, serving as attention sinks), (b) similarity buffer (events retrieved via k-NN), (c) contiguity buffer (temporally adjacent events, managed as a queue), and (d) local context (the most recent tokens, fitting within the model's native context window, receiving full softmax attention).

Information flows as follows: a long input sequence is processed in chunks of m tokens (typically 512). For each chunk, the LLM computes attention normally within the local context window. As each chunk completes, the Event Segmenter partitions it into events using surprise signals and refinement. These events are stored in the Episodic Memory Store with their KV pairs and representative tokens. When processing the next chunk, the Two-Stage Retriever selects which past events to load, and the Context Window Manager combines them with the local context for the next attention computation. This loop continues until the entire sequence is processed. Crucially, the LLM never attends to the full sequence — only to the local context plus a curated subset of retrieved events — yet all past information remains accessible via retrieval.

3.3 Roadmap for the Deep Dive

The explanation follows the architecture's pipeline order because each component's output feeds into the next:

  • First, the formal definition of episodic memory in LLM terms — this establishes the interface between the LLM's key-value cache and the memory system, clarifying what exactly is being stored and retrieved.

  • Second, the surprise-based event segmentation — this is the first stage of memory formation and the paper's most cognitively-motivated innovation. We'll walk through the Bayesian surprise formulation, the adaptive thresholding mechanism, and why this particular signal was chosen over alternatives.

  • Third, the boundary refinement algorithm — this takes surprise-based boundaries and optimizes them for within-event cohesion using graph-theoretic metrics (modularity, conductance). We'll cover the adjacency matrix construction, the metric definitions, and the sequential optimization procedure.

  • Fourth, the memory storage and representation scheme — how events are stored in the KV cache, how representative tokens are selected, and how the system manages memory across arbitrarily long sequences.

  • Fifth, the two-stage retrieval mechanism — the similarity-based k-NN stage and the temporal contiguity buffer, including how they interact, how buffer sizes are set, and how retrieval operates independently per layer.

  • Sixth, the context window assembly — how the four token groups (initial, similarity, contiguity, local) are combined into the final attention context, and the design rationale for each.

  • Seventh, the complexity analysis — a walk-through of why EM-LLM's attention and retrieval costs scale sub-quadratically with sequence length, making 10M-token contexts feasible.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and methods paper whose core idea is that LLM context windows can be extended to practically infinite lengths by organizing past KV pairs into semantically coherent episodes — segmented using the model's own prediction errors and refined with graph clustering — and retrieving them through a two-stage process that mimics human memory access patterns, all without fine-tuning the underlying model.

Defining Episodic Memory in LLM Terms

Before describing how events are formed and retrieved, we need to establish what exactly the memory system stores. The paper defines episodic memory in the LLM context as "the organised, event-based collection of past key-value pairs, analogous to the latent representations of personal experiences in human memory" (Section 3.2).

In a standard Transformer, each attention head at each layer computes, for every token position, a key vector $K_i$ and a value vector $V_i$ from the token's hidden representation. During autoregressive generation, the model stores all past $(K_i, V_i)$ pairs in a KV cache to avoid recomputing them for each new token. The standard attention output at position $t$ with query vector $Q_t$ is:

ut=∑i=1tatiVi,whereati=exp⁡(Qt⋅Ki/d)∑j=1texp⁡(Qt⋅Kj/d)u_t = \sum_{i=1}^{t} a_{ti} V_i, \quad \text{where} \quad a_{ti} = \frac{\exp(Q_t \cdot K_i / \sqrt{d})}{\sum_{j=1}^{t} \exp(Q_t \cdot K_j / \sqrt{d})}

where $d$ is the key/query dimensionality, $a_{ti}$ is the attention weight from query $t$ to key $i$, and $u_t$ is the output vector (the weighted sum of value vectors).

What this computes: a normalized weighted average of all past value vectors, where the weight for each past position is proportional to the exponentiated dot-product similarity between the current query and that position's key. This is the standard softmax attention mechanism.

The memory system's role: rather than summing over all $t$ past positions (which becomes both computationally expensive and representationally noisy as $t$ grows), EM-LLM approximates this by summing only over positions in (a) the local context window and (b) retrieved episodic events. The memory system determines which past positions to include. Formally, the paper invokes an approximate equivalence argument (Appendix F.1): if the softmax is sharply peaked (i.e., most attention weight concentrates on a few highly similar keys), then restricting the sum to the top-k most similar keys yields an output vector close to the full attention output. This justifies replacing full attention with retrieval-based attention.

What is stored per event: for each identified event (a contiguous span of tokens with start and end positions), the system stores:

  • The KV pairs $(K_i, V_i)$ for all tokens $i$ in the event.
  • The representative tokens — a small subset of tokens within the event whose keys receive the highest accumulated attention scores. These are used for retrieval queries (computing similarity between a new query and the event as a whole), following the approach of InfLLM (Xiao et al., 2024a).
  • The event's temporal position (its ordinal index in the sequence), used by the contiguity buffer.

Why this representation: storing KV pairs by event rather than by fixed-size block or individual token enables the system to retrieve semantically coherent information. When a query matches an event's representative tokens, the entire event (with its sequential context) is loaded into the attention window. This preserves local syntactic and semantic relationships within the event, which are essential for the Transformer to process the retrieved information meaningfully. Storing individual tokens would fragment this context; storing fixed-size blocks would create arbitrary divisions unrelated to content.

Surprise-Based Event Segmentation

This is the first stage of memory formation and the component most directly inspired by human cognition. The core idea, drawn from the event cognition literature (Zacks et al., 2007; 2011; Fountas et al., 2022), is that event boundaries occur at moments when the brain's (or model's) predictions are significantly violated by incoming information. In an autoregressive LLM, the model's prediction for the next token $x_t$ given the previous tokens $x_1, \ldots, x_{t-1}$ is $P(x_t \mid x_1, \ldots, x_{t-1}; \theta)$, where $\theta$ represents the model parameters. The surprise (or surprisal) of token $x_t$ is:

S(xt)=−log⁡P(xt∣x1,…,xt−1;θ)S(x_t) = -\log P(x_t \mid x_1, \ldots, x_{t-1}; \theta)

where $P(x_t \mid \ldots)$ is the probability the model assigned to the actual token $x_t$ before observing it, and $\theta$ represents all model parameters (fixed, since no fine-tuning occurs).

What it computes: the negative log-probability of the ground-truth token given all preceding tokens. This is a standard information-theoretic quantity: it measures how many bits of information the token $x_t$ carries that the model did not anticipate. A token the model assigns probability 0.5 has surprise $-\log(0.5) \approx 0.69$ nats; a token assigned probability $10^{-4}$ has surprise $-\log(10^{-4}) \approx 9.21$ nats. High values indicate the model was "surprised" — the token was highly unlikely according to the model's internal predictive distribution.

Why this form: surprise (negative log-likelihood) is the natural Bayesian measure of prediction error. It has several properties that make it suitable for event boundary detection: (a) it is computed for free during autoregressive inference (the model already produces token probabilities as part of decoding), (b) it is a local, per-token measure that doesn't require looking ahead, enabling online processing, and (c) it aligns with the neuroscientific theory that event boundaries correspond to prediction errors — the brain's generative model of the world being violated by sensory input. Alternative measures like absolute embedding change or attention pattern shift would require additional computation and lack the direct cognitive interpretation.

Boundary detection via adaptive thresholding:

Surprise values alone don't tell us where boundaries are — we need a criterion for deciding when surprise is "high enough" to mark an event boundary. The paper uses an adaptive threshold computed from recent surprise history:

T=μt−τ:t+γ⋅σt−τ:tT = \mu_{t-\tau:t} + \gamma \cdot \sigma_{t-\tau:t}

where $\mu_{t-\tau:t}$ is the mean surprise over the most recent $\tau$ tokens (a moving window), $\sigma_{t-\tau:t}$ is the standard deviation of surprise over the same window, and $\gamma$ is a scaling factor that controls the threshold's sensitivity.

What it computes: a token $x_t$ is classified as a potential event boundary if its surprise $S(x_t)$ exceeds $T$ — that is, if it is more than $\gamma$ standard deviations above the local mean surprise. The window $\tau$ (not explicitly specified in the main text but adjustable) ensures the threshold adapts to contextual shifts: if the model enters a region where everything is moderately surprising (e.g., a particularly technical passage), the threshold rises so that only unusually surprising tokens within that context trigger boundaries.

Why this form: using a moving mean and standard deviation makes the threshold context-adaptive — it automatically calibrates to the local difficulty or predictability of the text. This avoids the need for a manually tuned absolute surprise threshold, which would be fragile across different domains, models, and text types. The parameter $\gamma$ provides a single, interpretable knob: $\gamma = 1$ means "mark boundaries at tokens that are more than one standard deviation more surprising than average for this local context," $\gamma = 2$ is more conservative, and so on. The paper's hyperparameter sweep (Appendix D.1, Figure 12) found that $\gamma = 1$ performs best across most models and tasks, with some models preferring $\gamma = 2$ (specifically LLaMA-3-8B-Instruct). This single-parameter sensitivity is a practical advantage: it means the method works out-of-the-box without per-task or per-domain tuning.

The output of this stage: an ordered list of token indices $B = \{b_1, b_2, \ldots, b_k\}$ that mark potential event boundaries. The sequence from $b_i$ to $b_{i+1} - 1$ (inclusive) forms a candidate episodic event. These boundaries serve as the initialization for the refinement step — they are not the final segmentation.

A crucial implementation detail: surprise-based segmentation is applied per chunk as the sequence is processed. When the system ingests a new chunk of $m$ tokens, it computes surprise for each token in that chunk, applies the adaptive threshold, and identifies initial boundaries within the chunk only. Events from previous chunks have already been formed and stored; the new chunk may start a new event or continue the last event from the previous chunk, depending on whether the chunk boundary coincides with a surprise-based boundary. This chunk-by-chunk, online processing enables handling of streaming or arbitrarily long sequences without requiring the full sequence to be available upfront.

Graph-Theoretic Boundary Refinement

The initial surprise-based boundaries, while cognitively motivated, are not guaranteed to produce segments that are optimal for the downstream retrieval task. The key insight motivating the refinement step is that memory retrieval is most efficient when within-event token similarity (as measured by attention key vectors) is high and cross-event similarity is low — that is, events should be coherent clusters in the key embedding space. If two tokens that are semantically related end up in different events (because the surprise signal failed to detect that they belong together), then a query that matches one may fail to retrieve the other, fragmenting the retrieved context.

The paper formalizes this as a graph clustering problem and introduces a refinement step that adjusts the initial boundaries to optimize for within-event cohesion.

Constructing the adjacency matrix:

For a given attention head $h$ at a given layer, and for a sequence of tokens $x_1, x_2, \ldots, x_n$ within the current processing chunk, the system constructs a similarity matrix $A^h$:

Aijh=sim(Kih,Kjh)A^h_{ij} = \text{sim}(K^h_i, K^h_j)

where $K^h_i$ and $K^h_j$ are the key vectors for tokens $x_i$ and $x_j$ at head $h$, and the similarity function is the dot product: $\text{sim}(K^h_i, K^h_j) = K^{h^T}_i \cdot K^h_j$.

What it computes: a symmetric $n \times n$ matrix where entry $(i,j)$ represents how similar token $i$'s key is to token $j$'s key. High values indicate that, from this attention head's perspective, the two tokens should attend to each other strongly.

Why dot product similarity: this aligns with the self-attention mechanism itself — attention weights are computed as softmax over dot products of queries and keys. Using the same similarity measure ensures the graph structure directly reflects the model's own notion of token relatedness, rather than an external metric that might not correspond to what the attention mechanism cares about. This is a crucial design choice: using, say, cosine similarity or an external embedding model's similarity would produce a graph that doesn't necessarily reflect how the LLM's attention distributes across tokens.

Which head's keys to use: the paper states that refinement uses keys from "an attention head $h$ within the local context window" (Section 3.3). The implementation appears to aggregate across heads or use a representative head, as the refinement produces a single set of event boundaries for the entire KV cache (not per-head events). The exact head selection or aggregation procedure is not fully specified in the main text but the algorithm's structure is agnostic to this choice — one could use mean-pooled keys across heads, or a specific head, with potential performance differences.

The metric functions:

To evaluate the quality of a candidate set of event boundaries $B$, the paper defines two graph-theoretic metrics borrowed from community detection in network science.

Modularity (Newman and Girvan, 2004):

fM(Ah,B)=14m∑i,j[Aijh−12m(∑iAijh)(∑jAijh)]δ(ci,cj)f_M(A^h, B) = \frac{1}{4m} \sum_{i,j} \left[ A^h_{ij} - \frac{1}{2m} \left( \sum_i A^h_{ij} \right) \left( \sum_j A^h_{ij} \right) \right] \delta(c_i, c_j)

where $m$ is the total sum of all edge weights in the graph (i.e., $m = \frac{1}{2}\sum_{i,j} A^h_{ij}$), $c_i$ is the community (event) to which token $i$ is assigned according to boundaries $B$, and $\delta(c_i, c_j)$ is the Kronecker delta function — equal to 1 if tokens $i$ and $j$ are in the same event, 0 otherwise.

What it computes: modularity measures the density of edges within communities compared to what would be expected if edges were placed randomly (preserving the degree distribution of each node). The first term $A^h_{ij}$ is the actual edge weight between $i$ and $j$. The second term $\frac{1}{2m}(\sum_i A^h_{ij})(\sum_j A^h_{ij})$ is the expected edge weight under a null model where connections are proportional to node degrees. When the actual weight exceeds the null expectation and the two tokens are in the same community ($\delta(c_i, c_j) = 1$), modularity increases. When they are in different communities ($\delta = 0$), the term contributes zero regardless of edge weight.

Why this form: modularity directly operationalizes the "high within-event similarity, low cross-event similarity" desideratum. By subtracting the null-model expectation, it penalizes simply putting everything in one big event — the null model would also predict high within-community density for large communities. Maximizing modularity finds the partition that most deviates from random connectivity patterns. This is the standard objective for community detection in networks and has well-understood properties (though also known limitations, such as resolution limit — it struggles to detect very small communities in large graphs).

Conductance (alternative metric):

fC(Ah,B)=min⁡S∈V∑i∈S,j∉SAijhmin⁡(vo(S),vo(V∖S))f_C(A^h, B) = \min_{S \in V} \frac{\sum_{i \in S, j \notin S} A^h_{ij}}{\min(v_o(S), v_o(V \setminus S))}

with

vo(S)=∑i,j∈SAij,vo(V∖S)=∑i,j∉SAijv_o(S) = \sum_{i,j \in S} A_{ij}, \quad v_o(V \setminus S) = \sum_{i,j \notin S} A_{ij}

where $S = \{b_i, b_i+1, \ldots, b_{i+1}\}$ is the subset of nodes (tokens) belonging to a single event, and $V$ is the set of all nodes in the graph.

What it computes: for each event $S$, conductance measures the total weight of edges leaving the event (connecting tokens in $S$ to tokens outside $S$) divided by the minimum of the total internal edge weight of $S$ and the total internal edge weight of its complement $V \setminus S$. Lower conductance means the event has strong internal connections (large denominator) and weak external connections (small numerator) — it is a well-separated cluster.

Why this form: conductance penalizes partitions where events "leak" — where high-similarity token pairs are split across boundaries. The $\min$ in the denominator prevents the metric from being trivially minimized by making one event tiny (which would have small internal weight, making the denominator small and thus the ratio large — i.e., high conductance, which is penalized in minimization). The paper reports using both modularity (maximization) and conductance (minimization) for refinement, with modularity giving the best results (Table 1, SM vs. S+C variants).

The refinement algorithm:

The refinement procedure (Algorithm 1 in the paper) takes the initial surprise-based boundaries $B$ and sequentially adjusts each one to optimize the chosen metric:

  1. Start with $B = \{b_1, b_2, \ldots, b_k\}$, the initial boundaries from surprise thresholding.
  2. For each consecutive pair of boundaries $(\alpha, \beta) = (b_i, b_{i+1})$, search over all possible boundary positions $\hat{\beta}$ in the range $(\alpha, \beta]$ (i.e., between the current left boundary and up to and including the current right boundary).
  3. For each candidate $\hat{\beta}$, evaluate the metric function $f(A^h, \{\alpha, \hat{\beta}, b_{i+2}, \ldots\})$ — i.e., the metric computed with the new proposed boundary replacing $\beta$ while keeping all other boundaries fixed.
  4. Set $b_{i+1}$ to the $\hat{\beta}$ that maximizes modularity or minimizes conductance (depending on which metric is used).
  5. Continue through all boundaries sequentially.

What this computes: a greedy, sequential optimization that tries to improve the quality of each boundary independently. For each boundary, it tests every possible alternative position between the previous and current boundary and picks the one that gives the best community structure score, holding all other boundaries fixed.

Why this sequential, greedy approach: jointly optimizing all boundaries would require evaluating an exponential number of possible partitions ($k$ boundaries with $n$ possible positions each). The sequential greedy approach has complexity $O(k \cdot (n/k)^2) = O(n^2/k)$ per chunk (see Appendix C.1 for the full derivation including the adjacency matrix construction cost), which is dominated by $O(nm)$ in practice because the chunk size $m$ bounds the search. The algorithm is guaranteed to never worsen the metric compared to the initial surprise-based boundaries because it only moves a boundary if the metric improves (or leaves it unchanged if no improvement exists).

A theoretical connection worth noting: the paper observes that this algorithm can be seen as "a single pass of Phase 1 of the heuristic Louvain method initialized with surprise-based segmentation" (Appendix E.3). The Louvain method (Blondel et al., 2008) is a widely-used community detection algorithm that iteratively moves nodes between communities to maximize modularity. EM-LLM's refinement is a simplified, single-pass version that only considers moving boundaries (not reassigning individual nodes arbitrarily) and only in one direction (merging/splitting at the boundary). This simplification is computationally necessary for online processing but means the refinement converges to a local optimum, not necessarily the global optimum.

Performance impact: the paper demonstrates in Table 2 and Figure 4 that refinement significantly improves within-event similarity metrics over surprise-only segmentation. Across LLaMA-2, LLaMA-3, and Mistral, SM (surprise + modularity refinement) achieves 18.7–39.9% higher modularity, 24.6–30.6% lower conductance, and 25.0–35.3% higher intra/inter-similarity ratio compared to random segmentation — consistently outperforming surprise-only (S) and fixed-segmentation with refinement (FM, FC). Notably, fixed-segmentation with refinement (FM, FC) underperforms surprise-based counterparts, corroborating the paper's claim that surprise provides a crucial initialization — refinement amplifies a good starting point but cannot rescue a poor one.

Memory Storage and Representative Token Selection

Once events are segmented (after surprise detection and boundary refinement), they need to be stored in a way that enables efficient retrieval. The storage scheme follows InfLLM's approach but adapts it to variable-length events.

Event storage in the KV cache:

For each event, the system stores the full KV pairs for all tokens within that event. In a long-context scenario, most of these KV pairs reside in CPU memory or on disk (see Appendix C.3.2 for the memory management strategy), with only actively retrieved events loaded onto GPU memory. The storage is organized as a list of events, each with:

  • Start and end token indices.
  • The KV cache tensor for all tokens in the event.
  • A set of representative tokens and their corresponding key vectors.

Representative token selection:

For an event containing $L$ tokens, the system identifies a small subset of tokens that best represent the event for retrieval purposes. Following InfLLM's method, the representative tokens are those that receive the highest accumulated attention scores during processing. Specifically, while processing the chunk, the system tracks, for each token, the sum of attention weights it receives from all subsequent tokens within the local context window. The tokens with the largest accumulated attention are selected as representatives for their event.

What this computes: for each token $i$ in the event, compute $\sum_{j > i} a_{ji}$ where $a_{ji}$ is the attention weight from token $j$ to token $i$ during chunk processing. The top-ranked tokens by this sum become the event's representatives.

Why this selection criterion: tokens that receive high cumulative attention are those that subsequent tokens consistently "look back at" — they act as information hubs or summary points for the surrounding context. When a new query needs to find relevant past information, matching against these hub tokens is a good proxy for matching against the event as a whole. Using all tokens in the event for retrieval would be computationally expensive and would dilute the retrieval signal with tokens that carry little discriminative information (e.g., function words, punctuation). Using a single centroid vector (e.g., mean-pooled keys) would lose the ability to match different aspects of the event.

The number of representative tokens per event is determined by the event's length and a configurable parameter. In practice, with chunk size $m = 512$ and typical event sizes of a few dozen to a few hundred tokens, a small constant number (e.g., 4–8) of representatives per event is sufficient. The exact number is not specified in the main text but is inherited from InfLLM's configuration.

Memory management across very long sequences:

For sequences extending to millions of tokens, the total KV cache becomes too large to keep in GPU memory. The paper implements a tiered storage strategy (Appendix C.3.2):

  • GPU memory: holds the KV cache for the current local context and actively retrieved events.
  • CPU memory: holds the full KV cache for all past events, organized in a least-recently-used (LRU) eviction scheme based on allocation slots (pre-allocated memory regions that are overwritten rather than freed, avoiding memory fragmentation).
  • Disk: when CPU memory is exhausted, LRU events are offloaded to disk storage.

The representative token vectors, used for k-NN search, are kept on GPU memory for sequences up to approximately 1M tokens; beyond that, they are offloaded to CPU memory (with a corresponding increase in retrieval latency). The paper reports that with 2GB of CPU memory per inference instance, the system can handle arbitrarily long sequences as long as sufficient disk space is available.

Two-Stage Memory Retrieval

When generating output for a new token or processing a new chunk, the system must select which past events to load into the active context window. The retrieval process has two stages, operating independently at each Transformer layer.

Stage 1: Similarity-based k-NN retrieval

For the current query token(s) at layer $l$, the system computes the dot-product similarity between the query vector $Q^l$ and the representative key vectors of all stored events. The $k_s$ events with the highest similarity scores are selected — these form the similarity buffer.

S(e,Ql)=max⁡r∈R(e)(Ql⋅Krl)S(e, Q^l) = \max_{r \in R(e)} (Q^l \cdot K^l_r)

where $R(e)$ is the set of representative tokens for event $e$, $K^l_r$ is the key vector for representative token $r$ at layer $l$, and the $\max$ over representatives means an event's similarity to the query is determined by its single most-similar representative.

What it computes: for each stored event, find which of its representative tokens has the highest dot product with the current query, and use that maximum as the event's relevance score. Sort events by this score and take the top $k_s$.

Why max-pooling over representatives: using the maximum similarity (rather than, say, the mean) ensures that an event is retrieved if any of its representative tokens strongly matches the query. This is appropriate because the representatives span different aspects of the event's content — a query about one aspect should retrieve the event even if other representatives are irrelevant. Mean-pooling would dilute the signal, making events harder to distinguish.

Why dot-product similarity: this is identical to the pre-softmax attention score in the Transformer. Using the same similarity measure for retrieval as for attention ensures that the retrieved events would naturally receive high attention weights if they were in the full context — the retrieval is selecting events that the LLM would attend to if it could attend to everything.

Implementation of k-NN search: for large memory stores (millions of events), exact nearest-neighbor search over all representative vectors would be prohibitively expensive. The paper uses the FAISS library (Johnson et al., 2019; Douze et al., 2024) for approximate k-NN search, which provides sub-linear query time through indexing structures (inverted file indices, product quantization, or hierarchical navigable small world graphs). The paper does not specify the exact FAISS index configuration but notes that the approximate nature of the search is acceptable because the retrieval is used to approximate softmax attention, which is itself a soft selection mechanism — exact nearest neighbors are not required.

Per-layer retrieval: a critical architectural feature is that retrieval is performed independently at every Transformer layer. Each layer has its own query, key, and value projections, and thus its own notion of token similarity. Layer 5 might retrieve events relevant to syntactic structure; layer 20 might retrieve events relevant to semantic content; layer 30 might retrieve events relevant to factual recall. The paper demonstrates this visually in Appendix Figure 5: the ratio of blocks retrieved uniquely by a single layer (not retrieved by any other layer) is substantial, showing that layer-wise retrieval provides complementary information that a single retrieval step (as in RAG) would miss.

Stage 2: Temporal contiguity buffer

The similarity-based retrieval identifies events that are topically relevant to the current query, but it ignores temporal relationships. In human episodic memory, retrieving one memory often triggers recall of temporally adjacent memories — this is the temporal contiguity effect (Howard and Kahana, 2002; see Figure 3A). The paper's second retrieval stage is designed to replicate this effect in the LLM's context.

The contiguity buffer is a queue of size $k_c$. When an event $e$ is retrieved via similarity (in stage 1), the system also enqueues its neighboring events — specifically, events at positions $e-1$ and $e+1$ in the original temporal sequence (the paper uses $n=1$, meaning one event on each side). These temporally adjacent events are added to the contiguity buffer.

The queue mechanism: the buffer operates as a first-in-first-out (FIFO) queue with eviction of oldest entries when capacity is exceeded. When a neighboring event is already in the buffer (because it was previously enqueued from another retrieval), it is not duplicated. The queue's capacity $k_c$ is typically set to a fraction of the total retrieval budget; the paper's hyperparameter sweep (Appendix D.1, Figure 13) found $k_r = k_c / (k_s + k_c) = 0.3$ to perform best overall, meaning 30% of the retrieval budget goes to the contiguity buffer and 70% to the similarity buffer.

What this computes: the contiguity buffer ensures that the LLM's context window contains not just the topically relevant events but also the temporal context surrounding them — what happened just before and just after each relevant event. This is crucial for tasks requiring temporal reasoning (understanding event order, causality, narrative progression) and for enabling the "induction" attention heads that Ji-An et al. (2024) showed exhibit the same contiguity and asymmetry effects as human memory.

Why decaying via a queue: as new events are retrieved and their neighbors enqueued, older contiguity entries are evicted. This creates a recency-weighted temporal context — events that were recently retrieved (and are thus likely still relevant to the ongoing context) have their temporal neighbors available, while neighbors of events retrieved long ago gracefully fall out of the context window. This mirrors the temporal decay observed in human free recall experiments.

Interaction between the two stages: the two buffers together occupy the full retrieval budget $k = k_s + k_c$ events. The retrieved events (from the similarity buffer) and the contiguity events (from the queue) are concatenated and provided to the LLM's attention mechanism as a single set of KV pairs, albeit at fixed position embeddings (see Context Window Assembly below). The paper notes that certain tasks benefit more from contiguity than others — retrieval-intensive tasks (PassageRetrieval, Retrieve.KV) show large gains from the contiguity buffer, while tasks where temporal relationships are less important may perform better without it (allowing more similarity-based events). This task-dependent optimality is discussed in Section 4.4.

Context Window Assembly

The final step in each processing cycle is to assemble the tokens that will participate in the attention computation for the next chunk. The context window is divided into four groups, each with a specific role:

1. Initial tokens (128 tokens): the first 128 tokens of the entire sequence are always kept in the context window, regardless of their relevance to the current query. This practice, introduced by Xiao et al. (2024b) and Han et al. (2024b), addresses the observation that initial tokens often serve as attention sinks — they accumulate disproportionately high attention weights during processing, and removing them degrades performance on window attention. By always including these tokens, the model maintains a stable reference point that helps recover the performance lost when using sparse attention patterns.

Why 128: this number is inherited from prior work (Xiao et al., 2024a; 2024b) and was not independently optimized in this paper. It represents a small, fixed overhead on the context window.

2. Similarity buffer (events from k-NN retrieval): the $k_s$ events retrieved by similarity-based k-NN search across all layers. Since retrieval is per-layer, the similarity buffer is assembled by taking the union of events retrieved by any layer — an event retrieved by any layer's attention head is included in the context window for all layers. This is a design choice: while retrieval is per-layer, the event loading is shared, meaning each layer sees the same set of retrieved events but can attend to different parts of them based on its own attention patterns.

3. Contiguity buffer (events from temporal queue): the $k_c$ events currently in the contiguity buffer, populated by neighbors of similarity-retrieved events.

4. Local context (recent tokens within native window): the most recent tokens, fitting within the LLM's original training context window size. For a model like LLaMA-3.1-8B with a 128K training context, the local context might be set to 4K tokens to balance recency with retrieval capacity. The local context receives full softmax attention across all its tokens — no retrieval or sparsification is applied here. This ensures the model has perfect access to immediate context for tasks like short-term syntactic processing, coreference resolution, and maintaining conversation flow.

Position embeddings for retrieved tokens:

A technical challenge arises because retrieved events are discontinuous — they come from arbitrary positions in the past, not necessarily sequential with the local context. If they were assigned their original positional indices, the model would see a non-monotonic sequence of positions (e.g., tokens at positions 1, 2, 3, then suddenly 50,000, 50,001, 50,002, then back to 4,997, 4,998...), which would confuse the positional encoding, especially for relative encodings like RoPE.

The paper follows the solution from Raffel et al. (2020) and InfLLM: retrieved tokens are assigned a fixed, shared position embedding regardless of their original position. This means all retrieved events are treated as if they occupy a single, constant position relative to each other, with the only positional information being their relative order within each event (which is preserved by using the original positional differences within the event but shifted to a common base). This approach has become standard in retrieval-based long-context methods and, while it discards absolute positional information for retrieved tokens, it avoids the representation damage that would result from exposing the model to wildly out-of-distribution position indices.

The total context length is $L_{\text{total}} = 128 + \sum_{e \in \text{retrieved}} |e| + L_{\text{local}}$, where $|e|$ is the length (in tokens) of event $e$ and $L_{\text{local}}$ is the local context size. The paper configures this based on the base model's capacity; for Mistral-7B, the configuration is typically 4K local + 2K retrieved tokens (written as "4K+2K" in Table 1); for LLaMA-3 and 3.1, 4K+4K; for Phi-3, 1K+3K. These choices follow InfLLM's settings to enable direct comparison.

Why four distinct groups: this structure separates different types of information by their retrieval mechanism and temporal properties. Initial tokens provide stability. Local context provides high-fidelity recent information. The similarity buffer provides topically relevant but temporally distant information. The contiguity buffer provides temporal context around relevant information. A simpler approach — e.g., retrieving only by similarity — would miss the temporal relationships that the contiguity buffer captures and the stability that the initial tokens provide.

Complexity Analysis: Why This Scales to 10M Tokens

A central claim of the paper is that EM-LLM handles "practically infinite context lengths while maintaining computational efficiency" (Abstract). The complexity analysis (Appendix C.2) explains why this is the case.

Standard full attention complexity: for a sequence of length $n$, computing the full attention matrix costs $O(n^2)$ in both time and memory because every query attends to every key. For $n = 10^7$, this would require $10^{14}$ operations per layer per token — completely infeasible.

EM-LLM's attention complexity: the sequence is processed in chunks of size $m$ (typically 512). For each chunk, attention is computed over:

  • The $n_l$ tokens in the local context.
  • The $n_r$ tokens in the retrieved events (from similarity + contiguity buffers).

The per-chunk attention cost is $O(m \cdot (n_l + n_r)^2)$ if computing full attention over the assembled context, or $O(m \cdot (n_l + n_r))$ if using memory-efficient attention implementations. Since $n_l$ and $n_r$ are constants (not growing with sequence length $n$), the attention cost per chunk is constant in $n$. Processing the full sequence requires $n/m$ chunks, yielding:

  • Attention complexity: $O(n \cdot (n_l + n_r))$ — linear in $n$, not quadratic.

Note: the paper's Appendix C.2 writes this as $O(n(n_l + n_r))$ after accounting for the $n/m$ chunks, each with $O(m(n_l + n_r))$ attention. This is correct: the $m$ cancels, leaving linear scaling in sequence length.

k-NN retrieval complexity: for each chunk, the system must retrieve the top-k events from the memory store. With $N_{\text{events}} \approx n/b$ events (where $b$ is the minimum event size), the retrieval cost per chunk using FAISS approximate k-NN is sub-linear in $N_{\text{events}}$ (typically $O(\log N_{\text{events}})$ for tree-based indices or $O(1)$ amortized for graph-based indices). Processing all chunks gives retrieval complexity $O((n/m) \cdot \log(n/b))$ — polylogarithmic in $n$.

Overall complexity: $O(n \cdot (n_l + n_r + \frac{\log n}{m}))$. The dominant term is linear in $n$, with a constant factor determined by the local + retrieved context size. For $n_l + n_r \approx 6000$ (4K local + 2K retrieved), this is approximately 6,000 operations per token, compared to $n$ operations per token for full attention. At $n = 10^7$, EM-LLM's per-token cost is $10^7 / 6000 \approx 1667\times$ lower than full attention.

Boundary refinement complexity: as derived in Appendix C.1, the refinement step costs $O(nm)$ for the full sequence, dominated by $O(m^2)$ per chunk (to compute the adjacency matrix) times $n/m$ chunks. With $m = 512$, this is $O(512 \cdot n)$ — linear in $n$ with a small constant, making it negligible compared to the attention cost for long sequences.

Visual confirmation: Figures 9 and 10 in Appendix C.2 plot the scaling of EM-LLM's attention complexity vs. full-context attention for sequences up to 100K and 10M tokens, respectively. On a log scale, EM-LLM's cost is nearly flat compared to the quadratic explosion of full attention. The empirical passkey retrieval experiment (Figure 1, bottom; Section 4.1) confirms this scaling holds in practice: EM-LLM achieves 100% accuracy on passkey retrieval at 10.2M tokens, a length at which full-context models either cannot run (due to memory constraints) or fail (the paper's Figure 1 shows baseline methods degrading well before this scale).

Putting It All Together: The Full Inference Loop

To make the processing pipeline fully concrete, here is the step-by-step flow for processing a long sequence:

  1. Initialization: load the base LLM. Set the local context size $n_l$, retrieval budget $k$, contiguity ratio $k_r$, surprise threshold parameter $\gamma$, and chunk size $m$.

  2. Process first chunk ($m$ tokens): the LLM computes full softmax attention over these tokens (no retrieval yet, since there is no past to retrieve from). The initial 128 tokens are identified and stored. For each token in the chunk, compute surprise $S(x_t) = -\log P(x_t \mid x_{<t})$.

  3. Segment first chunk: apply the adaptive threshold (Equation 1) to identify initial event boundaries within the chunk. For each identified event, construct the key similarity adjacency matrix and run boundary refinement (Algorithm 1) to produce the final event boundaries. Store each event in the episodic memory with its KV pairs and representative tokens.

  4. Process next chunk: before computing attention, perform retrieval:

    • For each layer, compute query-key similarities with all stored event representatives, retrieve top events via approximate k-NN.
    • Union the retrieved events across layers to form the similarity buffer.
    • For each retrieved event, enqueue its temporal neighbors ($\pm 1$) into the contiguity buffer.
    • Assemble the context window: initial tokens + similarity buffer events + contiguity buffer events + local context (last $n_l$ tokens).
    • Assign fixed position embeddings to retrieved events.
    • Compute attention over the assembled context.

  5. Segment current chunk: same as step 3 — compute surprise, detect initial boundaries, refine boundaries, store events. The first event of the chunk may be merged with the last event of the previous chunk if no surprise boundary exists between them.

  6. Repeat steps 4–5 until the entire sequence is processed.

  7. For generation tasks: when generating new tokens (rather than encoding), the local context includes the previously generated tokens, and retrieval is performed at each generation step using the current query vector.

This loop achieves what the paper claims: all past information is stored and retrievable, but only a constant-sized window is attended to at any step, making the system's computational cost linear in sequence length while its accessible memory grows without bound.

4. Key Insights and Innovations

Innovation 1: Difficulty-Conditioned Compute-Optimal Test-Time Scaling

What's distinctive at the idea level: The paper's most fundamental contribution is not any single method but rather the meta-strategy of adaptively allocating test-time compute based on prompt difficulty. Prior work treated test-time compute as a uniform knob: turn it up (more samples, more search) and performance improves. This paper demonstrates that the relationship between compute and performance is qualitatively different depending on problem difficulty, and that ignoring this heterogeneity leaves enormous efficiency on the table.

What makes this genuinely novel — rather than an obvious observation — is that the difficulty-dependent behavior is often counterintuitive. The strongest optimizer (beam search) actually hurts performance on easy problems at high budgets due to verifier over-optimization (Figure 3, right), while it helps substantially on medium-difficulty problems. Similarly, sequential revisions dominate on easy problems but a balanced sequential-parallel ratio is optimal on hard ones (Figure 7, right). These are not monotonic relationships where "more powerful = better." The compute-optimal policy exploits these non-monotonicities to achieve 4× better efficiency than best-of-N (Figures 4 and 8), which is a significant practical gain.

Comparison to prior work: Before this paper, the dominant paradigm was uniform allocation: apply the same test-time strategy (typically best-of-N) to every problem regardless of its characteristics. The scaling laws community (Hoffmann et al., 2022) had established principled frameworks for allocating pretraining compute between model size and data quantity, but no analogous framework existed for inference-time compute. This paper's compute-optimal scaling is best understood as an inference-time analog of the Chinchilla scaling laws — the conceptual parallel is direct (optimize a resource allocation under a budget constraint), but the underlying mechanism is entirely different (discrete strategy selection conditioned on difficulty vs. continuous optimization over parameters and tokens).

Significance beyond raw performance: The innovation is as much a diagnostic concept as an engineering contribution. By demonstrating that different methods have complementary, difficulty-dependent strengths, the paper provides a unified explanation for contradictory prior findings: Huang et al. (2023) found "LLMs cannot self-correct reasoning" because their problem distribution skewed hard; Madaan et al. (2023) found self-refinement helped because their distribution skewed easier. The compute-optimal framework reveals that both findings were right, just on different subsets of the difficulty spectrum. This converts a confusing set of contradictory results into a coherent picture with clear boundary conditions, enabling future researchers to design experiments with appropriate difficulty controls.

Fundamental vs. incremental: This is a fundamental conceptual shift — not because the individual pieces (best-of-N, beam search, revisions) are new, but because the meta-strategy of difficulty-conditioned allocation transforms how we think about test-time compute. It reframes inference from a "one-size-fits-all" paradigm to an adaptive resource allocation problem where the optimal strategy is prompt-dependent. The paper's empirical evidence — that the compute-optimal policy recovers 4× efficiency gains using only predicted (non-oracle) difficulty bins (Figure 4, Figure 8) — demonstrates that this reframing is not merely theoretical but practically deployable.

Tie to evidence: The core evidence for this innovation is the difficulty-bin breakdown in Figures 3 (right) and 7 (right), which show qualitatively different — and sometimes opposite — scaling behavior across difficulty quintiles for the same method, and the compute-optimal curves in Figures 4 and 8, which show the efficiency gains from the adaptive policy. The fact that predicted difficulty bins nearly match oracle bins (curves overlapping in Figure 4, slight gap at high budgets in Figure 8) confirms deployability without ground-truth labels.

Innovation 2: The Proposal Distribution and Verifier as Complementary, Independent Scaling Axes

What's distinctive at the idea level: The paper's unifying framework in Section 2 — decomposing all test-time compute methods into modifications to the proposal distribution (what the model generates) versus the verifier (how outputs are selected) — is not itself technically novel. It echoes the proposer-scorer decomposition familiar from MCMC and reinforcement learning. What is novel is the paper's empirical demonstration that these two axes have complementary, difficulty-dependent strengths and that combining them yields gains neither achieves alone.

The key insight is that revisions (proposal modification) and search (verifier optimization) are not just two ways to spend the same compute — they solve different parts of the problem. Revisions excel at local refinement: taking a roughly-correct answer and fixing specific errors. Search excels at global exploration: finding a correct solution strategy among many candidates when the model's initial attempts are unreliable. These complementary roles manifest in the difficulty-dependent results: revisions work best on easy problems where initial answers are close to correct (bin 1–2, Figure 7 right), while search works best on medium problems where the model needs to explore qualitatively different approaches (bin 3–4, Figure 3 right).

Comparison to prior work: Prior work studied these mechanisms in isolation, often reaching pessimistic conclusions about one or the other. Huang et al. (2023) concluded self-correction doesn't work. Various search papers found beam search underperforms best-of-N at scale. This paper reveals that both conclusions were artifacts of testing on implicit (and biased) difficulty distributions. The reframing as complementary axes with difficulty-dependent effectiveness explains when each mechanism works and why prior studies found conflicting results. This is a genuine reconciliation, not just a "both are good" compromise.

Significance beyond raw performance: The paper does not fully combine revisions and PRM tree-search (Section 8 acknowledges this as future work), yet the framework it provides is arguably more valuable than any individual combination would be. By establishing that these are independent scaling dimensions, the paper provides the intellectual scaffolding for a research program where future systems deploy both mechanisms adaptively — revisions to refine, search to explore, with difficulty estimation deciding the allocation. This is a research agenda, not just a system.

Fundamental vs. incremental: The framework is fundamentally clarifying; the individual methods are incremental refinements of prior work (revisions build on Qu et al., 2024; search builds on Lightman et al., 2023 and InfLLM; Xiao et al., 2024a). The paper's contribution is not a new algorithm but a new taxonomy that makes the landscape navigable. By showing that the proposal-verifier axes are complementary, the paper changes the question from "should we use revisions or search?" to "how should we combine revisions and search per prompt?" — a more productive framing.

Tie to evidence: The complementary difficulty-dependent behavior is shown in Figure 3 (right) for search and Figure 7 (right) for revisions, where the optimal strategy shifts from one mechanism to the other as difficulty changes. The FLOPs-matched analysis (Section 7, Figure 9) further shows that revisions and search have different strengths in the pretraining-vs-inference tradeoff, with revisions outperforming search in most regimes.

Innovation 3: Empirical Evidence That Test-Time Compute Can Substitute for Pretraining — With Sharp Boundaries

What's distinctive at the idea level: The FLOPs-matched comparison in Section 7 provides, to the authors' knowledge, the first demonstration in a realistic setting (no ground-truth access at inference) that a smaller model with additional test-time compute can outperform a ~14× larger model on problems within its capability range. What distinguishes this from prior work on training-inference tradeoffs (Jones, 2021; Villalobos and Atkinson, 2023; Sardana and Frankle, 2023) is the specificity of the finding: the paper doesn't claim a universal substitution but precisely characterizes where it works and where it fails.

The finding has three boundary conditions that are each independently significant:

  • Difficulty boundary: On the hardest problems (difficulty bin 5), test-time compute provides essentially zero benefit regardless of budget, meaning some capabilities can only be acquired through pretraining. This establishes that test-time compute amplifies existing capability but does not create it from nothing — a sharp boundary condition that prior work had not empirically characterized.

  • Inference-to-pretraining ratio boundary: The variable R=Dinference/DpretrainR = D_{\text{inference}} / D_{\text{pretrain}} determines how much inference budget the smaller model gets. When R≪1R \ll 1 (few inference tokens relative to pretraining), the case for test-time compute is strong because the pretraining savings dominate. When R≫1R \gg 1 (many inference tokens), the case weakens because the larger model's per-token inference cost dominates anyway. This adds practical nuance that prior analyses missed.

  • Mechanism boundary: Revisions outperform search in the FLOPs-matched comparison (Figure 9), suggesting that improving the proposal distribution is more FLOPs-efficient than improving candidate selection, at least for the specific models and budgets tested.

Comparison to prior work: Prior FLOPs-matched comparisons in the language modeling domain largely assumed access to ground-truth answers (Sardana and Frankle, 2023), making them less practically relevant — you can't use an oracle verifier when you don't know the answer. This paper operates in the realistic setting where correctness is unknown, using trained verifiers, making the comparison directly applicable to deployment decisions. The caveat is that the 14× larger model uses greedy decoding with no test-time augmentation — a weak baseline that makes the comparison somewhat favorable to test-time compute. The paper acknowledges this limitation (Section 8).

Significance beyond raw performance: This finding has direct implications for how organizations allocate compute budgets. It suggests a regime where it is more cost-effective to train a smaller model and invest the savings in smarter inference — a regime shift from the prevailing "train the largest model you can afford, then deploy with greedy decoding" paradigm. But the finding's significance is equally in its negative results: on hard problems and at high RR, pretraining is unambiguously better. This prevents over-claiming and establishes that test-time and pretraining compute are not 1-to-1 exchangeable — they have different strengths.

Fundamental vs. incremental: This is an empirical finding with fundamental implications, not a method. The contribution is the careful characterization of when the substitution works, using a FLOPs-matched framework that makes the comparison rigorous. The 4× efficiency gains in Figures 4 and 8 are the practical headline, but the difficulty-dependent boundaries in Figure 9 are the intellectual contribution — they establish the limits of the approach.

Tie to evidence: Figure 9 and the bar charts in Figure 1 show the FLOPs-matched results across difficulty bins and RR values. The key pattern: easy questions (bin 1) favor test-time compute across all RR values; hard questions (bin 5) favor pretraining across all RR values; medium questions (bins 2–4) show sensitivity to RR, with test-time compute winning at low RR and pretraining at high RR.

Innovation 4: Verifier Over-Optimization as a First-Class Phenomenon in Test-Time Scaling

What's distinctive at the idea level: While reward hacking and over-optimization are well-documented in the RLHF literature (e.g., reward models being exploited during RL fine-tuning), this paper provides some of the first clear evidence that the same phenomenon governs test-time search scaling and is the primary bottleneck preventing unbounded improvements from additional compute. The paper doesn't just document the phenomenon — it shows that over-optimization is difficulty-dependent and uses this understanding to design mitigation strategies.

The evidence is concrete and multi-pronged: beam search degrades easy-problem performance at high budgets (Figure 3, right) because the PRM's small errors get amplified by aggressive optimization; lookahead search — the most powerful optimizer — paradoxically performs worst overall (Figure 3, left) because it over-optimizes the most; and qualitative examples in Appendix M show search producing degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM but are actually incorrect. This is not a failure of search per se — it's a failure of the verifier to remain calibrated under adversarial optimization pressure.

Comparison to prior work: Prior work on test-time compute largely treated verifiers as reliable oracles whose quality improved monotonically with training. The search literature focused on developing more sophisticated algorithms (e.g., MCTS variants) without questioning whether the verifier could support them. This paper inverts the priority: rather than developing ever-more-sophisticated search algorithms, the bottleneck is verifier robustness. Lookahead search, the most sophisticated algorithm tested, underperforms simpler methods because it amplifies verifier errors. This redirects research attention from search algorithm design to verifier training and calibration.

Significance beyond raw performance: The identification of verifier over-optimization as the primary bottleneck changes the research agenda. It suggests that improving verifier robustness — through better training data, adversarial training, ensemble methods, or calibration techniques — would likely yield larger gains than developing more complex search algorithms. The paper's compute-optimal policy can be understood partly as a way to stay below the over-optimization threshold per difficulty level: use weaker optimization (best-of-N) where the verifier is reliable (easy problems) and stronger optimization (beam search) only where the verifier signal has more room to provide genuine guidance (medium problems). This is a principled mitigation, not just an observation.

The phenomenon also explains why prior work found negative results for sophisticated search methods: those studies likely pushed past the over-optimization threshold. This connects to a broader lesson for the field: more optimization is not always better, and the optimal degree of optimization depends on the reliability of the signal being optimized against.

Fundamental vs. incremental: The phenomenon identification is a fundamental diagnostic contribution. The individual pieces of evidence (beam search degrading on easy problems, lookahead underperforming) are empirical results, but the framing of verifier over-optimization as the central bottleneck — and the demonstration that it affects test-time search in difficulty-dependent ways — is a conceptual contribution that reframes how we think about scaling inference compute.

Tie to evidence: Figure 3 (right) shows the degradation on easy problems; Figure 3 (left) shows lookahead search underperforming; Appendix M provides qualitative examples; the compute-optimal policy's success (Figures 4, 8) implicitly validates the over-optimization diagnosis by showing that adaptive strategy selection (which avoids aggressive optimization where the verifier is unreliable) recovers substantial gains.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses two primary benchmarks. LongBench (Bai et al., 2023) is a bilingual, multi-task benchmark designed for long-context understanding, with 15 tasks grouped into six categories: single-document QA (NarrativeQA, Qasper, MultiFieldQA), multi-document QA (HotpotQA, 2WikiMQA, Musique), summarization (GovReport, QMSum, MultiNews), few-shot learning (TREC, TriviaQA, SAMSum), retrieval (PassageRetrieval), and code (LCC, RepoBench-P). ∞-Bench (Zhang et al., 2024) extends evaluation beyond 100K tokens with tasks including coding (Code.Debug), multiple choice (En.MC), math (Math.Find), and several retrieval variants (Retrieve.KV, Retrieve.PassKey, Retrieve.Number). For the extreme-length passkey experiment, an extended ∞-Bench PassKey.Retrieval is used, with sequences up to 10.2M tokens.

  • Base model(s). Experiments span five open-weight LLM families: Mistral-7B-Instruct-v0.2, LLaMA-3-8B-Instruct, LLaMA-3.1-8B-Instruct, Phi-3-Mini-4K-Instruct, and Phi-3.5-mini-Instruct. These represent the dominant 7–8B parameter class of models available at the time of writing, with varying native context window lengths (4K for Phi-3-Mini, 32K for Mistral, 128K for LLaMA-3.1). The paper deliberately selects models across architectures and training regimes to demonstrate that EM-LLM's benefits are not model-specific. All models are used in their pretrained/fine-tuned state without any additional training — EM-LLM is applied as a training-free inference-time modification.

  • Metrics. The primary metric is task-level accuracy or F1 score as defined by each benchmark's standard evaluation protocol, with scores aggregated into task-group averages and an overall average. LongBench reports accuracy for most tasks and F1 for a subset; ∞-Bench reports accuracy for retrieval tasks and F1 for code tasks. The paper reports maximum relative improvement over InfLLM for each task across all EM-LLM variants. Statistical significance at the benchmark level is assessed via two-tailed z-test with p < 0.05 as the threshold (Appendix A.1). For the human correlation experiments (Section 4.2), the metrics are cohesiveness (modularity, conductance, intra/inter-similarity ratio from Equation 5) and Wasserstein distance between human-annotated event boundary distributions and model-generated boundaries.

  • Baselines. The paper compares against three categories of baselines. InfLLM (Xiao et al., 2024a) is the state-of-the-art KV-retrieval method and the most direct predecessor — it uses fixed-size memory units with representative-token-based k-NN retrieval, matching EM-LLM's retrieval budget to enable direct comparison. RAG baselines use two different retrievers: NV-Embed-v2 (Lee et al., 2024), a SOTA LLM-based retriever ranking first on MTEB as of September 2024, and all-mpnet-base-v2 (Reimers, 2022), a smaller 110M-parameter model. For RAG, the context is split into chunks of 300 words with top-5 retrieval, following Li et al. (2024c). An additional RAG-S variant uses EM-LLM's surprise-based segmentation for chunking instead of fixed-size chunks. Full-context baselines process all tokens directly within the LLM's softmax attention without retrieval. For the extreme-length passkey experiment (Figure 1, bottom), comparison is against five long-context methods: YaRN (Peng et al., 2024), LongLoRA (Chen et al., 2024b), LongRoPE (Ding et al., 2024), and the extended-context versions of Mistral-7B and CodeLLaMA.

  • Generation budget / compute accounting. For retrieval-based methods (InfLLM and EM-LLM), the budget is measured as the total number of tokens in the active context window, partitioned into local context and retrieved tokens. Configurations are chosen to match InfLLM's settings: Mistral uses 4K local + 2K retrieved (4K+2K); LLaMA-3 and 3.1 use 4K+4K; Phi-3 and Phi-3.5 use 1K+3K. The retrieval budget within the retrieved portion is further divided between the similarity buffer (k_s events) and contiguity buffer (k_c events), with a contiguity ratio k_r = k_c / (k_s + k_c) = 0.3 selected from hyperparameter sweeps (Appendix D.1). For the extreme-length passkey experiment, the retrieved buffer size is held constant regardless of total sequence length — the key metric is whether the model can locate the correct answer at all, not generation quality per se.

  • Cross-validation / statistical protocol. The paper does not employ cross-validation in the traditional machine learning sense — the models are not trained. Instead, hyperparameter selection for γ (surprise threshold sensitivity) and k_r (contiguity ratio) is performed via grid search on the LongBench benchmark (Appendix D.1). For γ, values in {1.0, 1.5, 2.0, 2.5, 3.0, 3.5} are evaluated with Mistral on LongBench, then the best value per model is selected by evaluating each base LLM with γ ∈ {1, 2, 3} on the full benchmark and choosing the value yielding highest overall average. Selected values: γ = 1 for Mistral, LLaMA-3.1, Phi-3, Phi-3.5; γ = 2 for LLaMA-3-8B. For k_r, values in {0.3, 0.5, 0.7} are evaluated with Mistral on LongBench at γ ∈ {1, 2}, with slight overall preference for 0.3. Statistical significance at the benchmark level uses two-tailed z-test; the paper acknowledges that significance does not hold for the majority of individual tasks due to small per-task sample sizes but argues that the consistency and frequency of improvements, plus benchmark-level significance, justify the overall improvement claim (Appendix A.1).

Main Quantitative Results

Comparison with InfLLM (KV-Retrieval Baseline)

The headline result appears in Table 1 and the detailed task-by-task breakdowns in Appendix A.1 (Tables 3–7). EM-LLM improves on InfLLM across all five base LLMs, on 80% of individual task groups, and on the overall LongBench and ∞-Bench averages. The magnitude of improvement varies by base model, EM-LLM variant, and task type, but the direction of improvement is remarkably consistent.

LongBench aggregate results (Table 1): The best-performing EM-LLM variant per model achieves:

  • Mistral-7B: 43.71% (SM+C, surprise + modularity + contiguity) vs. InfLLM's 41.90% → +4.3% relative improvement
  • LLaMA-3-8B: 47.24% (SM, surprise + modularity) vs. InfLLM's 46.95% → +0.6% relative
  • LLaMA-3.1-8B: 51.58% (S, surprise only) vs. InfLLM's 51.05% → +1.0% relative
  • Phi-3-Mini: 35.46% (S+C, surprise + contiguity) vs. InfLLM's 34.46% → +2.9% relative
  • Phi-3.5-Mini: 34.86% (S, surprise only) vs. InfLLM's 34.22% → +1.9% relative

Task-type breakdown (Tables 3–7): The largest and most consistent gains occur in two task categories:

  • Retrieval tasks: Up to 32.69% relative improvement on PassageRetrieval (Mistral, S+C variant: 84.92% vs. InfLLM's 64.00%), and up to 19.51% on Retrieve.KV (LLaMA-3.1, SM+C: 96.80% vs. InfLLM's 81.00%). The paper attributes these gains to the dynamic segmentation keeping relevant information together within events rather than splitting it across fixed-size blocks, making the retrieval signal cleaner.
  • Multi-document QA tasks: Up to 10.45% on HotpotQA (Phi-3-Mini, S: 36.05% vs. 32.64%), 10.16% on 2WikiMQA (Phi-3-Mini, S: 28.74% vs. 27.08%), and 29.70% on Musique (Phi-3-Mini, S+C: 19.52% vs. 15.05%). Multi-document QA requires synthesizing information across documents, benefiting from the contiguity buffer's ability to maintain temporal context around relevant passages.

∞-Bench aggregate results (Table 1): EM-LLM similarly improves on ∞-Bench averages:

  • Mistral: 66.16% (SM+C) vs. InfLLM's 65.78% → +0.6%
  • LLaMA-3-8B: 49.00% (SM+C) vs. InfLLM's 50.31% → −2.6% (the one aggregate where InfLLM outperforms; note LLaMA-3's InfLLM baseline uses γ = 2, and the S variant achieves 48.81%, suggesting the retrieval-heavy ∞-Bench tasks benefit less from the paper's specific configuration for this model)
  • LLaMA-3.1-8B: 66.69% (S+C) vs. InfLLM's 64.00% → +4.2%

Variant comparison (Appendix A.1, Tables 3–7): Across all tasks and models, no single EM-LLM variant dominates universally. The SM+C variant (surprise + modularity refinement + contiguity buffer) achieves the best performance in the most settings overall, but S-only and S+C variants win on specific task-model combinations. The paper interprets this as evidence that the three components (surprise segmentation, refinement, contiguity) are complementary and their optimal combination depends on the task's retrieval and temporal reasoning demands. This finding is explored in detail in the ablation analysis (Section 4.4, Figures 12–13 in Appendix D.1).

Statistical significance: Using two-tailed z-tests, the benchmark-level improvements are statistically significant (p < 0.05) for all models except Phi-3.5 (p = 0.23 on LongBench). The paper acknowledges that individual task improvements generally do not reach significance due to small per-task sample sizes but argues that the frequency and consistency of improvements across tasks and models, combined with benchmark-level significance, support the overall improvement claim (Appendix A.1).

Comparison with RAG and Full-Context Models

These comparisons (Table 9, Appendix A.2; also Figure 1, top) position EM-LLM against two prominent paradigms for handling long contexts: embedding-based retrieval (RAG) and brute-force full attention (full-context).

LongBench (LLaMA-3.1-8B, Table 9):

  • EM-LLM (S variant, 4K+4K): 51.58% average
  • Full-context (128K window): 39.30% average
  • RAG with NV-Embed-v2: 36.44% average
  • RAG with all-mpnet-base-v2: 30.75% average (from earlier LLaMA-3 comparison, Table 8)
  • RAG-S (surprise-based chunking with NV-Embed-v2): 25.89% average

EM-LLM outperforms full-context by 31.3% relative (51.58 vs. 39.30) and the best RAG system by 41.5% relative (51.58 vs. 36.44). The full-context model performs worse than EM-LLM on most tasks, which the paper attributes to diluted attention — when all tokens compete for attention weight in a 128K context, the model struggles to focus on the relevant ones. EM-LLM's explicit retrieval mechanism avoids this by pre-selecting a small set of highly relevant events.

Per-task analysis reveals task-dependent patterns (Table 9):

  • Retrieval tasks show the most dramatic gaps. PassageRetrieval: EM-LLM achieves 99.50% vs. full-context's 100% vs. RAG's 65.50%. The full-context model theoretically sees everything but the retrieval task is designed to test precise information location — full-context models with diluted attention can still find explicitly queried information because the query provides a strong attention bias. The RAG gap is substantial, reflecting the single-retrieval-step limitation.
  • Code tasks show EM-LLM's strongest advantage over full-context: LCC 67.45% vs. 19.30%, RepoBench-P 64.33% vs. 18.33%. Code tasks require precise reference to specific functions, variables, and imports scattered throughout the context — exactly the scenario where retrieval outperforms diluted attention.
  • Few-shot learning (TREC) shows a striking pattern: EM-LLM 71.50%, full-context 4.50%, RAG (NV-Embed-v2) 22.50%. The full-context model's near-total failure on TREC is likely due to the task requiring retrieval of specific examples from throughout the context, which diluted attention fails at completely. The RAG improvement over full-context but large gap to EM-LLM demonstrates the value of layer-wise retrieval.
  • Summarization and QA show smaller but consistent advantages for EM-LLM over full-context (e.g., GovReport 35.04% vs. 34.49%, MultiFieldQA 52.52% vs. 54.98% — note full-context slightly outperforms here), suggesting that tasks requiring broad coverage of the full document benefit less from selective retrieval.

Why RAG underperforms (Appendix A.2, Appendix Figure 5): The paper identifies two structural reasons for RAG's deficit. First, single-step vs. layer-wise retrieval: RAG retrieves chunks once before generation using an external embedding model, whereas EM-LLM retrieves events at every Transformer layer using the LLM's own key representations. Appendix Figure 5 shows that the ratio of events retrieved uniquely by a single layer (not by any other layer) is substantial — roughly 60–80% of retrieved events in some layers — meaning layer-wise retrieval provides complementary information that a single retrieval step cannot capture. Second, retriever bottleneck: the NV-Embed-v2 retriever, while SOTA, has 7B parameters and uses embeddings that may not perfectly align with the downstream LLaMA-3.1's internal notion of relevance. EM-LLM's retrieval uses the LLM's own query and key projections, which are by definition aligned with its attention mechanism.

∞-Bench (LLaMA-3.1-8B, Table 9):

  • EM-LLM (S variant): 66.66% average
  • Full-context: 66.33% average
  • RAG (NV-Embed-v2): 58.97% average

On ∞-Bench, the gap between EM-LLM and full-context narrows substantially (66.66 vs. 66.33, essentially tied), while RAG maintains a clear deficit. The paper attributes the smaller full-context gap to ∞-Bench's task composition — several tasks (Retrieve.PassKey, Retrieve.Number) are explicitly designed to test retrieval at scale, and full-context models perform near-perfectly on them when the context fits in memory. EM-LLM matches this performance at a fraction of the computational cost.

RAG-S variant (Table 9): The surprise-based RAG variant underperforms standard RAG (25.89% vs. 36.44% on LongBench), which the paper hypothesizes is due to disjointed retrieved context — surprise-segmented chunks may be smaller and more numerous, leading to either highly fragmented retrieved passages or insufficient context per chunk. Incorporating contiguity into RAG's retrieval mechanism could potentially close this gap, but the paper does not test this.

Extreme-Length Passkey Retrieval (10M Tokens)

This experiment (Figure 1, bottom) tests whether EM-LLM can retrieve specific information from contexts far longer than any full-context model can handle. The task is a variant of ∞-Bench's Retrieve.PassKey: a random passkey (e.g., a 5-digit number) is inserted at a random position in a long sequence of filler text, and the model must output the passkey when queried.

Results: EM-LLM achieves 100% accuracy at 10.2M tokens (the maximum tested). By comparison, the paper's Figure 1 (bottom) shows that existing long-context methods degrade well before this scale:

  • Full-context models (with context window extensions): LongRoPE extends to 2048K (2M tokens) but performance is not reported at this extreme; the longest reported successful retrieval in prior work is LongRoPE at 2M with non-trivial but sub-100% accuracy.
  • KV-cache eviction methods: not tested at this scale in the paper's comparison but are structurally disadvantaged because they permanently discard information, making passkey retrieval at 10M effectively impossible if the passkey is evicted before the query.

The key enabling factor is EM-LLM's event-based storage: the passkey is stored as part of whichever event contains it, and when the query asks for the passkey, similarity-based retrieval can locate that event regardless of its absolute position. The event itself is loaded into the active context window, where the LLM can attend to the full passkey within its local event context. The retrieval budget (number of events loaded) remains constant regardless of total sequence length, so the computational cost per query is independent of whether the passkey is at position 1,000 or 10,000,000.

The limiting factor for even longer sequences (beyond the tested 10.2M) would be CPU/disk memory for storing the KV cache and retrieval index, not computational complexity. The paper's memory management strategy (Appendix C.3.2) uses disk offloading, making the practical limit disk capacity rather than algorithmic scaling.

Human Correlation Analysis (Section 4.2)

This experiment validates the cognitive motivation underlying EM-LLM's surprise-based segmentation by comparing LLM-generated event boundaries against human-annotated boundaries on podcast transcripts from Kumar et al. (2023).

Setup: Three short podcasts (7–30 minutes) with human-annotated event boundaries are used. The podcast transcripts are processed through LLaMA-2, LLaMA-3, and Mistral (the results for LLaMA-2 are in the main Figure 4; Mistral and LLaMA-3 results are in Appendix B, Figures 7–8). Event boundaries are generated using seven segmentation methods: random (baseline), fixed (InfLLM's fixed-size chunks, denoted F), fixed + modularity refinement (FM), fixed + conductance refinement (FC), surprise only (S), surprise + modularity refinement (SM), and surprise + conductance refinement (SC). For each method, the boundaries are compared against human annotations using two metrics.

Metric 1: Cohesion/separation quality (Figure 4A, Table 2, Appendix Figures 6–8). This measures how well the segmentation groups semantically related tokens together, using the key-similarity adjacency matrix from attention heads (Equation 5). The metrics are computed as the difference from random segmentation:

  • Modularity (↑ higher is better): Human annotations achieve the highest scores across most layers. Among computational methods, SM (surprise + modularity refinement) achieves 18.7–39.9% higher modularity than random across the three LLMs (Table 2), consistently outperforming surprise-only (S: 8.7–18.6%) and fixed-segmentation variants (FM: 0.6–29.2%). The SM variant approaches, and in some layers exceeds, human-level modularity.
  • Conductance (↓ lower is better): SM achieves 24.6–30.6% lower conductance than random, again approaching human-level performance. Surprise-only (S) achieves 23.6–29.7%.
  • Intra/inter-similarity ratio (↑ higher is better): SM achieves 25.0–35.3% higher I/IS than random; surprise-only achieves 15.0–17.9%.

The consistent pattern across all three metrics and all three LLMs is: SM > S > random ≈ fixed > FM/FC. This demonstrates that (a) surprise-based initialization captures meaningful semantic structure, (b) refinement further improves clustering quality, and (c) fixed-size segmentation (InfLLM's approach) performs no better than random — it is actively worse than random on some metrics, meaning it fragments coherent semantic units.

Metric 2: Wasserstein distance to human annotations (Figure 4B, Appendix Figures 6–8). This measures how close the distribution of model-generated event boundaries is to the distribution of human-annotated boundaries (lower distance = closer to humans). Across all three podcasts and all three LLMs, the surprise-based methods (S, SM, SC) produce boundaries significantly closer to human annotations than fixed-based methods (F, FM, FC) or random segmentation. For example, on the "Monkey" podcast with LLaMA-2 (Figure 4B): SM achieves Wasserstein distance ~0.13 (×10) vs. S at ~0.14 vs. human-to-human (individual subjects vs. consensus) at ~0.08 vs. fixed-based methods at ~0.33–0.50. The surprise-based methods are within a factor of 2 of inter-human variability, while fixed methods are 4–6× farther from humans than inter-human variability.

Takeaway from this experiment: EM-LLM's surprise-based segmentation (with or without refinement) produces event boundaries that (1) create more cohesive token groupings than fixed-size or random segmentation, and (2) correlate strongly with human-perceived event boundaries. This validates the paper's central cognitive hypothesis: that LLM next-token prediction error serves as a proxy for the prediction-error signals that drive human event segmentation, and that organizing LLM context windows according to these signals produces more natural, retrievable memory units than arbitrary fixed-size chunking.

Segmentation Method Comparison on PG-19 (Section 4.3)

The PG-19 dataset (Rae et al., 2020), comprising a diverse corpus of English books, is used to compare segmentation methods across a broader text distribution than the podcast transcripts.

Results (Table 2): Evaluated across Mistral-7B, LLaMA-2-7B, and LLaMA-3-8B with γ = 10⁻³, using modularity, conductance, and I/IS as metrics (all reported as difference from random segmentation):

  • SM (surprise + modularity refinement) consistently achieves the highest scores across all three LLMs and all three metrics. For LLaMA-3-8B: modularity +27.0 (×10⁵), conductance −30.6, I/IS +28.1 (×10³).
  • SC (surprise + conductance refinement) is second-best: conductance −33.9 (best conductance across all methods), modularity +18.3, I/IS +16.4.
  • S (surprise only) achieves intermediate scores: modularity +13.1, conductance −29.7, I/IS +15.7.
  • Fixed-based methods with refinement (FM, FC) underperform their surprise-based counterparts: FM achieves modularity +18.9 (much lower than SM's +27.0 but higher than S's +13.1), while fixed-only (F) performs near or slightly below random: modularity −1.6, conductance +11.3, I/IS −3.8.

The key finding: initialization with surprise matters more than refinement alone. Fixed-segmentation with modularity refinement (FM) achieves better metrics than fixed-only (F) but substantially worse than surprise + refinement (SM). This supports the paper's two-stage approach: surprise provides a cognitively meaningful initialization, and refinement amplifies it. Refinement applied to random or fixed initializations cannot recover the quality achieved by surprise-based initialization.

Cross-LLM consistency: The relative ordering of methods is identical across Mistral, LLaMA-2, and LLaMA-3, suggesting the surprise signal captures a property of the text (mediated by the model's internal representations) that generalizes across model architectures and training regimes. The absolute metric values differ (LLaMA-3 generally shows higher modularity improvements than LLaMA-2), but the method ranking is preserved.

Effect of Surprise Threshold Sensitivity (γ) (Appendix D.1, Figure 12)

This ablation evaluates how the surprise threshold parameter γ (Equation 1) affects EM-LLM's performance on LongBench with Mistral-7B. The parameter controls threshold sensitivity: larger γ means fewer, larger events (higher threshold, fewer boundaries detected).

Results (Figure 12): Evaluating γ ∈ {1.0, 1.5, 2.0, 2.5, 3.0, 3.5}:

  • Smaller γ (more boundaries, smaller events) generally performs best. γ = 1 achieves the highest average score across LongBench tasks (approximately 43.5% overall) and is best or within 1 percentage point of best for most individual tasks.
  • Larger γ degrades performance on most tasks, with γ = 3.5 showing notably lower scores on retrieval tasks (PassageRetrieval drops from ~84% at γ = 1 to ~63% at γ = 3.5) and QA tasks.
  • Refinement (SM) is particularly beneficial at larger γ. At γ = 3.5, the gap between S and SM variants is larger than at γ = 1, because surprise-only segmentation with a high threshold produces very coarse events that refinement can meaningfully improve by splitting at natural cohesion breakpoints.
  • The relationship between mean event size and performance is task-dependent. Retrieval tasks (PassageRetrieval) favor smaller events (more granular segmentation), while summarization tasks (GovReport, QMSum) are less sensitive to event size.

Per-model γ selection: The paper evaluates γ ∈ {1, 2, 3} on LongBench for each base LLM using surprise-only segmentation. Selected values: γ = 1 for Mistral, LLaMA-3.1, Phi-3, Phi-3.5; γ = 2 for LLaMA-3-8B. The LLaMA-3 preference for γ = 2 may reflect differences in the model's prediction confidence calibration — more confident models produce less variable surprise signals, requiring a lower threshold to detect meaningful boundaries.

Effect of Contiguity Ratio (k_r) (Appendix D.1, Figure 13)

This ablation evaluates the contiguity buffer size relative to the similarity buffer, parameterized as the contiguity ratio k_r = k_c / (k_s + k_c), on LongBench with Mistral-7B.

Results (Figure 13): Evaluating k_r ∈ {0.3, 0.5, 0.7} at γ ∈ {1, 2}:

  • k_r = 0.3 performs best overall, yielding the highest average score (approximately 43.5% at γ = 1). This means 30% of the retrieval budget goes to the contiguity buffer and 70% to the similarity buffer.
  • Task-dependent optimality: Some tasks benefit from higher contiguity (e.g., 2WikiMQA shows best performance at k_r = 0.5 or 0.7), while others favor lower contiguity or are insensitive. Retrieval tasks (PassageRetrieval) show significant sensitivity, with k_r = 0.3 clearly outperforming higher ratios.
  • Larger k_r reduces similarity buffer capacity: at k_r = 0.7, the similarity buffer holds only 30% of the retrieved events, limiting the model's ability to find topically relevant information. The paper interprets this as evidence that similarity-based retrieval is the primary driver of performance, with contiguity providing a secondary, task-dependent benefit.
  • Interaction with γ: The γ = 1 configuration consistently outperforms γ = 2 across all k_r values, confirming the earlier finding that lower surprise thresholds are generally better.

The paper selects k_r = 0.3 for all subsequent experiments based on the slight overall preference and the intuition that the contiguity buffer should provide temporal context without crowding out relevance-based retrieval.

Retrieved Buffer Size Ablation (Appendix D.3, Table 12, Figure 14)

This ablation examines how the number of retrieved tokens affects performance on LongBench's summarization tasks with Mistral-7B and EM-LLM (S variant).

Results (Table 12): Testing buffer sizes of 1K, 2K, 4K, and 6K retrieved tokens (with 4K local tokens):

  • GovReport: Performance is essentially flat across buffer sizes (31.26–31.44%), suggesting the task is not sensitive to retrieval quantity beyond a minimum threshold.
  • QMSum: Performance increases with buffer size from 23.24% at 1K to 24.47% at 4K, then plateaus at 24.30% at 6K. This task benefits from more retrieved context but with diminishing returns.
  • MultiNews: Flat at 26.59–26.67%, similar to GovReport.
  • SAMSum: Modest variation (42.48–43.38%) with no clear trend.

Context-length interaction (Figure 14): When plotted as a function of example context length, QMSum shows that longer examples benefit from larger retrieved buffers — the 6K buffer outperforms 2K on examples >8K tokens. On other tasks, the relationship is weaker or reversed (shorter examples sometimes prefer smaller buffers). This suggests that retrieval budget should ideally scale with total context length, but the paper uses a fixed budget for simplicity and fair comparison with InfLLM.

Critical Assessment

Does EM-LLM Outperform InfLLM? (Claim: "superior performance, consistently outperforming the state-of-the-art retrieval model InfLLM across various baseline LLMs")

What the experiments demonstrate: Across five base LLMs, the best EM-LLM variant improves over InfLLM on LongBench averages (Table 1), with relative improvements ranging from +0.6% (LLaMA-3-8B) to +4.3% (Mistral-7B). On ∞-Bench, improvements are also present for four of five models, though LLaMA-3-8B shows a −2.6% regression. The improvements are consistent in direction but modest in magnitude — this is not a transformative leap over InfLLM but a consistent edge.

What requires qualification: The claim of "consistently outperforming" should be understood at the benchmark-average level, not the per-task level. The paper acknowledges (Appendix A.1) that individual task improvements largely do not reach statistical significance (p > 0.05) due to small per-task sample sizes. The consistency across tasks and models is the evidence, not per-task significance. This is a reasonable interpretation — the alternative (concluding no improvement because per-task tests are underpowered) would discard a clear signal — but it is a weaker form of evidence than statistically significant per-task improvements would be.

Missing comparisons: The paper does not compare against KV-cache eviction methods (H2O, Dynamic Memory Compression) or against linear-attention models on long-context benchmarks. It also does not compare against InfLLM with tuned hyperparameters — InfLLM's fixed chunk size and representative token count are used as-is, and it's possible that InfLLM with, say, smaller chunk sizes would close some of the gap. InfLLM is the most natural and direct baseline, but the paper's framing as "superior to the SOTA" implies a broader comparison than is actually performed.

Concern about the InfLLM configuration: For LLaMA-3-8B, InfLLM uses 4K+4K configuration while the model's native context is 8K (reported as 8K in some sources, though the paper treats it as compatible with 4K+4K). This is InfLLM's default configuration and the paper matches it for fairness, but it's possible that InfLLM would benefit from different settings that EM-LLM's dynamic segmentation partly compensates for.

Does EM-LLM Outperform RAG and Full-Context Models? (Claim: "surpasses full-context models in most tasks, while successfully performing retrieval across 10 million tokens")

What the experiments demonstrate: On LongBench with LLaMA-3.1-8B (Table 9), EM-LLM (S variant, 51.58%) substantially outperforms full-context (39.30%) and the best RAG system (NV-Embed-v2, 36.44%). On ∞-Bench, EM-LLM (66.66%) essentially ties full-context (66.33%) while outperforming RAG (58.97%). The 10M-token passkey retrieval at 100% accuracy is a genuine demonstration of scalability that no full-context model can currently match.

What requires qualification (full-context comparison): The full-context model's poor performance on LongBench (39.30% vs. 51.58%) is striking but raises a methodological concern: LLaMA-3.1-8B has a 128K native context window, yet it underperforms a retrieval method with a 4K+4K effective context. This suggests the model is suffering from the "lost in the middle" problem — but it may also reflect that the full-context model was tested in a configuration that doesn't maximize its performance (e.g., no special prompt formatting for long-context, no position interpolation, etc.). The paper does not report whether any mitigation techniques (e.g., prompting strategies, attention modulation) were attempted for the full-context baseline.

More importantly, the "full-context" comparison confounds two different models: EM-LLM's full-context baseline in Table 9 uses the same LLaMA-3.1-8B model but processes the full sequence with standard softmax attention. This is a fair comparison in that it uses the same underlying LLM, but it is unfair in that EM-LLM gets the benefit of an explicit retrieval mechanism while full-context gets none. A more informative comparison would give both systems the same total context budget (e.g., full-context with a 128K window vs. EM-LLM with 128K of retrieved events), but this is not tested. The current comparison tests EM-LLM's retrieval strategy against a model that is overloaded by its own context.

What requires qualification (RAG comparison): The RAG setup uses 300-word chunks and top-5 retrieval, following a standard protocol (Li et al., 2024c). However, RAG performance could potentially be improved by:

  • Tuning chunk size to match the optimal event size EM-LLM discovers automatically.
  • Using more sophisticated pre/post-retrieval techniques (query expansion, reranking) not tested here.
  • Using a retriever model of comparable scale to the LLM (NV-Embed-v2 is 7B, matching LLaMA-3.1-8B, so this concern is lessened).

The paper's RAG comparison is nonetheless fair as a test of the standard RAG pipeline most practitioners would deploy, and the 15+ percentage point gap on LongBench is large enough that even optimized RAG would unlikely close it entirely. The paper's explanation — layer-wise retrieval provides complementary information across attention heads — is compelling and supported by Appendix Figure 5's visualization of per-layer retrieval uniqueness.

What requires qualification (10M-token result): The 10M-token passkey retrieval is an impressive scalability demonstration, but it tests only a single, simple capability: can the model find a specific string in a long context? It does not test whether EM-LLM can reason over information distributed across events at this scale, or maintain coherence in multi-turn interactions at this length. The result establishes that retrieval works at scale, not that understanding works at scale.

Does Surprise-Based Segmentation Correlate with Human Event Perception? (Claim: "strong correlations between EM-LLM's event segmentation and human-perceived events")

What the experiments demonstrate: Figure 4 and Appendix Figures 6–8 show that surprise-based segmentation methods (S, SM, SC) produce event boundaries significantly closer to human annotations (lower Wasserstein distance) than fixed-segmentation or random baselines, and achieve higher cohesion metrics (modularity, conductance, I/IS) that approach or match human-segmentation quality. This holds across three different LLMs (LLaMA-2, LLaMA-3, Mistral) on three different podcasts.

What requires qualification: The human data comes from a single study (Kumar et al., 2023) with three short podcasts (7–30 minutes each). This is a small sample — three texts, each with human annotations from a limited number of participants (the exact number is not specified in this paper). The correlation shown is genuine but we do not know whether it generalizes to other text types (technical documents, code, dialogue, non-narrative text) or to longer texts where event structures become more complex and nested.

Furthermore, the paper's Method for converting human annotations to discrete boundaries for comparison (Appendix B.1) involves selecting "as many of the most likely positions in the [Gaussian-smoothed human] distribution as our initial surprise-based event segmentation had identified." This means the human baseline is matched to the quantity of EM-LLM's boundaries, not evaluated independently. If EM-LLM produces many boundaries, the human comparison uses the same number of most-likely human boundary positions; if it produces few, fewer human positions are used. This matching procedure could artificially inflate the similarity between methods if the number of boundaries strongly affects the distance metric. The paper acknowledges this (the discrete human positions are converted to a Mixture of Gaussians for Wasserstein distance computation), but the initial boundary count matching step is a potential source of bias.

Does Boundary Refinement Improve Performance? (Claim: "refinement... enhances efficient information recall")

What the experiments demonstrate: In the human correlation analysis (Figure 4, Table 2), refinement (SM, SC) consistently improves cohesion metrics over surprise-only (S) across all three LLMs and all three metrics. In the benchmark results (Tables 3–7), refinement variants (SM, SM+C) achieve the best performance in 60% of tasks across LongBench and ∞-Bench (Section 4.4). The evidence that refinement improves retrieval-oriented metrics is clear.

What requires qualification: Despite refinement improving similarity metrics, the benchmark performance improvement from adding refinement to surprise-only is inconsistent. In several model-task combinations, surprise-only (S) outperforms surprise+refinement (SM). For example, on LLaMA-3.1-8B LongBench average (Table 5): S achieves 51.58% vs. SM at 51.28%. On Mistral-7B PassageRetrieval (Table 3): SM at 78.92% vs. S at 82.67%. The improvement from refinement is task-dependent and sometimes negative. This suggests that while refinement optimizes the cohesion metric, the relationship between cohesion and downstream task performance is not monotonic — overly aggressive cohesion optimization might merge semantically distinct but attentionally similar events, reducing retrieval precision. The paper does not explore this tradeoff.

Gaps and Missing Experiments

1. No comparison against LLMs with native long-context training. The paper compares against full-context with the same base LLM, but not against models specifically trained for long contexts (e.g., LLaMA-3.1's 128K version is used, but models like GPT-4-128K, Claude-3-200K, or Gemini-1.5-Pro-1M are not tested). These commercial models have demonstrated strong long-context performance that might reduce or eliminate the advantage of retrieval-based methods. The paper's focus on open-weight 7–8B models is consistent with its training-free, drop-in philosophy, but limits the strength of the "surpasses full-context models" claim.

2. No latency or throughput measurements. The paper reports wall-clock time per chunk (Table 10, Appendix C.3.1): EM-LLM (S) takes 1.12× InfLLM's time, EM-LLM (SM) takes 1.62×. These are modest overheads, but the paper does not report end-to-end latency for the full LongBench or ∞-Bench benchmarks, nor throughput (tokens/second) for generation tasks. For deployment, these metrics matter as much as accuracy, and the boundary refinement step's O(nm) cost could become significant at very long sequences despite the paper's asymptotic analysis showing it is dominated by attention.

3. No test of per-layer independent segmentation. The paper mentions as future work "extending our segmentation processes to operate at each layer of the Transformer independently" (Section 5), which could create a hierarchical event structure. The current implementation uses a single event segmentation shared across all layers, potentially missing layer-specific semantic groupings. No ablation tests whether per-layer segmentation would help.

4. Dynamic buffer sizing is not explored. The retrieval buffer size is fixed (e.g., 4K+2K for Mistral), matching InfLLM. The paper's own ablation (Appendix D.3) suggests that longer contexts benefit from larger buffers, but the main experiments use a one-size-fits-all configuration. An adaptive policy that allocates more retrieval budget to longer sequences could improve performance but is not tested.

5. The 10M-token experiment tests only passkey retrieval. This demonstrates that the indexing and retrieval mechanism functions at scale, but it does not demonstrate that EM-LLM can perform complex reasoning (multi-hop QA, summarization, code understanding) at this length. A million-token codebase understanding task or book-length summarization would be more informative about practical capabilities but is not tested.

6. No combination with positional encoding extension methods. EM-LLM assigns fixed position embeddings to retrieved events (Section 3.4). This discards positional information that might be valuable for certain tasks (e.g., "what happened between event X and event Y?"). Combining EM-LLM with positional interpolation or RoPE scaling could allow retrieved events to carry approximate positional information without causing out-of-distribution issues. This is not explored.

Summary Assessment

The experiments provide solid evidence that surprise-based event segmentation with graph-theoretic refinement produces more coherent memory units than fixed-size chunking, that the resulting retrieval mechanism outperforms InfLLM on most benchmarks (with modest but consistent margins), and that the approach scales to lengths far beyond full-context models. The human correlation analysis provides genuine validation of the cognitive motivation, though on a small and specific dataset.

The evidence for outperforming RAG and full-context models is strong but narrower than claimed: the comparisons use specific configurations that may not be optimal for the baselines, and the full-context comparison confounds model scale with context processing strategy. The 10M-token result is an impressive scalability proof but tests only a single, simple capability.

The most robust finding in the paper is that surprise-based segmentation captures semantically meaningful event structure that both correlates with human perception and improves retrieval quality over fixed-size chunking. This finding is replicated across three LLMs, three cohesion metrics, and two evaluation paradigms (human correlation and retrieval benchmarks). The refinement and contiguity contributions are also well-supported but show more task-dependent benefits, suggesting they are useful but secondary to the core surprise-based segmentation innovation.

6. Limitations and Trade-offs

Limitation 1: Dynamic Segmentation Depends on the Quality and Domain-Specificity of the LLM's Surprise Signal

The assumption or constraint: EM-LLM's entire memory formation pipeline rests on the hypothesis that the base LLM's next-token prediction errors — its Bayesian surprise $-\log P(x_t \mid x_{<t})$ — correspond to meaningful semantic or structural shifts in the content, and that these shifts define event boundaries that are useful for downstream retrieval. The paper validates this hypothesis on narrative text (podcast transcripts, Section 4.2; English books, Section 4.3), where the correlation with human-perceived event boundaries is strong. However, there is no guarantee that the surprise signal behaves similarly on other text types. On highly technical, repetitive, or non-narrative content — legal boilerplate, API documentation, log files, structured data serialized as text — the LLM's prediction errors may reflect superficial features (rare tokens, formatting changes, tokenization artifacts) rather than genuine semantic transitions. The paper does not evaluate surprise-based segmentation on any non-narrative, non-book text domain.

The consequence: If surprise fails to identify meaningful event boundaries on a given domain, EM-LLM's segmentation degrades toward random or worse-than-random chunking. The refinement step can improve cohesion within whatever initial boundaries it receives, but it cannot rescue a fundamentally poor initialization. Table 2 demonstrates this: fixed-segmentation with modularity refinement (FM) achieves only +18.9 modularity vs. +27.0 for surprise + refinement (SM) on LLaMA-3-8B, showing that refinement amplifies initialization quality but cannot overcome a bad starting point. In a domain where surprise produces boundaries no better than fixed (or actively misleading), the refinement step would converge to a local optimum around an incorrect segmentation, and EM-LLM would reduce to something approximating InfLLM's fixed-size chunking — losing the claimed advantages without necessarily signaling failure to the user. There is no confidence score or quality check on the segmentation itself; the system silently degrades.

What evidence exists in the paper: The only domains tested for segmentation quality are narrative podcasts (Kumar et al., 2023; Section 4.2) and English books (PG-19; Rae et al., 2020; Section 4.3). Both are narrative, human-authored text with clear event structure. The benchmarks (LongBench, ∞-Bench) include diverse text types — academic papers, Wikipedia articles, code, meeting transcripts, legal documents — but the paper reports only end-task accuracy on these benchmarks, not segmentation quality metrics (modularity, conductance, I/IS) by domain. A practitioner deploying EM-LLM on a new domain has no way to predict whether the surprise signal will be informative without running the same kind of human-annotation study that Section 4.2 performs, which is impractical for most deployment scenarios. The paper does not report any negative results where surprise-based segmentation fails to improve over fixed-size chunking on a specific benchmark subtask, which would be diagnostically useful.

Mitigation status: The paper does not address domain sensitivity directly. The boundary refinement step provides partial mitigation — it can correct surprise-based boundaries that produce poor within-event coherence — but only to the extent that the initial boundaries are approximately correct. The paper acknowledges that "numerous other methods for graph clustering and sequence segmentation could potentially be applied" (Section 5), suggesting the current refinement is one of many possibilities, but does not propose alternative boundary detection mechanisms for domains where surprise fails. No domain-adaptation technique (e.g., calibrating the surprise threshold per domain, using a different segmentation signal) is explored.

Limitation 2: The Difficulty Estimation Overhead Is Not Accounted for in the Efficiency Claims

The assumption or constraint: The paper's difficulty estimation procedure requires generating 2048 samples per question and either checking them against ground-truth answers (oracle) or averaging the PRM's predicted final-answer correctness (predicted). The authors explicitly acknowledge this cost in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

In a deployment scenario, a prompt arrives, the system knows nothing about it, and it must decide how to allocate the test-time compute budget. The current method for making that decision — which is what makes the compute-optimal policy possible — costs more than the largest test-time compute budget studied (2048 samples vs. the maximum 256–512 generation budget in the experiments). The paper frames this as an "exploration-exploitation tradeoff" (Section 3.2) but does not include the exploration cost in any reported efficiency metric.

The consequence: The headline 4×4\times efficiency gain (e.g., compute-optimal at 16 generations matching best-of-N at 64; Figure 4) is computed after difficulty is known. The true deployment efficiency, including the cost of learning difficulty, would be dramatically lower. If difficulty estimation costs 2048 generations per question, then the total compute for compute-optimal at a 16-generation budget is 2048+16=20642048 + 16 = 2064 generations, which exceeds the 64-generation best-of-N baseline that the 4×4\times figure claims to match — completely negating the efficiency advantage. For the gains to be real in practice, difficulty must be estimated far more cheaply, but the paper does not provide such a method. The predicted-difficulty variant removes the need for ground-truth labels but does not reduce the sample count — it still requires 2048 generations per question to estimate PRM score distributions.

This limitation is not merely a missing optimization; it undermines the practical interpretation of the paper's central claim. The 4×4\times gain is an upper bound on achievable efficiency under the assumption that a zero-cost difficulty oracle exists, not a realized deployment gain. The paper's demonstrations that predicted difficulty bins nearly match oracle bins (Figure 4, Figure 8; "the two curves largely overlap") show that the PRM can substitute for ground-truth labels, but they do not address the cost of computing those predicted bins. A practitioner reading the abstract and seeing "4×4\times better efficiency" might reasonably conclude that their system will use 4×4\times less compute at deployment — the paper should be understood as showing that such efficiency is possible in principle given cheap difficulty estimation, not that it is achieved with the current method.

What evidence exists in the paper: The difficulty estimation procedure is described in Section 3.2, with the 2048-sample count stated explicitly. The acknowledgment that "our experiments do not account for this cost" appears in the same section. The paper also notes that "predicting difficulty with the PRM uses zero ground truth data but incurs significant computation cost" (Section 3.2). The Figures 4, 8, and 9 all plot compute-optimal scaling curves that exclude difficulty estimation cost — the x-axes show only the test-time strategy budget, not the total inference cost including difficulty assessment.

Mitigation status: The paper flags cost-efficient difficulty estimation as "a key avenue for future work" (Section 3.2) and suggests two directions: "pretraining or finetuning models to directly predict difficulty of a question" (Section 8), and starting inference with "a few parallel samples, assessing the score distribution, and deciding in real-time" (Section 3.2). Neither is implemented or evaluated. Until such a method exists and is validated, the 4×4\times figure should be treated as a proof-of-concept efficiency ceiling, not a deployment-ready claim.

Limitation 3: No Combination of the Two Complementary Axes (Revisions + PRM Search)

The assumption or constraint: The paper studies two independent test-time compute mechanisms — revisions (modifying the proposal distribution) and PRM-guided search (optimizing against a verifier) — and demonstrates they have complementary, difficulty-dependent strengths (Section 3.2: revisions best on easy problems, search best on medium problems). However, the paper never combines them. Section 8 explicitly acknowledges:

"we did not experiment with PRM tree-search techniques in combination with revisions"

The experiments treat them as separate pipelines: compute-optimal search selects the best search algorithm per difficulty bin (Section 5.3), and compute-optimal revisions selects the best sequential-to-parallel ratio per difficulty bin (Section 6). No experiment involves using a revision model as the proposal distribution within a PRM-guided beam search, or using the PRM to score and select among revision chains, or adaptively switching from revisions to search mid-computation.

The consequence: The paper's results represent a lower bound on what an integrated system could achieve. The two mechanisms have genuinely complementary strengths: revisions improve the quality of generated candidates (better proposal), while search improves candidate selection (better verifier). Applied together, the revision model would produce higher-quality candidate steps for beam search to evaluate, potentially pushing correct solutions above the PRM's detection threshold on problems where neither mechanism alone suffices. Conversely, the PRM could guide which revision branch to pursue, avoiding the 38% correct-to-incorrect reversion rate (Section 6.1) by detecting when a revision is making things worse.

The difficulty-dependent patterns in the paper suggest where combined benefits would be largest: on medium-difficulty problems (bins 3–4), where both mechanisms individually outperform best-of-N (search: Figure 3 right; revisions: less pronounced but the balanced sequential-parallel ratio in Figure 7 right shows benefits from both exploration and refinement). A combined system could use the revision model to generate candidate solution steps, the PRM to score them, beam search to allocate compute toward promising branches, and the contiguity-like sequential context from the revision chain to maintain coherence — none of which requires fundamentally new components beyond what the paper already implements.

What evidence exists in the paper: The evidence for complementarity is in the per-difficulty-bin analyses: Figure 3 (right) shows search beating best-of-N on bins 3–4 but degrading on bin 1; Figure 7 (right) shows revisions with a balanced ratio beating fully parallel or fully sequential on bins 3–4; Table 1 row "S+C" (surprise + contiguity) shows contiguity's value in retrieval tasks. The paper never measures the interaction effect of combining these mechanisms, nor does it report any experiment attempting to do so. The acknowledgment in Section 8 is the only direct statement about this gap.

Mitigation status: The paper frames the combination as a natural next step and "a promising direction for future research" (Section 8). It does not speculate on expected gains or propose a specific integration architecture. For a practitioner, this means that the individual mechanisms' demonstrated improvements are likely conservative — a fully integrated system might exceed the reported numbers — but the integration risk (potential negative interactions, additional hyperparameters, increased complexity) is entirely uncharacterized.

Limitation 4: The Revision Model's Correct-to-Incorrect Reversion and Sensitivity to Training Methodology

The assumption or constraint: The revision model is fine-tuned solely on sequences where all in-context answers are incorrect, followed by a correct target (Section 6.1). This training objective teaches the model to correct errors but provides no signal about what to do when the current answer is already correct. At test time, revision chains inevitably produce correct answers at intermediate steps (Figure 6, left: pass@1 starts at ~18.2% and rises to ~24–25% through the chain), and when the model encounters these correct answers in its context, it has no learned behavior for preserving them. The paper reports:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach" (Section 6.1)

The mitigation — selecting the best answer from the entire chain via majority voting or verifier rather than always taking the last revision — is a post-hoc patch that does not prevent the reversion from occurring. It salvages the chain's best output but wastes the computation spent on subsequent (degrading) revision steps.

Furthermore, Appendix K (Figure 16) shows that attempting to optimize the revision model further using ReSTEM^{EM} (Singh et al., 2024) causes performance to degrade substantially with sequential revisions — "additional sequential revisions substantially hurt performance" (Appendix K). The paper hypothesizes that "on-policy data collection in ReSTEM^{EM} exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly." This negative result reveals that the revision training procedure is fragile with respect to training methodology: small changes to the data generation process can break the revision capability entirely.

The consequence: The 38% correct-to-incorrect reversion rate means that, in the fully sequential setting, more than a third of correctly solved problems are unsolved by the next revision step. While the within-chain selection mechanism recovers these correct answers (so the final output can still be correct), the computation spent on the degrading revisions is wasted — it contributes nothing to output quality and could have been allocated to other parallel chains. This partially explains why the optimal sequential-to-parallel ratio is not purely sequential even on easy problems (Figure 7, right): the reversion problem makes long chains increasingly inefficient.

More concerning for practitioners is the sensitivity to training methodology. The ReSTEM^{EM} failure (Appendix K) suggests that the paper's specific training recipe — offline data construction with edit-distance-based incorrect-correct pairing — is not trivially improvable or generalizable. A practitioner attempting to replicate or extend the revision model on a different base LLM, dataset, or training pipeline may encounter unexpected degradation, and the paper provides no diagnostic tools or stability guarantees.

What evidence exists in the paper: The 38% reversion rate is reported in Section 6.1. The ReSTEM^{EM} degradation is detailed in Appendix K (Figure 16), showing that at 256 generations, fully sequential performance drops to ~33.5% vs. ~38.5% at the optimal ratio. The paper does not quantify how much of the revision chain's computational budget is wasted on reverting correct answers, nor does it report the distribution of when in the chain reversions occur (early vs. late).

Mitigation status: The paper identifies the reversion problem and partially mitigates it with within-chain selection (majority voting or verifier-based selection across the chain, Section 6.1). This patch works — it recovers the correct answer that was later revised away — but it does not prevent the reversion from occurring, nor does it reduce the wasted computation. The paper does not explore training the revision model to recognize when no revision is needed (i.e., adding correct-to-correct or correct-to-unchanged trajectories to the training data), which would be a more principled solution. The ReSTEM^{EM} failure is presented as a negative result without a proposed fix, and the paper does not investigate what specifically about on-policy data collection breaks the revision capability.

Limitation 5: The Contiguity Buffer's Importance Is Task-Dependent and Its Optimal Configuration Is Not Predictable a Priori

The assumption or constraint: The contiguity buffer is introduced as a key innovation — a mechanism for replicating the temporal contiguity and asymmetry effects observed in human free recall (Howard and Kahana, 2002) and recently demonstrated in Transformer attention heads (Ji-An et al., 2024). The paper argues that this buffer "promotes temporal relationships in retrieval" and "enables the LLM's induction attention heads to exhibit contiguity and asymmetry effects" (Section 3.4). However, the experiments reveal that the contiguity buffer's benefit is highly task-dependent: it improves performance on retrieval-intensive tasks (PassageRetrieval gains up to 40% over InfLLM with S+C; Table 3) and multi-document QA (up to 29.7% on Musique with S+C; Table 6), but provides little or negative benefit on tasks where temporal relationships are less important, such as summarization (GovReport: S+C 31.10% vs. S 31.40% on Mistral; Table 3) and code (LCC: S+C 54.90% vs. SM 57.03% on Mistral; Table 3).

The paper acknowledges this task-dependence (Section 4.4): "the fact that certain tasks still appear to benefit more from either surprise-only, refinement, or contiguity, is an interesting result. This is likely due to the nature of the tasks and the varying importance of contiguity across these tasks." However, it does not provide a way to predict which tasks will benefit from contiguity, nor a mechanism for dynamically enabling or disabling the contiguity buffer based on task characteristics.

The consequence: A practitioner deploying EM-LLM must choose whether to include the contiguity buffer and at what ratio $k_r$ — but the paper provides no principled way to make this choice beyond the specific benchmarks tested. The hyperparameter sweep in Appendix D.1 (Figure 13) evaluates $k_r \in \{0.3, 0.5, 0.7\}$ on LongBench and finds $k_r = 0.3$ best overall, but this is an average across 15 tasks with heterogeneous requirements. A deployment on a single task type (e.g., legal document QA, which resembles multi-document QA, or meeting summarization, which resembles GovReport/QMSum) might benefit from a very different $k_r$. The paper's recommended $k_r = 0.3$ could be suboptimal for a specific deployment, and the only way to determine the optimal value is to run a hyperparameter sweep on that deployment's task distribution — exactly what most practitioners cannot do (they lack labeled evaluation data for their specific use case).

Furthermore, the contiguity buffer competes with the similarity buffer for the fixed retrieval budget $k$. Increasing $k_r$ reduces $k_s$, crowding out similarity-based retrieval. The paper notes that "a contiguity buffer that is as big or smaller than the similarity buffer yields the best results... suggesting that the similarity buffer is still the most crucial part of our approach" (Section 4.4). This means the contiguity buffer provides a secondary, conditional benefit that trades off against the primary retrieval mechanism — getting this tradeoff wrong hurts performance, and the optimal tradeoff varies by task in ways the paper does not characterize.

What evidence exists in the paper: Table 3 (Mistral) shows tasks where contiguity helps (PassageRetrieval: S+C 84.92% vs. S 82.67%; Musique: S+C 17.98% vs. S 17.97%) and tasks where it hurts (NarrativeQA: S+C 21.10% vs. S 21.77%; GovReport: S+C 31.10% vs. S 31.40%). Table 4 (LLaMA-3) shows similar inconsistency: S+C helps on NarrativeQA (24.66% vs. S 24.47%) but hurts on LCC (58.55% vs. S 58.49%). The hyperparameter sweep (Figure 13, Appendix D.1) shows that the optimal $k_r$ varies across tasks, with tasks like 2WikiMQA benefiting from higher contiguity and tasks like PassageRetrieval strongly preferring lower contiguity. The paper does not report per-task $k_r$ optimal values or attempt to predict them from task metadata.

Mitigation status: The paper identifies the task-dependence as "an interesting result" (Section 4.4) but does not propose a solution. The compute-optimal framework developed for search and revisions (Sections 5–6) — where the optimal strategy is selected per difficulty bin — could in principle be extended to select the optimal contiguity ratio per task type, but no such extension is implemented or suggested. The paper treats contiguity as a fixed architectural choice rather than an adaptive parameter, which limits its practical applicability. A dynamic contiguity policy that adjusts $k_r$ based on real-time retrieval patterns (e.g., increasing contiguity when the similarity-retrieved events' temporal neighbors are frequently relevant) would address this limitation but is not explored.

Limitation 6: Single Benchmark Family, Single Model Scale, and Unverified Generalization Claims

The assumption or constraint: All experiments are conducted on the MATH benchmark (Hendrycks et al., 2021) for the search and revision experiments, and on LongBench and ∞-Bench for the EM-LLM experiments. The base models are all in the 7–8B parameter class (PaLM 2-S* for the search/revision paper; LLaMA-3, Mistral, Phi for EM-LLM). The paper acknowledges these constraints but frames them as sufficient:

"we believe this model is representative of the capabilities of many contemporary LLMs" (Section 4)

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute" (Section 7)

These are assumptions of representativeness that the paper does not verify empirically.

The consequence: Several aspects of the findings could be specific to the benchmark or model scale used:

  • PRM quality and over-optimization behavior (Section 5.3) depend on the base model's output distribution. A model with better-pass@1 on MATH (e.g., GPT-4, Claude, or a fine-tuned math specialist) would shift the difficulty distribution upward, potentially changing which difficulty bins benefit from search vs. revisions. A model with worse calibration might exhibit different (possibly worse) verifier over-optimization patterns.

  • The revision model's ability to learn from incorrect in-context examples (Section 6.1) depends on the base model's in-context learning capacity and its tendency to produce structurally similar but wrong answers (for edit-distance-based pairing). Models with different failure modes — e.g., models that tend to produce completely unrelated wrong answers rather than near-misses — might not benefit from the edit-distance-based training data construction.

  • The difficulty quintile bins (Section 3.2) are defined relative to the base model's pass@1 distribution. A substantially more capable model would have a different bin distribution, with fewer problems in bins 4–5 (hard) and more in bins 1–2 (easy). The compute-optimal policy learned on PaLM 2-S* would not transfer to a different model, and the shape of the difficulty-dependent curves (Figure 3 right, Figure 7 right) could differ qualitatively — a more capable model might show larger benefits from search on what it considers "medium" difficulty, or might show over-optimization at different budget thresholds.

  • For EM-LLM, the human correlation result (Section 4.2) is demonstrated on three short podcasts and the PG-19 book dataset. Whether LLM surprise correlates with human event perception on non-narrative domains (code, legal text, structured data, dialogue, instruction manuals) is untested. The benchmarks (LongBench, ∞-Bench) include such domains but measure only end-task accuracy, not segmentation quality per domain.

  • The 7–8B parameter scale means all experiments are on models that are "capable but not saturated" on the tested benchmarks. A much larger model (70B, 405B) might have pass@1 high enough that test-time compute provides minimal gains on most problems, or might have attention patterns sufficiently different that the retrieval mechanisms behave differently.

What evidence exists in the paper: The MATH results are from a single model family (PaLM 2-S*; Section 4) on a single benchmark. The LongBench/∞-Bench results span five base LLMs but all are 7–8B models from the same generation of open-weight releases. The paper does not report results on different model scales within the same family (e.g., LLaMA-3-70B or LLaMA-3-405B), different model architectures (encoder-decoder vs. decoder-only), or different task domains (code generation benchmarks like HumanEval, scientific QA, dialogue). The acknowledgment that the models are "representative" is stated as belief, not demonstrated.

Mitigation status: The paper does not claim to have tested generalization across scales or domains, and the statements about representativeness are appropriately hedged as belief. However, the abstract and introduction make strong claims ("enabling them to handle practically infinite context lengths," "superior performance, consistently outperforming the state-of-the-art") that imply broader applicability than the experiments support. A more accurate framing would specify that the results are demonstrated on open-weight 7–8B decoder-only Transformers on English-language benchmarks, and that generalization to other model families, scales, languages, and domains is plausible but unverified. The paper does not provide any theoretical argument or preliminary evidence for why the findings should transfer — the human cognition inspiration suggests some universality, but the empirical link between human event perception and LLM surprise has only been shown on narrative text.

7. Implications and Future Directions

How This Work Changes the Landscape

EM-LLM introduces a conceptual reframing of how we think about context in LLMs: from a fixed-size window that must contain everything the model can access, to an unbounded memory store where only the relevant fraction enters active processing. This is not a paradigm shift in the sense of replacing the Transformer — the underlying architecture is untouched — but it is a genuine reframing of the inference-time context problem from one of compression (how to squeeze more into the window) to one of organization and retrieval (how to structure information so the right parts are accessible when needed).

The paper bridges two previously disconnected research communities — computational cognitive science (event segmentation, episodic memory, free recall) and LLM engineering (KV-cache management, retrieval-augmented generation, context window extension) — and demonstrates that the bridge is load-bearing: cognitive principles (prediction error as event boundaries, temporal contiguity in recall) translate directly into measurable improvements on standard long-context benchmarks. This is the paper's deepest conceptual contribution: it provides empirical evidence that the information structures relevant to the human brain's episodic memory system are also relevant to an LLM's attention mechanism, and that organizing context according to these structures improves retrieval. The human correlation results (Section 4.2, Figure 4) — showing that LLM surprise-based segmentation closely matches human-perceived event boundaries, and that fixed-size chunking (InfLLM's approach) performs worse than random on cohesion metrics — make this connection difficult to dismiss as metaphor. Something real in the LLM's internal representations aligns with human event perception.

What this reframing changes in practice:

  • Research priority shifts from context window extension to context organization. The paper's results challenge the prevailing assumption that the primary bottleneck for long-context LLMs is the size of the context window. For most tasks on LongBench, EM-LLM with a 4K+4K effective context (8K total tokens attended at any step) outperforms a full-context model with a 128K window (Table 9: 51.58% vs. 39.30% on LLaMA-3.1-8B). This is striking: the model with 16× less active context performs 31% better because the context it does attend to is curated and organized. The implication is that attention dilution — not context window capacity — is the dominant failure mode for long-context processing, and that retrieval-based attention with smart organization addresses it more effectively than expanding the window. This should redirect research investment: rather than pursuing ever-larger training context windows (which are exponentially expensive to train), the field should invest in better segmentation, retrieval, and memory organization strategies that make existing context windows more effective.

  • Training-free methods gain credibility as a serious alternative to fine-tuning for long contexts. Prior to EM-LLM, the dominant approaches for extending context beyond the training length required some form of fine-tuning — positional interpolation methods (YaRN, NTK-aware scaling), continued pretraining on longer sequences (LongLoRA, PoSE), or architectural modifications (linear attention, RingAttention). EM-LLM demonstrates that a training-free, inference-time architecture — applied as a drop-in modification to any pretrained model — can match or exceed these approaches on standard benchmarks while scaling to lengths (10M tokens) that no fine-tuned model currently handles. This is significant because training-free methods are universally applicable: they work with any existing model without requiring access to training data, training compute, or model weights modifications. For practitioners with a specific pretrained model and a long-context problem, EM-LLM represents a zero-training-cost solution pathway that was not previously credible.

  • The InfLLM → EM-LLM transition establishes that how you segment context matters more than whether you segment it. InfLLM already demonstrated that group-based KV retrieval outperforms per-token retrieval methods. EM-LLM shows that the segmentation method — whether chunks are fixed-size or content-adaptive — accounts for a substantial fraction of the remaining performance gap. The fixed-size segmentation that InfLLM uses (and that dominates the RAG literature) is revealed as a pessimal baseline: on the human podcast data, fixed segmentation performs worse than random on cohesion metrics (Figure 4A: F, FM, FC consistently below the random baseline on modularity, conductance, and I/IS). This is a strong negative result that should influence how the field designs chunking strategies going forward. Arbitrary boundary placement is not just suboptimal — it actively fragments semantically coherent units and degrades retrieval quality below what random chunking would achieve.

  • The contiguity buffer establishes temporal structure as a first-class retrieval signal, separate from semantic similarity. Most retrieval systems (RAG, InfLLM, dense passage retrieval) retrieve solely by topical relevance — cosine similarity between query and document embeddings. EM-LLM demonstrates that adding a second retrieval stage based on temporal adjacency (the contiguity buffer) provides complementary benefits, particularly on tasks requiring narrative or multi-step reasoning (multi-document QA, retrieval across passages). This challenges the implicit assumption that semantic similarity is the only retrieval signal worth optimizing, and opens the door for richer retrieval objectives that incorporate recency, temporal order, causal relationships, and other structural properties of the original sequence.

Reconciling prior contradictions:

The paper indirectly resolves a tension in the long-context literature that is analogous to the one resolved by the compute-optimal test-time scaling paper in its own domain. Prior work showed mixed results: some studies found retrieval-based methods (InfLLM, Memorizing Transformers) outperformed full-context models on specific tasks, while others found full-context models with expanded windows performed better on different tasks. EM-LLM's results suggest both findings were right — but on different definitions of "retrieval quality." InfLLM's fixed-size chunking provided retrieval benefits for some tasks but fragmented semantic units in ways that hurt others (narrative tasks, multi-hop reasoning). EM-LLM's dynamic segmentation improves retrieval quality across the board by adapting to content structure, closing the gap between retrieval and full-context methods on tasks where InfLLM struggled. The takeaway is that retrieval-based methods can match or exceed full-context performance, but only if the retrieval units respect the semantic structure of the content — arbitrary chunking is the bottleneck, not retrieval per se.

Research directions that become more attractive:

  • Adaptive and learned segmentation strategies: EM-LLM uses surprise as a fixed, cognitively-motivated heuristic. The strong performance of this simple heuristic — coupled with the demonstration that refinement further improves it — suggests that learning segmentation directly (via fine-tuning a boundary detector, training a dedicated segmentation model, or using reinforcement learning to optimize retrieval quality) could yield further gains. The paper makes this direction credible by establishing that better segmentation causes better downstream performance, which was not obvious before.

  • Multi-level, hierarchical memory architectures: The paper's demonstration that event-based organization improves retrieval at the level of individual events naturally extends to hierarchical organizations — events within episodes, episodes within sessions, sessions within a lifetime of interaction. The human episodic memory literature strongly suggests hierarchical structure (Baldassano et al., 2017), and EM-LLM provides the first LLM-based evidence that surprise can detect boundaries at multiple temporal scales.

  • Memory consolidation and continual learning: EM-LLM stores everything and retrieves on demand. This is efficient but lacks the consolidation mechanisms that biological memory uses to extract gists, forget irrelevant details, and integrate new information with existing knowledge. The paper's architecture provides a natural substrate for adding consolidation: periodically processing stored events to summarize, merge, or compress them, creating a multi-resolution memory that sacrifices detail for efficiency while preserving essential information.

Research directions that become less attractive:

  • Training ever-larger context windows as the primary solution. If retrieval-based methods with smart organization can match or exceed full-context models at a fraction of the computational cost, the economic case for training models with 1M+ context windows weakens substantially. This is especially true given that pretraining costs scale quadratically with context length (for standard attention), while EM-LLM's inference costs scale linearly.

  • Fixed-size chunking as a default in RAG and retrieval systems. The paper's demonstration that fixed segmentation underperforms random on cohesion metrics (Figure 4A) should give pause to any system that chunks text into fixed-size segments without considering content boundaries. The field has treated chunking as a preprocessing detail — EM-LLM shows it is a first-order determinant of retrieval quality.

Follow-Up Research This Work Enables

1. Learning to segment: training a boundary detector from LLM surprise signals, human annotations, and retrieval success.

EM-LLM uses a fixed, hand-designed segmentation pipeline: compute surprise from the LLM's next-token probabilities, threshold adaptively, refine with graph modularity. This works well on narrative text but has unknown properties on other domains and adds O(nm) computational overhead for the refinement step. A natural extension is to train a lightweight boundary detector that predicts event boundaries directly from token-level features (hidden states, attention patterns, positional information), supervised by a combination of EM-LLM's surprise signal (distillation), human event boundary annotations (where available, from datasets like Kumar et al., 2023), and downstream retrieval success (treating segmentation quality as measured by modularity or retrieval accuracy as a reward signal, trainable via policy gradient or differentiable boundary relaxation). A strong follow-up would compare a trained boundary detector against EM-LLM's surprise-based heuristic on a diverse set of domains — narrative text (PG-19, BookSum), code (The Stack, CodeSearchNet), legal text (ECHR, CUAD), and structured data (TabFact, ToTTo) — measuring both segmentation quality (modularity, conductance, Wasserstein distance to human annotations where available) and downstream performance on domain-specific long-context tasks. The key question is whether learned boundaries consistently outperform surprise-based boundaries, and whether a single trained detector generalizes across domains or requires domain-specific fine-tuning. This work is newly tractable because EM-LLM provides both a strong baseline (surprise-based segmentation) and a clear optimization objective (modularity of attention-key similarity).

2. Hierarchical event segmentation that captures nested temporal structure.

Human event perception is hierarchical: a conversation contains topics, topics contain exchanges, exchanges contain sentences (Baldassano et al., 2017; Zacks, 2020). EM-LLM currently segments at a single granularity, determined by the surprise threshold γ. The surprise signal likely contains structure at multiple scales — a large prediction error at a topic shift, smaller ones at subtopic boundaries, even smaller ones at sentence transitions. A natural extension is to construct a hierarchy of events by applying the surprise thresholding and refinement at multiple γ values simultaneously (e.g., γ = 0.5 for fine-grained, γ = 2 for coarse-grained), or by applying the segmentation recursively within each event, producing a tree-structured memory where retrieval can operate at the appropriate granularity for the query. A strong follow-up would implement such a hierarchy, evaluate it on tasks requiring multi-scale reasoning (e.g., book-level summarization where the model must retrieve chapter-level events and paragraph-level details), and measure whether hierarchical retrieval improves efficiency (retrieving coarse events first, then drilling into fine-grained sub-events only for the most relevant ones) without sacrificing accuracy. The connection to human event perception (Figure 4) suggests that hierarchical segmentation may better match human memory organization, potentially improving performance on narrative understanding tasks. This is tractable because EM-LLM's surprise-based boundary detection naturally produces boundaries at different thresholds, and the modularity-based refinement can be applied at each level independently.

3. Stress-testing the surprise signal: when does LLM surprise fail to detect meaningful event boundaries, and what are the alternatives?

The paper demonstrates surprise-based segmentation works on narrative text (podcasts, books), but its performance on non-narrative domains is entirely uncharacterized. Critical stress tests include: code (where surprise might spike at rare tokens or syntax errors rather than semantic boundaries like function definitions), legal/regulatory text (where formulaic, repetitive structure might produce uniformly low surprise despite clear section boundaries), multilingual text (where tokenization artifacts and varying model confidence across languages might distort the surprise signal), and adversarial text (where an attacker could craft sequences that produce misleading surprise patterns). A strong follow-up would evaluate EM-LLM with surprise-based segmentation on a systematically varied set of text types, measuring both segmentation quality (modularity vs. fixed-size baselines) and downstream task performance, and identifying the domains where surprise-based segmentation fails to outperform or even underperforms fixed-size chunking. For failure domains, the follow-up would test alternative segmentation signals: attention-pattern change, hidden-state change between consecutive tokens, or an external segmentation model (e.g., a fine-tuned BERT for boundary detection). The outcome would be a domain-to-segmentation-method mapping that guides practitioners on when to trust surprise-based segmentation and when to fall back to alternatives. This is tractable because LongBench and ∞-Bench already include diverse text types (academic papers, code, meeting transcripts, legal documents), and the paper's segmentation quality metrics (modularity, conductance, I/IS) provide a domain-agnostic evaluation framework that does not require human annotations.

4. Combining EM-LLM with positional encoding extension methods to recover temporal order in retrieved events.

EM-LLM assigns fixed position embeddings to all retrieved events, discarding their absolute and relative temporal positions (Section 3.4). This prevents out-of-distribution positional encoding failures but loses information about event order and temporal distance that could be crucial for tasks requiring temporal reasoning ("what happened between event A and event B?," "did X occur before Y?"). The contiguity buffer partially addresses this by maintaining local temporal context, but it does not provide global temporal ordering across all retrieved events. A promising extension is to combine EM-LLM's event-based retrieval with a positional encoding extension method (e.g., NTK-aware RoPE scaling, positional interpolation, or YaRN) applied selectively to retrieved events. Rather than assigning all retrieved events the same fixed position, assign them approximate positions based on their original temporal indices (scaled down to fit within the model's training range), or assign relative positional offsets between retrieved events while keeping absolute positions within the training window. A strong follow-up would compare EM-LLM with and without positional information on tasks explicitly requiring temporal reasoning (e.g., temporal QA tasks from TimeQA or TempReason, event ordering tasks, narrative cloze tests), measuring whether recovered positional information improves accuracy without reintroducing the attention dilution or out-of-distribution issues that retrieval avoids. This is tractable because the positional encoding extension literature provides well-characterized methods (YaRN, NTK scaling, positional interpolation) that can be applied as a lightweight inference-time modification, and EM-LLM's architecture cleanly separates retrieved events from local context, making selective position assignment straightforward.

5. Dynamic retrieval budget allocation: scaling the similarity and contiguity buffers with context length and query complexity.

EM-LLM uses a fixed retrieval budget (e.g., 2K or 4K retrieved tokens) regardless of total sequence length or query complexity. The paper's own ablation (Appendix D.3, Figure 14) hints that this is suboptimal: on QMSum, longer contexts benefit from larger retrieved buffers, while on other tasks, the relationship is flat or negative. This suggests an adaptive policy: allocate more retrieved tokens to longer sequences or more complex queries, and fewer to shorter or simpler ones. A natural extension is a compute-adaptive retrieval policy that dynamically adjusts the similarity buffer size k_s based on (a) total context length (more context → more retrieval budget, to maintain coverage), (b) query complexity as estimated by the LLM's own uncertainty or the number of distinct events retrieved by different layers (high diversity of retrieved events across layers suggests a complex query requiring more context), and (c) the contiguity ratio k_r based on task type or detected narrative structure (higher contiguity for narrative/sequential tasks, lower for factual retrieval). A strong follow-up would implement this adaptive policy, compare it against fixed-budget EM-LLM on LongBench and ∞-Bench with varied context lengths, and measure whether the adaptive policy recovers the performance of the per-task optimal fixed budget without requiring per-task tuning. The key metric is whether dynamic allocation closes the gap between the average-case k_r = 0.3 configuration and the per-task optimal configurations shown in Figure 13. This is tractable because context length is known at inference time, and per-layer retrieval diversity is computable from the retrieval step itself (Appendix Figure 5 provides the visualization framework).

6. EM-LLM as a computational model of human episodic memory: testing predictions from cognitive science in LLM attention patterns.

The paper's human correlation analysis (Section 4.2) demonstrates that LLM surprise correlates with human-perceived event boundaries — but this is a correlation at the group level (averaged across participants and podcasts). A deeper question is whether EM-LLM's retrieval mechanisms exhibit the same individual-level memory phenomena documented in the cognitive psychology literature: the serial position effect (better recall of items at the beginning and end of an event), the temporal contiguity effect (Figure 3A), the asymmetry effect (forward vs. backward recall probability), and the event boundary advantage (better memory for items near event boundaries; Michelmann et al., 2023a). A strong follow-up would design controlled experiments inspired by human free recall paradigms: present EM-LLM with structured sequences (lists of items, narrative stories with embedded facts), manipulate event boundary positions (using surprise-based segmentation vs. fixed boundaries vs. random boundaries), and measure whether the model's retrieval accuracy and attention patterns reproduce the characteristic biases of human memory (e.g., higher attention to boundary-adjacent tokens, contiguity effects in retrieval order, primacy and recency within events). The outcome would be a computational model of episodic memory phenomena implemented in a working LLM system, which could test cognitive theories (e.g., the Event Horizon Model, contextual drift models) at a scale impossible with human subjects (millions of tokens, precise control over event structure). This is tractable because EM-LLM provides explicit event boundaries, a contiguity buffer that implements temporal context, and a retrieval mechanism that can be instrumented to measure attention patterns — the cognitive phenomena are directly measurable in the model's behavior.

Practical Applications and Downstream Use Cases

1. Long-document analysis and question-answering over complete knowledge bases in resource-constrained settings.

EM-LLM is designed for scenarios where the context far exceeds what can fit in GPU memory or the model's training window. A concrete deployment: a legal document review system that must answer questions across thousands of pages of case law, contracts, and regulatory filings. With a 7B-parameter model like LLaMA-3.1-8B and EM-LLM's memory management (Appendix C.3.2), the full document corpus is stored as organized episodic events, and queries retrieve only the relevant events into a 4K+4K context window. The system requires minimal GPU memory (a single 32GB GPU for the model, plus CPU/disk for the KV cache), making it deployable on hardware that cannot run even a 32K-token full-context model. The paper's results suggest this would outperform a full-context approach: on LongBench's multi-document QA tasks (HotpotQA, 2WikiMQA, Musique), EM-LLM achieves up to 10.45% improvement over InfLLM and significantly outperforms full-context models (Table 9: 54.02% vs. full-context's 54.01% on HotpotQA for LLaMA-3.1, with larger gaps on other tasks). The 10M-token passkey retrieval result (Figure 1, bottom) provides the scalability guarantee that the system will not degrade as the document corpus grows — the computational cost per query remains constant.

2. Continuous, long-running conversational agents with persistent memory.

A second deployment: personal assistant LLMs that maintain conversation history across weeks or months of interaction, accumulating context far beyond any practical context window. EM-LLM's event-based storage naturally handles this: each conversation session or topic shift becomes an event, stored and retrievable when future queries reference past interactions. The contiguity buffer enables the model to maintain temporal context around retrieved memories (e.g., retrieving not just the specific past conversation where the user mentioned a preference, but the surrounding context of that conversation). The similarity buffer ensures new queries find relevant past events even if they occurred months ago. The training-free nature of EM-LLM is particularly valuable here: the LLM can be updated or replaced without rebuilding the memory store, and the memory persists independently of the model. The paper's results on LongBench's few-shot learning tasks (TREC: EM-LLM 71.50% vs. RAG 22.50% vs. full-context 4.50%; Table 9) demonstrate EM-LLM's ability to retrieve and utilize scattered information from long contexts — a capability directly relevant to maintaining coherent conversation history where relevant facts are distributed across many past sessions.

3. Codebase understanding and repository-scale reasoning without fine-tuning.

A third deployment: code assistants that must reason over entire codebases (hundreds of files, millions of tokens) to answer questions, implement features, or debug issues. EM-LLM's per-layer retrieval is well-suited to code because different attention heads can specialize in different aspects of code structure — some retrieving type definitions, others retrieving function implementations, others retrieving usage examples. The paper's results on LongBench's code tasks show EM-LLM's strong advantage over alternatives: on LCC, EM-LLM achieves 67.45% vs. full-context's 19.30% vs. RAG's 13.16% (Table 9, LLaMA-3.1-8B); on RepoBench-P, EM-LLM achieves 64.33% vs. full-context's 18.33% vs. RAG's 18.66%. These are not incremental improvements — EM-LLM more than triples the performance of full-context and RAG on code tasks, suggesting that selective retrieval of relevant code sections (function definitions, import statements, class hierarchies) is dramatically more effective than exposing the model to the entire codebase. The training-free property means developers can use their preferred code LLM (CodeLLaMA, DeepSeek-Coder, StarCoder) without modification, applied to any codebase.

4. Processing and reasoning over multi-modal interaction histories.

A fourth deployment enabled by EM-LLM's architecture: multi-modal memory for applications that accumulate images, audio transcripts, sensor readings, and text over time. The paper does not evaluate multi-modal inputs, but its architecture treats embeddings generically — the KV cache stores key-value pairs from any modality, and the surprise-based segmentation operates on the LLM's next-token predictions regardless of modality. A multi-modal deployment (e.g., a robot accumulating visual observations and task instructions over hours of operation, or a medical record system integrating imaging reports, clinical notes, and lab results over a patient's history) would organize cross-modal events, retrieve relevant multi-modal context for each new query, and maintain temporal coherence through the contiguity buffer. The paper's efficiency claims — linear scaling in sequence length, constant per-query cost — make this feasible for interaction histories that would be computationally prohibitive for full-context multi-modal models. The per-layer retrieval property is particularly relevant for multi-modal inputs, where different layers might specialize in different modalities or different aspects of cross-modal integration.