掌握AI尚方寶劍：注意力機制
加入技術(shù)交流群
1 前言
經(jīng)過上一期的范例和解說，您對于相似度的計算，已經(jīng)建立良好的基礎(chǔ)了。就可以輕易地來理解和掌握注意力(Attention) 機制。這項機制在許多大語言模型( 如ChatGPT、Gemma等) 里，都扮演了極為關(guān)鍵性的角色。再看看最近聲勢非常浩大的Sora，其關(guān)鍵技術(shù)——DiT(Diffusion Transformer) 的核心也是注意力機制。
于是，本文就從上一期介紹的相似度(Similarity) 基礎(chǔ)，繼續(xù)延伸到注意力機制。此外，更重要的是：此項機制也是可以學習的(Learnable)，于是就來把它包裝于NN模型里，成為可以訓練的注意力模型(Attention model)。
典型的Attention 模型， 包括兩種： 交叉注意力(CrossAttention) 和自注意力(SelfAttention)。本文就先來說明SelfAttention 模型的計算邏輯，及其訓練方法。
2 以“企業(yè)經(jīng)營”來做比喻
首先來做個比喻。例如，一個公司有三個部門，其投資額( 以X表示)，經(jīng)過一年的經(jīng)營績效比率( 以W表示)，其營收額( 以V 表示)，如圖1所示。
圖1
這三部門投資額是：X=[10, 6, 2.5]，其單位是---百萬元。經(jīng)過一年的經(jīng)營，其營收比率是：W=[2.0]，就可以計算出營收金額是：V=[20, 12, 5]。
接下來，公司的經(jīng)營團隊開始規(guī)畫下一年度的投資方案，針對未來新的商業(yè)投資獲利注意點，擬定一個投資預算分配表( 即注意力表)，然后計算出新年度的投資預算金額( 單位：百萬元)，如圖2所示。
圖2
其中的預算分配表，可以是相似度矩陣(Similarity matrix)，亦即經(jīng)由相似度的計算而來?，F(xiàn)在，就來理解上圖的計算邏輯，請觀摩一個Python的實現(xiàn)代碼：
# ax01.py
import numpy as np
import torch
X = torch.tensor([[10.0],[6.0],[2.5]]) # 投資額
W = torch.tensor([[2.0]]) # 經(jīng)營績效
V = X.matmul(W) # 計算營收
A = torch.tensor(
[[1.0, 0., 0.],
[0.9, 0.1, 0.],
[0.6, 0.3, 0.1]]) # 預算分配表
Z = A.matmul(V) # 計算分配額
print(‘n 投資預算額Z：’)
print(Z) #np.round(Z.detach().numpy()))
#END
接著，就執(zhí)行這個程序。此時就輸入X和W，計算出V值。然后輸入相似度表A，計算出新年度的投資預算額，并輸出如下：
3 使用Attention計算公式
在上一期里，已經(jīng)說明了，相似度矩陣是直接計算向量的點積(Dot-product)，即將兩向量的對應(yīng)元素相乘再相加。然后，這相似度矩陣再除以它們的歐氏長度的乘積，將相似度的值正規(guī)化，就得到余弦(Cosine)相似度。而且，如果將上述的相似度矩陣，在經(jīng)由Softmax() 函數(shù)的運算，就得到注意力矩陣(Attention weights) 了。例如，有兩個矩陣：Q 和K，就能計算出注意力矩陣，如圖3 所示。
圖3
那么，就可以繼續(xù)思考一個重要問題，就是：如何計算出Q和K矩陣呢? 答案是：可以由SelfAttention模型來預測出來。也就是，由輸入數(shù)據(jù)X來與SelfAttention模型的權(quán)重Wq相乘而得到Q。同時，也由輸入數(shù)據(jù)X來與這模型的權(quán)重Wk 相乘而得到K，如圖4所示。
圖4
當我們把上圖里的Wq、Wk和Wv權(quán)重都放入SelfAttention模型里， 就能進行機器學習(Machine learning) 來找出最佳的權(quán)重值( 即Wq、Wk 和Wv)，就能預測出Q、K 和V 了。并且可繼續(xù)計算出A 了。
4 訓練SelfAttention模型
現(xiàn)在就把Wq、Wk 和Wv 都放入SelfAttention 模型里。請觀摩這個SelfAttention 模型的代碼范例，如下：
# ax02.py
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module): # 定義模型
def __init__(self):
super(SelfAttention, self).__init__()
self.Wq = nn.Linear(1, 2, bias=False)
self.Wk = nn.Linear(1, 2, bias=False)
self.Wv = nn.Linear(1, 1, bias=False)
def forward(self, x):
Q = self.Wq(x)
K = self.Wk(x)
V = self.Wv(x)
Scores = Q.matmul(K.T)
A = F.softmax(Scores, dim=-1) # Attention_weights
Z = A.matmul(V) # 計算Z
return Z, A, V
model = SelfAttention() # 建立模型
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.004)
# 輸入X
X = torch.tensor([[10.0],[6.0],[2.5]])
# 設(shè)定Target Z
target_attn = torch.tensor([[20.0],[19.0],[16.0]])
print(‘展開訓練1800 回合...’)
for epoch in range(1800+1):
Z, A, V = model(X) # 正向傳播
loss = criterion(Z, target_attn) # 計算損失
optimizer.zero_grad() # 反向傳播和優(yōu)化
loss.backward()
optimizer.step()
if(epoch%600 == 0):
print(‘ep=’, epoch,‘loss=’, loss.item())
# 進行預測
Z, A, V = model(X)
print(‘n----- 預算分配表A -----’)
print(np.round(A.detach().numpy(), 1))
print(‘n----- 投資預算額Z -----’)
print(np.round(Z.detach().numpy()))
#END
然后就執(zhí)行這個程序，此時會展開1800 回合的訓練。在訓練過程中，回持續(xù)修正模型里的權(quán)重( 即Wq、Wk和Wv)，并且其損失(Loss) 值會持續(xù)下降，如下：
一旦訓練完成了，就可以展開預測(Prediction)。此時，就計算出Q、K 和V，然后繼續(xù)計算出A 和Z 值。
5 結(jié)束語
本期基于相似度計算，繼續(xù)說明注意力機制的計算邏輯，建立SelfAttention模型，并且訓練1800 回合，然后進行預測。
從這范例中，可以領(lǐng)會到SelfAttention模型能順利捕捉到企業(yè)的經(jīng)營規(guī)律，并進行準確的預測。
（本文來源于《EEPW》2024.4）
評論