Self-Attention的学习记录以及代码

理论学习部分

步骤 1.1

通过矩阵运算,得出相应的q,k,v,其中$w^q,w^k,w^v$是对应的权重，可以通过模型的训练来不断迭代更新。其计算公式如下所示：

q : query(to match others)

$$q^i=w^q{a}^i$$

k : key(to be matched)

$$k^i=w^k{a}^i$$

v :value(information to be extracted)

$$v^i=w^v{a}^i$$
其中，q表示为查询(query),k表示为键(key),v表示为值(value)。$a^i$为输入的第i个数据。其图解如下：

步骤 1.2

  计算注意力得分(哪些数据更值得模型去关注),有两种计算策略

1）缩放点积注意力(scaled dot-product attention)(重点)

  使用点积可以得到计算效率更高的评分函数,但是点积操作要求q和k具有相同的长度$d$。 假设q和k的所有元素都是独立的随机变量， 并且都满足零均值和单位方差， 那么两个向量的点积的均值为$0$，方差为$d$。 为确保无论向量长度如何， 点积的方差在不考虑向量长度的情况下仍然是$1$， 我们将点积除以$\sqrt{d}$， 则缩放点积注意力（scaled dot-product attention）评分函数为:

$$\alpha (\mathbf{Q},\mathbf{K})=\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d_k}}$$

  上述公式中，Q和K分别表示为$q^1,q^2...,q^n$和$k^1,k^2...,k^n$堆叠而成的矩阵，当具体到计算第i个数据与第j个数据的注意力得分时(其中 i 是可以等于 j 的)，其公式如下：

$${\alpha }_{i,j}=\frac{{\mathbf{q}}^{\mathbf{i}}{{\mathbf{k}}^{\mathbf{i}}}^{\top }}{\sqrt{d_k}}$$

2）加性注意力(additive attention)（了解）

  一般来说，当q和k是不同长度的矢量时， 我们可以使用加性注意力作为评分函数。 给定$\mathbf{q}\in {\mathbb{R}}^q$和$\mathbf{k}\in {\mathbb{R}}^k$, 加性注意力（additive attention）的评分函数为:
$$\alpha (\mathbf{q},\mathbf{k})={\mathbf{w}}_v^{\top }\text{tanh}({\mathbf{W}}_q\mathbf{q}+{\mathbf{W}}_k\mathbf{k})\in \mathbb{R}$$

  其中可学习的参数是${\mathbf{W}}_q\in {\mathbb{R}}^{h\times q}$,${\mathbf{W}}_k\in {\mathbb{R}}^{h\times k}$和${\mathbf{w}}_v\in {\mathbb{R}}^h$。如上述公式所示， 将q和k输入到一个多层感知机（MLP）中， 感知机包含一个隐藏层，其隐藏单元数是一个超参数h(控制矩阵W的形状)。 通过使用tanh作为激活函数，并且禁用偏置项(bias)。

步骤 1.3

SoftMax归一化

  使用softmax作为归一化函数,能让数值差异更明显,具体公式如下:

$${\hat{\alpha }}_{1,i}=\frac{exp({\alpha }_{1,i})}{{\sum }_{n}exp({\alpha }_{1,n})}$$

  分子为第1个数据与第i个数据的注意力得分取指数，分母为第1个数据与其他所有数据的注意力得分取指数，然后求和。

步骤 1.4

加权求和

  将归一化后的注意力得分进行加权求和 ，其中$b^1$是由输入数据$a^1$ 。其公式如下所示：

$$b^1=\sum _i{\hat{\alpha }}_{1,i}{v}^i$$

  最后，上述所有公式合成：

$$Attention(Q,K,V)=softmax(\frac{\mathbf{Q}{\mathbf{K}}^{\top }}{\sqrt{d_k}})V$$

简单例子

Weights for query:      Weights for key:     Weights for value:        input_A:
  [[1, 0, 1],              [[0, 0, 1],          [[0, 2, 0],        a1  [[1, 0, 1, 0],
   [1, 0, 0],               [1, 1, 0],           [0, 3, 0],        a2   [0, 2, 0, 2],
   [0, 0, 1],               [0, 1, 0],           [1, 0, 3],        a3   [1 ,1, 1, 1]]
   [0, 1, 1]]               [1, 1, 0]]           [1, 1, 0]]

