Skip to content

Length Extrapolation & Long Context

Overview

Length Extrapolation is the ability to generalize to sequences longer than training data. Long Context handling enables models to maintain quality over very long sequences (10K-1M tokens).

  • Challenge: Train on 4K tokens, test on 8K or 32K
  • Naive approach: Breaks badly (20-50% quality drop)
  • Solutions: RoPE with interpolation, ALiBi, Sliding Window, Recurrent models
  • State-of-the-art: 1M token context windows (Llama 2 100K, GPT-4 128K)

The Problem: Training Length vs Test Length Mismatch

Why Extrapolation Fails

Training:
- Model trained on sequences of length 4096
- Learns to handle up to position 4096
- Attention patterns designed for this scale

Testing on 8192 tokens:
- Positions 4097-8192 are never seen!
- Model must extrapolate to unseen position indices
- Absolute position embeddings: Can't extrapolate (out of learned range)
- Relative position embeddings: Treat all distances > bucketing limit as "far"

Result:
 - Accuracy drops significantly
 - Performance: 90% → 70% (20% drop)
 - Not production-ready for longer sequences

Why It's Hard

Challenge 1: Position embeddings
 - Absolute: Learn embedding for position i
 - At inference: Position 4097 is unknown!
 - Can't generalize to unseen indices

Challenge 2: Attention pattern changes
 - Short sequences: Can attend to everything
 - Long sequences: Attention bottleneck, must be selective
 - Model never learned selective attention patterns!

Challenge 3: Long-range dependencies
 - Training length: 4K
 - Actual need: 32K tokens back
 - Model never learned to maintain info over 32K

Challenge 4: Numerical stability
 - Softmax on large matrices can overflow
 - Gradient flow over many layers gets noisy
 - Optimization becomes harder

-

Solutions: Enabling Length Extrapolation

1. RoPE with Position Interpolation

Standard RoPE:
 - θ_{m,j} = m × Θ^{-2j/d} where Θ = 10,000
 - For position m=4000: θ ≈ 4000 × base^(...)
 - For position m=8000: θ ≈ 8000 × base^(...)
 - Frequencies scale with position linearly

Problem:
 - Model learned frequencies for positions 0-4096
 - Positions 4097+ are "unlearned" frequency space

Solution: Position Interpolation (PI)
 - Shrink position indices by scaling factor α
 - α = training_length / test_length
 - Example: trained on 4K, want 8K
 - α = 4096 / 8192 = 0.5
│
 - Use: θ_{m,j} = (m × α) × Θ^{-2j/d}
 - For position m=8000: θ ≈ 4000 × Θ^(...)
 - Falls in learned range (0-4096 scale)
│
 - Effectively: compress 8K tokens into 4K space
 - Tokens closer together, but model learns to handle it

Result:
 - Can extrapolate 1.5-2x training length
 - Loss: ~5-10% quality
 - Better than naive extrapolation (20%+ loss)!

2. Frequency Scaling / Base Adjustment

Alternative to position interpolation:

Idea: Adjust base to expand frequency spectrum

Standard: Base = 10,000
Modified: Base = 10,000 × (test_length / train_length)

Example:
 - Training: 4K length, base = 10,000
 - Testing: 8K length, base = 10,000 × 2 = 20,000
 - Effect: Frequencies stretched to 2x spectrum
 - Positions 0-8K now map to learned range!

Pros vs PI:
 - Simpler to implement
 - Works in both directions (extrapolate and compress)

Cons:
 - Needs calibration for each length
 - Less effective than PI for extreme extrapolation

3. ALiBi (Attention with Linear Biases)

Different approach: Don't encode position in embeddings
 Encode in attention scores!

Standard attention:
 - Score_{i,j} = Q_i @ K_j / √d

ALiBi modification:
 - Score_{i,j} = Q_i @ K_j / √d + b × (i - j)
 - Linear bias based on relative distance!

Key insight:
 - Doesn't need position embeddings at all!
 - Relative distance (i-j) encoded directly
 - Generalizes naturally to longer sequences!

Training:
 - Train on 4K sequences
 - Attention scores: i - j ranges from -4096 to +4096

Testing on 8K:
 - Attention scores: i - j ranges from -8192 to +8192
 - Already seen range -4096 to +4096
 - Unseen range extrapolates naturally
 - Quality: ~95% (vs 70% for standard RoPE)

Advantage:
 - Excellent length extrapolation
 - Tested up to 32K with <5% loss

4. Sliding Window Attention

Alternative for long context:

Idea: Don't attend to everything
 Only attend to recent window of tokens

