自注意力简介

在注意力机制中，每个查询都会关注所有的键值对并生成一个注意力输出。如果查询q，键k和值v都来自于同一组输入，那么这个注意力就被称为是自注意力（self-attention）。自注意力这部分理论，我觉得台大李宏毅老师的课程讲得最好。

自注意力就是输入一堆向量，假设称为a1,a2,a3,a4，那么这四个向量都会参与自注意力机制的运算，得到的结果仍然是四个输出，这四个输出再去做全连接运算。而每一个自注意力机制的输出都用到了a1~a4四个向量来进行运算，也就是说每个输出都是观察了所有的输入之后才得到的。

首先，输入a1需要和a2,a3,a4分别计算相关性，这个相关性可以由缩放点积方式来计算，也就称作缩放点积注意力，也可以由两个两个输入向量相加后再做非线性处理得到，称为加性注意力。

缩放点积的计算方法如下：

输入一个向量v1和一个向量v2，v1去乘上一个可训练矩阵Wq得到q，v2去乘上一个可训练矩阵Wk得到k，再把这个q和k做一个点积运算，得到的就是α，类似于相似度。

回到前面的例子中，a1这里既作为q，又作为k，又作为v。其中，Wq*a1就是q1，Wk*a1就是k1，q1和k1的点积就是α11，相当于a1自己和自己的相似度，同样的，a1和a2,a3,a4分别计算得到α12，α13，α14，然后将α11，α12，α13，α14经过softmax得到最终的四个输出，如下图所示：

然后再用一个可训练矩阵Wv去乘以a1得到v1，用计算得到的相似度α'11去乘以v1，得到一个值temp11；同样的，用可训练矩阵Wv去乘以a2得到v2，用计算得到的相似度α‘12去乘以v2，得到temp12；类似的得到temp13,temp14，然后把temp11+temp12+temp13+temp14得到b1，这个b1就是自注意力机制的第一个输出。

我们刚刚是以a1的视角做的运算，得到b1，同样可以以a2,a3,a4的视角做运算，得到b2,b3,b4。这次就得到了自注意力机制的输出。光看最后这个结构图，有点类似全连接，只是里面的运算过程比全连接要复杂。

下面，我们来看一下如何用代码实现自注意力的计算。

import torch  
import torch.nn as nn  
import torch.nn.functional as F  class SelfAttention(nn.Module):# embed_size代表输入的向量维度，heads代表多头注意力机制中的头数量def __init__(self, embed_size, heads): super(SelfAttention,self).__init__()self.embed_size = embed_sizeself.heads = headsself.head_dim = embed_size // heads # 每个头的维度# 用assert断言机制判断assert (self.head_dim * heads == embed_size), "Embedding size needs to be divisible by heads" # 没有偏置项，其实这个线性层本质上就是为了计算值Wv*a = Vself.values = nn.Linear(self.head_dim, self.head_dim, bias=False)self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)# 最后的全连接操作，输出仍是输入的向量维度，也就是说大小是不变的self.fc_out = nn.Linear(heads*self.head_dim, self.embed_size)def forward(self, values, keys, query, mask):# 这个mask也很关键，它用于控制模型在处理序列数据时应该关注哪些部分，以及忽略哪些部分N = query.shape[0] # 获取输入的批量个数print("N:",N)value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # 获取输入序列的长度# Split the embedding into self.heads different pieces  # 把k,q,v都切分为多个组values = values.reshape(N, value_len, self.heads, self.head_dim)keys = keys.reshape(N, key_len, self.heads, self.head_dim)queries = query.reshape(N, query_len, self.heads, self.head_dim)# 计算k,q,vvalues = self.values(values)keys = self.keys(keys)queries = self.queries(queries)energy = torch.einsum("nqhd,nkhd->nhqk",[queries, keys]) # 格式转化print("queries.shape:", queries.shape)print("keys.shape:", keys.shape)print("energy.shape:", energy.shape)if mask is not None:energy = energy.masked_fill(mask==0, float("-1e20"))attention = torch.softmax(energy/(self.embed_size**(1/2)), dim=3) # softmax内部是缩放点积print("attention.shape:", attention.shape)print("values.shape:", values.shape)out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads*self.head_dim)out = self.fc_out(out)return out

embed_size = 512
heads = 8
attention = SelfAttention(embed_size, heads)# batch size 1, seq length 60
values = torch.rand(1,60,embed_size)
keys = torch.rand(1,60,embed_size)
queries = torch.rand(1,60,embed_size)
mask = None # 假设没有maskout = attention(values, keys, queries, mask)
print(out.shape)# 输出
N: 1
queries.shape: torch.Size([1, 60, 8, 64])
keys.shape: torch.Size([1, 60, 8, 64])
energy.shape: torch.Size([1, 8, 60, 60])
attention.shape: torch.Size([1, 8, 60, 60])
values.shape: torch.Size([1, 60, 8, 64])
torch.Size([1, 60, 512])

通过这个程序，我们可以看到，自注意力机制是不改变输入和输出的形状的，输入的Q,K,V格式是[1,60,512]，输出的结果的仍然是[1,60,512]。

下面是几点说明：

1. 这里的embed_size代表的是输入到自注意力层中的每个元素的向量维度。在Transformer模型中，输入数据首先会被转换成一个固定长度的向量，这个向量的长度就称为embed_size。

2. mask表示的是模型在处理序列数据时，应该忽略掉哪部分，我这里设置为None，也就是全部参与计算。

3. einsum，称为爱因斯坦求和，起源是爱因斯坦在研究广义相对论时，需要处理大量求和运算，为了简化这种繁复的运算，提出了求和约定，推动了张量分析的发展。einsum 可以计算向量、矩阵、张量运算，如果利用得当，sinsum可完全代替其他的矩阵计算方法。

例如，C = einsum('ij,jk->ik', A, B)，就相当于两个矩阵求内积：cik = Σj AijBjk。

通过输出可以看到，在计算前queries的形状是[1,60,8,64]，keys的形状是[1,60,8,64]，在表达式"nqhd,nkhd->nhqk"中，n=1,q=60,h=8,d=64,k=60,两个矩阵进行内积，因此得到的结果是nhqk，也就是[1,8,60,60]。