RoFormer: Mô hình Transformer cải tiến với Rotary Position Embedding

Hướng dẫn đọc

Vấn đề và động lực

Cơ chế self-attention trong Transformer vốn có tính bất biến với hoán vị (permutation-invariant), đòi hỏi phải có thông tin vị trí rõ ràng. Các phương pháp trước đây nhìn chung chia thành hai nhóm: nhúng vị trí tuyệt đối được cộng trực tiếp vào biểu diễn từ (các vector học được hoặc vector dạng sin cố định) và nhúng vị trí tương đối được tích hợp vào ma trận attention (ví dụ: Shaw et al., Transformer-XL, T5, DeBERTa). Tuy nhiên, các phương pháp tương đối hiện tại phần lớn dựa vào các phép phân rã cộng làm thay đổi logit attention trung gian, khiến chúng không tương thích với các cơ chế linear self-attention dựa trên kernel vốn đòi hỏi biểu diễn có thể phân tích nhân tử. Hơn nữa, các mô hình thường gặp khó khăn trong việc mở rộng quy mô một cách liền mạch trên các độ dài chuỗi biến thiên.

Các ý tưởng cốt lõi

• Khoảng cách tương đối qua phép quay vector: Thay vì cộng các vector vị trí, RoPE áp dụng ma trận quay phụ thuộc vào vị trí lên các query và key sau biến đổi afin. Tích vô hướng giữa query đã quay tại chỉ số $m$ và key đã quay tại chỉ số $n$ tự nhiên chỉ phụ thuộc vào khoảng cách tương đối $m-n$.
• Tính trực giao và bảo toàn chuẩn: Do toán tử quay là trực giao, nó bảo toàn chuẩn Euclidean của query và key, giúp ổn định các trạng thái ẩn.
• Sự suy giảm theo khoảng cách dài (Long-term Decay): Tuân theo tỷ lệ tần số dạng sin, tương tác giữa các token tự nhiên suy giảm khi khoảng cách tương đối tăng lên, phù hợp với trực giác ngôn ngữ rằng các token ở xa có mức độ phụ thuộc yếu hơn.
• Tương thích với Linear Attention: Vì thông tin vị trí được áp dụng trực tiếp lên các vector token trước khi tính tích vô hướng, RoPE có thể kết hợp với linear attention dựa trên kernel (chẳng hạn như Performer).

Phương pháp (với các phương trình chính)

Mục tiêu là tìm các hàm $f_q({\mathbf{x}}_m,m)$ và $f_k({\mathbf{x}}_n,n)$ sao cho tích vô hướng của chúng chỉ phụ thuộc vào các embedding và khoảng cách tương đối: $$⟨f_q({\mathbf{x}}_m,m),f_k({\mathbf{x}}_n,n)⟩=g({\mathbf{x}}_m,{\mathbf{x}}_n,m-n)$$

Trong không gian con 2D, coi các vector thực như các số phức mang lại nghiệm: $$f_q({\mathbf{x}}_m,m)=({\mathbf{W}}_q{\mathbf{x}}_m)e^{im\theta },f_k({\mathbf{x}}_n,n)=({\mathbf{W}}_k{\mathbf{x}}_n)e^{in\theta }$$

Với chiều không gian nhúng chẵn $d$, không gian được chia thành $d/2$ không gian con 2D với các tần số $\Theta =\{{\theta }_i=10000^{-2(i-1)/d},i\in [1,\dots ,d/2]\}$. Phép biến đổi tổng quát áp dụng ma trận quay đường chéo khối ${\mathbf{R}}_{\Theta ,m}^d$: $$f_{\{q,k\}}({\mathbf{x}}_m,m)={\mathbf{R}}_{\Theta ,m}^d{\mathbf{W}}_{\{q,k\}}{\mathbf{x}}_m$$

