Long Context Transformers

Chapter Overview

Extending transformer context length beyond standard limits (512-2048 tokens) enables processing long documents, books, and extended conversations. This chapter covers techniques for scaling to 32K, 100K, and even 1M+ token contexts.

Learning Objectives

1. Understand context length limitations and bottlenecks
2. Implement position interpolation and extrapolation
3. Apply memory-augmented transformers
4. Use retrieval-augmented generation (RAG)
5. Implement recurrent transformers (Transformer-XL)
6. Compare long-context methods and trade-offs

Context Length Limitations

The Quadratic Memory Bottleneck

Standard transformer architectures face fundamental limitations when processing long sequences due to the quadratic scaling of self-attention with respect to sequence length. The self-attention mechanism computes pairwise interactions between all tokens in a sequence, requiring the materialization of an attention matrix of size $n \times n$ where $n$ is the sequence length. This quadratic scaling manifests in three critical bottlenecks: computational complexity, memory consumption, and position encoding limitations. Understanding these bottlenecks quantitatively is essential for appreciating why long context processing requires specialized techniques and architectural modifications.

The computational complexity of self-attention is $O(n^2 d)$ where $d$ is the model dimension. For each of the $n$ queries, the model computes attention scores with all $n$ keys through dot products of dimension $d$, requiring $n^2 d$ multiply-accumulate operations. The subsequent softmax normalization adds $O(n^2)$ operations, and the weighted sum over values requires another $n^2 d$ operations. While the feed-forward network has complexity $O(n d^2)$, for long sequences where $n > d$, the attention computation dominates. For example, with $n = 16384$ tokens and $d = 768$, attention requires approximately 206 billion FLOPs per layer while the feed-forward network requires only 19 billion FLOPs, making attention the primary computational bottleneck.

The memory bottleneck is even more severe than the computational one. During the forward pass, the attention matrix must be fully materialized in memory before applying softmax, requiring $n^2$ memory locations per attention head. During the backward pass for training, these attention matrices must be stored for gradient computation, effectively doubling the memory requirement. For multi-head attention with $h$ heads, the total memory for attention matrices is $h \times n^2$ floating-point values per layer. In FP32 format, each value requires 4 bytes, while FP16 requires 2 bytes. This memory grows quadratically with sequence length, quickly exceeding available GPU memory for long sequences.

The third fundamental limitation involves position encodings. Standard transformers use position encodings trained on sequences of a fixed maximum length, typically 512 to 2048 tokens. When these models encounter sequences longer than their training length, the position encodings must extrapolate to unseen positions. Absolute position embeddings, which assign a learned vector to each position index, cannot extrapolate at all—positions beyond the training length have no corresponding embedding. Even sinusoidal position encodings, which use deterministic trigonometric functions, exhibit degraded performance when extrapolating beyond training lengths due to the model's learned attention patterns being calibrated to the training distribution of position encodings.

This extrapolation failure manifests as rapidly degrading perplexity for tokens beyond the training context length. A model trained on 2048-token sequences might achieve perplexity of 15 on positions 0-2048, but perplexity can increase to 25 or higher for positions 2048-4096 without specialized position encoding schemes. This degradation occurs because the model's attention patterns have learned to interpret specific position encoding values as corresponding to specific relative distances, and these learned patterns break down when position encodings take on values outside the training distribution.

Position Encoding for Long Context

The Extrapolation Challenge

Position encodings enable transformers to incorporate sequential order information into their otherwise permutation-invariant architecture. However, different position encoding schemes exhibit dramatically different behaviors when processing sequences longer than those seen during training. This extrapolation capability is critical for long context applications, where retraining on maximum-length sequences is often computationally prohibitive. The choice of position encoding scheme can determine whether a model trained on 2048-token sequences can successfully process 8192-token sequences with minimal fine-tuning, or whether it requires extensive retraining from scratch.

Absolute position embeddings assign a learned vector to each position index, with the position encoding for position $i$ being a trainable parameter $\vp_i \in \R^d$. These embeddings are added to token embeddings before the first transformer layer. While simple and effective within the training length, absolute embeddings cannot extrapolate beyond the maximum training position. A model trained with positions 0 through 2047 has no learned embedding for position 2048 or beyond. Attempting to extend such a model requires either initializing new position embeddings (which perform poorly without extensive fine-tuning) or using position interpolation techniques to map longer sequences into the trained position range.

