Long-context architectures

Modern AI systems need to process increasingly long sequences—entire documents, code repositories, or conversation histories. Traditional Transformers face quadratic complexity $O(n^2)$ in sequence length, making them impractical for contexts beyond ~2k-8k tokens. Long-context architectures solve this by redesigning attention mechanisms and memory structures to handle100k+ token contexts efficiently.

Why Long Context Matters

Standard Transformer attention computes similarity between every token pair: $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

For $n$ tokens, this requires $O(n^2)$ memory and compute. At 100k tokens, that's 10 billion comparisons per layer.

Architectural Approaches

1. Sparse Attention Patterns

Deriving Longformer's Sliding Window:

Start with full attention cost for token $i$: $C_i = n \text{ comparisons}$

Restrict to local window of size $w$: $C_i^{\text{local}} = w \text{ comparisons}$

For $n$ tokens: $C_{\text{total}} = n \cdot w$

If $w$ is constant (e.g., 512), complexity becomes: $O(n \cdot w) = O(n) \text{ when } w \ll n$

But local-only attention loses global context. Longformer's solution: add global attention tokens that attend to everything:

• Most tokens: local window ($w$ comparisons)
• $g$ special tokens: full attention ($n$ comparisons each)
• Total: $O(nw + gn)$

If $g \ll n$, this is $O(n)$ with global information preserved.

Token "sat" attends to: ["cat", "sat", "on"] ← local Token [CLS] attends to: ALL tokens ← global

Cost per token:

• Regular token: 3 comparisons (local window)
• Global token [CLS]: 11 comparisons (whole sequence)

There are 10 regular tokens + 1 global token: Total: 10×3 + 1×11 = 41 vs. 11×11 = 121 full attention

**Why this step?** Global tokens (like [CLS]) aggregate document-level information, while local windows capture nearby dependencies. This trades a small amount of full attention for mostly-local attention.

### 2. Linear Attention Mechanisms

> [!formula] Kernel Approximation Derivation
> Standard attention's softmax creates the bottleneck:
> $$\text{Attention}(Q, K, V) = \text{softmax}(QK^T)V$$

Expand softmax:
$$\text{softmax}(QK^T)_{ij} = \frac{e^{q_i \cdot k_j}}{\sum_{j'} e^{q_i \cdot k_{j'}}}$$

This requires computing all $q_i \cdot k_j$ dot products → $O(n^2)$.

==Linear attention== approximates $e^{q \cdot k}$ with explicit feature maps $\phi$:
$$e^{q \cdot k} \approx \phi(q)^T \phi(k)$$

Then:
$$\text{Attention}(Q, K, V)_i = \frac{\sum_j \phi(q_i)^T \phi(k_j) v_j}{\sum_j \phi(q_i)^T \phi(k_j)}$$

**Key insight**: Rearrange the sum using associativity:
$$= \frac{\phi(q_i)^T \left(\sum_j \phi(k_j) v_j^T\right)}{\phi(q_i)^T \left(\sum_j \phi(k_j)\right)}$$

Now we compute (let $d_\phi$ = feature map dimension, $d_v$ = value dimension):
1. $S = \sum_j \phi(k_j) v_j^T$ once for all tokens → $O(n \, d_\phi d_v)$
2. $Z = \sum_j \phi(k_j)$ once → $O(n \, d_\phi)$
3. For each query: numerator $\phi(q_i)^T S$ and denominator $\phi(q_i)^T Z$ → $O(d_\phi d_v)$ **per query**, i.e. independent of $n$

**Total across $n$ queries: $O(n \, d_\phi d_v)$ — linear in sequence length!**

**Why this step?** By precomputing the key-value aggregation $S$ and normalizer $Z$, each query only touches these fixed-size summaries. The per-query cost depends on $d_\phi, d_v$, **not** on $n$. The feature map $\phi$ must approximate the softmax kernel well enough to maintain quality.

