tansformer 从入门到上手只需要看这一篇

tansformer 架构解析

架构总览

Alt text

词嵌入层 Embeddings

输入 token 的索引处在维度为词表大小 d_vocab 的空间里,每个 token 要经过词嵌入层变成维度为 d_model 的特征向量。

1
2
3
4
5
6
7
8
class Embeddings(nn.Module):
def __init__(self, d_model, d_vocab):
super(Embeddings, self).__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)

在原论文,编码器输入、解码器输入以及解码器输出是使用相同参数的 Embeddings,但是在有些地方,发现解码器输出是使用单独的生成器 Generator 结构,原因是因为解码器输出在功能上是获取一个从特征向量到词表空间的概率映射。

1
2
3
4
5
6
7
8
9
class Generator(nn.Module):
"Define standard linear + softmax generation step."

def __init__(self, d_model, d_vocab):
super(Generator, self).__init__()
self.proj = nn.Linear(d_model, vocab)

def forward(self, x):
return torch.log_softmax(self.proj(x), dim=-1)

位置编码 Positional Encoding

为了让模型利用序列的次序,必须注入一些关于序列中的相对位置或绝对位置的信息。为此,在 encoder stack 和 decoder stack 底部的 input embedding 中添加了 positional encoding 。positional encoding 与 embedding 具有相同的维度 ,因此可以将二者相加。

Pj=(PE(j,1),PE(j,2)...PE(j,dmodel))PE(j,2i)=sin(j100002i/dmodel),PE(j,2i+1)=cos(j100002i/dmodel)P_j=(PE_{(j,1)}, PE_{(j,2)}...PE_{(j,d_{model})}) \\ PE_{(j,2i)}=\sin(\frac{j}{10000^{2i/d_{model}}}), PE_{(j,2i+1)}=\cos(\frac{j}{10000^{2i/d_{model}}})

其中,j 表示 position,i 表示维度。即 positional encoding 的每个维度对应于一个正弦曲线,正弦曲线的波长从 2π2\pi (当维度 i=0i=0 时)到 10000×2π10000 \times 2\pi (当维度 2i=dmodel2i=d_{model} 时)。

对于任意固定的偏移量 k,Pj+kP_{j+k} 可以表示为 PjP_j 的线性函数,且正弦版本的 positional embedding 可以应用到训练期间 unseen 的位置,从而可以让模型推断出比训练期间遇到的序列长度更长的序列。

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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class PositionalEncoding(nn.Module):
"""实现Positional Encoding功能"""

def __init__(self, d_model, dropout=0.1, max_len=5000):
"""
位置编码器的初始化函数
:param d_model: 词向量的维度,与输入序列的特征维度相同,512
:param dropout: 置零比率
:param max_len: 句子最大长度,5000
"""
# 如果d_model不是偶数那么后面pe[:, 1::2]会少一个维度,赋值报错。
assert d_model % 2 == 0
super(PositionalEncoding, self).__init__()
# 初始化一个nn.Dropout层,设置给定的dropout比例
self.dropout = nn.Dropout(p=dropout)

# 初始化一个位置编码矩阵
# (5000,512)矩阵,保持每个位置的位置编码,一共5000个位置,每个位置用一个512维度向量来表示其位置编码
pe = torch.zeros(max_len, d_model)
# 偶数和奇数在公式上有一个共同部分,使用log函数把次方拿下来,方便计算
# position表示的是字词在句子中的索引,如max_len是128,那么索引就是从0,1,2,...,127
# 论文中d_model是512,2i符号中i从0取到255,那么2i对应取值就是0,2,4...510
# (5000) -> (5000,1)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 计算用于控制正余弦的系数,确保不同频率成分在d_model维空间内均匀分布
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
# 根据位置和div_term计算正弦和余弦值,分别赋值给pe的偶数列和奇数列
pe[:, 0::2] = torch.sin(
position * div_term
) # 从0开始到最后面,补长为2,其实代表的就是偶数位置
pe[:, 1::2] = torch.cos(
position * div_term
) # 从1开始到最后面,补长为2,其实代表的就是奇数位置
# 上面代码获取之后得到的pe:[max_len * d_model]
# 下面这个代码之后得到的pe形状是:[1 * max_len * d_model]
# 多增加1维,是为了适应batch_size
# (5000, 512) -> (1, 5000, 512)
pe = pe.unsqueeze(0)
# 将计算好的位置编码矩阵注册为模块缓冲区(buffer),这意味着它将成为模块的一部分并随模型保存与加载,但不会被视为模型参数参与反向传播
self.register_buffer("pe", pe)

