ALiBi — Phạt khoảng cách trực tiếp, generalize dài hơn training

Bạn muốn deploy model với context window 8K tokens, nhưng training dài như vậy tốn GPU gấp 4 lần và thời gian gấp bội. Có cách nào để "train ngắn, test dài" không? ALiBi là câu trả lời: thay vì dạy model "bạn đang ở tọa độ nào", ta dạy model "càng xa càng ít nhìn", và điều kỳ diệu là model tự extrapolate lên sequence dài chưa từng thấy.

Vấn đề

Cái bẫy của positional encoding truyền thống: Sinusoidal hay learned absolute position embeddings giống như dán mã vạch GPS lên mỗi token. Khi train với max 2048 tokens, model học rằng "tọa độ 3000" không tồn tại. Đưa vào inference 4096 tokens, model nhìn thấy mã vạch lạ hoắc và bối rối — perplexity tăng vọt, attention pattern tan rã.

RoPE (Rotary Position Embedding) giải quyết bằng cách xoay vector, nhưng khi cần extend context, bạn phải làm Position Interpolation phức tạp — giãn vòng xoay, calibrate lại tần số, rồi fine-tune thêm vài ngàn steps để model quen với "vũ trụ bị giãn nở".

Vấn đề sâu hơn: training context dài tốn KV Cache khủng khiếp. Với 7B model, 8K context cần ~40GB VRAM chỉ cho cache. Nếu có thể train ngắn (rẻ) mà dùng dài (đắt) sau đó, ta tiết kiệm được cả núi tiền compute.

Ý tưởng cốt lõi

ALiBi không thêm position embedding vào input. Token embeddings giữ nguyên "sạch", không bị nhiễu bởi thông tin vị trí tuyệt đối. Thay vào đó, nó phạt trực tiếp trên attention scores: càng xa thì càng bị trừ điểm nặng.

Cụ thể, trước khi softmax, ALiBi thêm vào attention score một hằng số âm tuyến tính theo khoảng cách:

$\text{Attention}(Q, K) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}} - m \cdot |i - j|\right)$

Trong đó $m$ là slope (độ dốc) học được cho mỗi head, thường khởi tạo theo cấp số nhân từ $2^{-8}$ đến $2^{-1}$.

Tại sao điều này tạo ra "soft masking"? Softmax của một số âm lớn $\rightarrow$ 0. Token cách xa 10 bước bị penalty $-10m$, qua softmax thì attention weight $\propto e^{-10m}$ — thực tế bằng 0 nếu $m$ đủ lớn. Bạn không cần hard-coded window size như Sliding Window Attention; penalty tuyến tính tự động tạo ra "vùng mờ" exponential decay mà không cần mask cứng.

Đây là phần quan trọng nhất: ALiBi chỉ quan tâm khoảng cách tương đối, không quan tâm vị trí tuyệt đối. Vì penalty chỉ phụ thuộc $|i-j|$ (khoảng cách giữa token $i$ và $j$), khi bạn dịch cả sequence đi 1000 bước, pattern attention giữ nguyên. Model không cần học "tọa độ mới", nó chỉ cần áp dụng quy tắc "càng xa càng mờ" vô hạn định. Đây là lý do extrapolation hoạt động — bạn train trên 1024 tokens, model tự động generalize lên 2048, 4096, thậm chí 1M tokens mà không cần đụng vào code.

Hierarchical Receptive Fields: Mỗi attention head học slope $m$ khác nhau. Head với $m$ lớn (gần 1) chỉ nhìn local — như ống nhòm zoom cao. Head với $m$ nhỏ nhìn global — như camera góc rộng. Model không cần quyết định "local hay global", nó có cả hai: population coding cho context length. Đây là lý do ALiBi không cần Grouped Query Attention hay Multi-Query Attention để tiết kiệm memory, nhưng kết hợp cả hai thì càng tốt.

Nhiều người nói: "ALiBi là linear bias, nên attention sẽ linear decay, không đủ mạnh để chặn long-range dependencies hoàn toàn." — Thực ra misleading. Softmax biến linear penalty thành exponential decay trong weight space. Điều này đủ mạnh để tạo inductive bias local nhưng không cứng rắn như hard cut-off, giữ lại khả năng "nhìn xa" khi cần thiết.

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

Toán học đằng sau sự "mờ": Khi attention score bị trừ đi $m|i-j|$, qua softmax ta có:

$\alpha_{ij} \propto \exp\left(\frac{q_i^T k_j}{\sqrt{d_k}} - m|i-j|\right) = \exp\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right) \cdot \exp(-m|i-j|)$

