Chuyển đến nội dung chính
Kiến thức nền AI Infra

3.5 Hiểu sâu positional encoding trong Transformer

Positional encoding trông nhỏ trong kiến trúc Transformer nhưng lại rất quan trọng

Transformer Positional encoding RoPE Sinusoidal ALiBi

Positional encoding trông nhỏ trong kiến trúc Transformer nhưng lại rất quan trọng. Không có nó, Transformer không phân biệt được “mèo ăn cá” và “cá ăn mèo”. Bài này xuất phát từ câu hỏi gốc vì sao cần thông tin vị trí, lần lượt giải thích nguyên lý và cách triển khai của Sinusoidal encoding, learned encoding, RoPE, ALiBi và các phương án phổ biến khác, rồi đi sâu vào kỹ thuật ngoại suy positional encoding khi mở rộng ngữ cảnh dài, cuối cùng nối sang thực hành AI Infra: fuse CUDA kernel và quản lý KV Cache.

📑 Mục lục


1. Vì sao Transformer cần thông tin vị trí

1.1 Trật tự từ quyết định nghĩa

Ngôn ngữ tự nhiên là hệ thống ký hiệu phụ thuộc rất mạnh vào thứ tự. Cùng vài từ, chỉ đổi chỗ cũng có thể mang nghĩa hoàn toàn khác.

Vài nhóm ví dụ:

  • “Anh ấy cho tôi mượn một cuốn sách” vs “Tôi cho anh ấy mượn một cuốn sách” — người cho mượn và người mượn đổi chỗ
  • “Chỉ có bạn hiểu tôi” vs “Chỉ có tôi hiểu bạn” — chủ và khách đảo ngược hoàn toàn
  • “Anh ấy không biết cô ấy đã đến” vs “Cô ấy không biết anh ấy đã đến” — ai biết và ai không biết bị đảo

Ngay ở mức tinh hơn, đổi thứ tự từ cũng làm thay đổi quan hệ bổ nghĩa và sắc thái:

  • “kết quả rất tốt” vs “tốt kết quả rất” — câu sau không hợp ngữ pháp
  • “Tôi hôm nay đặc biệt vui” vs “Hôm nay tôi đặc biệt vui” — sắc thái nhấn tinh tế (câu trước nhấn “tôi”, câu sau nhấn “hôm nay”)

Các ví dụ này nói lên một sự thật cơ bản: trật tự từ là một phần của nghĩa. Mọi mô hình ngôn ngữ bỏ qua trật tự từ đều không hiểu đúng ngôn ngữ tự nhiên.

1.2 Tính equivariant với hoán vị của Attention: góc nhìn toán

Giờ ta nêu chặt chẽ vì sao Self-Attention thô không cảm nhận được vị trí.

Nhắc lại phép tính cốt lõi của Self-Attention. Cho ma trận đầu vào X∈RN×dX \in \mathbb{R}^{N \times d} (N token, mỗi token d chiều), quá trình là:

Q=XWQ,K=XWK,V=XWVQ = XW_Q, \quad K = XW_K, \quad V = XW_V Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

Giờ xét một phép hoán vị bất kỳ trên đầu vào. Gọi Π\Pi là ma trận hoán vị N×NN \times N (mỗi hàng và mỗi cột đúng một số 1, còn lại là 0). ΠX\Pi X xáo hàng của X. Q, K, V sau hoán vị là:

Q′=ΠXWQ=ΠQ,K′=ΠK,V′=ΠVQ' = \Pi X W_Q = \Pi Q, \quad K' = \Pi K, \quad V' = \Pi V

Khi đó đầu ra Attention sau hoán vị là:

softmax(Q′K′Tdk)V′=softmax(ΠQKTΠTdk)ΠV\text{softmax}\left(\frac{Q' K'^T}{\sqrt{d_k}}\right) V' = \text{softmax}\left(\frac{\Pi Q K^T \Pi^T}{\sqrt{d_k}}\right) \Pi V

Vì softmax chạy theo từng hàng, còn Π\Pi chỉ xếp lại hàng, nên:

softmax(ΠQKTΠTdk)=Π⋅softmax(QKTdk)⋅ΠT\text{softmax}\left(\frac{\Pi Q K^T \Pi^T}{\sqrt{d_k}}\right) = \Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) \cdot \Pi^T

Kết quả cuối:

Π⋅softmax(QKTdk)⋅ΠT⋅ΠV=Π⋅softmax(QKTdk)V\Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) \cdot \Pi^T \cdot \Pi V = \Pi \cdot \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V

Đó là ý nghĩa của permutation equivariance (đẳng biến với hoán vị): xáo đầu vào thì đầu ra chỉ bị xáo cùng cách, còn nội dung thông tin mỗi token gom lại không đổi. Nói cách khác, Attention thô coi đầu vào như một tập hợp chứ không phải chuỗi — token ở vị trí 0 và vị trí 99 với nó không khác gì nhau.

Đó là lý do phải thêm positional encoding từ bên ngoài: bản thân Attention là “mù vị trí”, cần bơm tín hiệu vị trí từ ngoài vào mới phân biệt được các trật tự từ khác nhau.


2. Absolute positional encoding vs relative positional encoding

Trước khi vào từng phương án cụ thể, nên tách hai hướng thiết kế.

2.1 Absolute positional encoding

Absolute positional encoding gán mỗi vị trí trong chuỗi một vector mã hóa cố định pip_i, rồi cộng (hoặc nối) vào word embedding của token tương ứng:

xi′=xi+pix_i' = x_i + p_i

Ý tưởng rất thẳng: vị trí 0 có một “thẻ căn cước” p0p_0, vị trí 1 có thẻ khác p1p_1, và cứ thế. Mô hình học quan hệ giữa các thẻ này để gián tiếp cảm nhận vị trí tương đối.

Phương án đại diện: Sinusoidal encoding (cố định), learned encoding của BERT/GPT-2 (học được).

