自注意力機制允許模型根據輸入序列不同部分的重要性動態調整權重,從而提升模型對序列資料的理解能力。本文首先以一個簡化的自注意力機制為例,逐步講解如何計算注意力分數、注意力權重以及上下文向量。接著,引入了可訓練的權重矩陣,並使用 PyTorch 實作了更通用的自注意力模組。程式碼中使用 torch.softmax 函式進行注意力分數的歸一化,並闡述了 dim 引數的作用。同時,也解釋了使用縮放點積注意力的好處,即避免梯度消失問題,並透過矩陣乘法計算上下文向量。最後,文章提供了程式碼範例,展示瞭如何使用 PyTorch 的 nn.Linear 層簡化和最佳化自注意力模組的實作,並以圖表形式清晰地展現了自注意力機制的運作流程。

3.3 自注意力機制中對輸入的不同部分進行處理

在圖 3.12 的步驟 2 中,我們對每一行進行歸一化,使得每行的值加總為 1:

attn_weights = torch.softmax(attn_scores, dim=-1)
print(attn_weights)

這將傳回以下注意力權重張量,與圖 3.10 中的值相符:

tensor([[0.2098, 0.2006, 0.1981, 0.1242, 0.1220, 0.1452],
        [0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581],
        [0.1390, 0.2369, 0.2326, 0.1242, 0.1108, 0.1565],
        [0.1435, 0.2074, 0.2046, 0.1462, 0.1263, 0.1720],
        [0.1526, 0.1958, 0.1975, 0.1367, 0.1879, 0.1295],
        [0.1385, 0.2184, 0.2128, 0.1420, 0.0988, 0.1896]])

在 PyTorch 中,像 torch.softmax 這樣的函式中的 dim 引數指定了將在其上計算函式的輸入張量的維度。透過設定 dim=-1,我們指示 softmax 函式沿著 attn_scores 張量的最後一個維度進行歸一化。如果 attn_scores 是一個二維張量(例如,具有 [rows, columns] 的形狀),它將跨列進行歸一化,使得每行的值(對列維度求和)加總為 1。

內容解密:

  1. 歸一化過程:使用 torch.softmax 對注意力分數進行歸一化,確保每行的值加總為 1。
  2. dim 引數的作用:指定 softmax 操作的維度,-1 表示最後一個維度,即列維度。
  3. 結果驗證:檢查每行的加總是否為 1,以驗證歸一化是否正確。

我們可以驗證各行的加總是否確實為 1:

row_2_sum = sum([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
print("Row 2 sum:", row_2_sum)
print("All row sums:", attn_weights.sum(dim=-1))

結果如下:

Row 2 sum: 1.0
All row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

在圖 3.12 的第三步也是最後一步中,我們使用這些注意力權重透過矩陣乘法計算所有上下文向量:

all_context_vecs = attn_weights @ inputs
print(all_context_vecs)

在產生的輸出張量中,每一行包含一個三維上下文向量:

tensor([[0.4421, 0.5931, 0.5790],
        [0.4419, 0.6515, 0.5683],
        [0.4431, 0.6496, 0.5671],
        [0.4304, 0.6298, 0.5510],
        [0.4671, 0.5910, 0.5266],
        [0.4177, 0.6503, 0.5645]])

內容解密:

  1. 上下文向量計算:使用注意力權重和輸入向量透過矩陣乘法計算上下文向量。
  2. 輸出結構:每個上下文向量代表輸入序列中對應元素的所有相關資訊的加權和。
  3. 結果驗證:透過比較第二行的上下文向量與之前手動計算的結果,驗證程式碼的正確性。

3.4 使用可訓練權重的自注意力機制實作

接下來,我們將實作原始 Transformer 結構、GPT模型和其他流行的大語言模型(LLM)中使用的自注意力機制,也稱為縮放點積注意力(scaled dot-product attention)。圖3.13展示了這種自注意力機制如何融入實作LLM的更廣泛背景中。

圖表說明

此圖示展示了在實作大語言模型(LLM)的過程中,不同的注意力機制所處的位置。我們已經實作了一個簡化的注意力機制,現在將透過新增可訓練的權重來擴充套件自注意力機制。

圖表解說

  • 簡化自注意力:我們已經實作了簡化的注意力機制,以瞭解其基本原理。
  • 自注意力:現在,我們將透過新增可訓練權重來擴充套件自注意力機制。
  • 因果注意力多頭注意力:後續章節將進一步介紹這些擴充套件技術。
  • 資料準備和取樣、LLM架構、預訓練和微調:這些是構建LLM的不同階段,其中注意力機制是核心元件。

使用可訓練權重的自注意力機制的逐步計算

我們將透過引入三個可訓練的權重矩陣 $W_q$、$W_k$ 和 $W_v$,逐步實作自注意力機制。這三個矩陣用於將嵌入的輸入令牌 $x^{(i)}$ 投影到查詢(query)、鍵(key)和值(value)向量,如圖3.14所示。

圖表說明

此圖示展示了具有可訓練權重矩陣的自注意力機制的第一步,我們計算查詢(q)、鍵(k)和值(v)向量,用於輸入元素x。

@startuml
skinparam backgroundColor #FEFEFE
skinparam componentStyle rectangle

title 自注意力機制實作與程式碼解析

package "自注意力機制" {
    package "注意力計算" {
        component [Query 向量] as query
        component [Key 向量] as key
        component [Value 向量] as value
    }

    package "權重計算" {
        component [點積運算] as dot
        component [縮放因子] as scale
        component [Softmax 歸一化] as softmax
    }

    package "輸出生成" {
        component [注意力權重] as attn_weight
        component [上下文向量] as context
        component [PyTorch 實作] as pytorch
    }
}

query --> dot : Q·K^T
key --> dot : 計算相似度
dot --> scale : 除以 √d_k
scale --> softmax : 機率分佈
softmax --> attn_weight : 權重矩陣
value --> context : 加權求和
attn_weight --> context : 應用權重

note bottom of scale
  避免梯度消失
  穩定訓練過程
end note

collect --> clean : 原始資料
clean --> feature : 乾淨資料
feature --> select : 特徵向量
select --> tune : 基礎模型
tune --> cv : 最佳參數
cv --> eval : 訓練模型
eval --> deploy : 驗證模型
deploy --> monitor : 生產模型

note right of feature
  特徵工程包含:
  - 特徵選擇
  - 特徵轉換
  - 降維處理
end note

note right of eval
  評估指標:
  - 準確率/召回率
  - F1 Score
  - AUC-ROC
end note

@enduml

圖表解說

  • 查詢、鍵和值的計算:透過將輸入元素x與相應的權重矩陣 $W_q$、$W_k$ 和 $W_v$ 相乘,分別獲得查詢、鍵和值向量。
  • 權重矩陣的作用:這些矩陣是在模型訓練過程中更新的,從而使模型能夠學習產生“好”的上下文向量。

程式碼實作與解析

首先,我們定義輸入序列和三個可訓練的權重矩陣 $W_q$、$W_k$ 和 $W_v$:

import torch

# 定義輸入序列和權重矩陣
inputs = torch.tensor([
    [0.4, 0.1],
    [0.3, 0.7],
    [0.1, 0.8],
    [0.4, 1.1],
    [0.3, 1.0],
    [0.3, 0.9]
])

W_q = torch.tensor([[1.1, 2.], [3.,4]])
W_k = torch.tensor([[5.,6.], [7.,8]])
W_v = torch.tensor([[9.,10.], [11.,12]])

# 分別計算查詢、鍵和值向量
queries = inputs @ W_q.T
keys = inputs @ W_k.T
values = inputs @ W_v.T

print("Queries:", queries)
print("Keys:", keys)
print("Values:", values)

程式碼解密:

  • 輸入與權重矩陣定義:定義輸入序列張量 inputs 和三個權重矩陣張量 W_qW_kW_v
  • 查詢、鍵和值的計算:透過矩陣乘法,將輸入序列投影到查詢、鍵和值空間,分別得到 querieskeysvalues 張量。
  • @運算元的使用:表示矩陣乘法,用於計算查詢、鍵和值向量。

實作具備可訓練權重的自注意力機制

在前一節中,我們使用簡化的自注意力機制來計算上下文向量。在本文中,我們將實作具備可訓練權重的自注意力機制,這是 Transformer 架構中的關鍵元件。

初始化輸入與權重矩陣

首先,我們定義輸入張量 inputs,並初始化三個權重矩陣 W_queryW_keyW_value,用於將輸入元素轉換為查詢(query)、鍵(key)和值(value)向量。

x_2 = inputs[1]
d_in = inputs.shape[1]
d_out = 2

torch.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)

內容解密:

  • x_2 代表第二個輸入元素。
  • d_ind_out 分別代表輸入和輸出的維度。
  • W_queryW_keyW_value 是三個權重矩陣,用於將輸入元素轉換為查詢、鍵和值向量。
  • requires_grad=False 表示在計算過程中不計算梯度,這樣可以減少輸出的雜訊;如果需要進行模型訓練,則應設為 True

計算查詢、鍵和值向量

接下來,我們使用權重矩陣將輸入元素轉換為查詢、鍵和值向量。

query_2 = x_2 @ W_query
key_2 = x_2 @ W_key
value_2 = x_2 @ W_value
print(query_2)

內容解密:

  • query_2key_2value_2 分別代表第二個輸入元素的查詢、鍵和值向量。
  • 矩陣乘法 @ 用於計算查詢、鍵和值向量。

計算所有鍵和值向量

為了計算上下文向量,我們需要所有輸入元素的鍵和值向量。

keys = inputs @ W_key
values = inputs @ W_value
print("keys.shape:", keys.shape)
print("values.shape:", values.shape)

內容解密:

  • keysvalues 分別代表所有輸入元素的鍵和值向量。
  • 矩陣乘法 @ 用於計算所有鍵和值向量。

計算注意力分數

注意力分數是用於衡量查詢向量與鍵向量之間的相似度。

attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

attn_scores_2 = query_2 @ keys.T
print(attn_scores_2)

內容解密:

  • attn_score_22 代表第二個輸入元素的查詢向量與其鍵向量之間的注意力分數。
  • attn_scores_2 代表第二個輸入元素的查詢向量與所有鍵向量之間的注意力分數。
  • .dot() 方法用於計算兩個向量之間的點積。

計算注意力權重

注意力權重是透過對注意力分數進行縮放和 softmax 運算得到的。

d_k = keys.shape[-1]
attn_weights_2 = torch.softmax(attn_scores_2 / d_k**0.5, dim=-1)
print(attn_weights_2)

內容解密:

  • d_k 代表鍵向量的維度。
  • attn_weights_2 代表第二個輸入元素的查詢向量對應的所有注意力權重。
  • torch.softmax() 函式用於計算 softmax 值。

為何使用縮放的點積注意力?

使用縮放的點積注意力的原因是為了避免在訓練過程中出現梯度消失的問題。當嵌入維度較大時,點積運算的結果可能會很大,從而導致 softmax 函式的梯度變得很小。縮放因子 d_k**0.5 有助於減緩這個問題。

計算上下文向量

最後,我們透過將注意力權重與值向量相乘並求和來計算上下文向量。

# 省略具體實作程式碼,因為圖 3.17 已經給出了示意圖

內容解密:

  • 上下文向量是透過將注意力權重與值向量相乘並求和得到的。
  • 這一步驟是自注意力機制的最後一步,用於生成最終的輸出。

自注意力機制實作詳解

在深度學習與自然語言處理領域,自注意力機制(Self-Attention)已成為理解與處理序列資料的關鍵技術。本章節將探討自注意力機制的實作細節,並透過 Python 類別封裝實作可重複使用的自注意力模組。

自注意力機制基礎

自注意力機制的核心思想是讓模型能夠動態地關注輸入序列的不同部分,並根據其重要性進行加權平均。首先,我們來實作基本的自注意力計算步驟。

計算注意力權重

給定輸入序列 $X = [x^{(1)}, x^{(2)}, …, x^{(T)}]$,我們需要計算查詢(Query)、鍵(Key)和值(Value)矩陣。這些矩陣是透過將輸入 $X$ 與可訓練的權重矩陣 $W_q$、$W_k$ 和 $W_v$ 相乘得到的。

import torch
import torch.nn as nn

class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))
        self.W_key = nn.Parameter(torch.rand(d_in, d_out))
        self.W_value = nn.Parameter(torch.rand(d_in, d_out))
    
    def forward(self, x):
        keys = x @ self.W_key
        queries = x @ self.W_query
        values = x @ self.W_value
        
        # 計算注意力評分並進行歸一化
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        
        # 計算上下文向量
        context_vec = attn_weights @ values
        return context_vec

