my-attention
考虑别的token对当前token的语义影响
第一种情况， 维度缩减
输入x= [4x6]
dk=3
wq = [6x3]
wk = [6x3]
wv = [6x3]
q = x@ wq = [4x6]*[6x3] = [4x3]
k = x@ wk = [4x6]*[6x3] = [4x3]
r=q@k.T = [4x3]*[3x4]=[4x4]
缩放
r = r/sqrt(dk)=[4x4]
a = softmax®=[4x4]
v = x@ wv = [4x6]*[6x3] = [4x3]
out = a@v = [4x4] * [4x3] = [4x3]
第二种情况， 维度不缩减
输入x= [4x6]
输出维度为6
dk=6
随机生成qkv
wq = [6x6]
wk = [6x6]
wv = [6x6]
q = x@ wq = [4x6]*[6x6] = [4x6]
k = x@ wk = [4x6]*[6x6] = [4x6]
r=q@k.T = [4x6]*[6x4]=[4x4]
缩放
r = r/sqrt(dk)=[4x4]
归一化
a = softmax®=[4x4]
原始值增加权重
v = x@ wv = [4x6]*[6x6] = [4x6]
out = a@v = [4x4] * [4x6] = [4x6]
保证输出结果的维度和要求要一致
下面是用代码实现了一下自注意力机制
import math
import torch
from torch import nn
x = torch.randn(16, 64, 512)
d_model = 512
h_num = 8
class Self_Attention(nn.Module):
def __init__(self, d_model, h_num):
# 调用父类构造函数
super(Self_Attention, self).__init__()
self.d_model = d_model
self.h_num = h_num
self.softmax = nn.Softmax(dim=-1)
self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.w_o = nn.Linear(d_model, d_model)
def forward(self, x):
B, L, D = x.shape
h_d = self.d_model // self.h_num
q, k, v = self.w_q(x), self.w_k(x), self.w_v(x)
q = q.view(B, L, self.h_num, h_d).transpose(1, 2)
k = k.view(B, L, self.h_num, h_d).transpose(1, 2)
v = v.view(B, L, self.h_num, h_d).transpose(1, 2)
r = q @ k.transpose(2, 3) / math.sqrt(h_d)
mask = torch.tril(torch.ones(L, L, dtype=bool))
r = r.masked_fill(~mask, -10000)
a = self.softmax(r)
o = a @ v
o = o.transpose(1, 2).contiguous().view(B, L, self.d_model)
return self.w_o(o)
attention = Self_Attention(d_model, h_num)
y = attention(x)
print(y.shape)
print(y)