您本可以設計出最先進的位置編碼

加爾定律
一個能工作的複雜系統，總是從一個能工作的簡單系統演變而來。
約翰·加爾

本文將逐步引導您瞭解Transformer模型中最先進的位置編碼的發現過程。我們將透過迭代改進位置編碼方法來實現這一目標，最終得到最新LLama 3.2版本和大多數現代Transformer中使用的Rotary Positional Encoding (RoPE)。本文旨在限制所需的數學知識，但需要一些基本的線性代數、三角學和自注意力機制的理解。

問題陳述

知其所伴，便知其義
約翰·魯珀特·弗斯

與所有問題一樣，最好先精確地理解我們試圖實現的目標。Transformer中的自注意力機制用於理解序列中詞元之間的關係。自注意力是一種集合操作，這意味著它具有置換等變性。如果我們不為自注意力機制提供位置資訊，許多重要的關係將無法確定。

這透過一個例子可以很好地說明。

啟發性示例

考慮以下句子，其中同一個詞出現在不同位置

$$\text{The dog chased another dog}$$

直觀上，“dog”指的是兩個不同的實體。讓我們看看如果先對它們進行詞元化，對映到Llama 3.2 1B的真實詞元嵌入，並透過torch.nn.MultiheadAttention會發生什麼。

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
    for param in mha.parameters():
        nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True

正如我們所看到的，在沒有任何位置資訊的情況下，（多頭）自注意力操作的輸出對於不同位置的相同詞元是相同的，儘管這些詞元顯然代表著不同的實體。現在，讓我們開始設計一種方法，用位置資訊增強自注意力，使其能夠根據詞的位置確定它們之間的關係。

理想的位置編碼方案應該如何表現？

理想特性

讓我們嘗試定義一些理想的特性，這將使最佳化過程儘可能簡單。

特性1 - 每個位置的唯一編碼（跨序列）

每個位置都需要一個唯一的編碼，無論序列長度如何，該編碼都應保持一致——在位置5的詞元，無論當前序列長度是10還是10,000，都應該具有相同的編碼。

特性2 - 兩個編碼位置之間的線性關係

位置之間的關係應該在數學上是簡單的。如果我們知道位置$p$的編碼，那麼計算位置$p+k$的編碼就應該很簡單，這使得模型更容易學習位置模式。

如果您考慮我們如何在數軸上表示數字，就很容易理解5距離3有2步，或者10距離15有5步。我們的編碼中應該存在相同的直觀關係。

特性3 - 推廣到比訓練中遇到的序列更長的序列

為了增加模型在現實世界中的實用性，它們應該能夠泛化到訓練分佈之外的資料。因此，我們的編碼方案需要足夠靈活，能夠處理意外的輸入長度，而不會違反任何其他理想特性。

特性4 - 由模型可以學習的確定性過程生成

如果我們的位置編碼能夠從確定性過程中獲取，那將是理想的。這將允許模型有效地學習我們的編碼方案背後的機制。

特性5 - 可擴充套件到多個維度

隨著多模態模型成為常態，我們的位置編碼方案能夠自然地從$1D$擴充套件到$nD$至關重要。這將使模型能夠處理影像或腦部掃描等資料，它們分別是$2D$和$4D$的。

現在我們知道了理想的特性（以下稱為$Pr_n$），讓我們開始設計和迭代我們的編碼方案。

整數位置編碼

首先想到的方法可能是簡單地將詞元位置的整數值新增到詞元嵌入的每個分量中，其值範圍從$0 \rightarrow L$，其中$L$是當前序列的長度。

