Transformer解读（附pytorch代码）

作者: Cingti | 来源:发表于2020-05-08 13:05 被阅读0次

Transformer早在2017年就出现了，直到BERT问世，Transformer开始在NLP大放光彩，目前比较好的推进就是Transformer-XL（后期附上）。这里主要针对论文和程序进行解读，如有不详实之处，欢迎指出交流，如需了解更多细节之处，推荐知乎上川陀学者写的。本文程序的git地址在这里。程序如果有不详实之处，欢迎指出交流~

前言

2017年6月，Google发布了一篇论文《Attention is All You Need》，在这篇论文中，提出了 Transformer 的模型，其旨在全部利用Attention方式来替代掉RNN的循环机制，从而通过实现并行化计算提速。在Transformer出现之前，RNN系列网络以及seq2seq+attention架构基本上铸就了所有NLP任务的铁桶江山。由于Attention模型本身就可以看到全局的信息， Transformer实现了完全不依赖于RNN结构仅利用Attention机制，在其并行性和对全局信息的有效处理上获得了比之前更好的效果。

Transformer的整体结构

图1：Transformer整体结构
Transformer的整体结构就是分成编码器和解码器两部分，并且两部分之间是有联系的，可以注意到编码器的输出是解码器第二个Multi-head Attention中和的输入，这里，我们把编码器的输出称为state用来初始化解码器的状态，而实际上对于解码器而言，每一层的解码器的state是一样的(都是编码器的输出)，并不会像RNN中的state一样改变。对应的pytorch程序如下：

class transformer(nn.Module):
    def __init__(self, enc_net, dec_net):
        super(transformer, self).__init__()
        self.enc_net = enc_net   # TransformerEncoder的对象   
        self.dec_net = dec_net   # TransformerDecoder的对象
    
    def forward(self, enc_X, dec_X, valid_length=None, max_seq_len=None):
        """
        enc_X: 编码器的输入
        dec_X: 解码器的输入
        valid_length: 编码器的输入对应的valid_length,主要用于编码器attention的masksoftmax中，
                      并且还用于解码器的第二个attention的masksoftmax中
        max_seq_len:  位置编码时调整sin和cos周期大小的，默认大小为enc_X的第一个维度seq_len
        """
        
        # 1、通过编码器得到编码器最后一层的输出enc_output
        enc_output = self.enc_net(enc_X, valid_length, max_seq_len)
        # 2、state为解码器的初始状态，state包含两个元素，分别为[enc_output, valid_length]
        state = self.dec_net.init_state(enc_output, valid_length)
        # 3、通过解码器得到编码器最后一层到线性层的输出output，这里的output不是解码器最后一层的输出，而是
        #    最后一层再连接线性层的输出
        output = self.dec_net(dec_X, state)
        return output

纵观图1整个Transformer的结构，其核心模块其实就是三个：Multi-Head attention、Feed Forward 以及 Add&Norm。这里关于Multi-Head attention部分只讲程序的实现，关于更多细节原理，请移至简书开头推荐的知乎链接。

Multi-Head Attention实现

Transformer中的attention采用的是多头的self-attention结构，并且在编码器中，由于不同的输入mask的部分不一样，因此在softmax之前采用了mask操作，并且解码时由于不能看到t时刻之后的数据，同样在解码器的第一个Multi-Head attention中采用了mask操作，但是二者是不同的。因为编码器被mask的部分是需要在输入到Transformer之前事先确定好，而解码器第一个Multi-Head attention被mask的部分其实就是从t=1时刻开始一直到t=seq_len结束，对应于图2。在图2中，横坐标表示解码器一个batch上的输入序列长度(也就是t)，紫色部分为被mask的部分，黄色部分为未被mask的部分，可以看出，随着t的增加，被mask的部分逐一减少。而解码器第二个Multi-Head attention的mask操作和编码器中是一样的。

图2：解码器第一个Multi-Head attention中的mask操作

mask+softmax程序如下：

