Focused Transformer: Contrastive Training for Context Scaling

🎯 Pitch

Standard long-context models get 'distracted' as irrelevant documents pile up—their attention flatlines across everything. By training attention layers to actively repel keys from unrelated documents, LONGLLAMA extrapolates to 256k tokens, hitting 73% passkey retrieval while fine-tuned on just 8k context.

1. Executive Summary

This paper introduces the Focused Transformer (FOT), a method that fine-tunes language models to handle longer contexts by improving the structure of the (key, value) space in a subset of attention layers. The core mechanism is a crossbatch training procedure inspired by contrastive learning — during training, selected attention layers are exposed to keys from both the current document (positive examples) and other unrelated documents (negative examples), incentivizing the model to distinguish relevant from irrelevant information and thereby addressing what the authors term the distraction issue (where attention mass becomes uniformly spread across relevant and irrelevant documents as the number of documents grows). The approach is demonstrated by fine-tuning 3B and 7B OpenLLaMA checkpoints, producing models called LONGLLAMA that achieve 94.5% accuracy on passkey retrieval at 100k tokens and 73% at 256k tokens, despite being fine-tuned with only an 8k context length. The method enables meaningful gains on few-shot in-context learning tasks like TREC (improving from 67.0% to 73.3% accuracy for the 3B model when expanding context from 2k to 8k tokens), while establishing that this extrapolation capability emerges only when the crossbatch training exposes the model to negative keys from unrelated documents — a standard fine-tuning baseline without this contrastive mechanism fails to utilize contexts beyond its training length.

2. Context and Motivation

The Core Problem: Language Models Cannot Effectively Use Long Contexts

The fundamental capability this paper targets is in-context learning at scale — the ability of a language model to incorporate new information by simply including it in the prompt, without any weight updates. This capability is enormously valuable for applications ranging from repository-level code generation (where the model needs awareness of thousands of lines of code across hundreds of files) to long-document question answering (e.g., answering questions about an entire research paper or book) to multi-document reasoning (synthesizing information across many sources). Yet this capability is severely bottlenecked by what the paper calls the effective context length — the maximum number of tokens over which the model can reliably attend to and utilize information.

The paper introduces a specific diagnosis of why extending context length is hard: as the number of documents or information sources in the context grows, the ratio of relevant to irrelevant tokens diminishes, and standard training fails to equip the model with the ability to distinguish between them. The paper gives this phenomenon a name — the distraction issue — and argues it is a primary obstacle that prior work has not directly addressed.

This matters because in-context learning is fundamentally more practical than fine-tuning for incorporating new knowledge. Fine-tuning requires collecting and curating training data, managing compute resources, and dealing with issues like catastrophic forgetting and distribution shift. It also operates indirectly — fine-tuning on a text teaches the model to predict tokens from that distribution, not to explicitly answer questions about that text. In-context learning bypasses all of this: you simply put the information in the prompt and ask your question. The barrier is purely architectural: existing models cannot handle prompts long enough to contain all the information relevant to complex tasks.

The Distraction Issue: A Precise Diagnosis

The paper's central conceptual contribution is to isolate and formalize the distraction issue. Here is what happens, concretely:

In a standard transformer, when a query token attends over a set of keys, the softmax operation distributes attention mass across all available keys. If the set of keys contains entries from many different documents, and the model has not been trained to discriminate between document boundaries, the attention mass becomes uniformly distributed across relevant and irrelevant documents. The paper measures this directly: they define positive attention mass as the fraction of attention that falls on keys from the same document as the query, and find that for a standard transformer, this value is approximately $1/d$ where $d$ is the number of documents in context. In other words, the model is equally distracted by every document — it treats a key from an unrelated document identically to a key from the current document.

This is illustrated in Figure 3. A standard Transformer (blue line) shows $r_d \approx 1/d$ — the positive attention mass decreases exactly in proportion to the number of documents, meaning the model learns nothing about which keys are relevant. The consequence is that as memory scales, the attention signal from genuinely useful keys gets drowned out by noise from irrelevant keys. The model's effective context length hits a ceiling not because of architectural limitations per se, but because the key-value representations lack the structure to support discrimination between semantically similar and dissimilar sources.

The paper explicitly connects this to a known phenomenon in language model representation spaces: the degeneracy problem identified by Gao et al. (2019), where embeddings become tightly packed in a narrow cone, making cosine similarities uniformly high and discrimination difficult. The distraction issue can be understood as a manifestation of this degeneracy in the context of multi-document attention: keys from different documents are too similar in representation space for the attention mechanism to meaningfully differentiate them.

Why Prior Approaches Fall Short

Architectural Efficiency Approaches (Sparse Attention, Compression)

A large body of work addresses the computational efficiency of long contexts — the quadratic complexity of full self-attention. Examples include Transformer-XL (Dai et al., 2019), which caches previous segments and attends across segment boundaries; Longformer (Beltagy et al., 2020) and BigBird (Zaheer et al., 2020), which use sparse attention patterns to reduce computation; hierarchical approaches like Hourglass Transformers (Nawrot et al., 2021) that downsample intermediate activations; and COLT5 (Ainslie et al., 2023) which applies conditional computation to save memory.

Where they fall short: These methods solve the computational bottleneck — they make it feasible to process long sequences — but they do not solve the representational bottleneck. Making attention sparse or hierarchical does not, by itself, teach the model which sparse connections are useful or how to structure key-value representations so that relevant information can be found. As the paper notes, the distraction issue persists regardless of how efficiently you compute attention — a sparse attention mechanism that spreads its limited budget uniformly across documents will be just as ineffective as dense attention. Kaddour et al. (2023) is cited as showing that efficiency-focused approaches lead to limited gains, suggesting the bottleneck is not (only) computational.

Positional Encoding Extensions (Position Interpolation, Landmark Attention)

More recently, several works have approached context length extension through modifications to positional encodings. Position Interpolation (Chen et al., 2023; kaiokendev, 2023) rescales rotary position embeddings so that positions beyond the training length are mapped into the range the model was trained on, enabling fine-tuning for up to 32k contexts. Landmark Attention (Mohtashami and Jaggi, 2023) compresses context into landmark tokens that serve as retrieval anchors, extending LLaMA-7B to 32k tokens.

Where they fall short: These methods are fundamentally bounded by the positional encoding scheme — they can extend context but only by reinterpreting or compressing position information. The paper explicitly contrasts its approach: "our method does not rely on positional encodings, following the findings from Haviv et al. (2022)." By removing positional encoding from the additional context (memory attention layers use no positional information), FOT achieves theoretically unbounded context length because there is no positional encoding scheme to saturate or extrapolate. Empirically, this translates to 256k token passkey retrieval — far beyond the 32k ceiling of position-based methods.

Moreover, position-based methods don't address the distraction issue at all — they make it possible to attend to more tokens, but don't ensure the model uses that attention effectively. The paper shows that standard long-context fine-tuning (even with the same architecture and context length) fails to extrapolate beyond its training length (Table 2, Figure 1), while FOT succeeds — the difference is not the architectural capacity but the representational structure of the key-value space, which FOT explicitly optimizes.

Memory-Augmented Transformers (Memorizing Transformer, kNN-LM)

The most directly relevant prior work is the Memorizing Transformer (MT) (Wu et al., 2022), which FOT extends. MT augments a subset of attention layers with access to an external memory of (key, value) pairs, retrieved via k-nearest neighbors (kNN) lookup. At inference time, each query in a memory layer attends to both the local context and the top-k matching keys retrieved from memory. This allows the model to incorporate far more context than fits in the local attention window.

Where MT falls short: The paper identifies two critical limitations. First, MT's training procedure is not differentiable through the memory keys and values. During training, MT retrieves keys from a non-differentiable memory store — the keys and values used for retrieval are frozen snapshots, and gradients only flow through the local context. This means the model cannot jointly optimize its key, query, and value representations for the retrieval task. The key structure is whatever the base model happened to learn, not what would be optimal for retrieval.

Second, and more fundamentally, MT was designed for and trained on single-document scenarios — its memory contains only tokens from the current document. It was never exposed to the distraction issue because all keys in its memory are relevant by construction (they all come from the same document). When you deploy MT in a multi-document setting, where the memory contains keys from many unrelated documents, the model has never learned to distinguish relevant from irrelevant keys. The distraction issue hits MT just as hard as a standard transformer.

FOT addresses both limitations simultaneously: the crossbatch training is fully differentiable through all (key, value) pairs (both positive and negative), and the inclusion of negative examples from unrelated documents directly trains the model to handle distractions. The ablation in Figure 5 shows the concrete impact: MT and FOT both trained with 512-token local context, but FOT with $d=8$ achieves substantially better perplexity in multi-document evaluation, with the gap widening as memory grows.

The paper also contrasts with kNN-LM (Khandelwal et al., 2019), which interpolates between the base LM's output distribution and a distribution derived from nearest neighbors in an external datastore. kNN-LM is a zero-shot method — it modifies the output probabilities without any training. The paper positions FOT differently: "we extend the model context in a subset of attention layers, potentially allowing for reasoning within this extended context." The distinction is that kNN-LM influences only the final token probabilities, while FOT integrates retrieved information into the internal attention computation of specific layers, enabling the model to reason over the retrieved information rather than just adjust its output distribution.

Contrastive Learning for Language Models

The paper draws explicit inspiration from contrastive learning, particularly the insight that exposing models to negative examples helps them learn discriminative representations. In vision, SimCLR (Chen et al., 2020) and CLIP (Radford et al., 2021) demonstrated that contrasting positive pairs against negative pairs yields state-of-the-art representations, with larger batch sizes (more negatives) improving quality (Gao et al., 2021b).

In language modeling, TRIME (Zhong et al., 2022) is the closest prior work. TRIME trains language models with memory augmentation by incorporating negative examples to improve representation quality for retrieval. However, the paper identifies key differences: "we incorporate negatives into the chosen subset of attention layers instead of interpolating in the output layer and use the standard language modeling loss." TRIME modifies the output probability distribution, while FOT modifies the internal attention mechanism. TRIME also focuses on retrieval from large databases, while FOT focuses on extending context.

ContraCLM (Jain et al., 2023) applies contrastive losses at both token and sequence levels to promote more uniformly distributed, isotropic representations, improving performance on semantic similarity benchmarks. The paper acknowledges this as related but positions FOT differently: "While ContraCLM focuses on improving the general expressiveness of representations, our work introduces contrastive-inspired techniques designed specifically for training the attention mechanism to handle longer context lengths."

The paper's key insight is that contrastive learning is not just a general-purpose representation improvement tool — it is specifically the right tool for addressing the distraction issue, because the distraction issue is fundamentally a discrimination failure: the model cannot tell relevant keys from irrelevant ones. Contrastive learning directly trains this discrimination by forcing the model to distinguish positive keys (same document) from negative keys (different documents) during training.

How This Paper Positions Itself

The paper occupies a specific, previously unfilled position in the long-context landscape. It is:

Not an architecture paper: FOT adds no new parameters beyond the base transformer (the memory attention layer uses the same attention mechanism, just with expanded key-value sources). The contribution is a training method, not an architectural innovation.
Not an efficiency paper: While FOT uses kNN retrieval for inference-time efficiency, the paper does not claim computational efficiency gains as a primary contribution. The focus is on representational quality — making the model actually use long contexts effectively, not making it cheap to process them.
Not a from-scratch training method: The paper explicitly demonstrates that FOT can be applied to existing, pre-trained models (OpenLLaMA 3B and 7B) through fine-tuning, making it a plug-and-play extension rather than a recipe requiring full retraining.
A representational solution to a representational problem: The distraction issue is a problem of key-value space structure. The crossbatch training procedure is a contrastive-inspired method for shaping that structure. The paper's framing is that the bottleneck in context scaling is not computational or architectural but representational, and that contrastive training is the natural remedy.

The paper positions itself as building directly on the Memorizing Transformer (same memory layer design, same kNN retrieval mechanism) but fundamentally rethinking the training objective. It also positions itself as complementary to positional encoding methods — FOT achieves its extrapolation by removing positional information from the additional context rather than interpolating it, making the approaches potentially combinable (though the paper does not explore this combination).

A crucial distinction the paper draws is between single-document context scaling (the typical benchmark, where all context comes from one long document) and multi-document context scaling (where context contains many unrelated documents, as in repository-level code or multi-source QA). The paper argues that multi-document scaling is the more realistic and challenging setting because it introduces the distraction issue, and that most prior work has implicitly or explicitly focused on single-document scenarios where all tokens in context are at least potentially relevant. FOT's crossbatch training is specifically designed to handle the multi-document case, and the paper shows that this training generalizes to improve single-document performance as well (Section 5.4) — an example of a harder training objective yielding better representations for an easier evaluation setting.

3. Technical Approach

3.1 Reader orientation (approachable technical breakdown)

The Focused Transformer (FOT) is a training procedure — not a new architecture — that you apply to an existing transformer language model during fine-tuning to make its internal key-value representations more discriminative, so that when you later give the model access to a much larger context at inference time (via a memory of previously seen key-value pairs), it can effectively focus on the relevant parts and ignore the irrelevant ones. The system solves the distraction issue: the problem that a standard transformer's attention mechanism, when faced with keys from many different documents, spreads its attention mass uniformly across all documents regardless of relevance, because it was never trained to tell the difference between "this key is from my current document" and "this key is from some random other document." The solution is surprisingly simple in concept: during fine-tuning, deliberately show the model's attention layers both relevant keys (from the same document, earlier in the context) and irrelevant keys (from other unrelated documents, pulled from other positions in the same training batch), and let gradients flow through all of them, so the model learns to embed keys in a space where similarity corresponds to relevance.

3.2 Big-picture architecture (diagram in words)

The FOT system has four major components that interact during training and inference:

A base transformer language model (e.g., OpenLLaMA 3B or a custom 184M-parameter decoder-only model) — this is the model being fine-tuned or trained. Most of its layers operate exactly as in a standard transformer. A designated subset of layers (e.g., layer 8 in the 12-layer models) are memory attention layers.
The memory attention layer(s) — during inference, these layers have access to an external memory of (key, value) pairs that have been accumulated from all previous tokens the model has processed (not just the current document or the current local context window). For each query, the layer retrieves the $k$ most similar keys from this memory via exact k-nearest-neighbors (kNN) search, and the query attends to both the local context tokens and these $k$ retrieved memory tokens. During training, the memory is not used; instead, a crossbatch mechanism provides the additional context.
The crossbatch training data pipeline — this is the core innovation. The training corpus is organized so that each position in a batch corresponds to a different document. For a given document, the model sees (a) the current local context window (the tokens being predicted), (b) the previous local context window from the same document (positive keys — relevant), and (c) the previous local context windows from $d-1$ other documents in the batch (negative keys — irrelevant). All of these are concatenated into the key-value set that the memory attention layer attends to during the forward pass, meaning gradients flow through the keys and values of all $d$ documents.
The inference-time kNN memory — after training, the crossbatch mechanism is replaced by a standard external memory. As the model processes tokens, the memory attention layer stores its (key, value) outputs into a FAISS index. For each subsequent query, the top- $k$ most similar keys are retrieved and their associated values are included in the attention computation alongside the local context.

Information flows as follows at training time: a batch of documents is loaded $\rightarrow$ for each document, the current local context and previous local context are embedded $\rightarrow$ for each document, the memory attention layer constructs an expanded key-value set by concatenating the local context's keys with the previous context's keys from the same document (positives) and $d-1$ other documents (negatives) $\rightarrow$ the standard transformer forward pass proceeds, with gradients flowing through all keys and values in this expanded set $\rightarrow$ the standard language modeling loss (next-token prediction cross-entropy) is computed and backpropagated.

