1.2.2 模型架構

在本小節中，我們將介紹Transformer模型的詳細架構，為了更容易理解，我們首先將模型簡化為一個黑盒子，然后一層層深入，并一一介紹相關概念。我們以Transformer最開始的應用——機器翻譯任務為例進行講解。在機器翻譯任務中，它的輸入是一種語言（比如漢語）的一個句子，輸出是表示相同意思的另一種語言（比如英語）的一個句子，如圖1-11所示。

圖1-11 Transformer模型在機器翻譯中的輸入與輸出示意

打開Transformer模型，可以發現它是由一個編碼器組件、一個解碼器組件以及它們之間的連接組成的。其中，編碼器組件由一系列編碼器組成，解碼器組件由一系列相同數量的解碼器組成。

所有的編碼器在結構上都是相同的，如圖1-12左圖所示，它分為兩個部分：自注意力層和前饋神經網絡。編碼器的輸入首先經過自注意力層，對輸入句子中的特定單詞向量計算注意力分數，結合所有單詞向量的注意力分數，編碼為一個新的隱藏狀態向量。自注意力層的輸出被發送到前饋神經網絡，進行線性映射。自注意力層根據輸入句子中不同的單詞得出不同的注意力分數，因此自注意力層的權重參數不同，而前饋神經網絡對輸入句子中不同的單詞應用完全相同的權重參數。

解碼器的結構與編碼器類似，如圖1-12右圖所示，不過分為了3個部分：自注意力層、注意力層和前饋神經網絡。其中自注意力層和前饋神經網絡的結構與編碼器中相同，在這兩層之間是一個“編碼器-解碼器”注意力層，使用來自編碼器壓縮的上下文向量，讓解碼器不僅可以關注輸入句子中的相關單詞，還可以關注輸出句子中的相關單詞。

圖1-12 Transformer模型中編碼器和解碼器的結構細節示意

1.自注意力層

正如前面提到的，編碼器接收向量列表作為輸入，然后將向量列表傳遞給自注意力層進行處理，之后傳遞給前饋神經網絡，最后將輸出發送給下一個編碼器。在編碼器計算的過程中，每個位置的單詞都會經過自注意力層的計算，然后通過一個完全相同的前饋神經網絡。

實現編碼器中的自注意力層需要四個步驟，如圖1-13所示。

圖1-13 自注意力層的計算流程示意

1）為每個單詞創建查詢（query）鍵（key）和值（value）三個向量。對于輸入序列中的每個單詞，將其表征向量分別和對應的權重矩陣（WQ，WK，WV）相乘，映射成查詢向量、鍵向量和值向量。

2）使用查詢向量對其他單詞的鍵向量進行評分。對于每個輸入單詞，使用其對應的查詢向量與其他所有單詞的鍵向量進行內積運算，獲得注意力分數，這個注意力分數就表示當前單詞與其他單詞之間的相關性。

3）對注意力分數進行標準化。每個單詞對于其他單詞的注意力分數除以鍵向量維度的平方根，然后通過softmax函數進行標準化，這樣有助于提高梯度的穩定性。

4）將值向量乘以標準化后的注意力分數然后加權求和。對于每個輸入單詞，將其對應的值向量與上一步得到的標準化后的注意力分數相乘，然后將所有乘積結果相加，得到經過自注意力機制修正后的表征。

為什么需要三個向量來計算注意力分數呢？query、key和value的概念來自信息檢索系統。例如，當我們想要搜索某一個問題時，會打開瀏覽器在搜索欄輸入關鍵字，這個關鍵字就是查詢問題的query，搜索結果可能包含很多文章，這個文章的標題就是key，文章的內容就是value，對搜索結果進行排序時，其實就是用問題query與標題key進行匹配，計算相關度，即注意力分數。

為什么要除以進行縮放呢？假設query向量和key向量中的元素都是相互獨立、均值為0、方差為1的隨機變量，那么這兩個向量的內積的期望為0，方差為向量的維度d_k。因此，當d_k較大時，點積的方差也較大，這樣不同的查詢向量與不同的鍵向量計算出來的分數會相差較大。softmax函數的公式為

對x_i求偏導得

對x_j求偏導得

這里的。當方差較大時，可能會出現某個x_i遠大于或者遠小于其他的情況，進而可能導致梯度消失和梯度爆炸。

以如下代碼為例，我們先創建一個從-10到9的數組x₁，然后在其尾部添加上一個值100，表示可能出現的遠大于其他數字的值，之后x₁將除以100進行縮放，定義為x₂，最后分別繪制兩個數組的softmax值，結果如圖1-14所示。可以發現，不進行縮放的藍色的曲線比較陡峭，在兩端可能會出現梯度接近于0或梯度過大的值；而進行縮放的紅色的曲線相對平緩，梯度的變化不會太劇烈，這樣有助于保持模型訓練的穩定性。

python

import numpy as np

import matplotlib.pyplot as plt

#定義softmax函數

def softmax(x):

　　e_x=np.exp(x-np.max(x))　# for numerical stability