def masked_softmax(X, valid_length, value=-1e6):
    # 如果valid_length是一维的:valid_length的维度等于batch_size的大小
    # 对每一个batch去确定一个valid_length,因此valid_length的维度与batch_size大小相同
    # 再将valid_length内的元素通过repeat操作将valid_length内的元素repeat seq_len（X.size()[1]）次
    # 结果就是对每一个batch上的X根据valid_length输出相应的attention weights，因此一个batch上的attention weights是一样的

    # 如果valid_length是二维的:valid_length的维度等于[batch_size, seq_length]
    # 此时是针对每一个batch的每一句话都设置了seq_length
    if valid_length is None:
        return F.softmax(X, dim=-1)
    else:
        X_size = X.size()
        device = valid_length.device
        if valid_length.dim() == 1:
            valid_length = torch.tensor(valid_length.cpu().numpy().repeat(X_size[1], axis=0),
                                        dtype=torch.float, device=device) if valid_length.is_cuda \
                else torch.tensor(valid_length.numpy().repeat(X_size[1], axis=0),
                                  dtype=torch.float, device=device)
        else:
            valid_length = valid_length.view([-1])
        X = X.view([-1, X_size[-1]])
        max_seq_length = X_size[-1]
        valid_length = valid_length.to(torch.device('cpu'))
        mask = torch.arange(max_seq_length, dtype=torch.float)[None, :] >= valid_length[:, None]
        X[mask] = value
        X = X.view(X_size)
        return F.softmax(X, dim=-1)

mask操作其实就是对于无效的输入，用一个负无穷的值代替这个输入，这样在softmax的时候其值就是0。而在attention中(attention操作见下式)，softmax的操作出来的结果其实就是attention weights，当attention weights为0时，表示不需要attention该位置的信息。
$softmax(\frac{QK^{T}}{\sqrt{d}})V$
对于Multi-Head attention的实现，其实并没有像论文原文写的那样，逐一实现多个attention，再将最后的结果concat，并且通过一个输出权重输出。下面通过程序和公式讲解一下实际的实现过程，这里假设 $Q$ , $K$ , $V$ 的来源是一样的，都是 $X$ ，其维度为[batch_size, seq_len, input_size]。(需要注意的是在解码器中第二个Multi-Head的输入中 $Q$ 与 $KV$ 的来源不一样)

图3：论文原文中的attention操作

class DotProductAttention(nn.Module):
    # 经过DotProductAttention之后，输入输出的维度是不变的，都是[batch_size*h, seq_len, d_model//h]
    def __init__(self, dropout,):
        super(DotProductAttention, self).__init__()
        self.drop = nn.Dropout(dropout)

    def forward(self, Q, K, V, valid_length):
        # Q, K, V shape:[batch_size*h, seq_len, d_model//h]
        d_model = Q.size()[-1]  # int
        # torch.bmm表示批次之间（>2维）的矩阵相乘
        attention_scores = torch.bmm(Q, K.transpose(1, 2))/math.sqrt(d_model)
        # attention_scores shape: [batch_size*h, seq_len, seq_len]
        attention_weights = self.drop(masked_softmax(attention_scores, valid_length))
        return torch.bmm(attention_weights, V)  # [batch_size*h, seq_len, d_model//h]

