Llama-3-8B train với context 8k tokens. Sau vài tháng, cộng đồng có Llama-3-8B-128k chạy được 128k tokens, một số biến thể tới 1M. Không phải retrain từ đầu, chỉ tweaking nhỏ.

Cùng lúc, Gemini-1.5 Pro quảng cáo 2M context, GPT-4 Turbo 128k, Claude 200k. Có thật là model “đọc hiểu” toàn bộ 2 triệu token? Hay chỉ là attention dense kéo dài?

Bài này mở câu chuyện position encoding, vì đó là chỗ context length thực sự được quyết định. Hiểu RoPE, RoPE scaling, YaRN, ALiBi để biết khi nào context dài là thật, khi nào là marketing.

Dành cho: dev đã hiểu attention (bài 9-11) và muốn hiểu tại sao “context length” lại là một con số kỳ lạ và biến động.

Mental model: position là gì trong Transformer

Attention không quan tâm thứ tự token. "chó cắn người" và "người cắn chó" cho cùng output nếu chỉ có embedding mà không có position.

Để Transformer “biết” token nào đứng trước, cần thêm thông tin position vào input. Đây là position encoding.

Lịch sử ngắn:

  • GPT-1, BERT (2018): learned absolute position embedding. Mỗi position 0, 1, …, max_len học một vector. Vấn đề: không generalize ngoài max_len train.
  • Sinusoidal (Vaswani 2017): dùng sin/cos pattern, lý thuyết extrapolate. Thực tế cũng không tốt lắm.
  • RoPE (Su 2021, 2024 phổ biến): rotate query và key theo position. Hiện là default cho hầu hết LLM hiện đại (Llama, Mistral, DeepSeek, Qwen).
  • ALiBi (Press 2022): bias attention score theo khoảng cách. Đơn giản, extrapolate tốt nhưng không phổ biến bằng RoPE.

Phần 1: RoPE, position bằng rotation

Ý tưởng RoPE (Rotary Position Embedding): thay vì thêm position vector vào embedding, rotate query và key vector theo position trước khi tính attention.

Hình dung 2D đơn giản: vector [1, 0] ở position 0. Ở position 1, rotate theo góc θ. Position 2: rotate 2θ. Position N: rotate Nθ.

position 0: vector quay 0 độ
position 1: vector quay 10 độ
position 2: vector quay 20 độ
...
position N: vector quay N * θ độ

Trong 4096 chiều (Llama-3-8B), chia vector thành 2048 cặp, mỗi cặp rotate với góc khác nhau:

cặp 1: rotate với góc θ_1 = 1.0 (chậm nhất)
cặp 2: rotate với θ_2 = 0.1
...
cặp 2048: rotate với θ_2048 = 1/10000 (nhanh nhất)

Khi tính attention Q · K^T, dot product giữa 2 vector đã rotate cho kết quả phụ thuộc vào khoảng cách giữa 2 position, không phải position tuyệt đối. Đây là tính chất quan trọng nhất.

Q_pos_5 · K_pos_2 = f(5 - 2) = f(3)
Q_pos_100 · K_pos_97 = f(100 - 97) = f(3)

Hai pair token cùng khoảng cách 3 cho attention score tương tự, dù ở đầu hay cuối context. Relative position tự nhiên.

Code conceptual:

def rope_rotate(x, position):
    # x shape: [dim], chia thành các cặp (x_even, x_odd)
    # theta_i = 1 / 10000^(2i/dim)
    theta = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
    angle = position * theta
    cos = torch.cos(angle)
    sin = torch.sin(angle)
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    rotated_even = x_even * cos - x_odd * sin
    rotated_odd = x_even * sin + x_odd * cos
    return interleave(rotated_even, rotated_odd)

Trong Transformer thực, RoPE apply lên Q và K trước attention:

Q' = rotate(Q, position)
K' = rotate(K, position)
Attention = softmax(Q' · K'^T / sqrt(d)) · V

Phần 2: Tại sao RoPE break khi context dài hơn train

Vấn đề: tham số theta_base = 10000 được chọn cho train context (vd 8192). Khi inference với context dài hơn, các cặp dimension thấp (rotate chậm) chưa “kịp” hoàn thành chu kỳ; các cặp dimension cao (rotate nhanh) thì đã quay nhiều vòng, mất ý nghĩa.

Cụ thể: với position 16384 (gấp đôi train 8192), một số cặp đã rotate đến mức mà attention không còn match được pattern đã học.

Hậu quả: model “lost” với token xa hơn 8192. Đây không phải bug, mà là extrapolation không tốt với RoPE default.

