《從零構建大語言模型》學習筆記4，自注意力機制1

---

---

前言

本書原項目地址：https://github.com/rasbt/LLMs-from-scratch
我們進入第三章，探討自注意力機制——這是大語言模型的核心基礎算法。自注意力中的"自"表示該機制能夠分析輸入序列內部不同位置之間的關聯，動態計算注意力權重。它通過學習輸入元素（如句子中的單詞或圖像中的像素）之間的相互關系和依賴模式來實現這一功能。

---

一、實現一個簡單的無訓練權重的自注意力機制

比如我們有6個詞元的embeddings的向量，接下來看下怎么計算第二個詞元與其它詞元之間的注意力分數。

計算方法也很簡單，就是把第二個詞元向量分別與其它詞元向量做點積，原理圖如下：

點積后得到的就是每一個向量與第二向量的自注意分數，然后進行歸一化就得到了注意力權重。
計算代碼如下：

import torchinputs = 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 = inputs[1]  # 2nd input token is the queryattn_scores_2 = torch.empty(inputs.shape[0])
#建立一個未初始化的張量來記錄注意力得分
for i, x_i in enumerate(inputs):attn_scores_2[i] = torch.dot(x_i, query) # 相似性度量計算attention分數# 從公式上看也就是點乘
print(attn_scores_2)attn_weights_2 = torch.softmax(attn_scores_2, dim=0)
print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())
#用torch優化過的softmax對邊緣值也挺友好的

然后把得到的注意力權重又分別與對應的向量相乘后再累加，就得到了上下文向量，這個就是我們最后要求的輸出。

代碼如下：

query = inputs[1] # 2nd input token is the query
context_vec_2 = torch.zeros(query.shape)
#創造一個內容的零向量
for i,x_i in enumerate(inputs):context_vec_2 += attn_weights_2[i]*x_i#把不同內容的向量+起來
print(context_vec_2)

以上就是實現一個簡單的無訓練權重的自注意力機制

二、實現具有可訓練權重的自注意力機制

1. 分步計算注意力權重

上述注意力計算過程可拆解為三個步驟，每個步驟都需要使用詞向量。為此，我們為每個步驟的詞向量分別乘以可訓練的參數矩陣（Wq、Wk、Wv），從而得到對應的query、key和value向量。這樣計算過程就如下圖所示：


同樣使用上面的例子，先用第二個詞元向量點乘對應的Wq矩陣得到q2，然后用q2去點乘其它所有詞元的key,就得到對應詞元的注意力分數，同樣歸一化后得到對應的注意力權重。
最后用對應的注意力權重與value想成后累加，就得到了上下文向量。當然應用可以訓練的參數矩陣，所以后面可以根據上下文向量結果來訓練參數矩陣。
實現代碼如下：

x_2 = inputs[1] # second input element
d_in = inputs.shape[1] # the input embedding size, d=3
d_out = 2 # the output embedding size, d=2torch.manual_seed(123)
#固定隨機種子確保可復現性W_query = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_key   = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.rand(d_in, d_out), requires_grad=False)
#初始化三個矩陣來存放
#不要求梯度降低了復雜度query_2 = x_2 @ W_query # _2 because it's with respect to the 2nd input element
key_2 = x_2 @ W_key 
value_2 = x_2 @ W_value
#點積計算
print(query_2)keys = inputs @ W_key 
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("values.shape:", values.shape)
#中途檢驗下keys_2 = keys[1] # Python starts index at 0
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)attn_scores_2 = query_2 @ keys.T # All attention scores for given query
print(attn_scores_2)
#計算注意力跟query值d_k = keys.shape[1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
#壓縮函數, 有利于儲存與比較
print(attn_weights_2)context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

2.實現自注意力Python類

把上面的過程集中到一個類里面，并且按照pytorch里面構建神經網絡的方式來重寫這個類，__init__函數是初始化一些參數，其中就包括用nn模塊里面的線性層來初始化W_query，W_key ，W_value 這三個參數矩陣。forward函數是這個網絡的前向運算過程，就是用初始化后的K，Q，V矩陣按照上面講到的順序進行矩陣相乘，最后得到上下文向量矩陣。
代碼如下（示例）：

class SelfAttention_v2(nn.Module):def __init__(self, d_in, d_out, qkv_bias=False):super().__init__()self.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)#權重初始化def forward(self, x):keys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.T #Query跟Key的計算 得出初始的分數傳遞到后面進行歸一化操作attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ values#直接基于注意力對于文本計算return context_vectorch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

這個類還需進一步優化，需要引入掩碼的概念。掩碼的作用是遮蓋注意力權重矩陣的特定部分，通常是對矩陣的上三角部分進行處理。因為在實際推理過程中，模型需要預測后續未知的詞元，所以在訓練階段就要通過掩碼將未來的權重信息遮蓋掉，讓模型學會對未知信息的合理擬合。如果不這樣做，可能會導致模型過快收斂。原理如下圖：

