[toc]

回顾 Attention

Q、K 与 V

Scaled-Dot-Product Attention

$$
Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V=AV
$$

${d_k}$是 key 和 value 的维度，当${d_k}$越大，则$QK^T$越大，经过$softmax$函数可能会在梯度极小的区域，为了避免梯度消失的问题，scaled-dot-product 的操作是将$QK^T$除以$\sqrt{d_k}$，但仍然$QK^T$保持方差为 1。

在 Transformer 的源码中，scale-dot-product attention 的实现如下：

1
2
3
4
5
6
7
8
9
10
11
def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

Self-Attention

自注意力机制（Self-attention）考虑的是输入元素之间的相关性，而非输入元素和输出元素之间的相关性，其目标序列和输入的原始序列相同。

图1 self-attention原理图

Multi-Head Self-Attention

图2 多头注意力机制

多头注意力机制让模型能够在不同的位置上共同关注来自不同 representation 的子空间的信息。多个平行的 attention 层分别获得各自的 Q、K、V：

$$
MultiHead(Q,K,V)=Concat(head_1,…,head_h)W^O\space where\space head_i=Attention(QW_i^Q,KW_i^K,VW_I^V)
$$

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h # heads论文中取8
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k
        query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]

        # 2) Apply attention on all the projected vectors in batch.
        # self-attention中query、key、value全部相同，size为[batch_size,len,d_model]
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        # 还原序列
        return self.linears[-1](x)

Transformer(vanilla Transformer)

Encoder

  • 层数$N=6$，每层有两个 sub-layers，分别是多头自注意力机制和根据位置的全连接前馈网络
  • 两个子层之前使用残差连接[11]和正则化
  • 生成 K、V 矩阵

Decoder

  • 层数$N=6$，在 encoder 的两个子层之间加入第三层

  • 对 encoder 的输出使用多头注意力机制

  • 对 decoder 的第一层多头注意力层进行了修改->masking multi-head attention 保证第一层注意力机制在预测当前位置的单词时只依赖当前位置之前的单词

  • 生成 Q 矩阵

Positional Encoding

  • $d_{model}$是词 embedding 的维度，论文中取 512。
  • 偶数位置：$PE_{(pos,2i)}=sin(pos/10000^{2i/d_{model}})$
  • 奇数位置：$PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}})$

绝对位置向量中蕴含着相对位置的信息，相对位置会在注意力机制中消失。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model) #d_model论文中取512
        position = torch.arange(0, max_len).unsqueeze(1) #维度增加
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model)) #相对位置公式
        pe[:, 0::2] = torch.sin(position * div_term) #奇数列
        pe[:, 1::2] = torch.cos(position * div_term) #偶数列
        pe = pe.unsqueeze(0) #增加维度
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)],
                         requires_grad=False)
        return self.dropout(x)

Positional-wise Feed-Forward Network

$$
FFN(x)=max(0,xW_1+b_1)W_2+b_2
$$

FFN 将每个位置的 attention 结果映射到一个更大维度的特征空间，然后用 ReLU 引入非线性层进行筛选，最后恢复原始维度。

1
2
3
4
5
6
7
8
9
10
class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

Residual Connection and Normalization

残差结构：输出的 Z 和经过位置编码的 X 对位相加，作为输出。

Layer Normalization：不需要像 Batch Normalization 一样考虑所有 batch，只需要考虑同一个 example 中的不同 feature 计算均值和方差，然后对 example 的向量进行 normalization。

以下是实现的代码：

在子层与子层之间进行了残差连接与归一化，子层的输出为$LayerNorm(x+Sublayer(x))$。tensor2tensor 的实现是$x+SubLayer(LayerNorm(x))$。

1
2
3
4
5
6
7
8
9
10
11
12
class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
1
2
3
4
5
6
7
8
9
10
11
12
13
class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))

Training

  • teacher forcing: 使用正确的 truth 作为 input。

Transformer 中的注意力机制

  • self-attention in encoder and decoder: $Q=K=V=X$
  • masked self-attention: 在 transformer 的 decoder 中，masked 是保证 attention 加权计算时忽略当前位置后面的单词，保证信息来源于当前位置以及之前的位置。
  • cross-attention: keys 和 values 来源于最后一层的 encoder 的输入。