At inference time: a document is processed token by token $\rightarrow$ at the memory attention layer, each (key, value) output is appended to the FAISS memory $\rightarrow$ for each query, the top- $k$ nearest keys are retrieved from memory $\rightarrow$ the query attends to both the local context keys and the $k$ retrieved memory keys $\rightarrow$ the output value is computed as the softmax-weighted sum of the associated values $\rightarrow$ the rest of the forward pass proceeds normally.

3.3 Roadmap for the deep dive

First, the distraction issue — how it is formally defined and measured, because understanding the problem is essential to understanding why the solution takes the specific form it does.
Second, the memory attention layer — how it integrates additional context during inference, because this is the mechanism that FOT's training procedure is designed to optimize for.
Third, the crossbatch training procedure — the core technical contribution, including exactly how the key-value sets are constructed, how $d$ controls the ratio of positives to negatives, and why differentiability through all keys matters.
Fourth, the specific design choices and hyperparameters — the concrete numbers and configurations used in the experiments, including model sizes, layer selections, learning rates, and the schedule for increasing $d$ during training.
Fifth, the application to LONGLLAMA — the specific modifications made to fine-tune OpenLLaMA checkpoints (including how positional encodings are handled and why kNN retrieval is replaced with dense attention for these models).

3.4 Detailed, sentence-based technical breakdown

This is primarily a training methodology paper whose core idea is that the bottleneck preventing transformers from effectively using long contexts is not computational cost but rather the representational quality of the key-value space, and that this quality can be dramatically improved by a contrastive-inspired training procedure that exposes attention layers to both relevant and irrelevant keys during training, with full differentiability through all of them.

The Distraction Issue — Formal Definition and Empirical Measurement

The distraction issue is the paper's diagnosis of why extending context length is hard, and it is defined with precise mathematical formalism in Section 3.3. The setup is as follows: consider a memory-augmented transformer where a specific attention layer is exposed to additional context beyond the local window. When the model processes multiple documents, the memory accumulates (key, value) pairs from all of them. The question is: how does the model's attention distribute itself across keys from the current document versus keys from other documents?

The paper formalizes this using the concept of positive attention mass. For a document $\delta$ , let the memory attention layer be exposed to a set of (key, value) pairs $\{p^\delta_i\}_{i = 1, \dots, d}$ , where $i = 1$ indexes pairs from the previous local context of the same document $\delta$ (the positives), and $i = 2, \dots, d$ index pairs from the previous local contexts of $d-1$ other documents (the negatives). Let $w_{ij}$ be the softmax attention weight assigned to the $j$ -th token in $p^\delta_i$ (the set from document $i$ ). The positive attention mass $r_d$ is defined as:

$r_d := \sum_j w_{1j} \Big/ \sum_{i=1}^d \sum_j w_{ij}$

where $w_{ij}$ is the softmax weight for the $j$ -th token in $p^\delta_i$ , the numerator sums attention weights over all tokens from the current document's previous context (positives), and the denominator sums attention weights over all tokens from all $d$ documents in the expanded context.

What it computes: the fraction of the total attention mass (summing to 1 by the softmax property) that falls on keys from the same document as the query, as opposed to keys from other documents. If the model is perfectly focused, $r_d \approx 1$ (all attention on the relevant document). If the model is perfectly distracted, $r_d \approx 1/d$ (attention uniformly distributed across all $d$ documents).

Why this form: the ratio $r_d$ is a direct normalized measure of selectivity. By comparing it to the uniform baseline $1/d$ , we can quantify how much better (or worse) than chance the model's attention allocation is. The denominator normalizes by the total mass, so $r_d$ is always in $[0, 1]$ regardless of the number of tokens per document, and it is directly comparable across different values of $d$ .

The empirical finding, shown in Figure 3, is stark: for a standard transformer (no crossbatch training, labeled d=1 in the figure but meaning standard training without negatives), $r_d$ tracks $1/d$ almost perfectly. As the number of documents $d$ increases from 1 to 64, $r_d$ drops from approximately 1.0 to approximately $1/64 \approx 0.016$ . The model is not just somewhat distracted — it is exactly as distracted as random chance would predict. It has learned nothing about which keys are relevant because it was never given a training signal that would incentivize such learning.

In contrast, FOT trained with crossbatch ( $d=2$ then increased to $d=64$ , the green line in Figure 3) maintains $r_d \approx 0.8$ even at $d=64$ — it dedicates roughly 80% of its attention mass to the relevant document and only 20% spread across the 63 irrelevant ones. The model has learned to discriminate. FOT trained with $d=8$ (the orange line) shows intermediate behavior, with $r_d$ declining from approximately 0.9 at $d=1$ to approximately 0.35 at $d=64$ , indicating that the number of negatives seen during training directly controls the robustness of the learned discrimination.

The paper connects this to the known phenomenon of representation degeneration in language models (Gao et al., 2019), where embeddings become concentrated in a narrow cone, making cosine similarities between any two embeddings uniformly high. The distraction issue is a specific manifestation: if all keys are in a narrow cone, then the inner products $\langle q, \text{key} \rangle$ are nearly identical for all keys, the softmax is nearly uniform, and attention becomes non-selective. The crossbatch training procedure directly counteracts this by creating a training objective that rewards the model for making keys from different documents distinguishable and penalizes it when keys from the same document are not distinguishable from keys from other documents.

The Memory Attention Layer — Inference-Time Mechanism

The memory attention layer is the architectural component that FOT's training procedure is designed to optimize. It is closely based on the Memorizing Transformer (Wu et al., 2022) but with one key simplification: the gating mechanism.

In a standard transformer layer, each query $q$ computes attention over keys from the local context $C_{<q}$ (all tokens preceding $q$ in the current sequence), producing an output value as the softmax-weighted sum of the associated values:

$v = \sum_{(\text{key}, \text{val}) \in C_{<q}} s(\text{key}) \cdot \text{val}$

where $s(\text{key})$ is the softmax score for that key, computed over all keys in $C_{<q}$ .

In a memory attention layer, the set of keys is expanded to include both the local context $C_{<q}$ and the top- $k$ most similar keys retrieved from an external memory $M$ . The memory $M$ is a growing store of all (key, value) pairs that this specific layer has produced for all previous tokens across all documents processed so far (unless explicitly cleared, as in single-document evaluation mode). Formally, let $M$ be the set of all previously stored pairs, and let $M_{\text{top}}$ be the subset of $k$ pairs that maximize the inner product with the query:

$M_{\text{top}} \subset M \quad \text{such that} \quad |M_{\text{top}}| = k \quad \text{and} \quad \langle q, \text{key} \rangle \text{ is maximized for pairs in } M_{\text{top}}$

The attention output is then computed over the union:

$v = \sum_{(\text{key}, \text{val}) \in M_{\text{top}} \cup C_{<q}} s(\text{key}) \cdot \text{val}$

where $s(\text{key})$ are softmax scores computed with a learnable temperature $\tau$ :

$s(\text{key}) = \text{softmax}\left( \frac{\langle q, \text{key} \rangle}{\tau} \right)$

where the softmax is taken over all keys in $M_{\text{top}} \cup C_{<q}$ , and $\tau$ is a per-head learnable scalar that controls the sharpness of the attention distribution.

What it computes: for each query at this layer, an attention-weighted combination of local context values and the values associated with the $k$ most similar keys retrieved from the external memory. The local context provides the standard autoregressive information, while the memory provides long-range retrieval from potentially millions of earlier tokens. The temperature $\tau$ allows the model to learn how "sharp" or "flat" the attention distribution should be in this layer, balancing the influence of the most similar retrieved key against the need to consider multiple candidates.

Why this form: the paper compares this simple concatenation approach to the gating mechanism used in the original Memorizing Transformer (Wu et al., 2022, Equation 2). In the gating approach, the memory output $v_M$ and local output $v_C$ are computed separately and then blended via a learned gate $g = \sigma(b_g)$ :

$v = v_M \cdot g + v_C \cdot (1 - g)$

The paper reports (Figure 4 in Appendix) that the simple concatenation approach performs equivalently to gating while being simpler and requiring no additional parameters. The paper attributes this to the crossbatch training procedure: "the crossbatch training backpropagates through the (key, value) pairs from the previous context $C_{\text{prev}}$ in contrast to MT that cannot backpropagate there and needs to rely on local context when computing gradients for keys and values." In other words, because FOT's training makes the memory keys and values fully differentiable, the model can learn to balance memory and local context through the key and query representations themselves, without needing an explicit gating parameter.

The kNN retrieval is implemented using FAISS (Johnson et al., 2017) with exact search — not approximate nearest neighbors — meaning the retrieval is lossless. The paper uses $k = 128$ for most experiments (the number of top keys retrieved from memory). For the LONGLLAMA models (Section 4), the kNN retrieval is replaced with dense attention over the full memory, because the authors found "only marginal performance differences and it is simpler to implement."

A critical design choice: no positional encodings are used in the memory attention layers for most FOT models. The reasoning follows Haviv et al. (2022), who showed that transformers can learn positional information without explicit positional encodings. Removing positional encodings from the additional context means there is no positional encoding scheme to saturate — the memory can theoretically grow unbounded, and the model's ability to use it is limited only by the representational quality of the keys, not by any architectural ceiling. The paper explicitly contrasts this with position-interpolation methods: "our method does not rely on positional encodings... Removing positional encoding in additional context allows us to extrapolate to 256k tokens, although the model was only trained on sequences up to 8K, yielding theoretically unbounded context length."

For the LONGLLAMA models specifically, positional encodings are retained in the local context (to maintain backward compatibility with the original LLaMA architecture), but memory keys are encoded as if they were at position 0 in the local context window — a simple fixed encoding that eliminates dependence on the absolute position of the memory tokens.

The memory is populated incrementally: as the model processes each token, the memory attention layer's (key, value) outputs for that token are appended to the memory store. In single-document evaluation mode, the memory is cleared at the start of each new document, so the model only attends to earlier tokens from the same document. In multi-document evaluation mode, the memory persists across documents, so the model must distinguish relevant from irrelevant keys on the fly — this is the harder and more realistic setting that FOT's crossbatch training is specifically designed for.

The Crossbatch Training Procedure — Core Technical Innovation

The crossbatch training procedure is the paper's central technical contribution. It replaces the non-differentiable memory retrieval used by the Memorizing Transformer during training with a fully differentiable expanded key-value set drawn from multiple documents within the same batch. This is described in Section 3.2 and illustrated in Figure 2.

The training data pipeline is organized as follows:

Each position in the training batch corresponds to a different document. This is crucial: the batch is not randomly shuffled at the token level; rather, each batch index is dedicated to a single document for the duration of that document's processing. The paper states: "the batch index occupied by each document is fixed from the moment we load the document till we finish processing it."
For each document in the batch, two consecutive context windows are embedded: the current local context $C_{\text{curr}}$ (the tokens currently being processed and predicted) and the previous local context $C_{\text{prev}}$ (the immediately preceding window of tokens from the same document). The previous context serves as the source of longer-range in-document information.
For a document $\delta$ at batch index $i$ , the memory attention layer constructs an expanded key-value set $p^\delta$ that contains: all tokens from the previous local context of document $\delta$ (the positives), and all tokens from the previous local contexts of documents at batch indices $i+1, i+2, \dots, i+d-1$ (modulo batch size) — these are the negatives.
The attention computation for a query $q$ from document $\delta$ is then:

$v = \sum_{(\text{key}, \text{val}) \in p^\delta \cup C^{\delta, <q}_{\text{curr}}} s(\text{key}) \cdot \text{val}$

where $C^{\delta, <q}_{\text{curr}}$ is the set of (key, value) pairs from the current local context that precede $q$ , and $s(\text{key})$ is the softmax score computed with temperature $\tau$ over all keys in the combined set.

where $p^\delta$ is the expanded set of (key, value) pairs from $d$ documents (one positive, $d-1$ negative), $C^{\delta, <q}_{\text{curr}}$ is the local context preceding query $q$ in the current window, and $s(\text{key})$ is the softmax score as defined above.

What it computes: the attention output in the memory layer is the softmax-weighted combination of values from three sources: immediate local context (standard autoregressive attention), the recent history of the same document (positive — these are tokens from slightly earlier in the same document, beyond the current local window), and the recent history of $d-1$ other documents (negatives — these are tokens from unrelated documents that happen to be nearby in the batch). The key insight is that the attention mechanism must learn to weight the positive keys highly and the negative keys lowly in order to minimize the language modeling loss on the current document's tokens, because only the positive keys provide relevant context for predicting the next token.

Why this form: this design makes the training objective fully differentiable through all keys and values in $p^\delta$ . Unlike the Memorizing Transformer, where the keys retrieved from memory are frozen snapshots and gradients only flow through the query and the local context, FOT's crossbatch procedure allows gradients to flow from the loss, through the attention weights, through the queries, through the inner products, and into the keys and values of all $d$ documents. This joint optimization of queries, keys, and values is what allows the model to shape the representation space to support retrieval.

The paper provides a concrete code sketch in Listing 1 (with a more detailed version in Listing 2 of the Appendix). The implementation is straightforward: for each element of the batch, compute a set of crossbatch indices that reference other positions in the batch, gather the keys from those positions, encode them as if they were at the beginning of the local context (position 0, meaning no positional encoding in most FOT variants), and concatenate them with the local context keys before computing attention. The paper notes: "the changes to the code are small; they are localized to the memory layer (the other layers follow the standard transformer protocol) and do not require any new trainable parameters."

The only new hyperparameter introduced by crossbatch is $d$ , which controls the ratio of positive to negative documents. The paper explicitly describes how $d$ interacts with the batch size $b_S$ : "the number of different documents is equal to $b_S$ (the batch size, i.e. each document has a separate index in the batch)" and "we include into $p^\delta$ all tokens from $C_{\text{prev}}$ with the batch indices in $\{i, (i+1) \bmod b_S, \dots, (i+d-1) \bmod b_S\}$ ."