2.2 Relative positional encoding

Relative positional encoding không gắn nhãn cố định cho từng vị trí, mà khi tính Attention score thì bơm thẳng thông tin “hai token cách nhau bao xa”.

Ví dụ: absolute encoding giống phát cho mỗi người một số phòng (“tôi ở 301, bạn ở 305”) rồi suy khoảng cách từ hiệu số; relative encoding thì nói thẳng khi hai người nói chuyện rằng “các bạn cách nhau 4 phòng”.

Phương án đại diện: RoPE (rotary position embedding), ALiBi (linear bias).

2.3 So sánh hai hướng

ChiềuAbsolute positional encodingRelative positional encoding
Điểm bơm thông tinTầng đầu vào (cộng vào embedding)Trong lúc tính Attention
Độ bền của thông tin vị tríCó thể loãng dần khi mạng sâu hơnMỗi tầng Attention bơm lại
Khả năng ngoại suyThường yếu hơnThường mạnh hơn
Tổng quát hóa độ dàiKhó xử lý độ dài chưa thấy khi trainTương đối dễ mở rộng chuỗi dài hơn

Hầu hết large model hiện nay dùng relative positional encoding (chủ yếu RoPE). Lý do sẽ phân tích ở phần sau.


3. Sinusoidal positional encoding

3.1 Công thức toán

Bài “Attention Is All You Need” đề xuất positional encoding không cần tham số học được. Với vị trí pospos và chỉ số chiều ii, giá trị mã hóa là:

PE(pos,2i)=sin⁡(pos100002i/d)PE(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d}}\right) PE(pos,2i+1)=cos⁡(pos100002i/d)PE(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d}}\right)

Trong đó pospos là vị trí token trong chuỗi (bắt đầu từ 0), ii là chỉ số cặp chiều (0≤i<d/20 \leq i < d/2), dd là hidden dimension của mô hình.

Mỗi cặp chiều kề (2i,2i+1)(2i, 2i+1) chia sẻ cùng tần số ωi=1/100002i/d\omega_i = 1 / 10000^{2i/d}, lần lượt dùng sine và cosine. Khi ii đi từ 0 đến d/2−1d/2 - 1, tần số giảm theo hàm mũ từ 1 (khi i=0i=0, mẫu số là 100000=110000^0 = 1) xuống 1/100001/10000 (khi i=d/2−1i = d/2 - 1).

3.2 Trực giác đa tần số: ẩn dụ chỉnh radio

Công thức trông trừu tượng lúc đầu, nhưng trực giác phía sau khá tự nhiên.

Tưởng tượng trước mặt bạn một hàng núm radio từ trái sang phải. Núm ngoài cùng bên trái chỉnh tần số rất cao — kim quay rất nhanh, mỗi bước một vị trí gần như một vòng. Núm ngoài cùng bên phải chỉnh rất thấp — kim quay cực chậm, đi trăm vị trí mới lệch một góc nhỏ. Các núm giữa giảm tần số lần lượt.

Giờ cho một số vị trí (ví dụ pos=42), bạn xoay mỗi núm theo góc tương ứng. Vì tần số mỗi núm khác nhau, họa tiết góc kim cho pos=42 là duy nhất. Đổi số vị trí (ví dụ pos=43), mọi kim lệch ở mức khác nhau và tạo họa tiết duy nhất khác.

Đó chính là cách Sinusoidal encoding hoạt động:

  • Chiều tần số cao (ii nhỏ) đổi nhanh, phân biệt tốt vị trí sát nhau (pos=42 và pos=43 lệch rõ trên các chiều này)
  • Chiều tần số thấp (ii lớn) đổi chậm, phân biệt vị trí xa nhau (pos=0 và pos=1000 lệch rõ, còn pos=42 và pos=43 gần như giống nhau)

Tất cả tần số cùng nhau tạo cho mỗi vị trí một “dấu vân phổ” riêng.

3.3 Python: sinh và vẽ Sinusoidal positional encoding

Đoạn mã dưới đây sinh ma trận Sinusoidal positional encoding và vẽ heatmap:

import numpy as np
import matplotlib.pyplot as plt

def sinusoidal_position_encoding(max_len, d_model):
    """
    Sinh ma trận Sinusoidal positional encoding.

    Tham số:
        max_len: độ dài chuỗi tối đa
        d_model: hidden dimension của mô hình

    Trả về:
        PE: ma trận positional encoding hình dạng (max_len, d_model)
    """
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]          # (max_len, 1)
    dim_pair_index = np.arange(0, d_model, 2)[np.newaxis, :]  # (1, d_model/2)

    # Tần số: omega_i = 1 / 10000^(2i/d_model)
    # Tương đương exp(-2i/d_model * ln(10000))
    div_term = np.exp(dim_pair_index * -(np.log(10000.0) / d_model))

    # Chiều chẵn dùng sin, chiều lẻ dùng cos
    PE[:, 0::2] = np.sin(position * div_term)
    PE[:, 1::2] = np.cos(position * div_term)

    return PE

# Sinh positional encoding
max_len = 128
d_model = 64
PE = sinusoidal_position_encoding(max_len, d_model)

# Trực quan hóa
fig, axes = plt.subplots(1, 2, figsize=(16, 6))

# Heatmap: toàn bộ ma trận mã hóa
im = axes[0].imshow(PE, aspect='auto', cmap='RdBu_r', interpolation='nearest')
axes[0].set_xlabel('Dimension Index')
axes[0].set_ylabel('Position')
axes[0].set_title('Sinusoidal Position Encoding Matrix')
plt.colorbar(im, ax=axes[0])

# Đường cong: giá trị mã hóa theo vị trí ở vài chiều chọn sẵn
dims_to_plot = [0, 1, 10, 11, 30, 31]
for dim in dims_to_plot:
    label = f'dim {dim} ({"sin" if dim % 2 == 0 else "cos"}, freq_idx={dim // 2})'
    axes[1].plot(PE[:, dim], label=label, alpha=0.8)