Sinusoidal position encodings, introduced in the original Transformer paper, use deterministic trigonometric functions rather than learned parameters. For position $i$ and dimension $j$, the encoding is defined as:

Position Interpolation

Position interpolation addresses the extrapolation problem by mapping longer sequences into the position range seen during training, rather than extending beyond it. Instead of asking the model to extrapolate to unseen position indices, interpolation compresses the position indices of a long sequence into the trained range, effectively treating the long sequence as a "compressed" version of a training-length sequence.

The key insight behind position interpolation is that it keeps position encodings within the distribution seen during training, avoiding the extrapolation problem entirely. For example, extending from $L = 2048$ to $L' = 8192$ maps position 8191 to interpolated position $8191 \times 2048/8192 = 2047.75$, which falls within the training range. The model's attention patterns, having been trained on positions 0 through 2047, can handle this interpolated position much more effectively than they could handle the raw position 8191.

Position interpolation has been successfully applied to extend LLaMA models from 2048 to 8192 tokens and beyond. The technique requires minimal fine-tuning—typically only 1000 to 10000 training steps on long sequences—compared to training from scratch. After fine-tuning with position interpolation, LLaMA 2 models maintain perplexity within 5-10\% of their original performance when extended from 4096 to 32768 tokens, whereas naive extrapolation without interpolation results in perplexity degradation of 50\% or more.

The computational cost of position interpolation is negligible, as it only affects the position encoding computation, not the attention mechanism itself. The primary cost is the fine-tuning required to adapt the model to the compressed position space. However, this fine-tuning is far less expensive than training from scratch: extending a 7B parameter model from 4K to 32K context requires approximately 100 GPU-hours of fine-tuning compared to 100,000+ GPU-hours for full pretraining.

Rotary Position Embedding (RoPE)

Rotary Position Embedding represents a fundamental advance in position encoding design, achieving excellent extrapolation properties by encoding relative position information directly into the attention computation through rotation operations. RoPE has become the position encoding of choice for modern large language models including GPT-NeoX, LLaMA, PaLM, and many others due to its combination of strong extrapolation, computational efficiency, and theoretical elegance.

The crucial property of RoPE is that the attention score between positions $m$ and $n$ depends only on their relative distance $m - n$, not their absolute positions. This can be verified through the rotation addition formula:

This relative position property means that the model learns attention patterns based on relative distances between tokens rather than absolute positions. When extrapolating to longer sequences, the model encounters the same relative distances it saw during training, just in different combinations. A model trained on sequences up to 2048 tokens has seen all relative distances from -2047 to +2047. When processing a 4096-token sequence, it encounters the same relative distances, enabling much better extrapolation than absolute position encodings.

RoPE's extrapolation capability can be further enhanced through position interpolation. By scaling the rotation frequencies $\theta_j$ by a factor $L'/L$ when extending from length $L$ to $L'$, the effective relative distances are compressed into the training range. This combination of RoPE's inherent relative position encoding with position interpolation enables extensions from 2048 to 32768 tokens or beyond with minimal quality degradation.

The computational overhead of RoPE is minimal compared to the attention computation itself. Applying rotations to queries and keys requires $O(nd)$ operations, which is negligible compared to the $O(n^2d)$ cost of attention. The rotation operations can be efficiently implemented using vectorized operations on modern GPUs, adding less than 5\% to the total attention computation time. Memory overhead is also minimal, as the rotation matrices are block-diagonal and can be computed on-the-fly rather than stored.

In practice, RoPE has enabled dramatic context length extensions. LLaMA models using RoPE have been successfully extended from 2048 to 32768 tokens through position interpolation with only 1000 fine-tuning steps. The perplexity degradation is typically less than 10\% even at 16× the training length, compared to 50-100\% degradation for absolute position embeddings. This extrapolation capability has made RoPE the de facto standard for new large language models designed for long context applications.

ALiBi: Attention with Linear Biases