代碼很簡單，就是給權重矩陣乘一個三角矩陣：

context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
#Mask矩陣,直接保留Diagonal下部分的,上部分掩蓋掉
print(mask_simple)masked_simple = attn_weights*mask_simple
print(masked_simple)
#簡單的效果圖row_sums = masked_simple.sum(dim=-1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)
#掩碼之后的softmax

記得掩碼后，要重新歸一化。
為了加大訓練難點，還會把卷積網絡中比較成熟的drop層也引用進來，就是隨機丟棄權重矩陣中的一些權重。

實現代碼如下：

torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5) 
# dropout rate of 50%丟包率doge
example = torch.ones(6, 6) 
# create a matrix of ones滿的6*6矩陣被1包圓了print(dropout(example))
#輸出需要被放大相應的倍數,為了維持恒定torch.manual_seed(123)
print(dropout(attn_weights))

把以上兩個技巧應用到類里面后，重寫的代碼如下：

class CausalAttention(nn.Module):def __init__(self, d_in, d_out, context_length,dropout, qkv_bias=False):#初始化定義網絡結構和參數super().__init__()self.d_out = d_outself.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.dropout = nn.Dropout(dropout) # Newself.register_buffer('mask', torch.triu(torch.ones(context_length, context_length), diagonal=1)) # New#定義QKV并對進行dropout防止過擬合#注冊mask向量,對未來進行負無窮的擬合def forward(self, x):b, num_tokens, d_in = x.shape # New batch dimension b#提取batch的大小、token的數量、跟寬度keys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)#進行運算計算attn_scores = queries @ keys.transpose(1, 2) # Changed transpose#通過點積來計算attention的數值attn_scores.masked_fill_(  # New, _ ops are in-placeself.mask.bool()[:num_tokens, :num_tokens], -torch.inf)  # `:num_tokens` to account for cases where the number of tokens in the batch is smaller than the supported context_sizeattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1## 縮放因子 √d，用于穩定梯度)#在時間順序上進行mask確保信息不會被泄露attn_weights = self.dropout(attn_weights) # New#防止過擬合的dropout處理方式context_vec = attn_weights @ values# 根據注意力權重計算上下文向量return context_vectorch.manual_seed(123)
context_length = batch.shape[1]
ca = CausalAttention(d_in, d_out, context_length, 0.0)context_vecs = ca(batch)print(context_vecs)
print("context_vecs.shape:", context_vecs.shape)

三、將單頭注意力擴展到多頭注意力

在卷積神經網絡(CNN)中，通常采用不同尺寸的卷積核（如3×3、5×5等）來捕獲圖像的多尺度特征。這些不同尺寸卷積核提取的特征圖經過通道維度的拼接（concat）后，能形成更全面的特征表示。類似地，注意力網絡通過初始化多組q、k、v參數來獲取不同的上下文向量并進行合并，這種機制被稱為多頭注意力。

比較簡單的實現方式是，使用for循環做多次單注意力計算，然后再拼接就可以了，代碼如下：

class MultiHeadAttentionWrapper(nn.Module):def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):super().__init__() #多個實例,每個都是一個頭self.heads = nn.ModuleList([CausalAttention(d_in, d_out, context_length, dropout, qkv_bias) for _ in range(num_heads)])def forward(self, x):return torch.cat([head(x) for head in self.heads], dim=-1)#模型的訓練torch.manual_seed(123)context_length = batch.shape[1] # This is the number of tokens
d_in, d_out = 3, 2
mha = MultiHeadAttentionWrapper(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)

然而，上述方法的效率較低，主要體現在以下兩個方面：首先，這種方法需要重復進行n次參數初始化和前向傳播計算（n代表注意力頭的數量），導致計算資源的浪費；其次，多個獨立的參數矩陣會導致內存訪問不連續，降低緩存命中率。更常見且高效的做法是在模型初始化階段就進行維度擴展。最后代碼如下：

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 dim#初始化頭的維度、數量self.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 outputs#頭的輸出結合線性層self.dropout = nn.Dropout(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.shapekeys = 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# 將掩碼縮減到當前 token 數量，并轉換為布爾型# 進而實現動態遮蔽,所以不用另開好幾個數組mask_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_dim#對上下文向量的形狀進行調整，確保輸出的形狀context_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)

---

總結

以上是關于注意力機制的講解，重點在于理解Q、K、V三個參數。我一直在思考為什么是三個參數，是否能構造更多參數。從理論上看，增加更多參數是可行的，但從數學角度來說，過多的線性相乘參數可能對后續求導沒有實質性幫助。此外，采用query、key、value的命名方式也使其含義更加直觀易懂。在撰寫過程中，已假設大家具備卷積神經網絡和PyTorch的基礎知識。若有任何表述不清或理解有誤之處，歡迎隨時指正交流。