大語言模型中常用的旋轉位置編碼RoPE詳解:為什么它比絕對或相對位置編碼更好?
自 2017 年發表“ Attention Is All You Need ”論文以來,Transformer 架構一直是自然語言處理 (NLP) 領域的基石。它的設計多年來基本沒有變化,隨著旋轉位置編碼 (RoPE) 的引入,2022年標志著該領域的重大發展。
旋轉位置嵌入是最先進的 NLP 位置嵌入技術。大多數流行的大型語言模型(如 Llama、Llama2、PaLM 和 CodeGen)已經在使用它。在本文中,我們將深入探討什么是旋轉位置編碼,以及它們如何巧妙地融合絕對位置嵌入和相對位置嵌入的優點。
位置編碼的需求
為了理解 RoPE 的重要性,我們首先回顧一下為什么位置編碼至關重要。Transformer 模型根據其固有的設計,不會考慮輸入標記的順序。
例如,像“the dog chases the pig ”和“the pig chases the dogs”這樣的短語雖然含義不同,但由于它們被視為一組無序的標記,因此被視為無法區分。為了維護序列信息及其含義,需要一個表示來將位置信息集成到模型中。
絕對位置編碼
在句子的上下文中,假設我們有一個代表一個單詞的嵌入。為了對其位置進行編碼,需要使用另一個具有相同維度的向量,其中每個向量唯一地代表句子中的一個位置。例如,為句子中的第二個單詞指定特定向量。所以每個句子位置都有其獨特的向量。然后通過將詞嵌入與其相應的位置嵌入求和來形成 Transformer 層的輸入。
有兩種主要方法來生成這些嵌入:
- 從數據中學習:在這里,位置向量是在訓練過程中學習的,就像其他模型參數一樣。我們為每個位置(例如從 1 到 512)學習一個唯一的向量。這引入了一個限制——最大序列長度受到限制。如果模型僅學習到位置 512,則它無法表示比該位置更長的序列。
- 正弦函數:此方法涉及使用正弦函數為每個位置構建唯一的嵌入。盡管這種構造的細節很復雜,但它本質上為序列中的每個位置提供了獨特的位置嵌入。實證研究表明,從數據中學習和使用正弦函數可以在現實世界模型中提供相當的性能。
絕對位置編碼的局限性
盡管使用廣泛但絕對位置嵌入也并非沒有缺點:
- 有限序列長度:如上所述,如果模型學習到某個點的位置向量,它本質上不能表示超出該限制的位置。
- 位置嵌入的獨立性:每個位置嵌入都是獨立于其他位置嵌入的。這意味著在模型看來,位置 1 和 2 之間的差異與位置 2 和 500 之間的差異相同。但是其實位置 1 和 2 應該比位置 500 相關性更密切,位置 500 距離明顯更遠。這種相對定位的缺乏可能會阻礙模型理解語言結構的細微差別的能力。
相對位置編碼
相對位置位置不是關注標記在句子中的絕對位置,而是關注標記對之間的距離。該方法不會直接向詞向量添加位置向量。而是改變了注意力機制以納入相對位置信息。
最經典得案例就是T5(Text-to-Text Transfer Transformer)是一種利用相對位置嵌入的著名模型。T5 引入了一種處理位置信息的微妙方式:
- 位置偏移的偏差: T5 使用偏差(浮點數)來表示每個可能的位置偏移。例如,偏差 B1 可能表示任意兩個相距一個位置的標記之間的相對距離,無論它們在句子中的絕對位置如何。
- 自注意力層中的集成:該相對位置偏差矩陣被添加到自注意力層中的查詢矩陣和關鍵矩陣的乘積中。這確保了相同相對距離的標記始終由相同的偏差表示,無論它們在序列中的位置如何。
- 可擴展性:該方法的一個顯著優點是其可擴展性。它可以擴展到任意長的序列,這比絕對位置嵌入有明顯的優勢。
相對位置編碼的局限性
盡管它們在理論上很有吸引力,但相對位置編碼得問題很嚴重
- 計算效率低下:必須創建成對的位置編碼矩陣,然后執行大量張量操作以獲得每個時間步的相對位置編碼。特別是對于較長的序列。這主要是由于自注意力層中的額外計算步驟,其中位置矩陣被添加到查詢鍵矩陣中。
- 鍵值緩存使用的復雜性:由于每個附加令牌都會改變每個其他令牌的嵌入,這使得 Transformer 中鍵值緩存的有效使用變得復雜。使用 KV 緩存的一項要求是已經生成的單詞的位置編碼, 在生成新單詞時不改變(絕對位置編碼提供)因此相對位置編碼不適合推理,因為每個標記的嵌入會隨著每個新時間步的變化而變化。
由于這些工程復雜性,位置編碼未得到廣泛采用,特別是在較大的語言模型中。
旋轉位置編碼 (RoPE)?
RoPE 代表了一種編碼位置信息的新方法。傳統方法中無論是絕對方法還是相對方法,都有其局限性。絕對位置編碼為每個位置分配一個唯一的向量,雖然簡單但不能很好地擴展并且無法有效捕獲相對位置;相對位置編碼關注標記之間的距離,增強模型對標記關系的理解,但使模型架構復雜化。
RoPE巧妙地結合了兩者的優點。允許模型理解標記的絕對位置及其相對距離的方式對位置信息進行編碼。這是通過旋轉機制實現的,其中序列中的每個位置都由嵌入空間中的旋轉表示。RoPE 的優雅之處在于其簡單性和高效性,這使得模型能夠更好地掌握語言語法和語義的細微差別。
旋轉矩陣源自我們在高中學到的正弦和余弦的三角性質,使用二維矩陣應該足以獲得旋轉矩陣的理論,如下所示!
我們看到旋轉矩陣保留了原始向量的大小(或長度),如上圖中的“r”所示,唯一改變的是與x軸的角度。
RoPE 引入了一個新穎的概念。它不是添加位置向量,而是對詞向量應用旋轉。旋轉角度 (θ) 與單詞在句子中的位置成正比。第一個位置的向量旋轉 θ,第二個位置的向量旋轉 2θ,依此類推。這種方法有幾個好處:
- 向量的穩定性:在句子末尾添加標記不會影響開頭單詞的向量,有利于高效緩存。
- 相對位置的保留:如果兩個單詞在不同的上下文中保持相同的相對距離,則它們的向量將旋轉相同的量。這確保了角度以及這些向量之間的點積保持恒定
RoPE 的矩陣公式
RoPE的技術實現涉及到旋轉矩陣。在 2D 情況下,論文中的方程包含一個旋轉矩陣,該旋轉矩陣將向量旋轉 Mθ 角度,其中 M 是句子中的絕對位置。這種旋轉應用于 Transformer 自注意力機制中的查詢向量和鍵向量。
對于更高維度,向量被分成 2D 塊,并且每對獨立旋轉。這可以被想象成一個在空間中旋轉的 n 維。聽著這個方法好好像實現是復雜,其實不然,這在 PyTorch 等庫中只需要大約十行代碼就可以高效的實現。
import torch
import torch.nn as nn
class RotaryPositionalEmbedding(nn.Module):
def __init__(self, d_model, max_seq_len):
super(RotaryPositionalEmbedding, self).__init__()
# Create a rotation matrix.
self.rotation_matrix = torch.zeros(d_model, d_model, device=torch.device("cuda"))
for i in range(d_model):
for j in range(d_model):
self.rotation_matrix[i, j] = torch.cos(i * j * 0.01)
# Create a positional embedding matrix.
self.positional_embedding = torch.zeros(max_seq_len, d_model, device=torch.device("cuda"))
for i in range(max_seq_len):
for j in range(d_model):
self.positional_embedding[i, j] = torch.cos(i * j * 0.01)
def forward(self, x):
"""
Args:
x: A tensor of shape (batch_size, seq_len, d_model).
Returns:
A tensor of shape (batch_size, seq_len, d_model).
"""
# Add the positional embedding to the input tensor.
x += self.positional_embedding
# Apply the rotation matrix to the input tensor.
x = torch.matmul(x, self.rotation_matrix)
return x
為了旋轉是通過簡單的向量運算而不是矩陣乘法來執行。距離較近的單詞更有可能具有較高的點積,而距離較遠的單詞則具有較低的點積,這反映了它們在給定上下文中的相對相關性。
使用 RoPE 對 RoBERTa 和 Performer 等模型進行的實驗表明,與正弦嵌入相比,它的訓練時間更快。并且該方法在各種架構和訓練設置中都很穩健。
最主要的是RoPE是可以外推的,也就是說可以直接處理任意長的問題。在最早的llamacpp項目中就有人通過線性插值RoPE擴張,在推理的時候直接通過線性插值將LLAMA的context由2k拓展到4k,并且性能沒有下降,所以這也可以證明RoPE的有效性。
代碼如下:
import transformers
old_init = transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__
def ntk_scaled_init(self, dim, max_position_embeddings=2048, base=10000, device=None):
#The method is just these three lines
max_position_embeddings = 16384
a = 8 #Alpha value
base = base * a ** (dim / (dim-2)) #Base change formula
old_init(self, dim, max_position_embeddings, base, device)
transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__ = ntk_scaled_init
總結
旋轉位置嵌入代表了 Transformer 架構的范式轉變,提供了一種更穩健、直觀和可擴展的位置信息編碼方式。
RoPE不僅解決了LLM context過長之后引起的上下文無法關聯問題,并且還提高了訓練和推理的速度。這一進步不僅增強了當前的語言模型,還為 NLP 的未來創新奠定了基礎。隨著我們不斷解開語言和人工智能的復雜性,像 RoPE 這樣的方法將有助于構建更先進、更準確、更類人的語言處理系統。