3 編碼注意力機制

本章內容包括：

• 探索在神經網絡中使用注意力機制的原因
• 介紹一個基本的自注意力框架，并逐步過渡到增強型自注意力機制
• 實現一個因果注意力模塊，使得LLM能夠一次生成一個token
• 使用dropout隨機屏蔽部分注意力權重以減少過擬合
• 將多個因果注意力模塊堆疊成一個多頭注意力模塊

在上一章中，你學習了如何為訓練LLM準備輸入文本。這包括將文本劃分為單詞和子詞的token，這些token可以被編碼為向量表示，即所謂的嵌入（embeddings），供LLM使用。

在本章中，我們將探討LLM架構中的一個核心部分——注意力機制，如圖3.1所示。

圖3.1 對LLM編碼的三個主要階段的思維模型：在通用文本數據集上對LLM進行預訓練，然后在有標簽數據集上進行微調。本章重點介紹注意力機制，它是LLM架構的一個核心組成部分。

注意力機制是一個內容廣泛的主題，這也是我們將整整一章專門用于它的原因。我們將在很大程度上將這些注意力機制單獨進行考察，并從機制層面進行深入探討。在下一章中，我們將編寫圍繞自注意力機制的LLM其余部分的代碼，以便看到它的實際運作，并創建一個用于生成文本的模型。

在本章的過程中，我們將實現四種不同變體的注意力機制，如圖3.2所示。

圖3.2 本圖展示了我們將在本章中編寫的不同注意力機制。從一個簡化版本的自注意力開始，然后逐步加入可訓練權重。因果注意力機制在自注意力的基礎上添加了一個掩碼，使得LLM能夠一次生成一個詞。最后，多頭注意力將注意力機制組織為多個頭，從而使模型能夠并行捕捉輸入數據的多個方面。

3.1 建模長序列的問題

在本章稍后我們深入探討LLM核心的自注意力機制之前，先來看一看：在沒有注意力機制、早于LLM的架構中，存在什么問題？假設我們要開發一個語言翻譯模型，用于將文本從一種語言翻譯為另一種語言。如圖3.3所示，我們不能僅僅逐詞翻譯文本，因為源語言和目標語言之間的語法結構不同。

圖3.3 當將文本從一種語言翻譯為另一種語言（例如從德語翻譯為英語）時，無法僅僅逐詞翻譯。相反，翻譯過程需要——

為了解決不能逐詞翻譯的問題，通常會使用一個包含兩個子模塊的深度神經網絡，即所謂的 編碼器-解碼器（encoder-decoder） 結構。編碼器的任務是首先讀入并處理整個文本，然后解碼器生成翻譯后的文本。

我們已經在第1章（第1.4節：將LLM用于不同任務）介紹Transformer架構時簡要討論過編碼器-解碼器網絡。在Transformer出現之前，循環神經網絡（RNN）是語言翻譯中最常見的編碼器-解碼器架構。

RNN是一種神經網絡，其前一步的輸出被作為當前步驟的輸入，這使得它非常適合處理文本等序列數據。如果你不熟悉RNN，不用擔心，你不需要了解RNN的詳細工作原理就能理解我們接下來的討論；我們在這里更關注編碼器-解碼器結構的一般概念。

在編碼器-解碼器RNN中，輸入文本被送入編碼器，編碼器按順序對其進行處理。編碼器在每一步都會更新其隱藏狀態（即隱藏層中的內部數值），并試圖在最終的隱藏狀態中捕捉整個輸入句子的含義，如圖3.4所示。隨后，解碼器使用這個最終的隱藏狀態開始逐詞生成翻譯后的句子。解碼器也會在每一步更新其隱藏狀態，該隱藏狀態被認為包含了下一詞預測所需的上下文信息。

圖3.4 在Transformer模型出現之前，編碼器-解碼器RNN是機器翻譯中的一種流行選擇。編碼器將源語言中的一系列token作為輸入，編碼器的隱藏狀態（即神經網絡的中間層）編碼了整個輸入序列的壓縮表示。然后，解碼器利用其當前的隱藏狀態開始逐個token地進行翻譯。

雖然我們不需要了解這些編碼器-解碼器RNN的內部工作機制，但這里的關鍵思想是：編碼器部分將整個輸入文本處理為一個隱藏狀態（記憶單元）。然后，解碼器接收這個隱藏狀態來生成輸出。你可以將這個隱藏狀態看作一個嵌入向量（embedding vector），這是我們在第2章中討論過的概念。

編碼器-解碼器RNN的主要問題和局限在于：在解碼階段，RNN無法直接訪問編碼器的早期隱藏狀態。因此，它只能依賴當前的隱藏狀態，而這個狀態封裝了所有相關的信息。這可能導致上下文信息的丟失，特別是在那些依賴關系跨越很長距離的復雜句子中。

對于不熟悉RNN的讀者來說，不必理解或學習這一架構，因為本書不會使用它。本節的核心信息是：編碼器-解碼器RNN存在一個缺陷，這個缺陷促使人們設計了注意力機制。

3.2 使用注意力機制捕捉數據依賴關系

在Transformer LLM出現之前，如前所述，通常使用RNN來處理語言建模任務，例如語言翻譯。RNN在翻譯短句時表現良好，但在處理較長文本時效果不佳，因為它們無法直接訪問輸入中的先前詞語。

這種方法的一個主要缺點是：RNN必須在將編碼信息傳遞給解碼器之前，將整個輸入編碼為一個單一的隱藏狀態，如上一節圖3.4所示。

因此，研究人員在2014年為RNN開發了所謂的Bahdanau注意力機制（以相關論文的第一作者命名），該機制對編碼器-解碼器RNN進行修改，使得解碼器在每一步解碼時可以選擇性地訪問輸入序列中的不同部分，如圖3.5所示。

圖3.5 通過使用注意力機制，網絡中生成文本的解碼器部分可以選擇性地訪問所有輸入token。這意味著，對于生成某個輸出token而言，一些輸入token比其他的更重要。其重要性由所謂的 注意力權重（attention weights） 決定，我們將在后文中計算它們。請注意，此圖展示的是注意力機制的一般思想，并未呈現Bahdanau機制的具體實現，后者是一種超出本書范圍的RNN方法。

有趣的是，僅僅三年之后，研究人員發現構建用于自然語言處理的深度神經網絡并不需要RNN架構，并提出了原始的Transformer架構（在第1章中已討論），該架構中的自注意力機制（self-attention） 受Bahdanau注意力機制的啟發。

自注意力是一種機制，在計算序列的表示時，它允許輸入序列中的每個位置關注同一序列中的所有位置。自注意力是以Transformer架構為基礎的當代LLM（如GPT系列）的關鍵組成部分。

本章將重點編寫和理解GPT類模型中使用的這個自注意力機制，如圖3.6所示。在下一章中，我們將編寫LLM的其余部分。

圖3.6 自注意力是Transformer中的一種機制，它通過允許序列中的每個位置與同一序列中的所有其他位置進行交互并評估其重要性，從而計算出更高效的輸入表示。在本章中，我們將從零開始編寫這個自注意力機制，然后在下一章中編寫GPT類LLM的其余部分。

3.3 通過自注意力關注輸入的不同部分

我們現在將深入探討自注意力機制的內部工作原理，并學習如何從零開始編寫它。自注意力機制是基于Transformer架構的所有LLM的基石。值得注意的是，這一主題可能需要你付出相當多的專注力（無意雙關），但一旦你掌握了其基本原理，就相當于征服了本書中最難的部分之一，也是實現LLM過程中最具挑戰性的部分之一。

自注意力中的“自”

在自注意力機制中，“自”指的是該機制通過關聯單個輸入序列中的不同位置來自行計算注意力權重的能力。它能夠評估并學習輸入自身各部分之間的關系和依賴，例如一句話中的詞語之間的關系，或者一幅圖像中像素之間的關系。這一點不同于傳統的注意力機制，后者通常關注的是兩個不同序列之間元素的關系，例如在序列到序列（sequence-to-sequence）模型中，注意力可能是在輸入序列與輸出序列之間建立聯系，就像圖3.5所展示的例子那樣。

由于自注意力機制在初次接觸時可能顯得復雜，我們將從下一小節開始介紹一個簡化版本的自注意力機制。之后，在第3.4節中，我們將實現帶有可訓練權重的自注意力機制，這也是LLM中實際使用的版本。

3.3.1 一個沒有可訓練權重的簡化自注意力機制

