設計位置編碼

Gall 定律
一個有效的複雜系統通常是從一個有效的簡單系統演化而來的
—— John Gall

本文將帶你一步步探究 Transformer 模型中先進的位置編碼技術。我們將透過迭代改進編碼位置的方法，最終得出 旋轉位置編碼 (Rotary Postional Encoding, RoPE)，這也是最新發布的 LLama 3.2 和大多數現代 transformer 模型所採用的方法。本文旨在儘量減少所需的數學知識，但理解一些基本的線性代數、三角學和自注意力機制是有幫助的。

問題陳述

你可以透過詞語與其他詞語的關係來理解一個詞語的意義
—— John Rupert Firth

在所有問題中，首先要做的是理解 我們到底在解決什麼問題。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\) 至關重要。這將使模型能夠處理影像、腦部掃描等二維和四維資料。

現在我們知道了理想的屬性 (以下稱為 \(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}} \]

為了找到這個變換矩陣，我們可以將其表示為一個通用的 2×2 矩陣，其係數為未知數 \(u_1\)、\(v_1\)、\(u_2\) 和 \(v_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 在 Attention is all you need 一文中設計的編碼方案早在 2017 年就已經透過旋轉來編碼相對位置了！但從正弦位置編碼到 RoPE (旋轉位置編碼) 卻花了整整 4 年，儘管旋轉的概念早已被提出……

絕對位置編碼 vs 相對位置編碼

認識到旋轉的重要性後，讓我們回到引導示例，並嘗試為下一次迭代探索一些靈感。

\[\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 chased another dog} \end{align*} \]

上圖展示了詞元的絕對位置以及從 \(\text{chased}\) 到其他詞元的相對位置。透過正弦位置編碼，我們生成了一個表示絕對位置的獨立向量，並利用一些三角技巧將相對位置編碼了進來。

當我們試圖理解這些句子時， this 單詞是這篇部落格中第 2157 個詞的重要性有多大？或者我們更關心它與周圍單詞的關係？單詞的絕對位置對於其意義來說很少重要——真正重要的是單詞之間的相互關係。

在上下文中理解位置編碼

從現在開始，我們需要將位置編碼放在 自注意力機制的上下文中 中分析。再強調一次，自注意力機制讓模型能夠評估輸入序列中不同元素的重要性，並動態調整它們對輸出的影響。

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

在之前的所有迭代中，我們生成了一個獨立的位置資訊向量，並在 \(Q\)、\(K\) 和 \(V\) 投影之前將其 加到 詞元嵌入上。這樣做實際上是將位置資訊與語義資訊混合在一起，對嵌入向量的範數進行了修改。為了避免汙染語義資訊，我們應該嘗試使用乘法進行編碼。

用詞典的類比來說，當我們查詢一個單詞 (query) 在詞典 (keys) 中的含義時，鄰近的單詞應該比遠處的單詞有更大的影響。這種影響是透過 \(QK^T\) 的點積來決定的，因此位置編碼應該集中在這裡！

\[\vec{a} \cdot \vec{b} = |\vec{a}| |\vec{b}| \cos \theta \]

點積的幾何解釋為我們提供了一個極好的靈感: 我們可以透過改變兩個向量之間的角度來調整點積的結果。此外，透過旋轉向量，我們不會改變向量的範數，從而不會影響詞元的語義資訊。

現在，我們知道應該把注意力放在哪裡，並從另一個視角看到了為什麼旋轉是一種合理的“通道”來編碼位置資訊——讓我們把一切整合起來吧！

旋轉位置編碼 (RoPE)

旋轉位置編碼 (Rotary Positional Encoding，簡稱 RoPE ) 是在 RoFormer 論文 中定義的 (Jianlin Su在他的部落格 這裡 和 這裡中獨立設計了這一方法)。

如果直接跳到最終結果，可能會覺得它像是某種“魔法”。但透過將正弦位置編碼放在自注意力 (尤其是點積) 的背景下思考，我們可以看出它的核心邏輯。

與正弦位置編碼類似，我們將向量 \(\mathbf{q}\) 或 \(\mathbf{k}\) (而非投影前的 \(\mathbf{x}\)) 分解為二維的對/塊。不同於透過將基於逐漸減小頻率的正弦函式生成的向量相加來直接編碼 絕對 位置，RoPE 透過 對每一對分量應用旋轉矩陣 來編碼 相對 位置。

令 \(\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} \]

與正弦位置編碼類似，\(\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}\) 的二維塊上巧妙地應用旋轉操作，並從加法切換為乘法，我們可以顯著提升評估效能 \([^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 中，位置編碼常被視為事後的補丁。我認為我們應當改變這種看法——自注意力機制有一個“致命弱點”，它需要被反覆修補。

希望這篇部落格能讓你明白，即便這種方法起初看起來難以直觀理解，你也可以發現最新的狀態 -of-the-art 位置編碼。在後續文章中，我將探討如何實踐 RoPE 的具體實現細節以最佳化效能。

這篇文章最初發布在 這裡

參考文獻

• Transformer Architecture: The Positional Encoding
• Rotary Embeddings: A Relative Revolution
• How positional encoding works in transformers?
• Attention Is All You Need
• Round and round we go! What makes Rotary Positional Encodings useful?
• RoFormer: Enhanced Transformer with Rotary Position Embedding
  \([^1]\): 正弦和二進位制動畫來源於 影片。
  \([^2]\): 使用 \(\theta = 10000\) 可生成 \(2 \pi \cdot 10000\) 個唯一位置，理論上下文長度上限約為 63,000。
  \([^3]\): 本文部分內容參考自 這篇精彩的文章 (作者: Amirhossein Kazemnejad。
  \([^4]\): 相關實驗資料請參見 EleutherAI 的這篇文章。)

---

原文連結: https://hf.co/blog/designing-positional-encoding

原文作者: Christopher Fleetwood

譯者: chenin-wang