在上面的動畫中，我們為詞元$\color{#699C52}\text{chased}$建立了位置編碼向量，並將其新增到我們的詞元嵌入中。這裡的嵌入值是Llama 3.2 1B真實值的一個子集。我們可以觀察到它們集中在0附近。這對於避免訓練期間梯度消失或爆炸是理想的，因此是我們在整個模型中希望保持的。

顯然，我們目前這種樸素的方法將導致問題。位置值的幅度遠遠超過我們輸入的實際值。這意味著信噪比非常低，模型很難將語義資訊與位置資訊分離。

有了這些新知識，一個自然的後續想法可能是將位置值標準化為$\frac{1}{N}$。這會將值限制在0到1之間，但引入了另一個問題。如果我們將$N$選為當前序列的長度，那麼對於不同長度的每個序列，位置值將完全不同，這違反了$Pr_1$。

有沒有更好的方法來確保我們的數字在0和1之間？如果我們對此深思熟慮一段時間，我們可能會想到從十進位制數切換到二進位制數。

二進位制位置編碼

我們不必將（可能已標準化的）整數位置新增到嵌入的每個分量中，而是可以將其轉換為二進位制表示，並將其拉伸以匹配我們的嵌入維度，如下所示。

我們已將感興趣的位置（252）轉換為其二進位制表示（11111100），並將每個位新增到詞元嵌入的相應分量中。最低有效位（LSB）將在每個後續詞元之間在0和1之間迴圈，而最高有效位（MSB）將在每$2^{n-1}$個詞元迴圈，其中$n$是位數。您可以在下面的動畫中看到不同索引的位置編碼向量$[^1]$。

我們已經解決了值範圍問題，現在我們有了在不同序列長度下保持一致的唯一編碼。如果我們將詞元嵌入的低維版本繪製出來，並可視化二進位制位置向量對於不同值的新增會發生什麼。

我們可以看到結果非常“跳躍”（正如我們可能從二進位制的離散性質中預期的那樣）。最佳化過程喜歡平滑、連續和可預測的變化。我們知道有哪些具有相似值範圍的函式是平滑和連續的嗎？

如果我們四處看看，我們可能會注意到$\sin$和$\cos$都符合要求！

正弦位置編碼

上述動畫可視化了我們的位置嵌入，其中每個分量都交替從$\sin$和$\cos$中提取，並且波長逐漸增加。如果將其與之前的動畫進行比較，您會發現驚人的相似之處！

我們現在得到了正弦嵌入；最初定義在Attention is all you need論文中。讓我們看看這些方程：

$$PE_{(pos,2i)} = \color{#58C4DD}\sin\left(\color{black}\frac{pos}{10000^{2i/d}}\color{#58C4DD}\right)\color{black} \\ \quad \\ PE_{(pos,2i+1)} = \color{#FC6255}\cos\left(\color{black}\frac{pos}{10000^{2i/d}}\color{#FC6255}\right)\color{black} \\$$

其中$pos$是詞元的位置索引，$i$是位置編碼向量中的分量索引，$d$是模型維度。$10,000$是基準波長（以下簡稱$\theta$），我們根據分量索引對其進行拉伸或壓縮。我鼓勵您代入一些實際值，以感受這種幾何級數。

這個方程中有幾部分乍一看令人困惑。作者是如何選擇$10,000$的？為什麼我們對偶數位置使用$\sin$，對奇數位置使用$\cos$呢？

似乎使用$10,000$作為基準波長是實驗確定的$[^2]$。解釋$\sin$和$\cos$共同使用的原因更為複雜，但對於我們理解的迭代方法至關重要。這裡的關鍵是我們希望兩個編碼位置之間存線上性關係$Pr_2$。要理解$\sin$和$\cos$如何協同產生這種線性關係，我們必須深入研究一些三角學。

考慮一個正弦和餘弦對序列，每個對都與頻率$\omega_i$相關聯。我們的目標是找到一個線性變換矩陣$\mathbf{M}$，它可以將這些正弦函式平移一個固定的偏移量$k$

$$\mathbf{M} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i(p + k)) \\ \cos(\omega_i(p + k)) \end{bmatrix}$$

頻率$\omega_i$遵循幾何級數，其隨維度索引$i$而減小，定義為

$$\omega_i = \frac{1}{10000^{2i/d}}$$

為了找到這個變換矩陣，我們可以將其表示為一個具有未知係數$u_1$、$v_1$、$u_2$和$v_2$的通用2×2矩陣

$$\begin{bmatrix} u_1 & v_1 \\ u_2 & v_2 \end{bmatrix} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i(p+k)) \\ \cos(\omega_i(p+k)) \end{bmatrix}$$

透過將三角和定理應用於右側，我們可以將其展開為

$$\begin{bmatrix} u_1 & v_1 \\ u_2 & v_2 \end{bmatrix} \cdot \begin{bmatrix} \sin(\omega_i p) \\ \cos(\omega_i p) \end{bmatrix} = \begin{bmatrix} \sin(\omega_i p)\cos(\omega_i k) + \cos(\omega_i p)\sin(\omega_i k) \\ \cos(\omega_i p)\cos(\omega_i k) - \sin(\omega_i p)\sin(\omega_i k) \end{bmatrix}$$

透過展開，我們可以透過匹配係數得到一個包含兩個方程的系統

$$\begin{align} u_1\sin(\omega_i p) + v_1\cos(\omega_i p) &= \cos(\omega_i k)\sin(\omega_i p) + \sin(\omega_i k)\cos(\omega_i p) \\ u_2\sin(\omega_i p) + v_2\cos(\omega_i p) &= -\sin(\omega_i k)\sin(\omega_i p) + \cos(\omega_i k)\cos(\omega_i p) \end{align}$$

透過比較兩邊包含 $\sin(\omega_i p)$ 和 $\cos(\omega_i p)$ 的項，我們可以解出未知係數

