你瞭解幾種Attention機制？

整理的初衷

在現代自然語言處理（NLP）和機器學習領域，Transformer架構已成為模型設計和大規模語言模型（LLM）微調的關鍵工具。自其問世以來，Transformer迅速主導了機器翻譯、文字生成、分類和問答系統等各類NLP任務。然而，面對各種變種和改進，如何為特定任務選擇合適的Transformer架構成為許多研究人員和工程師的困擾。

在模型設計中，我們不僅需理解Transformer的基本原理，還要了解不同變種的優劣及其適用場景。任務需求和資料特點會影響架構選擇：處理長序列文字時，可能傾向於高效記憶體管理的變種；實時性要求高的任務，則需要計算速度更快的版本。

爲了幫助大家更好地選擇合適的Transformer架構，我們有必要回顧其發展歷程。從最初的原始Transformer到BERT、GPT、RoBERTa、T5等改進版本，每個變種都解決了特定問題，並在特定場景下表現突出。透過回顧這些架構的演進，我們能更好地理解它們的優勢和適用場景，從而在實際專案中做出更明智的選擇。

在本篇技術部落格中，我們將深入探討Transformer架構的演進歷程，解析各個變種的核心思想和適用場景，幫助大家在複雜任務中高效地選擇最合適的Transformer模型。希望本文能為大家在模型設計和LLM微調中提供有價值的參考。

重溫下Transformer的架構圖

Transformer 模型自從2017年提出以來，迅速成為自然語言處理（NLP）領域的主流模型，憑藉其強大的效能和靈活的結構，推動了多個領域的進步。我們一起來回顧下Transformer的發展歷史及其關鍵論文：

1. Transformer的提出（2017）：

• 關鍵論文：Vaswani, A., et al., "Attention is All You Need," 2017.

• 貢獻：Transformer模型首次提出，完全基於注意力機制（Attention Mechanism），摒棄了傳統的迴圈神經網路（RNN）結構。Transformer模型的核心創新在於自注意力機制（Self-Attention），能夠更好地捕捉長距離的依賴關係。

2. BERT的誕生（2018）：

• 關鍵論文：Devlin, J., et al., "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding," 2018.

• 貢獻：BERT（Bidirectional Encoder Representations from Transformers）模型透過雙向訓練來理解上下文，並在多個NLP任務上取得了顯著的效能提升。BERT的預訓練-微調（Pre-training and Fine-tuning）正規化成爲了後續大多數NLP模型的標準流程。

3. GPT 系列的發展（2018-2020）：

• Radford, A., et al., "Improving Language Understanding by Generative Pre-Training," 2018.（GPT-1）

• Radford, A., et al., "Language Models are Unsupervised Multitask Learners," 2019.（GPT-2）

• Brown, T., et al., "Language Models are Few-Shot Learners," 2020.（GPT-3）

• 關鍵論文：

• 貢獻：OpenAI推出的GPT系列模型，從GPT-1到GPT-3，展示了生成式預訓練模型在文字生成和理解任務上的強大能力，尤其是GPT-3，憑藉其1750億引數，展現了少樣本學習（Few-shot Learning）的驚人能力。

4. Transformer在影象處理上的應用（2020）：

• 關鍵論文：Dosovitskiy, A., et al., "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale," 2020.

• 貢獻：Vision Transformer (ViT) 模型首次將Transformer架構引入計算機視覺領域，透過將影象分割成固定大小的塊（類似於單詞的處理方式），並使用自注意力機制來進行影象分類任務，取得了與傳統摺積神經網路（CNN）相媲美的效能。

5. 多模態Transformer的提出（2021）：

• 關鍵論文：Radford, A., et al., "Learning Transferable Visual Models From Natural Language Supervision," 2021.

• 貢獻：CLIP（Contrastive Language–Image Pretraining）模型能夠透過自然語言監督進行視覺模型的訓練，將文字和影象的表示空間對齊，增強了多模態任務的處理能力。

6. 高效Transformer的研究（2020-2021）：

• Tay, Y., et al., "Efficient Transformers: A Survey," 2020.

• Kitaev, N., et al., "Reformer: The Efficient Transformer," 2020.

• 關鍵論文：

• 貢獻：針對Transformer的計算和記憶體效率問題，提出了多種變體和最佳化方法，如Reformer、Linformer、Performer等，顯著降低了計算複雜度和記憶體佔用。

7. 最新Transformer模型（2022及以後）：

• 關鍵論文：Chowdhery, A., et al., "PaLM: Scaling Language Modeling with Pathways," 2022.

• 貢獻：PaLM（Pathways Language Model）是Google提出的大規模語言模型，透過更高效的架構和訓練策略，進一步提升了Transformer模型在語言理解和生成任務上的效能。