在本節中，我們將實現一個不含任何可訓練權重的簡化版本的自注意力機制，該機制在圖3.7中進行了概括。本節的目標是在引入可訓練權重（第3.4節）之前，先闡明自注意力中的幾個關鍵概念。

圖3.7 自注意力機制的目標是為每個輸入元素計算一個上下文向量，該向量結合了來自所有其他輸入元素的信息。在本圖所示的示例中，我們計算的是上下文向量$z^{(2)}$。每個輸入元素在計算$z^{(2)}$時的重要性或貢獻程度由注意力權重${\alpha }_{21}$到${\alpha }_{2T}$決定。在計算$z^{(2)}$時，注意力權重是相對于輸入元素$x^{(2)}$和所有其他輸入元素來計算的。這些注意力權重的具體計算方式將在本節后文中進行討論。

圖3.7展示了一個輸入序列，記作$x$，由$T$個元素組成，表示為$x^{(1)}$到$x^{(T)}$。這個序列通常表示一段文本，例如一個句子，并已經被轉換為token嵌入，如第2章所解釋的那樣。

例如，考慮如下輸入文本：“Your journey starts with one step.” 在這種情況下，序列中的每個元素，例如$x^{(1)}$，對應一個$d$維的嵌入向量，表示某個特定的token，比如“Your”。在圖3.7中，這些輸入向量被表示為三維嵌入。

在自注意力機制中，我們的目標是為輸入序列中的每個元素$x^{(i)}$計算上下文向量$z^{(i)}$。上下文向量可以理解為一種增強后的嵌入向量。

為了說明這個概念，我們聚焦于第二個輸入元素$x^{(2)}$（它對應token“journey”）的嵌入向量，以及圖3.7底部所示的對應上下文向量$z^{(2)}$。這個增強后的上下文向量$z^{(2)}$是一個包含了$x^{(2)}$以及所有其他輸入元素$x^{(1)}$到$x^{(T)}$信息的嵌入。

在自注意力機制中，上下文向量發揮著至關重要的作用。它們的目的是通過整合序列中所有其他元素的信息，為輸入序列（如一個句子）中的每個元素創建增強表示，如圖3.7所示。這在LLM中至關重要，因為LLM需要理解句子中各個詞語之間的關系與相關性。稍后，我們將引入可訓練的權重，幫助LLM學習如何構造這些上下文向量，使其對于生成下一個token而言是有意義的。

在本節中，我們將實現一個簡化版的自注意力機制，逐步計算這些權重和對應的上下文向量。

考慮如下輸入句子，它已經按照第2章所討論的方法被嵌入為三維向量。我們選擇一個較小的嵌入維度進行示例，以確保內容能在頁面中展示而不換行：

from importlib.metadata import versionprint("torch version:", version("torch"))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)
)

實現自注意力機制的第一步是計算中間值$\omega$，也稱為注意力分數（attention scores），如圖3.8所示。

圖3.8 本節的總體目標是說明如何使用第二個輸入序列$x^{(2)}$作為查詢（query）來計算上下文向量$z^{(2)}$。此圖展示了第一個中間步驟，即通過點積計算查詢$x^{(2)}$與所有其他輸入元素之間的注意力分數$\omega$。（注意：圖中的數值被截斷為小數點后一位，以減少視覺干擾。）

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) # dot product (transpose not necessary here since they are 1-dim vectors)print(attn_scores_2)

tensor([0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865])

理解點積運算

點積本質上是一種將兩個向量按元素相乘后再將乘積求和的簡潔方式，我們可以通過以下方式進行演示：

res = 0.for idx, element in enumerate(inputs[0]):res += inputs[0][idx] * query[idx]print(res)
print(torch.dot(inputs[0], query))

輸出結果表明，逐元素相乘后求和與點積的結果是相同的：

tensor(0.9544)
tensor(0.9544)

除了將點積視為一種將兩個向量組合為一個標量的數學工具之外，點積也是一種相似度的度量方式，因為它量化了兩個向量之間的對齊程度：點積越大，表示這兩個向量越相似或越對齊。在自注意力機制的背景下，點積決定了序列中各元素彼此關注的程度：點積越大，表示兩個元素之間的相似度越高，注意力分數也越高。

在下一步中，如圖3.9所示，我們將對前面計算出的注意力分數進行歸一化處理。

圖3.9 在計算了關于輸入查詢$x^{(2)}$的注意力分數${\omega }_{21}$到${\omega }_{2T}$之后，下一步是通過歸一化這些注意力分數，得到注意力權重${\alpha }_{21}$到${\alpha }_{2T}$。

圖3.9中所示的歸一化操作的主要目標是獲得注意力權重，這些權重的總和為1。這種歸一化是一種約定，它有助于解釋模型的行為，并在LLM的訓練中保持穩定性。

以下是實現這一歸一化步驟的一個簡單方法：

attn_weights_2_tmp = attn_scores_2 / attn_scores_2.sum()print("Attention weights:", attn_weights_2_tmp)
print("Sum:", attn_weights_2_tmp.sum())

如輸出所示，注意力權重的總和現在為1：

Attention weights: tensor([0.1455, 0.2278, 0.2249, 0.1285, 0.1077, 0.1656])
Sum: tensor(1.0000)

在實踐中，更常見且更推薦的做法是使用softmax函數進行歸一化。這種方法在處理極端數值時效果更好，并在訓練過程中具有更優的梯度性質。下面是一個用于歸一化注意力分數的softmax函數的基礎實現：

def softmax_naive(x):return torch.exp(x) / torch.exp(x).sum(dim=0)attn_weights_2_naive = softmax_naive(attn_scores_2)print("Attention weights:", attn_weights_2_naive)
print("Sum:", attn_weights_2_naive.sum())

如輸出所示，softmax函數同樣實現了目標，并將注意力權重歸一化為總和為1：

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

此外，softmax函數還能確保所有注意力權重始終為正數。這使得輸出可以被解釋為概率或相對重要性，權重越高表示該輸入越重要。

需要注意的是，這個樸素的softmax實現（softmax_naive）在處理很大或很小的輸入值時可能會遇到數值不穩定的問題，例如上溢和下溢。因此，在實際應用中，建議使用PyTorch優化過的softmax實現，該實現經過廣泛優化以提高性能：

attn_weights_2 = torch.softmax(attn_scores_2, dim=0)print("Attention weights:", attn_weights_2)
print("Sum:", attn_weights_2.sum())

在這種情況下，我們可以看到它與之前的softmax_naive函數給出了相同的結果：

Attention weights: tensor([0.1385, 0.2379, 0.2333, 0.1240, 0.1082, 0.1581])
Sum: tensor(1.)

現在我們已經計算出了歸一化后的注意力權重，接下來就可以進行圖3.10所示的最后一步：通過將嵌入的輸入token$x^{(i)}$與相應的注意力權重相乘，然后對結果向量求和，來計算上下文向量$z^{(2)}$。

圖3.10 在計算并歸一化注意力分數以獲得關于查詢$x^{(2)}$的注意力權重后，最后一步是計算上下文向量$z^{(2)}$。這個上下文向量是所有輸入向量$x^{(1)}$到$x^{(T)}$的加權組合，其中權重即為注意力權重。

圖3.10所示的上下文向量$z^{(2)}$是通過對所有輸入向量進行加權求和計算得到的。這一過程包括將每個輸入向量與其對應的注意力權重相乘：

query = inputs[1] # 2nd input token is the querycontext_vec_2 = torch.zeros(query.shape)
for i,x_i in enumerate(inputs):context_vec_2 += attn_weights_2[i]*x_iprint(context_vec_2)

該計算的結果如下所示：

tensor([0.4419, 0.6515, 0.5683])

在下一節中，我們將對這一用于計算上下文向量的過程進行泛化，以便同時計算所有上下文向量。

3.3.2 為所有輸入token計算注意力權重

在上一節中，我們計算了輸入2的注意力權重和上下文向量，如圖3.11中高亮顯示的行所示。現在，我們將擴展這一計算，來為所有輸入計算注意力權重和上下文向量。

圖3.11 高亮的那一行表示以第二個輸入元素作為查詢時的注意力權重，這是我們在上一節中計算的內容。本節將對該計算過程進行泛化，以獲得所有其他的注意力權重。

我們遵循與之前相同的三個步驟，如圖3.12所總結的那樣，唯一的不同是在代碼中做了一些修改，以便計算所有上下文向量，而不僅僅是第二個上下文向量$z^{(2)}$。

