Transformer 架構深度解析注意力機制

– 11 分鐘閱讀

Transformer 架構在深度學習領域，特別是自然語言處理方面，具有重要地位。其核心元件注意力機制透過計算查詢、鍵和值向量之間的關係，有效捕捉長距離的語義資訊。理解注意力機制的運作原理對於掌握 Transformer 架構至關重要。本文除了詳細解釋注意力機制的不同變體，如點積注意力和多頭自注意力，還將探討其計算複雜度，並提供 PyTorch 實作的程式碼範例，幫助讀者更好地理解和應用 Transformer。此外，文章也涵蓋了自適應注意力機制，闡述其如何最佳化注意力範圍，提升模型效率和效能。

Transformer 架構深度解析：注意力機制與計算複雜度

Transformer 架構自提出以來，不僅革新了深度學習社群，也對自然語言處理（NLP）領域產生了深遠的影響。本文將探討 Transformer 中的核心元件——注意力機制（Attention Mechanism），並分析其計算複雜度和實作細節。

Transformer 類別剖析

PyTorch 中的 Transformer 實作本質上是一個標準的 nn.Module，包含一個定義明確的 forward() 函式。該架構最關鍵的元件是 MultiheadAttention 層，其他層如 Dropout、Linear 和 LayerNorm 在非 Transformer 模型中也很常見。PyTorch 的實作與原論文中的架構規格一致，尤其是標題為「Attention Is All You Need」的論文。

注意力機制的工作原理

注意力機制是一種深度神經網路層，主要功能是學習長距離的「全域」特徵。它扮演著「資訊路由器」的角色，決定輸入序列中哪些嵌入向量對單一輸出向量有貢獻。具體而言，注意力機制透過將隱藏狀態向量轉換為三個獨立的向量：查詢（query）、鍵（key）和值（value），然後計算注意力權重。

點積注意力機制

點積注意力是一種簡單直觀的相似度衡量方法，透過計算兩個向量的點積來評估它們的相似程度。在 Transformer 中，點積注意力的計算過程如下：

  1. 將隱藏狀態向量 $h_i$ 分別乘以矩陣 $W_K$、$W_Q$ 和 $W_V$，得到鍵向量 $k_i$、查詢向量 $q_i$ 和值向量 $v_i$。

    $k_i = W_K h_i$, $q_i = W_Q h_i$, $v_i = W_V h_i$

  2. 計算查詢向量 $q_i$ 與所有鍵向量的點積，並進行 softmax 運算，得到注意力權重。

    $z_i = \text{softmax}(q_i K^T) V$

  3. 為了提高效率，可以將這些向量封裝成矩陣，進行平行計算。

    $Z = \text{softmax}(QK^T) V$

多頭自注意力機制

多頭自注意力是點積注意力的擴充套件，透過多個獨立的注意力頭來捕捉不同的特徵表示。這種機制允許模型平行處理多個時間步，提高了計算效率。

計算複雜度分析

在評估注意力機制的效率時，計算複雜度是一個重要的指標。雖然「Big Oh」符號在理論分析中很有用，但在深度學習中，實際效能往往受到多種因素的影響，如程式語言、CPU/GPU 架構和平行度等。

注意力機制的計算複雜度

