The Building Blocks of Transformers

Transformer完全改变了2017年后NLP领域的模型方向, 从某种意义上说,Bert,GPT等模型都是Transformer模型的变体, 虽然模型结构有各种改变, 但是其中的一些基本计算单元则变化较小.

Transformer几乎就是为了改善计算性能而专门设计的模型.

完全没有RNN之类的循环计算需求, 这就极大降低了计算过程中的顺序依赖, 可以极大提高并行性.
大量使用矩阵乘, 不使用卷积这种计算强度不够大的算子

正是由于Transformer使用到的基本计算单元非常简单, 几乎就只有 gemm, +-*/, layernorm, softmax, 也没有奇怪的计算流程, 所以原文的作者将其称为一个"简单"的模型是很有道理的.

本文就是简单记录Transfomer中使用到的基本计算单元.

Basics

想要理解Transfomer计算流程的话, 可以参考 Transformer’s Encoder-Decoder: Let’s Understand The Model Architecture

如果有不清楚的地方, 可以参考Github

我这里仅简单的对Transformer做一个计算流程上的小结, 为了简化, 这个例子里省略了很多细节.

推理

仅需要准备原始输入src, src是一个shape为1xM的token序列:

context = encode(src)
tgt = [<BOS>]
while tgt[-1] != <EOS>
    out = decode(tgt,context)
    new_tok = map_to_token(out[-1])
    tgt.append(new_tok)

从功能上说,一次decode会推理出输入参数tgt后面应该跟的下一个词, transformer的设计上把新增加的token放到out的最后.

例如,假如tgt是 ("I","Love"), 那么输出out对应的就可能是("Love","You"). 另外, out实际的shape是(2 , VOCAB_SIZE),每一行都是一个概率分布, 以输出("Love","You")为例,那么第一行中"Love"对应的概率就是最大的,第二行中"You"对应的概率就是最大的.

训练

需要准备原始输入src (1xM) 和对应的期待输出 expect_tgt (1xN), 具体的训练策略也有很多,其中一种容易理解的是teacher forcing.
具体而言,也就是说, 当我们通过src准备好context之后,期望下面这些映射关系在decode阶段都能成立