Transformer Family

### Transformer-XL（ACL 2019）

标准 Transformer 有一个固定有限的 attention 范围。在每一个更新的步中，模型只能够处理同一 segment 中的其他元素，信息无法在分离的 segment 中传递，也就意味着 Transformer 很难捕捉长期的依赖关系，在上下文较短的情况下，难以进行预测。同时，当 segment 每右移一位时，新的 segment 会重新处理，尽管有重叠的标记。

Transformer-XL (Dai et al., 2019)进行了如下的改进：

  • 重用segment 之间的隐藏状态。
  • 采用新的适合重用状态的位置编码。

隐藏状态重用

Transformer-XL 通过连续使用前一个 segment 的隐藏状态，将 segment 之间的递归引入到模型中。如下图，一个 segment 的长度为 4，可以看出 Transformer 和 Transformer-XL 的区别。

图1 Transformer和Transformer-XL的对比（图片来源：Dai et al.,2019）

第$(\tau + 1)$段的第 n 层的隐藏状态记为 $\mathbf{h}{\tau+1}^{(n)} \in \mathbb{R}^{L \times d}$，第 n-1 层的隐藏状态为 $\mathbf{h}{\tau+1}^{(n-1)}$。前一个段的第 n 层隐藏状态为$\mathbf{h}_{\tau}^{(n)}$，模型有如下的公式，通过合并之前的隐藏状态的信息，实现注意力范围的延长。

$$\begin{aligned} {\widetilde{\mathbf{h}}{\tau+1}^{(n-1)}} &= [\text{stop-gradient}(\mathbf{h}{\tau}^{(n-1)}) \circ \mathbf{h}{\tau+1}^{(n-1)}] \ \mathbf{Q}{\tau+1}^{(n)} &= \mathbf{h}{\tau+1}^{(n-1)}\mathbf{W}^q \ \mathbf{K}{\tau+1}^{(n)} &= {\widetilde{\mathbf{h}}{\tau+1}^{(n-1)}} \mathbf{W}^k \ \mathbf{V}{\tau+1}^{(n)} &= {\widetilde{\mathbf{h}}{\tau+1}^{(n-1)}} \mathbf{W}^v \ \mathbf{h}{\tau+1}^{(n)} &= \text{transformer-layer}(\mathbf{Q}{\tau+1}^{(n)}, \mathbf{K}{\tau+1}^{(n)}, \mathbf{V}_{\tau+1}^{(n)}) \end{aligned}$$

key 和 value 都依赖于扩展后的隐藏状态，而 query 只依赖于当前的隐藏状态。

相对位置编码

Transformer-XL 提出了一种新的位置编码来适应隐藏状态重用。在基础的 Transformer 中，编码绝对位置的话，前一段和当前段将会被相同的编码。在 Transformer-XL 中并不需要。为了保持段之间的位置信息的流动，Transformer-XL 提出了相对位置编码，它只要知道位置的偏移量就可以做出预测。

$i$位置的 query 和$j$位置的 key 的 attention score 为：

$$\begin{aligned} a_{ij} &= \mathbf{q}_i {\mathbf{k}_j}^\top = (\mathbf{x}_i + \mathbf{p}_i)\mathbf{W}^q ((\mathbf{x}_j + \mathbf{p}_j)\mathbf{W}^k)^\top \ &= \mathbf{x}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{x}_j^\top + \mathbf{x}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{p}_j^\top + \mathbf{p}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{x}_j^\top + \mathbf{p}_i\mathbf{W}^q {\mathbf{W}^k}^\top\mathbf{p}_j^\top \end{aligned}$$

整理可得：

$$a_{ij}^\text{rel} = \underbrace{ \mathbf{x}i\mathbf{W}^q { {\mathbf{W}_E^k}^\top } \mathbf{x}_j^\top }_\text{content-based addressing} + \underbrace{ \mathbf{x}_i\mathbf{W}^q { {\mathbf{W}_R^k}^\top } {\mathbf{r}{i-j}^\top} }\text{content-dependent positional bias} + \underbrace{ {\mathbf{u}} { {\mathbf{W}_E^k}^\top } \mathbf{x}_j^\top }_\text{global content bias} + \underbrace{ {\mathbf{v}} { {\mathbf{W}_R^k}^\top } {\mathbf{r}{i-j}^\top} }_\text{global positional bias}$$