Thành phần $\exp(-m|i-j|)$ là một hàm exponential decay theo khoảng cách. Với $m=0.5$, token cách 10 bước có weight giảm $e^{-5} \approx 0.0067$ lần so với token cách 0. Đây là soft locality — bạn không cần Sliding Window Attention với window size cố định.

Tính tương đối cho phép extrapolation: Vì $|i-j|$ bất biến khi dịch chuyển toàn bộ sequence, attention pattern giữ nguyên dù bạn đưa vào 1000 token hay 10000 token. Điều này khác biệt với sinusoidal encoding, nơi tần số cố định gây ra aliasing khi sequence vượt quá training length.

Code snippet minh họa (PyTorch-style pseudo-code):

# ALiBi attention bias
def get_alibi_slopes(n_heads):
    # Geometric progression from 2^-8 to 2^-1
    return 2.0 ** (-8 * torch.arange(1, n_heads + 1) / n_heads)

def alibi_attention(Q, K, V):
    # Q, K, V: [batch, heads, seq, dim]
    seq_len = Q.size(-2)
    
    # Create distance matrix |i - j|
    distances = torch.arange(seq_len, device=Q.device).unsqueeze(0) - \
                torch.arange(seq_len, device=Q.device).unsqueeze(1)
    distances = distances.abs().unsqueeze(0).unsqueeze(0)  # [1, 1, seq, seq]
    
    # Get slopes per head: [1, heads, 1, 1]
    slopes = get_alibi_slopes(Q.size(1)).view(1, -1, 1, 1).to(Q.device)
    
    # Compute attention with linear bias
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(Q.size(-1))
    scores = scores - slopes * distances  # ALiBi penalty
    probs = F.softmax(scores, dim=-1)
    
    return torch.matmul(probs, V)

Ý nghĩa thực tế

So sánh với các phương pháp khác:

Benchmarks thực tế:

• BLOOM (176B): Train với 2048 context, deploy với 8192 context không cần position interpolation, perplexity degradation tối thiểu.
• MPT-7B/30B (MosaicML): Sử dụng ALiBi để đạt context window 8K+ với training cost thấp hơn RoPE-based models 30-40%.
• WikiText-103: ALiBi outperform sinusoidal và learned embeddings trên các sequence dài hơn training context.

Ai đang dùng: BigScience (BLOOM), MosaicML (MPT series), Salesforce (CodeGen), Stability AI. Nhiều Vision Transformers cũng thử nghiệm ALiBi cho positional encoding trong ảnh (DINOv2 analysis).

Limitations cần biết:

• Fixed recency bias: Luôn penalize khoảng cách, có thể hurt tasks đòi hỏi long-range dependencies nhất định (ví dụ: reasoning cần nhìn lại đoạn mở đầu của văn bản dài).
• Trade-off với RoPE: Trên training length chuẩn, RoPE đôi khi đạt perplexity thấp hơn nhẹ; ALiBi trade absolute performance trên training length để lấy extrapolation.
• Very long contexts (>16k): Dù extrapolate tốt hơn APE, vẫn có thể cần Position Interpolation (Chen et al. 2023) để đạt chất lượng tối ưu ở 32k-128k tokens.

Đào sâu hơn

Paper gốc: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (Press et al., 2021) — Giới thiệu ALiBi và chứng minh khả năng extrapolation.

Paper liên quan: Position Interpolation Improves ALiBi Extrapolation (Chen et al., 2023) — Kết hợp ALiBi với interpolation để đạt context cực dài.

Cùng cụm (Attention Efficiency):

• Flash Attention — Cùng phép toán, nhanh 3x nhờ tiling trên SRAM
• Multi-Query Attention — Share K/V heads, KV cache nhỏ hẳn
• Grouped Query Attention — Goldilocks giữa MHA và MQA
• Sliding Window Attention — Chỉ nhìn N token gần nhất, đủ rồi
• RoPE — Xoay vector thay vì cộng position, elegant hơn hẳn

Đọc tiếp:

• KV Cache — Đừng tính lại cái đã tính, cache K và V (nền tảng để hiểu tại sao context dài tốn memory)
• Long Context Extrapolation (Level 2) — YaRN & LongRoPE: Kéo dài context window sau training (so sánh với ALiBi)
• Native Long-Context Training (Level 2) — Train dài từ đầu, reason tốt hơn extend (khi ALiBi không đủ)

External resources:

• SambaNova ALiBi Analysis — So sánh extrapolation vs interpolation cho production models.
• Position Encoding Comparison — Phân tích toàn diện sinusoidal, RoPE, và ALiBi.