注意力機制的核心運算是矩陣乘法和 softmax 運算。這些運算的計算複雜度取決於輸入序列的長度和向量的維度。

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        
        self.query_linear = nn.Linear(embed_dim, embed_dim)
        self.key_linear = nn.Linear(embed_dim, embed_dim)
        self.value_linear = nn.Linear(embed_dim, embed_dim)
        
        self.dropout = nn.Dropout(0.1)
    
    def forward(self, query, key, value):
        # 分別計算查詢、鍵和值的線性變換
        Q = self.query_linear(query)
        K = self.key_linear(key)
        V = self.value_linear(value)
        
        # 將 Q、K、V 分割成多個頭
        Q = Q.view(-1, Q.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(-1, K.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(-1, V.size(1), self.num_heads, self.head_dim).transpose(1, 2)
        
        # 計算注意力得分
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention_weights = F.softmax(scores, dim=-1)
        
        # 應用 dropout
        attention_weights = self.dropout(attention_weights)
        
        # 計算最終輸出
        output = torch.matmul(attention_weights, V).transpose(1, 2).contiguous().view(-1, query.size(1), self.embed_dim)
        
        return output

#### 程式碼內容解密：
1. **`MultiHeadAttention` 類別初始化**：定義多頭注意力層的引數，包括嵌入維度 (`embed_dim`) 和注意力頭數 (`num_heads`)。
2. **線性變換**：使用 `nn.Linear` 將輸入的查詢 (`query`)、鍵 (`key`) 和值 (`value`) 分別進行線性變換，以生成對應的查詢 (`Q`)、鍵 (`K`) 和值 (`V`) 向量。
3. **多頭分割**：將變換後的 `Q`、`K` 和 `V` 向量按頭數進行分割和重新排列，以便進行平行計算。
4. **注意力得分計算**：透過矩陣乘法計算 `Q` 和 `K` 的點積，並縮放結果以獲得注意力得分。
5. **Softmax 與 Dropout**：對注意力得分進行 Softmax 運算，以獲得注意力權重，並應用 Dropout 以防止過擬合。
6. **最終輸出**：將注意力權重與 `V` 相乘，得到最終輸出，並重新調整形狀以匹配原始輸入維度。

## 為什麼是三個向量？注意力機制的直觀理解

在閱讀 Transformer 論文時，我們不禁會問：為什麼要將輸入拆分成三個向量（查詢、鍵和值）？這種做法看似有些武斷和複雜。難道不能使用兩個或四個向量嗎？

### 從資料函式庫的角度理解注意力機制

從命名上來看，這種做法似乎與資料函式庫的概念有關。讓我們考慮一下如何在普通的 Python 字典中實作類別似的功能，而不涉及神經網路。

假設我們有一個句子對應的字典結構：
```python
sentence = {
    "word_1": "Squirtle",
    "word_2": "is",
    "word_3": "the",
    "word_4": "greatest",
    "word_5": "Pokemon",
    "word_6": "ever"
}

當我們想要從字典中取得值時，會使用類別似 sentence['word_3'] 的查詢。Python 會在後台將查詢 word_3 與字典中的所有鍵進行比較，然後傳回對應的值。

點積注意力機制的運作原理

點積注意力機制的工作原理與此類別似。查詢向量代表當前詞正在尋找什麼，而每個詞對應的鍵向量代表該詞能夠提供什麼。值向量則包含查詢向量正在尋找的資訊。

讓我們考慮以下句子：

Mario is short, but he can jump super high.

當 Transformer 處理到「he」這個詞時，它的查詢向量可能正在尋找一個名字或人物，以明確「he」指的是誰。Transformer 會計算「he」的查詢向量與句子中其他詞的鍵向量之間的點積，從而得到一個對齊分數，用於衡量查詢和鍵之間的匹配程度。

程式碼範例：計算點積

import numpy as np

# 生成兩個隨機向量
vector1 = np.random.randn(10)
vector2 = np.random.randn(10)

# 計算點積
dot_product = np.dot(vector1, vector2)
print(dot_product)

內容解密：

  1. np.random.randn(10) 用於生成一個包含 10 個隨機元素的向量，這些元素來自標準常態分佈。
  2. np.dot(vector1, vector2) 計算兩個向量的點積，衡量它們之間的相似度。
  3. 輸出結果為一個純量值，代表兩個向量之間的對齊程度。

自注意力機制的運作流程

在計算自注意力時，Transformer 會使用當前層的隱藏狀態來計算查詢、鍵和值向量。這些向量都來自同一層的序列。

對於解碼器中的注意力層，查詢向量來自解碼器的隱藏表示，而鍵和值向量則來自編碼器的隱藏表示。這使得解碼器能夠關注編碼器的所有先前隱藏表示。

縮放點積注意力機制

當向量維度較大時，點積可能會變得很大。為瞭解決這個問題，研究人員提出了縮放點積注意力機制。

程式碼範例：比較不同維度的點積

import numpy as np

# 生成兩個隨機向量（維度 10）
small_dots = [np.dot(np.random.randn(10), np.random.randn(10)) for _ in range(100)]
print(np.mean(np.absolute(small_dots)))

# 生成兩個隨機向量（維度 10000）
large_dots = [np.dot(np.random.randn(10000), np.random.randn(10000)) for _ in range(100)]
print(np.mean(np.absolute(large_dots)))

內容解密：

  1. 程式碼比較了不同維度（10 和 10000）向量的點積平均值。
  2. 結果顯示，當維度增大時，點積的平均值也會顯著增加。
  3. 這說明需要對點積進行縮放，以避免因維度過大而導致的數值問題。

多頭注意力機制與自適應注意力範圍

在前面的章節中，我們探討了注意力機制的基礎概念及其運作原理。現在，讓我們進一步深入瞭解多頭注意力機制（Multi-Head Self-Attention）以及自適應注意力範圍（Adaptive Attention Span）這兩個重要的概念。

多頭注意力機制

多頭注意力機制是 Transformer 架構中的一個關鍵組成部分。這一機制允許模型在處理輸入序列時，能夠同時關注不同的特徵和模式，從而提升模型的表現。

在傳統的注意力機制中，模型使用一套查詢（Query）、鍵（Key）和值（Value）矩陣來計算注意力權重。然而，這種方式可能無法有效地捕捉輸入序列中的多樣化特徵。因此，多頭注意力機制透過使用多套不同的權重矩陣來轉換輸入序列，從而產生多組查詢、鍵和值：

$$ Q_i = W_i^Q x, K_i = W_i^K x, V_i = W_i^V x $$

其中，$i$ 代表不同的頭（head），$W_i^Q$、$W_i^K$ 和 $W_i^V$ 是各自獨立的權重矩陣。

每個頭獨立地進行縮放點積注意力計算：

$$ \text{head}_i = \text{softmax}\left(\frac{Q_i K_i^T}{\sqrt{d_k}}\right) V_i $$

最後，將所有頭的輸出拼接起來，並透過另一個線性變換來調整維度：

$$ z = W^O [\text{head}_1; \text{head}_2; …; \text{head}_n] $$

這種多頭注意力機制使模型能夠同時捕捉不同的特徵和模式，從而提升了模型的表現。

內容解密：

  1. 多頭注意力機制的必要性：單一的注意力機制可能無法有效地捕捉輸入序列中的多樣化特徵，因此需要使用多頭注意力機制。
  2. 多套權重矩陣的使用：透過使用多套不同的權重矩陣，模型能夠產生多組查詢、鍵和值，從而實作多頭注意力機制。
  3. 獨立的注意力計算：每個頭獨立地進行縮放點積注意力計算，從而能夠捕捉不同的特徵和模式。
  4. 輸出的拼接和線性變換：將所有頭的輸出拼接起來，並透過另一個線性變換來調整維度，從而得到最終的輸出。

自適應注意力範圍

自適應注意力範圍是一種旨在降低 Transformer 模型計算成本的方法。在傳統的注意力機制中，計算注意力權重需要進行 $n^2$ 次點積運算，其中 $n$ 是輸入序列的長度。這種計算成本會隨著序列長度的增加而迅速增長。

自適應注意力範圍透過允許模型學習哪些 token 之間需要計算注意力權重，從而減少了不必要的計算。這種方法可以根據輸入序列的特性動態地調整注意力範圍，從而降低計算成本。

內容解密：

  1. 計算成本問題：傳統的注意力機制需要進行 $n^2$ 次點積運算，這種計算成本會隨著序列長度的增加而迅速增長。
  2. 自適應注意力範圍的提出：為了降低計算成本，自適應注意力範圍被提出，允許模型學習哪些 token 之間需要計算注意力權重。
  3. 動態調整注意力範圍：自適應注意力範圍可以根據輸入序列的特性動態地調整注意力範圍，從而降低計算成本。

自適應注意力機制（Adaptive Attention Span）與純注意力架構（All-Attention）

在深度學習的 Transformer 架構中，注意力機制（Attention Mechanism）扮演著至關重要的角色。然而，傳統的注意力機制存在計算效率和記憶體佔用的問題。為瞭解決這些問題，研究人員提出了自適應注意力機制和純注意力架構，分別針對注意力範圍的調整和模型架構的簡化進行了創新。

自適應注意力機制

傳統的 Transformer 模型在計算注意力時，需要對序列中的每個 token 進行全域的注意力計算。這種做法在處理長序列時，會導致計算量和記憶體佔用大幅增加。為瞭解決這個問題，Facebook 的研究團隊提出了自適應注意力機制（Adaptive Attention Span）。

自適應注意力機制的核心思想是允許不同的注意力頭（Attention Head）學習不同的注意力範圍（Attention Span）。這樣一來，部分注意力頭可以專注於區域性資訊，而其他頭則可以關注更廣泛的上下文資訊。

實作方法

為了實作自適應注意力機制，研究團隊引入了一個稱為「遮罩函式」（Masking Function）的機制。這個函式根據 token 之間的距離，輸出一個介於 0 和 1 之間的值，用於控制注意力權重的計算。

具體來說，遮罩函式定義如下：

[ m_z(x) = \min\left(\max\left(\frac{1}{R}(R + z - x), 0\right), 1\right) ]

其中，$z$ 是可學習的引數，控制著注意力的範圍；$R$ 是一個超引數，用於調整遮罩函式的平滑度；$x$ 是 token 之間的距離。

優勢

自適應注意力機制的優勢在於，它能夠根據任務需求自動調整注意力範圍，從而在保持模型效能的同時，減少不必要的計算和記憶體佔用。實驗結果表明，自適應注意力機制在 enwik8 資料集上達到了當時的最佳效能，同時減少了記憶體和 FLOPs 的消耗。

純注意力架構

另一個值得注意的研究是純注意力架構（All-Attention），由 FAIR 團隊提出。這個架構的目標是簡化 Transformer 模型，去掉傳統架構中的位置前饋網路（Position-wise Feed-Forward Network），並將其替換為純粹的注意力機制。

實作方法

在原始的 Transformer 架構中，每個層都包含一個自注意力機制和一個位置前饋網路。位置前饋網路對每個 token 的向量表示進行非線性變換。然而，研究人員發現，這些前饋網路佔用了模型的大部分引數。

純注意力架構的核心思想是，將位置前饋網路替換為額外的注意力層。這樣一來，整個模型就變成了純粹的注意力機制堆積疊。

理論依據

研究人員透過數學推導，證明瞭位置前饋網路的計算可以近似為一種特殊的注意力機制。他們指出，當去掉偏置項並將啟用函式替換為 softmax 時，前饋網路的計算就變成了加權求和的形式，這與注意力的計算非常相似。

優勢

純注意力架構的優勢在於簡化了模型的架構，使得整個模型更加統一和簡潔。雖然這個改變並沒有帶來顯著的效能提升或計算效率的提高，但它為 Transformer 架構的進一步簡化和創新提供了新的思路。

程式碼實作

以下是一個簡化的自適應注意力機制的程式碼範例，使用 PyTorch 實作：

import torch
import torch.nn as nn
import torch.nn.functional as F

class AdaptiveAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, attention_span):
        super(AdaptiveAttention, self).__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.attention_span = attention_span
        self.query_key_value = nn.Linear(hidden_size, hidden_size * 3)
        self.mask_func = nn.Parameter(torch.randn(num_heads,))

    def forward(self, x):
        qkv = self.query_key_value(x)
        q, k, v = qkv.chunk(3, dim=-1)
        attention_scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.hidden_size)
        
        # 計算遮罩函式
        distance = torch.arange(x.size(1), device=x.device).unsqueeze(0) - torch.arange(x.size(1), device=x.device).unsqueeze(1)
        mask = torch.clamp(1 - distance / self.attention_span, 0, 1)
        mask = mask.unsqueeze(0).unsqueeze(0)  # 增加頭部和批次維度
        
        # 將遮罩函式應用到注意力分數
        attention_scores = attention_scores * mask + (-1e9) * (1 - mask)
        attention_weights = F.softmax(attention_scores, dim=-1)
        
        output = torch.matmul(attention_weights, v)
        return output

內容解密：

  1. AdaptiveAttention類別定義：我們定義了一個名為AdaptiveAttention的類別，繼承自nn.Module，用於實作自適應注意力機制。
  2. __init__方法：初始化函式設定了隱藏層大小、注意力頭數量和注意力範圍等超引數，並建立了用於計算查詢、鍵和值的線性層，以及用於遮罩函式的引數。
  3. forward方法：前向傳播函式首先透過線性層計算查詢、鍵和值，然後計算注意力分數。
  4. 遮罩函式計算：根據 token 之間的距離計算遮罩函式，用於控制注意力範圍。
  5. 應用遮罩函式：將遮罩函式應用到注意力分數上，以實作自適應注意力範圍。
  6. 輸出計算：最後，根據注意力權重和值計算輸出。

這個範例展示瞭如何實作自適應注意力機制，並將其應用到 Transformer 模型中，以提高模型的效率和效能。

玄貓 BlackCat

玄貓 BlackCat

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。