ALiBi (Attention with Linear Biases) takes a radically different approach to position encoding by eliminating position embeddings entirely and instead adding position-dependent biases directly to attention scores. This simple modification achieves remarkable extrapolation properties, enabling models trained on 1024-token sequences to process 10,000+ token sequences at inference time with no fine-tuning whatsoever.

The head-specific slopes are set geometrically: for $h$ attention heads, the slopes are $m_1, m_2, \ldots, m_h$ where $m_i = 2^{-8i/h}$. For example, with 8 heads, the slopes are $2^{-1}, 2^{-2}, \ldots, 2^{-8}$, giving values from 0.5 to 0.0039. This geometric spacing ensures that different heads have different receptive field sizes: heads with large slopes focus on nearby tokens, while heads with small slopes can attend to distant tokens with less penalty.

ALiBi's extrapolation capability stems from its use of relative distances rather than absolute positions, combined with the linear form of the bias. During training on sequences up to length $L$, the model encounters biases ranging from 0 (attending to the same position) to $-m \cdot L$ (attending to the most distant position). When extrapolating to length $L' > L$, the model encounters biases up to $-m \cdot L'$, which are simply larger values along the same linear scale. The attention patterns learned during training—which balance content-based attention (from $\vq_i\transpose \vk_j$) against distance-based penalties (from $-m \cdot |i-j|$)—continue to work effectively at these larger distances.

Empirical results demonstrate ALiBi's exceptional extrapolation. Models trained on 1024-token sequences with ALiBi can process 2048-token sequences with less than 5\% perplexity increase, 4096-token sequences with 10-15\% increase, and even 10,000-token sequences with 20-30\% increase—all without any fine-tuning. In contrast, the same models with sinusoidal position encodings show 50\% perplexity increase at 2048 tokens and become essentially non-functional beyond 4096 tokens. This zero-shot extrapolation capability makes ALiBi particularly attractive for applications where the maximum sequence length is unknown at training time or varies widely across use cases.

ALiBi has been adopted by several prominent models including BLOOM (176B parameters) and MPT (7B-30B parameters). BLOOM was trained with ALiBi on sequences up to 2048 tokens but can effectively process sequences of 4096 tokens or longer at inference time. MPT models trained with ALiBi on 2048-token sequences have been successfully deployed on tasks requiring 8192-token contexts with minimal quality degradation.

The computational overhead of ALiBi is negligible. Computing the bias $-m \cdot |i-j|$ for all $n^2$ attention scores requires $O(n^2)$ operations, which is dominated by the $O(n^2d)$ cost of computing $\mQ\mK\transpose$. The bias can be precomputed once per sequence and reused across all layers and heads (with different slopes $m$), further reducing overhead. Memory overhead is also minimal, as the bias matrix can be computed on-the-fly or stored once and reused.

The primary limitation of ALiBi is that it assumes a monotonic relationship between distance and relevance—more distant tokens are always penalized more heavily. This assumption holds well for many natural language tasks where local context is indeed more important than distant context. However, for tasks with long-range dependencies that are not distance-dependent (such as matching opening and closing brackets in code, or resolving coreferences across document sections), ALiBi's linear bias may be suboptimal compared to learned position encodings that can capture more complex position-dependent patterns.

Efficient Attention for Long Context

Sparse Attention Patterns

While position encoding improvements enable better extrapolation, they do not address the fundamental quadratic scaling of attention computation and memory. Efficient attention mechanisms reduce this quadratic bottleneck by restricting which tokens can attend to which other tokens, computing attention only over a subset of the $n^2$ possible connections. These sparse attention patterns can reduce complexity from $O(n^2)$ to $O(n \times w)$ where $w$ is a fixed window size, enabling processing of sequences that would be impossible with full attention.

The key insight behind sparse attention is that not all token pairs require attention computation. In many domains, particularly natural language, most relevant information comes from nearby tokens, with occasional long-range dependencies. By carefully designing sparsity patterns that preserve important connections while eliminating redundant ones, sparse attention can maintain model quality while dramatically reducing computational and memory requirements.

Longformer and BigBird