axes[1].set_xlabel('Position')
axes[1].set_ylabel('Encoding Value')
axes[1].set_title('Encoding Values at Selected Dimensions')
axes[1].legend(fontsize=8)
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('sinusoidal_pe.png', dpi=150)
plt.show()

Chạy đoạn này bạn sẽ thấy: trên heatmap, cột trái (tần số cao) là sọc dày, cột phải (tần số thấp) là dải màu rộng, phản ánh rõ các tần số khác nhau. Trên đường cong, chiều chỉ số thấp dao động mạnh, chiều chỉ số cao đổi chậm.

3.4 Vị trí tương đối có thể lấy bằng biến đổi tuyến tính

Sinusoidal encoding có một tính chất toán đẹp: độ lệch giữa hai vị trí bất kỳ có thể biểu diễn bằng một biến đổi tuyến tính cố định. Nghĩa là mô hình có khả năng học được thông tin vị trí tương đối.

Cụ thể, với mọi độ lệch cố định kk, tồn tại ma trận MkM_k chỉ phụ thuộc kk sao cho:

PE(pos+k)=Mk⋅PE(pos)PE(pos + k) = M_k \cdot PE(pos)

Lý do nằm ở công thức cộng góc của hàm lượng giác. Với một tần số ωi\omega_i, xét cặp chiều (2i,2i+1)(2i, 2i+1):

[sin⁡(ωi(pos+k))cos⁡(ωi(pos+k))]=\begin{bmatrix} \sin(\omega_i (pos + k)) \\ \cos(\omega_i (pos + k)) \end{bmatrix} = [cos⁡(ωik)sin⁡(ωik)−sin⁡(ωik)cos⁡(ωik)][sin⁡(ωi⋅pos)cos⁡(ωi⋅pos)]\begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix} \begin{bmatrix} \sin(\omega_i \cdot pos) \\ \cos(\omega_i \cdot pos) \end{bmatrix}

Ma trận xoay bên phải chỉ phụ thuộc độ lệch kk và tần số ωi\omega_i, không phụ thuộc vị trí tuyệt đối pospos. Ghép mọi cặp chiều lại được ma trận block-diagonal MkM_k.

Nghĩa là dù Sinusoidal encoding là “absolute positional encoding”, cấu trúc lượng giác bên trong vẫn cho mô hình khả năng rút vị trí tương đối bằng cách học biến đổi tuyến tính. Nhưng khả năng này là gián tiếp — mô hình phải tự học từ dữ liệu để dùng tính chất toán đó, chứ không được bơm tường minh tín hiệu vị trí tương đối.

3.5 Hạn chế

Dù thiết kế tinh tế, Sinusoidal encoding dần bị thay trong kỷ nguyên large model, chủ yếu vì:

  1. Tín hiệu bị loãng: positional encoding được cộng vào word embedding ở tầng đầu vào. Sau hàng chục tầng biến đổi phi tuyến, tín hiệu cộng này có thể bị “át”, khiến mạng sâu khó dùng hiệu quả thông tin vị trí.

  2. Khó ngoại suy: về lý thuyết có thể sinh encoding cho độ dài bất kỳ, nhưng mô hình khi train chỉ thấy chuỗi hữu hạn. Với vị trí vượt độ dài train, giá trị encoding vẫn có định nghĩa, nhưng mô hình chưa từng train trên các vị trí đó, nên Attention pattern thường tổng quát hóa kém.

  3. Vị trí tương đối gián tiếp: như mục 3.4, mô hình phải tự học rút vị trí tương đối qua biến đổi tuyến tính. Điều đó tăng gánh học so với việc bơm thẳng vị trí tương đối vào Attention.


4. Learned positional encoding

4.1 Ý tưởng cơ bản

Vì encoding lượng giác cố định có hạn chế, ý tưởng tự nhiên là: để mô hình tự học vector mã hóa cho mỗi vị trí. Đó là Learned Positional Embedding.

Cách làm cụ thể: giữ một bảng embedding học được hình dạng (Lmax,d)(L_{max}, d), trong đó LmaxL_{max} là độ dài chuỗi tối đa hỗ trợ, dd là chiều mô hình. Encoding ở vị trí pospos là hàng thứ pospos của bảng, rồi cộng vào word embedding của token:

xi′=xi+Epos[i]x_i' = x_i + E_{pos}[i]

Cơ chế giống hệt word embedding — word embedding tra theo token ID, position embedding tra theo position ID.

4.2 Mô hình đại diện

BERT: dùng learned positional encoding, độ dài tối đa 512. Ma trận position embedding có 512 x 768 = 393,216 tham số, rất nhỏ so với khoảng 110 triệu tham số của BERT-base.

GPT-2: cũng dùng learned positional encoding, độ dài tối đa 1024. Ma trận là 1024 x 768 (GPT-2 Small) hoặc 1024 x 1600 (GPT-2 XL).

4.3 Ưu và nhược

Ưu:

  • Triển khai cực đơn giản: một lần tra bảng embedding rồi cộng
  • Mô hình tự do học mọi pattern vị trí, không bị khóa bởi dạng hàm cố định
  • Trong phạm vi độ dài train, hiệu năng thường không kém, thậm chí hơi tốt hơn Sinusoidal encoding

Nhược:

  • Giới hạn độ dài cứng: độ dài tối đa đã cố định khi định nghĩa mô hình. Muốn chuỗi dài hơn phải train lại hoặc nội suy phức tạp
  • Không ngoại suy được: vị trí ngoài bảng không có vector mã hóa tương ứng
  • Chung khuyết điểm với absolute encoding: tín hiệu vẫn có thể loãng ở tầng sâu, vị trí tương đối vẫn phải học gián tiếp