def forward(self, x):transformer1
"""
x: [seq_len, batch_size, d_model] 经过词向量的输入
"""
x = (
x + self.pe[:, : x.size(1)].clone().detach()
) # 经过词向量的输入与位置编码相加
# clone().detach() 组合起来的主要作用是创建一个与原张量数据相同,但不参与梯度计算且内存独立的新张量
# Dropout层会按照设定的比例随机“丢弃”(置零)一部分位置编码与词向量相加后的元素,
# 以此引入正则化效果,防止模型过拟合
return self.dropout(x)

多头注意力

Alt text

注意力机制 Scaled Dot-Product Attention

原论文叫 Scaled Dot-Product Attention (放缩点积注意力),架构如上图左侧所示。

注意力机制包括 Query ( Q ),Key ( K )和 Value ( V )三个组成部分,这三个部分由WQ,WK,WVW_Q,W_K,W_V三个线性层生成,这部分原论文架构图没有给出,是我自己补充的,公式如下:

Q(len_q,d_k)=XQ(len_q,dmodel)WQ(dmodel,d_k)K(len_v,d_k)=XV(len_v,dmodel)WK(dmodel,d_k)V(len_v,d_v)=XV(len_v,dmodel)WV(dmodel,d_v)Q(len\_q, d\_k) = X_Q(len\_q, d_model) \cdot W_Q(d_model, d\_k) \\ K(len\_v, d\_k) = X_V(len\_v, d_model) \cdot W_K(d_model, d\_k) \\ V(len\_v, d\_v) = X_V(len\_v, d_model) \cdot W_V(d_model, d\_v)

其中 Q 和 K 有相同的特征维度 d_k,而 K 和 V 有相同的长度维度 len_v,它们由XQ,XVX_Q,X_V两个输入生成,当XQ=XVX_Q=X_V时计算的是自注意力,当它们不同时计算的就是交叉注意力。

在很多代码中WQ,WK,WVW_Q,W_K,W_V三个线性层是默认有偏置项的,这里为了便于理解没有写出来。

注意力计算公式为:

Attention(Q,K,V)=softmax(QKTdk)VQRlenq×dk,KRlenv×dk,vRlenv×dv,QKTRlenq×lenv,Attention(Q,K,V)Rlenq×dvAttention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V \\ Q \in \R ^{len_q \times d_k},K \in \R ^{len_v \times d_k},v \in \R ^{len_v \times d_v}, QK^T \in \R ^{len_q \times len_v}, Attention(Q,K,V) \in \R ^{len_q \times d_v}

内积注意力使用了 1/dk1/\sqrt{d_k} 的缩放因子,因为维度越大,则内积中累加和的项越多,内积结果越大。很大的数值会导致 softmax 函数位于饱和区间,梯度几乎为零。

因为输出特征维度为 d_v,所以为了添加残差,通常要求 d_v=d_model 或者也可以引入一个dvdmodeld_v \to d_model的线性层。

多头注意力 Multi-Head Attention

