Long Context Extrapolation: NTK-Aware Scaling + YaRN + LongRoPE — Journey from 8K to 1M Tokens

Ders Haritası (10 Bölüm)#

1. Context length problem — niye naive RoPE long context'te break
2. NTK-aware scaling intuition — frequency rescaling
3. Position interpolation — Chen 2023, simpler approach
4. YaRN (Peng 2023) — comprehensive solution
5. YaRN math — frequency-dependent scaling + temperature
6. LongRoPE (Microsoft 2024) — evolutionary search
7. Fine-tune protocol — kaç step, ne kadar data
8. Llama-3-8B 8K → 128K — production recipe
9. Gemini 1.5 1M token — Google'ın yaklaşımı
10. Future — 10M+ context görünür mü

1-5. NTK-Aware + YaRN Deep Dive#

1.1 Naive RoPE long context#

• Position indices 8K beyond → frequencies model görmemiş
• Attention pattern broken
• Perplexity patlar (10x+ degradation)

1.2 Position Interpolation (Chen 2023)#

position_scaled = position / scale_factor (e.g., 4 for 4x stretch)

1.3 NTK-aware scaling (LocalLlama community 2023)#

base_new = base × scale_factor^(d/(d-2))

1.4 YaRN (Peng 2023)#

1. Frequency-dependent scaling: low/mid/high frequencies farklı strategy
2. Temperature scaling: attention logits scale (compensates for longer context softmax dilution)
3. Ramp function: smoothing between scaling regions

λ_d = piecewise function of dimension d
θ_d_new = θ_d × λ_d
Attention scaled by sqrt(temperature)

1.5 YaRN empirical (Peng 2023)#

• 16K: 200 fine-tune steps
• 32K: 400 steps
• 64K: 600 steps
• 128K: 1000 steps

1.6 Why YaRN works#

1.7 LongRoPE (Microsoft 2024)#

• Population: candidate frequency rescalings
• Fitness: validation perplexity on long context
• Selection: top-K survives
• Crossover/mutation: explore search space

1.8 Frequency profile insight#

7-8. Production Extension Recipe#

7.1 Llama-3-8B 8K → 128K#

1. Base model: Llama-3-8B trained on 8K context, RoPE base 500000
2. Stage 1 (4K → 32K): 400 fine-tune steps, batch=4, lr=2e-5, ~100M tokens
3. Stage 2 (32K → 128K): 600 fine-tune steps, batch=2 (memory), ~200M tokens
4. Validation: needle-in-haystack benchmark

7.2 YaRN-based extension code#

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
    rope_scaling={
        "type": "yarn",
        "factor": 16.0,  # 8K → 128K (16x)
        "original_max_position_embeddings": 8192,
        "beta_fast": 32,
        "beta_slow": 1,
    },
)
# Fine-tune with long-context data
# ...

7.3 Open-source extensions#

• Yi-34B-200K: YaRN-based 200K
• Llama-3-Long: community-driven 256K extensions
• Mistral-Nemo-128K: native YaRN-compatible

7.4 Needle-in-haystack benchmark#

• Random information in long document ('Joe's favorite color is purple')
• Ask question depending on info
• Measure recall accuracy across context length

• Native 8K: 100% recall at 8K
• YaRN 32K extension: 99% at 32K
• YaRN 128K extension: 92% at 64K, 80% at 128K (degradation)

7.5 Gemini 1.5 1M token#

• Mixture of Experts (sparse compute)
• Custom attention pattern (sliding window + dense)
• Massive long-context training data
• TPU-specific optimization

Sıradaki Ders: Capstone — Llama-3 RoPE Implementation#