圖3.12

首先，在圖3.12所示的步驟1中，我們添加了一個額外的for循環，用于計算所有輸入對之間的點積：

attn_scores = torch.empty(6, 6)for i, x_i in enumerate(inputs):for j, x_j in enumerate(inputs):attn_scores[i, j] = torch.dot(x_i, x_j)print(attn_scores)

得到的注意力分數如下：

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

上述張量中的每個元素表示每一對輸入之間的注意力分數，如圖3.11所示。請注意，圖3.11中的值是歸一化后的，因此與上述未歸一化的注意力分數不同。我們將在之后處理歸一化問題。

在計算上述注意力分數張量時，我們使用了Python的for循環。然而，for循環通常較慢，我們可以使用矩陣乘法來實現相同的結果：

attn_scores = inputs @ inputs.T
print(attn_scores)

我們可以直觀地驗證結果與之前相同：

tensor([[0.9995, 0.9544, 0.9422, 0.4753, 0.4576, 0.6310],[0.9544, 1.4950, 1.4754, 0.8434, 0.7070, 1.0865],[0.9422, 1.4754, 1.4570, 0.8296, 0.7154, 1.0605],[0.4753, 0.8434, 0.8296, 0.4937, 0.3474, 0.6565],[0.4576, 0.7070, 0.7154, 0.3474, 0.6654, 0.2935],[0.6310, 1.0865, 1.0605, 0.6565, 0.2935, 0.9450]])

在步驟2中，如圖3.12所示，我們現在對每一行進行歸一化，使得每一行的值之和為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]])

在進入圖3.12所示的步驟3之前，我們可以快速驗證每一行的確都歸一化為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])

在第三步也是最后一步中，我們現在使用這些注意力權重通過矩陣乘法來計算所有上下文向量：

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]])

我們可以再次檢查代碼是否正確，通過將第二行與我們在第3.3.1節中計算的上下文向量$z^{(2)}$進行比較：

print("Previous 2nd context vector:", context_vec_2)

根據結果，我們可以看到之前計算的context_vec_2與上面張量中的第二行完全一致：

Previous 2nd context vector: tensor([0.4419, 0.6515, 0.5683])

這就完成了對一個簡單自注意力機制的代碼講解。

在下一節中，我們將引入可訓練權重，使LLM能夠通過數據學習，從而提升其在特定任務上的表現。

3.4 實現帶有可訓練權重的自注意力機制

在本節中，我們將實現自注意力機制，它被用于原始的Transformer架構、GPT模型以及大多數流行的LLM中。這種自注意力機制也被稱為縮放點積注意力（scaled dot-product attention）。圖3.13提供了一個思維模型，用以說明該自注意力機制如何融入實現LLM的整體背景中。

圖3.13 該思維模型展示了我們在本節中所編寫的自注意力機制在本書和本章更廣泛內容中的位置。在上一節中，我們編寫了一個簡化的注意力機制，以理解注意力機制背后的基本原理。而在本節中，我們將在該注意力機制中加入可訓練權重。在接下來的各節中，我們將通過添加因果掩碼和多頭機制進一步擴展這一自注意力機制。

如圖3.13所示，帶有可訓練權重的自注意力機制建立在之前的概念基礎之上：我們希望為特定的輸入元素，計算其上下文向量，即對輸入向量的加權求和。正如你將看到的，這一機制與我們在第3.3節中編寫的基礎自注意力機制相比，只有一些細微的區別。

最顯著的區別是引入了在模型訓練過程中不斷更新的權重矩陣。這些可訓練的權重矩陣至關重要，因為它們使得模型（具體來說是模型內部的注意力模塊）能夠學習如何生成“優質”的上下文向量。（請注意，我們將在第5章中對LLM進行訓練。）

我們將通過兩個子小節來處理這個自注意力機制。首先，我們將像以前一樣逐步編寫其實現代碼。其次，我們會將代碼組織成一個簡潔的Python類，可以將其導入到LLM架構中，而這個架構將在第4章中實現。

3.4.1 逐步計算注意力權重

我們將通過引入三個可訓練的權重矩陣$W_q$、$W_k$和$W_v$，逐步實現自注意力機制。這三個矩陣用于將嵌入后的輸入token$x^{(i)}$投影為查詢向量（query）、鍵向量（key）和值向量（value），如圖3.14所示。

圖3.14 在帶有可訓練權重矩陣的自注意力機制的第一步中，我們為輸入元素$x$計算查詢（$q$）、鍵（$k$）和值（$v$）向量。與前幾節類似，我們將第二個輸入$x^{(2)}$指定為查詢輸入。查詢向量$q^{(2)}$是通過輸入$x^{(2)}$與權重矩陣$W_q$進行矩陣乘法得到的。同樣地，鍵向量和值向量分別通過與權重矩陣$W_k$和$W_v$進行矩陣乘法得到。

在第3.3.1節中，我們將第二個輸入元素$x^{(2)}$定義為查詢（query），用于計算簡化的注意力權重以求得上下文向量$z^{(2)}$。隨后在第3.3.2節中，我們對這個過程進行了泛化，從而計算出六詞輸入句子 “Your journey starts with one step.” 的所有上下文向量$z^{(1)}$ 到 $z^{(T)}$。

類似地，這里我們也將從計算一個上下文向量$z^{(2)}$開始進行說明。在下一節中，我們將修改這段代碼以計算所有的上下文向量。

我們首先定義幾個變量：

x_2 = inputs[1] # A second input element
d_in = inputs.shape[1] # B the input embedding size, d=3
d_out = 2 # C the output embedding size, d=2

請注意，在GPT類模型中，輸入維度和輸出維度通常是相同的。但為了更清晰地說明計算過程，這里我們選擇不同的輸入維度（$d_{in}=3$）和輸出維度（$d_{out}=2$）。

接下來，我們初始化圖3.14中展示的三個權重矩陣$W_q$、$W_k$和$W_v$：

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)

請注意，這里我們將requires_grad=False，是為了在輸出中減少干擾，便于演示。但如果我們希望使用這些權重矩陣進行模型訓練，就應將requires_grad=True，以便在訓練過程中更新這些矩陣。

然后我們根據圖3.14的說明，計算查詢（query）、鍵（key）和值（value）向量：

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_valueprint(query_2)
print(key_2)
print(value_2)

如輸出所示，查詢向量是一個二維向量，因為我們通過$d_{out}=2$設置了對應權重矩陣的列數為2：

tensor([0.4306, 1.4551])
tensor([0.4433, 1.1419])
tensor([0.3951, 1.0037])

權重參數 vs 注意力權重

請注意，權重矩陣$W$中的“權重”一詞是“權重參數（weight parameters）”的縮寫，指的是神經網絡中在訓練期間被優化的值。這一點不要與注意力權重（attention weights）混淆。正如我們在前一節中看到的，注意力權重決定了上下文向量在多大程度上依賴輸入中的不同部分，也就是說，網絡關注輸入各部分的程度。

總結來說，權重參數是神經網絡中定義連接關系的基礎可學習系數，而注意力權重是動態的、依賴上下文的值。

盡管我們當前的目標是僅計算一個上下文向量$z^{(2)}$，我們仍然需要所有輸入元素的key和value向量，因為它們在根據查詢$q^{(2)}$計算注意力權重時是必要的，如圖3.14所示。

我們可以通過矩陣乘法獲得所有的key和value：

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

從輸出中可以看出，我們成功地將6個輸入token從三維嵌入空間投影到了二維嵌入空間：

keys.shape: torch.Size([6, 2])
values.shape: torch.Size([6, 2])

接下來的第二步是計算注意力分數，如圖3.15所示。

圖3.15 注意力分數的計算是點積計算，和我們在第3.3節中使用的簡化自注意力機制類似。這里的新特點在于，我們不再直接對輸入元素進行點積運算，而是使用通過權重矩陣變換得到的查詢和鍵（query和key）來進行計算。

首先，我們計算注意力分數${\omega }_{22}$：

keys_2 = keys[1] # A Python starts index at 0
attn_score_22 = query_2.dot(keys_2)
print(attn_score_22)

輸出結果是未歸一化的注意力分數：

tensor(1.8524)

同樣地，我們可以通過矩陣乘法將該計算泛化為所有注意力分數的計算：

attn_scores_2 = query_2 @ keys.T  # 所有針對給定查詢的注意力分數
print(attn_scores_2)