Phần 3: RoPE scaling, sửa nhanh không retrain

Cách đơn giản nhất để mở rộng context: giảm tốc rotation. Nếu train 8192 mà muốn chạy 32768 (gấp 4 lần), giảm tần số rotation đi 4 lần.

Có vài flavors:

Position Interpolation (PI, Meta 2023): chia position cho factor.

def rope_pi(x, position, scale=4):
    return rope_rotate(x, position / scale)  # position effective bị nén

Position 16384 effective được tính như position 4096. Model đã quen với 4096 -> hoạt động được.

Nhược: chất lượng giảm. Cần fine-tune ngắn (~1B tokens) để recover.

NTK-aware scaling (Reddit /r/LocalLLaMA 2023): thay vì scale đều tất cả dimensions, scale base theta thay đổi.

def rope_ntk(x, position, original_max=8192, new_max=32768):
    factor = (new_max / original_max) ** (dim / (dim - 2))
    new_base = 10000 * factor
    theta = 1.0 / (new_base ** (torch.arange(0, dim, 2) / dim))
    angle = position * theta
    # ... rotate như cũ

Idea: dim cao (rotate nhanh) cần scale ít, dim thấp (rotate chậm) scale nhiều. Smooth transition.

Tốt hơn PI, không cần fine-tune. Default cho Llama-3.1 long context.

YaRN (2023): kết hợp NTK với thêm vài trick (attention temperature scaling, gentle interpolation for high freq). Cho extrapolation lên 128k+ ổn định.

# RoPE scaling config trong HF
rope_scaling:
  type: "yarn"
  factor: 16   # extend 8k -> 128k
  original_max_position_embeddings: 8192

YaRN là default cho hầu hết model long-context post-2024.

So sánh nhanh:

MethodCần fine-tune?Quality dropMax extension
RoPE defaultKhôngHỏng > 8k1x
PICó (~1B tokens)1-3%4-8x
NTK-awareKhông<1%4x
YaRNCó (vài trăm M tokens)<1%16-32x

Phần 4: ALiBi, lối đi đơn giản hơn

ALiBi (Attention with Linear Biases, Press 2022) chọn cách khác hoàn toàn. Không rotate gì cả. Chỉ thêm bias âm vào attention score, tỉ lệ với khoảng cách:

attention_score[i, j] = (Q_i · K_j) - m * |i - j|

m là slope, mỗi head có một slope khác. Token xa nhau bị penalty nặng hơn.

Tính chất:

  • Đơn giản: chỉ thêm 1 phép trừ trong attention
  • Extrapolation tốt: train 4k chạy được 16k mà quality không drop nhiều
  • Không cần fine-tune cho extension

Nhược:

  • Bias linear “đều”, không học được pattern phức tạp về position
  • Một số benchmark cho thấy RoPE outperform ALiBi ở quality general

ALiBi được dùng trong MPT (MosaicML), BLOOM, một số model khác. Không phổ biến bằng RoPE nhưng có triết lý đẹp.

Phần 5: Long context có thật là dùng được?

Quảng cáo 128k, 200k, 1M sounds great. Thực tế?

Test “needle in haystack”: nhét một fact ngẫu nhiên vào context dài, hỏi LLM về fact đó.

Llama-3.1-128k: pass ~85% (đặc biệt yếu ở 30-70% giữa context). Claude-3.5: pass ~95%. Gemini-1.5 Pro 2M: pass 90-95% lên tới 1M token, drop dần sau đó.

Vấn đề chung: lost in the middle. Model attention vào đầu và cuối context tốt nhất. Giữa context (30-70% position) thường bị bỏ qua.

Attention quality theo position:
  0%-10%:    cao
  10%-30%:   trung bình
  30%-70%:   thấp (lost in middle)
  70%-90%:   trung bình
  90%-100%:  cao

Đây là lý do RAG vẫn tốt hơn “stuff everything into context”: top-K chunks có khả năng được attend cao hơn nhiều so với đặt giữa context dài.

Phần 6: Cost của long context

Long context không free, dù model hỗ trợ.

Compute cost: attention là O(N²) với N là context length. Context 128k tốn 256x compute so với 8k. Trong prefill phase, chờ đợi rất lâu.

Memory cost: KV cache scale tuyến tính với N. Context 128k cho Llama-3-70B FP16: ~80 GB KV cache cho 1 sequence. Cần A100 80GB chỉ để serve 1 user.

Quality cost: đã nói, lost in middle.

Latency cost: first-token latency có thể 10-30 giây với context 128k.