Longformer combines local windowed attention ($O(nw)$ for window size $w$) with global attention on task-specific tokens, enabling efficient processing of documents up to 4096+ tokens. BigBird extends this with random connections that create small-world shortcuts, proving that sparse attention with $O(n \log n)$ connections is a universal approximator for sequence-to-sequence functions. Both reduce attention memory by 6--8$\times$ compared to full attention. See Chapter~[ref] for the detailed mechanism definitions, complexity analysis, and benchmark results.

In the context of long-context processing, the key trade-off is between sparsity pattern and information propagation depth. With window size $w$, information requires $\lceil n/w \rceil$ layers to propagate across the full sequence. BigBird's random connections reduce the expected propagation path to $O(\log n)$ layers, which is particularly valuable for tasks requiring long-range reasoning. For sequences up to 2048 tokens, optimized full attention (e.g., Flash Attention) is typically faster; for 2048--8192 tokens, sparse attention becomes beneficial; beyond 8192 tokens, it becomes essential.

Comparison of Sparse Attention Methods

Different sparse attention patterns offer different trade-offs between efficiency, model quality, and implementation complexity. Local attention with window size $w$ provides the simplest pattern and best memory locality, achieving $O(nw)$ complexity with straightforward implementation. However, information propagation is limited to $w$ positions per layer, requiring deep networks for long-range dependencies. Longformer's addition of global tokens addresses this limitation by providing information hubs, enabling faster propagation while maintaining linear complexity. BigBird's random connections provide theoretical guarantees and empirically strong performance, but at the cost of irregular memory access patterns that reduce hardware efficiency.

For sequences up to 2048 tokens, the overhead of sparse attention often outweighs its benefits—full attention with optimized implementations like Flash Attention is typically faster and simpler. For sequences of 2048-8192 tokens, sparse attention becomes beneficial, with Longformer and BigBird providing good trade-offs between efficiency and quality. For sequences beyond 8192 tokens, sparse attention becomes essential, as full attention exceeds available memory on most hardware. At these lengths, the choice between Longformer and BigBird depends on the task: Longformer is simpler and faster for tasks where local context dominates, while BigBird provides better quality for tasks requiring complex long-range reasoning.

Recurrent Transformers

Transformer-XL

Relative position encodings: Modified for segment-level recurrence

Retrieval-Augmented Generation

RAG Architecture

Retrieval methods:

• BM25 (sparse)
• Dense retrieval (BERT embeddings + nearest neighbors)
• Hybrid (combine sparse and dense)

Context 1: The Eiffel Tower was constructed from 1887 to 1889...
Context 2: Gustave Eiffel designed the tower for the 1889 World's Fair...
Context 3: The tower is 330 meters tall and was the tallest...
Question: When was the Eiffel Tower built?
Answer:

RETRO: Retrieval-Enhanced Transformer

Architecture:

• Chunk input into segments (64 tokens)
• Retrieve neighbors for each chunk
• Cross-attend to retrieved chunks
• Chunked cross-attention layers

Performance: 25× fewer parameters with retrieval achieves same performance as larger model without retrieval!

Memory-Augmented Transformers

Compressive Transformer

Effective context: Active + Recent + Compressed

Memorizing Transformers

Key innovation: $k$-NN attention over entire history

Architecture:

• Store all past $(key, value)$ pairs in memory
• For each query, retrieve $k$ nearest neighbors
• Attend to local context + retrieved keys/values

Benefits:

• Effectively infinite context (limited by storage)
• Constant-time attention (with approximate $k$-NN)
• Improves perplexity on long documents

Long Context Models in Practice

LongT5: Efficient Encoder-Decoder

LongT5 extends the T5 encoder-decoder architecture to handle sequences up to 16,384 tokens by applying efficient attention mechanisms to both the encoder and decoder. Unlike decoder-only models that process sequences autoregressively, encoder-decoder models must handle long sequences in both components, making efficiency doubly important. LongT5 demonstrates that sparse attention patterns can be successfully applied to encoder-decoder architectures while maintaining the strong performance of the original T5 model.