我們可以快速驗證：輸出中的第二個元素與我們之前計算的${\omega }_{22}$一致：

tensor([1.2705, 1.8524, 1.8111, 1.0795, 0.5577, 1.5440])

接下來的第三步是將注意力分數轉換為注意力權重，如圖3.16所示。

圖3.16 在計算出注意力分數$\omega$之后，下一步是使用softmax函數對這些分數進行歸一化，從而得到注意力權重$\alpha$。

接下來，如圖3.16所示，我們通過縮放注意力分數并使用之前介紹的softmax函數來計算注意力權重。

與之前不同的是，我們現在將注意力分數除以鍵向量嵌入維度的平方根進行縮放（注意，開平方在數學上等價于乘方0.5）：

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

得到的注意力權重如下所示：

tensor([0.1500, 0.2264, 0.2199, 0.1311, 0.0906, 0.1820])

縮放點積注意力的原理

將注意力分數按嵌入維度進行歸一化的原因是：通過避免梯度過小來提升訓練效果。例如，當我們增加嵌入維度（在GPT類LLM中通常大于一千）時，較大的點積會在softmax函數作用下導致非常小的梯度。在點積變大時，softmax函數的行為會越來越像階躍函數，進而使梯度趨近于0。這些微小的梯度可能嚴重減緩學習進程，甚至導致訓練停滯。

用嵌入維度平方根進行縮放，正是這種自注意力機制被稱為 縮放點積注意力（scaled-dot product attention） 的原因。

現在，最后一步是計算上下文向量，如圖3.17所示。

圖3.17 在自注意力計算的最后一步，我們通過將所有值向量根據注意力權重進行加權組合，來計算上下文向量。

類似于第3.3節中我們將上下文向量計算為輸入向量的加權和的做法，我們現在將上下文向量計算為**值向量（value vectors）**的加權和。在這里，注意力權重充當加權因子，用于衡量每個值向量的重要性。和第3.3節一樣，我們可以使用矩陣乘法一步獲得輸出：

context_vec_2 = attn_weights_2 @ values
print(context_vec_2)

得到的向量內容如下：

tensor([0.3061, 0.8210])

到目前為止，我們只計算了一個上下文向量$z^{(2)}$。在下一節中，我們將對代碼進行泛化，計算輸入序列中所有的上下文向量，即$z^{(1)}$到$z^{(T)}$。

---

為什么要使用query、key和value？

在注意力機制中，術語“key”（鍵）、“query”（查詢）和“value”（值）源自信息檢索和數據庫領域，在這些領域中有類似的概念被用于存儲、搜索和檢索信息。

“query”（查詢）類似于數據庫中的搜索查詢。它表示當前模型關注或試圖理解的元素（例如句子中的某個詞或token）。查詢用于探查輸入序列中的其他部分，以確定對它們需要關注多少。

“key”（鍵）就像數據庫中的索引鍵，用于查找和匹配。在注意力機制中，輸入序列中的每一項（例如句子中的每個詞）都有一個對應的key。這些key用于與query進行匹配。

“value”（值）在這里類似于數據庫中鍵值對中的value。它表示輸入項的實際內容或表示。當模型確定了哪些key（也就是輸入中哪些部分）與query最相關之后，它就會檢索出對應的value。

3.4.2 實現一個簡潔的自注意力 Python 類

在前幾節中，我們逐步完成了自注意力輸出的計算。這主要是為了說明目的，以便我們能夠一步一步地理解整個過程。而在實踐中，特別是考慮到下一章即將實現的LLM，將這些代碼組織成一個Python類會更加有用，代碼如下：

代碼清單 3.1 一個簡潔的自注意力類

import torch.nn as nnclass 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_keyqueries = x @ self.W_queryvalues = x @ self.W_valueattn_scores = queries @ keys.T # omegaattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn context_vec

在這段PyTorch代碼中，SelfAttention_v1 是繼承自 nn.Module 的類，nn.Module 是PyTorch模型的基本構建模塊，它提供了用于創建和管理模型層所需的功能。

__init__ 方法初始化了用于查詢（query）、鍵（key）和值（value）的可訓練權重矩陣（$W_{query}$、$W_{key}$ 和 $W_{value}$），這些矩陣將輸入維度$d_{in}$變換為輸出維度$d_{out}$。

在前向傳播階段（即forward方法中），我們先通過查詢和鍵計算注意力分數（attn_scores），然后通過softmax對這些分數進行歸一化。最后，我們使用這些歸一化后的注意力分數對值向量進行加權，生成上下文向量。

我們可以如下使用這個類：

torch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
print(sa_v1(inputs))

由于inputs包含六個嵌入向量，因此輸出是一個包含六個上下文向量的矩陣：

tensor([[0.2996, 0.8053],[0.3061, 0.8210],[0.3058, 0.8203],[0.2948, 0.7939],[0.2927, 0.7891],[0.2990, 0.8040]], grad_fn=<MmBackward0>)

可以快速驗證：第二行（[0.3061, 0.8210]）與前一節中context_vec_2的內容一致。

圖3.18總結了我們剛剛實現的自注意力機制。

圖3.18 在自注意力機制中，我們使用三個權重矩陣$W_q$、$W_k$和$W_v$對輸入矩陣$X$中的向量進行變換。然后，我們基于生成的查詢$Q$和鍵$K$計算注意力權重矩陣。接著，結合注意力權重和值$V$，我們計算出上下文向量$Z$。（為了視覺清晰，圖中聚焦于一個具有$n$個token的單個輸入文本，而不是一個包含多個輸入的batch。因此，三維輸入張量在此上下文中簡化為二維矩陣。這種處理方式有助于更直觀地可視化和理解相關過程。）

如圖3.18所示，自注意力機制涉及三個可訓練的權重矩陣$W_q$、$W_k$和$W_v$。這些矩陣將輸入數據轉換為查詢（query）、鍵（key）和值（value），它們是注意力機制中的關鍵組成部分。隨著模型在訓練過程中接觸到更多的數據，它會不斷調整這些可訓練權重，這一點我們將在后續章節中看到。

我們可以通過使用PyTorch的nn.Linear層來進一步改進SelfAttention_v1的實現。nn.Linear在禁用偏置項的情況下可以高效地執行矩陣乘法。此外，使用nn.Linear而不是手動實現nn.Parameter(torch.rand(...))的一個重要優勢是，nn.Linear具備優化過的權重初始化策略，這有助于訓練過程更加穩定和高效。

代碼清單 3.2 使用 PyTorch 的 Linear 層實現自注意力類

class SelfAttention_v2(nn.Module):def __init__(self, d_in, d_out, 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)def forward(self, x):keys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.Tattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=1)context_vec = attn_weights @ valuesreturn context_vec

你可以像使用SelfAttention_v1那樣使用SelfAttention_v2：

torch.manual_seed(789)
sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

輸出結果為：

tensor([[-0.0739,  0.0713],[-0.0748,  0.0703],[-0.0749,  0.0702],[-0.0760,  0.0685],[-0.0763,  0.0679],[-0.0754,  0.0693]], grad_fn=<MmBackward0>)

請注意，SelfAttention_v1和SelfAttention_v2的輸出不同，因為它們使用了不同的初始權重。nn.Linear使用的是更加復雜的權重初始化方案。

---

練習 3.1 對比 SelfAttention_v1 和 SelfAttention_v2

請注意，SelfAttention_v2中的nn.Linear使用的權重初始化方式不同于SelfAttention_v1中使用的nn.Parameter(torch.rand(d_in, d_out))。這導致兩種機制產生了不同的結果。

為了驗證這兩個實現除了初始化不同之外是等價的，我們可以將SelfAttention_v2對象中的權重矩陣轉移到SelfAttention_v1對象中，這樣兩個對象就會輸出相同的結果。

你的任務是：正確地將一個SelfAttention_v2實例中的權重賦值到一個SelfAttention_v1實例中。為此，你需要理解兩個版本中權重的存儲方式之間的關系。（提示：nn.Linear中存儲的權重矩陣是轉置形式。）

