Self-attention(自注意力機制)

17 min readApr 1, 2025

同步在HackMD(文章有精美的LaTex)上。

Transformer 架構由 Encoder 和 Decoder 兩部分組成，其中 Encoder 部分發展出了 BERT 這類主要用於表徵學習的雙向模型，而 Decoder 部分則發展出 GPT 等主要用於自回歸生成的模型。目前我們熟知的許多大型生成模型（如 GPT 系列、LLaMA）主要基於 Decoder 架構，但也有一些生成式模型（如 T5、BART）採用了 Encoder-Decoder 結構。

Single-head attention(單頭自注意力)

單頭注意力 (Single-Head Attention) 只使用一個注意力頭 (Attention Head) 來計算權重，相較於多頭注意力，計算量較小，但仍然保留了注意力機制的核心思想。
針對輸入序列，注意力機制計算 Token 與 Token 之間的關聯性，透過 Query 和 Key 的相似度 來決定注意力權重，然後對 Value 向量 進行加權求和 (weighted sum)，從而為該 Token 生成新的表示 (representation)。
Note: 這裡的“詞”實際上指的是 Token，而非傳統 NLP 中的詞彙 (word)。
Eaxmple: 中和有一個永和路，可以切成以下四個Token

假設Attention matrix算出來如下

詞得到新的表示(representation):

中和 = 0.5×中和 + 0×有 + 0×一個 + 0.5×永和路
有 = 0×中和 + 1×有 + 0×一個 + 0×永和路
一個 = 0×中和 + 0×有 + 1×一個 + 0×永和路
永和路 = 0.5×中和 + 0×有 + 0×一個 + 0.5×永和路

Self-Attention (自注意力機制) 透過計算 Query (Q) 和 Key (K) 之間的相似度 來確定不同 Token 之間的關聯性。相關的 Token 會被分配較高的注意力權重，而不相關的 Token 則被分配較低的權重。這些權重透過 softmax 正規化，使其總和為 1，然後用這些權重對 Value (V) 向量 進行加權求和 (weighted sum)，最終產生新的詞向量表示，使得該 Token 能夠融合來自其他 Token 的信息。
以上範例
中和←→永和路，有很高相關
“有”、”一個” 屬於贅詞就跟其他詞獨立。

Single-head attenttion公式

所有介紹的文章都會放這張圖，我也放一下，這是Single-head attenttion的計算方式

Q: Query → 查詢: 其表示目前需要關注的內容。
K: Key → 關鍵: 其表示與查詢相符的內容。
V: Value → 值: 表示最終要提取的信息，通常和Key對應。

我們不解釋運算原因，我們先介紹有什麼運算

剛剛介紹我們針對每個句子來算(中和有一個永和路)，這跟Q、K、V三個輸入有什麼關係

ANS: 需要透過投影將每個詞投影到Q、K、V。

要先將”中和有一個永和路”向量化(透過embedding layer)，假設句子我們已經做過(1)Tokenization和(2)向量化。

所以這段一段句子(中和有一個永和路)的向量表示

這個Tokeizer後我設定為4個詞，句子長度依據不同模型都不太一樣，假設句子的詞有n個，我們用n來表示。

整個Single-Head Attention計算如上圖，我們把流程切成五步驟，分別介紹

(1)計算Query, Key, Value 矩陣所以要先將X投影到Q、K、V，投影方式很簡單設定三個要學習的參數矩陣

其中d_q=d_k

(2) 計算點積(MatMul)得到score
計算Query和Key的點積，得到score：

類似統計上的Correlation matrix計算，所以說這是在算QK的相關係數矩陣

(3)縮放(Scale)：

在 Self-Attention 中，點積結果會除以 √dk，主要目的是防止數值過大影響 softmax 的分佈，並確保梯度穩定。

具體來說，score 矩陣的每個元素是 Q 的一個向量與 Kᵀ 的一個向量的點積，這涉及 dk 個數值相乘再相加。如果 dk 很大，則點積結果也會變得很大，進而使得 softmax 產生極端分佈（接近 one-hot），這會影響梯度流動，使模型難以學習長距離關係。

為了解決這個問題，Transformer 論文中選擇 將點積結果除以 √dk，讓數值保持在合理範圍內。這樣既能防止 softmax 變得過於極端，也能確保梯度不會過小導致學習困難。因此，這個縮放 (scaling) 步驟的作用是 在大維度時防止數值過大，在小維度時避免數值過小，確保模型能夠有效學習關係。

(4)softmax：對縮放後的點積結果套用softmax函數，得到注意力權重矩陣(Attention Score)，主要原因是希望attention socre可以類似機率的概念，讓美的詞的權重分布在0~1，且總和(每個row)是1。

(5)加權求和：將Attention Score與Value相乘，得到加權求和的結果

程式碼也很簡單，如下操作

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