下面是Q，K，V的具体计算过程，依据公式 $A{W}^q=Q，A{W}^k=K，A{W}^v=V$ ,
其输入input的形状为(seq_len , input_dim)，其中seq_len为序列长度，input_dim为数据维度。Q，K的权重矩阵的形状为(input_dim , dim_k)，因为点积操作要求q和k具有相同的长度d(dim_k)。V的权重矩阵的形状为(input_dim , dim_v)。

Q的计算过程：

               [1, 0, 1]
[1, 0, 1, 0]   [1, 0, 0]   [1, 0, 2]         q1
[0, 2, 0, 2] x [0, 0, 1] = [2, 2, 2]         q2
[1, 1, 1, 1]   [0, 1, 1]   [2, 1, 3]         q3

形状变化 (seq_len , input_dim) x (input_dim ,  dim_k)=(seq_len , dim_k)

K的计算过程：

               [0, 0, 1]
[1, 0, 1, 0]   [1, 1, 0]   [0, 1, 1]         k1
[0, 2, 0, 2] x [0, 1, 0] = [4, 4, 0]         k2
[1, 1, 1, 1]   [1, 1, 0]   [2, 3, 1]         k3

形状变化 (seq_len , input_dim) x (input_dim ,  dim_k)=(seq_len , dim_k)

V的计算过程：

               [0, 2, 0]
[1, 0, 1, 0]   [0, 3, 0]   [1, 2, 3]         v1
[0, 2, 0, 2] x [1, 0, 3] = [2, 8, 0]         v2
[1, 1, 1, 1]   [1, 1, 0]   [2, 6, 3]         v3

形状变化 (seq_len , input_dim) x (input_dim ,  dim_v)=(seq_len , dim_v)

计算score,依据公式 $\mathbf{Q}{\mathbf{K}}^{\top }$ 或者 ${\alpha }_{i,j}={\mathbf{q}}^{\mathbf{i}}{{\mathbf{k}}^{\mathbf{i}}}^{\top }$，上面公式的${\alpha }_{i,j}$可以这样理解：第 i 个数据的 q 与第 j 个数据的 k 计算得到的注意力分数。在计算中，K矩阵需要转置，才能进行矩阵运算。注意，此处为了方便计算，没有除去$\sqrt{d_k}$。

                    k1 k2 k3    
  q1   [1, 0, 2]   [0, 4, 2]   [2,  4,  4]            [α11, α12, α13]
  q2   [2, 2, 2] x [1, 4, 3] = [4, 16, 12]            [α21, α22, α23]
  q3   [2, 1, 3]   [1, 0, 1]   [4, 12, 10]            [α31, α32, α33]
  
形状变化 (seq_len , dim_k) x (dim_k,  seq_len )=(seq_len , seq_len)

softmax计算

[2,  4,  4]  -->  [0.0634,         0.4683,     0.4683]
[4, 16, 12]  -->  [6.0337e-06, 9.8201e-01, 1.7986e-02]
[4, 12, 10]  -->  [2.9539e-04, 8.8054e-01, 1.1917e-01]

最后加权求和,下面式子单拎出${\alpha }^{1j}$

                                        [1, 2, 3]  
α1j [0.0634,     0.4683,     0.4683]  x [2, 8, 0]  = b1
                                        [2, 6, 3]   
                                        
整体矩阵形状变化 (seq_len , seq_len) x (seq_len ,  dim_v)=(seq_len , dim_v)            

图解

代码实现

import torch
import torch.nn as nn

class Self_Attention(nn.Module):
    def __init__(self,input_dim,dim_k,dim_v):
        super(Self_Attention,self).__init__()
        self.q = nn.Linear(input_dim,dim_k,bias=False)#在线性层中可以设置 bias ，是否采用偏置值
        self.k = nn.Linear(input_dim,dim_k,bias=False)
        self.v = nn.Linear(input_dim,dim_v,bias=False)
        self._norm_fact = 1 / (dim_k**0.5) #缩放系数

    def forward(self,x):
        # input shape: batch_size * seq_len * ipnut_dim
      
        Q = self.q(x) # Q shape: batch_size * seq_len * dim_k
        K = self.k(x) # K shape: batch_size * seq_len * dim_k
        V = self.v(x) # V shape: batch_size * seq_len * dim_v

        #torch.bmm只支持3维矩阵，torch.matmul支持高维的矩阵计算
        #permute相当于可以同时操作于tensor的若干维度，transpose只能同时作用于tensor的两个维度
        
        #atten = torch.softmax(torch.bmm(Q,K.permute(0,2,1)) * self._norm_fact,dim=-1) 
        atten = torch.softmax(torch.matmul(Q,K.transpose(-2, -1)) * self._norm_fact,dim=-1)
        
        output=torch.matmul(atten,V)
        return output

