簡介

因為RoPE的優異性能，其已成為各種大模型中位置編碼的首選，包括多模態模型；在一些多模態模型或視頻理解模型中，甚至會用到多維度RoPE。雖然RoPE已廣泛應用，之前也看了不少針對其原理解析的文章，特別是蘇神的推導帖，可能是我個人能力問題，也可能是太過于追求原理的問題，對RoPE一直沒有理解透徹，總感覺差那么一點。本文會適當弱化RoPE為什么會有效的原理推導，將重點放在RoPE應該如何實現以及具體的實現步驟和一些不同實現方式的解析。

RoPE的核心思想是通過旋轉操作將位置信息融入到詞向量。在復數域中，一個復數由實部和虛部組成，故一個復數可以看作復數域中的向量。將一個復數與復數 $e^{i\theta }=\cos \theta +i\sin \theta$相乘時，相當于在復平面上繞原點旋轉了 $\theta$度角。假設有一個復數 $v=x+iy$，按上述說明，將 $v$與 $e^{i\theta }$相乘就是將 $v$繞原點旋轉了 $\theta$度角，可以通過公式推導驗證：
$$\begin{array}{rl}{v}^{\prime }& =v?e^{i\theta }\\ & =(x+iy)?(\cos \theta +i\sin \theta )\\ & =\cos \theta x+i\sin \theta x+i\cos \theta y?\sin \theta y\\ & =(\cos \theta x?\sin \theta y)+i(\sin \theta x+\cos \theta y)\end{array}\text{\hspace{1em}(1)?}$$

通過公式(1)有 $x^{\prime }=\cos \theta x?\sin \theta y,y^{\prime }=\sin \theta x+\cos \theta y$。從二維向量的角度看，變換前后的復數 $v,v^{\prime }$其實均可以看作列向量 $\left[\begin{array}{c}x\\ y\end{array}\right]，\left[\begin{array}{c}{x}^{\prime }\\ y^{\prime }\end{array}\right]$，公式(1)表示 $v^{\prime }$是由 $v$旋轉而來，且恰好對應了二維向量空間的旋轉矩陣，如下所示：
$$\begin{array}{c}\left[\begin{array}{c}{x}^{\prime }\\ y^{\prime }\end{array}\right]=\left[\begin{array}{cc}\cos \theta & ?\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]\left[\begin{array}{c}x\\ y\end{array}\right]\end{array}\text{\hspace{1em}(2)}$$

如公式(2)所示，在二維向量空間中，將一個向量繞原點旋轉 $\theta$度就是乘以一個旋轉矩陣，即 $R(\theta )=\left[\begin{array}{cc}\cos \theta & ?\sin \theta \\ \sin \theta & \cos \theta \end{array}\right]$。通過上述推導可知，二維向量空間中的旋轉和復數乘法在進行旋轉時是等價的。

RoPE是將隱向量序列中不同位置的token的特征向量與不同的旋轉角對應的旋轉矩陣相乘，即將不同位置token的特征向量旋轉不同的角度。注意力計算的主要部分是token向量之間的點積運算，可以通過推導展示RoPE是如何在自注意力計算中引入相對位置信息的。

假設目前token的維度為2，有序列中m處的查詢向量 $q_m=\left[\begin{array}{c}{q}_{m,1}\\ q_{m,2}\end{array}\right]$和n處的鍵向量 $k_n=\left[\begin{array}{c}{k}_{n,1}\\ k_{n,2}\end{array}\right]$，設旋轉基礎角度為 $\theta$，序列中不同位置token的旋轉角由位置索引與基礎角度相乘得到，即上述兩個向量的旋轉角度為別是 $m\theta ,n\theta$，RoPE就是給不同位置的token向量與對應位置的旋轉角的旋轉矩陣相乘，基于前文的推理，有，
$$\begin{array}{c}{q}_m^{\prime }=\left[\begin{array}{cc}\cos (m\theta )& ?\sin (m\theta )\\ \sin (m\theta )& \cos (m\theta )\end{array}\right]\left[\begin{array}{c}{q}_{m,1}\\ q_{m,2}\end{array}\right]=R(m\theta )q_m\end{array}\text{\hspace{1em}(3)}$$
$$\begin{array}{c}{k}_n^{\prime }=\left[\begin{array}{cc}\cos (n\theta )& ?\sin (n\theta )\\ \sin (n\theta )& \cos (n\theta )\end{array}\right]\left[\begin{array}{c}{k}_{n,1}\\ k_{n,2}\end{array}\right]=R(n\theta )k_m\end{array}\text{\hspace{1em}(4)}$$