Quy tắc thực tế:

  • Chat thông thường: 4k-8k đủ
  • Document analysis: 16k-32k thường đủ
  • Codebase analysis: 32k-128k tuỳ project
  • 1M context: thí nghiệm, không production

Phần 7: Pitfall thực tế

Pitfall 1: Set max_model_len = max model nhưng không dùng hết.

vLLM --max-model-len 128000 allocate VRAM cho KV cache theoretical, giảm batch size khả thi. Đặt theo use case thực (4k, 8k), không max model.

Pitfall 2: Trông cậy long context thay vì RAG.

“Stuff 50 documents vào context 128k” sounds easy. Thực tế: lost in middle, retrieval của RAG vẫn tốt hơn cho most cases. Long context bổ sung RAG, không thay thế.

Pitfall 3: Fine-tune model long-context với data ngắn.

Bạn fine-tune Llama-3-8B-128k với data 2k. Sau fine-tune, position embeddings cho 8k-128k mất calibration vì không thấy data dài. Model trở lại “8k effective”.

Fix: fine-tune long-context model cần data dài tương ứng. Hoặc dùng QLoRA với rank thấp để giảm shift.

Pitfall 4: So sánh model long-context bằng benchmark ngắn.

MMLU dài ~2k token mỗi câu. Test trên MMLU không cho biết model có long-context tốt không. Dùng LongBench, RULER, hoặc test “needle in haystack” thủ công.

Phần 8: Hands-on YaRN với llama.cpp

llama.cpp hỗ trợ YaRN qua flag:

./main -m llama-3-8b-instruct.gguf \
    -c 32768 \
    --rope-scale 4.0 \
    --yarn-orig-ctx 8192 \
    --yarn-ext-factor 1.0 \
    -p "<long context 30k tokens>" \
    -n 256

Flags:

  • -c 32768: context length tối đa
  • --rope-scale 4.0: scale factor (8192 * 4 = 32768)
  • --yarn-orig-ctx 8192: original context của model
  • --yarn-ext-factor 1.0: YaRN extension factor

Đo thực nghiệm:

  1. Tạo file context 30000 tokens (vd: copy nhiều bài blog).
  2. Nhét “needle”: câu fact ngẫu nhiên ở position 5k, 15k, 25k.
  3. Hỏi LLM về fact đó. Đo accuracy.

Bạn sẽ thấy: position 5k và 25k pass tốt, position 15k (giữa) miss thường xuyên hơn.

Cheatsheet

Khái niệmBản chất
RoPERotate Q, K theo position
Position InterpolationChia position cho factor để fit context cũ
NTK-aware scalingScale base theta, dim cao ít, dim thấp nhiều
YaRNNTK + attention temperature, extension 16-32x ổn
ALiBiBias attention score theo khoảng cách
Lost in middleModel attention yếu ở 30-70% context
Needle in haystackTest retrieval fact từ context dài
KV cache scalingO(N) memory, O(N²) compute attention

Quy tắc:

  • RoPE default: train context length là max usable
  • RoPE + YaRN: extend 4-32x với fine-tune nhẹ
  • ALiBi: free extension nhưng quality general kém RoPE
  • Long context bổ sung RAG, không thay thế
  • Context cost = memory + compute + latency + quality drop

Lời kết

Position encoding là một trong những phần “khoa học hơn nghệ thuật” nhất của LLM. Một thay đổi nhỏ trong cách encode position có thể mở khoá khả năng context dài 16x mà không retrain.

Hiểu RoPE và scaling giúp bạn:

  • Chọn đúng model long-context cho use case
  • Debug khi context dài cho output kém
  • Biết “context 1M” trên paper có nghĩa gì thực tế
  • Tránh trap “stuff everything into context” thay vì RAG

Hands-on cho bạn:

  1. Download Llama-3-8B-Instruct và Llama-3-8B-Instruct-128k. Test needle-in-haystack trên cả hai với context 30k. So sánh accuracy.
  2. Đọc YaRN paper (Peng et al. 2023), tập trung section 3 (method). Code không khó, ý tưởng đẹp.
  3. Thử fine-tune một model 8k với LoRA + data dài 16k. Quan sát position embedding behavior.
  4. Build benchmark needle-in-haystack đơn giản: 10 needle position khác nhau, đo accuracy. Test với model bạn quan tâm.

Bài tiếp theo: Reasoning models: o1, R1, chain-of-thought training. 2024 là năm reasoning model bùng nổ. OpenAI o1, DeepSeek R1, Claude với extended thinking. Chúng có gì khác model thường, và tại sao “suy nghĩ trước khi trả lời” lại là breakthrough chứ không phải prompt trick.