其中：

  • $\mathbf{p}j$ 被替换为相对位置编码 $\mathbf{r}{i-j} \in \mathbf{R}^{d}$;
  • $\mathbf{p}_i\mathbf{W}^q$ 被替换为两个可训练的参数 ${u}$ 和 ${v}$ ，分别用于计算 content 和 location 信息
  • $\mathbf{W}^k$ 被分成两个矩阵， $\mathbf{W}^k_E$ 处理 content 信息， $\mathbf{W}^k_R$ 处理 location 信息。

Reformer（ICLR 2020）

Reformer(Kitaev, et al. 2020)主要解决了 Transformer 以下的几个问题：

  • 具有$N$层的模型中的内存是单层模型中的$N$倍。
  • 中间 Feed Forward 层通常相当大。
  • 长度为$L$的序列上的 attention matrix 通常需要$O（L^2）$的内存和时间。

Reformer 主要的改进为：

  • 将 dot-product attention 替换为 locality-sensitive hashing (LSH) attention，将复杂度从 $O(L^2)$ 降低到$O(L\log L)$。
  • 将残差模块替换为reversible 的残差层，在训练期间仅存储一次激活（activation），而不是$N$次。
图2. LSH attention原理 (图片来源：Kitaev, et al. 2020）
图3. LSH attention的计算步骤 (图片来源：Kitaev, et al. 2020）

Sparse Transformer

(Child et al., 2019) 通过稀疏矩阵分解，提出了factorized self-attention，使得在 16384 的序列长度上训练数百层的密集注意力网络成为可能，否则在硬件上是不可行的。

给定一组注意力连接模式（attention connectivity pattern） $\mathcal{S} = {S_1, \dots, S_n}$, 其中每个 $S_i$ 记录了第 i 个 query vector 的注意到的一组 key position，有：

$$\begin{aligned} \text{Attend}(\mathbf{X}, \mathcal{S}) &= \Big( a(\mathbf{x}i, S_i) \Big){i \in {1, \dots, L}} \text{ where } a(\mathbf{x}i, S_i) = \text{softmax}\Big(\frac{(\mathbf{x}_i \mathbf{W}^q)(\mathbf{x}_j \mathbf{W}^k){j \in S_i}^\top}{\sqrt{d_k}}\Big) (\mathbf{x}j \mathbf{W}^v){j \in S_i} \end{aligned}$$

在自回归模型中，一个注意力范围（attention span ）被定义为 $S_i = {j: j \leq i}$ ，因为它允许每个 token 关注过去的所有位置。

在 factorized self-attention 中, $S_i$ 会被分解成一个依赖关系的树，对每一对 $(i, j)$ 其中$j \leq i$, 都有一条将$i$ 连接到 $j$ 的路径 ，$i$ 可以直接或间接地注意到 $j$。

$S_i$被分成 p 个非重叠的子集，第 m 个子集表示为$A^{(m)}_i \subset S_i, m = 1,\dots, p$，这样输出位置 i 和任意一个 j 的最大距离为$p+1$。

参考文献

Transformer/NMT 相关项目、源码

项目 来源 文档 Github 描述
Transformers hugging face [doc] [source code]
OpenNMT Harvard NLP group and SYSTRAN [doc] [source code]
Tensor2Tensor Google Brain team [Readme] [source code] 
1
pip install -U sacremoses

参考文献

[1] Li Hungyi. “Self-Attention” lecture PPT,2021.

[2] Ashish Vaswani, et al. “Attention is all you need.” NIPS,2017.

[3] Alexander Rush. “The Annotated Transformer” 2018.

1
onmt_translate -model toy-ende/trans9_step_5000.pt -src toy-ende/src-test.txt -output toy-ende/pred_trans9.txt -gpu 0 -verbose