LongT5 uses a combination of local and global attention patterns similar to Longformer, but adapted for the encoder-decoder structure. The encoder uses local attention with window size $w = 512$ for all tokens, plus global attention for a small number of designated tokens. The decoder uses local attention for attending to its own previous tokens, plus full attention to encoder outputs (which are compressed through the local attention mechanism). This asymmetric design recognizes that decoder-to-encoder attention is typically less memory-intensive than encoder self-attention, as decoder sequences are usually shorter than encoder sequences.

The memory savings from LongT5's sparse attention are substantial. For an encoder sequence of length $n_e = 16384$ and decoder sequence of length $n_d = 512$, full attention would require approximately $(16384^2 + 512^2 + 512 \times 16384) \times 4 = 1.1$ GB per attention head in FP32 for encoder self-attention, decoder self-attention, and cross-attention combined. With LongT5's sparse patterns using $w = 512$, the requirement reduces to approximately $(16384 \times 1024 + 512^2 + 512 \times 16384) \times 4 = 71$ MB per head, a 15× reduction. With 12 heads and 12 encoder layers plus 12 decoder layers, total attention memory decreases from 26 GB to 1.7 GB.

LongT5 has been successfully applied to long-document summarization tasks where input documents exceed 10,000 tokens. On the arXiv summarization dataset with papers averaging 6000 tokens, LongT5 achieves ROUGE-L scores of 48.3 compared to 44.1 for T5 with truncated 512-token inputs, demonstrating that access to full document context significantly improves summary quality. On the PubMed summarization task with medical papers averaging 3000 tokens, LongT5 outperforms T5 by 3-4 ROUGE points across all metrics.

Production Long Context Systems

As of early 2025, production language models support context lengths of 128K--1M+ tokens, enabling processing of entire codebases, books, and extended conversations.\footnote{Context length capabilities evolve rapidly. Check current documentation for the latest specifications.} These systems combine the techniques covered in this chapter---efficient attention, position interpolation, and memory-efficient implementations---to make long-context processing practical. The fact that such contexts are served at interactive latencies indicates the use of multiple optimizations simultaneously: sparse or approximate attention to reduce the quadratic bottleneck, Flash Attention for memory efficiency, and model parallelism across multiple GPUs.

Practical Considerations for Long Context

Deploying long context models in production requires careful consideration of when the additional context is actually beneficial versus when alternative approaches might be more effective. Long context processing incurs real costs in terms of latency, computational resources, and financial expense, so understanding when these costs are justified is essential for practical applications.

Long context is most valuable when the task requires reasoning over or synthesizing information from multiple parts of a long document. Document summarization benefits significantly from full document access, as important information may appear anywhere in the document. Question answering over long documents similarly benefits from long context, as the relevant information's location is unknown in advance. Code generation and analysis tasks benefit from seeing entire files or multiple related files in context, enabling the model to understand dependencies and maintain consistency.

However, many tasks that initially seem to require long context can be effectively addressed with shorter contexts and retrieval. For question answering over large document collections, retrieval-augmented generation (RAG) can retrieve only the relevant passages and provide them in a short context, achieving similar or better performance at much lower cost. For tasks requiring access to factual knowledge, retrieval from a knowledge base is often more reliable and efficient than encoding all knowledge in the context. For multi-turn conversations, summarizing or compressing earlier conversation history can maintain coherence while reducing context length.

The cost-benefit analysis depends on several factors. Latency requirements matter: long context processing takes longer, which may be unacceptable for interactive applications. Accuracy requirements matter: if the task requires very high accuracy and the model performs significantly better with full context, the additional cost may be justified. Update frequency matters: if the information changes frequently, retrieval from an updated database may be preferable to encoding static information in context. Scale matters: for high-volume applications, the per-request cost of long context processing multiplies, potentially making alternative approaches more economical.

A practical strategy is to use a hybrid approach: employ retrieval or summarization to reduce context length when possible, but fall back to full long context processing when the task genuinely requires it. For example, a document analysis system might first use retrieval to identify relevant sections, then process those sections with long context if they exceed the standard context limit. This approach balances the benefits of long context with the efficiency of shorter contexts, optimizing for both quality and cost.

Comparison and Trade-offs

Method Comparison

