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