旋轉后的兩個向量進行以下點積計算：
$$\begin{array}{rl}{q}_m^{{}^{\prime }T}{k}_n^{\prime }& =(R(m\theta )q_m{)}^T?R(n\theta )k_n\\ & =q_m^T?R(m\theta {)}^T?R(n\theta )?k_n\\ & =q_m^T?\left[\begin{array}{cc}\cos (m\theta )& \sin (m\theta )\\ ?\sin (m\theta )& \cos (m\theta )\end{array}\right]\left[\begin{array}{cc}\cos (n\theta )& ?\sin (n\theta )\\ \sin (n\theta )& \cos (n\theta )\end{array}\right]?k_n\\ & =q_m^T?\left[\begin{array}{cc}\cos (m\theta )\cos (n\theta )+\sin (m\theta )\sin (n\theta )& ?\cos (m\theta )\sin (n\theta )+\sin (m\theta )\cos (n\theta )\\ ?\sin (m\theta )\cos (n\theta )+\cos (m\theta )\sin (n\theta )& \sin (m\theta )\sin (n\theta )+\cos (m\theta )\cos (n\theta )\end{array}\right]?k_n\\ & =q_m^T?\left[\begin{array}{cc}\cos ((n?m)\theta )& ?\sin ((n?m)\theta )\\ \sin ((n?m)\theta )& \cos ((n?m)\theta )\end{array}\right]?k_n\\ & =q_m^T?R((n?m)\theta )?k_n\end{array}\text{\hspace{1em}(5)?}$$

上述公式(5)推導過程要使用高中數學知識–三角函數和差定理。可以看到，兩個旋轉不同角度后向量的點積與它們之間旋轉的角度之差有關系，是一種相對位置的體現，即體現了RoPE給序列中不同位置token引入了相對位置信息。

注意：上述旋轉均是逆時針旋轉！

拆解

上述只是對RoPE的原理進行了簡短解釋，盡可能直白地說明RoPE原理，但上述討論局限在二維空間內，而LLMs模型中token向量的維度都很大，接下來對常規的RoPE實現進行拆解，將整個實現過程解剖出來。

1. 為了沿用二維旋轉矩陣的性質，RoPE會將一個token的高維向量看作多個復數，假設向量特征維度為 $d$(基本所有LLMs中的特征維度數是偶數)，那么會將一個$d$維的token特征向量視為 $d/2$個復數

• 假設一個token向量為 $[1,2,3,4]$，那么 $[1,2]$是第一個復數，$[3,4]$是第二個復數
• 同一個token中的不同復數對的旋轉角度不同
• 基礎旋轉角 ${\theta }_i=10000^{?2i/d}=\frac{1}{10000^{\frac{i}{\frac{d}{2}}}}$；注意，同一token向量中不同位置復數的旋轉角是由 $i$來決定的。上文已經說到，一個token向量會分成 $d/2$個復數， $i\in [0,d/2)$，所以 ${\theta }_i$就對應第 $i$對復數的旋轉角
  • 10000是一個超參數

2. 同一序列中不同位置的token向量的旋轉角也不同，就是基礎旋轉角與token對應的序列索引乘積，即一個token在序列中的索引為m，那么其對應的最終的旋轉角為 ${\theta }_i?m$。