Different approaches to long context processing involve fundamentally different trade-offs between computational efficiency, memory usage, model quality, and implementation complexity. Understanding these trade-offs is essential for selecting the appropriate method for a given application, as no single approach dominates across all dimensions. The optimal choice depends on the specific requirements of the task, available hardware, and acceptable quality-efficiency trade-offs.

Standard full attention with optimized implementations like Flash Attention remains the gold standard for quality and simplicity when sequence lengths permit. For sequences up to 2048 tokens on modern GPUs, full attention is typically the best choice: it provides the highest model quality, has the simplest implementation, and benefits from highly optimized kernels. Flash Attention reduces memory bandwidth requirements through kernel fusion, enabling batch sizes 2-4× larger than naive implementations while maintaining identical outputs to standard attention. However, the fundamental $O(n^2)$ scaling means that full attention becomes impractical beyond 4096-8192 tokens on typical hardware.

Sparse attention methods like Longformer and BigBird reduce complexity to $O(n \times w)$ where $w$ is a fixed window size, enabling sequences of 4096-16384 tokens on standard GPUs. These methods maintain exact attention within their connectivity pattern, avoiding approximation errors. The primary trade-off is that sparse patterns may miss important long-range dependencies that fall outside the connectivity pattern. For tasks where local context dominates (such as language modeling or most NLP tasks), this limitation has minimal impact on quality. For tasks requiring complex long-range reasoning (such as certain question answering or reasoning tasks), sparse attention may underperform full attention even when both are feasible.

Linear attention methods like Performer and Linformer achieve $O(n)$ complexity through mathematical approximations, enabling very long sequences of 32768 tokens or more. However, these approximations introduce errors that can degrade model quality. Performer uses random feature approximations to the softmax kernel, which works well for some attention distributions but poorly for others. Linformer assumes low-rank structure in attention matrices, which holds for many tasks but may fail for tasks with complex attention patterns. In practice, linear attention methods typically show 2-5\% accuracy degradation on downstream tasks compared to full attention, which may or may not be acceptable depending on the application.

Recurrent methods like Transformer-XL process sequences in segments with recurrent connections, enabling unlimited context length with constant memory per segment. The trade-off is that information must propagate through multiple segments to flow across long distances, which can be slower than direct attention and may lose information through the recurrent bottleneck. Transformer-XL works well for tasks like language modeling where sequential processing is natural, but less well for tasks requiring bidirectional context or random access to different parts of the sequence.

Retrieval-augmented generation (RAG) sidesteps the context length problem entirely by retrieving only relevant information and providing it in a short context. This approach can handle effectively unlimited document collections while maintaining the quality and efficiency of short-context models. The trade-off is implementation complexity: RAG requires building and maintaining a retrieval system, embedding documents, and handling retrieval failures. Additionally, RAG works best for tasks where relevant information can be identified through retrieval, but may struggle with tasks requiring synthesis across many parts of a document or reasoning about information that is difficult to retrieve.

Hardware and Memory Considerations

The practical feasibility of different long context methods depends critically on available hardware and memory constraints. Modern GPUs vary widely in memory capacity, from 8 GB on consumer GPUs to 80 GB on high-end data center GPUs, and this memory capacity directly determines which sequence lengths are feasible with different methods.

For a BERT-base scale model with $d = 768$, $h = 12$ heads, and $L = 12$ layers, the memory requirements for different methods and sequence lengths are as follows. Full attention with $n = 2048$ requires approximately 2.4 GB for attention matrices across all layers, which fits comfortably on any modern GPU. At $n = 4096$, full attention requires 9.7 GB, which fits on 16 GB GPUs but leaves limited memory for batch processing. At $n = 8192$, full attention requires 38.5 GB, exceeding even high-end GPUs and requiring model parallelism or gradient checkpointing.

Sparse attention dramatically improves these numbers. Longformer with $w = 512$ and $n = 4096$ requires only 2.4 GB for attention matrices, enabling batch sizes 4× larger than full attention on the same hardware. At $n = 8192$, Longformer requires 4.8 GB, which fits comfortably on 16 GB GPUs. At $n = 16384$, Longformer requires 9.6 GB, still feasible on standard hardware. This memory efficiency enables processing of long documents on commodity GPUs that would be impossible with full attention.