完成賦值后，你應該可以觀察到兩個實例產生了相同的輸出。

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)
)d_in, d_out = 3, 2
import torch.nn as nnclass SelfAttention_v1(nn.Module):def __init__(self, d_in, d_out):super().__init__()self.d_out = d_outself.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_keyqueries = x @ self.W_queryvalues = x @ self.W_valueattn_scores = queries @ keys.T # omegaattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)context_vec = attn_weights @ valuesreturn context_vectorch.manual_seed(123)
sa_v1 = SelfAttention_v1(d_in, d_out)
class SelfAttention_v2(nn.Module):def __init__(self, d_in, d_out):super().__init__()self.d_out = d_outself.W_query = nn.Linear(d_in, d_out, bias=False)self.W_key   = nn.Linear(d_in, d_out, bias=False)self.W_value = nn.Linear(d_in, d_out, bias=False)def forward(self, x):keys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.Tattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=1)context_vec = attn_weights @ valuesreturn context_vectorch.manual_seed(123)
sa_v2 = SelfAttention_v2(d_in, d_out)
sa_v1.W_query = torch.nn.Parameter(sa_v2.W_query.weight.T)
sa_v1.W_key = torch.nn.Parameter(sa_v2.W_key.weight.T)
sa_v1.W_value = torch.nn.Parameter(sa_v2.W_value.weight.T)
print(sa_v1(inputs))
print(sa_v2(inputs))

tensor([[-0.5337, -0.1051],[-0.5323, -0.1080],[-0.5323, -0.1079],[-0.5297, -0.1076],[-0.5311, -0.1066],[-0.5299, -0.1081]], grad_fn=<MmBackward0>)
tensor([[-0.5337, -0.1051],[-0.5323, -0.1080],[-0.5323, -0.1079],[-0.5297, -0.1076],[-0.5311, -0.1066],[-0.5299, -0.1081]], grad_fn=<MmBackward0>)

---

在下一節中，我們將對自注意力機制進行增強，重點添加因果性（causal）和多頭機制（multi-head）。

“因果性”方面，涉及修改注意力機制，以防止模型在處理序列時訪問未來的信息。這一點對于語言建模等任務至關重要，因為每一個詞的預測應當只依賴于它之前的詞。

“多頭機制”方面，則是將注意力機制拆分為多個“頭”。每一個頭學習數據的不同方面，使得模型可以在不同的位置同時關注來自不同表示子空間的信息。這種機制能顯著提升模型在復雜任務中的表現。

3.5 使用因果注意力隱藏未來詞

在本節中，我們將對標準自注意力機制進行修改，構建出因果注意力機制（causal attention mechanism），這是在后續章節中開發LLM所必需的。

因果注意力，也被稱為掩碼注意力（masked attention），是一種特殊形式的自注意力機制。它限制模型在處理某個token時只能考慮該序列中當前位置及其之前的輸入。這與標準的自注意力機制形成對比，后者允許一次性訪問整個輸入序列。

因此，在計算注意力分數時，因果注意力機制確保模型只考慮當前token及其之前出現的token。

為了在GPT類LLM中實現這一點，對于每一個被處理的token，我們會屏蔽掉位于當前token之后的所有未來token，如圖3.19所示。

圖3.19 在因果注意力中，我們會屏蔽掉注意力權重矩陣中主對角線以上的部分，這樣在計算上下文向量時，LLM就無法訪問未來的token。例如，在第二行中的單詞“journey”，我們僅保留對當前位置（“journey”）及其之前位置（“Your”）的注意力權重。

如圖3.19所示，我們屏蔽掉注意力權重矩陣中主對角線以上的部分，然后對未被屏蔽的注意力權重進行歸一化，使得每一行的注意力權重之和為1。在下一節中，我們將通過代碼實現這一掩碼與歸一化的過程。

3.5.1 應用因果注意力掩碼

在本節中，我們將通過代碼實現因果注意力掩碼。我們從圖3.20中總結的過程開始。

圖3.20 在因果注意力中，獲得被掩碼的注意力權重矩陣的一種方法是：先對注意力分數應用softmax函數，屏蔽掉主對角線以上的元素（即將其置零），然后對結果矩陣進行歸一化。

為了實現圖3.20中總結的因果注意力掩碼步驟，從而獲得被掩碼的注意力權重，我們將基于上一節的注意力分數和權重來編寫因果注意力機制的代碼。

在圖3.20所示的第一步中，我們像之前幾節那樣使用softmax函數計算注意力權重：

# Reuse the query and key weight matrices of the
# SelfAttention_v2 object from the previous section for convenience
queries = sa_v2.W_query(inputs)
keys = sa_v2.W_key(inputs) 
attn_scores = queries @ keys.Tattn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
print(attn_weights)

這會產生如下注意力權重：

tensor([[0.1921, 0.1646, 0.1652, 0.1550, 0.1721, 0.1510],[0.2041, 0.1659, 0.1662, 0.1496, 0.1665, 0.1477],[0.2036, 0.1659, 0.1662, 0.1498, 0.1664, 0.1480],[0.1869, 0.1667, 0.1668, 0.1571, 0.1661, 0.1564],[0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.1585],[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],grad_fn=<SoftmaxBackward0>)

我們可以使用PyTorch的tril函數實現圖3.20中的第二步，構造一個主對角線以上為零的掩碼矩陣：

context_length = attn_scores.shape[0]
mask_simple = torch.tril(torch.ones(context_length, context_length))
print(mask_simple)

得到的掩碼如下所示：

tensor([[1., 0., 0., 0., 0., 0.],[1., 1., 0., 0., 0., 0.],[1., 1., 1., 0., 0., 0.],[1., 1., 1., 1., 0., 0.],[1., 1., 1., 1., 1., 0.],[1., 1., 1., 1., 1., 1.]])

現在我們可以將這個掩碼與注意力權重相乘，從而將主對角線以上的值置零：

masked_simple = attn_weights*mask_simple
print(masked_simple)

可以看到，主對角線以上的元素已經被成功置零：

tensor([[0.1921, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.2041, 0.1659, 0.0000, 0.0000, 0.0000, 0.0000],[0.2036, 0.1659, 0.1662, 0.0000, 0.0000, 0.0000],[0.1869, 0.1667, 0.1668, 0.1571, 0.0000, 0.0000],[0.1830, 0.1669, 0.1670, 0.1588, 0.1658, 0.0000],[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],grad_fn=<MulBackward0>)

圖3.20中的第三步是重新對每一行進行歸一化，使得每一行的注意力權重之和為1。我們可以通過將每個元素除以其所在行的總和來實現：

row_sums = masked_simple.sum(dim=-1, keepdim=True)
masked_simple_norm = masked_simple / row_sums
print(masked_simple_norm)

最終結果是一個注意力權重矩陣，其中對角線以上的值被置零，并且每一行的總和為1：

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],[0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],[0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],[0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],grad_fn=<DivBackward0>)

---

信息泄露問題（Information leakage）

當我們應用掩碼并重新歸一化注意力權重時，可能初看起來會認為：未來token的信息（我們本意是要屏蔽的）仍然會影響當前token，因為它們的值本來是參與softmax計算的。

然而，關鍵在于：在掩碼之后進行重新歸一化時，實質上就是在一個更小的子集上重新計算softmax，因為被屏蔽的位置不會對softmax的值產生任何貢獻。

softmax的數學優雅性在于：盡管一開始的分母中包含了所有位置，但在掩碼與歸一化之后，被屏蔽的位置在實際計算中完全不會起作用——它們不會以任何有意義的方式影響softmax得分。

換句話說，掩碼與歸一化之后，注意力權重的分布就好像一開始只在未被屏蔽的位置上計算一樣。這確保了我們原本希望屏蔽的未來token的信息不會泄露。

盡管我們在技術上已經完成了因果注意力的實現，但我們還可以利用softmax函數的數學性質，用更少的步驟更高效地計算被掩碼的注意力權重，如圖3.21所示。

圖3.21 在因果注意力中，更高效的一種方式是：在應用softmax之前，將注意力分數中主對角線以上的位置設置為負無窮，從而實現掩碼效果。

softmax函數會將其輸入轉換為一個概率分布。當某一行中存在負無窮（$?\infty$）時，softmax函數會將這些位置視為概率為0（在數學上，這是因為$e^{?\infty }$趨近于0）。

我們可以通過一個更高效的掩碼“技巧”來實現這一點：先創建一個主對角線以上全為1的掩碼，然后將這些1替換為負無窮（$?\infty$）值：

mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)
masked = attn_scores.masked_fill(mask.bool(), -torch.inf)
print(masked)

這將產生如下掩碼后的注意力分數：