　　return e_x/e_x.sum()

#創建一個從-10到9的數組，并在末尾添加一個值100

x1=np.arange(-10, 10)

x1=np.append(x1, 100)

#計算softmax值

y1=softmax(x1)

#創建另一個從-10到9的數組，在末尾添加一個值100，然后除以100進行縮放

x2=np.arange(-10, 10)

x2=np.append(x2, 100)

x2=x2/100

#計算softmax值

y2=softmax(x2)

plt.figure(figsize=(10, 6))

plt.plot(range(1, len(x1)+1), y1, marker="o", color="blue", label="Original")

plt.plot(range(1, len(x2)+1), y2, marker="o", color="red", label="Scaled")

plt.xlabel("Index")

plt.ylabel("Softmax Value")

plt.title("Softmax Values of Two Vectors")

plt.grid(True)

plt.legend()

plt.show()

圖1-14 關于softmax的輸入是否進行縮放的函數值對比

為了更直觀地理解注意力分數，我們來看一個例子。假如我們要翻譯一句話“The animal didn’t cross the street because it was too tired.”，其中文意思是“動物沒有過馬路，因為它太累了。”，那么在翻譯的時候，“it”指的到底是什么呢，是animal還是street？對于人類來說，我們可以很輕易地意識到“it”指的是動物（animal）。然而，對于算法來說，確定“it”指的是什么可能是一個復雜的問題。引入了注意力機制之后，在處理這個單詞時，模型就可以使“it”與“animal”建立更強的聯系。事實上，如圖1-15所示，我們可以通過谷歌提供的Tensor2Tensor Notebook來交互查看注意力分數，訓練完成后，結果確實如我們前面所述。

圖1-15 注意力機制中“it”的注意力分數可視化示意

最后需要注意的是，每個單詞都會對應查詢、鍵和值三個向量，但是在實際的代碼中使用的是整個輸入序列的矩陣。輸入X矩陣，它的每一行就對應輸入句子中的每一個單詞，然后分別乘以WQ、WK、WV這3個矩陣，得到Q、K、V三個矩陣，接著計算注意力分數矩陣，之后經過softmax函數，再乘以V矩陣，就得到了Z矩陣，這個Z矩陣就是要發送給前饋神經網絡的矩陣。整個流程如圖1-16所示。

圖1-16 自注意力機制的計算流程

2.多頭注意力機制

在Transformer模型的論文中還使用了一種加強注意力的設計——多頭注意力（MHA）機制，進一步完善了注意力層。多頭注意力機制可以讓模型同時關注不同表示子空間的信息。例如，當我們看一篇文章的時候，會更注意標題和粗體的文字，而不是正文那種又小又密集的文字，也會更注意顏色鮮艷的文字，比如紅色的標題。這里的字體和顏色就是兩個表示子空間，我們如果能同時關注字體和顏色，那么就可以更有效地定位文章中強調的內容。同理，多頭注意力機制就是綜合利用各個方面的信息，從多個表示子空間中匯總重要特征。

要使用多頭注意力機制，就要有多組query、key、value權重矩陣。在標準的Transformer模型中使用了8個注意力頭，因此每個編碼器和解碼器的注意力層都有8個權重集合，這些集合中的每一個參數矩陣在開始訓練時都是隨機初始化的。模型訓練完成后，每個權重集合可以將輸入特征向量投影到不同的表示子空間中，模型可以通過Q、K、V在不同的空間去學習特征，從而避免使用同一種注意力機制時可能產生的偏執，讓語義擁有更多元的表達。

多頭注意力要為每個頭維護單獨的WQ、WK、WV權重矩陣，首先用X乘以WQ、WK、WV矩陣以產生Q、K、V矩陣，然后進行與前面提到的注意力機制相同步驟的計算，只是使用的權重矩陣不同，最終會得到8個不同的Z矩陣。但是前饋神經網絡不能處理8個矩陣，因此我們需要再做一個矩陣運算以將這8個矩陣變換成1個矩陣。整個流程如圖1-17所示。

圖1-17 多頭注意力機制流程

在結構上，多頭注意力就是由多個點積注意力模塊組合而成的，可以表示為，多頭中的h個點積注意力模塊是可以并行的，它們之間沒有依賴關系，因此可以進一步提升效率。多頭注意力機制還可以提升模型效果，不過論文中并沒有給出更清晰的理論支持，只是在實際訓練時發現效果更好。這也是AI科研論文的一個特點，研究員經常憑借強烈的科研意識和敏銳性，發現一些新的研究方向，并通過實驗證實其有效性。而這些方向并不一定能夠得到完美的理論支持，這也為后續研究者提供了改進的空間。MQA和GQA就是其中的兩種優化方案。

如圖1-18a所示，多頭注意力機制的每個頭都有自己獨立的查詢（Q）、鍵（K）和值（V），在進行計算時，每個頭都要依據自己的Q、K、V進行計算，這會占用大量的存儲空間，并且其占用空間規模是隨著模型隱藏層維度的增加而成倍增加的。