class SingleHeadAttention(nn.Module):
    def __init__(self, embed_dim):
        """
        :param embed_dim: embedding dimension for Query、Key and Value
        """
        super(SingleHeadAttention, self).__init__()
        self.embed_dim = embed_dim

        # Linear project for Query,Key, Value
        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.scale = torch.sqrt(torch.FloatTensor([self.embed_dim]))

    def forward(self, x):
        """
        :param x: input data, Shape: (batch_size, seq_len, embed_dim)
        :return: Output, Shape: (batch_size, seq_len, embed_dim)
        """
        # project to Query, Key, Value
        Q = self.query_linear(x)
        K = self.key_linear(x)
        V = self.value_linear(x)

        # matmul and scale
        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) 
        # scale
        attention_scores = attention_scores / self.scale
        # softmax
        attention_weights = F.softmax(attention_scores, dim=-1)
        # matmul
        output = torch.matmul(attention_weights, V)

        return output, attention_weights

batch_size = 1
seq_len = 4 # 輸入的序列長度
embed_dim = 6 # 假設word embedding的dimension是6

# generate data
x = torch.randn(batch_size, seq_len, embed_dim)

# init
attention = SingleHeadAttention(embed_dim)

# forward
output, attention_weights = attention(x)

print("Output:\n", output)
print("Output shape:\n", output.shape)
print("Attention Weights:\n", attention_weights)
print("Attention Weights shape:\n", attention_weights.shape)

輸出結果如下: 可以看到attention matrix是4*4的大小，也就是我們預測輸入序列長度的大小，可以看到這就是表示每個詞彙之間的相關係數，然後每個row總和是1。

Multi-Head Attention(多頭自注意力)

所有介紹的文章都會放這張圖，我也放一下，Multi-Head attenttion的計算方式

Multi-Head Attention (MHA) 的核心思想 是透過 多組獨立的注意力機制 來學習不同的關係模式，而不是單純將輸入拆成不同子空間。具體來說，輸入 X 會經過 三組線性變換 (Wq, Wk, Wv)，分別得到 Q, K, V，然後將它們拆分成多個 head，每個 head 會獨立計算注意力權重，最後將所有 head 的輸出拼接並通過一個線性層投影回原始維度。

我之前誤以為 QKV 是額外的 Linear 運算來投影到不同空間，但實際上，這些 Linear 變換 (Wq, Wk, Wv) 是 Multi-Head Attention 的標準做法，而不是額外的設計。另外，有些實作會直接對 feature vector 進行切割來模擬多頭注意力，但這並不是 Transformer 原始設計的方式。理論上，也可以用一個大矩陣來學習所有注意力模式，但這樣參數量會增加，計算成本也會變高。

這邊一樣我們將運算拆成四個步驟介紹。

(1) 計算Query, Key, Value 矩陣
如同Single-Head Attention一樣，前面的輸入要先投影到Q、K、V，運算和Signle-Head Attention一樣就不多介紹了

(2) 將維度切成多個Head使用
將Q、K和V的維度切成多個head使用，
假設輸入的embedding維度是ddim，然後設定h個head
將維度切成h組，

範例切成2個head

(3) 計算每個Head的attention
如同Single-head Attention的計算

(4) Concat.所有head的輸出，然後Linear Project學習合併後的結果

程式碼也很簡單，如下操作

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

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        """
        :param embed_dim: embedding dimension for Query、Key and Value
        :param num_heads: number of head
        """
        super(MultiHeadAttention, 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, "Embed size needs to be divisible by heads"

        # linear project to Query、Key 和 Value
        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)

        # Linear layer for output
        self.out = nn.Linear(embed_dim, embed_dim)

        # scale
        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim]))

    def forward(self, x):
        """
        :param x: input data, Shape: (batch_size, seq_len, embed_dim)
        :return: Output, Shape: (batch_size, seq_len, embed_dim)
        """
        batch_size = x.shape[0]

        # project to Query, Key, Value
        Q = self.query_linear(x)  # [batch_size, seq_len, embed_dim]
        K = self.key_linear(x)      # [batch_size, seq_len, embed_dim]
        V = self.value_linear(x)  # [batch_size, seq_len, embed_dim]

        # split to multi-heads
        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]

        # matmul and scale
        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale  # [batch_size, num_heads, seq_len, seq_len_k]
        # Softmax
        attention_weights = F.softmax(attention_scores, dim=-1)  # [batch_size, num_heads, seq_len, seq_len_k]
        # matmul
        output = torch.matmul(attention_weights, V)  # [batch_size, num_heads, seq_len, head_dim]
        # concat.
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)  # [batch_size, seq_len, embed_dim]
        # liear for outpur
        output = self.out(output)  # [batch_size, seq_len, embed_dim]

        return output, attention_weights

# 
batch_size = 1
seq_len = 5
embed_dim = 4
num_heads = 2

# generate data
x = torch.randn(batch_size, seq_len, embed_dim)
# init
attention = MultiHeadAttention(embed_dim, num_heads)
# forward
output, attention_weights = attention(x)

print("Output:\n", output)
print("Output shape:\n", output.shape)
print("Attention Weights:\n", attention_weights)
print("Attention Weights shape:\n", attention_weights.shape)

輸出結果如下: 可以看到attention matrix是輸出2個5*5的大小，也就是我們g設定2個head，所以對應的每個head都有個各自的attentation matrix，然後每個row總和是1。(因為只是宣告沒有訓練，所以目前的attentation score出來沒有意義。)

備註: 這份資料文字和文字內容有利用ChatGPT來修改和生成。

Self-attention(自注意力機制)

Single-head attention(單頭自注意力)

Single-head attenttion公式

Multi-Head Attention(多頭自注意力)

Written by Tommy Huang

Responses (1)