tensor([[0.2899,   -inf,   -inf,   -inf,   -inf,   -inf],[0.4656, 0.1723,   -inf,   -inf,   -inf,   -inf],[0.4594, 0.1703, 0.1731,   -inf,   -inf,   -inf],[0.2642, 0.1024, 0.1036, 0.0186,   -inf,   -inf],[0.2183, 0.0874, 0.0882, 0.0177, 0.0786,   -inf],[0.3408, 0.1270, 0.1290, 0.0198, 0.1290, 0.0078]],grad_fn=<MaskedFillBackward0>)

現在我們只需要對這些已掩碼的結果應用softmax函數即可完成：

attn_weights = torch.softmax(masked / keys.shape[-1]**0.5, dim=-1)
print(attn_weights)

從輸出可以看出，每一行的值之和為1，因此不需要進一步歸一化處理：

tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.5517, 0.4483, 0.0000, 0.0000, 0.0000, 0.0000],[0.3800, 0.3097, 0.3103, 0.0000, 0.0000, 0.0000],[0.2758, 0.2460, 0.2462, 0.2319, 0.0000, 0.0000],[0.2175, 0.1983, 0.1984, 0.1888, 0.1971, 0.0000],[0.1935, 0.1663, 0.1666, 0.1542, 0.1666, 0.1529]],grad_fn=<SoftmaxBackward0>)

現在，我們可以使用修改后的注意力權重通過如下方式計算上下文向量，就像在第3.4節中那樣：context_vec = attn_weights @ values 不過，在下一節中，我們將先介紹對因果注意力機制的另一個小改進，該方法在訓練LLM時對減少過擬合非常有幫助。

3.5.2 使用dropout對額外的注意力權重進行掩碼

在深度學習中，dropout是一種技術，它會在訓練過程中隨機忽略（即“丟棄”）部分隱藏層單元。這種方法通過防止模型過度依賴特定的隱藏層單元集合，從而有助于避免過擬合。需要強調的是，dropout只在訓練時啟用，在推理（測試）階段是禁用的。

在Transformer架構中（包括GPT等模型），注意力機制中的dropout通常應用于兩個具體位置：

1. 在計算注意力分數之后；
2. 或在將注意力權重應用于值向量之后。

在這里，我們將在計算完注意力權重之后應用dropout掩碼，如圖3.22所示，因為這在實踐中是更常見的做法。

圖3.22 使用因果注意力掩碼（左上角）之后，我們應用了一個額外的dropout掩碼（右上角），以將更多的注意力權重置零，從而在訓練過程中進一步降低過擬合的風險。

在下面的代碼示例中，我們使用了50%的dropout率，這意味著將屏蔽掉一半的注意力權重。（在后面的章節中訓練GPT模型時，我們將使用更低的dropout率，例如0.1或0.2。）

在以下代碼中，我們首先將PyTorch的dropout實現應用于一個由全1組成的$6\times 6$張量，以便進行說明：

torch.manual_seed(123)
dropout = torch.nn.Dropout(0.5) # A dropout rate of 50%
example = torch.ones(6, 6) # B create a matrix of onesprint(dropout(example))

可以看到，大約一半的值被置零了：

tensor([[2., 2., 2., 2., 2., 2.],[0., 2., 0., 0., 0., 0.],[0., 0., 2., 0., 2., 0.],[2., 2., 0., 0., 0., 2.],[2., 0., 0., 0., 0., 2.],[0., 2., 0., 0., 0., 0.]])

當我們以50%的dropout率將其應用于一個注意力權重矩陣時，矩陣中的一半元素會被隨機設置為0。為了補償有效元素的減少，剩余元素的值會被放大一個因子$1/0.5=2$。這種縮放對于維持注意力權重的整體平衡非常關鍵，它可以確保注意力機制在訓練和推理階段的平均影響力保持一致。

現在，我們將dropout直接應用于注意力權重矩陣：

torch.manual_seed(123)
print(dropout(attn_weights))

結果是一個注意力權重矩陣，其中一些額外的元素被置為0，其余元素被縮放：

tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.0000, 0.8966, 0.0000, 0.0000, 0.0000, 0.0000],[0.0000, 0.0000, 0.6206, 0.0000, 0.0000, 0.0000],[0.5517, 0.4921, 0.0000, 0.0000, 0.0000, 0.0000],[0.4350, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],[0.0000, 0.3327, 0.0000, 0.0000, 0.0000, 0.0000]],grad_fn=<MulBackward0>)

請注意，由于不同操作系統的實現機制不同，最終的dropout輸出可能會略有差異；你可以在這個PyTorch問題跟蹤頁面了解更多關于該不一致性的討論。

在理解了因果注意力和dropout掩碼的工作原理之后，我們將在下一節中開發一個簡潔的Python類，用于高效地結合并應用這兩種技術。

3.5.3 實現一個簡潔的因果注意力類

在本節中，我們將把因果注意力和dropout機制整合進第3.4節中開發的SelfAttention Python類中。該類將作為下一節開發多頭注意力機制的模板，而多頭注意力將是本章中我們實現的最后一個注意力模塊。

但在開始之前，我們還需要確保代碼能夠處理包含多個輸入的批量（batch），以便CausalAttention類可以兼容第2章中實現的數據加載器所產生的批量輸出。

為簡化起見，為了模擬這種批量輸入，我們將輸入文本示例復制兩次：

batch = torch.stack((inputs, inputs), dim=0)
print(batch.shape) # 2 inputs with 6 tokens each, and each token has embedding dimension 3

這將生成一個三維張量，表示兩個輸入文本，每個文本有6個token，每個token是一個三維的嵌入向量：

torch.Size([2, 6, 3])

下面的CausalAttention類與我們之前實現的SelfAttention類相似，不同之處在于我們添加了dropout機制和因果掩碼，這些改動在以下代碼中已高亮顯示：

---

代碼清單 3.3：一個簡潔的因果注意力類

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)) # Newdef forward(self, x):b, num_tokens, d_in = x.shape # New batch dimension b# 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 forward method. keys = self.W_key(x)queries = self.W_query(x)values = self.W_value(x)attn_scores = queries @ keys.transpose(1, 2) # Changed transposeattn_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)attn_weights = self.dropout(attn_weights) # Newcontext_vec = attn_weights @ valuesreturn context_vec

雖然所有新增的代碼行你應該都在前幾節中見過，但這里我們新增了一個self.register_buffer()調用，這是在__init__方法中完成的。PyTorch中的register_buffer并不是所有場景都必須使用，但在這里它有幾個優點：

例如，當我們在LLM中使用CausalAttention類時，緩沖區（buffer）會隨著模型一起自動遷移到合適的設備（CPU或GPU），這在后續訓練LLM時尤為重要。這樣我們就不需要手動確保這些張量和模型參數在同一個設備上，從而避免設備不一致錯誤（device mismatch errors）。

---

我們可以如下使用CausalAttention類，其方式與先前的SelfAttention類似：

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

輸出的上下文向量是一個三維張量，其中每個token現在被表示為一個二維的嵌入向量：

