Rotary Position Embedding (RoPE)

Mã hóa vị trí xoay (RoPE) cung cấp cho transformer thông tin về thứ tự token bằng cách xoay từng vector query và key một góc tỷ lệ thuận với vị trí của token. Bản thân self-attention không thể phân biệt thứ tự từ, vì vậy model phải đưa thông tin vị trí vào theo một cách nào đó; với RoPE, điểm attention giữa query ở vị trí m và key ở vị trí n chỉ phụ thuộc vào độ lệch m − n. Được giới thiệu năm 2021 cùng model RoFormer và phổ biến rộng rãi nhờ GPT-J và GPT-NeoX của EleutherAI, RoPE hiện là mã hóa vị trí tiêu chuẩn của các model ngôn ngữ lớn (LLM) mã nguồn mở, bao gồm Llama, Mistral, Qwen, Gemma và DeepSeek.

Cách RoPE hoạt động#

RoPE chia mỗi vector query và key có chiều d thành d/2 cặp đặc trưng, rồi xem mỗi cặp như một điểm trên mặt phẳng. Ở vị trí m, cặp i được xoay một góc m·θᵢ, với θᵢ = base^(−2i/d) và base bằng 10,000 trong bài báo gốc. Các cặp đầu xoay nhanh, giúp phân biệt thứ tự cục bộ tinh vi; các cặp cuối xoay chậm và vẫn phân biệt được các token cách nhau hàng nghìn vị trí. Xoay query một góc mθ và key một góc nθ khiến tích vô hướng của chúng chỉ phụ thuộc vào hiệu số (m − n)θ; phép xoay không bao giờ làm thay đổi độ dài vector. Phép xoay được áp dụng trong từng lớp attention, sau các phép chiếu query và key; vector value không bị xoay, và không có gì được cộng vào embedding token ở đầu vào:

Đoạn code dưới đây triển khai phép xoay bằng số phức, giống như code tham chiếu Llama của Meta, rồi kiểm tra tính chất vị trí tương đối:

import torch

torch.manual_seed(0)
dim = 64
freqs = 10000 ** (-torch.arange(0, dim, 2) / dim)  # one rotation frequency per pair of features


def rope(x, pos):
    """Rotate each consecutive feature pair of x by the angle pos * frequency."""
    pairs = torch.view_as_complex(x.reshape(-1, 2))
    return torch.view_as_real(pairs * torch.polar(torch.ones_like(freqs), pos * freqs)).flatten()


q, k = torch.randn(2, dim)
# The same query and key, three tokens apart, at two different places in the sequence
print(f"positions 5 and 2:     {rope(q, 5) @ rope(k, 2):.4f}")
print(f"positions 105 and 102: {rope(q, 105) @ rope(k, 102):.4f}")
print(f"positions 5 and 4:     {rope(q, 5) @ rope(k, 4):.4f}")  # one token apart
print(f"Vector length kept: {rope(q, 5).norm():.4f} vs {q.norm():.4f}")

Kết quả in ra là -8.2811 cho cả hai cặp cách nhau ba token, dù chúng nằm ở vị trí 5 và 2 hay 105 và 102; kết quả là -11.4563 khi cùng các vector đó cách nhau một token. Vector query sau khi xoay vẫn giữ độ dài 8.3705.

RoPE so với các mã hóa vị trí khác#

• Mã hóa tuyệt đối hình sin: Transformer ban đầu cộng các vector sin và cos cố định vào embedding đầu vào đúng một lần, trước lớp đầu tiên. Khoảng cách tương đối phải được suy ra từ tín hiệu đã trộn vào.
• Embedding tuyệt đối có thể học: BERT và GPT-2 học một vector cho mỗi vị trí, tối đa lần lượt 512 và 1,024 vị trí; chúng không có embedding cho vị trí vượt quá giới hạn đó.
• Độ lệch vị trí tương đối: T5 và ALiBi thêm độ lệch theo khoảng cách có thể học hoặc cố định vào điểm attention. RoPE mã hóa khoảng cách tương đối thông qua chính phép xoay và không thêm tham số.

Mở rộng cửa sổ ngữ cảnh#

Model được huấn luyện với RoPE trên các chuỗi có độ dài L chưa từng thấy góc xoay của những vị trí vượt quá L, và attention suy giảm khi gặp các vị trí đó. Các phương pháp mở rộng ngữ cảnh điều chỉnh lại góc thay vì huấn luyện lại từ đầu. Position Interpolation nén các vị trí mới vào phạm vi đã huấn luyện, mở rộng model Llama lên cửa sổ ngữ cảnh 32,768 token chỉ trong 1,000 bước fine-tune. NTK-aware scaling và YaRN kéo giãn các cặp tần số thấp, xoay chậm nhiều hơn các cặp xoay nhanh; YaRN mở rộng Llama 2 lên 128K token với số token huấn luyện ít hơn 10 lần so với các phương pháp trước đó. Llama 3 tăng base từ 10,000 lên 500,000 trong quá trình tiền huấn luyện, khiến các cặp xoay chậm nhất quay ít hơn trên mỗi token và giữ được khả năng phân biệt các vị trí xa nhau.

RoPE trong model thị giác#

Ảnh cần hai tọa độ, vì vậy Vision Transformer sử dụng RoPE 2D, hay RoPE trục: một nửa số cặp đặc trưng trong mỗi head xoay theo cột của patch, nửa còn lại xoay theo hàng. Qwen2-VL bổ sung trục thời gian thứ ba cho video trong M-RoPE đa phương thức. SAM 2 của Meta, có trong gói Ultralytics, áp dụng RoPE trục trong bốn lớp memory attention, nơi đặc trưng của khung hình video hiện tại attention đến các đặc trưng đã lưu từ những khung hình trước. Key của mọi khung hình đã lưu nhận cùng một lưới góc như query của khung hình hiện tại, vì vậy attention giữa chúng phụ thuộc vào độ lệch không gian giữa các patch.