The Building Blocks of Transformers

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

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

  1. 完全没有RNN之类的循环计算需求, 这就极大降低了计算过程中的顺序依赖, 可以极大提高并行性.
  2. 大量使用矩阵乘, 不使用卷积这种计算强度不够大的算子
正是由于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阶段都能成立

SrcLable
[<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()

  1. Tokenizie: 将原始序列中的文本变成token, 从而得到一个新的集合 tok_seqs = tokenize(seqs)
  2. Batching: 将tok_seq放到同一个矩阵中src = stack(tok_seqs,dim=0),得到一个形状为(M,MAX_SEQ_LEN)的矩阵.
  3. Embedding: 将原始的src映射到词向量空间,X = embedding(src), 得到一个形状为(M,MAX_SEQ_LEN,DIM_EMBED)的矩阵, Embedding会引入一个可训练的参数,用于提供Embedding矩阵
  4. Positional Encoding: 为X附加一个位置信息X = X + PE.slice(X.shape)

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) 

在这个阶段, 我们还会有一个额外产物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}

我们只需要补充一些要点

  1. 引入了三个可训练参数W_Q,W_K,W_V
    • 这三个矩阵的M维方向都因为要参与矩阵乘而被固定了, 和对应矩阵的词向量维度一致
    • W_K的N方向维度一般被称为d_kW_V的N方向维度一般被称为d_v, 都是可以随模型设计而改变的参数.
  2. 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,这个将会在后面介绍
  3. 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函数计算损失.

Read more

Coroutine

从一般概念上说, 协程是特殊的函数调用: 被调用的函数可以在可控的位置被中断,然后在下一次调用时,继续从上次中断的位置继续执行。 本文主要通过Python的协程来介绍协程, 这是我唯一熟悉的一种协程实现. Classic Coroutine 下面的python代码很好的说明了协程的核心功能 def co_routine(): recv0 = yield 996 # hangs here after first coro.send assert recv0 == "Second" yield 711 # hangs here after second coro.send return def main(): coro = co_routine() # Create a new coroutine object value = coro.send(None)

By Edimetia3D

GDB with Python

这篇文章的主要应用场景是调试Python的C/C++ Extension 1. 同时使用pdb / gdb 进行调试. 通俗点说, 既可以break在 .py 文件中,也可以break在 .cc 文件中 2. 在gdb中不但可以获得常规的调试信息, 还可以获得python VM 的调试信息, 例如获得python的调用栈, 访问Python局部变量等. 这将会在调试exception时(如Segmentfalut)非常有用, 这种场景下, 定位 Python VM 正运行到哪一行代码往往可以提供一些直观的重要信息. 第一步: 编译源码以获得一些辅助数据. 我们并不真的需要使用从源码编译的Python, 但是一些调试相关的辅助文件需要从源码中获得, 包括 python-gdb.py及debug symbol等. 在 https://www.python.org/ftp/python/ 或 https://github.com/python/cpython

By Edimetia3D

Bazel Notes

这是一篇2019年左右的记录, 内容可能过时, 也不太全面 杂谈 Bazel是Google为Monorepo服务而开发的构建工具. 首先是巨大,当问题的规模变大,事情总是会变得更复杂. 而Google面对的"巨大Monorepo",应该是世间罕有的. 然后是Monorepo,这极大的影响了代码的组织风格.例如,你要写一个操作系统内核ProjectOS,还要写一个游戏ProjectGame.在传统的开发习惯中,这两个项目会组织到两个不同的Repo里,PorjectOS和ProjectGame之间无法直接相互引用,例如,你在ProjectOS里写了一个高级的数据结构,想要在Game里也使用,要么直接复制粘贴,要么是创建一个新的CommonRepo,把可公用的代码都放在Common里,然后两个项目各自引入Common作为依赖. 使用MonoRepo则不存在这个问题,Game可以直接依赖OS内的组件,按照Bazel的语法描述,就是在Game中可以直接使用@ProjectOS//path/to/package:AdvancedStruct.当然,你仍然可以选择重构一

By Edimetia3D

Unix related things

这是一篇2017年左右的记录, 仅用作分享 杂 * 在shell内能干的事,我们都可以比较简单地通过系统调用实现. * `称为反引号,^称为脱字符,常用来表示CTRL * windows的系统调用是不开放的,windows下只能直接使用windows.h里的windows API. * /dev目录下的设备是供用于程序直接使用的,主要由block,char,pipe,socket类型 * 并不是所有设备都能映射为这种形式 * /sys/device/目录称为sysfs,他下面存放了所有设备的信息.(不能直接从/dev获得任何设备信息) * udevadm info --query=all --name="/dev/sda1"可以用于查询/dev下某个设备对应的sysfs路径 权限系统 * 权限系统由两部分组成 * 文件属性:用于标注文件owner,所属组,以及权限的设定(默认只有owner和root可以修改权限设置) *

By Edimetia3D