LLMs from Scratch  ·  Part 8 of 12: Long Context and Memory ← Part 7
LLMs · 11 min read

Long Context and Memory in LLMs: How Models Handle More

Arjun Jaggi · September 17, 2026 · LLMs from Scratch, Part 8 of 12

The original transformer had a fixed context window determined by learned absolute position embeddings. Extending that window requires rethinking how position information is encoded and how attention is computed across thousands or millions of tokens. This part covers the key techniques that made modern long-context models possible and what their practical limits actually are.

O(n²)
the computational complexity of standard full attention with respect to sequence length n; doubling sequence length quadruples attention compute cost
RoPE
Rotary Position Embedding: the position encoding scheme used in LLaMA 2, Mistral, and most modern long-context models; allows extension beyond training length via interpolation (arXiv:2104.09864)
Lost-in-the-middle
the empirical finding that models perform better on information near the beginning or end of a long context than on information in the middle (arXiv:2307.03172)

The Context Window Problem

In the original transformer architecture, position was encoded using learned absolute position embeddings: a lookup table mapping each position index to a learned vector that was added to the token embedding before attention was computed. These learned embeddings worked well at the sequence lengths the model was trained on, but they could not be extrapolated to longer sequences without retraining, because the model had never seen the position indices beyond its training length.

This created a hard cap on how much context a model could process. Early GPT models were trained at 512 to 2048 tokens. For many use cases, this was sufficient. But for tasks like processing legal contracts, analyzing medical records, reasoning over long code repositories, or maintaining coherent conversation across many turns, the limit became a significant constraint.

Rotary Position Embeddings

Rotary Position Embeddings (RoPE, Su et al., arXiv:2104.09864) encode position as a rotation applied to query and key vectors before the attention dot product. Instead of adding a position vector, RoPE multiplies by a rotation matrix parameterized by the position index. The key property is that the inner product between a query at position m and a key at position n depends only on the relative displacement m minus n, not on the absolute positions themselves. This makes position encoding relative by construction.

RoPE rotation: q_m = R(m) * q_m Attention score: q_m^T k_n = q^T R(m-n) k Where R(theta) is a block-diagonal rotation matrix with angle theta

RoPE's rotational structure enables extrapolation beyond training length using a technique called position interpolation: instead of using the original position angles for positions beyond the training length, the position indices are scaled down to fit within the trained range. This allows a model trained at 4096 tokens to be used at 8192 or 16384 tokens without complete failure, though with some degradation that is partially addressed by a small amount of fine-tuning at the extended length.

ALiBi and Learned Bias Approaches

ALiBi (Attention with Linear Biases, Press et al., arXiv:2108.12409) takes a different approach. Instead of encoding position in the token representation, it adds a linear penalty to the attention score based on the distance between query and key positions. The penalty is proportional to the distance: the further apart two positions are, the more the attention score is reduced. Different attention heads use different penalty slopes, learned as parameters.

The advantage of ALiBi is that the penalty is unbounded and monotone: the model naturally pays less attention to tokens that are further away, which generalizes beyond the training length without interpolation. The disadvantage is that ALiBi's strong recency bias may not suit all tasks, and it performs somewhat worse than RoPE-based models on tasks requiring attention to distant context.

Sliding Window Attention

Mistral's sliding window attention (Jiang et al., arXiv:2310.06825) addresses the O(n²) compute cost of full attention by restricting each token to attend only to a fixed window of the most recent tokens. In Mistral 7B, each token attends to a window of 4096 tokens, regardless of the total sequence length. For sequences longer than the window, information flows between non-adjacent positions through intermediate tokens across multiple layers.

The effective receptive field grows with depth: after L layers with window size W, a token can indirectly incorporate information from a span of W times L tokens. For a 32-layer model with window size 4096, the effective receptive field at the final layer is over 130,000 tokens, even though no single attention operation spans that range. This makes sliding window attention computationally linear in sequence length while still enabling information to propagate across long distances through the layer stack.

