掌握AI尚方寶劍：注意力機制

1 前言

經(jīng)過上一期的范例和解說，您對于相似度的計算，已經(jīng)建立良好的基礎(chǔ)了。就可以輕易地來理解和掌握注意力(Attention) 機制。這項機制在許多大語言模型( 如ChatGPT、Gemma等) 里，都扮演了極為關(guān)鍵性的角色。再看看最近聲勢非常浩大的Sora，其關(guān)鍵技術(shù)——DiT(Diffusion Transformer) 的核心也是注意力機制。

于是，本文就從上一期介紹的相似度(Similarity) 基礎(chǔ)，繼續(xù)延伸到注意力機制。此外，更重要的是：此項機制也是可以學(xué)習(xí)的(Learnable)，于是就來把它包裝于NN模型里，成為可以訓(xùn)練的注意力模型(Attention model)。

典型的Attention 模型， 包括兩種： 交叉注意力(CrossAttention) 和自注意力(SelfAttention)。本文就先來說明SelfAttention 模型的計算邏輯，及其訓(xùn)練方法。

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è)投資獲利注意點，擬定一個投資預(yù)算分配表( 即注意力表)，然后計算出新年度的投資預(yù)算金額( 單位：百萬元)，如圖2所示。

圖2

其中的預(yù)算分配表，可以是相似度矩陣(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]])         # 預(yù)算分配表

Z = A.matmul(V)          # 計算分配額

print(‘n 投資預(yù)算額Z：’)

print(Z) #np.round(Z.detach().numpy()))

#END

接著，就執(zhí)行這個程序。此時就輸入X和W，計算出V值。然后輸入相似度表A，計算出新年度的投資預(yù)算額，并輸出如下：

3 使用Attention計算公式

在上一期里，已經(jīng)說明了，相似度矩陣是直接計算向量的點積(Dot-product)，即將兩向量的對應(yīng)元素相乘再相加。然后，這相似度矩陣再除以它們的歐氏長度的乘積，將相似度的值正規(guī)化，就得到余弦(Cosine)相似度。而且，如果將上述的相似度矩陣，在經(jīng)由Softmax() 函數(shù)的運算，就得到注意力矩陣(Attention weights) 了。例如，有兩個矩陣：Q 和K，就能計算出注意力矩陣，如圖3 所示。

圖3

那么，就可以繼續(xù)思考一個重要問題，就是：如何計算出Q和K矩陣呢? 答案是：可以由SelfAttention模型來預(yù)測出來。也就是，由輸入數(shù)據(jù)X來與SelfAttention模型的權(quán)重Wq相乘而得到Q。同時，也由輸入數(shù)據(jù)X來與這模型的權(quán)重Wk 相乘而得到K，如圖4所示。

圖4

當我們把上圖里的Wq、Wk和Wv權(quán)重都放入SelfAttention模型里， 就能進行機器學(xué)習(xí)(Machine learning) 來找出最佳的權(quán)重值( 即Wq、Wk 和Wv)，就能預(yù)測出Q、K 和V 了。并且可繼續(xù)計算出A 了。

4 訓(xùn)練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(‘展開訓(xùn)練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())

# 進行預(yù)測

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í)行這個程序，此時會展開1800 回合的訓(xùn)練。在訓(xùn)練過程中，回持續(xù)修正模型里的權(quán)重( 即Wq、Wk和Wv)，并且其損失(Loss) 值會持續(xù)下降，如下：

一旦訓(xùn)練完成了，就可以展開預(yù)測(Prediction)。此時，就計算出Q、K 和V，然后繼續(xù)計算出A 和Z 值。

5 結(jié)束語

本期基于相似度計算，繼續(xù)說明注意力機制的計算邏輯，建立SelfAttention模型，并且訓(xùn)練1800 回合，然后進行預(yù)測。

從這范例中，可以領(lǐng)會到SelfAttention模型能順利捕捉到企業(yè)的經(jīng)營規(guī)律，并進行準確的預(yù)測。

（本文來源于《EEPW》2024.4）