• 實現時一般會先將 ${\theta }_i$計算出來，再基于序列長度獲取不同位置token索引的一維向量，然后兩者相乘，得到不同位置token中不同復數維度的準確旋轉角度

3. 將輸入沿著特征維度，按順序每兩個為一組作為一對復數進行拆分，向量維度從[batch_size, seq_len, dim]–>[batch_size, seq_len, dim//2, 2]

• 最后一個維度中，第一組向量為所有復數的實部，第二組向量為所有復數的虛部

4. 按照二維矩陣中的計算方式，進行旋轉矩陣乘法，然后將最終的實部和虛部重新合并為原來的維度[batch_size, seq_len, dim]，即完成了RoPE的應用

上述對RoPE的具體操作過程進行了敘述，以下是一種較原始的樸素實現：

import torchdef get_rotary_matrix(seq_len, dim, base=10000):"""生成RoPE的旋轉矩陣"""# 生成不同頻率的正弦和余弦值theta = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))  # shape為[dim//2]# 生成位置索引position = torch.arange(seq_len).float()  # shape為[seq_len]# 計算每個位置和維度對應的角度theta = torch.outer(position, theta)  # 計算外積，其中第(i, j)個元素是position[i] * theta[j]；shape為[seq_len, dim//2]# 計算正弦和余弦值cos = torch.cos(theta)  # shape為[seq_len, dim//2]sin = torch.sin(theta)  # shape為[seq_len, dim//2]return cos, sindef apply_rotary_embedding(x, cos, sin):"""應用旋轉位置編碼"""# 假設x的形狀為[batch_size, seq_len, dim]# 將向量視為復數，每兩個維度一組x_reshape = x.view(*x.shape[:-1], -1, 2)  # shape為[batch_size, seq_len, dim//2, 2]，即沿著特征維度拆分# 構建正弦和余弦矩陣，使其與x_reshape形狀匹配cos_expanded = cos.view(1, cos.shape[0], cos.shape[1], 1)  # shape為[1, seq_len, dim//2, 1]sin_expanded = sin.view(1, sin.shape[0], sin.shape[1], 1)  # shape為[1, seq_len, dim//2, 1]# 旋轉操作（復數乘法）# [x_real, x_imag] * (cos + i*sin) = [x_real*cos - x_imag*sin, x_real*sin + x_imag*cos]x_out_1 = x_reshape[:, :, :, 0:1] * cos_expanded - x_reshape[:, :, :, 1:2] * sin_expandedx_out_2 = x_reshape[:, :, :, 0:1] * sin_expanded + x_reshape[:, :, :, 1:2] * cos_expanded# 合并結果x_out = torch.cat([x_out_1, x_out_2], dim=-1)  # shape為[batch_size, seq_len, dim//2, 2]return x_out.view(*x.shape)# 示例用法
def apply_rope(x):"""對輸入向量應用RoPE位置編碼"""batch_size, seq_len, dim = x.shapecos, sin = get_rotary_matrix(seq_len, dim)return apply_rotary_embedding(x, cos, sin)# 測試代碼
if __name__ == "__main__":# 創建一個隨機輸入張量batch_size, seq_len, dim = 2, 10, 512x = torch.randn(batch_size, seq_len, dim)# 應用RoPEx_with_rope = apply_rope(x)print(f"輸入形狀: {x.shape}")print(f"輸出形狀: {x_with_rope.shape}")

一些tricks

• 在對輸入的嵌入維度上應用RoPE時，當特征維度很長，或者希望節約資源資源時，可以不對所有維度全部應用，而是可以設置一個比例，即只在前 $d_{rope}$維度上應用RoPE，后面保持原始數值不變
• 在計算 ${\theta }_i$時原始論文使用的base為10000，但可以對其進行縮放，實現長度外推；如增加base，可以處理訓練過程中未見過的長序列