对自注意力(self-attention)的理解以及基于pytorch的简易示例

简介

自注意力(self-attention)：一个seq2seq的映射运算，具体而言，也就是将输入向量通过映射(或者说某种函数运算)输出对应的结果. 向量的维度都为。

对于每个输出，就是用自注意力运算生成的，而这个运算原理，其实就是对输入向量进行加权平均罢了，公式为：

在此处是整个序列的索引(范围即[1,k])，并使权重相加为1。注意这个权重并不是某个参数，因为它是从和计算而来。实现这个权重最简单的方法就是先对输入向量作点积(dot product)运算：

当然，经过这次点积运算，各值域为负无穷到正无穷，所以在这之后使用softmax将其映射到的范围，以确保它们在整个序列上的和为1。

以上为自注意力的基本运算。

代码

下代码展示了基本的自注意力运算，注意此处

import torch
import torch.nn.functional as F
#固定随机数以获得相同的结果
torch.manual_seed(0)
if __name__=='__main__':

    #假设输入x是一个张量, batch size为b, 矩阵size为 txk
    #为了简便，我们在这里创建一个batch size 为1，size为2x3的张量
    #为了体现下面softmax的负值转正，此处使用标准正态分布来生成含负数的矩阵
    input_tensor= torch.randn((1, 2, 3))
    x = input_tensor
    print('X.shape:{}\nDatas:\n{}\n'.format(x.shape, x))

    # 对于原始的点积权重, 我们只需要将x与它的转置矩阵相乘即可
    raw_weights = torch.bmm(x, x.transpose(1, 2))
    
    print('raw_weights.shape:{}\n{}\n'.format(raw_weights.shape, raw_weights))

    #之后, 为了将该权重的值全转换为正值，并使它们的和为1，我们接着使用一个row-wise(对行进行操作)的softmax函数：
    #再次注意这里是基于行的softmax，进行验算时按行相加会发现和是为1的
    weights = F.softmax(raw_weights, dim=2)
    print('weights.shape:{}\n{}\n'.format(weights.shape, weights))

    #最后，为了计算输出序列y，我们只需要将权重与x相乘即可。
    #y的结构也为(b ,t, k)，它的每个结果由权重矩阵的每行与x的每列相乘所得
    y = torch.bmm(weights, x)
    print('y.shape:{}\n{}\n'.format(y.shape, y))

结果：

X.shape:torch.Size([1, 2, 3])
Datas:
tensor([[[ 1.5410, -0.2934, -2.1788],  
         [ 0.5684, -1.0845, -1.3986]]])

raw_weights.shape:torch.Size([1, 2, 2])
tensor([[[7.2079, 4.2414],
         [4.2414, 3.4554]]])

weights.shape:torch.Size([1, 2, 2])    
tensor([[[0.9510, 0.0490],
         [0.6870, 0.3130]]])

y.shape:torch.Size([1, 2, 3])
tensor([[[ 1.4934, -0.3322, -2.1406],  
         [ 1.2366, -0.5411, -1.9346]]])

此处设定的txk为2x3，即有两个向量，以上由输入x到输出y的大致流程如下图：

可以看到，对于每个结果，它们都结合了输入里的每个向量进行加权。对于输出，如果我希望的重要性降低，那么它对应的点积结果就应该是个很小甚至是负值。

参考：Transformers from scratch | peterbloem.nl