Transformer模型的發展史展示了其在NLP、計算機視覺和多模態任務中的廣泛應用和持續創新。透過不斷的改進和最佳化，Transformer模型將繼續推動人工智慧領域的前沿研究。

幾個關鍵的演進

在 Transformer 模型中，Attention 機制是一個關鍵的組成部分，極大地提升了模型在自然語言處理任務中的表現。

1. 自注意力機制（Self-Attention）

介紹

Transformer 模型最初引入了自注意力機制，這是一種能夠在編碼器和解碼器中捕捉序列內部依賴關係的方法。在自注意力機制中，每個輸入序列元素都會與其他元素進行互動，並根據其重要性進行加權求和，從而生成新的表示。

原因和思考

傳統 RNN 和 LSTM 模型在處理長依賴關係時表現不佳，因為它們需要逐步地處理序列，容易導致梯度消失或爆炸問題。自注意力機制能夠並行化處理序列中的所有元素，且可以直接建模任意長度的依賴關係，從而解決了 RNN 的一些侷限性。

2. 多頭自注意力（Multi-Head Self-Attention）

介紹

多頭自注意力是一種改進，它透過引入多個並行的注意力頭（Attention Heads）來捕捉不同的特徵子空間。每個注意力頭都獨立地計算注意力權重，並將結果進行拼接和線性變換。

原因和思考

單一的自注意力機制可能無法充分捕捉到序列中的多種不同依賴關係。多頭自注意力允許模型在不同的子空間中關注不同的資訊，從而增強模型的表達能力和魯棒性。

3. 位置編碼（Positional Encoding）

介紹

由於 Transformer 模型沒有像 RNN 那樣的順序處理能力，需要引入位置編碼來提供序列中的位置資訊。這些編碼被加到輸入的嵌入向量中，允許模型識別輸入中的順序資訊。

原因和思考

自注意力機制本質上是無序的，即它並不考慮輸入序列的順序。爲了讓模型理解序列的順序資訊，必須顯式地新增這些位置資訊。位置編碼解決了這一問題，使得 Transformer 可以處理和理解順序依賴。

4. 縮放點積注意力（Scaled Dot-Product Attention）

介紹

縮放點積注意力透過計算輸入的點積來衡量不同元素之間的相似性，並根據這些相似性進行加權求和。爲了避免點積值過大導致梯度消失，加入了縮放因子。

原因和思考

點積注意力是計算效率較高且易於理解的一種注意力機制。然而，在高維向量的點積計算中，可能會導致數值問題。引入縮放因子可以穩定梯度，防止數值不穩定，從而提高計算的穩定性和模型的訓練效果。

5. 層歸一化（Layer Normalization）

介紹

層歸一化在每個注意力子層和前饋神經網路子層之後進行，以對輸入進行歸一化處理，減少訓練中的內部協變數偏移。

原因和思考

層歸一化幫助模型更快地收斂，提高了訓練的穩定性和效率。透過標準化每一層的輸入，使得模型可以在更穩定的環境中學習特徵，提高了模型的泛化能力。

6. 注意力掩碼（Attention Masking）

介紹

在解碼器中，爲了保證自迴歸性質，需要引入注意力掩碼，遮蔽未來詞的注意力權重，只允許模型關注已經生成的詞。

原因和思考

解碼器的自迴歸性質要求在生成時僅依賴於已經生成的部分，而不應考慮未來的資訊。注意力掩碼確保了這一點，從而使得模型在解碼過程中保持正確的資訊流。

透過這些技術的逐步演進，Attention機制在Transformer模型中的表現得到了顯著提升，使其成為自然語言處理任務中的一個強大工具。每一個改進都是爲了克服之前方法的不足，進一步增強模型的表現和穩定性。

接下來，我們針對基礎的Multi-Head Attention、Multi-Query Attention、Grouped-Query Attention進行展開說明。

多頭注意力（MHA）

多頭注意力機制（Multi-Head Attention）是Transformer模型中的核心元件之一。它透過並行多個注意力頭來捕捉不同的特徵和關係，從而增強模型的表達能力。每個注意力頭都有自己的查詢（Query）、鍵（Key）和值（Value）矩陣，並分別計算注意力，然後將各個頭的輸出拼接並透過線性變換得到最終的輸出。

機制詳解

1. 輸入線性變換

首先，對輸入進行線性變換以得到查詢（Q）、鍵（K）和值（V）：

$Q = X*W_q$Q=X∗Wq
$K = X*W_k$K=X∗Wk
$V = XW_v$V=XWv

其中，( $W_q$Wq )，( $W_k$Wk )，( $W_v$Wv ) 是可學習的權重矩陣。