為了解決這個問題，有一種新提出的改進模型，稱為多查詢注意力（MQA）機制。MQA的設計思想是，讓查詢保持原來的多頭設計，但鍵和值則只有一個頭，如圖1-18c所示。在這種設計中，所有的Q頭共享一組KV頭，因此得名“多查詢”注意力。盡管這種設計在某些情況下可能會對模型性能產生一定的影響，但基于其各種優勢，這種微小的性能降低是可以接受的。實驗發現，MQA模型通常可以提高30%到40%的處理效率。

MQA模型為何能提高處理效率呢？主要原因在于它降低了KV緩存的大小。雖然從運算量上看，MQA和MHA的計算復雜度是差不多的，但由于MQA模型只需要讀取一組KV頭然后供所有Q頭使用，因此，這種設計在內存和計算之間存在不對稱性的情況下具有明顯的優勢。具體來說，MQA模型的這種設計可以減少需要從內存中讀取的數據量，從而縮短了計算單元的等待時間，提高了計算效率。同時，由于KV緩存的大小減小了，顯存中需要保存的張量的大小也相應減小，這為增大批處理量留出了空間，進一步提高了顯存利用率。

除了MQA模型外，還有一種折中的解決方案，即分組查詢注意力（GQA）機制，如圖1-18b所示。GQA是MHA和MQA的一種混合模型，其目的是在不過分損失性能的前提下，盡可能地獲取MQA模型的推理加速優勢。在GQA模型中，不是所有的Q頭共享一組KV，而是分組的一定數量的Q頭共享一組KV。這種設計既兼顧了性能，也考慮了推理速度。

圖1-18 MHA、GQA和MQA的對比

在后面我們將介紹的Llama 2的論文中，給出了三種多頭注意力機制的效果對比，如圖1-19所示，可以發現，GQA的效果相對來說是比較好的，MHA的效果次之，MQA的效果則要差一些。

圖1-19 Llama 2論文中MHA、GQA和MQA在不同任務上的效果

3.注意力層

Transformer模型的解碼器使用的是注意力層，基本的計算步驟與自注意力層相似，只不過其中的鍵矩陣和值矩陣是由編碼器提供的上下文向量。編碼器首先處理輸入序列，然后將頂部編碼器的輸出轉換為一組注意力矩陣——K和V，這兩個矩陣將在每個解碼器的注意力層中使用。

在解碼階段，我們需要給定一個初始的隨機向量，用于表示開始生成序列。類似于我們在編碼器階段對輸入序列進行處理的方式，解碼器在生成序列的過程中也會添加位置編碼向量，用于指示每個單詞在序列中的位置。解碼器每一步生成序列時只輸出一個元素，而該元素會作為下一步的輸入再次輸入到解碼器中。這樣，解碼器會逐步生成完整的序列。如圖1-20所示，在注意力層解碼器每次都會根據輸入元素構建Q矩陣，與編碼器發送過來的K矩陣和V矩陣共同進行注意力計算，這樣有助于解碼器在生成時注意到輸入序列中的元素位置。解碼器不斷重復該生成過程，直到出現一個表示生成結束的特殊符號，表明解碼器已完成輸出。

圖1-20 Transformer模型解碼器輸出步驟示意

解碼器中的自注意力層的操作方式與編碼器中的自注意力層的計算方式略有不同：在解碼器中，自注意力層只允許關注輸出序列中出現較早的位置。這個限制是通過在計算注意力分數時，在應用softmax函數之前，對未來位置的分數設置為負無窮（-inf）來實現的。這樣一來，在softmax步驟中，未來位置的注意力分數會趨于0，從而使得解碼器只關注過去和當前的位置。

輸入矩陣和掩碼矩陣具有相同的維度，在掩碼矩陣中，遮擋位置的值被設置為0。如圖1-21所示，對于單詞0，它只能使用自身的信息；而對于單詞1，它可以使用單詞0和自身的信息。換句話說，掩碼矩陣的作用是限制在注意力計算中使用的信息范圍。當某個位置的掩碼值為0時，該位置的相關信息應被排除在注意力計算之外。因此，掩碼矩陣確保了每個位置只能關注它之前的位置，不會使用未來位置的信息。

圖1-21 輸入矩陣與掩碼矩陣示意

通過輸入矩陣X，我們計算得到了Q、K和V三個矩陣。然后我們將Q矩陣與K的轉置矩陣KT進行矩陣乘法，得到了Q·KT矩陣。接下來，我們對Q·KT矩陣與掩碼矩陣按位相乘，得到了Mask Q·KT矩陣。然后，我們對Mask Q·KT矩陣進行softmax操作，使得Mask Q·KT矩陣的每一行的值相加為1。最后，我們將完成softmax操作的Mask Q·KT矩陣與V矩陣進行矩陣乘法運算，得到最終的輸出矩陣Z。整個流程如圖1-22所示。

圖1-22 注意力機制中的Q、K、V矩陣計算與掩碼操作流程