Linear attention methods like Linformer with $k = 256$ require even less memory. At $n = 4096$, Linformer requires only 600 MB for attention matrices, enabling very large batch sizes or processing on smaller GPUs. At $n = 16384$, Linformer requires 2.4 GB, comparable to full attention at $n = 2048$. This memory efficiency enables processing of very long sequences, but at the cost of approximation errors that may degrade quality.

The memory requirements extend beyond attention matrices to include activations, gradients, model parameters, and optimizer states. For training, the total memory requirement is typically 4-6× the attention matrix memory when accounting for all these components. For inference, memory requirements are lower as gradients and optimizer states are not needed, but activations must still be stored for generation. These additional memory requirements mean that the feasible sequence length for training is typically 2-4× shorter than for inference on the same hardware.

Recommendations by Use Case

Selecting the appropriate long context method requires matching the method's characteristics to the specific requirements of the application. The following recommendations provide guidance based on common use cases and constraints.

For general NLP tasks with sequences up to 2048 tokens, use standard full attention with Flash Attention optimization. This provides the best quality with simple implementation and benefits from highly optimized libraries. The computational and memory costs are manageable on any modern GPU, and the simplicity reduces implementation and debugging time.

For document processing tasks with sequences of 2048-8192 tokens, use sparse attention methods like Longformer or BigBird. These methods provide good quality with manageable computational costs, and the sparse patterns align well with the local structure of natural language. Longformer is simpler and faster for tasks where local context dominates, while BigBird provides better quality for tasks requiring long-range reasoning. Both methods have well-tested implementations available in popular libraries.

For very long sequences of 8192-32768 tokens where quality is critical, consider using full attention with model parallelism or gradient checkpointing if hardware permits, or sparse attention if hardware is limited. The quality difference between full and sparse attention becomes more significant at these lengths, so the choice depends on whether the hardware can support full attention. If full attention is infeasible, BigBird typically provides better quality than Longformer at these lengths due to its random connections.

For extremely long sequences beyond 32768 tokens, or when processing large document collections, use retrieval-augmented generation (RAG) rather than attempting to fit everything in context. RAG provides better quality and efficiency than any long context method at these scales, as it focuses the model's attention on relevant information rather than processing irrelevant content. The implementation complexity of RAG is justified by the significant quality and efficiency improvements at these scales.

For streaming or online processing tasks, use Transformer-XL or similar recurrent methods that can process sequences incrementally without recomputing previous segments. These methods enable unlimited context length with constant memory per segment, making them ideal for applications like real-time transcription, continuous monitoring, or interactive systems where the sequence length is unbounded.

For tasks requiring frequent updates to the knowledge base or document collection, prefer RAG over long context methods. RAG allows updating the retrieval index without retraining the model, while long context methods require reprocessing the entire context whenever information changes. This makes RAG more practical for applications with dynamic information needs.

Exercises

Solutions

Full solutions for all exercises are available at \url{https://deeplearning.hofkensvermeulen.be}.

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
import numpy as np

# Part (a): Load pretrained GPT-2
model_name = "gpt2"  # Trained on 1024-token contexts
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

print(f"Original max position: {model.config.n_positions}")  # 1024

# Part (b): Implement position interpolation
def extend_position_embeddings(model, new_max_length=4096):
    """
    Extend position embeddings using interpolation
    """
    old_max_length = model.config.n_positions
    scale_factor = old_max_length / new_max_length  # 1024/4096 = 0.25
    
    # Get original position embeddings
    old_pos_emb = model.transformer.wpe.weight.data  # (1024, 768)
    
    # Create new position embeddings
    new_pos_emb = torch.zeros(new_max_length, old_pos_emb.shape[1])
    
    # Interpolate positions
    for new_pos in range(new_max_length):
        # Map new position to old position space
        old_pos_float = new_pos * scale_factor
        
        # Linear interpolation between floor and ceil
        old_pos_floor = int(np.floor(old_pos_float))
        old_pos_ceil = min(int(np.ceil(old_pos_float)), old_max_length - 1)
        
        if old_pos_floor == old_pos_ceil:
            new_pos_emb[new_pos] = old_pos_emb[old_pos_floor]
        else:
            # Interpolation weight
            weight = old_pos_float - old_pos_floor
            new_pos_emb[new_pos] = (
                (1 - weight) * old_pos_emb[old_pos_floor] +
                weight * old_pos_emb[old_pos_ceil]
            )
    
    # Update model
    model.transformer.wpe = nn.Embedding(new_max_length, old_pos_emb.shape[1])
    model.transformer.wpe.weight.data = new_pos_emb
    model.config.n_positions = new_max_length
    
    return model

# Extend to 4096 tokens
extended_model = extend_position_embeddings(model, new_max_length=4096)
print(f"Extended max position: {extended_model.config.n_positions}")  # 4096

# Part (c): Fine-tune on long sequences
from torch.utils.data import DataLoader, Dataset

class LongTextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=4096):
        self.encodings = []
        for text in texts:
            tokens = tokenizer.encode(text, max_length=max_length, truncation=True)
            if len(tokens) >= max_length // 2:  # Only use long sequences
                self.encodings.append(tokens)
    
    def __len__(self):
        return len(self.encodings)
    
    def __getitem__(self, idx):
        return torch.tensor(self.encodings[idx])