多头注意力将 query, key, value 线性投影 h 次,其中每次线性投影都是不同的并且将 query, key, value 分别投影到 dk,dk,dvd_k,d_k,d_v 维。然后,在每个 query, key, value 的投影后的版本上,并行执行注意力函数,每个注意力函数产生 dvd_v 维的 output 。这些 output 被拼接起来并再次投影,产生 final output ,如上图右侧所示。

原论文说,多头注意力允许模型在每个位置联合地关注来自不同子空间的信息。如果只有单个注意力头,那么平均操作会抑制这一点。

MultiHead(Q,K,V)=Concat(head1,head2...headh)WOheadi=Attention(QWiQ,KWiK,VWiV)WiQRdmodel×dk,WiKRdmodel×dk,WiVRdmodel×dv,WOR(hdv)×dmodelMultiHead(Q,K,V)=Concat(head_1,head_2 ... head_h)W^O \\ head_i=Attention(QW^Q_i,KW^K_i,VW^V_i) \\ W^Q_i \in \R ^{d_{model} \times d_k},W^K_i \in \R ^{d_{model} \times d_k},W^V_i \in \R ^{d_{model} \times d_v},W^O \in \R ^{(hd_v) \times d_{model}}

公式中给出的是更一般的形式,实际上这里的输入 Q 对应上文注意力机制中的XQX_Q,输入 K 和 V 对应上文注意力机制中的XVX_V,另外原论文中 d*model=512, dk=dv=dmodel/h=64d_k=d_v=d*{model}/h=64

在 Transformer 中以三种不同的方式使用多头注意力:

  1. 在 encoder-decoder attention 层中,query 来自于前一个 decoder layer,key 和 value 来自于 encoder 的输出。这允许 decoder 中的每个位置关注 input 序列中的所有位置。这模仿了 sequence-to-sequence 模型中典型的 encoder-decoder attention 注意力机制。
  2. encoder 包含自注意力层。在自注意力层中,所有的 query, key, value 都来自于同一个地方(在这个 case 中,就是 encoder 中前一层的输出)。encoder 中的每个位置都可以关注 encoder 上一层中的所有位置。
  3. 类似地,decoder 中的自注意力层允许 decoder 中的每个位置关注 decoder 中截至到当前为止(包含当前位置)的所有位置。我们需要防止 decoder 中的信息向左流动,从而保持自回归特性。我们通过在 scaled dot-product attention 内部屏蔽掉 softmax input 的某些 value 来实现这一点(将这些 value 设置为 -1e9,一个很小的数),这些 value 对应于无效连接 illegal connection。

代码实现 pytorch

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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def clones(module, N):
"Produce N identical layers."
return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

def attention(query, key, value, mask=None, dropout=None):
'''
query: (len_q, d_k)
key: (len_v, d_k)
value: (len_v, d_v)
mask: (len_q, len_v) fill true
'''
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, -1e9)
p_attn = torch.softmax(scores, dim = -1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
'''
Take in model size and number of heads.
We assume d_v always equals d_k.
'''
super(MultiHeadedAttention, self).__init__()
assert d_model % h == 0
self.d_k = d_model // h
self.h = h
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):
'''
query: (len_q, d_k)
key: (len_v, d_k)
value: (len_v, d_k)
mask: (len_q, len_v) fill true
'''
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 = [
lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for lin, x in zip(self.linears, (query, key, value))
]

# 2) Apply attention on all the projected vectors in batch.
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)
)
del query
del key
del value
return self.linears[-1](x)

编码器掩码 length mask

Alt text
对于输入 x,其原始序列长度是长短不一的,但是为了放入张量中进行训练,会对短序列填补到最大长度,使输入序列的长度统一。这样在计算注意力时,填补的维度应该是无效的不能参与计算。

注意力计算中 socre 的维度为(len_q, len_v),因为 torch.softmax(scores, dim = -1)计算最后一个维度列,对不同列之间进行 softmax,所以实际上是对每行数据 softmax。这样要保证每行不能全被 mask,防止计算出现 NaN。这样对 len_v 列维度进行 mask 就行了,len_q 对应部分最终计算 loss 时会被过滤掉,没有影响。