#用法示例
torch.random.manual_seed(420)
X=torch.randn(1,3,4)  # input : batch_size * seq_len * input_dim
Attention=Self_Attention(4,7,2)  # input_dim , dim_k , dim_v
Y=Attention(X)

Multi-Head Self-Attention

在理解的Self-Attention（自注意力）的原理后之后，我们可以衍生出其变种形式Multi-Head Self-Attention（多头自注意力）。该衍生的核心如下图所示，对于每个数据的q、k、v的数量都乘2，使得每个数据都有2个q、k、v，对每个都进行如Self-Attention中的计算，那么就会得到2个b，也就是输出结果。

其整体的计算过程如下图所示：

实现代码1

class Muti_Head_Self_Attention(nn.Module):
    def __init__(self,input_dim,dim_k,dim_v,nums_head):
        super(Muti_Head_Self_Attention,self).__init__()
        #在实现多头注意力机制中，需要判断dim_k和dim_v能否被nums_head均分
        assert dim_k % nums_head == 0   
        assert dim_v % nums_head == 0
        
        self.q = nn.Linear(input_dim,dim_k,bias=False)    
        self.k = nn.Linear(input_dim,dim_k,bias=False)   
        self.v = nn.Linear(input_dim,dim_v,bias=False)    
        self.fc = nn.Linear(dim_v, dim_v, bias=False)
        self.nums_head = nums_head
        self._norm_fact = 1 / ((dim_k//nums_head)**0.5)
        '''
可以采用以下的初始化方式，来规避 assert 操作。

        self.nums_head=nums_head
        self.q = nn.Linear(input_dim,dim_k*nums_head,bias=False)    
        self.k = nn.Linear(input_dim,dim_k*nums_head,bias=False)   
        self.v = nn.Linear(input_dim,dim_v*nums_head,bias=False)    
        self.fc = nn.Linear(dim_v*nums_head, dim_v*nums_head, bias=False)
        self._norm_fact = 1 / ((dim_k//nums_head)**0.5)
'''
    def forward(self,x):
        # input : batch_size * seq_len * input_dim
        
        #两次形状变化，1.对Q、K、V按最后一个维度进行切分 2.对切分后的中间两维进行交换
        #(batch_size,seq_len,dim_k)->(batch_size,seq_len,nums_head,dim_k//nums_head)->(batch_size,nums_head,seq_len,dim_k//nums_head)   ,其中dim_k同dim_v
        Q=self.q(x).reshape(x.shape[0], x.shape[1], self.nums_head, -1).permute(0, 2, 1, 3)
        K=self.k(x).reshape(x.shape[0], x.shape[1], self.nums_head, -1).permute(0, 2, 1, 3)
        V=self.v(x).reshape(x.shape[0], x.shape[1], self.nums_head, -1).permute(0, 2, 1, 3)

        atten = torch.softmax(torch.matmul(Q,K.transpose(-2, -1)) * self._norm_fact,dim=-1)
        
        #对于atten的三次形状变化，1.将atten与V进行矩阵相乘 2.将数据的中间两维进行交换 3.恢复到切分前的形状
        #(batch_size,nums_head,seq_len,seq_len)->(batch_size,nums_head,seq_len,dim_v//nums_head)->(batch_size,seq_len,nums_head,dim_v//nums_head)->(batch_size,seq_len,dim_v)
        output= torch.matmul(atten,V).permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[1],-1)
        return self.fc(output)

#使用示例
torch.random.manual_seed(420)
nums_head=2
X=torch.randn(1,3,4)
Attention1=Muti_Head_Self_Attention(4,6,6,nums_head)
Y1=Attention1(X)

如有错误之处,欢迎留言指正。

参考资料

浅谈Transformer的初始化、参数化与标准化 为什么除于k**0.5
超详细图解Self-Attention
李宏毅b站课程视频
动手学深度学习教材
The Annotated Transformer
注意力机制——注意力评分函数（代码+详解）
多头注意力机制 +代码解读
Self-Attention自注意可视化介绍
工具Markdown使用