2. 計算注意力

對每個頭，計算注意力：

其中，( $d_k$dk) 是鍵的維度。

3. 多頭拼接

將所有頭的輸出拼接起來：

$\text{Concat}(head_1, head_2, \ldots, head_h)$Concat(head1,head2,…,headh)

4. 輸出線性變換

最後，再透過一個線性變換得到最終輸出：

$\text{Output} = \text{Concat}(head_1, head_2, \ldots, head_h) W_o$Output=Concat(head1,head2,…,headh)Wo

程式碼示例

以下是一個簡化的 PyTorch 實現：

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

class MultiQueryAttention(nn.Module):
def **init**(self, embed\_dim, num\_heads):
super(MultiQueryAttention, self).**init**()
self.embed\_dim = embed\_dim
self.num\_heads = num\_heads
self.head\_dim = embed\_dim // num\_heads

        assert self.head_dim * num_heads == embed_dim, "Embedding dimension must be divisible by number of heads"
        
        # Linear layers to generate queries, keys, and values
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, self.head_dim)  # Shared keys
        self.v_linear = nn.Linear(embed_dim, self.head_dim)  # Shared values
        self.out_linear = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, x):
        batch_size, seq_len, embed_dim = x.size()
        
        # Generate queries, keys, values
        queries = self.q_linear(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = self.k_linear(x).view(batch_size, seq_len, self.head_dim).transpose(0, 1)
        values = self.v_linear(x).view(batch_size, seq_len, self.head_dim).transpose(0, 1)
        
        # Calculate attention scores
        attn_scores = torch.matmul(queries, keys.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = F.softmax(attn_scores, dim=-1)
        
        # Apply attention to values
        attn_output = torch.matmul(attn_weights, values)
        
        # Concatenate heads and pass through final linear layer
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        output = self.out_linear(attn_output)
        
        return output

# Example usage

batch\_size = 2
seq\_len = 5
embed\_dim = 16
num\_heads = 4

x = torch.randn(batch\_size, seq\_len, embed\_dim)
mqa = MultiQueryAttention(embed\_dim, num\_heads)
output = mqa(x)

print(output.shape)  # Output shape should be (batch\_size, seq\_len, embed\_dim)

程式碼解釋

1. 初始化：

• embed_dim 是輸入和輸出的嵌入維度。

• num_heads 是注意力頭的數量。

• head_dim 是每個頭的維度，它是 embed_dim 除以 num_heads。

2. 線性變換：

• q_linear、k_linear 和 v_linear 分別生成查詢、鍵和值。

• fc_out 是最後的線性層。

3. 前向傳播：

• 首先對 x 進行線性變換得到查詢、鍵和值，並將其拆分成多個頭。

• 然後進行縮放點積注意力計算。

• 最後將各個頭的輸出拼接並透過線性層得到最終輸出。

這個實現是多頭注意力機制的基礎版本，實際使用中還可能涉及到掩碼（masking）等其他操作。

多Query注意力（MQA）

Multi-Query Attention (MQA) 是一種改進的注意力機制，旨在提升計算效率並減少記憶體需求。與標準的多頭自注意力機制（Multi-Head Self-Attention, MHSA）不同，MQA 透過共享所有注意力頭的 Keys 和 Values 進行計算，僅為每個查詢計算獨立的注意力權重。這樣做不僅可以減少計算複雜度，還可以降低記憶體佔用。

機制解釋

1. 將輸入序列對映到查詢（Query）、鍵（Key）和值（Value）。

2. 共享所有注意力頭的鍵和值。

3. 為每個查詢計算獨立的注意力權重，然後應用這些權重到共享的值上，得到輸出。

程式碼示例

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

class MultiQueryAttention(nn.Module):
def **init**(self, embed\_dim, num\_heads):
super(MultiQueryAttention, self).**init**()
self.embed\_dim = embed\_dim
self.num\_heads = num\_heads
self.head\_dim = embed\_dim // num\_heads

        assert self.head_dim * num_heads == embed_dim, "Embedding dimension must be divisible by number of heads"
        
        # Linear layers to generate queries, keys, and values
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, self.head_dim)  # Shared keys
        self.v_linear = nn.Linear(embed_dim, self.head_dim)  # Shared values
        self.out_linear = nn.Linear(embed_dim, embed_dim)
        
    def forward(self, x):
        batch_size, seq_len, embed_dim = x.size()
        
        # Generate queries, keys, values
        queries = self.q_linear(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = self.k_linear(x).view(batch_size, seq_len, self.head_dim).transpose(0, 1)
        values = self.v_linear(x).view(batch_size, seq_len, self.head_dim).transpose(0, 1)
        
        # Calculate attention scores
        attn_scores = torch.matmul(queries, keys.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn_weights = F.softmax(attn_scores, dim=-1)
        
        # Apply attention to values
        attn_output = torch.matmul(attn_weights, values)
        
        # Concatenate heads and pass through final linear layer
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        output = self.out_linear(attn_output)
        
        return output

# Example usage

batch\_size = 2
seq\_len = 5
embed\_dim = 16
num\_heads = 4

x = torch.randn(batch\_size, seq\_len, embed\_dim)
mqa = MultiQueryAttention(embed\_dim, num\_heads)
output = mqa(x)

print(output.shape)  # Output shape should be (batch\_size, seq\_len, embed\_dim)

請注意以下幾點：

1. embed_dim 必須是 num_heads 的整數倍。

2. q_linear 用於生成查詢；k_linear 和 v_linear 用於生成共享的鍵和值。

3. 透過對 queries, keys 和 values 進行維度調整和矩陣乘法來計算注意力分數和加權值。

4. 最終的輸出連線並透過一個線性層進行對映。

這種實現方式有效地減少了計算複雜度，使得注意力機制在處理長序列時更加高效。

分組注意力（GA）

Grouped Attention是一種改進注意力機制的方法，旨在提升處理長序列或高維資料時的計算效率和效果。其核心思想是將輸入資料劃分成若干組，然後在每一組內分別應用注意力機制，再將結果合併以獲得最終的輸出。這種方法能夠減少計算複雜度，同時保持或提升模型的效能。

機制詳解

1. 輸入劃分

將輸入序列或資料劃分成若干組。例如，對於一個長度為N的序列，可以將其劃分成G組，每組包含N/G個元素。

2. 組內注意力計算

在每一組內分別應用標準注意力機制（如自注意力）。這涉及計算查詢（Query）、鍵（Key）和值（Value）的投影，然後基於查詢和鍵之間的相似性來加權求和值。這一階段的計算複雜度較低，因為注意力計算僅在較小的組內進行。

3. 組間合併

將各組內的注意力輸出合併起來，形成最終的輸出。合併方法可以是簡單的拼接（Concatenation）或某種形式的聚合（如加權求和）。

4. 可選的跨組注意力

在某些變體中，還可以在合併之前或之後引入跨組注意力機制，以捕捉組之間的依賴關係。這進一步增強了模型的表達能力，但也會增加一些計算複雜度。

程式碼示例

以下是一個簡化的PyTorch程式碼示例，展示了Grouped Attention的基本實現：

import torch
import torch.nn as nn

class GroupedAttention(nn.Module):
    def __init__(self, input_dim, num_heads, group_size):
        super(GroupedAttention, self).__init__()
        self.num_heads = num_heads
        self.group_size = group_size
        self.attention = nn.MultiheadAttention(input_dim, num_heads)

    def forward(self, x):
        N, L, D = x.shape  # Batch size (N), Sequence length (L), Embedding dimension (D)
        assert L % self.group_size == 0, "Sequence length must be divisible by group size"
        
        # Reshape input into groups
        num_groups = L // self.group_size
        x = x.view(N, num_groups, self.group_size, D)
        
        # Apply attention within each group
        x = x.permute(1, 0, 2, 3).contiguous()  # (num_groups, N, group_size, D)
        x = x.view(num_groups, N * self.group_size, D)  # (num_groups, N * group_size, D)
        
        attn_output, _ = self.attention(x, x, x)
        
        # Reshape back to original dimensions
        attn_output = attn_output.view(num_groups, N, self.group_size, D)
        attn_output = attn_output.permute(1, 0, 2, 3).contiguous()  # (N, num_groups, group_size, D)
        attn_output = attn_output.view(N, L, D)
        
        return attn_output

# Example usage
batch_size = 2
seq_length = 8
embedding_dim = 16
num_heads = 2
group_size = 4

x = torch.rand(batch_size, seq_length, embedding_dim)
grouped_attention = GroupedAttention(embedding_dim, num_heads, group_size)
output = grouped_attention(x)
print(output.shape)  # Should be (2, 8, 16)

優點

1. 計算效率高：由於注意力計算在較小的組內進行，計算複雜度顯著降低。

2. 適用長序列：更適合處理長序列或高維資料，減少了記憶體和計算資源的佔用。

3. 靈活性：可以根據具體應用需求調整組大小和注意力機制的引數。

缺點

1. 資訊丟失風險：如果組之間的關聯性較強，簡單的組內注意力可能會丟失一些跨組資訊。可透過引入跨組注意力機制來緩解這一問題。

2. 引數選擇複雜：需在實際應用中仔細選擇組大小和其他超引數，以平衡計算效率和模型效能。