YaRN & LongRoPE — Kéo dài context window sau training

Bạn có một LLM chỉ được train để đọc đoạn văn 4,000 tokens, nhưng cần nó xử lý cả cuốn sách dài 2 triệu tokens. Train lại từ đầu tốn hàng triệu đô và hàng tháng trời. Giải pháp? YaRN và LongRoPE — kỹ thuật "kéo giãn" context window chỉ qua vài trăm bước fine-tune, biến model ngắn hạn thành hệ thống trí nhớ dài hạn mà vẫn giữ được độ chính xác cục bộ.

Vấn đề

Context window của LLM giống như "tầm nhìn" của model. Một model 4K context có thể nhìn lại 4,000 tokens trước đó để dự đoán token tiếp theo. Khi đầu vào vượt quá giới hạn này, model bắt đầu "quên" — không phải vì thiếu bộ nhớ, mà vì hệ thống mã hóa vị trí (positional encoding) bị méo mó.

RoPE (Rotary Position Embedding), chuẩn mực hiện nay, mã hóa vị trí bằng cách xoay vector embedding theo các tần số khác nhau — như một bộ kim đồng hồ đa tần số. Kim quay nhanh (tần số cao) bắt được chi tiết cục bộ (token gần nhau), kim quay chậm (tần số thấp) định vị toàn cục (vị trí tuyệt đối trong văn bản).

Vấn đề xảy ra khi ta đưa vào sequence dài hơn lúc train. Các góc xoay bắt đầu "quấn vòng" (aliasing) hoặc mất độ phân giải. Cách đơn giản nhất là Position Interpolation (PI) — nén đều mọi vị trí xuống theo tỷ lệ. Nhưng PI làm mờ chi tiết cục bộ: khoảng cách 1 token trở thành 0.25 token, khiến model không còn phân biệt được "hôm nay" và "hôm qua".

Thêm nữa, khi context dài, attention scores trở nên "loãng" — variance của dot-product Q·K sụt giảm, khiến softmax mất entropy và attention mechanism ngừng hoạt động hiệu quả.

Ý tưởng cốt lõi

YaRN (Yet another RoPE extensioN) nhận ra rằng: không phải tất cả "kim đồng hồ" đều cần được kéo giãn như nhau.

NTK-by-parts: Chia band tần số

Thay vì nén đều, YaRN chia các chiều RoPE thành 3 band theo bước sóng λᵢ:

• Tần số cao (bước sóng ngắn): Giữ nguyên hoặc nén rất ít. Đây là kim quay nhanh, phụ trách mối quan hệ cục bộ (2-3 token cạnh nhau). Nếu nén band này, "hôm nay" và "ngày mai" sẽ nằm gần nhau đến mức model không phân biệt được.
• Tần số thấp (bước sóng dài): Nén theo hệ số s = L'/L (tỷ lệ context mới/cũ). Đây là kim quay chậm, phụ trách vị trí toàn cục. Chúng cần được "kéo" để mở rộng phạm vi nhìn xa.
• Tần số giữa: Nội suy tuyến tính theo hàm ngưỡng γ(r) dựa trên tỷ lệ r = L/λᵢ.

Đây chính là "aha moment": model không cần học lại cách đọc từng từ (local), chỉ cần điều chỉnh cách định vị mình trong tài liệu dài (global). Bằng cách bảo vệ band tần số cao, YaRN giữ nguyên khả năng xử lý ngôn ngữ cục bộ, trong khi mở rộng tầm nhìn xa qua band tần số thấp.

Attention Temperature Scaling

Khi kéo dài context, dot-product Q·K bị giảm variance do các vector xoay dàn trải ra. Điều này khiến softmax trở nên "lạnh" — phân bố attention quá flat, mất khả năng tập trung.

YaRN thêm hệ số nhiệt độ (temperature scaling): rescale Q/K vectors bằng 1/√t, với √t ≈ 0.1·ln(s) + 1. Điều này "hâm nóng" lại attention scores, duy trì entropy của softmax ở mức tối ưu dù context có dài đến đâu.

LongRoPE: Kỹ thuật progressive scaling

LongRoPE đẩy giới hạn lên 2 triệu tokens bằng chiến lược progressive extension:

1. Bắt đầu từ 256K tokens, fine-tune nhẹ
2. Mở rộng lên 512K với tỷ lệ nội suy không đều (non-uniform interpolation)
3. Tiếp tục lên 2048K (2M) tokens

Thay vì một bước nhảy vọt từ 4K lên 2M (gây sốc cho model), LongRoPE "tập cho model quen dần" với không gian vị trí mở rộng, giảm thiểu perplexity cliff.

Tại sao nó hoạt động

Bảo toàn cấu trúc cục bộ là then chốt. Các mối quan hệ ngắn hạn (như cấu trúc câu, ngữ pháp liền kề) phụ thuộc vào sự khác biệt nhỏ trong không gian embedding. Nếu ta nén tần số cao (high-frequency components), khoảng cách Euclidean giữa các token liền kề co lại, khiến attention mechanism không còn phân biệt được "token A bên cạnh token B".

Về mặt toán học, RoPE sử dụng tính chất: qₘ · kₙ = f(|m-n|) — chỉ phụ thuộc khoảng cách tương đối. YaRN điều chỉnh hệ số xoay θᵢ theo band tần số, biến đổi hàm f này sao cho:

• Với |m-n| nhỏ (local): f giữ nguyên độ dốc → giữ nguyên độ nhạy cảm
• Với |m-n| lớn (global): f được "kéo dãn" → mở rộng phạm vi

Temperature scaling đảm bảo rằng khi không gian mở rộng, độ "sắc nét" của attention không bị phân tán. Nếu không có bước này, attention weights sẽ đồng đều nhau, biến model thành "mù" — nhìn thấy toàn bộ context nhưng không tập trung vào điểm nào.

Ý nghĩa thực tế

Hiệu quả thực tế:

• YaRN: Mở rộng LLaMA 7B từ 4K lên 32K tokens với < 10% tăng perplexity (so với >50% của naive Position Interpolation).
• LongRoPE: Đạt 2 triệu tokens trên LLaMA2, với >90% accuracy trên bài test "passkey retrieval" (tìm số bí mật ở vị trí 1M+).
• Hiệu quả training: Chỉ cần ~400 bước fine-tune (tương đương 10× ít token hơn full pre-training), tiết kiệm 2.5× compute so với cách tiếp cận thô sơ.

Giới hạn cần biết:

• Degradation ở cực dài: Trên 100K tokens, perplexity vẫn suy giảm so với model được train native từ đầu với context dài. YaRN là giải pháp "tạm thời tốt", không thay thế hoàn toàn native long-context training.
• Trade-off ngắn hạn: Nếu cấu hình sai (nén quá nhiều high-frequency), model có thể chững lại trên task 4K-8K tokens thông thường.
• Hard limit: Giá trị base của RoPE (thường là 10,000) tạo giới hạn lý thuyết cho việc mở rộng, dù trong thực tế thường không chạm tới.
• Attention vẫn O(n²): YaRN giải quyết vấn đề position encoding, không giải quyết chi phí tính toán. Context 2M tokens vẫn cần Ring Attention và KV Cache Quantization để chạy được trên hardware thực tế.

Đào sâu hơn

Paper gốc:

• YaRN: Efficient Context Window Extension of Large Language Models (Peng et al., 2023) — Giới thiệu NTK-by-parts và temperature scaling
• LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens (Xiao et al., 2024) — Progressive scaling và non-uniform interpolation
• LongRoPE2: Near-Lossless LLM Context Window Scaling (2025) — Cải tiến mới nhất giảm thiểu tổn thất ở cực dài

Cùng cụm (Long Context):

• Ring Attention — Phân tán attention computation qua nhiều GPU để xử lý context triệu tokens, giải quyết vấn đề O(n²) compute mà YaRN không động tới
• Native Long-Context Training — So sánh giữa "kéo dài" sau training (YaRN) và "train dài từ đầu", khi nào nên dùng cái nào

Đọc tiếp (Foundation & Inference):

• Flash Attention — Cơ chế tiling để tính attention hiệu quả, nền tảng cho mọi kỹ thuật context dài
• KV Cache — Hiểu về bộ nhớ cache K/V, bottleneck chính khi mở rộng context mà YaRN phải đối mặt