这里预期的输入 batch_mask (N, len) 是已经处理好的输入序列的掩码(True 代码填充部分),只需要在批次维度N和序列长度维度len之间填充一个维度,那么在 scores.masked_fill(mask, -1e9) 时,最后维度 len_v 的掩码会自动在前面维度 len_q 上广播,而不用考虑 len_q 的长度变化,这样这个掩码既可以在编码器计算自注意力时用到,又可以用在计算交叉注意力时。

1
2
3
4
5
6
def len_mask(batch_mask) -> torch.Tensor:
return batch_mask.clone().unsqueeze(-2)

# 这里给出了第二中把 len_v 维度的掩码广播到新增的 len_q 维度的方法。
def len_mask2(batch_mask, len_q) -> torch.Tensor:
return batch_mask.clone().unsqueeze(-2).expand(*batch_mask.shape, len_q).transpose(-1, -2)

解码器掩码 causal mask

我们给模型的完整输入为 (x1,x2...xn)(x_1, x_2...x_n) ,但是预测 xix_i 时只用到了前 i-1 个输入的信息,也就是说后面的信息我们应当 mask 掉,只能根据历史信息来预测当前位置的输出。这很好理解,假如模型的输入为“早上好”,那么应该通过“早”来预测“上”,通过“早上”来预测“好”,而不是通过“早上好”来预测“上”或者“好”,因为这相当于直接把 ground truth 告诉模型了。

Alt text
在解码器的自注意力层,socre 的维度为 (len_q, len_v),我们只需要如上图生成一个上三角矩阵,将对角线以上的部分 mask 就行了。

1
2
3
4
5
6
7
def subsequent_mask(size):
"Mask out subsequent positions."
attn_shape = (1, size, size)
subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
torch.uint8
)
return subsequent_mask == 1

标准化

这里对末尾维度也就是特征维度进行标准化,标准化后还通过了一个线性层。

这里分别对线性层的权重初始化为 1,偏置初始化为 0。而 nn.liner 默认使用 Kaiming 初始化权重,均匀分布初始化偏置。

1
2
3
4
5
6
7
8
9
10
11
12
13
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

残差连接 residual connections

这里对传入的子层标准化后添加残差。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
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, d_model, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(d_model)
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)))

逐位置前馈网络 Position-wise Feed-Forward Networks

FFN(x)=max(0,xW1+b1)W2+b2FFN(x)=max(0, xW_1+b_1)W_2+b_2

先将 x 从 d_model=512 维投影到 dff=2048 维,再投影回 d_model=512 维。

1
2
3
4
5
6
7
8
9
10
11
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(self.w_1(x).relu()))

编码器 Ecoder

对于单个编码器,先通过残差连接包裹的 self_attn 自注意力层,再通过残差连接包裹的 feed_forward 逐位置前馈网络层。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class EncoderLayer(nn.Module):
"Encoder is made up of self-attn and feed forward (defined below)"

def __init__(self, d_model, self_attn, feed_forward, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(d_model, dropout), 2)
self.size = d_model

def forward(self, x, mask):
"Follow Figure 1 (left) for connections."
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
return self.sublayer[1](x, self.feed_forward)

完整编码器就是把上面的编码重复 N 层,然后通过一个标准化层。

1
2
3
4
5
6
7
8
9
10
11
12
13
class Encoder(nn.Module):
"Core encoder is a stack of N layers"

def __init__(self, layer, N):
super(Encoder, self).__init__()
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)

def forward(self, x, mask):
"Pass the input (and mask) through each layer in turn."
for layer in self.layers:
x = layer(x, mask)
return self.norm(x)

解码器 Decoder