Src	Lable
[`<BOS>`]	[expect_tgt[1]]
[`<BOS>`,expect_tgt[1]	[expect_tgt[1],expect_tgt[2]]
[`<BOS>`,expect_tgt[1],expect_tgt[2]]	[expect_tgt[1],expect_tgt[2],expect_tgt[3]]
…	…

对应到代码,就是

N = expect_tgt.size(-1)
for i in range(N-1)
    context = encode(src)
    tgt = expect_tgt[0:(i+1)]
    tgt_y = expect_tgt[1:(i+2)]
    out = decode(tgt,context)
    loss = loss_fn(out,tgt_y)
    loss.backward()

Computation Details

Get X

X是Encoder阶段的输入, 它是一个(M,MAX_SEQ_LEN,DIM_EMBED)形状的矩阵，每一行都对应了一个词向量.

为了支持batch, 我们先假定原始输入是一系列长短不一的序列, 存放在 seqs = set()中

Tokenizie: 将原始序列中的文本变成token, 从而得到一个新的集合 tok_seqs = tokenize(seqs)

Padding: 为了能将所有tok_seq都放在同一个矩阵中,我们需要进行padding

max_seq_len = max([len(seq) for seq in tok_seqs])
for seq in tok_seqs:
seq.resize(size=(1,max_seq_len),fill=PAD_TOK)

Batching: 将tok_seq放到同一个矩阵中src = stack(tok_seqs,dim=0),得到一个形状为(M,MAX_SEQ_LEN)的矩阵.
Embedding: 将原始的src映射到词向量空间,X = embedding(src), 得到一个形状为(M,MAX_SEQ_LEN,DIM_EMBED)的矩阵, Embedding会引入一个可训练的参数,用于提供Embedding矩阵
Positional Encoding: 为X附加一个位置信息X = X + PE.slice(X.shape)

在这个阶段, 我们还会有一个额外产物src_mask,它的形状为(M,MAX_SEQ_LEN), 用于标记src矩阵中不为PAD_TOK的部分,即src_mask = src != PAD_TOK

注意,在这个阶段,我们并没有约束用户输入的原始seq数量, 因此, 对于最后得到的X矩阵, 其shape (M,MAX_SEQ_LEN,DIM_EMBED)中, M和MAX_SEQ_LEN都是可以随用户输入可变的, DIM_EMBED 则是一个模型参数, 需要在设计模型的阶段中确定.

事实上, Transofmer的Encoder-Decoder架构中, 模型中的所有参数都不依赖于seq_len, encoder的输入src_seq_len以及decoder的tgt_seq_len都是动态可变的,不受模型约束.

关于Embedding

逻辑上, embedding实际是有两步的, 对于某个 tok 序列, 如 (8667, 1362, 106), 我们需要先将它映射成(3,VOCAB_SZ)的一个one-hot矩阵A(每一行都是one-hot向量), 然后通过矩阵乘法X=A@Embed变成(3,DIM_EMBED)的矩阵.

显然

Embed的形状为 (VOCAB_SZ, DIM_EMBED).
由于A实际上是one-hot的, 所以矩阵乘法实际上相当于是在抽取Embed的行形成新的矩阵,所以实际上在执行embedding操作时,并不会有转换成one-hot,再矩阵乘的操作, 实际的实现就是简单的对Embed矩阵按行做一个indexed slicing

关于 PE

这里PE被称为Positional Encoding, 它主要的目的是通过加法给词向量矩阵X的每一行打上一个标记, PE矩阵的大小一般为(BIG_ENOUGH,DIM_EMBED),其构造算法如下


\begin{array}{l}PE[pos,2i] = \sin (\frac{{pos}}{{{{10000}^{\frac{{2i}}{{DIM\_EMBED}}}}}})\\PE[pos,2i + 1] = \cos (\frac{{pos}}{{{{10000}^{\frac{{2i}}{{DIM\_EMBED}}}}}})\;\end{array}\

一般来说, 我们会构造一个充分大的常量PE矩阵,也就是让BIG_ENOUGH取一个充分大的值,然后当需要与X叠加时,直接将前seq_len行抽取出来即可.

至于它为什么有效, 为什么这么设计, 应该是另外一个问题了

Get Y

decoder的输入Y也是需要遵循相同的过程,从tgt来构造的, 不同的是, Y会需要一个subsequent_mask, 不过它就是一个下三角矩阵,功能和构造方法都很简单, 可以在后面再单独说.

Attention (Single Head)

从公式上说, 因为我们不必去了解Attention背后的算法原理, 所以只看计算而言, Attention是非常简单的


\begin{array}{l}
Attention(Y,X,Mask)\\
 = softmax (maskfill (\frac{{Y{W_Q}W_K^T{X^T}}}{{\sqrt {{d_k}} }},Mask),axis=-1)X{W_v}\\
 = Attention(Q,K,V,Mask)\\
 = softmax (maskfill (\frac{{Q{K^T}}}{{\sqrt {{d_k}} }},Mask),axis=-1)V\\
where,Q = Y{W_Q},K = X{W_K},V = X{W_V}
\end{array}

我们只需要补充一些要点

引入了三个可训练参数W_Q,W_K,W_V
- 这三个矩阵的M维方向都因为要参与矩阵乘而被固定了, 和对应矩阵的词向量维度一致
- W_K的N方向维度一般被称为d_k, W_V的N方向维度一般被称为d_v, 都是可以随模型设计而改变的参数.
maskfill (\frac{{Y{W_Q}W_K^T{X^T}}}{{\sqrt {{d_k}} }},Mask)这个矩阵是有直观意义的, 假设Y对应了M个token,X对应了N个Token,那么这个矩阵的形状就是(M,N),位于(i,j)的值将对应着Y中第i个token对X中第j个token的关注度
- Mask可以用于屏蔽掉特定的关注信息, 例如, 假如X的末尾有2个PAD_TOK,那么(M,N)矩阵的最后两列就应被标记为一个无效的数,表示Y中的所有token都不应该关注这X的最后两个token
- Mask的源头有两个, 一个是padding, 另一个是decoder执行self-attention时会用到的subsquent_mask,这个将会在后面介绍
Q,K,V的思想是, 输入X应该提供key和value的信息, Q实际上仅仅是一个Y发出的query请求. 计算的过程就是先通过这个请求去和key进行匹配, 然后映射到相关的value.

Self Attention

当Attention(Y,X)的两个输入相同时, 就是self-attention.

Encode和Decode阶段都会有self-attention的参与, 逻辑上说, 在decode阶段时, 我们是一个一个生成token的, 也就是说, tgt中先出现的token应该是看不到后出现的token的, 所以不应该对后出现的token产生关注, 为了体现这个约束, 就需要提供一个subsequent_mask 矩阵, 它是一个下三角矩阵,上半部分全都是0,下半部则全都是1

Cross Attention

当Attention(Y,X)的两个输入不同时, 就是cross-attention.
在Decode阶段, 逻辑上Y将和tgt对应,而X则和src转换得到的context对应, 通过cross-attention, 就能将decode阶段和encode阶段连接起来.

MultiHead Attention

从实现上讲, 他就是并行的使用了多个Attention, 最后再把每个Attention的结果汇总起来. 原文作者希望这样可以改善模型的表达能力, 有更大的搜索空间

MultiHead(Y,X)=Concat([Head_0(Y,X),Head_1(Y,X),...],axis=1)W_O

最后的W_O主要用于变形,从而将Concat得到的矩阵变成需要的形状. 在实践中, 一般希望Multihead(Y,X)的输出形状和Y相同, 具体的做法一般是将每个Head的 d_v 设为 Y.shape[-1] / head_num,这样一来, 就可以直接concat到最终形状, 而避免使用W_O

FFN

FFN的计算很简单

FFN(x)=ReLU(xW_1+b_1)W_2+b_2

LayerNorm

在Transformer的场景中,LayerNorm的意思就是对矩阵中每一个词向量单独做Norm,并且所有词向量共享同组参数.

Residual Connection

在 Transformer中, 为了保证Positional Encoding的信息不丢失, 如果x中附加了 Positional Encoding 数据, 那么涉及x的计算,都应该按照y = x + f(x)来执行, 这就是 Residual Connection.

Output: Linear

Decoder的最终输出需要重新映射到VOCAB_SIZE的字典空间中,得到响应的概率值, 这就是Linear的作用.

在训练阶段, 有可能还需要把结果再进行一次softmax, 以便于供后续的loss函数计算损失.

TinyEdi

Sharing & Fun