# Simulate long text dataset
long_texts = ["..." * 1000 for _ in range(100)]  # Replace with actual data
dataset = LongTextDataset(long_texts, tokenizer, max_length=4096)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

# Fine-tuning
optimizer = torch.optim.AdamW(extended_model.parameters(), lr=1e-5)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
extended_model.to(device)

extended_model.train()
for step, batch in enumerate(dataloader):
    if step >= 1000:
        break
    
    batch = batch.to(device)
    outputs = extended_model(batch, labels=batch)
    loss = outputs.loss
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if step 
        print(f"Step {step}, Loss: {loss.item():.4f}")



# Part (d): Evaluate perplexity at different lengths
def evaluate_perplexity(model, tokenizer, text, max_length):
    """Evaluate perplexity on a long sequence"""
    model.eval()
    tokens = tokenizer.encode(text, max_length=max_length, truncation=True)
    tokens = torch.tensor(tokens).unsqueeze(0).to(device)
    
    with torch.no_grad():
        outputs = model(tokens, labels=tokens)
        loss = outputs.loss
        perplexity = torch.exp(loss)
    
    return perplexity.item()

# Test sequences
test_text = "..." * 2000  # Long test text

lengths = [1024, 2048, 3072, 4096]
perplexities_extended = []
perplexities_original = []

for length in lengths:
    # Extended model
    ppl_ext = evaluate_perplexity(extended_model, tokenizer, test_text, length)
    perplexities_extended.append(ppl_ext)
    
    # Original model (will fail on >1024)
    if length <= 1024:
        ppl_orig = evaluate_perplexity(model, tokenizer, test_text, length)
        perplexities_original.append(ppl_orig)
    else:
        perplexities_original.append(float('inf'))  # Cannot handle
    
    print(f"Length {length}: Extended={ppl_ext:.2f}, Original={perplexities_original[-1]:.2f}")

# Part (e): Plot perplexity vs position
import matplotlib.pyplot as plt

def compute_position_perplexity(model, tokens, window_size=100):
    """Compute perplexity at each position"""
    model.eval()
    perplexities = []
    
    for i in range(0, len(tokens) - window_size, window_size):
        window = tokens[i:i+window_size]
        window_tensor = torch.tensor(window).unsqueeze(0).to(device)
        
        with torch.no_grad():
            outputs = model(window_tensor, labels=window_tensor)
            ppl = torch.exp(outputs.loss).item()
            perplexities.append(ppl)
    
    return perplexities

# Analyze 4096-token sequence
long_tokens = tokenizer.encode(test_text, max_length=4096, truncation=True)
position_ppls = compute_position_perplexity(extended_model, long_tokens)

plt.figure(figsize=(10, 6))
plt.plot(range(len(position_ppls)), position_ppls)
plt.xlabel('Position (in windows of 100 tokens)')
plt.ylabel('Perplexity')
plt.title('Perplexity vs Position for Extended GPT-2')
plt.grid(True)
plt.savefig('position_perplexity.png')