Transformer
由论文Attention Is All You Need提出,该模型舍弃了RNN和CNN,完全基于注意力机制实现了高效、并行化的训练,成为了NLP领域的里程碑式模型。本篇文章介绍Transformer
的模型结构并基于PyTorch
完成模型的搭建。
本文目录
- [0 引言](#0 引言)
- [1 输入部分](#1 输入部分)
-
- [1.1 文本嵌入层](#1.1 文本嵌入层)
- [1.2 位置编码器](#1.2 位置编码器)
- [2 编码器部分](#2 编码器部分)
-
- [2.1 掩码张量](#2.1 掩码张量)
- [2.2 注意力机制](#2.2 注意力机制)
- [2.3 多头注意力机制](#2.3 多头注意力机制)
- [2.4 前馈全连接层](#2.4 前馈全连接层)
- [2.5 规范化层](#2.5 规范化层)
- [2.6 子层连接结构(残差连接)](#2.6 子层连接结构(残差连接))
- [2.7 编码器层](#2.7 编码器层)
- [2.8 编码器](#2.8 编码器)
- [3 解码器部分](#3 解码器部分)
-
- [3.1 解码器层](#3.1 解码器层)
- [3.2 解码器](#3.2 解码器)
- [4 输出部分](#4 输出部分)
- [5 模型搭建](#5 模型搭建)
0 引言
Transformer
基于seq2seq
架构实现NLP领域的典型任务,如机器翻译、文本生成等。
在论文中,作者以机器翻译任务为主,详细介绍了Transformer
的结构,其总体框架如图1(来源于论文)所示,一共包含四个部分:
-
输入部分:输入词向量预处理
- 源文本嵌入层及位置编码器
- 目标文本嵌入层及位置编码器
-
编码器部分
- 多头自注意力机制子层和规范层以及一个残差连接
- 前馈前连接子层和规范层以及一个残差连接
-
解码器部分
- 多头自注意力机制子层和规范层以及一个残差连接
- 多头注意力机制子层和规范层以及一个残差连接
- 前馈前连接子层和规范层以及一个残差连接
-
输出部分:输出向量预处理
-
线性层
-
softmax层
-
图1 Transformer模型框架
1 输入部分
1.1 文本嵌入层
文本嵌入层
将文本中词汇数字表示转变为高维向量表示
,更便于捕捉词汇间的关系。其代码实现如下:
python
class Embedding(nn.Module):
def __init__(self, d_model, vocab):
'''
文本嵌入层
:param d_model: 词嵌入的维度
:param vocab: 词表的大小
'''
super().__init__()
# 嵌入层
self.lut = nn.Embedding(vocab, d_model)
self.d_model = d_model
def forward(self, x):
return self.lut(x) * math.sqrt(self.d_model)
1.2 位置编码器
位置编码器
在词向量中添加每个词在句子中的位置信息,使模型具有学习词序信息的能力。论文中采用正余弦交替
的位置编码,公式如下:
P E t ( i ) = { s i n ( w k t ) , i = 2 k c o s ( w k t ) , i = 2 k + 1 PE_t^{(i)}= \begin{cases} sin(w_kt),i = 2k \\ cos(w_kt),i = 2k+1\ \end{cases} PEt(i)={sin(wkt),i=2kcos(wkt),i=2k+1
w k = 1 1000 0 2 k / d _ m o d e l , k = 0 , 1 , 2 , . . . , d _ m o d e l 2 − 1 w_k={1 \over 10000^{2k/d\_model}},k=0,1,2,...,{d\_model \over 2}-1 wk=100002k/d_model1,k=0,1,2,...,2d_model−1
其中 i i i表示词向量 t t t所处的位置, d _ m o d e l d\_model d_model表示词向量 t t t的维度。
其代码实现如下:
python
class PositionEncoding(nn.Module):
def __init__(self, d_model, dropout=0.2, max_len=5000):
'''
位置编码层:正余弦交替编码
:param d_model: 词向量维度
:param dropout: 置零比例,正则化,防止模型过拟合
:param max_len: 每个句子的最大长度
'''
super().__init__()
# dropout正则化层
self.dropout = nn.Dropout(p=dropout)
# 初始化一个位置编码矩阵 [max_len, d_model]
pe = torch.zeros(max_len, d_model)
# 初始化一个位置矩阵(句子中每个词向量的位置) [max_len, 1]
position = torch.arange(0, max_len).unsqueeze(1)
# 构造位置编码中正余弦频率(共用频率) [d_model // 2,]
w = 1 / torch.pow(10000, torch.arange(0, d_model, 2) / d_model)
# 构造最终的位置编码:偶数位置使用正弦编码,奇数位置使用余弦编码 [max_len, d_model]
pe[:, 0::2] = torch.sin(position * w)
pe[:, 1::2] = torch.cos(position * w)
# [max_len, d_model] -> [1, max_len, d_model]
pe = pe.unsqueeze(0)
# 位置编码注册为模型buffer,其在优化过程中不会更新,在模型保存后重加载时和模型一起加载
self.register_buffer('pe', pe)
def forward(self, x):
'''
:param x: [bs(句子数量), length(句子长度), d_model] 序列的嵌入(embedding)表示
:return x_pos:[bs, length, d_model] 加上词序信息的词向量
'''
length = x.size(1) # 句子的长度
x = x + Variable(self.pe[:, :length, :], requires_grad=False)
return self.dropout(x)
2 编码器部分
编码器部分结构如图2所示,由N个编码器层组成,每个编码器层由两个子层连接而成:
- 第一个子层包括一个多头自注意力机制和规范层以及一个残差连接
- 第二个子层包括一个前馈全连接子层和规划层以及一个残差连接
图2 编码器结构 >
2.1 掩码张量
在Transformer
中,生成的attention信息中,可能会包含未来的信息,为了避免未来的信息被提前使用,会使用掩码张量
对部分信息进行遮掩。其构建代码如下:
python
def subsequent_mask(size):
attn_shape = (1, size, size)
# 首先构建一个全1矩阵, 再将其形成上三角矩阵(对角线元素也舍弃)
subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
# 构造一个下三角矩阵(包含对角线元素)
return torch.from_numpy(1 - subsequent_mask)
2.2 注意力机制
我们观察事物时,大脑能够快速的把
注意力
放在事物最具辨识度的部分从而作出判断。基于这种形式,在模型中引入注意力机制
,帮助模型聚焦于对当前任务更为关键的信息,降低对其他信息的关注,提高模型对特定任务处理的效率和准确率。
论文中使用的注意力计算方法如图3所示,计算公式如下:
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V Attention(Q,K,V)=softmax({{Q{K^T}} \over {\sqrt {{d_k}} }})V Attention(Q,K,V)=softmax(dk QKT)V
其中 Q ( q u e r y ) , K ( k e y ) , V ( v a l u e ) Q(query),K(key),V(value) Q(query),K(key),V(value)由输入词向量 X X X通过线性变化
得到,当 Q = K = V Q=K=V Q=K=V时称为自注意力机制
, d k d_k dk为词嵌入维度。
图3 注意力机制结构
注意力计算的实现代码如下:
python
def attention(query, key, value, mask=None, dorpout=None):
'''
:param query: [bs, head, length, d_k] Q
:param key: [bs, head, length, d_k] K
:param value: [bs, head, length, d_k] V
:return:
'''
d_k = query.size(-1) # d_model/d_k 词嵌入维度
# 计算Qk^T,得到注意力张量scores [bs, head, length, length]
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
# 判断是否使用掩码张量
if mask is not None:
scores = scores.mask_fill(mask == 0, -1e9)
# 对scores的最后一维度进行softmax操作 [bs, head, length, length]
p_attn = F.softmax(scores, dim=-1)
# 判断是否使用dropout
if dorpout is not None:
p_attn = dorpout(p_attn)
# 得到最终的注意力表达式 [bs, head, length, d_k]
attn = torch.matmul(p_attn, value)
return attn, p_attn
2.3 多头注意力机制
多头注意力机制结构如图4所示,其将输入词向量在词嵌入维度上d\_model
进行切分,分为 h h h个头,每个头的词嵌入维度为 d _ m o d e l / h d\_model / h d_model/h,分别进行自注意力机制运算后再对结果在通道上进行拼接。
多头注意力机制让每个注意力去优化每个词汇的不同特征部分,从而均衡同一种注意力机制可能产生的偏差,让词义信息具有更丰富的表达。
图4 多头注意力机制结构
多头注意力机制实现代码如下:
python
def clones(module, N):
'''用于生成相同网络层的克隆函数,其参数不共享'''
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
'''
多头注意力机制
:param h: 词向量在词嵌入维度上划分的头数
:param d_model: 词嵌入维度
:param dropout:
'''
super().__init__()
# 每个头的分割词向量维度d_k
assert d_model % h == 0
self.d_k = d_model // h
# 分割头数
self.h = h
# 得到线性变化层,分别用于Q,K,V矩阵以及对拼接矩阵的线性操作
self.linears = clones(nn.Linear(d_model, d_model), 4)
# 注意力张量
self.attn = None
# dropout正则化
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
# 样本数
batch_size = query.size(0)
# 多头处理环节:首先对QKV进行线性操作, 然后按照头数进行划分
# q,k,v:[bs, length, head, d_k] -> [bs, head, length, d_k]
query, key, value = \
[model(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
for model, x in zip(self.linears, (query, key, value))]
# 对每个头进行注意力计算
# x: 注意力张量[bs, head, length, d_k]
# attn: 注意力张量权重[bs, head, length, length]
x, self.attn = attention(query, key, value, mask=mask, dorpout=self.dropout)
# 将多头的注意力张量进行合并 [bs, head, length, d_k] -> [bs, length, head, d_k] -> [bs, length, head*d_k(d_model)]
x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
# 多头注意力机制返回结果 [bs, length, d_model]
return self.linears[-1](x)
2.4 前馈全连接层
前馈全连接层由两个线性操作和ReLU激活函数组成,其公式如下:
F F N ( x ) = W 2 m a x ( 0 , W 1 x + b 1 ) + b 2 FFN(x) = W_2max(0, W_1x+b_1)+b_2 FFN(x)=W2max(0,W1x+b1)+b2
其代码实现如下:
python
class PositionwiseFeedForward(nn.Module):
def __init__(self, d_model, d_hidden, dropout=0.1):
'''
前馈全连接层
:param d_model: 词嵌入维度
:param d_hidden: 隐藏层维度
:param dropout:
'''
super().__init__()
self.w1 = nn.Linear(d_model, d_hidden)
self.w2 = nn.Linear(d_hidden, d_model)
self.act = F.relu
self.dropout = nn.Dropout(p=dropout)
def forward(self, x):
'''
:param x: 词向量 (bs, length, d_model)
:return:
'''
y = self.dropout(self.act(self.w1(x)))
return self.w2(y)
2.5 规范化层
规范化层
实现对神经网络数值的标准化,使其特征数值满足均值为0,方差为1的分布,有利于模型的收敛。其代码实现如下:
python
class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
'''
规范化层
:param features: 输入特征的维度
:param eps:防止分母为0
'''
self.a2 = nn.Parameter(torch.ones(features)) # 随着模型训练
self.b2 = nn.Parameter(torch.zeros(features)) # 随着模型训练
self.eps = eps
def forward(self, x):
'''
:param x: 词向量 [bs, length, d_model]
:return:
'''
mean = x.mean(-1, keepdim=True) # 均值
std = x.std(-1, keepdim=True) # 方差
# 标准化 [bs, length, d_model]
return self.a2 * (x - mean) / (std + self.eps) + self.b2
2.6 子层连接结构(残差连接)
对子层结构进行封装,实现代码如下:
python
class SublayerConnection(nn.Module):
def __init__(self, size, dropout=0.1):
'''
:param size: 词嵌入维度
:param dropout:
'''
self.norm = LayerNorm(size) # 标准化
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, sublayer):
'''
:param x: 词向量 [bs, length, d_model]
:sublayer: MultiHeadedAttention / PositionwiseFeedForward
:return:
'''
return x + self.dropout(sublayer(self.norm(x)))
2.7 编码器层
结合以上函数,实现编码器层:
python
class EncoderLayer(nn.Module):
def __init__(self, size, self_attn, feed_forward, dropout):
'''
:param size: 词嵌入维度
:param self_attn: 多头注意力机制
:param feed_forward: 前馈全连接层
:param dropout:
'''
self.self_attn = self_attn
self.feed_forward = feed_forward
# 构造子层连接
self.sublayer = clones(SublayerConnection(size, dropout), 2)
self.size = size
def forward(self, x, mask):
y = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第一个子层(多头注意力+标准化+残差连接)
y = self.sublayer[1](y, self.feed_forward) # 第二个子层(前馈全连接+标准化+残差连接)
return y
2.8 编码器
基于编码器层,构造编码器:
python
class Encoder(nn.Module):
def __init__(self, layer, N):
'''
:param layer: 编码器层
:param N: 编码器层数
'''
super().__init__()
self.layers = clones(layer, N) # 构造N层编码器
self.norm = LayerNorm(layer.size) # 构造标准化层
def forward(self, x, mask):
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)
3 解码器部分
解码器部分结构如图5所示,由N个编码器层组成,每个编码器层由三个子层连接而成:
- 多头自注意力机制子层和规范层以及一个残差连接
- 多头注意力机制子层和规范层以及一个残差连接
- 前馈前连接子层和规范层以及一个残差连接
图5 解码器结构
3.1 解码器层
结合编码器部分中实现的模块,解码器层代码如下:
python
class DecoderLayer(nn.Module):
def __init__(self, size, self_attn, src_attn, feed_forward, dropout=0.1):
'''
:param size: 词嵌入维度
:param self_attn: 多头自注意力机制(Q=K=V)
:param src_attn: 多头注意力机制(Q!=K=V)
:param feed_forward: 前馈全连接层
:param dropout:
'''
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
# 构造子连接层
self.sublayers = clones(SublayerConnection(size, dropout), 3)
def forward(self, x, memory, source_mask, target_mask):
'''
:param x: 上一层的输出张量 [bs, length, d_model]
:param memory: 编码器层的输出 [bs, length, d_model]
:param source_mask: 多头注意力机制的掩码张量
:param target_mask: 多头自自注意机制的掩码张量
'''
m = memory
# 第一个子层:多头自注意力机制 + 标准化 + 残差连接
# 目标数据掩码张量:避免对未来信息的使用
x = self.sublayers[0](x, lambda x:self.self_attn(x, x, x, target_mask))
# 第二个子层:多头注意力机制 + 标准化 + 残差连接
# 源数据掩码张量:遮挡对结果无效的信息
x = self.sublayers[1](x, lambda x:self.src_attn(x, m, m, source_mask))
# 第三个子层:前馈全连接层 + 标准化 + 残差连接
x = self.sublayers[2](x, self.feed_forward)
return x
3.2 解码器
基于解码器层,构造解码器:
python
class Decoder(nn.Module):
def __init__(self, layer, N):
'''
:param layer: 解码器层
:param N: 数量
'''
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def forward(self, x, memory, source_mask, target_mask):
'''
:param x: 数据的嵌入表示 [bs, length, d_model]
:param memory: 编码器输出 [bs, length, d_model]
:param source_mask:
:param target_mask:
'''
for layer in self.layers:
x = layer(x, memory, source_mask, target_mask)
return self.norm(x)
4 输出部分
输出层包括线性层和softmax层,其代码如下:
python
class Generator(nn.Module):
def __init__(self, d_model, vocab_size):
'''
:param d_model: 词嵌入维度
:param vocab_size: 词表大小
'''
super().__init__()
# 线性层,将词向量转换为词表预测结果
self.project = nn.Linear(d_model, vocab_size)
def forward(self, x):
'''
:param x: 解码器的最终输出 [bs, length, d_model]
:return: 最终预测结果 [bs, length, vocab_size]
'''
return F.log_softmax(self.project(x), dim=-1)
5 模型搭建
python
class EncoderDecoder(nn.Module):
def __init__(self, encoder, decoder, source_embed, target_embed, generator):
'''
编码器 - 解码器结构
:param encoder: 编码器
:param decoder: 解码器
:param source_embed: 源数据嵌入函数
:param target_embed: 目标数据嵌入函数
:param generator: 类别生成器
'''
self.encoder = encoder
self.decoder = decoder
self.src_embed = source_embed
self.tgt_embed = target_embed
self.generator = generator
def forward(self, source, target, source_mask, target_mask):
'''
:param source: 源数据 [bs, length]
:param target: 目标数据 [bs, length]
:param source_mask: 源数据掩码张量
:param target_mask: 目标数据掩码张量
'''
return self.decode(self.encode(source, source_mask), source_mask,
target, target_mask)
def encode(self, source, source_mask):
'''编码器'''
return self.encoder(self.src_embed(source), source_mask)
def decode(self, memory, source_mask, target, target_mask):
return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)
def make_model(source_vocab, target_vocab, N=6,
d_model=512, d_hidden=2048, head=8, dropout=0.1):
'''
构建模型
:param source_vocab: 源数据特征(词汇)总数
:param target_vocab: 目标数据特征(词汇)总数
:param N: 解码器和编码器个数
:param d_model: 词嵌入维度
:param d_hidden: 前馈全连接层隐藏层层数
:param head: 多头注意力机制头数
:param dropout:
'''
# 深拷贝
c = copy.deepcopy
# 实例化多头注意力机制类
attn = MultiHeadedAttention(head, d_model)
# 实例化前馈全连接层类
ff = PositionwiseFeedForward(d_model, dropout)
# 实例化位置编码类
position = PositionEncoding(d_model, dropout=dropout)
# 构造模型
model = EncoderDecoder(
encoder=Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
decoder=Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
source_embed=nn.Sequential(Embedding(d_model, source_vocab), c(position)),
target_embed=nn.Sequential(Embedding(d_model, target_vocab), c(position)),
generator=Generator(d_model, target_vocab)
)
# 初始化参数
for p in model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform(p)
return model