Điểm attention thu được giữa các vị trí $m$ và $n$ thỏa mãn: $${\mathbf{q}}_m^{⊺}{\mathbf{k}}_n=({\mathbf{R}}_{\Theta ,m}^d{\mathbf{W}}_q{\mathbf{x}}_m{)}^{⊺}({\mathbf{R}}_{\Theta ,n}^d{\mathbf{W}}_k{\mathbf{x}}_n)={\mathbf{x}}_m^{⊺}{\mathbf{W}}_q{\mathbf{R}}_{\Theta ,n-m}^d{\mathbf{W}}_k{\mathbf{x}}_n$$

Để tính toán hiệu quả mà không cần phép nhân ma trận thưa, RoPE được triển khai thông qua các phép toán theo từng phần tử (element-wise): $${\mathbf{R}}_{\Theta ,m}^d\mathbf{x}=\left(\begin{array}{c}{x}_1\\ x_2\\ ⋮\\ x_{d-1}\\ x_d\end{array}\right)\otimes \left(\begin{array}{c}\cos m{\theta }_1\\ \cos m{\theta }_1\\ ⋮\\ \cos m{\theta }_{d/2}\\ \cos m{\theta }_{d/2}\end{array}\right)+\left(\begin{array}{c}-x_2\\ x_1\\ ⋮\\ -x_d\\ x_{d-1}\end{array}\right)\otimes \left(\begin{array}{c}\sin m{\theta }_1\\ \sin m{\theta }_1\\ ⋮\\ \sin m{\theta }_{d/2}\\ \sin m{\theta }_{d/2}\end{array}\right)$$

Kết quả (với các số liệu quan trọng nhất)

• Dịch máy: Trên tập WMT 2014 Anh-Đức, RoFormer đạt 27.5 BLEU, vượt trội hơn mô hình cơ sở Transformer-base (27.3 BLEU).
• Độ hội tụ tiền huấn luyện: Khi huấn luyện BERT-base với RoPE trên BookCorpus và Wikipedia, RoFormer hội tụ nhanh hơn và đạt loss masked language modeling (MLM) thấp hơn so với BERT thông thường.
• Fine-tuning trên GLUE: RoFormer tiền huấn luyện cải thiện đáng kể điểm số hạ nguồn trên MRPC (89.5 so với 88.9 F1), STS-B (87.0 so với 85.8 Spearman) và QQP (86.4 so với 71.2 F1), dù nó đạt kết quả thấp hơn BERT trên SST-2 (90.7 so với 93.5), QNLI (88.0 so với 90.5) và MNLI-m/mm (80.2/79.8 so với 84.6/83.4).
• Linear Attention (Performer): Trên tập dữ liệu Enwik8, việc thêm RoPE vào Performer cấp ký tự 12 lớp đã giúp đẩy nhanh đáng kể tốc độ hội tụ và giảm loss huấn luyện.
• Khớp ngữ nghĩa văn bản dài: Trên CAIL2019-SCM, RoFormer-512 đạt 68.29% độ chính xác kiểm tra (so với 67.77% của BERT-512 và 68.10% của WoBERT-512). Khi mở rộng lên RoFormer-1024, độ chính xác kiểm tra tăng lên 69.79% (cải thiện tuyệt đối 1.5% so với WoBERT-512).

Hạn chế và các câu hỏi mở

• Bài báo lưu ý việc thiếu lời giải thích lý thuyết cho lý do tại sao công thức quay 2D lại hội tụ nhanh hơn trong quá trình tiền huấn luyện so với các phương án dạng cộng.
• Mặc dù đặc tính suy giảm khoảng cách dài đã được chứng minh về mặt toán học, lý do tại sao RoFormer tổng quát hóa tốt hơn rõ rệt so với các mô hình cùng loại trên các chuỗi dài vẫn là một câu hỏi mở.
• Việc đánh giá RoFormer qua quá trình tiền huấn luyện quy mô lớn đòi hỏi lượng tài nguyên GPU tính toán đáng kể, để lại các đặc tính mở rộng quy mô rộng hơn cho các nghiên cứu tương lai.