> [!example] Linear Attention Calculation
> ```
> Given 4 tokens, d=2 dimensions, feature map ϕ(x) = [1, x₁, x₂] (so d_ϕ = 3):

Q = [[1, 0], [0, 1]]              # 2 queries
K = [[1, 0], [0, 1], [1, 1], [0, 0]]  # 4 keys  
V = [[1], [2], [3], [4]]         # 4 scalar values (d_v = 1)

Feature maps of keys (each is 3-dim):
ϕ(k₁) = [1, 1, 0]
ϕ(k₂) = [1, 0, 1]
ϕ(k₃) = [1, 1, 1]
ϕ(k₄) = [1, 0, 0]

Step 1: Precompute S = Σⱼ ϕ(kⱼ) vⱼ   (3-dim vector since d_v=1)
S = [1,1,0]·1 + [1,0,1]·2 + [1,1,1]·3 + [1,0,0]·4
  = [1,1,0] + [2,0,2] + [3,3,3] + [4,0,0]
  = [10, 4, 5]

Step 2: Precompute normalizer Z = Σⱼ ϕ(kⱼ)   (3-dim vector)
Z = [1,1,0] + [1,0,1] + [1,1,1] + [1,0,0]
  = [4, 2, 2]

Step 3: For q₁ = [1,0], ϕ(q₁) = [1, 1, 0]
numerator   = ϕ(q₁)·S = [1,1,0]·[10,4,5] = 10 + 4 + 0 = 14
denominator = ϕ(q₁)·Z = [1,1,0]·[4,2,2]  = 4 + 2 + 0 = 6
Attention₁  = 14 / 6 ≈ 2.33   ← normalization matters!

Compare to full attention: would need 4 separate q·k dot products per query.

Why this step? The linear scan over keys happens once globally (building $S$ and $Z$), not once per query. Each query then reuses those summaries in constant time w.r.t. $n$, and the division by $\phi(q)^TZ$ enforces the softmax-style normalization.

3. State Space Models (SSMs)

Discretization for tokens: Replace continuous time with discrete steps (tokens): $h_k = \bar{A}h_{k-1} + \bar{B}x_k$ $y_k = Ch_k$

where $\bar{A} = e^{\Delta A}$ and $\bar{B} = (e^{\Delta A} - I)A^{-1}B$ for step size $\Delta$.

Why SSMs for long context?

• State $h_k$ is fixed dimension (e.g., 256), not growing with sequence length
• Update is $O(d^2)$ per token where $d$ is state dimension
• Total: $O(nd^2)$ with $d \ll n$ → effectively $O(n)$

Choose $A$ to be diagonal plus low-rank: $A = \Lambda - pq^T$

This allows fast matrix exponential: $e^{\Delta A} \approx I + \Delta A + \frac{(\Delta A)^2}{2} + \ldots$

With diagonal $\Lambda$, the exponential is: $e^{\Delta \Lambda} = \text{diag}(e^{\Delta \lambda_1}, \ldots, e^{\Delta \lambda_d})$

The low-rank correction allows efficient updates via Woodbury identity.

Recurrent form: $h_k = \bar{A}h_{k-1} + \bar{B}x_k$

Convolutional form (for training parallelization): $y = \bar{K} * x$

where $\bar{K}_k = C\bar{A}^k\bar{B}$ is a learnable convolution kernel.

Transformer attention:

• Memory: 100k × 100k = 10B attention scores
• Compute: O(100k²) per layer

S4 SSM:

• Memory: state vector 256-dim + convolution kernel ~16k
• Compute: O(100k × 256²) = O(100k) effective per layer

At 100k length, SSM is ~1000× more memory efficient. But: SSM state is fixed-size, may lose information that attention retains.

**Why this step?** The state space formulation compresses the entire history into a fixed vector, similar to RNN hidden states but with better training stability. The structured matrices make it computationally feasible.

### 4. Memory-Augmented Architectures

> [!intuition] External Memory Banks
> Instead of attending to all tokens, maintain a ==compressed memory== of past context. New tokens attend to memory, and memory is updated selectively.

**Memorizing Transformers** architecture:
1. **Local attention**: Attend to recent $w$ tokens (sliding window)
2. **Memory attention**: Attend to $k$ nearest neighbors from external memory
3. **Memory update**: Store key-value pairs from current layer in memory bank

Memory lookup uses ==approximate nearest neighbor== (ANN) search:
$$\text{Retrieve}(q) = \text{TopK}_{m \in \text{Memory}}(\text{similarity}(q, m_k))$$

Using FAISS or similar, this is $O(\log M)$ where $M$ is memory size.

**Total complexity**: $O(nw) + O(n \log M)$ ≈ $O(n)$ when $w, \log M \ll n$

> [!example] Memory-Augmented Lookup
> ```
> Processing token 50,000 in a document:

Local window (w=512):
- Attend to tokens 49,488-50,000

Memory bank (M=10k compressed key-values from earlier tokens):
- Query current token embedding q₅₀ₖ
- FAISS finds top-k=64 most similar keys from memory
- Attend to those 64 retrieved values

Total attention: 512 + 64 = 576 tokens
vs. 50,000 for full attention

Memory stores: representative moments, chapter boundaries, 
key facts from earlier sections

Why this step? Not all past tokens are equally relevant. Retrieval focuses computation on what matters for the current token, mimicking human selective attention.

Comparison Table

Why it's wrong: The feature map $\phi$ approximates the softmax kernel. For tasks requiring precise token interactions (e.g., copying exact words, arithmetic), the approximation degrades quality. Empirically, linear attention often underperforms standard attention on tasks needing exact attention patterns.

The fix: Use linear attention for encoder layers processing long context, but keep standard attention for decoder layers generating output where precision matters. Or use hybrid approaches (local + linear).

Why it's wrong: Transformers have multiple layers. Information propagates through the network. With a sliding window of size $w$, the receptive field grows by about $w$ tokens per layer, so a token can reach distant context after roughly $n/w$ layers of aggregation—still far cheaper than full attention.

Example: Token at position 1000 with window $w=100$ (attends ±50 each side)

• Layer 1: receptive field ≈ 100 tokens (950–1050)
• Layer 2: each of those saw ±50 → field ≈ 200 tokens (900–1100)
• Layer $L$: receptive field ≈ $L \cdot w$ tokens
• To cover the whole 1000-token span you need $\approx n/w \approx 10$ layers, not 4–5.

The fix: Provide enough depth ($\sim n/w$ layers) for information to propagate, or add global tokens / dilated windows to shortcut long-range access in fewer layers.

Key Formulas Summary

Full attention is like for every new word you read, you flip back through the entire book to check how it relates to every previous word. That's exhausting! If the book is 1,000 pages, you'd do a million page flips.

Sparse attention says: "Only look at the last few pages (local), plus the table of contents and chapter summaries (global tokens)." Way faster! You still get the gist without reading everything again. (One catch: to connect very far-apart pages using only local peeks, you need many "passes" through the book.)

Linear attention is like making a cheat sheet as you read—you write down key points that update automatically. When you see a new word, you just check your one-page cheat sheet instead of the whole book. Checking the cheat sheet takes the same tiny time no matter how long the book is.

Memory banks are like highlighting the 100 most important sentences, then only looking at those highlights when you need to remember something.

Each method finds a smart shortcut to avoid re-reading everything, so you can handle really, really long books (or in AI, really long conversations and documents).

Connections

• Transformer Architecture - foundation being optimized
• Attention Mechanisms - core operation being re-engineered
• Computational Complexity - understanding $O(n^2)$ bottleneck
• Retrieval-Augmented Generation - related memory approach
• Recurrent Neural Networks - SSMs revive recurrent ideas with better scaling
• Kernel Methods - linear attention uses kernel approximation
• Information Bottleneck - state space models compress via fixed capacity

#flashcards/ai-ml

Why does standard Transformer attention have O(n²) complexity? :: Because it computes similarity (QKᵀ) between every pair of tokens - n queries each comparing to n keys.

What is the tradeoff of memory-augmented architectures? :: They achieve unlimited context by retrieving from external memory, but add retrieval overhead (O(log M) per query) and require managing what to store.

Concept Map

Hinglish (regional understanding)