為什麼使用查詢、鍵和值?

查詢、鍵和值的概念源自資訊檢索和資料函式庫領域。在注意力機制中,查詢代表當前模型關注的專案(如句子中的一個詞),鍵用於比對查詢,而值則代表輸入專案的實際內容或表示。一旦模型確定哪些鍵與查詢最相關,它就會檢索對應的值。

實作最佳化:使用 PyTorch 的 Linear 層

我們可以進一步最佳化 SelfAttention_v1 的實作,利用 PyTorch 的 nn.Linear 層來執行矩陣乘法。與手動實作 nn.Parameter 相比,nn.Linear 具有最佳化的權重初始化方案,有助於更穩定和有效的模型訓練。

class SelfAttention_v2(nn.Module):
    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()
        self.query_layer = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.key_layer = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.value_layer = nn.Linear(d_in, d_out, bias=qkv_bias)
    
    def forward(self, x):
        keys = self.key_layer(x)
        queries = self.query_layer(x)
        values = self.value_layer(x)
        
        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        
        context_vec = attn_weights @ values
        return context_vec

自注意力機制視覺化

自注意力機制涉及可訓練的權重矩陣 $W_q$、$W_k$ 和 $W_v$,它們將輸入資料轉換為查詢、鍵和值。這些元件對於注意力機制至關重要。在訓練過程中,隨著模型接觸到更多的資料,它會調整這些可訓練的權重。

圖示說明

此圖示展示了自注意力機制的運作流程。輸入向量 $X$ 經過三個權重矩陣 $W_q$、$W_k$ 和 $W_v$ 的轉換,分別得到查詢(Q)、鍵(K)和值(V)矩陣。然後,根據 Q 和 K 計算注意力權重矩陣,並利用該權重矩陣對 V 進行加權求和,從而得到上下文向量(Z)。

內容解密:
  1. 自注意力機制的核心:自注意力機制允許模型動態地對輸入序列的不同部分進行加權平均,從而更好地理解序列資料。
  2. SelfAttention_v1 實作:透過手動定義可訓練的權重矩陣 $W_q$、$W_k$ 和 $W_v$,並利用矩陣乘法計算查詢、鍵和值。
  3. SelfAttention_v2 最佳化:使用 PyTorch 的 nn.Linear 層替代手動定義的權重矩陣,提供最佳化的權重初始化方案。
  4. 視覺化與理解:透過圖示展示自注意力機制的運作流程,有助於深入理解其內部工作原理。
  5. 結論與應用:自注意力機制的實作對於開發和最佳化大語言模型具有重要意義,為自然語言處理領域的研究提供了有力支援。