A long context window and good long-context performance are not the same thing. Models with million-token windows often fail to reliably retrieve information from the middle of that range.

The Lost-in-the-Middle Problem

Liu et al. (arXiv:2307.03172) showed that models with long context windows exhibit a U-shaped performance curve when relevant information is placed at different positions within the context. Performance is highest when the relevant information is at the very beginning or very end of the context, and substantially lower when it is in the middle. This pattern was robust across model families, context lengths, and task types.

The practical implication is that simply providing more context does not guarantee that a model will use it well. For retrieval-augmented generation systems, placing the most important retrieved passages at the beginning or end of the prompt rather than in the middle can substantially improve performance. For multi-document reasoning tasks, the order in which documents are arranged in the prompt affects output quality independently of their content.

KV Cache and Inference Memory

At inference time, long context creates a memory constraint through the key-value (KV) cache. During autoregressive generation, the model caches the key and value matrices computed for all previous tokens so they do not need to be recomputed at each step. The size of this cache grows linearly with sequence length and with the number of attention layers and heads.

For a model with 32 layers, 32 attention heads, and a head dimension of 128, each token in the context requires roughly 32 times 2 times 32 times 128 times 2 bytes (for bfloat16) of KV cache storage, which is approximately 512 KB per token. At 100,000 tokens, this amounts to roughly 50 GB of KV cache alone. Managing this memory efficiently is one of the primary engineering challenges in deploying long-context models at scale.

Grouped-Query Attention (GQA, Ainslie et al., arXiv:2305.13245) reduces KV cache size by sharing key and value projections across multiple query heads. Instead of each of the 32 query heads having its own key and value matrices, GQA groups query heads and assigns one key-value pair per group. With 8 key-value heads serving 32 query heads, the KV cache is reduced by a factor of 4 with minimal impact on model quality. Most modern long-context models use GQA for this reason.

Memory Budget Constraints in Practice

The interaction between context length and inference cost creates a resource allocation problem for enterprise deployments. BudgetBench (Rao and Jaggi, arXiv:2609.13149) provides a protocol for systematically evaluating how different memory and context budget constraints affect retrieval accuracy in local LLM agents. The core finding is that performance under context budget constraints is highly task-dependent: some retrieval tasks degrade gracefully as the budget shrinks, while others exhibit sharp drop-offs at specific thresholds. This means that performance measured at full context length is a poor predictor of performance at the constrained budgets that organizations actually operate under.

Practical long-context deployment requires characterizing your specific workload's behavior under the budget constraints you face, not just evaluating at theoretical maximum context length. A model that performs well at 128K tokens may perform adequately at 32K tokens for your task or may fail at a much higher threshold depending on where in the context the relevant information sits and how the task requires the model to integrate it.

Retrieval-Augmented Generation as a Memory Layer

For many enterprise tasks, the practical alternative to extending the context window is retrieval-augmented generation (RAG). Rather than placing an entire knowledge base into the context, a retrieval system selects the most relevant passages at query time and inserts them into a shorter context window. This sidesteps the computational cost and quality degradation of very long contexts while allowing the model to draw on a much larger knowledge base than any context window could contain.

RAG and long-context processing are complementary rather than competing approaches. For tasks where exact retrieval is reliable and the relevant information is localized, RAG at standard context lengths is more compute-efficient. For tasks requiring the model to synthesize across many documents simultaneously, or where the relevant information cannot be reliably retrieved in isolation because its significance depends on other parts of the document, long context is preferable. Choosing between them requires characterizing your specific task's retrieval properties, not applying a blanket policy.

Fig 1 · Attention cost growth with sequence length (stylized, log scale)
Full O(n2) Sliding O(n) 4K 16K 64K 256K 1M Context Length (log scale) Compute Cost
LLMs from Scratch Series

References

Excited about AI, innovation, and growth?

Start a conversation