從零開始大模型之編碼注意力機制
- 1 長序列建模中的問題
- 2 使用注意力機制捕捉數據依賴關系
- 3 自注意力機制
- 4 實現帶可訓練權重的自注意力機制
- 5 利用因果注意力隱藏未來詞匯
- 6 將單頭注意力擴展到多頭注意力
- 7 Pytorch附錄
- 7.1 torch.nn.Linear
多頭+掩碼+可訓練權重的注意力機制。
為什么要自注意力機制?為什么要帶訓練權重的自注意力機制?為什么要增加掩碼功能?為什么要增加多頭功能?
1 長序列建模中的問題
無注意力機制存在語法不對齊,以及沒有聯系上下文的缺點。Transformer相比傳統的RNN能夠解決較長距離的依賴。
2 使用注意力機制捕捉數據依賴關系
自注意力機制是transformer機制中一種重要機制,它通過允許一個序列種的每個位置與同一序列中其他所有位置進行交互并權衡其重要性,能夠更高效的輸入表示。
3 自注意力機制
自注意力機制是序列token中,當前token與前后上下文的token之間的線性組合關系構成,構成的這個向量叫做上下文向量。上下文向量 = 各個詞元嵌入向量的加權和。
之所以引入注意力機制,就是為了進行語義的連貫,對于每個輸入的token詞元來說,不是獨立的單獨詞元;這就像我們做閱讀理解,總是要結合上下文要推敲某句話的含義是一樣的,通過自注意力機制,讓我們能夠關注到當前文章內容與上下文之間的聯系。
自注意力機制詳細的操作步驟如下:
1.計算選定token的詞元嵌入向量與所有token序列的詞元嵌入向量點積就得到注意力分數
2.softmax歸一化注意力分數
3.用歸一化的注意力分數對詞元嵌入向量加權求和,得到該位置的上下文向量,依次類推計算所有位置的上下文向量。
相關實現代碼:
import torch
import torch.nn.functional as F #該類提供了卷積,池化,激活等函數
# 假設對應詞序列的token對應的詞元嵌入向量為下面張量所示
inputs = torch.tensor([[0.43, 0.15, 0.89], # Your (x^1)[0.55, 0.87, 0.66], # journey (x^2)[0.57, 0.85, 0.64], # starts (x^3)[0.22, 0.58, 0.33], # with (x^4)[0.77, 0.25, 0.10], # one (x^5)[0.05, 0.80, 0.55]] # step (x^6)
)
#計算當前query與其他的token的點積后的點積矩陣,直接考慮矩陣乘法
att_scores = torch.matmul(inputs,inputs.T)#形成6x6的點積矩陣,每一行都是該token與其他token的點積注意力分數
# 對于每一行進行softmax歸一化
att_weights = F.softmax(att_scores, dim=1) # 計算每一行的softmax歸一化,dim=1表示對行進行歸一化
# 計算加權和,每一行即為一個上下文向量
context_vector = torch.matmul(att_weights,inputs)
4 實現帶可訓練權重的自注意力機制
帶可訓練權重的自注意力機制是在簡單的自注意力機制上引入了可訓練權重矩陣Wq,Wk,Wv,通過反向傳播算法進行更新,同時帶可訓練權重的自注意力機制又被稱為是縮放點積注意力。引入帶可訓練權重的自注意力機制一方面是為了能夠動態的聯系上下文,這是因為簡單的自注意力機制只能通過點積靜態計算相關權重,另一方面受到人類在數據庫中檢索信息的啟發,其本質是模仿人類搜索查詢,其中Q代表query(查詢),K代表key(鍵),V代表value(值)。
實現帶可訓練權重的自注意力機制仍然是三個核心步驟,計算權重系數,系數歸一化,系數線性疊加。
第一步,對于每個詞元嵌入向量,將q(query)查詢向量依次定為第一個token,第二個token的嵌入向量,在確定某token下的查詢向量,然后計算別的token的嵌入向量與權重矩陣Wk和Wv矩陣的乘積,可以得到k(key)鍵向量,v(value)值向量。
第二步將得到的當前的查詢向量q2,與其他所有token的k向量進行點積就得到每個token在上下文向量的占比分數,將該分數進行softmax歸一化后即為權重系數。
最后,用權重系數對每個token的值向量進行線性疊加就能得到上下文向量,這個上下文向量能夠動態理解語義,因為Wq,Wk,Wv是動態訓練可更新得到的。
總的來說,上面最后的公式可以總結為:
Attention(Q,K,V)=softmax(QKTdk)VAttention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_{k}}})VAttention(Q,K,V)=softmax(dk??QKT?)V
以下是相關步驟的簡練總價:
引入可訓練權重矩陣,Wq,Wk,Wv。
1.對當前位置的詞元嵌入向量分別與Wq,Wk,Wv相乘得到q,k,v,拿著這個當前位置的q,所有位置的k向量分別相乘,得到權重系數
2.對權重系數進行softmax歸一化處理
3.用權重系數對所有v向量進行加權求和得到該位置的上 下文向量,依次類推,計算所有上下文向量。
相關代碼如下:
#%% 帶訓練權重的自注意力機制
import torch
import torch.nn.functional as F #該類提供了卷積,池化,激活等函數
# 假設對應詞序列的token對應的詞元嵌入向量為下面張量所示
inputs = torch.tensor([[0.43, 0.15, 0.89], # Your (x^1)[0.55, 0.87, 0.66], # journey (x^2)[0.57, 0.85, 0.64], # starts (x^3)[0.22, 0.58, 0.33], # with (x^4)[0.77, 0.25, 0.10], # one (x^5)[0.05, 0.80, 0.55]] # step (x^6)
)
torch.manual_seed(123)#設置隨機種子,確保重復實驗能夠產生相同的效果
# 初始化Wq,Wk,Wv權重矩陣
Wq_matrix = torch.nn.Parameter(torch.rand(3, 2), requires_grad=False) #假設權重矩陣是3x2
Wk_matrix = torch.nn.Parameter(torch.rand(3, 2), requires_grad=False)
Wv_matrix = torch.nn.Parameter(torch.rand(3, 2), requires_grad=False)
# 計算所有token構成的q,k,v矩陣,所有嵌入詞元向量矩陣inputs與權重矩陣對應矩陣相乘
q = torch.matmul(inputs, Wq_matrix)
k = torch.matmul(inputs, Wk_matrix)
v = torch.matmul(inputs, Wv_matrix)
# 結合注意力的計算公式
context_vector = torch.matmul(F.softmax(torch.matmul(q, k.T) / torch.sqrt(torch.tensor(k.shape[1])), dim=-1), v)
為了后續的方便,同時考慮到nn.Linear相比于nn.Parameter,提供了優化的初始化方案,在模型上定義了一個抽象的層:
#%% 帶訓練權重的自注意力機制的抽象類
import torch
import torch.nn.functional as F #該類提供了卷積,池化,激活等函數
class SelfAttention(nn.Module):def __init__(self,d_in=3,d_out=2,qkv_vias=False):super().__init__() #繼承神經網絡# Wq,Wk,Wv的初始化self.Wq_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wk_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wv_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣def foward(self,inputs):#前向傳播計算q,k,vq = self.Wq_matrix(inputs) #inputs 與 Wq_matrix矩陣相乘k = self.Wk_matrix(inputs) # inputs 與 Wk_matrix矩陣相乘v = self.Wv_matrix(inputs) # inputs 與 Wv_matrix矩陣相乘attention_scores = q @ k.T # 注意力分數attention_weights = torch.softmax(attention_scores / k.shape()**0.5,dim=1) #softmax歸一化權重系數context_vec_matrix = attention_weights @ v #線性疊加求上下文向量矩陣return context_vec_matrix
5 利用因果注意力隱藏未來詞匯
權重系數矩陣右上角進行遮擋,然后歸一化,使用dropout隨機丟棄(隨機置零),防止過擬合。
事實上,考慮到邏輯的時空的連貫性,我們說話的時候是因果關系的,比如當我心里想說“我很想吃西瓜”,但是當我說出“我很想…”的時候,后續的的詞你是不知道的,這就是因果關系,只知道當前及過去,當前的詞是“想”,過去的詞是“我很”,所以理論上來說,對于“我很想吃西瓜”這句話,應該逐個詞掩碼,說==“我”的時候,掩碼”很想吃西瓜“,當說”我很“的時候,掩碼”想吃西瓜“==,依次類推,所以因果注意力的本質上是更接近真實自然的真實語境,基于此,事實上我們只要基于可訓練權重自注意力機制上,對softmax歸一化權重矩陣進行上面部分掩碼,然后重新歸一化計算系數。
考慮采用的思路是創建上三角不包含對角線的上三角塊,對權重系數矩陣的上三角部分掩碼為?∞-\infty?∞,這樣再次歸一化的時候,由于e?∞e^{-\infty}e?∞為零,就能實現因果+帶訓練權重的注意力機制,其相關代碼如下:
#%% 掩碼+帶訓練權重的自注意力機制的抽象類
import torch
import torch.nn.functional as F #該類提供了卷積,池化,激活等函數
class SelfAttention(nn.Module):def __init__(self,d_in=3,d_out=2,qkv_vias=False):super().__init__() #繼承神經網絡# Wq,Wk,Wv的初始化self.Wq_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wk_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wv_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣def foward(self,inputs):#前向傳播計算q,k,vq = self.Wq_matrix(inputs) #inputs 與 Wq_matrix矩陣相乘k = self.Wk_matrix(inputs) # inputs 與 Wk_matrix矩陣相乘v = self.Wv_matrix(inputs) # inputs 與 Wv_matrix矩陣相乘attention_scores = q @ k.T # 注意力分數attention_weights = torch.softmax(attention_scores / k.shape()**0.5,dim=1) #softmax歸一化權重系數mask = torch.triu(torch.ones(6,6),diagonal=1)#上三角(不包含對角線)矩陣創建masked = attention_weights.masked_fill(mask.bool(),-torch.inf)#對矩陣中為1的上半部分不包含對角線進行掩碼為-infattention_weights = torch.softmax(masked/k.shape()**0.5,dim=-1)#再次使用softmax歸一化context_vec_matrix = attention_weights @ v #線性疊加求上下文向量矩陣return context_vec_matrix
為了防止過擬合,也就是數據只在訓練集上表現較好的效果,事實上是模型過于復雜,舉個簡單列子來理解,就好比原來是三次函數的趨勢,你要了7次函數,雖然也可以擬合訓練集,但是模型更加復雜了,這樣預測的時候,容易陷入局部更加詳細的部分,為了解決這個問題,使用了dropout來隨機丟棄權重矩陣系數中的部分系數,或者說叫置零,通常設置隨機丟棄率來合理調整丟棄的比列,以下是增加了dropout后的代碼:
#%% dropout掩碼+帶訓練權重的自注意力機制的抽象類
import torch
import torch.nn.functional as F #該類提供了卷積,池化,激活等函數
class SelfAttention(nn.Module):def __init__(self,d_in=3,d_out=2,qkv_vias=False):super().__init__() #繼承神經網絡# Wq,Wk,Wv的初始化self.Wq_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wk_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.Wv_matrix = nn.Linear(d_in, d_out, bias=qkv_vias) #3x2全連接層的抽象層,默認會有一個初始化權重矩陣self.dropout = nn.Dropout(0.1) #dropout隨機丟棄10%def foward(self,inputs):#前向傳播計算q,k,vq = self.Wq_matrix(inputs) #inputs 與 Wq_matrix矩陣相乘k = self.Wk_matrix(inputs) # inputs 與 Wk_matrix矩陣相乘v = self.Wv_matrix(inputs) # inputs 與 Wv_matrix矩陣相乘attention_scores = q @ k.T # 注意力分數attention_weights = torch.softmax(attention_scores / k.shape()**0.5,dim=1) #softmax歸一化權重系數mask = torch.triu(torch.ones(6,6),diagonal=1)#上三角(不包含對角線)矩陣創建masked = attention_weights.masked_fill(mask.bool(),-torch.inf)#對矩陣中為1的進行掩碼為-infattention_weights = torch.softmax(masked/k.shape()**0.5,dim=-1)#再次使用softmax歸一化attention_weights = self.dropout(attention_weights) # 先對歸一化得權重系數進行隨機丟棄dropout,防止過擬合context_vec_matrix = attention_weights @ v #線性疊加求上下文向量矩陣return context_vec_matrix
6 將單頭注意力擴展到多頭注意力
將多個來源的上下文向量進行線性疊加。
class MultiHeadAttention(nn.Module):def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):super().__init__()assert (d_out % num_heads == 0), \"d_out must be divisible by num_heads"self.d_out = d_outself.num_heads = num_headsself.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dimself.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)self.out_proj = nn.Linear(d_out, d_out) # Linear layer to combine head outputsself.dropout = nn.Dropout(dropout)self.register_buffer("mask",torch.triu(torch.ones(context_length, context_length),diagonal=1))def forward(self, x):b, num_tokens, d_in = x.shape# As in `CausalAttention`, for inputs where `num_tokens` exceeds `context_length`, # this will result in errors in the mask creation further below. # In practice, this is not a problem since the LLM (chapters 4-7) ensures that inputs # do not exceed `context_length` before reaching this forwarkeys = self.W_key(x) # Shape: (b, num_tokens, d_out)queries = self.W_query(x)values = self.W_value(x)# We implicitly split the matrix by adding a `num_heads` dimension# Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) values = values.view(b, num_tokens, self.num_heads, self.head_dim)queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)# Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)keys = keys.transpose(1, 2)queries = queries.transpose(1, 2)values = values.transpose(1, 2)# Compute scaled dot-product attention (aka self-attention) with a causal maskattn_scores = queries @ keys.transpose(2, 3) # Dot product for each head# Original mask truncated to the number of tokens and converted to booleanmask_bool = self.mask.bool()[:num_tokens, :num_tokens]# Use the mask to fill attention scoresattn_scores.masked_fill_(mask_bool, -torch.inf)attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)attn_weights = self.dropout(attn_weights)# Shape: (b, num_tokens, num_heads, head_dim)context_vec = (attn_weights @ values).transpose(1, 2) # Combine heads, where self.d_out = self.num_heads * self.head_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)context_vec = self.out_proj(context_vec) # optional projectionreturn context_vectorch.manual_seed(123)batch_size, context_length, d_in = batch.shape
d_out = 2
mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads=2)context_vecs = mha(batch)print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)
7 Pytorch附錄
7.1 torch.nn.Linear
定義了一個全連接層,torch.nn.Linear(in_features,out_features,bias)
torch.nn.Linear(in_features,out_features,bias)是一個抽象全連接神經網絡層,輸入x,經過y=xwT+by=xw^{T}+by=xwT+b,輸出y,特別地當偏置項為0的時候y=xwTy=xw^{T}y=xwT,事實上就是一個矩陣的乘法,輸入x與wTw^{T}wT相乘,而當這個抽象全連接神經網絡層被定義的時候,權重矩陣wTw^{T}wT以及偏置項b就會自動初始化。