RoPE giải mã: Cơ chế hoạt động của nhúng vị trí xoay (kèm mã PyTorch tối ưu hóa GPU)

Giới thiệu Hãy hình dung việc đọc một cuốn sách mà tất cả các từ được viết trên những mảnh giấy riêng biệt, sau đó được bỏ vào một chiếc mũ và trộn lẫn. Để hiểu câu chuyện, bạn sẽ phải rút từng từ ra, đoán vị trí của nó và tái tạo lại các câu trong đầu. Đây chính xác là cách một mô hình Transformer cơ bản nhìn nhận ngôn ngữ của con người. Khi bài báo mang tính bước ngoặt “Attention Is All You Need” ra mắt vào năm 2017, nó đã thay đổi cơ bản bối cảnh AI bằng cách giới thiệu cơ chế Self-Attention (Tự chú ý). Cơ chế này cho phép các mạng thần kinh xử lý tất cả các token trong một chuỗi đồng thời theo cách song song, phá vỡ hoàn toàn các nút thắt cổ chai tuần tự, chậm chạp của các mạng thần kinh hồi quy (RNN) và mạng bộ nhớ dài ngắn hạn (LSTM) cũ hơn. Tuy nhiên, bước nhảy vọt lớn về tốc độ huấn luyện này đi kèm với một cái giá cấu trúc nghiêm trọng: các Transformer vốn dĩ không nhận biết được thứ tự từ. Hạn chế của mô hình cơ sở năm 2017 Để khắc phục sự mù mờ về không gian này, các tác giả gốc đã giới thiệu Sinusoidal Absolute Positional Encodings (Mã hóa vị trí tuyệt đối hình sin). Khái niệm này rất đơn giản: tạo một vector sóng tĩnh duy nhất cho mỗi chỉ số vị trí (pos trong {0, 1, 2, …]) và thêm trực tiếp vào các nhúng từ ngữ nghĩa: Đối với các chiều tọa độ chẵn: Đối với các chiều tọa độ lẻ: Mặc dù bản vá toán học này đã giúp các mô hình ban đầu như BERT và GPT-2 hiểu được thứ tự văn bản, nhưng nó đã tạo ra hai nút thắt cổ chai kỹ thuật lớn: Hỏng hóc ngữ nghĩa: Các không gian vector chiều cao là những cấu trúc hình học cực kỳ tinh tế. Việc ép một tín hiệu vị trí vào các tọa độ từ ngữ nghĩa thông qua phép cộng thô làm thay đổi hình học cơ bản của các token. Mạng phải dành các tham số trọng số nội bộ quý giá chỉ để “tách rời” ý nghĩa của một từ khỏi vị trí của nó. Giới hạn ngữ cảnh: Các nhúng tuyệt đối được mã hóa cứng vào một hệ tọa độ cố định. Nếu bạn huấn luyện một mô hình trên một chuỗi có độ dài 512 token, vị trí 1024 là một bí ẩn toán học hoàn toàn đối với nó. Mô hình không thể tổng quát hóa sang các văn bản dài hơn trong thời gian suy luận mà không cần huấn luyện lại hoàn toàn từ đầu. Giới thiệu RoPE (Rotary Position Embeddings) Để thu hẹp khoảng cách giữa thứ tự tuyệt đối và khoảng cách tương đối linh hoạt, Jianlin Su và các đồng tác giả đã giới thiệu Rotary Position Embeddings (RoPE) vào năm 2021. Thay vì coi vị trí là một khoản phí bổ sung, RoPE coi vị trí là một phép quay hình học. Nó ánh xạ các cặp số trong một vector từ lên một mặt phẳng phức 2D và nhẹ nhàng xoay chúng một góc tỷ lệ với chỉ số vị trí của token. Ngày nay, RoPE đã trở thành tiêu chuẩn vàng không thể tranh cãi cho các mô hình ngôn ngữ lớn (LLM) mã nguồn mở tiên tiến. Nếu bạn tìm hiểu sâu bên trong LLaMA 3 của Meta, Mistral, Qwen 2.5 hoặc Gemma của Google, bạn sẽ không tìm thấy một nhúng vị trí bổ sung nào. Tất cả đều chạy trên RoPE. Trong hướng dẫn toàn diện này, chúng ta sẽ phân tích cơ chế toán học chính xác đằng sau phép thuật quay này, loại bỏ các công thức ma trận khối chéo phức tạp và xem xét một thủ thuật Interleaved Slicing Trick được tối ưu hóa cao trong PyTorch từ đầu – tối đa hóa việc sử dụng băng thông bộ nhớ và song song hóa phần cứng trên mỗi token. Trực giác: Một đồng hồ đa tần số Để hiểu cách RoPE mã hóa cả vị trí tuyệt đối và khoảng cách tương đối đồng thời, hãy ngừng suy nghĩ về các vector chiều cao trong giây lát. Thay vào đó, hãy tưởng tượng một căn phòng đầy đồng hồ treo tường truyền thống. Mỗi vector từ trong một Transformer được chia thành các cặp số. Trong RoPE, mỗi cặp tọa độ được gán một đồng hồ riêng biệt. Nếu mô hình của bạn có kích thước đầu chú ý là 64, mỗi token mang một bộ 32 đồng hồ khác nhau. Nhúng Token: [ x1, y1, x2, y2, x3, y3, ... ] │ │ │ │ │ │ └───┘ └───┘ └───┘ Đồng hồ 1 Đồng hồ 2 Đồng hồ 3 Khi một token đi vào mô hình, chỉ số vị trí của nó xác định số phút trôi qua trên các đồng hồ của nó: Một token ở vị trí 1 làm kim đồng hồ tiến thêm 1 vạch. Một token ở vị trí 5 làm kim đồng hồ tiến thêm 5 vạch. Đây là cách vị trí tuyệt đối được ghi nhận. Nhưng đây là điểm mấu chốt: mỗi đồng hồ trong bộ chạy với một tốc độ (tần số) hoàn toàn khác nhau. Hệ thống bánh răng (Tần số) Đồng hồ 1 (Đồng hồ nhanh): Tích tắc nhanh chóng. Với mỗi chỉ số vị trí, kim của nó quay một góc lớn (ví dụ: 45° mỗi token). Nó rất nhạy cảm với ngữ cảnh cục bộ, tức thời. Đồng hồ 2 (Đồng hồ trung bình): Tích tắc chậm hơn. Nó có thể chỉ quay 15° mỗi token. Đồng hồ 32 (Đồng hồ ốc sên): Tích tắc cực kỳ chậm. Nó hầu như không di chuyển một phần nhỏ của độ mỗi token. Nó được thiết kế để nắm bắt các mối quan hệ tầm xa trên hàng nghìn token. Cách giải quyết "Bức tường ngữ cảnh" Bởi vì kim đồng hồ di chuyển mượt mà dọc theo một vòng tròn liên tục thay vì nhảy dọc theo một lưới tọa độ tuyệt đối cố định, mô hình không quan tâm đến việc độ dài chuỗi là 512 hay 8192. Nếu mô hình được huấn luyện trên độ dài ngữ cảnh 512, nó đã học được cách các kim đồng hồ hoạt động tương đối với nhau. Nếu nó đột nhiên gặp vị trí 1024 trong quá trình suy luận, các kim đồng hồ chỉ tiếp tục quay mượt mà quanh mặt đồng hồ. Mối quan hệ hình học vẫn hoàn toàn nguyên vẹn. Tại sao điều này tạo ra sự hiểu biết về "Khoảng cách tương đối" Hãy tưởng tượng Token A ở vị trí 2 và Token B ở vị trí 5. Khoảng cách giữa chúng là 3. Khi cơ chế tự chú ý so sánh Token A và Token B (thông qua tích vô hướng), về cơ bản nó đang đo góc giữa các kim đồng hồ của chúng. Bởi vì một vòng quay 5 vạch trừ đi một vòng quay 2 vạch luôn để lại một sự khác biệt ròng là 3 vạch, mô hình có thể tự nhiên suy ra chính xác khoảng cách giữa các từ, bất kể chúng xuất hiện ở đâu trong câu. Từ phép quay 2D đến RoPE Hãy tập trung kỹ vào toán học cơ bản về cách chúng ta quay các vectơ trong không gian 2D. Như đã lưu ý, RoPE không coi một vectơ từ là một thực thể lớn. Thay vào đó, nó chia một nhúng d chiều thành d/2 khối 2D độc lập (các cặp tọa độ như x và y). Đối với bất kỳ cặp nào, chúng ta biểu diễn nó dưới dạng một vectơ z trên một mặt phẳng phức với bán kính ban đầu r và góc α: Khi một token xuất hiện ở một chỉ số vị trí m cụ thể, chúng ta