tensor([[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]],[[-0.4519,  0.2216],[-0.5874,  0.0058],[-0.6300, -0.0632],[-0.5675, -0.0843],[-0.5526, -0.0981],[-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

---

圖3.23 總結了我們目前為止所完成的所有工作。

圖3.23 本圖展示了我們在本章中編碼的四種不同注意力模塊的心智模型。我們從一個簡化的注意力機制開始，隨后加入了可訓練權重，然后加入了因果注意力掩碼。在本章的剩余部分中，我們將進一步擴展因果注意力機制，并實現多頭注意力機制（multi-head attention），這將是我們在下一章實現LLM時所使用的最終模塊。

3.6 從單頭注意力擴展到多頭注意力

在本章的最后一節中，我們將把之前實現的因果注意力類擴展為多頭（multi-head）注意力。這一機制也稱為多頭注意力機制。

術語“多頭”指的是將注意力機制劃分為多個**“頭”**，每個頭獨立運行。在這個上下文中，一個單獨的因果注意力模塊可以被視為 單頭注意力（single-head attention），即僅使用一組注意力權重順序地處理輸入。

在接下來的小節中，我們將逐步實現從因果注意力到多頭注意力的擴展。第一個小節將通過堆疊多個CausalAttention模塊來直觀地構建一個多頭注意力模塊，用于說明原理。第二個小節則會以一種更復雜但計算效率更高的方式實現相同的多頭注意力機制。

3.6.1 堆疊多個單頭注意力層

從實踐角度來看，實現多頭注意力就是創建多個自注意力機制的實例（如第3.4.1節中圖3.18所示），每一個都有自己的權重，然后將它們的輸出組合起來。雖然使用多個自注意力機制實例會帶來一定的計算開銷，但這種機制對于Transformer類LLM所擅長的復雜模式識別任務來說至關重要。

圖3.24 展示了一個多頭注意力模塊的結構，它由多個單頭注意力模塊組成，這些模塊就像圖3.18中展示的那樣，被堆疊在一起。

圖3.24 圖中所示的多頭注意力模塊包含兩個堆疊的單頭注意力模塊。因此，相較于單頭注意力機制中僅使用一個值矩陣權重$W_v$，在一個具有兩個頭的多頭注意力模塊中，我們現在有兩個值權重矩陣：$W_{v1}$ 和 $W_{v2}$。對查詢矩陣$W_q$和鍵矩陣$W_k$也是一樣的道理。最終我們會得到兩組上下文向量，分別是$Z_1$和$Z_2$，它們可以組合成一個統一的上下文向量矩陣$Z$。

如前所述，多頭注意力機制背后的主要思想是：使用不同的、可學習的線性投影權重，將注意力機制并行地運行多次。這些線性投影的結果，類似于將輸入數據（如注意力機制中的query、key和value向量）乘以一個權重矩陣。

在代碼中，我們可以通過實現一個簡單的MultiHeadAttentionWrapper類來實現這一點，該類會堆疊多個我們之前實現的CausalAttention模塊實例：

---

代碼清單 3.4 多頭注意力的封裝類

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)

例如，如果我們使用該MultiHeadAttentionWrapper類，并設定兩個注意力頭（即num_heads=2），且CausalAttention的輸出維度為$d_{out}=2$，那么最終的上下文向量維度將是4（$d_{out}\times num\_heads=4$），如圖3.25所示。

圖3.25 使用MultiHeadAttentionWrapper時，我們可以通過參數num_heads來指定注意力頭的數量。如果我們設定num_heads=2（如圖所示），則會得到一個包含兩個上下文向量矩陣的張量。在每個上下文向量矩陣中，行表示對應于token的上下文向量，列表示通過$d_{out}=4$指定的嵌入維度。我們將這些上下文向量矩陣沿著列的維度（即最后一個維度）進行拼接。由于我們有兩個注意力頭，每個頭的嵌入維度為2，因此最終的嵌入維度為$2\times 2=4$。

為了進一步以具體示例說明圖3.25，我們可以像使用CausalAttention類那樣使用MultiHeadAttentionWrapper類：

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)

這會生成如下的張量，表示上下文向量：

tensor([[[-0.4519,  0.2216,  0.4772,  0.1063],[-0.5874,  0.0058,  0.5891,  0.3257],[-0.6300, -0.0632,  0.6202,  0.3860],[-0.5675, -0.0843,  0.5478,  0.3589],[-0.5526, -0.0981,  0.5321,  0.3428],[-0.5299, -0.1081,  0.5077,  0.3493]],[[-0.4519,  0.2216,  0.4772,  0.1063],[-0.5874,  0.0058,  0.5891,  0.3257],[-0.6300, -0.0632,  0.6202,  0.3860],[-0.5675, -0.0843,  0.5478,  0.3589],[-0.5526, -0.0981,  0.5321,  0.3428],[-0.5299, -0.1081,  0.5077,  0.3493]]], grad_fn=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 4])

context_vecs張量的第一個維度是2，因為我們有兩個輸入文本（輸入被復制了兩次，因此它們的上下文向量是完全相同的）。第二個維度表示每個輸入中的6個token。第三個維度表示每個token的4維嵌入向量。

---

練習 3.2：返回二維的嵌入向量

將MultiHeadAttentionWrapper(..., num_heads=2)調用中的輸入參數修改，使得輸出的上下文向量為2維而不是4維，同時保持num_heads=2不變。提示：你不需要修改類的實現，只需更改其它輸入參數中的一個即可。

torch.manual_seed(123)d_out = 1
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)