Implementation:
 - Attention window size: W (e.g., 256)
 - Each token attends only to previous W tokens
 - Reduces complexity from O(N²) to O(N×W) = O(N)

Extrapolation:
 - Window size W works for any sequence length N!
 - Can handle 4K, 8K, 32K, 128K same way
 - Only depends on local pattern understanding
 - Generalizes perfectly!

Trade-off:
 - Long-range dependencies limited
 - But with many layers, can cover full context
 - Quality: ~95-98% (very good!)
 - Speed: Much faster than dense attention!

Adoption:
 - Mistral 7B uses this
 - Trained on 32K contexts natively

-

Practical Solutions for Different Scales

4K → 8K (1.5x to 2x extension)

Recommended: Position Interpolation (simplest)

Implementation:
 - Use standard RoPE
 - Apply position scaling: m_new = m × (4096/8192) = m × 0.5
 - Training time: No change (compatible with existing checkpoints)
 - Testing: Automatic extrapolation
 - Quality: 90-95% (5-10% loss acceptable)

Code:
```python
def apply_rope_with_pi(q, k, cos, sin, scale_factor=0.5):
 """Position interpolation for RoPE"""
 # Original positions were 0 to seq_len
 # Scale them down to [0, seq_len * scale_factor]
 seq_len = q.shape[-2]
 positions = torch.arange(seq_len) * scale_factor

 # Get cos/sin for scaled positions
 cos_scaled = cos[positions.long()]
 sin_scaled = sin[positions.long()]

 # Apply standard rotation
 q_rotated = apply_rotation(q, cos_scaled, sin_scaled)
 k_rotated = apply_rotation(k, cos_scaled, sin_scaled)

 return q_rotated, k_rotated

4K → 32K (8x extension)

Recommended: ALiBi or Sliding Window (better quality)

ALiBi approach:
 - No retraining needed
 - Just apply attention bias
 - Quality: 95%+ (better than PI for large extensions)
 - Speed: Same (no overhead)
 - Simplest high-quality solution!

Sliding window approach:
 - Requires architecture change
 - But: Perfect generalization
 - Quality: 98%+ with large window
 - Speed: Much faster!
 - Best for production at scale

4K → 128K+ (32x+ extension)

Recommended: Recurrent models + Sliding Window (for extreme lengths)

Approach 1: Sliding Window
 - Fixed window (e.g., 4K)
 - Multi-layer stacking handles receptive field
 - Can scale to any length
 - Quality: 98%+

Approach 2: Sparse Attention
 - Mix dense + sparse patterns
 - Local dense + long-range sparse
 - Complexity: O(N log N)
 - Quality: 97%+

Approach 3: Recurrent Models
 - Process sequence in chunks
 - Maintain compressed state between chunks
 - Mamba, State Space Models
 - Quality: 97%+ with lower compute!

State-of-the-Art Systems

Llama 2 Extended Context (100K)

Method: Position Interpolation + Fine-tuning

Training:
 - Original: 4K context
 - Extended: Fine-tune with position interpolation
 - Data: 32K → 100K token sequences
 - Duration: Few thousand steps on 100K sequences

Results:
 - Quality on 100K: 95%+
 - No catastrophic forgetting on 4K
 - Long-doc QA: Near-perfect performance
 - Breakthrough: Made long context practical!

GPT-4 with Vision (128K)

Architecture:
 - Dense attention (may use sparse in practice)
 - Sliding window or similar for efficiency
 - Position embedding: Not disclosed

Capability:
 - 128K context window
 - Maintains quality across full range
 - Practical for: Long documents, multi-file code, transcripts
 - Trade: Slower (longer sequence = slower inference)

Best Practices

Training for Long Context

Step 1: Start with shorter context
 - Train on 4K or 8K (standard)
 - Convergence is easier
 - Cost is lower

Step 2: Extend context (if needed)
 - Use position interpolation
 - Fine-tune on longer sequences
 - Only last 10-20% of training budget
 - Can extend 2x training length with minimal data

Step 3: Final evaluation
 - Test on longest needed length
 - Benchmark on long-doc tasks
 - Verify quality holds

Inference Optimization for Long Context

Don't just enable longer sequences!

Performance optimization:
 - Use KV cache (always!)
 - PagedAttention (for batching multiple long sequences)
 - Flash Attention (faster computation)
 - Sliding window (if appropriate)
 - Combined: 10-100x speedup possible!

Key Takeaways

Length extrapolation: Can generalize to unseen sequence lengths RoPE + PI: Simple, effective for 2x extension ALiBi: Better for large extensions (8x+) Sliding window: Perfect generalization, efficient Modern: 128K context windows feasible with right techniques

-