class MultiHeadAttention(nn.Module):
    def __init__(self, input_size, hidden_size, num_heads, dropout,):
        super(MultiHeadAttention, self).__init__()
        # 保证MultiHeadAttention的输入输出tensor的维度一样
        assert hidden_size % num_heads == 0
        # hidden_size => d_model
        self.num_heads = num_heads
        # num_heads => h
        self.hidden_size = hidden_size
        # 这里的d_model为中间隐层单元的神经元数目,d_model=h*d_v=h*d_k=h*d_q
        self.Wq = nn.Linear(input_size, hidden_size, bias=False)
        self.Wk = nn.Linear(input_size, hidden_size, bias=False)
        self.Wv = nn.Linear(input_size, hidden_size, bias=False)
        self.Wo = nn.Linear(hidden_size, hidden_size, bias=False)
        self.attention = DotProductAttention(dropout)

    def _transpose_qkv(self, X):
        # X的输入维度为[batch_size, seq_len, d_model]
        # 通过该函数将X的维度改变成[batch_size*num_heads, seq_len, d_model//num_heads]
        self._batch, self._seq_len = X.size()[0], X.size()[1]
        X = X.view([self._batch, self._seq_len, self.num_heads, self.hidden_size//self.num_heads])  # [batch_size, seq_len, num_heads, d_model//num_heads]
        X = X.permute([0, 2, 1, 3])  # [batch_size, num_heads, seq_len, d_model//num_heads]
        return X.contiguous().view([self._batch*self.num_heads, self._seq_len, self.hidden_size//self.num_heads])

    def _transpose_output(self, X):
        X = X.view([self._batch, self.num_heads, -1, self.hidden_size//self.num_heads])
        X = X.permute([0, 2, 1, 3])
        return X.contiguous().view([self._batch, -1, self.hidden_size])

    def forward(self, query, key, value, valid_length):
        Q = self._transpose_qkv(self.Wq(query))
        K = self._transpose_qkv(self.Wk(key))
        V = self._transpose_qkv(self.Wv(value))
        # 由于输入的valid_length是相对batch输入的，而经过_transpose_qkv之后,
        # batch的大小发生了改变,Q的第一维度由原来的batch改为batch*num_heads
        # 因此,需要对valid_length进行复制,也就是进行np.title的操作
        if valid_length is not None:
            device = valid_length.device
            valid_length = valid_length.cpu().numpy() if valid_length.is_cuda else valid_length.numpy()
            if valid_length.ndim == 1:
                valid_length = np.tile(valid_length, self.num_heads)
            else:
                valid_length = np.tile(valid_length, [self.num_heads, 1])
            valid_length = torch.tensor(valid_length, dtype=torch.float, device=device)
        output = self.attention(Q, K, V, valid_length)
        output_concat = self._transpose_output(output)
        return self.Wo(output_concat)

首先，对于输入 $X$ ，通过三个权重变量得到 $Q$ , $K$ , $V$ ，此时三者维度相同，都是[batch, seq_len, d_model]，然后对其进行维度变换：[batch, seq_len, h, d_model//h]==>[batch, h, seq_len, d]==>[batch×h, seq_len, d]，其中d=d_model//h，因此直接将变换后的 $Q$ , $K$ , $V$ 直接做DotProductAttention就可以实现Multi-Head attention，最后只需要将DotProductAttention输出的维度依次变换回去，然后乘以输出权重就可以了。关于程序中的参数valid_length已在程序中做了详细的解读，这里不再赘述，注意的是输入的valid_length是针对batch这个维度的，而实际操作中由于X的batch维度发生了改变(由batch变成了batch×h)，因此需要对valid_length进行复制。

PositionWiseFFN的实现

FFN的实现是很容易的，其实就是对输入进行第一个线性变换，其输出加上ReLU激活函数，然后在进行第二个线性变换就可以了。

class PositionWiseFFN(nn.Module):
    # y = w*[max(0, wx+b)]x+b
    def __init__(self, input_size, fft_hidden_size, output_size,):
        super(PositionWiseFFN, self).__init__()
        self.FFN1 = nn.Linear(input_size, fft_hidden_size)
        self.FFN2 = nn.Linear(fft_hidden_size, output_size)

    def forward(self, X):
        return self.FFN2(F.relu(self.FFN1(X)))

Add&Norm的实现

Add&norm的实现就是利用残差网络进行连接，最后将连接的结果接上LN，值得注意的是，程序在Y的输出中加入了dropout正则化。同样的正则化技术还出现在masked softmax之后和positional encoding之后。

class AddNorm(nn.Module):
    def __init__(self, hidden_size, dropout,):
        super(AddNorm, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.LN = nn.LayerNorm(hidden_size)

    def forward(self, X, Y):
        assert X.size() == Y.size()
        return self.LN(self.drop(Y) + X)

positional encoding

positional encoding的实现很简单，其实就是对输入序列给定一个唯一的位置，采用sin和cos的方式给了一个位置编码，其中sin处理的是偶数位置，cos处理的是奇数位置。但是，这一块的工作确实非常重要的，因为对于序列而言最主要的就是位置信息，显然BERT是没有去采用positional encoding(尽管在BERT的论文里有一个Position Embeddings的输入，但是显然描述的不是Transformer中要描述的位置信息)，后续BERT在这一方面的改进工作体现在了XLNet中(其采用了Transformer-XL的结构），后续的简书中再介绍该部分的内容。

class PositionalEncoding(nn.Module):
    def __init__(self, dropout,):
        super(PositionalEncoding, self).__init__()

    def forward(self, X, max_seq_len=None):
        if max_seq_len is None:
            max_seq_len = X.size()[1]
        # X为wordEmbedding的输入,PositionalEncoding与batch没有关系
        # max_seq_len越大,sin()或者cos()的周期越小,同样维度
        # 的X,针对不同的max_seq_len就可以得到不同的positionalEncoding
        assert X.size()[1] <= max_seq_len
        # X的维度为: [batch_size, seq_len, embed_size]
        # 其中: seq_len = l, embed_size = d
        l, d = X.size()[1], X.size()[-1]
        # P_{i,2j}   = sin(i/10000^{2j/d})
        # P_{i,2j+1} = cos(i/10000^{2j/d})
        # for i=0,1,...,l-1 and j=0,1,2,...,[(d-2)/2]
        max_seq_len = int((max_seq_len//l)*l)
        P = np.zeros([1, l, d])
        # T = i/10000^{2j/d}
        T = [i*1.0/10000**(2*j*1.0/d) for i in range(0, max_seq_len, max_seq_len//l) for j in range((d+1)//2)]
        T = np.array(T).reshape([l, (d+1)//2])
        if d % 2 != 0:
            P[0, :, 1::2] = np.cos(T[:, :-1])
        else:
            P[0, :, 1::2] = np.cos(T)
        P[0, :, 0::2] = np.sin(T)
        return torch.tensor(P, dtype=torch.float, device=X.device)

编码器实现和解码器的实现

无论是编码器还是解码器，其实都是用上面说的三个基本模块堆叠而成，具体的实现细节大家可以看简书开头的git地址，这里需要强调的是以下几点：

无论是编码器还是解码器，都在word embedding后面乘上 $\sqrt{d_{model}}$ ，防止其值过小；
论文里面提到了他们用的优化器，是以 $\beta_1=0.9$ , $\beta_2=0.98$ 和 $\epsilon=10^{-9}$ 的Adam为基础，而后使用一种warmup的学习率调整方式来进行调节。具体公式如下：基本上就是先用一个固定warmup_steps进行学习率的线性增长，而后到达warmup_steps之后会随着step_num的增长而逐渐减小。
$l_{rate}=d_{model}^{-0.5}*min(step\_num^{-0.5},step\_num*warmup\_steps^{-1.5})$

class NoamOpt:
    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer    # 优化器
        self._step = 0                # 步长
        self.warmup = warmup          # warmup_steps
        self.factor = factor          # 学习率因子(就是学习率前面的系数)
        self.model_size = model_size  # d_model
        self._rate = 0                # 学习率

    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()

    def rate(self, step=None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * \
               (self.model_size ** (-0.5) *
                min(step ** (-0.5), step * self.warmup ** (-1.5)))

简书中出现的程序都在简书开头的git中了，直接执行main.ipynb就可以运行程序，如有不详实之处，还请指出~~~

Transformer解读（附pytorch代码）

前言

Transformer的整体结构

Multi-Head Attention实现

PositionWiseFFN的实现

Add&Norm的实现

positional encoding

编码器实现和解码器的实现

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读