tensor([[[-0.5740,  0.2216],[-0.7320,  0.0155],[-0.7774, -0.0546],[-0.6979, -0.0817],[-0.6538, -0.0957],[-0.6424, -0.1065]],[[-0.5740,  0.2216],[-0.7320,  0.0155],[-0.7774, -0.0546],[-0.6979, -0.0817],[-0.6538, -0.0957],[-0.6424, -0.1065]]], grad_fn=<CatBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

---

在本節中，我們實現了一個MultiHeadAttentionWrapper，它組合了多個單頭注意力模塊。然而需要注意的是，在forward方法中，這些模塊是順序執行的：[head(x) for head in self.heads] 我們可以通過并行處理多個注意力頭來改進這個實現。其中一種方式是通過矩陣運算同時計算所有注意力頭的輸出，我們將在下一節探索這一點。

3.6.2 使用權重切分實現多頭注意力

在上一節中，我們通過堆疊多個CausalAttention模塊創建了一個MultiHeadAttentionWrapper，從而實現了多頭注意力機制。

現在，我們不再維護兩個獨立的類（MultiHeadAttentionWrapper和CausalAttention），而是將這兩個概念整合進一個統一的MultiHeadAttention類中。除了將兩個類合并，我們還會進行其它修改，以便以更高效的方式實現多頭注意力。

在MultiHeadAttentionWrapper中，多頭注意力是通過創建一個CausalAttention對象列表（self.heads）來實現的，每個對象代表一個注意力頭。每個CausalAttention模塊獨立執行注意力機制，然后將所有頭的輸出拼接起來。

與之相對，下面的MultiHeadAttention類在一個類中集成了多頭機制。它通過reshape已投影的query、key和value張量來切分出多個注意力頭，然后計算注意力后再將結果組合回來。

代碼清單 3.5：一個高效的多頭注意力類

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 forwar'''每個輸出形狀為 (2, 6, 2)（d_out 為 2）。'''keys = self.W_key(x) # Shape: (b, num_tokens, d_out)queries = self.W_query(x)values = self.W_value(x)'''將 keys、values、queries 的最后一維從 d_out（2）重塑成 (num_heads, head_dim)。這里 num_heads=2, head_dim=1，所以原先 (2,6,2) 重塑為 (2,6,2,1)。'''# 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, queries, values 的形狀均為 (2, 2, 6, 1)。這屬于把 “頭” 的維度提前，這樣后續計算不會影響 “頭” 的維度'''keys = keys.transpose(1, 2)queries = queries.transpose(1, 2)values = values.transpose(1, 2)'''keys.transpose(2, 3) 將 keys 的最后兩個維度交換，變為 (b, num_heads, head_dim, num_tokens)(2, 2, 1, 6)。矩陣乘法 queries @ keys^T 結果形狀為 (b, num_heads, num_tokens, num_tokens)。(2, 2, 6, 1) @ (2, 2, 1, 6) -> (2, 2, 6, 6)得到每個“token”之間的關系'''# 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)'''計算注意力輸出。先用 attn_weights 與 values 相乘，得到每個 token 經過加權后的表示。  (2, 2, 6, 6) @ (2, 2, 6, 1) -> (2, 2, 6, 1)把關系還原到對應的 “頭” 維度上，計算出每個 token 的上下文表示。結果形狀為 (b, num_heads, num_tokens, head_dim)；然后轉置維度，將形狀轉換為 (b, num_tokens, num_heads, head_dim)。(2, 2, 6, 1) -> (2, 6, 2, 1)'''# Shape: (b, num_tokens, num_heads, head_dim)context_vec = (attn_weights @ values).transpose(1, 2) '''將多頭的輸出（當前形狀 (b, num_tokens, num_heads, head_dim)）重新合并成最后的輸出，形狀為 (b, num_tokens, d_out)。(2, 6, 2, 1) -> (2, 6, 2)這里 d_out = num_heads * head_dim = 2。'''# Combine heads, where self.d_out = self.num_heads * self.head_dimcontext_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)'''對合并后的輸出再進行一次線性投影，通常用來混合各頭的信息，輸出的形狀仍然為 (b, num_tokens, d_out)。(2, 6, 2)'''context_vec = self.out_proj(context_vec) # optional projectionreturn context_vec

盡管在MultiHeadAttention類中涉及了大量.view和.transpose操作使代碼看起來很復雜，但在數學層面，它實現的邏輯與我們之前的MultiHeadAttentionWrapper是完全一樣的。

從整體上看，在之前的MultiHeadAttentionWrapper中，我們是顯式堆疊多個單頭注意力層來構建多頭注意力層；而在MultiHeadAttention類中，我們采取了更緊湊的實現方式：從一開始就作為一個多頭注意力層來設計，然后在內部通過reshape和轉置操作來切分成多個注意力頭，如圖3.26所示。

圖3.26 在MultiHeadAttentionWrapper類中（上圖），我們初始化了兩個權重矩陣$W_{q1}$和$W_{q2}$，并分別計算出兩個query矩陣$Q_1$和$Q_2$。而在MultiHeadAttention類中（下圖），我們只初始化了一個更大的權重矩陣$W_q$，只對輸入執行一次矩陣乘法來得到一個大的query矩陣$Q$，然后再將其分割成$Q_1$和$Q_2$。對key和value的處理方式也是類似的（未畫出以簡化圖示）。

如圖3.26所示，對query、key和value張量的切分是通過PyTorch中的.view和.transpose方法對張量進行reshape和轉置操作來實現的。輸入首先通過用于query、key和value的線性層進行變換，然后被reshape為代表多個注意力頭的結構。

關鍵操作是將$d_{out}$維度切分為$nu{m}_{heads}$和$hea{d}_{dim}$，其中$hea{d}_{dim}=d_{out}/nu{m}_{heads}$。這種切分是通過.view方法實現的：一個維度為$(b,num\_tokens,d_{out})$的張量被reshape為$(b,num\_tokens,num\_heads,head\_dim)$。

隨后這些張量會進行轉置操作，將$nu{m}_{heads}$維度置于$nu{m}_{tokens}$之前，從而得到形狀為$(b,num\_heads,num\_tokens,head\_dim)$的張量。這種轉置對于正確對齊不同頭之間的query、key和value，并高效地進行批量矩陣乘法是至關重要的。

為了說明這種批量矩陣乘法，假設我們有如下示例張量：

# (1,2,3,4)
a = torch.tensor([[[[0.2745, 0.6584, 0.2775, 0.8573],[0.8993, 0.0390, 0.9268, 0.7388],[0.7179, 0.7058, 0.9156, 0.4340]],[[0.0772, 0.3565, 0.1479, 0.5331],[0.4066, 0.2318, 0.4545, 0.9737],[0.4606, 0.5159, 0.4220, 0.5786]]]])

現在我們對該張量與其最后兩個維度（$num\_tokens$和$head\_dim$）轉置后的版本執行批量矩陣乘法：

print(a @ a.transpose(2, 3)) # (1,2,3,4) @ (1,2,4,3)

結果如下：

# (1,2,3,3)
tensor([[[[1.3208, 1.1631, 1.2879],[1.1631, 2.2150, 1.8424],[1.2879, 1.8424, 2.0402]],[[0.4391, 0.7003, 0.5903],[0.7003, 1.3737, 1.0620],[0.5903, 1.0620, 0.9912]]]])

在這種情況下，PyTorch中的矩陣乘法實現會處理四維張量，在最后兩個維度（$num\_tokens$，$head\_dim$）上進行矩陣乘法，并對每個注意力頭重復該過程。

例如，上述操作實際上是對每個注意力頭分別計算矩陣乘法的更緊湊表達。等價實現如下：

first_head = a[0, 0, :, :]
first_res = first_head @ first_head.T
print("First head:\n", first_res)second_head = a[0, 1, :, :]
second_res = second_head @ second_head.T
print("\nSecond head:\n", second_res)

輸出結果與之前批量矩陣乘法得到的完全一致：

First head:tensor([[1.3208, 1.1631, 1.2879],[1.1631, 2.2150, 1.8424],[1.2879, 1.8424, 2.0402]])Second head:tensor([[0.4391, 0.7003, 0.5903],[0.7003, 1.3737, 1.0620],[0.5903, 1.0620, 0.9912]])

繼續討論MultiHeadAttention：在計算完注意力權重和上下文向量之后，來自所有頭的上下文向量會被轉置回形狀$(b,num\_tokens,num\_heads,head\_dim)$。然后這些向量會被reshape（扁平化）為$(b,num\_tokens,d_{out})$，從而有效地將所有注意力頭的輸出組合起來。

此外，我們在MultiHeadAttention中添加了一個輸出投影層（self.out_proj），用于在合并所有頭之后進一步處理結果向量。這在CausalAttention類中并不存在。這個輸出投影層并非絕對必要（詳見附錄B中的參考資料），但由于它在許多LLM架構中都被廣泛使用，因此我們在這里為了完整性將其加入。

盡管MultiHeadAttention類由于引入了張量reshape和轉置操作，看起來比MultiHeadAttentionWrapper更復雜，但它實際上更高效。原因是我們只需一次矩陣乘法就能計算出keys，例如：

keys = self.W_key(x)

對于queries和values也是同理。而在MultiHeadAttentionWrapper中，我們為每個頭都要重復這一步驟，而這正是計算中最昂貴的一步之一。

MultiHeadAttention類的用法與我們之前實現的SelfAttention和CausalAttention類似：

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

從輸出可以看出，輸出維度直接由$d_{out}$參數控制：

tensor([[[0.3190, 0.4858],[0.2943, 0.3897],[0.2856, 0.3593],[0.2693, 0.3873],[0.2639, 0.3928],[0.2575, 0.4028]],[[0.3190, 0.4858],[0.2943, 0.3897],[0.2856, 0.3593],[0.2693, 0.3873],[0.2639, 0.3928],[0.2575, 0.4028]]], grad_fn=<ViewBackward0>)
context_vecs.shape: torch.Size([2, 6, 2])

在本節中，我們實現了MultiHeadAttention類，它將在后續章節中實現和訓練LLM時被使用。

請注意，雖然代碼是完整可運行的，但我們使用了較小的嵌入維度和注意力頭數量，以便使輸出結果更易讀。

作為對比，最小的GPT-2模型（具有1.17億參數）擁有：

• 12個注意力頭
• 上下文向量的嵌入維度為768

而最大的GPT-2模型（15億參數）擁有：

• 25個注意力頭
• 上下文嵌入維度為1600

需要注意的是，在GPT模型中，token輸入嵌入和上下文嵌入的維度是相同的，即$d_{in}=d_{out}$。

---

練習 3.3 初始化與GPT-2規模一致的注意力模塊

使用MultiHeadAttention類，初始化一個具有與最小GPT-2模型相同數量的注意力頭的多頭注意力模塊（即12個注意力頭）。同時，請確保你使用的輸入和輸出嵌入維度類似GPT-2（即768維）。注意，最小的GPT-2模型支持的上下文長度為1024個token。

上下文長度為1024，輸入輸出嵌入維度設置如下：

context_length = 1024
d_in, d_out = 768, 768
num_heads = 12mha = MultiHeadAttention(d_in, d_out, context_length, 0.0, num_heads)

可選地，可以使用如下函數來統計參數數量：

def count_parameters(model):return sum(p.numel() for p in model.parameters() if p.requires_grad)count_parameters(mha)

輸出結果為：

2360064  #（2.36M）

GPT-2模型總共有1.17億個參數，但正如我們所看到的，多頭注意力模塊本身只占其中的一小部分參數量。

---

3.7 小結

• 注意力機制將輸入元素轉換為增強的上下文向量表示，這些表示整合了關于所有輸入的信息。

• 自注意力機制將上下文向量表示計算為對輸入的加權和。

• 在簡化的注意力機制中，注意力權重是通過點積計算的。

• 點積只是對兩個向量按元素相乘后求和的一種簡潔寫法。

• 盡管不是嚴格必要的，但矩陣乘法能幫助我們更高效且更簡潔地實現這些計算，替代嵌套的for循環。

• 在LLM中使用的自注意力機制（也稱為縮放點積注意力）中，我們引入了可訓練的權重矩陣來對輸入進行中間變換：query、key 和 value。

• 當我們使用從左到右讀取和生成文本的LLM時，需要添加因果注意力掩碼以防止LLM訪問未來的token。

• 除了使用因果掩碼將注意力權重置零，我們還可以添加dropout掩碼，以減少LLM中的過擬合。

• 基于transformer的LLM中的注意力模塊通常包含多個因果注意力模塊的實例，這被稱為多頭注意力。

• 我們可以通過堆疊多個因果注意力模塊的實例來構建一個多頭注意力模塊。

• 更高效的多頭注意力模塊實現方式是使用批量矩陣乘法。