$$\begin{align} u_1 &= \cos(\omega_i k) & v_1 &= \sin(\omega_i k) \\ u_2 &= -\sin(\omega_i k) & v_2 &= \cos(\omega_i k) \end{align}$$

這些解給了我們最終的變換矩陣 $\mathbf{M_k}$

$$\mathbf{M_k} = \begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix}$$

如果你之前做過遊戲程式設計，你可能會注意到我們推導的結果出奇地熟悉。沒錯，它就是旋轉矩陣！ $[^3]$。

因此，Noam Shazeer 在2017年的論文《Attention is all you need》中設計的編碼方案，早在2017年就已經將相對位置編碼為旋轉了！儘管旋轉的概念早已出現，但從正弦編碼到RoPE，卻又花了**四年**的時間……

絕對位置編碼與相對位置編碼

瞭解旋轉在此處的重要性後，讓我們回到最初的例子，嘗試為下一次迭代尋找一些直覺。

$$\begin{align*} &\hspace{0.7em}0 \hspace{1.4em} 1 \hspace{2em} 2 \hspace{2.6em} 3 \hspace{2.4em} 4\\ &\text{The dog chased another dog} \\ \\ &\hspace{0.3em}\text{-2} \hspace{1.4em} \text{-1} \hspace{1.7em} 0 \hspace{2.6em} 1 \hspace{2.4em} 2\\ &\text{The dog \color{#699C52}chased \color{black}another dog} \end{align*}$$

上面，我們可以看到標記的絕對位置，以及從$\color{#699C52}\text{chased}$到其他每個標記的相對位置。透過正弦編碼，我們生成了一個單獨的向量來表示絕對位置，並使用一些三角函式技巧來編碼相對位置。

當我們試圖理解這些句子時，這個詞是這篇博文中的第2157個詞重要嗎？還是我們更關心它與周圍詞語的關係？一個詞的絕對位置很少影響其含義——重要的是詞語之間的相互關係。

上下文中的位置編碼

從現在開始，關鍵是要**在自注意力機制的上下文中**考慮位置編碼。重申一下，自注意力機制使模型能夠權衡輸入序列中不同元素的重要性，並動態調整它們對輸出的影響。

$$\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

在之前的所有迭代中，我們都生成了一個單獨的位置編碼向量，並在進行$Q$、$K$和$V$投影之前將其**新增**到標記嵌入中。透過將位置資訊直接新增到標記嵌入中，我們**汙染**了語義資訊與位置資訊。我們應該嘗試在不修改範數的情況下編碼資訊。轉向乘法是關鍵。

用字典類比來說，當我們在字典（鍵）中查詢一個詞（查詢）時，附近的詞應該比遠處的詞有更大的影響。一個標記對另一個標記的影響是由$QK^T$點積決定的——所以這正是我們應該關注位置編碼的地方！

$$\vec{a} \cdot \vec{b} = |\vec{a}| |\vec{b}| \cos \theta$$

上面所示點積的幾何解釋給了我們一個了不起的洞察力。我們可以透過單純增加或減少兩個向量之間的角度來調節點積的結果。此外，透過旋轉向量，我們完全不影響向量的範數，而向量的範數編碼了我們標記的語義資訊。

現在我們知道應該把注意力集中在哪裡了，也從另一個角度明白了為什麼旋轉可能是一個編碼位置資訊的“通道”，接下來我們把它們綜合起來！

旋轉**位**置**編**碼 (RoPE)

**旋轉位置編碼**（或稱RoPE）在RoFormer論文中定義（蘇劍林在他的部落格這裡和這裡獨立設計了它）。如果跳到最終結果，它可能看起來像巫術，但透過在自注意力機制（特別是點積）的背景下思考正弦編碼，我們可以看到它是如何協同工作的。

與正弦編碼非常相似，我們將向量$\mathbf{q}$或$\mathbf{k}$（而不是預投影的$\mathbf{x}$）分解成二維對/塊。我們不是透過新增由緩慢減小頻率的正弦函式生成的向量來直接編碼*絕對*位置，而是直接編碼*相對*位置，方法是**將每對與旋轉矩陣相乘**。

令$\mathbf{q}$或$\mathbf{k}$為我們在位置$p$的輸入向量。我們建立一個塊對角矩陣，其中$\mathbf{M_i}$是該分量對所需旋轉的相應旋轉矩陣

$$R(\mathbf{q}, p) = \begin{pmatrix} \mathbf{M_1} & & & \\ & \mathbf{M_2} & & \\ & & \ddots & \\ & & & \mathbf{M_{d/2}} \end{pmatrix} \begin{pmatrix} q_1 \\ q_2 \\ \vdots \\ q_d \end{pmatrix}$$

與正弦編碼（Sinusoidal Encoding）非常相似，$\mathbf{M_i}$ 簡單來說就是

$$\mathbf{M_i} = \begin{bmatrix} \cos(\omega_i p) & \sin(\omega_i p) \\ -\sin(\omega_i p) & \cos(\omega_i p) \end{bmatrix}$$

在實踐中，我們不使用矩陣乘法來計算 RoPE，因為對於如此稀疏的矩陣來說，計算效率會很低。相反，我們可以直接將旋轉獨立地應用於成對的元素，利用計算中的規律模式：

$$R_{\Theta,p}^d q = \begin{pmatrix} q_1 \\ q_2 \\ q_3 \\ q_4 \\ \vdots \\ q_{d-1} \\ q_d \end{pmatrix} \otimes \begin{pmatrix} \cos p\theta_1 \\ \cos p\theta_1 \\ \cos p\theta_2 \\ \cos p\theta_2 \\ \vdots \\ \cos p\theta_{d/2} \\ \cos p\theta_{d/2} \end{pmatrix} + \begin{pmatrix} -q_2 \\ q_1 \\ -q_4 \\ q_3 \\ \vdots \\ -q_d \\ q_{d-1} \end{pmatrix} \otimes \begin{pmatrix} \sin p\theta_1 \\ \sin p\theta_1 \\ \sin p\theta_2 \\ \sin p\theta_2 \\ \vdots \\ \sin p\theta_{d/2} \\ \sin p\theta_{d/2} \end{pmatrix}$$

就這麼簡單！透過巧妙地將旋轉應用於 $\mathbf{q}$ 和 $\mathbf{k}$ 的 2D 塊，並在點積之前從加法切換到乘法，我們可以大大提升評估效能 $[^4]$。

將 RoPE 擴充套件到 $n$ 維

我們已經探討了 RoPE 的 $1D$ 案例，我希望您現在對這個一開始不太直觀的 Transformer 元件有了直觀的理解。最後，讓我們探討如何將其擴充套件到更高維度，例如影像。

一個自然的初步直覺可能是直接使用影像中的 $\begin{bmatrix} x \\ y \end{bmatrix}$ 座標對。這可能看起來很直觀，畢竟我們之前幾乎是隨意地配對了我們的元件。然而，這將是一個錯誤！

在 $1D$ 情況下，我們透過旋轉輸入向量中的值對來編碼相對位置 $m - n$。對於 $2D$ 資料，我們需要獨立編碼水平和垂直相對位置，例如 $m - n$ 和 $i - j$。RoPE 的精妙之處在於它如何處理多個維度。它不是將所有位置資訊編碼在一個旋轉中，而是將同一維度內的元件配對並旋轉它們，否則我們會將 $x$ 和 $y$ 偏移資訊混合在一起。透過獨立處理每個維度，我們保持了空間的自然結構。這可以推廣到所需的任意多個維度！

位置編碼的未來

RoPE 是位置編碼的終極形態嗎？DeepMind 的這篇 最新論文 深入分析了 RoPE 並指出了一些基本問題。簡而言之：RoPE 並非完美解決方案，模型主要關注低頻部分，但該論文表明，移除（而非旋轉）最低頻率可以提高 Gemma 2B 的效能！

我預計未來會取得一些突破，或許會從訊號處理中汲取靈感，例如小波或分層實現等思想。隨著模型越來越多地被量化以進行部署，我也希望看到編碼方案的創新，使其在低精度算術下仍能保持穩健。

總結

位置編碼在 Transformer 中一直被視為一個事後考慮的因素。我認為我們應該以不同的眼光看待它——自注意力機制有一個阿喀琉斯之踵，它被反覆修補。

我希望這篇部落格文章能向你展示，儘管起初它可能不直觀，但你也可以發現最先進的位置編碼。在後續的文章中，我很樂意探討 RoPE 的實際實現細節，以最大限度地提高效能。

這篇帖子最初發佈於此地。

參考資料

• Transformer 架構：位置編碼
• 旋轉嵌入：一場相對的革命
• Transformer 中位置編碼的工作原理？
• Attention Is All You Need
• 週而復始！是什麼讓旋轉位置編碼如此有用？
• RoFormer: 帶有旋轉位置嵌入的增強型 Transformer

[^1]：二進位制和正弦動畫是此影片中包含的動畫的再現。

[^2]：使用 $\theta = 10000$ 可提供 $2 \pi \cdot 10000$ 個唯一位置，或理論上上下文長度的上限約為 63,000。

[^3]：此文章的部分內容基於 Amirhossein Kazemnejad 撰寫的這篇精彩文章。

[^4]：關於實證證據，請參閱 EleutherAI 的這篇精彩文章。