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

1 前言

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

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

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

2 以“企業(yè)經(jīng)營(yíng)”來(lái)做比喻

首先來(lái)做個(gè)比喻。例如，一個(gè)公司有三個(gè)部門(mén)，其投資額( 以X表示)，經(jīng)過(guò)一年的經(jīng)營(yíng)績(jì)效比率( 以W表示)，其營(yíng)收額( 以V 表示)，如圖1所示。

圖1

這三部門(mén)投資額是：X=[10, 6, 2.5]，其單位是---百萬(wàn)元。經(jīng)過(guò)一年的經(jīng)營(yíng)，其營(yíng)收比率是：W=[2.0]，就可以計算出營(yíng)收金額是：V=[20, 12, 5]。

接下來(lái)，公司的經(jīng)營(yíng)團隊開(kāi)始規畫(huà)下一年度的投資方案，針對未來(lái)新的商業(yè)投資獲利注意點(diǎn)，擬定一個(gè)投資預算分配表( 即注意力表)，然后計算出新年度的投資預算金額( 單位：百萬(wàn)元)，如圖2所示。

圖2

其中的預算分配表，可以是相似度矩陣(Similarity matrix)，亦即經(jīng)由相似度的計算而來(lái)?，F在，就來(lái)理解上圖的計算邏輯，請觀(guān)摩一個(gè)Python的實(shí)現代碼：

# ax01.py

import numpy as np

import torch

X = torch.tensor([[10.0],[6.0],[2.5]]) # 投資額

W = torch.tensor([[2.0]])                 # 經(jīng)營(yíng)績(jì)效

V = X.matmul(W)                            # 計算營(yíng)收

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ù)，就執行這個(gè)程序。此時(shí)就輸入X和W，計算出V值。然后輸入相似度表A，計算出新年度的投資預算額，并輸出如下：

3 使用Attention計算公式

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

圖3

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

圖4

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

4 訓練SelfAttention模型

現在就把Wq、Wk 和Wv 都放入SelfAttention 模型里。請觀(guān)摩這個(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) # 計算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]])

   # 設定Target Z

   target_attn = torch.tensor([[20.0],[19.0],[16.0]])

   print(‘展開(kāi)訓練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())

# 進(jìn)行預測

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

然后就執行這個(gè)程序，此時(shí)會(huì )展開(kāi)1800 回合的訓練。在訓練過(guò)程中，回持續修正模型里的權重( 即Wq、Wk和Wv)，并且其損失(Loss) 值會(huì )持續下降，如下：

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

5 結束語(yǔ)

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

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

（本文來(lái)源于《EEPW》2024.4）