Length Extrapolation & Long Context: Training for Generalization to Longer Sequences¶

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

Rotary Position Embeddings (Rope) - Position encoding method
Sliding Window Attention - Alternative attention pattern
00 Attention Mechanisms - Foundation
Flash Attention - Inference optimization
Llm Inference Optimization - Complete stack