self-attention是什么?
一個 self-attention 模塊接收 n 個輸入，然后返回 n 個輸出。自注意力機制讓每個輸入都會彼此交互（自），然后找到它們應該更加關注的輸入（注意力）。自注意力模塊的輸出是這些交互的聚合和注意力分數。
self-attention模塊包括以下步驟：
- 准備輸入
- 初始化權重
- 推導key, query 和 value
- 計算輸入1 的注意力得分
- 計算 softmax
- 將分數與值相乘
- 將權重值相加，得到輸出 1
- 對輸入 2 和輸入 3 重復步驟 4-7
第一步：准備輸入
圖1.1: 准備輸入
假設有 3 個輸入，每個輸入的維度為 4.
Input 1: [1, 0, 1, 0]
Input 2: [0, 2, 0, 2]
Input 3: [1, 1, 1, 1]
第二步：初始化權重
每個輸入必須有 3 個表征，分別被稱為鍵（key，橙色）、查詢（query，紅色）和值（value，紫色）。在此示例中，我們設這些表征的維度為 3。因為每個輸入的維度為 4，所以這意味着每組權重的形狀為 4×3。
圖1.2：從每個輸入得出鍵、查詢和值的表示
為了得到這些表征，每個輸入（綠色）都要與一組鍵的權重、一組查詢的權重、一組值的權重相乘。在這個示例中，我們按如下方式初始化這三個權重：
鍵的權重：
[[0, 0, 1], [1, 1, 0], [0, 1, 0], [1, 1, 0]]
查詢的權重：
[[1, 0, 1], [1, 0, 0], [0, 0, 1], [0, 1, 1]]
值的權重：
[[0, 2, 0], [0, 3, 0], [1, 0, 3], [1, 1, 0]]
備注：在神經網絡設置中，這些權重通常是較小的數值，初始化也是使用合適的隨機分布來實現，比如高斯分布、Xavier 分布、Kaiming 分布。
第三步：推導鍵、查詢和值
現在我們有三組權重了，我們來實際求取每個輸入的鍵、查詢和值的表征：
輸入 1 的鍵表征：
[0, 0, 1] [1, 0, 1, 0] x [1, 1, 0] = [0, 1, 1] [0, 1, 0] [1, 1, 0]
使用同樣一組權重求取輸入 2 的鍵表征：
[0, 0, 1] [0, 2, 0, 2] x [1, 1, 0] = [4, 4, 0] [0, 1, 0] [1, 1, 0]
使用同樣一組權重求取輸入 3 的鍵表征：
[0, 0, 1] [1, 1, 1, 1] x [1, 1, 0] = [2, 3, 1] [0, 1, 0] [1, 1, 0]
向量化以上運算能實現更快的速度：
[0, 0, 1] [1, 0, 1, 0] [1, 1, 0] [0, 1, 1] [0, 2, 0, 2] x [0, 1, 0] = [4, 4, 0] [1, 1, 1, 1] [1, 1, 0] [2, 3, 1]
圖1.3a：從每個輸入推導出鍵表示
通過類似的方式，我們求取每個輸入的值表征：
[0, 2, 0] [1, 0, 1, 0] [0, 3, 0] [1, 2, 3] [0, 2, 0, 2] x [1, 0, 3] = [2, 8, 0] [1, 1, 1, 1] [1, 1, 0] [2, 6, 3]
圖1.3b：從每個輸入推導出值表示
最后還有查詢表征：
[1, 0, 1] [1, 0, 1, 0] [1, 0, 0] [1, 0, 2] [0, 2, 0, 2] x [0, 0, 1] = [2, 2, 2] [1, 1, 1, 1] [0, 1, 1] [2, 1, 3]
圖1.3b：從每個輸入推導出查詢表示
第四步：計算輸入1 的 attention scores
圖1.4：從查詢1中計算注意力得分(藍色)
為了獲得注意力得分，我們首先在輸入1的查詢(紅色)和所有鍵(橙色)之間取一個點積。因為有3個鍵表示(因為有3個輸入)，我們得到3個注意力得分(藍色)。
[0, 4, 2] [1, 0, 2] x [1, 4, 3] = [2, 4, 4] [1, 0, 1]
注意這里僅使用了輸入 1 的查詢。后面我們會為其它查詢重復同一步驟。
備注：上面的運算也被稱為點積注意力（dot product attention），這是眾多評分函數中的一個，其它評分函數還包括擴展式點積和 additive/concat
第五步：計算 softmax
圖1.5：Softmax注意力評分(藍色)
在所有注意力得分中使用softmax(藍色)。
softmax([2, 4, 4]) = [0.0, 0.5, 0.5]
第六步：將分數與值相乘
圖1.6：由值(紫色)和分數(藍色)的相乘推導出加權值表示(黃色)
每個輸入的softmaxed attention 分數(藍色)乘以相應的值(紫色)。結果得到3個對齊向量(黃色)。在本教程中，我們將它們稱為加權值。
1: 0.0 * [1, 2, 3] = [0.0, 0.0, 0.0] 2: 0.5 * [2, 8, 0] = [1.0, 4.0, 0.0] 3: 0.5 * [2, 6, 3] = [1.0, 3.0, 1.5]
第七步：將加權值相加得到輸出1
圖1.7：將所有加權值(黃色)相加，得到輸出1(深綠色)
將所有加權值(黃色)按元素指向求和：
[0.0, 0.0, 0.0] + [1.0, 4.0, 0.0] + [1.0, 3.0, 1.5] ----------------- = [2.0, 7.0, 1.5]
所得到的向量 [2.0, 7.0, 1.5]（深綠色）是輸出 1，這是基於輸入 1 的查詢表征與所有其它鍵（包括其自身的）的交互而得到的。
第八步：為輸入 2 和 3 重復 4-7 步驟
現在已經完成了對輸出 1 的求解，我們再為輸出 2 和輸出 3 重復步驟 4-7。
圖1.8：對輸入2和輸入3重復前面的步驟
代碼
步驟1：准備輸入
import torch
x = [
[1, 0, 1, 0], # Input 1
[0, 2, 0, 2], # Input 2
[1, 1, 1, 1] # Input 3
]
x = torch.tensor(x, dtype=torch.float32)
步驟2：初始化權重
w_key = [
[0, 0, 1],
[1, 1, 0],
[0, 1, 0],
[1, 1, 0]
]
w_query = [
[1, 0, 1],
[1, 0, 0],
[0, 0, 1],
[0, 1, 1]
]
w_value = [
[0, 2, 0],
[0, 3, 0],
[1, 0, 3],
[1, 1, 0]
]
w_key = torch.tensor(w_key, dtype=torch.float32)
w_query = torch.tensor(w_query, dtype=torch.float32)
w_value = torch.tensor(w_value, dtype=torch.float32)
步驟3: 推導鍵、查詢和值
keys = x @ w_key
querys = x @ w_query
values = x @ w_value
print(keys)
# tensor([[0., 1., 1.],
# [4., 4., 0.],
# [2., 3., 1.]])
print(querys)
# tensor([[1., 0., 2.],
# [2., 2., 2.],
# [2., 1., 3.]])
print(values)
# tensor([[1., 2., 3.],
# [2., 8., 0.],
# [2., 6., 3.]])
步驟4：計算注意力得分
attn_scores = querys @ keys.T
# tensor([[ 2., 4., 4.], # attention scores from Query 1
# [ 4., 16., 12.], # attention scores from Query 2
# [ 4., 12., 10.]]) # attention scores from Query 3
步驟5：計算softmax
from torch.nn.functional import softmax
attn_scores_softmax = softmax(attn_scores, dim=-1)
# tensor([[6.3379e-02, 4.6831e-01, 4.6831e-01],
# [6.0337e-06, 9.8201e-01, 1.7986e-02],
# [2.9539e-04, 8.8054e-01, 1.1917e-01]])
# For readability, approximate the above as follows
attn_scores_softmax = [
[0.0, 0.5, 0.5],
[0.0, 1.0, 0.0],
[0.0, 0.9, 0.1]
]
attn_scores_softmax = torch.tensor(attn_scores_softmax)
步驟6：將得分和值相乘
weighted_values = values[:,None] * attn_scores_softmax.T[:,:,None]
# tensor([[[0.0000, 0.0000, 0.0000],
# [0.0000, 0.0000, 0.0000],
# [0.0000, 0.0000, 0.0000]],
#
# [[1.0000, 4.0000, 0.0000],
# [2.0000, 8.0000, 0.0000],
# [1.8000, 7.2000, 0.0000]],
#
# [[1.0000, 3.0000, 1.5000],
# [0.0000, 0.0000, 0.0000],
# [0.2000, 0.6000, 0.3000]]])
步驟7：求和加權值
outputs = weighted_values.sum(dim=0)
# tensor([[2.0000, 7.0000, 1.5000], # Output 1
# [2.0000, 8.0000, 0.0000], # Output 2
# [2.0000, 7.8000, 0.3000]]) # Output 3