Chính những hạn chế này đẩy nhà nghiên cứu tìm phương án mã hóa vị trí tương đối trực tiếp trong Attention — và RoPE là câu trả lời thành công nhất.


5. RoPE: rotary position embedding

RoPE (Rotary Position Embedding) do Jianlin Su (苏剑林) đề xuất năm 2021, hiện được gần như mọi large model chủ đạo dùng: LLaMA, Mistral, Qwen, DeepSeek, v.v. Ý tưởng cốt lõi gói trong một câu: xoay vector Q và K để mã hóa vị trí, khiến Attention score tự nhiên mang thông tin vị trí tương đối.

5.1 Bắt đầu từ xoay 2D: ẩn dụ la bàn

Trước khi vào suy luận toán, dựng trực giác hình học.

Tưởng tượng bạn cầm một la bàn, kim chỉ một hướng, biểu diễn bằng vector 2D (x,y)(x, y). Xoay la bàn một góc θ\theta, kim chỉ hướng mới (x′,y′)(x', y'). Phép xoay không đổi độ dài (norm), chỉ đổi hướng.

Biểu thức toán của xoay 2D:

[x′y′]=[cos⁡θ−sin⁡θsin⁡θcos⁡θ][xy]\begin{bmatrix} x' \\ y' \end{bmatrix} = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix}

Sáng kiến cốt lõi của RoPE: dùng vị trí token để quyết định góc xoay. Token vị trí 0 không xoay, vị trí 1 xoay θ\theta, vị trí 2 xoay 2θ2\theta, và cứ thế. Token vị trí mm xoay mθm\theta.

Xét hai token: vector Q ở vị trí mm bị xoay mθm\theta, vector K ở vị trí nn bị xoay nθn\theta. Khi tính tích trong (Attention score), vì xoay là biến đổi cứng, tích trong chỉ phụ thuộc hiệu góc (m−n)θ(m - n)\theta, không phụ thuộc góc xoay tuyệt đối của từng vector.

Đó là bản chất hình học giúp RoPE mã hóa vị trí tương đối: tích trong của hai vector sau khi mỗi cái bị xoay chỉ phụ thuộc hiệu góc xoay.

5.2 Suy luận toán cho trường hợp 2D

Giờ suy luận chặt. Xét trường hợp 2D, Q và K lần lượt là query và key ở vị trí mm và nn. RoPE áp xoay phụ thuộc vị trí:

q~=R(mθ)⋅q=[cos⁡(mθ)−sin⁡(mθ)sin⁡(mθ)cos⁡(mθ)][q0q1]\tilde{q} = R(m\theta) \cdot q = \begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix} \begin{bmatrix} q_0 \\ q_1 \end{bmatrix} k~=R(nθ)⋅k=[cos⁡(nθ)−sin⁡(nθ)sin⁡(nθ)cos⁡(nθ)][k0k1]\tilde{k} = R(n\theta) \cdot k = \begin{bmatrix} \cos(n\theta) & -\sin(n\theta) \\ \sin(n\theta) & \cos(n\theta) \end{bmatrix} \begin{bmatrix} k_0 \\ k_1 \end{bmatrix}

Tích trong sau xoay:

q~Tk~=qTR(mθ)TR(nθ)k\tilde{q}^T \tilde{k} = q^T R(m\theta)^T R(n\theta) k

Dùng tính chất R(α)T=R(−α)R(\alpha)^T = R(-\alpha):

R(mθ)TR(nθ)=R(−mθ)R(nθ)=R((n−m)θ)R(m\theta)^T R(n\theta) = R(-m\theta) R(n\theta) = R((n-m)\theta)

Cuối cùng:

q~Tk~=qTR((n−m)θ)k\tilde{q}^T \tilde{k} = q^T R((n-m)\theta) k

Kết quả chỉ phụ thuộc (n−m)(n - m), tức hiệu vị trí tương đối của hai token, không phụ thuộc giá trị tuyệt đối của mm, nn. Đó đúng là tính chất ta muốn.

5.3 Mở rộng lên chiều cao

Trong mô hình thật, head dimension dkd_k lớn hơn nhiều so với 2 (thường 64 hoặc 128). Cách mở rộng của RoPE rất tự nhiên: tách vector dkd_k chiều thành dk/2d_k / 2 không gian con 2D, xoay độc lập trong mỗi không gian con nhưng dùng tần số khác nhau.

Với không gian con 2D thứ ii (0≤i<dk/20 \leq i < d_k / 2), góc xoay là m⋅θim \cdot \theta_i, với:

θi=1100002i/dk\theta_i = \frac{1}{10000^{2i / d_k}}

Lưu ý tham số tần số này trùng khít với Sinusoidal encoding.

Ghép mọi ma trận xoay của các không gian con thành một ma trận block-diagonal:

Rm=[R(mθ0)R(mθ1)⋱R(mθdk/2−1)]R_m = \begin{bmatrix} R(m\theta_0) & & & \\ & R(m\theta_1) & & \\ & & \ddots & \\ & & & R(m\theta_{d_k/2 - 1}) \end{bmatrix}

Mỗi R(mθi)R(m\theta_i) là ma trận xoay 2×22 \times 2.

Phép RoPE đầy đủ là:

q~m=Rm⋅qm,k~n=Rn⋅kn\tilde{q}_m = R_m \cdot q_m, \quad \tilde{k}_n = R_n \cdot k_n

5.4 Vì sao tích trong chỉ phụ thuộc hiệu vị trí tương đối: chứng minh chiều cao

Ở chiều cao, chứng minh vẫn gọn. Tích trong sau xoay:

q~mTk~n=qmTRmTRnkn\tilde{q}_m^T \tilde{k}_n = q_m^T R_m^T R_n k_n

Vì RmR_m là block-diagonal, RmTR_m^T cũng vậy (mỗi khối lấy chuyển vị):

