Self-Attention

发布时间 2023-08-27 23:03:09作者: Netsharp

Self-Attention

在Attention is all you need这篇论文中,可以看到这样一个公式:
$Attention(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_k} } )V$

1.定义input

在进行Self - Attention之前,我们首先定义3个1×4的input。 pytorch代码如下:

import torch

x = torch.FloatTensor([
    [1, 0, 1, 0],  # input 1
    [0, 2, 0, 2],  # input 2
    [1, 1, 1, 1]   # input 3
    ])

2.初始化权重

每个input和三个权重矩阵分别相乘会得到三个新的矩阵,分别是key(橙色),query(红色),value(紫色)。我们已经令input的shape为1×4,key、query、value的shape为1×3,因此可以推出与input相乘的权重矩阵的shape为4×3。 代码如下:

w_key = torch.FloatTensor([
    [0, 0, 1],
    [1, 1, 0],
    [0, 1, 0],
    [1, 1, 0]
    ])

w_query = torch.FloatTensor([
    [1, 0, 1],
    [1, 0, 0],
    [0, 0, 1],
    [0, 1, 1]
    ])

w_value = torch.FloatTensor([
    [0, 2, 0],
    [0, 3, 0],
    [1, 0, 3],
    [1, 1, 0]
    ])

3.计算key,query和value

现在我们计算key,query和value矩阵的值,计算的过程也很简单,运用矩阵乘法即可:

  • $key = input \times w_{key}$
  • $query = input \times w_{query} $
  • $value = input \times w_{value} $

value = input * w_value

    keys = x @ w_key
    # tensor([[0., 1., 1.],
    #         [4., 4., 0.],
    #         [2., 3., 1.]])

    querys = x @ w_query
    # tensor([[1., 0., 2.],
    #         [2., 2., 2.],
    #         [2., 1., 3.]])

    values = x @ w_value
    # tensor([[1., 2., 3.],
    #         [2., 8., 0.],
    #         [2., 6., 3.]])

即得到了数学公式中的$Q,K,V$

4.计算attention scores

例如:为了获得input1的注意力分数(attention scores),我们将input1的query(红色)与input1、2、3的key(橙色)的转置分别作点积,得到3个attention scores(蓝色)。 同理,我们也可以得到input2和input3的attention scores。

    attn_scores = querys @ keys.T

    # tensor([[ 2.,  4.,  4.],  # attention scores from Query 1
    #         [ 4., 16., 12.],  # attention scores from Query 2
    #         [ 4., 12., 10.]]) # attention scores from Query 3

对应数学公式中的:$QK^{T}$

5.对attention scores作softmax

上一步得到了attention scores矩阵后,我们对attention scores矩阵作softmax计算。softmax的作用为归一化,使得其中各项相加后为1。这样做的好处是凸显矩阵中最大的值并抑制远低于最大值的其他分量。

    from torch.nn.functional import softmax

    attn_scores_softmax = softmax(attn_scores, dim=-1)
    # tensor([[6.3379e-02, 4.6831e-01, 4.6831e-01],
    #         [6.0337e-06, 9.8201e-01, 1.7986e-02],
    #         [2.9539e-04, 8.8054e-01, 1.1917e-01]])

    attn_scores_softmax = [
    [0.0, 0.5, 0.5],
    [0.0, 1.0, 0.0],
    [0.0, 0.9, 0.1]
    ]
    attn_scores_softmax = torch.tensor(attn_scores_softmax)
    # tensor([[0.0000, 0.5000, 0.5000],
    #         [0.0000, 1.0000, 0.0000],
    #         [0.0000, 0.9000, 0.1000]])    

对应数学公式中的:$softmax(QK^{T})$

6.将attention scores与values相乘

每个score(蓝色)乘以其对应的value(紫色)得到3个alignment vectors(黄色)。在本教程中,我们将它们称为weighted values(加权值)。

    weighted_values = values[:,None] * attn_scores_softmax.T[:,:,None]
    # tensor([[[0.0000, 0.0000, 0.0000],
    #          [0.0000, 0.0000, 0.0000],
    #          [0.0000, 0.0000, 0.0000]],

    #         [[1.0000, 4.0000, 0.0000],
    #          [2.0000, 8.0000, 0.0000],
    #          [1.8000, 7.2000, 0.0000]],

    #         [[1.0000, 3.0000, 1.5000],
    #          [0.0000, 0.0000, 0.0000],
    #          [0.2000, 0.6000, 0.3000]]])

对应数学公式中的:$softmax(QK^{T})V$

7.对weighted values求和得到output

从图中可以看出,每个input生成3个weighed values(黄色),我们将这3个weighted values相加,得到output(深绿)。图中一共有3个input,所以最终生成3个output。

    outputs = weighted_values.sum(dim=0)

    # tensor([[2.0000, 7.0000, 1.5000],  # Output 1
    #         [2.0000, 8.0000, 0.0000],  # Output 2
    #         [2.0000, 7.8000, 0.3000]]) # Output 3

7. 为什么要除以$\sqrt{d_k}$?

  • 首先,$d_k$表示的是词向量的维度。我们除以$\sqrt{d_k}$ 是为了防止$ QK^{T} $值过大,导致softmax计算时上溢出(overflow)。
  • 其次,使用$d_k$可以使$ QK^{T} $的结果满足期望为0,方差为1的分布。

8. 这个公式的意义?

先从 $QK^{T}$ 看起,从几何角度看,点积是两个向量的长度与它们夹角余弦的积。如果两向量夹角为90°,那么结果为0,代表两个向量线性无关。如果两个向量夹角越小,两向量在方向上相关性也越强,结果也越大。点积反映了两个向量在方向上的相似度,结果越大越相似。
图片名称
对 $QK^{T}$进行相似度的计算后,再使用softmax归一化。最后将归一化的结果与 $ V $ 作乘法,计算的结果就是输入经过注意力机制加权求和之后的表示。