单个解码器里有两个注意力层,第一个注意力层先对输出结果计算自注意力,再将结果输出到第二个注意力层与编码器的输出结果一起计算交叉注意力,随后进入一个逐位置前馈网络层,这三层都使用前面定义的残差连接包裹。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class DecoderLayer(nn.Module):
"Decoder is made of self-attn, src-attn, and feed forward (defined below)"

def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
super(DecoderLayer, self).__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3)

def forward(self, x, memory, src_mask, tgt_mask):
"Follow Figure 1 (right) for connections."
m = memory
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
return self.sublayer[2](x, self.feed_forward)

完整解码器就是把上面的解码层重复 N 层,然后通过一个标准化层。

1
2
3
4
5
6
7
8
9
10
11
12
class Decoder(nn.Module):
"Generic N layer decoder with masking."

def __init__(self, layer, N):
super(Decoder, self).__init__()
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)

def forward(self, x, memory, src_mask, tgt_mask):
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)

Transformer

Alt text
把以上组件组合到一起构成 Transformer 架构,其中编码器和解码器的掩码生成代码没有加进去,因为掩码是要随数据一起初始化的,并不和网络结构一起初始化。

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
33
34
class Transformer(nn.Module):
"""
A standard Transformer architecture.
"""

def __init__(
self, src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1
):
super(Transformer, self).__init__()
c = copy.deepcopy
attn = MultiHeadedAttention(h, d_model)
ff = PositionwiseFeedForward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
self.encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N)
self.decoder = Decoder(
DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N
)
self.src_embed = nn.Sequential(Embeddings(d_model, src_vocab), c(position))
self.tgt_embed = nn.Sequential(Embeddings(d_model, tgt_vocab), c(position))
self.generator = Generator(d_model, tgt_vocab)
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)

def forward(self, src, tgt, src_mask, tgt_mask):
"Take in and process masked src and target sequences."
out = self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
return self.generator(out)

def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)

def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

测试

测试一下是否能跑通,测试没有问题。

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
def inference_test():
test_model = Transformer(11, 11, 2)
# test_model.eval() 评估模式,比如 Dropout 层在评估和训练时表现不同。
test_model.eval()
src = torch.LongTensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])
src_mask = torch.ones(1, 1, 10)
src_mask = src_mask == 0

memory = test_model.encode(src, src_mask)
ys = torch.zeros(1, 1).type_as(src)

for i in range(9):
out = test_model.decode(memory, src_mask, ys, subsequent_mask(ys.size(1)))
prob = test_model.generator(out[:, -1])
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
ys = torch.cat(
[ys, torch.empty(1, 1).type_as(src.data).fill_(next_word)], dim=1
)

print("Example Untrained Model Prediction:", ys)


def run_tests():
for _ in range(10):
inference_test()

训练第一个模型

这里我在网上找到了阿里云的天池数据集打榜挑战赛,并选择了中文医疗信息处理评测基准 CBLUE 中的中文医学命名实体识别 V2(CMeEE-V2)任务。

下载数据

数据好像下载要申请,挺麻烦,我直接在 huggingface 上找到了相同数据集。

1
2
3
4
5
6
7
8
# 先配置 git lfs ,不然大文件不会克隆下来
conda install git-lfs
git lfs install
# 对于大于5G的huggingface数据,还需要
pip install huggingface_hub
huggingface-cli lfs-enable-largefiles .
# 克隆数据
git clone https://huggingface.co/datasets/Rosenberg/CMeEE-V2

训练

将之前的模型部分的代码整理到 transformer.py 文件中,然后导入里面的类和函数。
测试通过,可以正常训练。

我设置 N=2, d_model=128, d_ff=512, batch_size=5 ,已经挺小的了,训练一个但是还是需要4个G的显存,以及每训练一个epoch需要2.9小时的时间。时间和本地电脑资源实在不太够用。

训练结果和代码已上传git:
https://github.com/hs3434/CMeEE-V2

从loss来看,是随着训练下降的。
Alt text
Alt text