Pytorch 自注意力和位置编码

环境使用 Kaggle 里免费建立的 Notebook

教程使用李沐老师的 动手学深度学习 网站和 视频讲解

小技巧：当遇到函数看不懂的时候可以按 Shift+Tab 查看函数详解。

y i = f ( x i , ( x 1 , x 1 ) , … , ( x n , x n ) ) ∈ R d mathbf{y}_i = f(mathbf{x}_i, (mathbf{x}_1, mathbf{x}_1), ldots, (mathbf{x}_n, mathbf{x}_n)) in mathbb{R}^d yi​=f(xi​,(x1​,x1​),…,(xn​,xn​))∈Rd

• 给定序列 x 1 , … , x n mathbf{x}_1, ldots, mathbf{x}_n x1​,…,xn​，其中任意 x i ∈ R d ( 1 ≤ i ≤ n ) mathbf{x}_i in mathbb{R}^d (1 leq i leq n) xi​∈Rd(1≤i≤n)。
• 自注意力Pooling 将 x i mathbf{x}_i xi​ 当做 Key，Value，Query 来对序列抽取特征得到 y 1 , … , y n mathbf{y}_1, ldots, mathbf{y}_n y1​,…,yn​。

与 CNN，RNN 对比：

• 和 CNN/RNN 不同，自注意力并没有记录位置信息
• 位置编码将位置信息注入到输入里
  • 假设长度为 n n n 的序列 X ∈ R n × d mathbf{X} in mathbb{R}^{n times d} X∈Rn×d，那么使用位置编码矩阵 P ∈ R n × d mathbf{P} in mathbb{R}^{n times d} P∈Rn×d 来输出 X + P mathbf{X} + mathbf{P} X+P 作为自编码输入
• P mathbf{P} P 的元素如下计算：
  p i , 2 j = sin ⁡ ( i 1000 0 2 j / d ) p i , 2 j + 1 = cos ⁡ ( i 1000 0 2 j / d ) begin{aligned} p_{i, 2j} &= sinleft(frac{i}{10000^{2j/d}}right)\p_{i, 2j+1} &= cosleft(frac{i}{10000^{2j/d}}right)end{aligned} pi,2j​pi,2j+1​​=sin(100002j/di​)=cos(100002j/di​)​

在位置嵌入矩阵 P mathbf{P} P 中，行代表词元在序列中的位置，列代表位置编码的不同维度。第 6 6 6 列和第 7 7 7 列的频率高于第 8 8 8 列和第 9 9 9 列。第 6 6 6 列和第 7 7 7 列之间的偏移量是由于正弦函数和余弦函数的交替。
每个样本表示的向量不一样，一个样本不同维度有不同的值。

!pip install -U d2l
import math
import torch
from torch import nn
from d2l import torch as d2l

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                                   num_hiddens, num_heads, 0.5)
attention.eval()

batch_size, num_queries, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
attention(X, X, X, valid_lens).shape

class PositionalEncoding(nn.Module):
    """位置编码"""
    # num_hiddens:向量长度，max_len：最长序列长度
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)

forward() 函数中的 X 是单词的 Embedding。

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',
         figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])

Q：该位置编码是自动学习的吗？
A：不是，BERT 里面的位置编码是可学习的。

Q：自注意力机制是不是可以理解为一个网络层，就像 CNN 或 RNN 的一层来理解？
A：可以的。