文章目錄
- 前言
- 一、絕對位置編碼
- 二、相對位置編碼
- 三、旋轉位置編碼
前言
由于attetnion運算的特性,Transformer本身不感知順序,位置編碼是彌補這一缺陷的關鍵。
一、絕對位置編碼
絕對位置編碼的方式是通過將每個位置映射到一個高維空間中,該編碼采用了正弦和余弦函數的組合
周期性:正弦和余弦函數具有周期性,這使得編碼能夠容易地表示固定范圍的位置信息。這個特點允許模型有能力處理序列的循環性質,例如在語言中,某些詞可能在不同上下文中重復出現,但它們的相對位置仍然是重要的。
不同頻率:對于不同的維度 ( i ),使用不同的頻率來編碼位置。低維度的編碼(例如低 ( i ) 值)對應較大的周期,能夠捕捉較長的依賴關系;而高維度編碼對應較小的周期,能夠捕捉短距離的依賴關系。因此,模型可以通過不同的維度考慮不同尺度的位置信息。
任意長度的序列:這種方法能夠處理任意長度的輸入序列,因為正弦和余弦函數是定義于所有實數的,可以為任意的位置提供唯一的編碼。
盡管絕對位置編碼看似只提供了位置的信息,但模型在訓練過程中會學會捕捉相對位置信息。原因如下:1、位置之間的差異:通過將位置編碼加到輸入向量中,模型可以學習到位置之間的相對關系。例如,給定位置 ( i ) 和位置 ( j ),它們的編碼向量可以表示為:
PE(i)?PE(j)
這一差值可以在一定程度上反映這兩個位置之間的相對距離。
2、向量性質:在高維空間中,向量之間的方向和距離能夠也反映相對位置。例如,當兩個詞在序列中相隔一定距離時,它們的相應位置編碼的差異會隱含這種相對關系。
import torch
import torch.nn as nn
import mathclass PositionalEncoding(nn.Module):"""實現經典的基于正弦和余弦函數的絕對位置編碼。"""def __init__(self, d_model, max_len=5000, dropout=0.1):"""Args:d_model (int): 模型的維度(或詞嵌入的維度)。max_len (int): 預先計算編碼的最大序列長度。dropout (float): Dropout 的比例。"""super(PositionalEncoding, self).__init__()self.dropout = nn.Dropout(p=dropout)# 創建一個足夠長的位置編碼矩陣# 形狀為 (max_len, d_model)pe = torch.zeros(max_len, d_model)# 創建一個位置索引張量# position.shape: (max_len, 1)position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)# 計算除法項 1 / (10000^(2i/d_model))# div_term.shape: (d_model/2)# 這里的 log 是為了數值穩定性,等價于 1 / (10000^(2i/d_model))div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))# 使用廣播機制計算正弦和余弦值# 偶數索引使用 sinpe[:, 0::2] = torch.sin(position * div_term)# 奇數索引使用 cospe[:, 1::2] = torch.cos(position * div_term)# 增加一個 batch 維度,使其能夠與輸入張量 (batch_size, seq_len, d_model) 直接相加# pe.shape from (max_len, d_model) to (1, max_len, d_model)pe = pe.unsqueeze(0)# 將 pe 注冊為 buffer。# buffer 是模型的狀態的一部分,但不是參數 (parameters),因此不會被優化器更新。# 它會隨著模型一起移動(例如 .to(device)),并且會保存在 state_dict 中。self.register_buffer('pe', pe)def forward(self, x):"""將位置編碼添加到輸入張量中。Args:x (torch.Tensor): 輸入張量,形狀為 (batch_size, seq_len, d_model)。Returns:torch.Tensor: 添加了位置編碼的輸出張量,形狀與輸入相同。"""# x.shape: (batch_size, seq_len, d_model)# self.pe.shape: (1, max_len, d_model)# 截取所需長度的位置編碼并與輸入相加# self.pe[:, :x.size(1), :] 的形狀變為 (1, seq_len, d_model),可以與 x 廣播相加x = x + self.pe[:, :x.size(1), :]return self.dropout(x)
二、相對位置編碼
相對位置編碼(Relative Positional Encoding)是一種改善模型捕捉序列中詞語相對位置關系的技術。與絕對位置編碼(Absolute Positional Encoding)不同,相對位置編碼側重于表示詞與詞之間的相對距離,從而增強模型的學習能力。
但相對位置編碼的實現和絕對位置編碼的實現差距還是挺大的,相對位置編碼以矩陣的形式呈現,計算每個詞i 和 詞j 的相對位置,如果序列長度過大,會導致整個位置編碼十分龐大,計算成本巨大。
標準的相對位置編碼方法(如論文 “Self-Attention with Relative Position Representations” 中提出的)并不會直接預計算所有位置對的編碼。相反,它創建了一個可學習的相對位置嵌入——relative position embedding查找表。
實現:
import torch
import torch.nn as nn
import mathclass RelativeAttentionScore(nn.Module):def __init__(self, d_model, max_len=50):super(RelativeAttentionScore, self).__init__()# 假設 d_k = d_model,在多頭注意力中 d_k = d_model / num_headsself.d_k = d_model# 定義一個最大相對距離。超過這個距離的位置被視為相同距離。# 這有助于模型泛化到比訓練時更長的序列。self.max_relative_position = max_len# 創建一個可學習的嵌入查找表,用于相對位置。# 大小為 2 * max_len - 1,覆蓋從 -(max_len-1) 到 (max_len-1) 的所有相對位置。# +1 是因為我們需要一個位置來存儲被裁剪的距離self.relative_embeddings = nn.Embedding(2 * self.max_relative_position + 1, self.d_k)def forward(self, queries, keys):"""Args:queries (torch.Tensor): 查詢張量,形狀為 (batch_size, seq_len, d_model)keys (torch.Tensor): 鍵張量,形狀為 (batch_size, seq_len, d_model)Returns:torch.Tensor: 加上了相對位置偏置的注意力得分,形狀為 (batch_size, seq_len, seq_len)"""batch_size, seq_len, _ = queries.shape# 1. 計算基于內容的注意力得分content_score = torch.matmul(queries, keys.transpose(-2, -1))# 2. 計算基于位置的注意力得分# a. 生成相對位置矩陣# range_vec.shape: (seq_len)range_vec = torch.arange(seq_len, device=queries.device)# relative_matrix.shape: (seq_len, seq_len)# 每一行表示當前位置與所有其他位置的相對距離relative_matrix = range_vec[None, :] - range_vec[:, None]print(relative_matrix)# b. 裁剪相對距離并移動到非負索引# 將距離裁剪到 [-max_relative_position, max_relative_position] 范圍內# print(self.max_relative_position)clipped_relative_matrix = torch.clamp(relative_matrix, -self.max_relative_position, self.max_relative_position)# 將索引平移到 [0, 2 * max_relative_position] 范圍,以用于Embedding查找positive_indices = clipped_relative_matrix + self.max_relative_position# c. 查找相對位置嵌入# pos_embeddings.shape: (seq_len, seq_len, d_k)pos_embeddings = self.relative_embeddings(positive_indices)# d. 計算位置得分# 我們需要計算 query 和 相對位置嵌入 的點積。# (b, i, d) * (i, j, d) -> (b, i, j)# b=batch_size, i=query_pos, j=key_pos, d=d_k# torch.einsum 是實現這種復雜矩陣乘法的優雅方式position_score = torch.einsum('bid,ijd->bij', queries, pos_embeddings)# 3. 將內容得分和位置得分相加attention_scores = content_score + position_score# (可選)應用縮放scaled_attention_scores = attention_scores / math.sqrt(self.d_k)return scaled_attention_scores
--- 相對位置矩陣示例 (seq_len=5) ---
原始相對位置矩陣:tensor([[ 0, 1, 2, 3, 4],[-1, 0, 1, 2, 3],[-2, -1, 0, 1, 2],[-3, -2, -1, 0, 1],[-4, -3, -2, -1, 0]])裁剪后的矩陣 (max_len=2):tensor([[ 0, 1, 2, 2, 2],[-1, 0, 1, 2, 2],[-2, -1, 0, 1, 2],[-2, -2, -1, 0, 1],[-2, -2, -2, -1, 0]])用于嵌入查找的非負索引:tensor([[2, 3, 4, 4, 4],[1, 2, 3, 4, 4],[0, 1, 2, 3, 4],[0, 0, 1, 2, 3],[0, 0, 0, 1, 2]])
由于裁剪機制的存在,任何超出 [-max_relative_position, max_relative_position] 范圍的相對距離都會被“壓縮”到邊界值上。這個機制其實很好的幫助了模型:
1)假設用最大長度為512的句子訓練了一個模型,max_relative_position 也設為512。現在,想用這個模型去處理一個長度為1024的句子。如果沒有裁剪:模型在處理這個長句子時,會遇到它從未見過的相對距離,比如 +600 或 -800。由于嵌入表中沒有這些距離的位置,或者這些位置的嵌入向量從未被訓練過,模型的表現會變得非常不穩定,甚至完全崩潰。
有了裁剪:模型在訓練時已經學會了一個對于“非常遠”的距離(比如+512或-512)的通用表示。當它在推理時遇到一個新的、更遠的距離(如 +600)時,它會將其裁剪到 +512,然后使用那個它已經熟知的“非常遠”的嵌入。這使得模型能夠平滑地泛化到比訓練時更長的序列,而不會因為遇到未知的距離而失敗。
2)在自然語言中,詞語之間的關系強度通常與距離有關,但這種關系不是無限延伸的。模型通過裁剪,學會了一個“局部注意力窗口”(在 [-50, 50] 范圍內),并對這個窗口內的位置進行精細建模。對于窗口外的所有位置,它只學習一個統一的“遠距離”表示。這是一種非常合理的歸納偏置(inductive bias)。
三、旋轉位置編碼
好的,旋轉位置編碼(Rotary Positional Encoding, RoPE)是目前大型語言模型(如 LLaMA, PaLM)中非常流行且效果出色的一種位置編碼方案。它由蘇建林在論文《RoFormer: Enhanced Transformer with Rotary Position Embedding》中提出。
與傳統的加性位置編碼(Absolute PE)或在注意力分數上增加偏置(Relative PE from Shaw et al.)不同,RoPE 的思想極為巧妙:它通過旋轉查詢(Query)和鍵(Key)向量來注入位置信息。
其核心思想為:
絕對位置決定初始角度:一個詞在序列中的絕對位置 m 決定了它的查詢向量 q 和鍵向量 k 需要旋轉的角度 mθ。相對位置體現在角度差:當計算兩個詞(位置m的查詢q和位置n的鍵k)的注意力時,它們旋轉后的點積結果,神奇地只與它們的內容 (q, k) 和它們的相對位置 m-n 有關,而與它們的絕對位置 m 和 n 無關
數學原理為:
RoPE 的魔法在于復數運算。將 d 維的向量兩兩配對,看作 d/2 個復數。對一個位于位置 m 的向量 x(它可以是 q 或 k),其旋轉操作可以表示為:
x’_m = x_m * e^(i * m * θ) ; x_m 是原始向量(被看作復數)。 i 是虛數單位。 m 是絕對位置。θ 是一個預設的、與維度相關的常數(類似于傳統PE中的頻率)
當計算旋轉后的查詢 q’_m 和鍵 k’_n 的點積(在復數域中是取共軛后相乘再取實部)時:
Re[ (q_m * e^(imθ)) * (k_n * e(i*nθ)) ]
= Re[ q_m * k_n^* * e^(imθ) * e^(-inθ) ]
= Re[ q_m * k_n^* * e^(i*(m-n)θ) ]
最終結果僅依賴于m-n
import torch
import torch.nn as nn
import mathclass RotaryPositionalEncoding(nn.Module):def __init__(self, d_model, max_len=512):super().__init__()self.d_model = d_model# 計算旋轉角度 theta# theta_i = 10000^(-2(i-1)/d) for i in [1, 2, ..., d/2]inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))# 預先計算所有可能位置 m 的 m*thetat = torch.arange(max_len, dtype=inv_freq.dtype)freqs = torch.einsum('i,j->ij', t, inv_freq)# freqs 包含了所有位置的 m*theta, 形狀是 (max_len, d_model/2)# 將其擴展為 (max_len, d_model) 以便應用# emb.shape: (max_len, d_model)emb = torch.cat((freqs, freqs), dim=-1)# 注冊為 buffer,這樣它就不會被視為模型參數,但會隨模型移動 (e.g., .to(device))# self.cos_cached.shape: (1, 1, max_len, d_model)# self.sin_cached.shape: (1, 1, max_len, d_model)self.register_buffer("cos_cached", emb.cos()[None, None, :, :])self.register_buffer("sin_cached", emb.sin()[None, None, :, :])def forward(self, x):# x.shape: (batch_size, num_heads, seq_len, head_dim)# head_dim == self.d_modelseq_len = x.shape[-2]# 獲取預計算的 cos 和 sin 值cos = self.cos_cached[:, :, :seq_len, ...]sin = self.sin_cached[:, :, :seq_len, ...]# 執行旋轉# 1. 將 x 分為兩半# x1.shape, x2.shape: (batch_size, num_heads, seq_len, head_dim/2)x1 = x[..., 0::2] # 偶數維度x2 = x[..., 1::2] # 奇數維度# 2. 應用旋轉公式# x_rotated = (x1 + i*x2) * (cos + i*sin) = (x1*cos - x2*sin) + i*(x1*sin + x2*cos)rotated_x1 = x1 * cos[..., 0::2] - x2 * sin[..., 0::2]rotated_x2 = x1 * sin[..., 1::2] + x2 * cos[..., 1::2]# 3. 將旋轉后的兩半合并rotated_x = torch.cat([rotated_x1, rotated_x2], dim=-1)return rotated_x# --- 集成到多頭注意力中 ---
class RoPEMultiHeadAttention(nn.Module):def __init__(self, d_model, num_heads, max_len=512):super().__init__()assert d_model % num_heads == 0self.d_model = d_modelself.num_heads = num_headsself.head_dim = d_model // num_headsself.q_proj = nn.Linear(d_model, d_model)self.k_proj = nn.Linear(d_model, d_model)self.v_proj = nn.Linear(d_model, d_model)self.out_proj = nn.Linear(d_model, d_model)self.rotary_encoder = RotaryPositionalEncoding(self.head_dim, max_len)def forward(self, x, mask=None):batch_size, seq_len, _ = x.shape# 1. 線性投影q = self.q_proj(x)k = self.k_proj(x)v = self.v_proj(x)# 2. 改變形狀以適應多頭# shape: (batch_size, num_heads, seq_len, head_dim)q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)# 3. 對 Q 和 K 應用旋轉位置編碼q = self.rotary_encoder(q)k = self.rotary_encoder(k)# 4. 計算注意力得分# scores.shape: (batch_size, num_heads, seq_len, seq_len)scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)if mask is not None:scores = scores.masked_fill(mask == 0, float('-inf'))attention = torch.softmax(scores, dim=-1)# 5. 應用注意力到 V# context.shape: (batch_size, num_heads, seq_len, head_dim)context = torch.matmul(attention, v)# 6. 恢復形狀并進行最終投影context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)output = self.out_proj(context)return output
RoPE 作為目前最受歡迎的位置編碼,其優勢如下:
良好的長度外推性:由于其相對位置的性質,RoPE 能夠很好地泛化到比訓練時更長的序列。
隨距離增加而衰減的注意力:RoPE 的數學性質天然地使得隨著相對距離的增加,注意力得分會有一個衰減的趨勢,這符合語言直覺(離得越遠的詞關系越弱)。
高性能:它不引入額外的模型參數,并且計算非常高效,可以無縫集成到現有的自注意力框架中。