The paper finds that starting training with small $d$ (typically $d \leq 8$ ) and then switching to larger values ( $d \geq 64$ ) is important: "we find it beneficial to start with small $d \leq 8$ (otherwise, the model tends to ignore the previous local context) and later switch to bigger values, say $d \geq 64$ ." This curriculum is necessary because at the start of training, if $d$ is too large, the model sees too many negative keys and insufficient signal from the single positive document, causing it to learn to ignore all crossbatch keys entirely rather than learning to discriminate positives from negatives. By starting with a small $d$ where the positive signal is strong (the model can easily learn that the one additional document's keys are relevant), and then gradually increasing $d$ to introduce more negatives, the model first learns the basic task of attending to crossbatch keys and then learns to maintain focus as the distraction level increases.

The standard language modeling loss (next-token prediction cross-entropy) is used — no additional contrastive loss term is added. The paper emphasizes this: "it does not require any additional loss (i.e., uses the standard transformer training objective) and is done on the level of the data loading pipeline and a minor self-attention change." The contrastive effect is achieved implicitly: to predict the next token of document $\delta$ correctly, the model must attend to the positive keys (which contain useful context from the same document) and ignore the negative keys (which contain irrelevant information from other documents). The language modeling loss provides the training signal that shapes the representations to support this discrimination.

A subtle but important property: the crossbatch training naturally avoids the staleness problem that affects the Memorizing Transformer. In MT, during training, the memory contains (key, value) pairs that were computed on previous training steps using older model parameters — they are "stale" relative to the current parameters. The paper notes (Wu et al., 2022, Section 3.2) that this staleness required engineering workarounds in MT. In FOT, because $C_{\text{prev}}$ is embedded for each batch using the current model parameters, all keys and values in the expanded set are up-to-date, and gradients flow through them naturally.

Design Choices for the LONGLLAMA Models — Adapting FOT to Pre-Trained LLaMA Checkpoints

Section 4 and Appendix A.2 describe the specific modifications made to apply FOT to the OpenLLaMA 3B and 7B checkpoints. The architecture of OpenLLaMA is described in Appendix A.1: "decoder-only architecture with rotary positional embeddings, pre-normalization with RMSNorm (Zhang and Sennrich, 2019), and SiLU activation (Elfwing et al., 2017). A SentencePiece tokenizer (Kudo and Richardson, 2018) with 32k vocabulary size is used." Three key adaptations are made to the standard FOT procedure:

1. Positional encodings are retained in the local context. The paper states: "To achieve backward compatibility with the original LLaMA, we retain positional encodings in the local context." This is a departure from the standard FOT design, which removes positional encodings from memory attention layers entirely. The tradeoff is that "our checkpoint is backward compatible, i.e., can be used with any existing LLaMA inference code (both in Hugging Face and other implementations), albeit without long-context capabilities." The positional encoding for memory keys is handled by "encoding them as if they were at position 0 in the local context window" — essentially, giving all memory keys the same positional encoding, which removes position-dependent information from the retrieved keys while maintaining compatibility with the LLaMA architecture's expectation that all keys have rotary position embeddings.

2. Dense attention replaces kNN retrieval. The paper states: "we use dense attention instead of the kNN retrieval, as we found only marginal performance differences and it is simpler to implement." This is a practical engineering choice for the LONGLLAMA release: dense attention over the full memory is conceptually simpler (no dependency on FAISS, no separate kNN index maintenance) and the performance difference was empirically small. However, this choice means the LONGLLAMA models do not demonstrate the computational efficiency benefits of kNN-based retrieval — they compute full attention over the entire memory, which would not scale to millions of tokens. The paper's analysis models (Section 5) use kNN retrieval, confirming that the method works with both dense and sparse attention over the memory.

3. Fine-grained crossbatch configuration. Rather than using a uniform $d$ for all elements in the batch, LONGLLAMA uses a heterogeneous setup where different segments of the batch receive different crossbatch configurations. For the 3B model, memory layers are $L = \{6, 12, 18\}$ (three of the 26 layers), and the batch is divided into four equal segments with the following configurations (written as (positives, negatives)):

$\frac{1}{4}(0, 0)$ : one quarter of the batch only sees local context (no crossbatch at all — these elements behave like standard training, providing a regularization effect and ensuring the model doesn't become dependent on crossbatch keys)
$\frac{1}{4}(1, 1)$ : one quarter sees one positive and one negative document
$\frac{1}{4}(2, 1)$ : one quarter sees two positive and one negative document
$\frac{1}{4}(3, 0)$ : one quarter sees three positive documents (from the same document — multiple previous context windows) and no negatives

For the 7B model, memory layers are $L = \{8, 16, 24\}$ , and the configuration is $\frac{1}{4}(0, 0), \frac{1}{4}(1, 2), \frac{1}{4}(2, 5), \frac{1}{4}(3, 4)$ .

Why this heterogeneous setup: the paper explains that this provides "more fine-grained control over the number of additional contexts and the ratio of positive to negative samples." By exposing different parts of the batch to different ratios, the model learns to handle a range of distraction levels, preventing it from overfitting to a single $d$ value. The mixture of $(0,0)$ elements (no crossbatch) ensures the model retains the ability to function without crossbatch input — important for backward compatibility and for stability. The mixture of configurations with multiple positives and low negatives (e.g., $(3,0)$ or $(2,1)$ ) provides strong training signal for learning to use in-document context, while configurations with more negatives (e.g., $(1,2)$ or $(2,5)$ in the 7B model) explicitly train discrimination.

Fine-tuning hyperparameters (Appendix A.2): batch size of 256K tokens, constant learning rate of $2 \times 10^{-5}$ (lower than the learning rate at the end of OpenLLaMA training, which was $3 \times 10^{-5}$ after 1T tokens), weight decay of 0.01, same optimizer as OpenLLaMA. The 3B model is fine-tuned on 10B tokens and the 7B model on 3B tokens, using an 8k context length. The training data mixture is based on RedPajama (TogetherComputer, 2023) with additional Python code from The Stack (Kocetkov et al., 2022), with proportions given in Table 4 of the Appendix: arxiv 25%, python 25%, book 10%, common_crawl 29%, c4 5%, github 2%, stackexchange 2%, wikipedia 2%. Documents shorter than specified minimums are filtered out, and "in case one document is too short to span across several contexts for crossbatch, then we concatenate it with the next document from the dataset."

The Formal Compute-Optimal Objective Analogy (Explicit Design Justification)

While the paper does not formulate a formal optimization objective for FOT the way some works do, the implicit objective that the crossbatch training optimizes can be understood as shaping the key representation space $K$ and query representation space $Q$ such that for queries from document $\delta$ and keys from document $\gamma$ :

$\langle q^\delta, k^\gamma \rangle \text{ is large when } \delta = \gamma \text{ and small when } \delta \neq \gamma$

In other words, the inner product between a query and a key should be a reliable signal of document co-membership. A standard transformer trained only on single-document contexts never receives a gradient signal that distinguishes $\delta = \gamma$ from $\delta \neq \gamma$ — all keys it sees during training are from the same document (the current one), so the model has no reason to learn representations where document boundaries are meaningful.

The crossbatch procedure introduces this signal by construction: the expanded key-value set $p^\delta$ for document $\delta$ always contains keys from both $\delta$ and $\delta' \neq \delta$ . To minimize the language modeling loss, the attention mechanism must learn to assign higher weights to the $\delta$ -keys than to the $\delta'$ -keys, because only the $\delta$ -keys contain information that helps predict the next token of document $\delta$ . This is a contrastive objective implemented entirely through the forward pass of the attention mechanism, without any explicit contrastive loss term — the "contrast" is between the positive document (whose keys should receive high attention) and the negative documents (whose keys should receive low attention), and the language modeling loss provides the gradient signal that shapes representations to achieve this contrast.

Why not an explicit contrastive loss? The paper briefly mentions that "exploring other contrastive learning objectives could be beneficial for further improving the key structure in future work" (Section 6), implying that the choice of using only the standard LM loss was deliberate: it keeps the method simple, requires no additional hyperparameters or loss balancing, and ensures the representations are optimized for the exact task they will be used for (language modeling with retrieval-augmented attention), rather than being optimized for some auxiliary metric that may not align perfectly with downstream performance. The cost is that the contrastive signal is implicit — the model must discover document boundaries through the statistical structure of the data itself (tokens from different documents are less predictive of each other than tokens from the same document), which may be a weaker signal than an explicit same-document/different-document label.

Comparison to Standard Long-Context Fine-Tuning (Why FOT Extrapolates)

A crucial claim in the paper is that FOT enables extrapolation beyond the training context length, while standard long-context fine-tuning does not. Section 4.5 provides the direct comparison: two 3B models fine-tuned on 1B tokens with 4K context length, one using FOT and one using standard fine-tuning ("done similarly to MosaicML, 2023; Nijkamp et al., 2023"). Both models improve when context is expanded from 2K (the original OpenLLaMA training length) to 4K (the fine-tuning length). However, when evaluated at 6K and 8K — beyond the training length — the FOT model continues to improve (TREC: 55.6% at 2K → 60.9% at 4K → 61.7% at 6K → 62.5% at 8K), while the baseline is marked with "—" indicating it cannot handle these lengths at all.

The reason for this difference is not explicitly justified in the paper but follows from the architecture: standard fine-tuning trains the model to attend over positions within the training context length using positional encodings. Beyond that length, the positional encodings correspond to positions the model has never seen, and the attention mechanism fails. FOT's memory attention layers, by contrast, use no positional encodings on the retrieved keys — retrieval is purely content-based (via inner product similarity), not position-based. As long as the key representations are well-structured (which crossbatch training ensures), the model can retrieve relevant keys regardless of how far back in the sequence they occurred. The local context still uses positional encodings (for LONGLLAMA, up to 2K tokens), but the memory context has effectively infinite range.

This is the key architectural insight: positional encodings bottleneck context length; content-based retrieval does not. Standard transformers rely on positional encodings to distinguish tokens at different positions, and these encodings have a finite training range. Removing positional encodings from the retrieval mechanism removes this bottleneck, leaving the quality of the key representations as the sole limiting factor — which is exactly what crossbatch training optimizes.

Ablation Justifications: Why Differentiability and Negatives Both Matter

The paper provides two critical ablations in Appendix C that justify key design choices:

Differentiability (Appendix C.1): Comparing FOT with $d=1$ (no negatives, but differentiable keys and values through the previous context) against Memorizing Transformer (non-differentiable memory). In the multi-document setting at the same training context length (512 tokens), FOT with $d=1$ achieves substantially better perplexity, with the gap widening as memory grows (Figure 5). Even in the single-document setting (Table 8), FOT with $d=1$ outperforms MT at all context lengths tested (512, 1024, 2048). The interpretation: being able to backpropagate through the keys and values of the previous context allows the model to jointly optimize its key, query, and value representations for the retrieval task, while MT's non-differentiable memory means the key representations are whatever the base model happened to learn, frozen and suboptimal.

Negatives (Appendix C.2): Comparing FOT with $d=1$ (differentiable but no negatives — only the previous context of the same document) against FOT with $d=2$ then $d=64$ (differentiable with negatives). In the multi-document setting (Figure 6), as memory grows, FOT with $d=1$ degrades significantly (perplexity rises from approximately 14.0 to approximately 16.0 as memory grows from 0 to 500k tokens), while FOT with $d=2 \to 64$ degrades much more gracefully (perplexity rises only from approximately 14.0 to approximately 14.5). The interpretation: without negatives during training, the model never learns to distinguish relevant from irrelevant keys — it learns that all crossbatch keys are relevant (because they all came from the same document during training), and this assumption fails catastrophically in the multi-document setting where most memory keys are irrelevant. The negatives teach the model that not all keys are created equal.

These ablations together establish that both differentiability and negatives are necessary for FOT's performance — neither alone is sufficient. Differentiability without negatives fails in multi-document settings; negatives without differentiability (the MT baseline) fail to achieve good key structure even in single-document settings.

4. Key Insights and Innovations

Innovation 1: The Distraction Issue — A New Diagnostic Framework for Why Long Contexts Fail

The paper's most conceptually distinctive contribution is not a method but a diagnosis: it identifies, names, and formally quantifies the distraction issue as the primary bottleneck preventing transformers from effectively using long multi-document contexts. Before this work, the community understood that long-context performance degraded, but the degradation was attributed to various causes — computational inefficiency (quadratic attention cost), positional encoding limitations, or general optimization difficulty. The field lacked a precise, mechanistic account of why attending over more tokens produces worse results even when the architectural capacity exists.

The paper's diagnostic move is to decompose the problem into a ratio: as the number of documents in context grows, the proportion of relevant to irrelevant keys shrinks, and the model's attention mass becomes uniformly distributed — not because the model can't compute attention, but because the key-value representations lack the structure to support discrimination. The formalization via positive attention mass $r_d$ (Section 3.3, Figure 3) transforms this from a vague intuition into a measurable quantity, and the empirical finding that a standard transformer achieves $r_d \approx 1/d$ — exactly the random-chance baseline — is a crisp, falsifiable diagnostic result.

Why this is a fundamental contribution rather than incremental: The distraction issue reframes the long-context problem from an architectural challenge (how to make attention cheaper) to a representational challenge (how to make keys discriminable). This shift has implications that cascade through the paper: it explains why sparse attention and positional interpolation methods hit ceilings (they solve computational or positional bottlenecks but not the representational one), it motivates the specific form of the training solution (contrastive exposure to negatives, because discrimination is exactly what needs to be learned), and it predicts where the method will and won't work (multi-document settings with high document count are the stress test; single-document settings are trivially easier because all keys are potentially relevant).

This reframing also resolves a subtle tension in the literature. Prior work on representation degeneration (Gao et al., 2019) had identified that language model embeddings collapse into narrow cones, but the connection to context scaling had not been drawn. The Memorizing Transformer (Wu et al., 2022) demonstrated that kNN retrieval could extend effective context, but it did not diagnose why naive kNN fails in multi-document settings — it simply avoided multi-document evaluation. FOT's diagnostic framework explains both: the cone collapse makes inner-product similarities uniformly high across all documents (Gao et al., 2019 explains the mechanism), and this uniformity produces the $r_d \approx 1/d$ distraction pattern (Figure 3 demonstrates the consequence). The diagnosis unifies two previously separate observations into a single causal chain.

Evidence anchor: Figure 3 is the key exhibit. The near-perfect alignment of the standard transformer's positive attention mass with the $1/d$ line is not a gradual degradation — it is a precise quantitative match to the random baseline, demonstrating that the model has learned literally nothing about document-level relevance. The FOT curves in the same figure ( $d=8$ and $d=2 \to 64$ ) show that crossbatch training breaks this pattern, confirming the causal link between training procedure and the distraction phenomenon.

Innovation 2: Contrastive Learning as an Implicit Attention-Shaping Mechanism — No Auxiliary Loss Required

The crossbatch training procedure is structurally simple — it rearranges which keys a subset of attention layers can see during training — but its conceptual significance lies in how it achieves a contrastive training effect without any explicit contrastive loss term. The paper does not add a SimCLR-style or CLIP-style loss. It does not require same-document/different-document labels. The contrastive effect emerges entirely from the interaction between the expanded key-value set and the standard language modeling loss: the model must learn to discriminate positive keys (from the same document) from negative keys (from other documents) because only the positive keys contain information useful for predicting the next token of the current document. The language modeling loss is the contrastive signal, operating through the attention mechanism's softmax.

Comparison to prior work: Contrastive learning in vision (SimCLR, Chen et al., 2020; CLIP, Radford et al., 2021) and in language (ContraCLM, Jain et al., 2023) typically operates through explicit contrastive loss terms added to the training objective, requiring careful balancing of loss weights, explicit definition of positive and negative pairs, and often large batch sizes to provide sufficient negatives. TRIME (Zhong et al., 2022), the closest prior work in language modeling with memory augmentation, similarly interpolates an explicit contrastive objective. FOT's implicit approach is more elegant in two ways: (a) it requires no loss engineering — the standard LM loss automatically provides the right gradient signal because the model's own prediction error reveals which keys were helpful and which were not; (b) the "positive/negative" relationship is not an arbitrary label but is grounded in the actual predictive utility of the keys — a key from a different document could be useful if it happens to contain relevant information (e.g., two documents about the same topic), and the model can learn this from the data, rather than being forced to treat all cross-document keys as negatives.

This design choice also explains why FOT generalizes from multi-document training to single-document evaluation (Section 5.4, Figure 9): because the model learns a continuous notion of key relevance based on predictive utility rather than a binary same-document/different-document signal, the learned representations generalize to settings where all keys come from the same document but some are more contextually relevant than others.

Why this is a conceptual advance rather than just an implementation trick: Most contrastive learning methods treat the contrastive objective as an auxiliary task — learn representations that are good at discriminating positives from negatives, then hope those representations transfer to downstream tasks. FOT collapses the distinction between the contrastive task and the downstream task: the task is language modeling, and the contrastive structure is baked into the attention computation itself. The representations are optimized for exactly what they will be used for — retrieval-augmented attention for next-token prediction — without any transfer gap. This integration of contrastive structure into the primary objective, without auxiliary losses, is a design pattern that could generalize beyond context scaling to other settings where attention selectivity matters.

Evidence anchor: The ablation in Appendix C.2 (Figure 6) isolates the role of negatives by comparing $d=1$ (differentiable crossbatch but no negatives — only the same document's previous context) against $d=2 \to 64$ (negatives present). The performance gap widens dramatically as memory grows: without negatives, the model's perplexity degrades sharply in multi-document settings because it never learned to reject irrelevant keys. This confirms that the contrastive signal from negatives — delivered entirely through the LM loss — is the active ingredient, not the mere presence of additional differentiable context.

Innovation 3: Content-Based Retrieval as an Escape from the Positional Encoding Bottleneck — Empirically Demonstrated Extrapolation to 256k from 8k Training

The paper makes a strong empirical claim that has significant architectural implications: a model trained with an 8k context length can perform passkey retrieval at 256k tokens (a 32× extrapolation) when the long-range attention is based on content similarity rather than positional encoding. This is not merely a performance result — it constitutes evidence for a specific architectural hypothesis: that positional encodings are the primary bottleneck limiting context length extrapolation, and that removing them from long-range attention removes this bottleneck entirely.

Comparison to prior work: Position Interpolation (Chen et al., 2023; kaiokendev, 2023) achieves context extension by rescaling positional encodings so that positions beyond the training range are mapped into the training range — it extends context but remains bounded by the positional encoding scheme's representational capacity. Landmark Attention (Mohtashami and Jaggi, 2023) achieves 32k context on LLaMA-7B, also through positional encoding manipulation. These methods accept the positional encoding constraint and work around it. FOT's approach is fundamentally different: it eliminates the constraint by eliminating positional encodings from the long-range attention mechanism entirely. The theoretical implication — "theoretically unbounded context length" as the paper claims — follows directly: with no positional encoding to saturate, the model's context length is limited only by the quality of the key representations, not by any architectural ceiling.

Why this is more than a scaling result: The passkey retrieval result (Figure 1) — 94.5% at 100k, 73% at 256k — demonstrates that content-based retrieval works at scales far beyond training. But the deeper significance is what it reveals about the failure mode of position-dependent approaches. If positional interpolation methods cap at ~32k, and FOT reaches 256k with no sign of catastrophic failure (the decline from 100k to 256k is gradual), then positional encoding is not just a minor inconvenience — it is the rate-limiting factor, and removing it unlocks an order-of-magnitude improvement. This is a falsifiable claim that the paper supports with a clean ablation: standard long-context fine-tuning (which retains positional encodings throughout) fails to extrapolate at all beyond its training length (Table 2, Figure 1 baseline), while FOT succeeds at 32× its training length.

The connection to Haviv et al. (2022) — which showed that transformers can learn positional information without explicit positional encodings — is important context. FOT's result extends this finding from "position can be learned" to "removing position from long-range attention enables unbounded context." It's not just that positional encodings are unnecessary; it's that they are actively harmful for extrapolation, and removing them is a sufficient condition for achieving it.

Evidence anchor: Figure 1 (passkey retrieval accuracy vs. prompt length) juxtaposed with Table 2 (no extrapolation for standard fine-tuning). The combination demonstrates both that FOT extrapolates and that the baseline does not, isolating the architectural difference (content-based retrieval without positional encoding vs. position-dependent attention) as the causal factor.

Innovation 4: Memory-Augmented Transformers Need NOT Train on Long Sequences — The Counterintuitive Efficiency of Crossbatch

A surprising finding with practical implications is that FOT achieves its long-context capabilities without ever training on long sequences or even using the inference-time memory mechanism during training. The total differentiable context during FOT training is small — typically 2× the local context length (current window + one previous window), sometimes 3× or 4× when $w > 1$ — yet the model learns representations that support retrieval from memories containing hundreds of thousands of tokens at inference time. The dictionary lookup task (Appendix F, Figure 10) makes this particularly stark: a 37M-parameter model trained on documents of length 512 tokens, after only 5k training steps, achieves over 92% accuracy with a 16M-token memory — a 31,250× gap between training context length and inference memory size.

Comparison to prior work: The Memorizing Transformer (Wu et al., 2022) trains with the actual memory mechanism active — it retrieves from a non-differentiable memory during training, which requires (a) the training documents to be long enough to populate a meaningful memory (typically 16k–65k tokens), and (b) engineering workarounds for the staleness problem (memory keys computed with outdated parameters). This creates a practical barrier: many widely-used corpora consist of short documents (Wikipedia articles, social media posts, code snippets) that cannot support MT-style memory training. FOT eliminates both requirements — crossbatch only needs two consecutive context windows from the same document (a few thousand tokens total), and differentiability eliminates staleness. The result is that FOT can be applied to a much broader range of training data.

Why this is conceptually important beyond convenience: The finding challenges an implicit assumption in the memory-augmented transformer literature — that the model needs to practice retrieval during training in order to learn how to use memory at inference. FOT demonstrates that this is false: the model only needs to learn good key-value representations (which crossbatch provides), and the kNN retrieval mechanism will work automatically because it's just a content-based lookup. The separation between learning to represent (done during training) and learning to retrieve (not needed — retrieval is a deterministic, non-learned operation) is a cleaner decomposition than the MT approach, which conflates them.

The dictionary lookup result (Figure 10) provides the cleanest demonstration because it removes any confounds about document structure or natural language semantics — it's purely a test of whether the key representations support content-based retrieval at scale. The fact that a tiny model trained on 512-token synthetic documents can retrieve from 16M tokens suggests that crossbatch training is doing something fundamental to the representation space, not just learning document-level heuristics.

Evidence anchor: Table 3 (FOT fine-tuned with 2k training context achieves perplexity improvements at 64k evaluation context) and Figure 10 (512-token training → 16M-token retrieval). Both demonstrate the training-inference context gap, with the dictionary task providing the most extreme ratio and the language modeling tasks demonstrating that the effect transfers to natural data.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary evaluation datasets span four long-context language modeling domains: PG-19 (English books, Rae et al., 2019), arXiv (mathematical papers, obtained via arXiv Bulk Data Access where each paper is comparable in length to a PG-19 book), GitHub (permissively licensed code repositories with files concatenated within subdirectories), and Isabelle (formal mathematical proofs from The Archive of Formal Proofs and the Isabelle standard library, concatenated based on import dependencies). The C4 dataset (Raffel et al., 2019a) is used for pre-training and fine-tuning in the analysis experiments. For downstream tasks, the paper uses TREC question classification (Li and Roth, 2002; Hovy et al., 2001), WebQS question answering (Berant et al., 2013), Qasper (Dasigi et al., 2021) from the SCROLLS benchmark (Shaham et al., 2022), and a synthetic passkey retrieval task (Mohtashami and Jaggi, 2023). A synthetic dictionary lookup task is also used for controlled analysis of retrieval capability. All datasets except Qasper and the synthetic tasks are used for zero-shot language modeling evaluation (perplexity measurement on 1M tokens per dataset for test perplexity calculation, as stated in Appendix J).
Base model(s). Analysis experiments (Sections 5.1–5.5) use decoder-only transformers with 12 layers and 184M parameters (unless stated otherwise), trained from scratch. A smaller 37M-parameter variant is used for the dictionary lookup task. The LONGLLAMA experiments (Section 4) fine-tune OpenLLaMA-3B and OpenLLaMA-7B checkpoints (Geng and Liu, 2023), which are open-source reproductions of LLaMA (Touvron et al., 2023) trained for 1T tokens. The 3B model uses 26 layers, the 7B model uses 32 layers. The paper argues these models are "representative of the capabilities of many contemporary LLMs" and sit in a performance regime where long-context fine-tuning is meaningful.
Metrics. The primary metric for language modeling experiments is perplexity — the exponentiated cross-entropy loss on next-token prediction, computed over 1M test tokens per dataset. For passkey retrieval, accuracy is the fraction of prompts where the model correctly retrieves the hidden passkey. For TREC and WebQS, accuracy is the fraction of questions correctly classified/answered, averaged over 20 trials with different random samplings of in-context demonstrations (confidence intervals reported as smaller than 1% for TREC and 0.1% for WebQS). For Qasper, performance is measured as the exact match score from the LM Evaluation Harness implementation. For the dictionary lookup task, accuracy is the fraction of queries where the model correctly predicts all four value tokens. For short-context capability preservation, the LM Evaluation Harness benchmark (Gao et al., 2021a) reports standard task-specific metrics (accuracy, accuracy_norm, F1, exact match).
Baselines. The paper compares against three baseline methods: (1) Standard Transformer — a vanilla decoder-only transformer with the same architecture, trained with the standard language modeling objective and no memory augmentation. This baseline appears in all language modeling experiments (Table 3, Figures 3, 6, 8, 9). (2) Memorizing Transformer (MT) (Wu et al., 2022) — the model that FOT directly extends, which augments a subset of attention layers with kNN-retrieved memory but uses non-differentiable retrieval during training. MT baselines are configured with memory sizes of 2k–65k tokens, as specified per experiment. For the baseline in Table 3 and Section 5.2, MT is fine-tuned from the same base checkpoint as FOT using local context of 1K and memory size of 16K or as specified. (3) Transformer-XL (Dai et al., 2019) — which caches previous context across all layers and attends across segment boundaries. In Table 3, Transformer-XL is fine-tuned from the same base checkpoint with both local context and window length of 1K. For the LONGLLAMA experiments, the primary baseline is standard long-context fine-tuning ("done similarly to MosaicML, 2023; Nijkamp et al., 2023") where the model is fine-tuned on 4K context using the standard objective without crossbatch, compared to FOT fine-tuned on the same context length (Table 2). Additionally, LongChat-7B (Ma and Zhang, 2023) serves as a comparison for the Qasper task.
Generation budget / compute accounting. In the language modeling experiments, models are compared at equivalent context lengths (the number of tokens the model can attend to during evaluation) rather than equivalent computational budgets. For training, the batch size is measured in tokens (32K tokens for most analysis experiments, 128K tokens for Section 5.2/5.4 experiments, 256K tokens for LONGLLAMA fine-tuning) and training steps are reported (100k pre-training steps + 10k fine-tuning steps for analysis; 10B/3B tokens for 3B/7B LONGLLAMA fine-tuning). The crossbatch training procedure incurs additional computational cost "only in a subset of layers" (Section 6), but the paper does not provide a detailed FLOPs comparison between FOT training and standard training. For inference, the primary compute metric is the total context length (tokens in local context + tokens in memory). The kNN retrieval uses exact search (not approximate), implemented with FAISS (Johnson et al., 2017), which the paper notes "is not scalable to large memory" (Section 6). For LONGLLAMA models, dense attention over the full memory is used instead of kNN. The paper does not report wall-clock time or inference latency comparisons.
Cross-validation / statistical protocol. For the TREC and WebQS few-shot experiments, results are averaged over 20 trials with different random samplings of in-context demonstrations from the training set, with reported confidence intervals smaller than 1% and 0.1% respectively. For the arXiv baseline experiment in Appendix K, three runs with different random seeds were performed, yielding a standard deviation of 0.002 perplexity. The paper acknowledges that "due to resource constraints, we were unable to conduct multiple runs for all experiments" and states that preliminary findings "indicate that the observed variance was minimal compared to the impact observed from other factors under investigation" (Appendix J). For the dictionary lookup task, error bars in Figure 10 represent the minimum and maximum values over 10 random seeds.

Main Quantitative Results

Distraction Issue Quantification (Section 3.3)

The paper's foundational empirical claim is that standard transformers suffer from the distraction issue — their attention mass is uniformly distributed across documents regardless of relevance — and that crossbatch training mitigates this. Figure 3 provides the key evidence.

Standard Transformer (d=1, blue line): The positive attention mass $r_d$ tracks $1/d$ almost exactly as the number of documents $d$ increases from 1 to 64. At $d=64$ , $r_d$ falls to approximately 0.016, indistinguishable from the $1/64 \approx 0.0156$ random baseline. The model allocates no more attention to the relevant document than to any individual irrelevant document.

FOT with d=8 (orange line): $r_d$ declines from approximately 0.9 at $d=1$ to approximately 0.35 at $d=64$ . This is substantially better than random but still shows significant degradation as the number of documents increases — the model is learning to discriminate but not robustly enough to handle large numbers of distractors.

FOT with d=2 → d=64 (green line): $r_d$ starts close to 1.0 at $d=1$ and remains above 0.8 even at $d=64$ . The model dedicates roughly 80% of its attention to the single relevant document and distributes the remaining 20% across 63 irrelevant ones. This demonstrates that exposure to many negatives during training (via the d=64 phase) builds robust discrimination that scales to large numbers of distractors.

The interpretation: the crossbatch training procedure directly addresses the distraction issue, and the number of negatives seen during training ( $d$ ) controls the robustness of the learned discrimination. Starting with small $d$ and increasing — the curriculum — is necessary: "we find it beneficial to start with small $d \leq 8$ (otherwise, the model tends to ignore the previous local context) and later switch to bigger values, say $d \geq 64$ " (Section 3.2).

Language Modeling Performance with Extended Context (Section 5.2, Table 3)

The paper evaluates whether FOT fine-tuning of a pre-existing model enables it to utilize context lengths far beyond the training context. A standard transformer is pre-trained for 100k steps with 1K context on C4, then fine-tuned for 10k steps with FOT (crossbatch, $d=128$ , total differentiable context 2K) and evaluated zero-shot on four long-context language modeling datasets.

The headline result from Table 3: FOT achieves steadily decreasing perplexity as evaluation context grows from 2K to 64K tokens, despite being fine-tuned with only 2K total differentiable context. On arXiv, perplexity improves from 8.17 at 2K to 6.81 at 64K — a 16.6% reduction. On PG-19, the improvement is from 23.74 to 22.65 (4.6% reduction). On GitHub, from 6.72 to 5.32 (20.8% reduction). On Isabelle, from 5.63 to 4.44 (21.1% reduction).

Comparison to Memorizing Transformer (MT, Table 3): MT with 16K memory is also fine-tuned from the same base checkpoint and evaluated at the same context lengths. At 2K evaluation context, FOT substantially outperforms MT across all datasets: 6.72 vs. 8.10 on GitHub, 5.63 vs. 7.34 on Isabelle, 8.17 vs. 9.39 on arXiv, 23.74 vs. 24.03 on PG-19. The gap persists and widens at longer contexts — at 64K, FOT achieves 6.81 on arXiv vs. MT's 8.60, and 22.65 on PG-19 vs. MT's 23.24. MT does improve with longer context (8.10 at 2K → 7.26 at 64K on GitHub), but the improvement is smaller than FOT's and starts from a higher baseline.

Comparison to Transformer-XL (Table 3): At 2K context, FOT achieves comparable or slightly better perplexity than Transformer-XL: 6.72 vs 6.85 on GitHub, 5.63 vs. 5.76 on Isabelle, 8.17 vs. 8.21 on arXiv, 23.74 vs. 23.57 on PG-19. This is notable because Transformer-XL has access to previous context in all layers, unlike FOT and MT which only augment a single memory layer, yet FOT matches its performance at 2K and substantially surpasses it at longer contexts where Transformer-XL cannot be evaluated (it is bounded by its training window of 1K).

A critical methodological note: "Unlike MT, our method does not require training on long sequences, which is reflected by the lower perplexities of FOT when evaluated in the zero-shot setting." The implication is that FOT's training efficiency (small training context) does not come at the cost of evaluation-time performance — quite the opposite, FOT achieves better results than MT despite using less training context.

Context Length Extrapolation in Single-Document Setting (Section 5.4, Figure 9)

An unexpected finding: FOT's crossbatch training, designed for multi-document settings, also improves single-document context extrapolation. A model is FOT fine-tuned on C4 and evaluated on PG-19 in the single-document setting (memory cleared for each new document) with varying context lengths.

The key finding from Figure 9: Increasing the crossbatch dimension $d$ from 1 to 2 yields a significant perplexity improvement across all evaluation context lengths up to 64K. For $w=1$ (two total contexts: current + one previous), the $d=2$ curve (green) is substantially below the $d=1$ curve (blue) at all context lengths beyond the training length. For $w=2$ (three total contexts: current + two previous), the $d=2,w=2$ configuration (brown) achieves the best performance, with perplexity approximately 22.4 at 64K tokens compared to approximately 23.2 for $d=2,w=1$ and approximately 23.8 for $d=1,w=1$ .

The comparison between $d=1,w=2$ (red) and $d=2,w=1$ (green) reveals an interesting tradeoff: the former (longer training context but no negatives) slightly outperforms the latter (shorter training context with negatives). "This is natural, as the former has longer training context." However, the $d=2,w=2$ configuration (both longer context and negatives) achieves the best results, suggesting that negatives and longer training context are complementary.

Why this matters: The result demonstrates that learning to discriminate relevant from irrelevant keys — even when all keys during training are from different documents — transfers to the single-document setting where the "irrelevant" keys are simply less contextually useful parts of the same document. The model learns a more nuanced notion of relevance than simple document membership.

Handling Distractions in Multi-Document Language Modeling (Section 5.3, Figure 8)

The paper directly tests the relationship between the crossbatch dimension $d$ (which controls the number of negatives during training) and the model's ability to handle distractions during inference. The setup: FOT models trained with different $d$ values are evaluated on PG-19 in the multi-document setting, where memory persists across books and the proportion of irrelevant keys grows with memory size.

The headline from Figure 8: Higher $d$ during training leads to better perplexity scaling as memory size increases. The standard transformer (no memory, purple dashed line) serves as a lower bound. FOT with $d=1$ (light blue) degrades rapidly as memory grows from 0 to over 500K tokens — perplexity rises from approximately 14.0 to approximately 16.0, approaching the no-memory baseline. FOT with $d=2$ (dark blue) degrades more slowly, reaching approximately 15.2 at 500K tokens. FOT with $d=4$ (teal) and $d=8$ (orange) show progressively better scaling, with $d=8$ reaching approximately 14.8 at 500K. FOT with $d=2 \to 64$ (bold green line) achieves the best scaling: "the perplexity increases only by 0.18 when scaling to > 500k tokens" — from approximately 14.1 at 0 memory to approximately 14.28 at maximum memory.

A critical comparison: Single-document MT evaluated in single-document mode (dashed dark green line) serves as a "soft lower bound" — it represents the best perplexity achievable if the model could perfectly ignore all irrelevant keys, since it never sees any. FOT with $d=2 \to 64$ approaches this lower bound, indicating that the model has learned to effectively filter out distractions. In contrast, FOT with $d=1$ is far from this bound, confirming that lack of negative training examples leads to catastrophic degradation in multi-document settings.

Why this matters: This is the most direct evidence that crossbatch training addresses its stated problem. The monotonic relationship between training $d$ and multi-document inference performance confirms the paper's central hypothesis: exposure to negatives during training is necessary for handling distractions at scale.

Dictionary Lookup Task — Extreme Extrapolation (Appendix F, Figure 10)

The dictionary lookup task provides the most dramatic demonstration of context extrapolation: models trained on documents of length 512 tokens are evaluated with memories containing up to 16M tokens (a 31,250× gap). The task format: a document contains key-value definitions followed by queries, and the model must retrieve the correct value for a queried key from memory.

The headline from Figure 10: FOT achieves over 92% accuracy with 16M tokens in memory after only 5K training steps, while the baseline transformer fails completely at 16K tokens (accuracy drops to near zero). The FOT model maintains high accuracy across the full range: approximately 99% at 256 tokens, 99% at 1K, 98% at 4K, 97% at 16K, 96% at 64K, 95% at 256K, 94% at 1M, and over 92% at 16M. The baseline transformer achieves near-perfect accuracy at 256 and 1K tokens (where the full dictionary fits in its 512-token local context window) but collapses immediately beyond its training length, dropping to near-zero accuracy for 16K tokens and beyond.

The training configuration is notable: FOT uses local context of 256 tokens, so the model must use the memory attention layer to answer queries correctly (the definitions are in the first half of the 512-token document, beyond the 256-token local window). The crossbatch training starts with $d=1$ and switches to $d=128$ "as soon as the model is able to reach 98% training accuracy." During inference, $k=32$ keys are retrieved from memory. Error bars in Figure 10 represent the minimum and maximum over 10 random seeds, showing consistent performance across runs.

Why this matters: The dictionary lookup task removes all natural language confounds — it is a pure test of content-based retrieval. The fact that FOT succeeds at 16M tokens with only 512-token training documents demonstrates in the cleanest possible way that crossbatch training teaches the model to structure its key-value space for retrieval, and that this structure generalizes to retrieval scales orders of magnitude beyond training.

LONGLLAMA Passkey Retrieval (Section 4.2, Figure 1)

The passkey retrieval task (Mohtashami and Jaggi, 2023) measures the effective context length — the maximum distance over which a model can attend to and retrieve a specific piece of information. The model must find a hidden passkey (a random 5-digit number) placed at a random position within a long prompt filled with repeated filler text ("The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.") and answer "What is the passkey?"

The headline from Figure 1: LONGLLAMA 3B achieves 94.5% accuracy at 100k tokens and 73% at 256k tokens, despite being fine-tuned with only an 8K context length. The original OpenLLaMA 3B achieves near-perfect accuracy at 2K (its training length) but drops to near-zero beyond, confirming that standard LLaMA models cannot extrapolate beyond their training context. LONGLLAMA maintains 100% accuracy through 16K tokens, then declines gradually: approximately 99% at 32K, 97% at 50K, 94.5% at 100K, and 73% at 256K.

This near-perfect performance through 16K (2× training length) and substantial performance at 256K (32× training length) is the paper's most visually striking result and serves as the primary evidence for the "theoretically unbounded context length" claim.

Why this matters: The passkey task isolates a specific capability: can the model find and use a single piece of information placed arbitrarily far back in its context? Unlike language modeling perplexity (which can be improved by better local predictions without genuine long-range retrieval), passkey retrieval is a direct test of long-range attention. The result demonstrates that FOT's content-based retrieval mechanism works at extreme scales and that the gradual decline from 100k to 256k suggests the mechanism is not hitting a hard architectural ceiling.

No explicit comparison to position-interpolation methods is made in Figure 1, but the paper notes in Section 2 that Landmark Attention achieves 32K context on LLaMA-7B and Position Interpolation enables fine-tuning for 32K context, making LONGLLAMA's 256K a substantial quantitative advance over those approaches.

TREC and WebQS Few-Shot In-Context Learning (Section 4.4, Table 1)

The paper evaluates whether the extended context provided by LONGLLAMA translates to improved performance on downstream tasks that benefit from more in-context examples.

The headline from Table 1: LONGLLAMA 3B improves TREC accuracy from 67.0% at 2K context to 73.3% at 8K context (+6.3 percentage points). LONGLLAMA 7B shows an even larger gain: 63.2% at 2K → 75.9% at 8K (+12.7 percentage points). On WebQS, the gains are more modest: LONGLLAMA 3B improves from 21.2% at 2K to 22.4% at 8K (+1.2 points), and LONGLLAMA 7B from 25.5% to 27.7% (+2.2 points).

The experimental design follows Hao et al. (2022): "we few-shot prompt the models with as many demonstration examples as possible up to the given context length," but "we do not use structured prompting like in Hao et al. (2022) — instead, we directly provide all demonstrations in context." The TREC dataset has 50 fine-grained question classes, and only 100 examples fit in the standard 2K context, meaning many classes may have zero in-context demonstrations — "making the task impossible" for those classes. Extending context to 8K allows more examples and thus more complete class coverage. The authors average over 20 trials with different random samplings of in-context demonstrations from the training set, with confidence intervals smaller than 1% for TREC and 0.1% for WebQS.

Why this matters: This demonstrates that the extended context translates to genuine downstream task improvements, not just synthetic benchmarks. The TREC result in particular shows a compelling use case: tasks with large label spaces where the standard context length cannot fit enough examples to cover all classes. WebQS shows smaller improvements, which the paper does not explicitly explain but may be due to the task being less dependent on having many in-context examples or having a question format that benefits less from additional demonstrations.

Comparison to Standard Long-Context Fine-Tuning (Section 4.5, Table 2)

A crucial controlled comparison: FOT fine-tuning vs. standard long-context fine-tuning, both on 3B models with 4K context for 1B tokens.

The headline from Table 2: FOT outperforms the baseline at the training context length and extrapolates beyond it; the baseline does not. At 2K context (the original OpenLLaMA training length), FOT achieves 55.6% on TREC vs. 52.8% for the baseline. At 4K (the fine-tuning context length for both), FOT achieves 60.9% vs. 57.2% for the baseline. Critically, at 6K and 8K — beyond the fine-tuning context length — FOT continues to improve: 61.7% at 6K and 62.5% at 8K. The baseline is marked with "—" for 6K and 8K, indicating it cannot handle these context lengths at all.

On WebQS, FOT and baseline are comparable at 2K (20.8% vs. 20.7%), but FOT improves to 21.0% at 4K while the baseline declines to 18.7%. At 6K, FOT achieves 21.2% while the baseline is "—". At 8K, FOT achieves 20.7% (slightly below the 4K result) while the baseline remains "—".

Why this matters: This is the critical ablation that isolates the crossbatch mechanism as the causal factor. Both models have the same base architecture, same fine-tuning data, same context length, and same number of fine-tuning tokens. The only difference is the training procedure — crossbatch vs. standard. The baseline's failure to extrapolate beyond its training length is consistent with the positional encoding bottleneck: standard fine-tuning trains the model's positional encodings for up to 4K positions, and beyond that, the encodings are out-of-distribution. FOT's memory attention layers, with their content-based retrieval (no positional encoding on memory keys), bypass this bottleneck entirely.

The slight decline in FOT's WebQS performance at 8K (20.7% vs. 21.0% at 4K) suggests that simply adding more context does not guarantee monotonic improvement — the task may have a saturation point, or the additional demonstrations may introduce noise. The paper does not investigate this further.

Qasper Question Answering (Section 4.3, Table 6)

The paper evaluates zero-shot performance on the validation set of Qasper (questions about research papers) from SCROLLS, comparing LONGLLAMA 3B with increasing context lengths against OpenLLaMA 3B and two 7B models.

The headline from Table 6: LONGLLAMA 3B shows consistent improvement with longer context on Qasper. At 2K, LONGLLAMA matches OpenLLaMA 3B at 18.7. At 4K, LONGLLAMA reaches 20.7 (vs. OpenLLaMA at 18.7 — no improvement since it can't use longer context). At 6K, 23.2. At 8K, 26.6. This steady improvement (+7.9 percentage points from 2K to 8K) demonstrates that the model uses the extended context to better answer questions about research papers.

The comparison to 7B models: LLaMA 7B achieves 18.7 at 2K (matching the 3B models), while LongChat-7B (Ma and Zhang, 2023) achieves 19.4 at 2K, 21.2 at 4K, 25.0 at 6K, and 28.8 at 8K. LONGLLAMA 3B (at 26.6 for 8K) trails LongChat-7B (at 28.8) but outperforms LLaMA 7B at the same 2K context.

Why this matters: Qasper involves answering questions about full research papers, which is a realistic long-context application. The comparison to LongChat-7B — a model specifically designed for long contexts — shows that LONGLLAMA 3B is competitive (trailing by only 2.2 points at 8K) while being less than half the size.

Short-Context Performance Preservation (Section 4.6, Table 5)

A practical concern with long-context fine-tuning is catastrophic forgetting — the model might lose its original short-context capabilities. The paper evaluates LONGLLAMA 3B and 7B on the LM Evaluation Harness benchmark without context extension (i.e., as standard LLaMA models) and compares to the original OpenLLaMA checkpoints.

The headline from Table 5: LONGLLAMA models maintain their short-context performance essentially unchanged. Across 18 tasks/metrics spanning reasoning (ANLI, ARC), commonsense (HellaSwag, PIQA, WinoGrande), reading comprehension (BoolQ, OpenBookQA, RTE), and other capabilities (TruthfulQA, WiC), the average scores are identical: 0.53 for both OpenLLaMA 3B and LONGLLAMA 3B, and 0.55 for both OpenLLaMA 7B and LONGLLAMA 7B. Individual task differences are within expected variance (e.g., HellaSwag accuracy_norm: 0.67 → 0.65 for 3B, 0.72 → 0.71 for 7B; ANLI R2 accuracy: 0.32 → 0.33 for 3B, 0.36 → 0.37 for 7B).

Why this matters: This confirms the paper's claim that LONGLLAMA checkpoints are "a drop-in replacement for LLaMA checkpoints" — they can be used with existing inference code and maintain their original capabilities while additionally supporting long contexts. The paper explicitly states: "This also confirms that LONGLLAMAs could be used as a drop-in replacement of LLaMA models as they are compatible with the original LLaMA inference code."

FOT vs. Memorizing Transformer on arXiv — Apples-to-Apples Training Comparison (Appendix K, Tables 10 and 11)

The paper provides a more granular comparison between FOT and MT when both are trained and evaluated on the same dataset (arXiv) rather than the zero-shot transfer setting of Table 3. Both models are trained for 500K steps with local context of 2K and evaluated on arXiv in single-document mode.

The headline from Table 10 (MT): MT with 2K training memory struggles to utilize context beyond 32K tokens, with perplexity essentially flat from 32K to 128K (2.199 at 32K, 2.195 at 64K, 2.195 at 128K). MT trained with larger memories achieves better scaling: MT with 8K memory reaches 2.178 at 32K and 2.169 at 64K; MT with 16K memory reaches 2.177 at 32K and 2.166 at 64K; MT with 32K memory is comparable (2.181 at 32K, 2.168 at 64K). The paper notes "diminishing returns when scaling up the training memory length to 16K tokens and beyond."

The headline from Table 11 (FOT): FOT with $d=2,w=2$ (4K additional context, with negatives) achieves 2.148 perplexity at 128K tokens, compared to the best MT result of 2.164 with 16K memory (Table 10). This is the best perplexity achieved in the arXiv comparison. More generally, $d=2$ configurations consistently outperform $d=1$ configurations at the same training context length: $d=2,w=1$ achieves 2.171 at 128K vs. $d=1,w=1$ at 2.187; $d=2,w=2$ achieves 2.148 vs. $d=1,w=2$ at 2.152. The $w$ parameter (number of previous contexts) also matters: $w=2$ configurations outperform $w=1$ at the same $d$ : $d=1,w=2$ achieves 2.152 vs. $d=1,w=1$ at 2.187; $d=2,w=2$ achieves 2.148 vs. $d=2,w=1$ at 2.171.

Why this matters: This apples-to-apples comparison on the same dataset eliminates the zero-shot transfer confound present in Table 3 and isolates the effect of the training procedure. The finding that $d=2$ consistently beats $d=1$ at the same total differentiable context length confirms that the diversity of context (including negatives) matters above and beyond the quantity of context. The finding that FOT with 4K additional context ( $d=2,w=2$ ) beats MT with 16K memory suggests FOT's training is substantially more sample-efficient in terms of context utilization.

Ablation Studies and Robustness Checks

Differentiability of keys and values (Appendix C.1, Figures 5 and Table 8): Comparing FOT ( $d=1$ , fully differentiable through previous context) against Memorizing Transformer (non-differentiable memory) in multi-document evaluation (Figure 5) reveals that differentiability provides a substantial advantage, with FOT achieving perplexity approximately 14.4 vs. MT approximately 14.9 at 500K memory tokens. Even in the single-document setting (Table 8), FOT with $d=1$ outperforms MT at all context lengths: 14.18 vs 14.68 at 512, 14.17 vs 14.46 at 1024, 14.11 vs 14.43 at 2048. This confirms that gradient flow through the previous context's keys and values is important for learning retrieval-compatible representations, not just for handling negatives.

Importance of negatives (Appendix C.2, Figure 6): Comparing FOT with $d=1$ (differentiable, no negatives) against FOT with $d=2 \to 64$ (differentiable, with negatives) in multi-document evaluation shows that the lack of negatives causes catastrophic degradation as memory grows. FOT $d=1$ perplexity rises from approximately 14.0 to 16.0 as memory grows from 0 to 500K, while FOT $d=2 \to 64$ rises only to approximately 14.5. The single-document MT lower bound is approximately 14.1, showing that FOT $d=2 \to 64$ approaches the performance ceiling achievable with perfect distraction filtering.

Gating vs. concatenation for memory integration (Appendix B.2, Figure 4): Comparing the original Memorizing Transformer gating mechanism (sigmoid-weighted combination of memory and local values) against the simpler concatenation approach used in FOT (memory keys treated identically to local context keys) shows "no difference in performance between these two memory integration methods." Figure 4 demonstrates that both approaches produce nearly identical perplexity curves during training on PG-19 with 16K single-document memory. The paper adopts concatenation because it "does not require any architectural changes (and thus makes fine-tuning existing models easy)."

Crossbatch dimension d schedule (Appendix E.1): The paper reports that for models with $d \in \{1,2,4,8\}$ , a constant schedule is used. For $d=2 \to 64$ , the model trains with $d=2$ for 450K steps and switches to $d=64$ for the final 50K steps. For the arXiv experiments in Section 5.2, $d$ is randomly sampled from $\{2, 128\}$ in each training step — "to prevent the model from overfitting to a large additional context length during training." For the dictionary lookup task, training starts with $d=1$ until reaching 98% accuracy, then switches to $d=128$ . The consistency of the curriculum pattern (start small, increase to large) across experiments suggests that gradual exposure to negatives is a robust requirement.

Impact of training context length w (Section 5.4, Figure 9): Increasing the number of previous contexts $w$ from 1 to 2 (doubling the total differentiable context from 2×1024 to 3×1024) provides additional perplexity improvements. For $d=2$ , the $w=2$ curve is below $w=1$ at all evaluation context lengths. For $d=1$ , $w=2$ outperforms $w=1$ . An interesting comparison: $d=1,w=2$ (longer context, no negatives) slightly outperforms $d=2,w=1$ (shorter context, with negatives), suggesting that total context quantity and negative diversity are partially substitutable but both beneficial.

Number of top-k retrieved keys (Appendix F): For the dictionary lookup task, $k=32$ is used during inference. The paper does not systematically vary $k$ to measure sensitivity, but this is consistent with the $k=128$ used in most other experiments (scaled down proportionally for the smaller model/task).

Mixed crossbatch configurations for LONGLLAMA (Appendix A.2): Instead of a uniform $d$ across the batch, LONGLLAMA uses heterogeneous crossbatch configurations: segments of the batch receive different ratios of positives to negatives. For the 3B model, the configuration is $\frac{1}{4}(0,0)$ , $\frac{1}{4}(1,1)$ , $\frac{1}{4}(2,1)$ , $\frac{1}{4}(3,0)$ . The paper states this provides "more fine-grained control over the number of additional contexts and the ratio of positive to negative samples." No ablation comparing this heterogeneous setup to a uniform $d$ is provided, so the benefit of this complexity is not directly demonstrated.

Combining FOT and MT training (Appendix C.3, Figure 7): A proof-of-concept experiment trains a model for 499K steps with crossbatch and then fine-tunes with the MT objective (non-differentiable memory) for 1K steps. The resulting model achieves lower perplexity than MT trained with the same total step budget (500K steps of pure MT), with the gap being approximately 0.4 perplexity at 65K evaluation context. The paper speculates that "there may be benefits in blending these two approaches" and suggests that "MT is better at providing 'hard' negatives for the model," but leaves this for future work.

ReST $^{EM}$ revision training (Appendix K, negative result — not applicable to FOT directly but mentioned for completeness): An attempt to further optimize a revision model using ReST $^{EM}$ backfires, with "additional sequential revisions substantially hurt performance." At 256 generations, fully sequential performance drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio. This negative result highlights the sensitivity of revision training to the data generation procedure.

Variance across random seeds (Appendix J): For the arXiv baseline experiment, three runs with different random seeds yield a standard deviation of 0.002 perplexity. For the dictionary lookup task (Figure 10), error bars across 10 seeds show consistent performance with visible but small variance. The paper acknowledges that "due to resource constraints, we were unable to conduct multiple runs for all experiments" but states that "preliminary findings indicate that the observed variance was minimal compared to the impact observed from other factors under investigation."

Critical Assessment

Claim 1: The distraction issue is a primary obstacle to scaling context length, and FOT addresses it.

Assessment: This claim is well-supported for the specific experimental conditions tested, but its generality is unproven. The distraction issue is formally defined and measured (Figure 3), and the causal link between crossbatch training and mitigation is demonstrated (Figures 3, 6, 8). The monotonic relationship between training $d$ (number of negatives) and multi-document inference performance (Figures 3, 6, 8) provides converging evidence.

Limitations: The distraction issue is demonstrated only on PG-19 (books) and the synthetic setup of Figure 3. Whether the same phenomenon — and the same $r_d \approx 1/d$ pattern — holds for other domains (code, scientific text, dialogue) is not tested. The paper defines the distraction issue in terms of document boundaries, but real-world distractions are more nuanced — two documents might be topically related (making them partial positives), or a single document might contain both relevant and irrelevant sections. The paper's binary same-document/different-document framing may not capture these gradations.

Missing experiment: A measurement of positive attention mass on the arXiv, GitHub, or Isabelle datasets would strengthen the claim that the distraction issue generalizes beyond narrative text. Similarly, measuring attention mass on LONGLLAMA models (which use heterogeneous crossbatch configurations and dense attention) would verify that the mechanism transfers to the larger-scale setting.

Claim 2: FOT enables context length extrapolation far beyond training length (256K from 8K).

Assessment: Strongly supported for passkey retrieval (Figure 1: 94.5% at 100K, 73% at 256K). Supported with qualifications for language modeling (Table 3 shows perplexity improvements up to 64K from 2K training; Appendix F Figure 10 shows 16M-token retrieval from 512-token training). The extrapolation is genuine — these are not merely improvements at the training length but continued gains at lengths the model was never trained on.

Limitations: The passkey task is synthetic and tests only the most basic form of long-range retrieval: finding a single salient token in a sea of meaningless filler. It does not test whether the model can reason over long contexts, integrate information from multiple distant locations, or handle the kind of complex long-range dependencies found in real documents. The 73% accuracy at 256K, while impressive for a 32× extrapolation, also indicates a 27% failure rate — the model is not fully reliable at that scale.

The language modeling results (Table 3) show continued improvement from 2K to 64K, but the rate of improvement diminishes — on PG-19, perplexity improves from 23.74 to 22.65 (only 4.6%), compared to arXiv's 16.6% improvement. The paper does not test beyond 64K for language modeling, so we cannot tell whether the improvements saturate or continue. For the dictionary lookup task, the 92% accuracy at 16M tokens is impressive but represents a decline from the 99%+ accuracy at shorter lengths — something is degrading, even if gradually.

Missing experiment: Language modeling evaluation of LONGLLAMA at context lengths beyond 8K (say, 16K, 32K, 64K) would test whether the extrapolation observed in the analysis models scales to the larger fine-tuned checkpoints. Currently, LONGLLAMA's long-context language modeling capabilities are not evaluated — only passkey retrieval and few-shot classification. The paper also doesn't evaluate whether the model can perform multi-hop reasoning across long contexts (e.g., "what is the sum of the values in paragraph 1 and paragraph 500?"), which would be a much stronger test than single-key retrieval.

Claim 3: FOT can be applied to existing pre-trained models to extend their effective context length with minimal fine-tuning.

Assessment: Supported by the LONGLLAMA results. The 3B model is fine-tuned on 10B tokens (a small fraction of the 1T tokens used for pre-training) and achieves meaningful context extension. Table 5 demonstrates that short-context capabilities are preserved, and Table 2 shows FOT outperforms standard fine-tuning at the same budget. The backward compatibility claim is procedurally demonstrated (the checkpoints work with existing LLaMA inference code).

Limitations: The fine-tuning dataset mixture (Table 4) is complex and includes both long-context and short-context subsets with specific proportions and minimum document length filters. The paper does not ablate this mixture — we don't know whether the LONGLLAMA results depend on careful dataset curation or would transfer to any fine-tuning corpus. The 7B model is fine-tuned on only 3B tokens (vs. 10B for 3B), and the reasoning for this difference is not explained. The paper also doesn't report whether further fine-tuning would continue to improve performance or whether there are diminishing returns.

The fine-tuning uses a constant learning rate of $2 \times 10^{-5}$ , lower than OpenLLaMA's final training learning rate, but no learning rate sweep is reported. The paper doesn't investigate sensitivity to fine-tuning hyperparameters.

Missing experiment: Fine-tuning on a single dataset (e.g., only C4 or only arXiv) and measuring whether the crossbatch benefit transfers would isolate the method from the dataset mixture. The paper also doesn't test whether FOT fine-tuning benefits from longer training — the 3B model gets 10B tokens, but would 20B or 50B tokens yield further gains?

Claim 4: FOT's improvements come from the crossbatch training procedure, specifically from differentiable exposure to both positive and negative keys.

Assessment: Strongly supported by the ablation experiments. Differentiability matters (Figure 5, Table 8: FOT $d=1$ outperforms MT). Negatives matter (Figure 6: FOT $d=2 \to 64$ dramatically outperforms FOT $d=1$ in multi-document settings). Both components are necessary — neither alone achieves the full FOT performance. The ablation structure is clean and the conclusions are unambiguous.

Limitations: The paper only tests one specific implementation of each component. Differentiability is tested only through the crossbatch mechanism — there might be other ways to achieve differentiable training (e.g., backpropagating through the kNN lookup itself, which the paper doesn't explore). The negatives are always drawn from the same batch in a fixed pattern (consecutive batch indices) — the paper doesn't test whether random negatives, hard negatives, or negatives from a separate negative database would work differently.

The number of negatives $d$ is a coarse control — the paper doesn't explore whether the ratio of negatives to positives matters independently of the absolute number, or whether the token count per document (which varies) interacts with the effective $d$ . The heterogeneous crossbatch configurations for LONGLLAMA (mixing different ratios within a batch) are not ablated against a simpler uniform configuration.

Missing experiment: A comparison where negatives are provided through an explicit contrastive loss term (like TRIME or ContraCLM) rather than implicitly through the attention mechanism would clarify whether the specific crossbatch design is superior to explicit contrastive approaches, or whether both achieve similar effects through different mechanisms.

Claim 5: Content-based retrieval without positional encodings is the key architectural insight enabling unbounded context.

Assessment: Plausible and consistent with the results, but incompletely isolated. The passkey extrapolation (Figure 1) and the dictionary lookup extrapolation (Figure 10) are achieved by models with no positional encodings on memory keys. The standard fine-tuning baseline (Table 2), which retains positional encodings throughout, fails to extrapolate. However, the paper does not run a direct ablation: FOT with positional encodings on memory keys vs. FOT without. The causal role of positional encoding removal is inferred from comparison to fundamentally different training procedures (standard fine-tuning), not from an isolated variable change.

The paper cites Haviv et al. (2022) to support the claim that positional encodings are unnecessary, but Haviv et al. showed that position can be learned without explicit encodings — they didn't study context extrapolation. The paper's claim that removing positional encodings enables "theoretically unbounded context length" is a hypothesis, not a proven theorem — there could be other bottlenecks (key representation quality, attention score saturation, softmax normalization with very large key sets) that emerge at scales beyond 256K.

Missing experiment: Train FOT with and without positional encodings on the memory keys, keeping all other factors constant, and measure extrapolation. This would directly test the paper's claim. Also, evaluating at scales beyond 256K (e.g., 512K, 1M) on the passkey task would probe whether the gradual decline continues or whether there's a hard failure point.

Structural Assessment of the Experimental Design

Test set sizes and statistical reliability: The paper does not report the number of test examples for most language modeling experiments — only that "1M tokens" are used for perplexity calculation (Appendix J). For TREC, results are averaged over 20 trials with confidence intervals <1%. For Qasper, the validation set size is not specified (Qasper's validation set contains 500 questions, but the paper doesn't confirm this). The arXiv baseline experiment shows σ = 0.002 perplexity over 3 seeds, suggesting low variance, but this is only reported for one experiment. The dictionary lookup task with 10 seeds shows visible but small error bars (Figure 10). The overall pattern suggests results are reasonably robust, but the paper's transparency about statistical methodology is inconsistent.

Model scale and generality: The analysis experiments use 184M-parameter models (12 layers), which are notably small by modern standards. The LONGLLAMA experiments demonstrate scaling to 3B and 7B parameters, but only with the specific OpenLLaMA architecture. Whether FOT works as well on other model families (e.g., non-LLaMA architectures, encoder-decoder models), larger scales (13B, 70B), or models with different pre-training data distributions is not tested. The paper's claim that OpenLLaMA is "representative of the capabilities of many contemporary LLMs" is an assertion, not a demonstrated fact.

Missing comparisons: The paper does not compare FOT to several relevant long-context methods: RETRO (Borgeaud et al., 2022), which uses a two-stage retrieve-then-embed approach; Parallel Context Windows (Ratner et al., 2023), which extends context without training; or any of the sparse attention methods (Longformer, BigBird) on a long-context benchmark. A comparison on a standardized long-context benchmark like SCROLLS (beyond just the Qasper subset) or L-Eval would situate FOT more clearly in the landscape. The paper's only external comparison on a shared benchmark is Qasper vs. LongChat (Table 6).

The cost of difficulty estimation / memory maintenance is not factored into compute comparisons: FOT's inference-time kNN retrieval uses exact search (FAISS), which the paper acknowledges "is not scalable to large memory" (Section 6). The LONGLLAMA models avoid kNN entirely by using dense attention over the full memory, which would be computationally prohibitive at the 256K scale demonstrated for passkey retrieval if the memory actually contained 256K tokens worth of (key, value) pairs (as opposed to the memory accumulating during the single forward pass of the passkey task, where most tokens are repetitive filler). The paper does not report inference latency, memory overhead (storing key-value pairs for all previous tokens), or how these scale with context length — all critical practical considerations.

The paper does not demonstrate the multi-document capability of LONGLLAMA: All LONGLLAMA evaluations (passkey, TREC, WebQS, Qasper) are single-document tasks — the model processes one long prompt rather than multiple distinct documents. The multi-document language modeling evaluation that demonstrates the distraction issue and its mitigation (Section 5.3, Figures 6, 8) is only performed on the small analysis models, not on LONGLLAMA. This is a significant gap: the paper's central diagnostic framework (the distraction issue) and its primary large-model demonstration (LONGLLAMA) are evaluated on fundamentally different problem settings. We cannot tell from the presented results whether LONGLLAMA actually solves the multi-document distraction problem that motivated FOT, or whether it only achieves single-document extrapolation through a different mechanism (e.g., better key structure for long single-document retrieval, which is a weaker test of the distraction hypothesis).

The paper's contributions to long-context language modeling and the specific issue of distraction are clearly demonstrated, but the link between the problem diagnosis and the large-scale solution is incompletely bridged in the experimental evidence presented.

6. Limitations and Trade-offs

Multi-Document Distraction: Diagnosed but Not Demonstrated at Scale

The assumption or constraint. The paper's central diagnostic framework — the distraction issue — is defined, measured, and mitigated exclusively in the context of multi-document scenarios where memory persists across unrelated documents. Section 3.3 formalizes positive attention mass $r_d$ and shows that FOT training with large $d$ prevents $r_d$ from collapsing to $1/d$ . However, every LONGLLAMA evaluation (passkey retrieval, TREC, WebQS, Qasper, LM Evaluation Harness) is a single-document task — the model processes one prompt or document at a time, and there is no cross-document memory sharing. The multi-document language modeling experiments demonstrating the distraction issue and its mitigation (Section 5.3, Figures 6 and 8) are performed only on the 184M-parameter analysis models trained from scratch on C4, not on the 3B or 7B LONGLLAMA checkpoints.

The consequence. A practitioner deciding whether to deploy LONGLLAMA for a multi-document application (e.g., repository-level code generation where the model must attend to many files simultaneously, or multi-source question answering where context contains dozens of unrelated articles) has no direct evidence that the fine-tuned models actually solve the distraction problem that motivated FOT. The passkey task embeds a single salient token in homogeneous filler text — there are no document boundaries to respect and no semantic distractors. The TREC and WebQS tasks provide in-context examples from a single training set, not multiple unrelated documents. The Qasper task involves a single research paper. It is entirely possible that LONGLLAMA achieves its extrapolation through improved key structure for single-document retrieval (a weaker requirement than cross-document discrimination) and would still suffer from distraction collapse when faced with many semantically unrelated documents in memory — the very failure mode that Figure 3 diagnoses in standard transformers. The paper's headline claim about solving the distraction issue and its large-model demonstration remain evaluated on fundamentally different problem settings.

What evidence exists in the paper. The analysis-scale experiments in Section 5.3 (Figures 6 and 8) demonstrate that FOT's multi-document performance scales with training $d$ and that FOT with $d=2 \to 64$ approaches the single-document MT lower bound. This is rigorous evidence at the 184M-parameter scale on PG-19. Section 4 contains no multi-document evaluation. The LONGLLAMA training mixture (Table 4) includes multiple document sources (arxiv, python, book, common_crawl, c4, github, stackexchange, wikipedia), so the models were exposed to multi-domain data, but whether this translates to multi-document inference capability is not measured. The heterogeneous crossbatch configurations for LONGLLAMA (Appendix A.2: $\frac{1}{4}(0,0), \frac{1}{4}(1,1), \frac{1}{4}(2,1), \frac{1}{4}(3,0)$ for 3B) include segments with negatives ( $(1,1)$ and $(2,1)$ ), suggesting the training procedure should produce discrimination ability, but this is not verified at inference time.

Mitigation status. Not addressed. The paper does not acknowledge this evaluation gap as a limitation. Section 6 ("Limitations and future work") discusses scaling up context, scaling up crossbatch, exploring contrastive learning, and combining with other methods, but none of these address the disconnect between the diagnostic framework and the large-model evaluation. A multi-document evaluation of LONGLLAMA — for instance, measuring whether the model can answer questions about a specific document when the memory contains dozens of other documents, or measuring positive attention mass directly on the fine-tuned checkpoints — would close this gap but is left to future work.

Training Cost Overhead: Crossbatch Increases Compute in Memory Layers, Not Quantified

The assumption or constraint. FOT's crossbatch training procedure expands the key-value set that memory attention layers attend to from the local context only to the local context plus $d$ additional contexts (each containing up to a full context window of tokens). This increases the attention computation cost in those layers by a factor proportional to $(d+1)$ during training. The paper acknowledges this implicitly in Section 6: "crossbatch increases the training cost, but only in a subset of layers." However, the paper provides no quantification of this overhead — no FLOPs comparison between FOT training and standard training, no measurement of training throughput reduction, and no analysis of how the overhead scales with $d$ , batch size, or model size.

The consequence. A practitioner considering FOT fine-tuning for their own model cannot estimate the additional training cost, making it difficult to evaluate the cost-benefit tradeoff against alternatives like standard long-context fine-tuning (which has no crossbatch overhead) or position interpolation methods (which modify positional encodings without changing attention computation). The paper states that overhead is "only in a subset of layers" — for the 184M analysis models, this is 1 layer out of 12 (~8% of layers); for LONGLLAMA 3B, it is 3 layers out of 26 (~11.5%); for LONGLLAMA 7B, 3 layers out of 32 (~9.4%). But the attention cost increase within those layers is multiplicative: if $d=64$ and each additional context has the same token count as the local context, the attention computation in those layers grows ~65× relative to standard training for those layers. The net impact depends on what fraction of total training FLOPs the memory layers' attention represents, which the paper does not break down. The paper also does not report whether crossbatch training affects training stability, requiring more careful learning rate tuning or more training steps to converge.

What evidence exists in the paper. The paper reports only the high-level training recipe: batch sizes in tokens (32K for analysis, 128K for Section 5.2/5.4, 256K for LONGLLAMA), number of training steps, and learning rates. There is no training throughput measurement, no FLOPs comparison, and no wall-clock time comparison between FOT and baseline training runs. The maximum $d$ values used ( $d=64$ or $d=128$ ) are described as "the maximum value that fits into the memory of a single TPUv3/TPUv2 machine" (Section 6), indicating that memory constraints, not compute, are the binding constraint for the reported configurations. The paper notes in Section 6 that "in future work, we want to further increase $d$ as well as test on devices with bigger memory or utilize multi-node training," suggesting the current $d$ ceiling is hardware-imposed.

Mitigation status. Partially acknowledged but not quantified. Section 6 lists "Scaling up crossbatch" as a future work direction and notes that increasing $d$ requires more device memory, but the specific compute overhead of the current approach is not reported. The paper's Listing 1 and Listing 2 show that the crossbatch implementation adds modest code complexity, but code simplicity does not imply computational efficiency. A FLOPs comparison or throughput measurement for FOT training vs. standard training, even for a single representative configuration, would substantially improve the practical guidance.

Inference Infrastructure: kNN Retrieval Limits Scalability, Dense Attention Is Impractical

The assumption or constraint. FOT's inference-time architecture relies on a kNN retrieval mechanism to select the top- $k$ most similar keys from the external memory before computing attention. The paper uses exact kNN search implemented in FAISS for all analysis experiments, which computes the inner product between the query and every key in memory to find the true top- $k$ . This is computationally linear in memory size. Section 6 explicitly acknowledges this: "In our experiments, we use the exact kNN search, which is not scalable to large memory. Using approximate kNN search will require a lot of engineering effort, as well as careful evaluation of the impact of the approximation on the model performance."

For the LONGLLAMA models, the paper replaces kNN with dense attention over the full memory: every query attends to every key in memory. This is computationally quadratic in memory size and practically infeasible at the scale of the demonstrated capabilities (256K tokens means 256K key-value pairs per layer per head for three memory layers).

The consequence. Neither inference approach described in the paper scales to the context lengths that the paper's headline results (256K passkey retrieval, 16M dictionary lookup) demonstrate. The exact kNN approach used for analysis models would require computing and sorting millions of inner products per query at 16M memory size — feasible for a benchmark but impractical for production deployment with real-time latency requirements. The dense attention approach used for LONGLLAMA would require computing attention over 256K key-value pairs per query per head per memory layer, which is exactly the quadratic complexity that sparse attention methods were designed to avoid. The paper's "theoretically unbounded context length" claim is architectural (no positional encoding bottleneck) but the demonstrated results depend on computational approaches that do not scale to the claimed regime. A practitioner cannot deploy LONGLLAMA with 256K effective context in a latency-sensitive application using either of the described inference methods.

The dictionary lookup result (Figure 10) with 16M tokens at 92% accuracy is particularly telling: this was achieved with exact kNN ( $k=32$ ) on the analysis models, meaning each query required computing 16M inner products. The paper does not report the wall-clock time for this retrieval, but at 16M vectors of embedding dimension 512–1024, this is a substantial computation per query.

What evidence exists in the paper. Section 6 explicitly acknowledges the scalability limitation of exact kNN and the need for approximate kNN in future work. Appendix I describes the hardware used (TPUv2/TPUv3 with 64GB/128GB device memory) but does not report inference latency. The paper states for LONGLLAMA that "we use dense attention instead of the kNN retrieval, as we found only marginal performance differences and it is simpler to implement" (Appendix A.2), but does not discuss the computational implications of this choice for the 256K context length. There is no measurement of tokens-per-second throughput or memory overhead for storing the key-value pairs (which grow linearly with tokens processed).

Mitigation status. Acknowledged as future work. Section 6 lists "Scaling up context" as the "most important future research direction" and notes that "storing more than 16M (key, value) pairs will require a distributed multi-node system" and that "using approximate kNN search will require a lot of engineering effort, as well as careful evaluation of the impact of the approximation on the model performance." The paper does not provide even a preliminary experiment with approximate kNN (e.g., IVF indices, HNSW graphs) to bound the performance impact of approximation. The LONGLLAMA release uses dense attention, which the paper implicitly acknowledges is not the intended production inference method but rather a simplification for the research release.

Single Model Family, Single Fine-Tuning Recipe: Generality Is Asserted, Not Demonstrated

The assumption or constraint. All experiments demonstrating FOT's effectiveness use transformer architectures from a single family: custom decoder-only models (184M, 37M parameters) and OpenLLaMA (3B, 7B parameters). OpenLLaMA is a specific reproduction of LLaMA with rotary positional embeddings, RMSNorm pre-normalization, SiLU activations, and a SentencePiece tokenizer. The paper claims this model is "representative of the capabilities of many contemporary LLMs" (Section 4 introduction), but provides no evidence that FOT works on other architectures (encoder-decoder models like T5, non-LLaMA decoder-only models like GPT-NeoX or MPT, mixture-of-experts architectures), other model scales (13B, 70B, or larger), or models with substantially different pre-training data distributions. The fine-tuning recipe is also fixed: a specific dataset mixture (Table 4), a specific learning rate ( $2 \times 10^{-5}$ ), a specific heterogeneous crossbatch configuration, and a specific number of fine-tuning tokens (10B for 3B, 3B for 7B). No sensitivity analysis is performed on any of these choices.

The consequence. A practitioner with a non-LLaMA model (e.g., a T5-based system for long-document summarization, a GPT-NeoX model trained on a different data distribution, or a domain-specific model fine-tuned on medical or legal text) has no guidance on whether FOT will work for their architecture. The distraction issue is a general phenomenon — it should affect any transformer with an attention mechanism — but whether crossbatch training successfully mitigates it may depend on architectural details (pre-norm vs. post-norm, type of positional encoding, activation function, tokenizer properties). The specific FOT design choices (which layers to use as memory layers, how to handle positional encodings, what $d$ schedule to use, what ratio of positives to negatives in crossbatch) were tuned for the specific models tested — there is no evidence these choices transfer. The paper does not explore whether the number of fine-tuning tokens matters (would 1B tokens suffice? would 50B tokens yield substantial further gains?), whether the learning rate is critical, or whether the heterogeneous crossbatch configuration is superior to a simpler uniform $d$ .

What evidence exists in the paper. All analysis experiments use the same 12-layer, 184M-parameter architecture with the same hyperparameters (Appendix E, Table 9). All LONGLLAMA experiments use OpenLLaMA 3B and 7B with the same architectural family. The paper does not cite or conduct experiments on T5, GPT-NeoX, MPT, Falcon, or any other model family. The short-context capability preservation (Table 5) is measured only on LONGLLAMA, not on the analysis models. The fine-tuning dataset mixture (Table 4) is specific to LONGLLAMA — the analysis models use C4 for pre-training and fine-tuning, a different data distribution. There is no cross-family validation and no ablation of the fine-tuning recipe.

Mitigation status. Not addressed. The paper does not discuss architecture dependence as a limitation. The claim that OpenLLaMA is "representative" is stated without qualification or evidence in Section 4. The paper's focus on a single model family is understandable given resource constraints (fine-tuning 3B and 7B models is already computationally expensive), but the absence of any discussion of this limitation, or any suggestion for how practitioners should adapt FOT to other architectures, is a gap in practical guidance.

Difficulty Estimation and Dynamic Strategy Selection: No Mechanism for Adaptive Allocation at Inference Time

The assumption or constraint. FOT uses a fixed inference procedure: for each query in a memory attention layer, retrieve the top- $k$ most similar keys from memory and attend to them alongside the local context. This procedure is applied uniformly to every query, every token, and every document, regardless of content or difficulty. There is no mechanism for the model to dynamically decide whether to use memory for a particular query (e.g., skipping retrieval when the local context is sufficient), how many keys to retrieve (adapting $k$ based on query difficulty or memory relevance), or which subsets of memory to search (e.g., restricting retrieval to a specific document or topic when the query indicates it). The crossbatch training teaches the model to discriminate relevant from irrelevant keys, but at inference time the kNN retrieval is a hard, non-learned selection step that returns exactly $k$ keys regardless of whether fewer would suffice or more would be beneficial.

The consequence. In a deployment where the model processes diverse inputs — some queries that need long-range retrieval (e.g., looking up a definition provided 50 pages ago) and some that only need local context (e.g., continuing a sentence with obvious next words) — the uniform retrieval mechanism incurs a fixed computational cost per query regardless of need. For queries where no relevant keys exist in memory (e.g., the model is processing a completely new topic unrelated to anything in memory), the top- $k$ retrieved keys are distractors by construction, yet the model is forced to attend to them. The paper showed that FOT training helps the model ignore these distractors (high positive attention mass despite many irrelevant keys), but it does not eliminate them — the irrelevant keys still consume attention mass and computational resources. A dynamic mechanism that could skip retrieval when the query-to-key similarities are uniformly low (indicating no relevant memory content) could save computation and reduce distraction simultaneously, but such a mechanism is not developed or evaluated.

What evidence exists in the paper. The paper shows that FOT maintains high positive attention mass even as the number of distractors grows (Figure 3: $r_d \approx 0.8$ at $d=64$ for $d=2 \to 64$ ). This demonstrates that the model can ignore distractors, not that it avoids processing them. Figure 8 shows that perplexity increases with memory size even for FOT — the $d=2 \to 64$ curve rises by approximately 0.18 perplexity from 0 to 500K tokens, showing that distractors have a cost even when the model is well-trained to handle them. The paper does not report the distribution of retrieved key similarities (e.g., what fraction of queries have their top-1 retrieved key similarity below some threshold, indicating "no good match"), nor does it explore whether retrieval could be skipped when similarities are low. The $k$ parameter is fixed at 128 for most experiments and 32 for the dictionary lookup task, with no sensitivity analysis.

Mitigation status. Not addressed. The paper does not discuss adaptive retrieval or dynamic $k$ selection as a limitation or future direction. Section 6 focuses on scaling up context and crossbatch rather than on making retrieval more efficient or selective. This is a missed opportunity: the crossbatch training procedure produces key representations that are specifically designed to support discrimination between relevant and irrelevant keys, and the similarity scores produced by these representations could naturally support a learned or threshold-based gating mechanism for retrieval. The absence of such a mechanism means the model always pays the full computational cost of retrieval even when it provides no benefit.

Latency vs. Throughput: Sequential Memory Population Introduces Serial Dependency

The assumption or constraint. The inference-time memory mechanism is inherently sequential: as the model processes tokens, the memory attention layer's (key, value) outputs are appended to the memory store, and subsequent queries can retrieve from the updated memory. This means that token processing cannot be fully parallelized across a long sequence — each token's memory attention computation depends on the memory state after all previous tokens have been processed. This is the same serial dependency that affects all autoregressive generation, but FOT adds an additional serial step: the kNN index must be updated (or at minimum, the new key-value pair must be added to the searchable set) before the next query can use it. The paper's passkey retrieval task processes the entire long prompt in a single forward pass (since the prompt is known in advance, key-value pairs for all tokens can be computed and indexed before the query token is processed). But for autoregressive generation — the primary use case for language models — tokens are generated one at a time, and each new token's key-value pair must be added to memory before the next token can retrieve from it.

The consequence. In an autoregressive generation setting (e.g., the model is writing a long document and needs to retrieve from its own earlier output), the memory population creates a latency bottleneck that compounds with generation length. Each decoding step requires: (a) computing the query, key, and value for the new token, (b) inserting the (key, value) into the FAISS index (or appending to the dense attention key-value store), and (c) performing kNN search over the updated index for the memory attention layers. Steps (b) and (c) are serial dependencies that cannot be parallelized across tokens. For the dictionary lookup task (deterministic key-value retrieval, no generation), this serial dependency is not limiting because all definitions are known before queries are processed. For passkey retrieval, the entire prompt is available in advance, so memory can be populated in a single parallel encoding pass. But for the most common deployment scenario — the model generating text autoregressively while using its own previous output as extended context — the paper provides no latency measurements or analysis of how the memory mechanism interacts with autoregressive decoding.

The paper's Qasper evaluation (Table 6) approximates a real deployment scenario most closely: the model reads a research paper and answers questions. However, the specific inference protocol (is the paper encoded once and the memory reused across multiple questions? is retrieval performed during encoding of the paper or only during question answering?) is not described, making latency implications impossible to assess.

What evidence exists in the paper. The paper reports no inference latency, no tokens-per-second throughput, and no analysis of how memory population interacts with autoregressive decoding. The evaluation protocols for language modeling (perplexity over 1M tokens) measure next-token prediction accuracy but do not report computational cost. The passkey retrieval task uses a known prompt structure, allowing parallel encoding. Appendix I mentions hardware (TPUv2/TPUv3 with up to 128GB device memory, 96 CPU cores, 300GB RAM) but no latency or throughput measurements. The LONGLLAMA release uses dense attention over the full memory, which would have latency characteristics very different from kNN-based retrieval (dense attention is parallelizable over the memory dimension but quadratic in memory size).

Mitigation status. Not addressed. The paper does not discuss latency or throughput as a limitation. Section 6 lists "Scaling up context" as a future direction focused on memory capacity ("storing more than 16M (key, value) pairs will require a distributed multi-node system") and approximate kNN for search efficiency, but does not address the serial dependency of memory population in autoregressive settings. A common mitigation in retrieval-augmented models is to encode the entire context in a single parallel pass and only perform retrieval during generation (as in RETRO), but FOT's design — where memory is populated incrementally — makes this optimization non-trivial. The paper does not discuss whether asynchronous memory updates (allowing generation to proceed with a slightly stale memory) would be acceptable for performance.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a representational reframing of the long-context problem that shifts attention from computational efficiency to the quality of the key-value embedding space. Before FOT, the dominant narrative in long-context transformer research was that the bottleneck was computational — the quadratic complexity of self-attention needed to be reduced through sparse patterns, compression, or hierarchical architectures. The distraction issue provides a competing diagnosis: even if you could compute full attention over an infinite context, a standard transformer would not use it effectively because its keys from different documents collapse into indistinguishable representations. The $r_d \approx 1/d$ result in Figure 3 makes this diagnosis falsifiable and quantitative in a way that prior efficiency-focused work never achieved.

This is not a paradigm shift — FOT builds directly on the Memorizing Transformer architecture and uses a standard language modeling loss — but it is a conceptual pivot with practical consequences. The finding that crossbatch training, without any long-sequence training or explicit contrastive loss, can produce key representations that support retrieval across a 31,250× train-test context gap (Appendix F, Figure 10) challenges the implicit assumption that models need to practice retrieval at scale during training. The separation between learning to represent (which crossbatch accomplishes with modest training context) and learning to retrieve (which is handled by the non-learned kNN mechanism at inference time) is a cleaner decomposition than prior memory-augmented approaches and suggests that the training efficiency barrier for long-context models is lower than previously thought.

The paper also reconciles a tension between two observations in the literature. Haviv et al. (2022) showed that transformers can learn positional information without explicit positional encodings, but this was a representational curiosity rather than a practical architecture principle. The position-interpolation methods (Chen et al., 2023; kaiokendev, 2023) demonstrated that positional encodings could be manipulated to extend context, but only by a factor of ~4–8×. FOT bridges these findings: by removing positional encodings from long-range attention entirely (following Haviv et al.'s implication that they are unnecessary) and relying on content-based retrieval, the model achieves a 32× extrapolation (8K training → 256K evaluation) on passkey retrieval, an order of magnitude beyond position-interpolation methods. The resolution is that positional encodings are not just unnecessary — they are actively harmful for extrapolation, because they impose a hard ceiling on the representable context length. Content-based retrieval removes this ceiling.

The research directions this work makes more attractive include contrastive representation learning for attention mechanisms (rather than for output embeddings, as in most prior work), implicit contrastive objectives that achieve discrimination through architecture design rather than auxiliary loss terms, and decoupling representational quality from sequential training length in retrieval-augmented models. Directions that become relatively less attractive include pure positional-encoding-based context extension (which FOT surpasses by an order of magnitude) and memory-augmented training schemes that require long documents (FOT shows these are unnecessary for learning retrieval-compatible representations).

Follow-Up Research This Work Enables

Measuring positive attention mass on LONGLLAMA models in multi-document settings. The paper's central diagnostic — the distraction issue and its mitigation via crossbatch — is demonstrated only on 184M-parameter analysis models trained from scratch on C4. The 3B and 7B LONGLLAMA checkpoints, which are the paper's headline contribution and the models released to the community, are never evaluated in a multi-document setting where cross-document distraction actually occurs. A direct experiment would replicate Figure 3 (positive attention mass $r_d$ vs. number of documents $d$ ) on LONGLLAMA 3B, measuring whether the heterogeneous crossbatch configuration used during fine-tuning (Section A.2: $\frac{1}{4}(0,0), \frac{1}{4}(1,1), \frac{1}{4}(2,1), \frac{1}{4}(3,0)$ ) produces attention focusing comparable to the $d=2 \to 64$ curve in the analysis experiments. A negative result ( $r_d \approx 1/d$ at scale despite crossbatch training) would indicate that the fine-tuning recipe does not transfer the distraction mitigation from small models to large ones, potentially due to the heterogeneous crossbatch configuration, the smaller effective $d$ , or interference from the retained positional encodings in the local context. A positive result would validate the paper's central claim at a practically relevant scale.

Approximate kNN with performance-bounded degradation. The paper uses exact kNN search (FAISS) for all analysis experiments and acknowledges that this "is not scalable to large memory" (Section 6), while the LONGLLAMA release uses dense attention — both approaches are computationally infeasible for deployment at the 256K–16M scales the paper demonstrates. A critical follow-up would train a FOT model (e.g., the 184M architecture) and evaluate it with approximate nearest neighbor indices (IVF with varying numbers of centroids, HNSW with varying graph connectivity) across the same range of memory sizes used in Figures 6 and 8. The key question is whether FOT's trained key representations — which are specifically optimized for inner-product-based retrieval — are robust to the approximation error introduced by these indices. The paper's finding that FOT maintains $r_d \approx 0.8$ even at $d=64$ (Figure 3) suggests the key space has strong cluster structure, which should make approximate indices effective, but this is an empirical question. The experiment could identify the minimum index configuration (and corresponding speed-accuracy tradeoff) that maintains perplexity within 0.1 of the exact kNN baseline at 500K memory tokens, providing practical guidance for deployment.

Dynamic k-selection based on query-memory similarity distribution. FOT currently uses a fixed $k=128$ (or $k=32$ for dictionary lookup) for all queries, regardless of whether the retrieved keys are genuinely relevant or are distractors returned because no better matches exist. A natural extension would add a learned or heuristic mechanism to dynamically select $k$ per query based on the distribution of retrieved key similarities. For example, if the top-1 similarity is below a threshold, or if the top- $k$ similarities are tightly clustered (indicating no single strong match), the model could reduce $k$ or skip memory attention entirely for that query, falling back to local context only. This would build on the paper's own observation that "the irrelevant keys still consume attention mass and computational resources" even for well-trained FOT models — the $d=2 \to 64$ curve in Figure 8 still shows a small perplexity increase (+0.18) when memory grows from 0 to 500K tokens. A simple threshold-based ablation (varying the similarity threshold and measuring the tradeoff between perplexity and average $k$ ) on the PG-19 multi-document setup would quantify how much computation could be saved without meaningful accuracy loss. A more ambitious version could train a lightweight gating network that takes the top- $k$ similarity distribution as input and outputs a binary "use memory" decision, trained with a combined loss that penalizes both computation and perplexity degradation.

FOT applied to retrieval-augmented generation (RAG) for open-domain QA. The paper evaluates FOT exclusively on language modeling (perplexity) and synthetic retrieval (passkey, dictionary lookup), with one short-context QA task (Qasper). A direct extension would test whether FOT's improved key representations improve the retrieval quality in a standard RAG pipeline: encode a large corpus of documents (e.g., Wikipedia, or the RedPajama corpus used for LONGLLAMA training) into the FOT-trained key space of a memory attention layer, use kNN retrieval at inference time to select relevant documents for a query, and measure downstream QA accuracy (e.g., on Natural Questions or TriviaQA). The comparison would be against a baseline using the same model architecture but with standard (non-FOT) key representations for retrieval. This experiment tests a different claim than the paper makes — FOT is designed to extend context, not to improve retrieval for a separate QA task — but it probes whether the key structure improvements transfer to a fundamentally different retrieval setting. A positive result would substantially broaden FOT's applicability; a negative result would clarify that the key structure is optimized specifically for within-document or within-context retrieval, not for cross-document semantic search.

Scaling FOT to 13B+ models with controlled fine-tuning budget experiments. The paper demonstrates FOT on 3B and 7B models with a fixed fine-tuning budget (10B and 3B tokens respectively), but provides no scaling analysis — we do not know whether the benefit of FOT fine-tuning increases, decreases, or remains constant with model scale, nor whether the optimal fine-tuning token count scales with model size. A systematic experiment would fine-tune OpenLLaMA models at 1B, 3B, 7B, and 13B parameters with FOT, varying the fine-tuning token count (e.g., 1B, 3B, 10B, 30B tokens) and measuring passkey retrieval accuracy at 8K, 32K, 128K, and 256K context lengths. The output would be a set of scaling curves showing how FOT's extrapolation capability depends on both model scale and fine-tuning budget. This would address the paper's unexamined claim that OpenLLaMA is "representative" and provide practical guidance for practitioners deciding how much fine-tuning is worth the compute investment. A finding that FOT benefit saturates quickly (e.g., 3B tokens is as good as 30B tokens for the 7B model) would make the method substantially more attractive than if benefit continues to scale with budget.

Multi-document reasoning benchmark for LONGLLAMA. The paper currently evaluates LONGLLAMA only on single-document tasks, leaving the multi-document distraction mitigation unverified at scale. A strong follow-up would construct a controlled multi-document reasoning task: present the model with a memory containing $N$ documents ( $N \in \{2, 4, 8, 16, 32\}$ ), one of which contains information needed to answer a factoid question, and measure accuracy as a function of $N$ . The documents should be semantically similar (e.g., all are Wikipedia articles about related topics, or all are research paper abstracts in the same field) to avoid trivial distinguishability. The key measurement would be whether LONGLLAMA maintains accuracy above 90% as $N$ grows from 2 to 32, compared to the baseline OpenLLaMA. The paper's analysis-scale results (Figures 6 and 8) predict that FOT's accuracy should degrade much more gracefully than the baseline, but this has not been demonstrated at the 3B/7B scale. A negative result (LONGLLAMA degrades similarly to the baseline in multi-document settings) would indicate that the crossbatch configuration used for fine-tuning does not transfer multi-document discrimination to the larger models, potentially because the heterogeneous batch segments with no negatives dilute the contrastive signal.

Practical Applications and Downstream Use Cases

Long-document question answering with fine-grained retrieval. A system that needs to answer questions about very long documents — legal contracts (hundreds of pages), technical documentation (thousands of pages of API references), or scientific literature reviews (tens of papers synthesized) — can use LONGLLAMA to encode the entire document set into memory and then answer queries with attention over both the local question context and the retrieved relevant passages. The Qasper result (Table 6: LONGLLAMA 3B improves from 18.7 at 2K to 26.6 at 8K context when answering questions about research papers) provides a proof of concept, though at modest scale. At the demonstrated 256K-token passkey retrieval capability (Figure 1), a single document or document collection of ~200,000 words could be held in memory, enabling question answering without chunking, sliding windows, or separate retrieval pipelines. The key practical advantage over standard RAG approaches is that retrieval and reasoning share the same key-value representations — the model retrieves what it finds useful for next-token prediction, rather than what a separately-trained retriever thinks is relevant.

Repository-level code generation with cross-file context. A coding assistant that needs to generate or edit code in one file while being aware of definitions, imports, and conventions across hundreds of other files in a repository can use FOT's multi-document memory to maintain key-value pairs from all processed files. The GitHub dataset used in the paper's language modeling evaluation (Table 3) already concatenates files from the same repository, simulating this setting. The multi-document training with negatives is specifically designed for this scenario: the memory attention layer learns to retrieve keys from the current file (positive) while ignoring keys from unrelated repositories or library code (negatives). At the demonstrated 16K–64K context range from the analysis experiments, a memory containing tens of thousands of lines of code is feasible. The paper's finding that FOT with $d=2 \to 64$ approaches the single-document lower bound in multi-document perplexity (Figure 8: within 0.2 perplexity of the single-doc MT baseline at 500K tokens) suggests that cross-file distraction — where keys from unrelated files compete for attention — would be substantially mitigated compared to a standard model, though this specific use case is not directly evaluated in the paper.

Passkey-like information retrieval from very large structured knowledge bases. The dictionary lookup result (Figure 10: 92% accuracy with 16M key-value pairs) demonstrates FOT's capability for content-addressable memory at scale. A practical deployment could encode a large structured knowledge base — product catalogs with millions of SKUs and their specifications, legal or regulatory databases with hundreds of thousands of clause-requirement mappings, or medical coding databases with diagnosis-to-code mappings — into FOT's key-value memory. Queries could then retrieve the relevant entry by content similarity rather than exact key matching. The 31,250× train-test context gap (512 tokens → 16M tokens) provides evidence that the key structure generalizes to retrieval scales far beyond training, meaning the system could be fine-tuned on a modest subset of the knowledge base and deployed on the full dataset. The critical advantage over standard vector database approaches is that the key representations are optimized for the model's own retrieval mechanism and the language modeling objective, potentially producing more relevant retrievals than a separately trained embedding model. However, the paper does not evaluate retrieval precision/recall directly — only end-to-end accuracy on the dictionary task — so the standalone retrieval quality is uncharacterized.

When to Prefer This Method

The paper does not provide an explicit decision framework comparing FOT against named alternatives (position interpolation, landmark attention, sparse attention methods) with quantified tradeoffs. The only direct comparison is against standard long-context fine-tuning (Table 2), where FOT achieves both better in-distribution performance and extrapolation capability. The paper positions FOT as complementary to other methods ("we believe that some of these methods could be combined with FOT, resulting in mutually beneficial interactions," Section 6) rather than as a replacement for them. Given the absence of systematic comparison data against alternatives, a "Prefer A when... Prefer B when..." matrix would be speculative rather than paper-grounded, and is not included here.