RmTRn=[R(mθ0)TR(nθ0)⋱R(mθdk/2−1)TR(nθdk/2−1)]R_m^T R_n = \begin{bmatrix} R(m\theta_0)^T R(n\theta_0) & & \\ & \ddots & \\ & & R(m\theta_{d_k/2-1})^T R(n\theta_{d_k/2-1}) \end{bmatrix}

Mỗi khối đường chéo:

R(mθi)TR(nθi)=R(−mθi)R(nθi)=R((n−m)θi)R(m\theta_i)^T R(n\theta_i) = R(-m\theta_i) R(n\theta_i) = R((n-m)\theta_i)

Vậy:

RmTRn=Rn−mR_m^T R_n = R_{n-m}

Kết quả cuối:

q~mTk~n=qmTR_n−mkn=f(qm,kn,n−m)\tilde{q}_m^T \tilde{k}_n = q_m^T R\_{n-m} k_n = f(q_m, k_n, n - m)

Tích trong chỉ phụ thuộc bản thân vector và hiệu vị trí tương đối n−mn - m, không phụ thuộc vị trí tuyệt đối. Xong.

5.5 Triển khai hiệu quả: không dựng tường minh ma trận xoay

Trong kỹ thuật thật, ta không dựng ma trận block-diagonal dk×dkd_k \times d_k rồi nhân. Xoay 2D làm bằng nhân và cộng theo từng phần tử:

q~2i=q2icos⁡(mθi)−q2i+1sin⁡(mθi)\tilde{q}_{2i} = q_{2i} \cos(m\theta_i) - q_{2i+1} \sin(m\theta_i) q~2i+1=q2isin⁡(mθi)+q2i+1cos⁡(mθi)\tilde{q}_{2i+1} = q_{2i} \sin(m\theta_i) + q_{2i+1} \cos(m\theta_i)

Dưới đây là Python cho phép RoPE cốt lõi:

import numpy as np

def precompute_freqs(dim, max_seq_len, base=10000.0):
    """
    Tiền tính tham số tần số và giá trị lượng giác cho RoPE.

    Tham số:
        dim: chiều mỗi attention head (d_k)
        max_seq_len: độ dài chuỗi tối đa
        base: cơ số tần số, mặc định 10000

    Trả về:
        cos_cached: cosine hình dạng (max_seq_len, dim/2)
        sin_cached: sine hình dạng (max_seq_len, dim/2)
    """
    # Tần số: theta_i = 1 / base^(2i / dim), i = 0, 1, ..., dim/2 - 1
    freq_indices = np.arange(0, dim, 2, dtype=np.float32)  # [0, 2, 4, ..., dim-2]
    freqs = 1.0 / (base ** (freq_indices / dim))            # (dim/2,)

    # Dãy vị trí
    positions = np.arange(max_seq_len, dtype=np.float32)     # [0, 1, 2, ..., max_seq_len-1]

    # Tích ngoài: góc cho mỗi (position, freq_pair)
    angles = np.outer(positions, freqs)  # (max_seq_len, dim/2)

    cos_cached = np.cos(angles)
    sin_cached = np.sin(angles)

    return cos_cached, sin_cached

def apply_rope(q, cos, sin):
    """
    Áp xoay RoPE lên vector query hoặc key.

    Tham số:
        q: vector hình dạng (seq_len, dim); dim phải chẵn
        cos: cosine tiền tính hình dạng (seq_len, dim/2)
        sin: sine tiền tính hình dạng (seq_len, dim/2)

    Trả về:
        q_rotated: vector sau xoay, cùng hình dạng đầu vào
    """
    dim = q.shape[-1]
    assert dim % 2 == 0, "chiều phải là số chẵn"

    # Tách chiều chẵn và lẻ
    q_even = q[:, 0::2]   # (seq_len, dim/2) -- lấy q_0, q_2, q_4, ...
    q_odd  = q[:, 1::2]   # (seq_len, dim/2) -- lấy q_1, q_3, q_5, ...

    # Xoay từng không gian con 2D
    q_even_rotated = q_even * cos - q_odd * sin
    q_odd_rotated  = q_even * sin + q_odd * cos

    # Xen kẽ rồi ghép lại
    q_rotated = np.zeros_like(q)
    q_rotated[:, 0::2] = q_even_rotated
    q_rotated[:, 1::2] = q_odd_rotated

    return q_rotated

# Ví dụ dùng
dim = 128        # head dimension
seq_len = 64     # độ dài chuỗi

# Tiền tính lượng giác (chỉ làm một lần lúc khởi tạo)
cos_cached, sin_cached = precompute_freqs(dim, seq_len)

# Giả lập vector Q và K
q = np.random.randn(seq_len, dim).astype(np.float32)
k = np.random.randn(seq_len, dim).astype(np.float32)

# Áp RoPE
q_rotated = apply_rope(q, cos_cached, sin_cached)
k_rotated = apply_rope(k, cos_cached, sin_cached)

# Kiểm tra: tích trong sau xoay chỉ phụ thuộc vị trí tương đối
# Q ở 10 và K ở 15
dot1 = np.dot(q_rotated[10], k_rotated[15])

# Q ở 20 và K ở 25 (cùng hiệu vị trí tương đối 5)
# Lưu ý: giá trị gốc của Q và K ở đây khác nên tích trong khác
# Nhưng với cùng vector q, k đặt ở vị trí tuyệt đối khác nhau,
# miễn hiệu vị trí tương đối giống nhau thì tích trong giống nhau
print(f"dot(q_rotated[10], k_rotated[15]) = {dot1:.4f}")

Lưu ý: hai cách ghép cặp chiều. Triển khai trên ghép hai chiều kề (q2i,q2i+1)(q_{2i}, q_{2i+1}), đúng phong cách LLaMA gốc của Meta và GPT-J. Bản LLaMA trong HuggingFace transformers ghép nửa trước với nửa sau, tức (qi,qi+d/2)(q_i, q_{i+d/2}), qua rotate_half, thuộc phong cách GPT-NeoX:

def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return np.concatenate((-x2, x1), axis=-1)

# Phong cách HF: cần mở rộng cos/sin thành (seq_len, dim), tức lặp bảng dim/2 hai lần trên trục cuối
q_rotated = q * cos + rotate_half(q) * sin

Hai cách tương đương về toán, nhưng mỗi cách đòi hỏi chiều đầu ra của WqW_q, WkW_k sắp theo đúng thứ tự tương ứng. Script chuyển đổi của HF (convert_llama_weights_to_hf.py) khi đưa trọng số Meta sang định dạng HF sẽ permute một lần WqW_q, $W_k`. Vì vậy khi port mô hình hoặc tự viết RoPE kernel, phải xác nhận layout trọng số khớp cách ghép cặp: lệch nhau thì chương trình không báo lỗi, nhưng đầu ra mô hình sẽ sai hoàn toàn. Đây là bẫy phổ biến khi port mô hình.

5.6 Liên hệ và khác biệt giữa RoPE và Sinusoidal

RoPE và Sinusoidal encoding có liên hệ sâu, nhưng cách hoạt động rất khác.

Liên hệ:

  • Cả hai dùng cùng tham số tần số θi=1/100002i/d\theta_i = 1 / 10000^{2i/d}
  • Cả hai dựa vào tính tuần hoàn và công thức cộng góc của hàm lượng giác
  • Tính chất biến đổi tuyến tính của vị trí tương đối trong Sinusoidal encoding (mục 3.4) về bản chất là một phép xoay; RoPE có thể xem như làm tường minh tính chất đó

Khác biệt:

ChiềuSinusoidalRoPE
Cách bơmCộng vào word embeddingXoay vector Q và K
Vị trí bơmTầng đầu vào (chỉ một lần)Mỗi tầng Attention (mỗi tầng đều bơm)
Loại mã hóaAbsolute positional encodingRelative positional encoding
Ảnh hưởng tới VGián tiếp (đầu vào của V đã mang thông tin vị trí)Không ảnh hưởng (chỉ xoay Q và K)
Giữ được ở tầng sâuTín hiệu vị trí có thể suy giảm trong mạng sâuMỗi tầng bơm lại, không suy giảm
Khả năng ngoại suyYếu hơnMạnh hơn (kết hợp nội suy)

Một khác biệt đáng nhấn: RoPE không biến đổi V. Thông tin vị trí chỉ ảnh hưởng “ai nên chú ý ai” (tính Attention weight), không ảnh hưởng “sau khi chú ý thì truyền thông tin gì” (nội dung V). Thiết kế này sạch hơn — vị trí điều khiển định tuyến luồng thông tin, không phải bản thân thông tin.


6. ALiBi: attention với linear bias

ALiBi (Attention with Linear Biases) do Press et al. đề xuất năm 2022, đi theo hướng rất khác RoPE.

6.1 Ý tưởng cốt lõi

ALiBi hoàn toàn không sửa Q và K. Sau khi tính Attention score, nó cộng thẳng một bias âm tỷ lệ với khoảng cách:

Attention(Q,K,V)=softmax(QKTdk+m⋅Bias)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + m \cdot \text{Bias}\right) V

Phần tử của ma trận Bias là:

Bias[i][j]=−∣i−j∣\text{Bias}[i][j] = -|i - j|

mm là hệ số dốc vô hướng theo từng head. Các head khác nhau dùng mm khác nhau (thường là cấp số nhân như 1/2,1/4,1/8,…1/2, 1/4, 1/8, \ldots).

Nói thường: ALiBi đặt một “phạt khoảng cách” lên Attention score — hai token càng xa thì bị trừ càng nhiều, nên mô hình nghiêng về chú ý token gần. Các head suy giảm với tốc độ khác nhau: head có dốc lớn gần như chỉ nhìn token sát bên, head có dốc nhỏ vẫn nhìn được token xa.

6.2 Ưu và hạn chế

Ưu:

  • Triển khai cực đơn giản: chỉ cần một phép cộng trên Attention score
  • Không thêm tham số học được
  • Ngoại suy mạnh: bài báo chứng minh mô hình train ở độ dài 1024 có thể ngoại suy lên 2048 thậm chí dài hơn lúc inference
  • Chi phí tính gần như bằng không

Hạn chế:

  • Khả năng biểu diễn thông tin vị trí khá hạn chế (chỉ một kiểu: suy giảm tuyến tính theo khoảng cách)
  • Trong ngữ cảnh rất dài, phạt tuyến tính có thể át quá mạnh Attention tới token xa
  • Trong thực tế large model, tỷ lệ dùng ALiBi thấp hơn nhiều so với RoPE

ALiBi hiện chủ yếu được một số ít mô hình như BLOOM dùng. Trong các phương án positional encoding cho large model, RoPE với khả năng biểu đạt và mở rộng ngoại suy tốt đã trở thành chuẩn thực tế.


7. Kỹ thuật mở rộng ngữ cảnh dài

Large model khi train thường dùng cửa sổ ngữ cảnh cố định (ví dụ 4096 hoặc 8192 token), nhưng người dùng thường cần đầu vào dài hơn — một cuốn tiểu thuyết, một đoạn mã dài, cả buổi họp. Điều đó đặt ra thách thức kỹ thuật cốt lõi: làm sao để mô hình lúc inference xử lý chuỗi vượt độ dài train?

7.1 Thách thức ngoại suy positional encoding

Lấy RoPE làm ví dụ. Nếu độ dài chuỗi tối đa khi train là 4096 thì mm trong m⋅θim \cdot \theta_i tối đa là 4095. Khi inference gặp vị trí m=8000m = 8000, một số chiều tần số cao xoay vượt xa phạm vi từng thấy lúc train, khiến Attention pattern đổi mạnh.

Cụ thể, chiều tần số cao (ii nhỏ, θi\theta_i lớn) lúc train đã “quay nhiều vòng”, quay thêm một chút ít ảnh hưởng; còn chiều tần số thấp (ii lớn, θi\theta_i nhỏ) lúc train có thể chưa đầy một vòng, đột nhiên yêu cầu hai vòng thì bước vào “vùng chưa biết”. Ngoại suy không đều giữa các chiều tần số là gốc của vấn đề.

7.2 Position Interpolation

Giải pháp sơ đẳng nhất là Position Interpolation: thu nhỏ chỉ số vị trí theo tỷ lệ sao cho vị trí lớn nhất mới tương ứng với vị trí lớn nhất lúc train.

Muốn mở rộng ngữ cảnh từ LtrainL_{train} lên LtargetL_{target}, hệ số scale là s=Ltarget/Ltrains = L_{target} / L_{train}, góc RoPE ở vị trí mm trở thành:

m′⋅θi=ms⋅θim' \cdot \theta_i = \frac{m}{s} \cdot \theta_i

Trực giác: encoding vốn phủ 0 đến 4095 bị “nén” vào khoảng 0 đến 4095/s4095/s, giải phóng không gian mã hóa cho chuỗi dài hơn.

Vấn đề: scale đều như nhau với mọi chiều tần số, nhưng chiều tần số cao không cần scale (khả năng ngoại suy đã tốt), chiều tần số thấp mới cần scale mạnh. Nội suy đều sẽ làm hại khả năng phân biệt vị trí gần của chiều tần số cao.

7.3 NTK-aware Scaling

Insight cốt lõi của NTK-aware Scaling (Neural Tangent Kernel-aware Scaling) là: không nên scale đều mọi chiều tần số; mỗi tần số cần chiến lược scale khác nhau.

Cách làm cụ thể: sửa tham số base của RoPE (giá trị gốc 10000), phóng thành:

base′=base⋅sd/(d−2)\text{base}' = \text{base} \cdot s^{d/(d-2)}

trong đó ss là hệ số mở rộng. Hiệu ứng:

  • Chiều tần số cao (ii nhỏ): tần số gần như không đổi, giữ khả năng phân biệt vị trí sát
  • Chiều tần số thấp (ii lớn): tần số giảm rõ, bước sóng dài ra, phủ được phạm vi vị trí lớn hơn

Chữ “NTK” trong tên đến từ động cơ lý thuyết: trong khung Neural Tangent Kernel, phân bố tần số của đặc trưng Fourier ảnh hưởng quan trọng tới khả năng tổng quát hóa, và NTK-aware Scaling thiết kế điều chỉnh tần số từ góc đó.

7.4 Cải tiến của YaRN

YaRN (Yet another RoPE extensioN) thêm hai cải tiến quan trọng trên NTK-aware Scaling:

Cải tiến thứ nhất: chiến lược phân đoạn. YaRN chia chiều theo cao thấp tần số thành ba nhóm, mỗi nhóm xử lý khác:

  • Chiều tần số cao (bước sóng ngắn hơn độ dài train gốc / một hệ số nào đó): không scale gì, giữ nguyên hành vi RoPE gốc
  • Chiều tần số thấp (bước sóng dài hơn độ dài train × một hệ số nào đó): scale nội suy tuyến tính đầy đủ
  • Dải giữa: chuyển tiếp mượt giữa “không scale” và “scale đầy đủ”

Chiến lược phân đoạn này tinh hơn điều chỉnh liên tục của NTK-aware, vừa giữ độ phân giải gần vừa mở tốt khả năng xa.

Cải tiến thứ hai: hiệu chỉnh nhiệt độ phân bố Attention. Sau khi mở rộng ngữ cảnh, phân bố Attention score đổi (vì số token trong tổng tăng). YaRN đưa vào hệ số nhiệt độ gắn với tỷ lệ mở rộng để sửa lệch phân bố đó.

7.5 Dynamic NTK

Các phương pháp trên cần biết trước độ dài mục tiêu và đặt tham số scale trước khi inference. Nhưng khi triển khai thật, độ dài đầu vào thay đổi động — request ngắn có, dài có. Scale mở rộng độ dài cho cả request ngắn là thừa, thậm chí hại hiệu năng.

Ý tưởng Dynamic NTK: điều chỉnh base theo độ dài chuỗi thực tế hiện tại.

Cụ thể, trong lúc inference, độ dài hiện tại LcurrentL_{current} tăng dần khi sinh token. Nếu Lcurrent≤LtrainL_{current} \leq L_{train}, dùng base gốc (không scale); nếu Lcurrent>LtrainL_{current} > L_{train}, tính base mới theo tỷ lệ Lcurrent/LtrainL_{current} / L_{train}.

Ưu điểm:

  • Chuỗi ngắn không bị ảnh hưởng, giữ độ chính xác gốc
  • Chuỗi dài tự nhận scale phù hợp
  • Không cần biết trước độ dài tối đa

8. Thực hành AI Infra

Lựa chọn và triển khai positional encoding ảnh hưởng trực tiếp tới kỹ thuật AI Infra. Phần này bàn vài liên hệ kỹ thuật quan trọng.

8.1 Fuse RoPE trong CUDA kernel

Phép cốt lõi của RoPE là nhân-cộng theo phần tử trên Q và K (mục 5.5). Về lượng tính toán, phép này rất nhẹ (O(N⋅d)O(N \cdot d)), nhỏ hơn nhiều so với phép nhân ma trận trong Attention (O(N2⋅d)O(N^2 \cdot d)). Nhưng nếu chạy như một CUDA kernel độc lập, chi phí kernel launch và đọc/ghi VRAM thừa sẽ thành nút thắt.

Vì vậy trong engine inference hiệu năng cao, RoPE thường được fuse vào phần đầu của Attention kernel:

  1. Fuse với chiếu QKV: sau phép chiếu tuyến tính QKV, xoay Q và K ngay trong cùng kernel, tránh ghi kết quả trung gian về HBM rồi đọc lại
  2. Fuse với FlashAttention: trong vòng tiling của FlashAttention, khi Q, K của mỗi block đã vào SRAM, xoay RoPE tại chỗ rồi tính Attention ngay

Lợi ích sau khi fuse:

  • Bớt một lần kernel launch (tiết khoảng 5-10 µs; ở giai đoạn Decode mỗi bước chỉ xử lý 1 token thì tỷ trọng chi phí này không nhỏ)
  • Bớt một vòng đọc/ghi HBM (đặc biệt quan trọng với Decode vốn memory-bound)

Ngoài ra, giá trị lượng giác của RoPE (cos⁡(mθi)\cos(m\theta_i) và sin⁡(mθi)\sin(m\theta_i)) thường được tiền tính và cache thành bảng tra, tránh tính lại lượng giác mỗi lần inference. Trong triển khai phổ biến bảng này là một tensor thường nằm ở bộ nhớ toàn cục GPU (HBM) — ví dụ cos_sin_cache của vLLM, hình dạng (max_position, rotary_dim); kernel đọc hàng tương ứng từ global memory theo position của từng token.

8.2 Ảnh hưởng ngữ cảnh dài tới KV Cache

Khả năng ngoại suy của positional encoding quyết định trực tiếp mô hình hỗ trợ ngữ cảnh dài đến đâu, còn tăng độ dài ngữ cảnh tạo áp lực lớn lên quản lý KV Cache.

Lấy LLaMA-2-7B làm ví dụ (theo tính toán trong bài mở đầu Transformer):

Độ dài ngữ cảnhKV Cache mỗi request (FP16)Tổng batch=16
4K2 GB32 GB
32K16 GB256 GB
128K64 GB1024 GB

Có thể thấy, khi ngữ cảnh từ 4K lên 128K, nhu cầu VRAM của KV Cache tăng 32 lần. Nghĩa là:

  • Một request đã có thể hết VRAM: KV Cache một request 128K (64 GB) đã gần bằng toàn bộ dung lượng một A100
  • Khả năng đồng thời tụt mạnh: dù có đủ VRAM, số request xử lý cùng lúc cũng giảm mạnh
  • PagedAttention càng quan trọng: trong ngữ cảnh dài, KV Cache tăng động và phân mảnh nặng hơn, cần quản lý bộ nhớ tinh hơn

Vì vậy kỹ thuật mở rộng ngữ cảnh dài của positional encoding (mục 7.2-7.5) không chỉ là vấn đề năng lực mô hình, mà còn là thách thức cấp hệ thống. Về kỹ thuật cần giải đồng thời:

  1. Tầng positional encoding: chọn phương án ngoại suy/nội suy phù hợp (YaRN, Dynamic NTK, …), bảo đảm chất lượng Attention trên chuỗi dài
  2. Tầng quản lý KV Cache: paging, lượng tử hóa KV Cache (FP8/INT4), chiến lược eviction, v.v. để kiểm soát chi phí VRAM
  3. Tầng tính toán: đặc tính O(N)O(N) VRAM của FlashAttention cực kỳ quan trọng với ngữ cảnh dài (nếu không, ma trận Attention O(N2)O(N^2) sẽ nuốt lượng VRAM khổng lồ)

8.3 So sánh kỹ thuật các phương án positional encoding

Phương ánTính toán thêm lúc inferenceTham số thêmĐộ phức tạp triển khaiThân thiện ngữ cảnh dài
SinusoidalTra bảng tiền tính + cộng vectorKhôngThấpKém
Learned PETra embedding + cộng vectorLmax×dL_{max} \times dRất thấpKém (trần cứng)
RoPEXoay theo phần tử (có thể fuse)KhôngTrung bìnhTốt (kết hợp nội suy)
ALiBiCộng bias lên Attention scoreKhôngThấpTốt (ngoại suy tự nhiên)

Từ góc thực hành kỹ thuật, RoPE trở thành chuẩn thực tế không chỉ vì lý thuyết đẹp, mà còn vì thân thiện kỹ thuật: không thêm tham số (không ảnh hưởng model parallelism), fuse mượt vào Attention kernel sẵn có, và có hệ sinh thái mở rộng ngữ cảnh dài chín muồi.


🎯 Danh sách tự kiểm

Sau khi học bài này, tự kiểm xem bạn đã thực sự hiểu positional encoding chưa:

  • Giải thích được vì sao Self-Attention thô không phân biệt trật tự từ, và nêu ngắn gọn ý nghĩa toán của permutation equivariance
  • Phân biệt được hướng thiết kế absolute vs relative positional encoding, nêu phương án đại diện của từng bên
  • Viết thuộc công thức Sinusoidal encoding và giải thích vai trò của các chiều tần số khác nhau
  • Dùng công thức cộng góc lượng giác để giải thích vì sao vị trí tương đối trong Sinusoidal encoding có được qua biến đổi tuyến tính
  • Nói rõ phép cốt lõi của RoPE (xoay 2D) và suy ra vì sao tích trong sau xoay chỉ phụ thuộc hiệu vị trí tương đối
  • Viết được mã triển khai RoPE theo phần tử, giải thích vai trò của việc tiền tính cos/sin
  • Nói rõ khác biệt cốt lõi giữa bốn phương án RoPE, Sinusoidal, Learned PE, ALiBi
  • Giải thích vì sao NTK-aware Scaling tốt hơn Position Interpolation đều (giữ tần số cao, kéo tần số thấp)
  • Ước lượng được ảnh hưởng mở rộng ngữ cảnh dài tới VRAM của KV Cache, nêu ít nhất hai chiến lược đối phó
  • Giải thích vì sao RoPE thường được fuse vào Attention kernel để chạy

📚 Tài liệu tham khảo

Bài báo

Blog và tutorial