碎碎念:Transformer的细枝末节

碎碎念:Transformer的细枝末节

更新预告:dropout设置策略 | 激活函数 | 已经 2020 年了

目录(缓更中):

  1. weight tying(WT)
  2. Feed-Forward Networks(FFN)
  3. Positional Encoding
  4. Scaled Dot-Product Attention
  5. Layer Normalization(LN)【2019.05.13 更新】
  6. Label Smoothing(LSR) 【2019.05.21 更新】
  7. Warmup & Noam 学习率更新 【2019.06.06 更新】
  8. Xavier 初始化【2020.04.15 更新】

最近我被问到,“你是怎么理解 Transformer 模型的?”

当时有点语塞。我曾在2018年做过 Transformer 和 RNN 的对比实验,当时只关注到了self-attention 和 position-encoding。如果翻一翻相关的博客,基本也都是介绍这些内容。

到了2019年,我开始重新研究 Transformer 模型,想努力弄明白一些让人困惑的细节。我将一些【个人理解】写在这里作为记录,不照搬照抄,希望能帮助到像我一样四处搜不到答案的同学。

在我看来,Transformer 能够全面超越 RNN,self-attention 固然是最重要的,但是只有这个是无法复现 state-of-art 的结果的。Transformer 的模型设计实在是精彩,但为了能让这个庞然大物 work 起来,在论文里被一笔带过的 trick 实在是功不可没。

本文适合已经熟悉了 Transformer 模型总体结构的同学,帮助加深理解。

1. Weight Tying

论文的3.4小节Embeddings and Softmax,有这样一句话:

In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to [29].

粗读的我觉得迷迷糊糊的,后来追着去读了参考文献[29] 《Using the Output Embedding to Improve Language Models》,才觉得清晰一点。这篇文章主要介绍的是RNNLM中的Weight Tying技术,不仅能压缩LM的大小,还能显著改善PPL表现。作者为了证明算法的普适性,在NMT任务上也进行了扩展实验。在seq2seq模型中,decoder可以近似地看作RNNLM,它必不可少地有一个embedding矩阵( U \in R^{C \times H} )和一个pre-softmax矩阵( V \in R^{C \times H} ),来完成词表大小C到隐层大小H的尺度转换:

h_{in} = U ^T C ,......,h_{pre\_softmax}=Vh_{out} \\ Weight Tying在操作上非常简单,即令 U=V 。在OPEN-NMT的Pytorch实现版本中,仅仅一行代码:

if model_opt.share_decoder_embeddings:
    generator[0].weight = decoder.embeddings.word_lut.weight

Transformer在英法和英德上,混用了源语言和目标语言的词表,因此使用了升级版的TWWT(Three way weight tying),把encoder的embedding层权重,也加入共享:

if model_opt.share_embeddings:
    tgt_emb.word_lut.weight = src_emb.word_lut.weight

虽然weight共享了,但是embedding和pre-softmax仍然是两个不同的层,因为bias是彼此独立的。

在我个人的理解中,one-hot向量和对U的操作是“指定抽取”,即取出某个单词的向量行;pre-softmax对 V 的操作是“逐个点积”,对隐层的输出,依次计算其和每个单词向量行的变换结果。虽然具体的操作不同,但在本质上,UV都是对任一的单词进行向量化表示(H列),然后按词表序stack起来(C行)。因此,两个权重矩阵在语义上是相通的

也是由于上面两种操作方式的不同,U 在反向传播中不如V 训练得充分。将两者绑定在一起,缓和了这一问题,可以训练得到质量更高的新矩阵。另外,由于词表大小 C 通常比模型隐层大小 H 高出一个数量级,Weight Tying 可以显著减小模型的参数量。这个数据在论文中是28%~51%,我使用Transformer-base在非共享词表的英中方向上进行测试,模型参数量从 131,636,930 减小到 102,177,474,足足的22%。模型更小,收敛更快更容易。

下图是我自己用Transformer-base模型跑的实验结果,绿色是Weight Tying版本,蓝色是基础版本。X轴为step,Y轴为Appromix BLEU。可以看到,两者的收敛速度接近,曲线平稳后绿色明显优于蓝色,峰值相差0.17。

绿色为WT策略

2. Feed-Forward Networks

在每个子层中,Multi-Head Attention层都接了一个FFN层,公式是这样子:

FFN(x)=max(0, xW_1+b1)W_2+b2 \\ 顺序上是先线性变换,然后ReLU非线性,再线性变换。论文的3.3小节Position-wise Feed-Forward Networks中,对“Position-wise”做了注解:

...which is applied to each position separately and identically...While the linear transformations are the same across different positions, they use different parameters from layer to layer.

对于输入序列中的每个位置 x ,使用相同的变换矩阵,且每个子层使用的不用的参数。在计算完Multi-Head Attention后,FFN层的输入矩阵为 X\in R^{d_{input} \times d_{model}},可以看作是由每个输入位置( d_{input} 行)的attention结果( d_{model} 列)堆叠而成。这些行进行相同的线性变换后,维度改变,重新堆叠成FFN层的输出。行与行之间无交错,完全是“separately and identically”,按位置进行变换。如下图:

Position-Wise

矩阵 W_1\in R^{d_{model} \times d_{ff}}W_2\in R^{d_{ff} \times d_{model}} 的维度倒置,作者认为可以理解为“two convolutions with kernel size 1”。对于d_{model}=512d_{ff} \in [1024,2048,4096],最终的BLEU值分别为 [25.4,25.8,26.2] ,稳步提升。FFN 相当于将每个位置的Attention结果映射到一个更大维度的特征空间,然后使用ReLU引入非线性进行筛选,最后恢复回原始维度。需要说明的是,在抛弃了 LSTM 结构后,FFN 中的 ReLU成为了一个主要的能提供非线性变换的单元

3. Positional Encoding

RNN天然是有序的,而Transformer解除了时序依赖。位置编码因此被引入进来,它将词序信息向量化,是模型中不可或缺的一部分。

在实现方式上,分为Facebook版本(《Convolutional Sequence to Sequence Learning》)和Google版本。作者形容前者为“learned and fixed”,即“Postional Embedding”,它只能表征有限长度内的位置,无法对任意位置进行建模。后者直接上了公式,改“Embedding”为“Encode”,想要多长就有多长。实验结果表明,两种形式的模型没有效果差别,但毕竟通过公式来计算更简单、参数量也更小

PE_{(pos,2i)}=sin(pos/10000^{2i/d_{model}})  \\ PE_{(pos,2i+1)}=cos(pos/10000^{2i/d_{model}})

至于为什么要坐标分组,即同时使用正弦和余弦,论文的3.5小节中这样解释:

We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, PEpos+k can be represented as a linear function of PEpos.

相隔 k 个词的两个位置 pos 和 pos+k 的位置编码是由 k 的位置编码定义的一个线性变换,推导公式如下:

PE(pos+k,2i)=PE(pos,2i)PE(k,2i+1)+PE(pos,2i+1)PE(k,2i) \\ PE(pos+k,2i+1)=PE(pos,2i+1)PE(k,2i+1)-PE(pos,2i)PE(k,2i)

这组公式由三角函数的“和角”公式推导而来,要实现这种相对位置的建模,sincos 都必不可少。

不过,倒不是非要设置奇偶下标 2i2i+1。官方代码tensor2tensor的最初版本只是简单地分了两段(issue):

 return tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1)
// PE(pos, 0:channels/2) = sin(scaled_time)
// PE(pos, channels/2:)  = cos(scaled_time) 

两种方案的训练速度和模型精度毫无差异。作者解释说,后面的全连接层可以帮助重排坐标,所以sincos 的实际排布可以怎么简单怎么来:

I think this does not matter, since there is a fully connected layer after the encoding is added (and before) and it can permute the coordinates.

相关的代码我没有找到....希望明白的同学来交流一下^___^

最后,每个坐标位置的三角函数波长是不同的,论文中的范围从 2\pi\ \rightarrow 10000\cdot 2\pi\ 。变化即可,具体的步长和范围并不重要,tensortensor源码中可以进行不同的设置。

4. Scaled Dot-Product Attention

常用的attention主要有“Add-相加”和“Mul-相乘”两种:

score(h_j,s_i)=<v, tanh(W_1h_j+W_2s_i)>[Add] \\ score(h_j,s_i)=<W_1h_j, W_2s_i>  \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ [Mul]

矩阵加法的计算更简单,但是外面套着 tanhv,相当于一个完整的隐层。在整体计算复杂度上两者接近,但考虑到矩阵乘法已经有了非常成熟的加速算法,Transformer采用了Mul形式。

在模型效果上,《Massive Exploration of Neural Machine Translation Architectures》对不同Attention-Dimension( d_k )下的Add和Mul进行了对比,如下图:

Add和Mul的BLEU对比

可以看到,在 d_k 较小的时候,Add和Mul相差不大;随着 d_k 增大,Add明显超越了Mul。Transformer 设置 d_k=64 ,虽然不在表格的范围内,但是可以推测Add仍略优于Mul。作者认为,d_k 的增大将点积结果推向了softmax函数的梯度平缓区,影响了训练的稳定性。原话是这么说的:

We suspect that for large values of dk, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients.

因此,Transformer 为“dot-product attention”加了一个前缀“scaled”,即引入一个温度因子(temperature)\sqrt{d_k} ,中文全称“缩放的点积注意力网络”:

Attention(Q,K,V)=softmax(QK^T/\sqrt{d_k})V\\ 等等,Add也离不开softmax,怎么没有这个问题呢?首先,左侧v 的导数就是 tanh 的输出,后者本身就是 [-1,1] 之间的,一定不会超限。然后,右侧 W_1 的导数就是隐层h_j ,必然是sigmoid或者其他激活函数的输出,值也不会太大。

回到Mul。左侧W_1 的导数来自于 h_j(W_2s_i) 。如果 s_i 分布在 (0,1) ,那么 W_2s_i 就会扩展到 (0,d_k) ,即点积可能会造成一个非常大的梯度值

5. Layer Normalization

Transformer 使用了 2016 年的 Layer Normalization,简称 LN。它在论文中被一笔带过,但却是不可或缺的一部分。每个子层的输出值为LayerNorm(x + Sublayer(x)) ,这在网络结构图上非常明显:

“Norm” 即 LN

LN 是 Normalization(规范化)家族中的一员,由 Batch Normalization(BN)发展而来。基本上所有的规范化技术,都可以概括为如下的公式:

h_i = f(a_i) \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ [调整前] \\ h_{i}^{'}=f(\frac{g_i}{\sigma_i}(a_i-u_i)+b_i) \ \ [调整后]

对于隐层中某个节点的输出为对激活值 a_i 进行非线性变换 f() 后的h_i ,先使用均值 u_i 和方差 \sigma_ia_i 进行分布调整。如果将其理解成正态分布,就是把“高瘦”和“矮胖”的都调整回正常体型(深粉色),把偏离x=0的拉回中间来(淡紫色)。

不同均值和方差下的正态分布 | sigmoid函数

这样做的第一个好处(平移)是,可以让激活值落入 f()梯度敏感区间(红色虚线的中间段)。梯度更新幅度变大,模型训练加快第二个好处是,可以将每一次迭代的数据调整为相同分布(相当于“白化”),消除极端值,提升训练稳定性

然而,在梯度敏感区内,隐层的输出接近于“线性”,模型表达能力会大幅度下降。引入 gain 因子 g_i 和 bias 因子 b_i ,为规范化后的分布再加入一点“个性”。需要注意的是, g_ib_i 作为模型参数训练得到u_i\sigma_i 在限定的数据范围内统计得到。BN 和 LN 的差别就在这里,前者在某一个 Batch 内统计某特定神经元节点的输出分布(跨样本),后者在某一次迭代更新中统计同一层内的所有神经元节点的输出分布(同一样本下)。盗用一张 @gylight 的图:

那么,为什么要舍弃 BN 改用 LN 呢?朴素版的 BN 是为 CNN 任务提出的,需要较大的 BatchSize 来保证统计量的可靠性,并在训练阶段记录全局的 u\sigma 供预测任务使用。对于天然变长的 RNN 任务,需要对每个神经元进行在每个时序的状态进行统计。这不仅把原本非常简单的 BN 流程变复杂,更导致偏长的序列位置统计量不足。相比之下,LN 的使用限制就小很多,不需要在预测中使用训练阶段的统计量,即使 BatchSize = 1 也毫无影响。

个人理解,对于 CNN 图像类任务,每个卷积核可以看做特定的特征抽取器,对其输出做统计是有理可循的;对于 RNN 序列类任务,统计特定时序每个隐层的输出,毫无道理可言——序列中的绝对位置并没有什么显著的相关性。相反,同一样本同一时序同一层内,不同神经元节点处理的是相同的输入,在它们的输出间做统计合理得多。

从上面的分析可以看出,Normalization 通常被放在非线性化函数之前。以 GRU 为例,来看看 LN 是怎么设置的:

蓝色方框内为一个单独的LN单元

可以看到,总体的原则是在“非线性之前单独处理各个矩阵”。对于 Transformer,主要的非线性部分在 FFN(ReLU) 和 Self-Attention(Softmax) 的内部,已经没有了显式的循环,但这些逐个叠加的同构子层像极了 GRU 和 LSTM 等 RNN 单元。信息的流动由沿着时序变成了穿过子层,把 LN 设置在每个子层的输出位置,意义上已经不再是“落入sigmoid 的梯度敏感空间来加速训练”了,个人认为更重要的是前文提到的“白化”—— 让每个词的向量化数值更加均衡,以消除极端情况对模型的影响,获得更稳定的深层网络结构 —— 就像这些词从 Embdding 层出来时候那样,彼此只有信息的不同,没有能量的多少。在和之前的 TWWT 实验一样的配置中,删除了全部的 LN 层后模型不再收敛。LN 正如 LSTM 中的tanh,它为模型提供非线性以增强表达能力,同时将输出限制在一定范围内。 因此,对于 Transformer 来说,LN 的效果已经不是“有多好“的范畴了,而是“不能没有”

此外,在这些细节策略上,各个版本的实现代码和论文不完全一致。比如 OPEN-NMT 中 实现的版本为 x + Sublayer(LayerNorm(x)) ,且对编解码器的输出再各做一次 LN。最终模型性能的区别不大。

6. Label Smoothing

论文的 6.4 小节,集中介绍 Transformer 的正则化技术,Label Smoothing 就是其中的一部分。作者认为虽然 ppl 受到影响,但是 bleu 会提升:

This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.

Label smoothing 的全称是 Label Smoothing Regularization(LSR),是 2015 年的经典论文 《Rethinking the Inception Architecture for Computer Vision》的一个副产品。同样来自 Google,只不过被从 CV 搬到了 NLP 里。既然是正则化技术,那么核心功能就是防止过拟合。除此之外,作者认为 LSR 可以避免模型 too confident。关于这一点论文里没有详细的解释,不过对于 NMT 任务来说,鼓励模型产生多样化的译文确实会帮助提升 BLEU 值,毕竟标准译文并不是唯一的。

最著名的正则化技术,dropout,通过随机抹掉一些节点来削弱彼此之间的依赖,通常设置在网络内部隐层。作为 dropout 的配合策略,LSR 考虑的是 softmax 层。

假设目标类别为 y ,任意类别为 k ,ground-truth 分布为 q(k) ,模型预测分布为 p(k) 。显然,当 k=y 时, q(k)=1。当 k\ne y 时, q(k)=0 。LSR 为了让模型的输出不要过于贴合单点分布,选择在 gound-truth 中加入噪声。即削弱 y 的概率,并整体叠加一个独立于训练样例的均匀分布 u(k)

q'(k)=(1-\epsilon)q(k)+\epsilon u(k)=(1-\epsilon)q(k)+\epsilon/K \\ 其中, K 是 softmax 的类别数。拆开来写可以看得清楚一点:

q'(k)=1-\epsilon + \epsilon/K,k=y \\ q'(k)=\epsilon/K,\ \ \ \ \ \ \ \ \ \ \ \ \ k\ne y

所有类别的概率和仍然是归一的。说白了就是把最高点砍掉一点,多出来的概率平均分给所有人。调整之后,交叉熵(损失函数)也随之变化

H(q',p)=(1-\epsilon)H(q,p) + \epsilon H(u.p) \\ 对于两个完全一致的分布,其交叉熵为0。LSR 可以看作是在真题优化目标中加入了正则项 H(u,p) ,在模型输出偏离均匀分布时施以惩罚。Transformer 中设置权重 \epsilon=0.1 ,和原论文一致。LSR的具体效果可以如下图实验,绿色的曲线始终不如蓝色。

绿色为关闭 Label Smoothing

7. Warmup & Noam 学习率更新

论文的 5.3 小节介绍优化器时,提出了一个全新的学习率更新公式:

lrate=d_{model}^{-0.5}\cdot min(step\_num^{-0.5},step\_num\cdot warmup\_steps^{-1.5}) \\ 看着有点长,但其实如果把 min() 脱掉的话,就变成一个warmup\_steps 为分界点的分段函数。

在该点之后, lrate=d_{model}^{-0.5}\cdot step\_num^{-0.5} ,是 decay 的部分。常用方法有指数衰减(exponential)、分段常数衰减(piecewise-constant)、反时限衰减(inverse-time)等等。Transformer 采用了负幂的形式,衰减速度先快后慢。

在该点之前,lrate=d_{model}^{-0.5}\cdot step\_num\cdot warmup\_steps^{-1.5},是 warmup 的部分。Transformer 采用了线性函数的形式,warmup_steps 越大,斜率越小。

不同d_model 和 warm下的学习率变化

画在图上明显很多。Transformer 的学习率更新公式叫作“noam”,它将 warmup 和 decay 两个部分组合在一起,总体趋势是先增加后减小。据我所知,Transformer 是第一个使用类似设置的 NLP 模型。

第一个问题,为什么要额外设置一个 warmup 的阶段呢?虽然 NLP 中鲜有相关尝试,但是 CV 领域常常会这样做。在《Deep Residual Learning for Image Recognition》(残差网络出处)中,作者训练 110 层的超深网络时就用过类似策略:

In this case, we find that the initial learning rate of 0.1 is slightly too large to start converging. So we use 0.01 to warm up the training until the training error is below 80% (about 400 iterations), and then go back to 0.1 and continue training.

对于 Transformer 这样的大型网络,在训练初始阶段,模型尚不稳定,较大的学习率会增加收敛难度。因此,使用较小的学习率进行 warmup,等 loss 下降到一定程度后,再恢复回常规学习率。

第二个问题,为什么要用线性增加的 warmup 方式?《Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour》有专门讨论 warmup 的小节。这篇文章将 ImageNet 上的训练时间从 Batchsize 256 in 29 hours 缩短到 Batchsize 8192 in 1 hour,而完全不损失模型精度。在解决超大 Batchsize 难优化问题时,warmup 是一个极其重要的策略。其中,线性增加的方法称为 gradual warmup,使用固定值的方法称为 constant warmup。后者的缺点是在 warmup 阶段结束时,学习率会有一次跳动抬升。

large minibatch sizes are challenged by optimization difficulties in early training and if these difficulties are addressed, the training error and its curve can match a small minibatch baseline closely.

作者认为,为了实现超大Batchsize,需要保证“k 个 minibatch , size = n , lr = η” 和 “1 个 minibatch , size = kn , lr = kη”的梯度近似相等。但是在模型变化剧烈时,这个等式会被打破。warmup 可以有效缓解这个问题。

不同 warmup 方案下的训练 error

作者对比了三种 warmup 策略,可以看到在 gradual warmup 下蓝线和黄线最接近,即超大 Batchsize 的模型精度逼近普通 Batchsize。

8. Xavier 初始化

在官方 tensor2tensor 代码中,可以看到关于 weight 参数初始化的时候,采用了一种叫作xavier_uniform 的方法:

tensor2tensor 中 encoder 的构建部分

也可以简称为 Xavier(读作 [ˈzeɪvjər]) 或者 Glorot,来自于作者 Xavier Glorot 在 2010 年和 Bengio 大神一起发表的 《Understanding the difficulty of training deep feedforward neural networks》

Xavier 是一种均匀初始化。其基线是朴素版本,即对于包含 n_i 个输入单元的网络层 layer_i ,其 weight 的初始值均匀采样自 [-1/\sqrt(n_i), 1/\sqrt(n_i)] 。作者发现,在深层神经网络中,朴素版本的均匀初始化和激活函数(sigmoid 与 tanh)配合得不好。具体说来,就是不仅收敛得慢,训练平稳后的模型效果还差。如果使用预训练的模型来初始化网络,则能够明显改善训练过程。因此,作者认为,更换初始化方式是可取的。

各层激活值分布(sigmoid + 朴素的均匀初始化)

首先来看看 sigmoid。它的线性区在 0.5 附近,偏向 1 或 0 表示则表示饱和。在理想情况下,激活值应当落在表达能力最强的非线性区间,远离线性区与饱和区。然而,在实际训练中,输出层 layer-4 在很长一段时间内都接近下饱和(黑);其他层由上往下离线性区越来越近(蓝绿红)。虽然从 epoch-100 开始慢慢正常了,但这显然拖慢了收敛过程。

各层激活值分布(tanh + 朴素的均匀初始化)

然后是正负对称的 tanh。随着训练进行,自下而上(红绿蓝黑青),各层慢慢地都落入了饱和区。明明验证集的效果还达不到预期,模型却“训不动”了。

那么,我们需要什么样的激活值分布呢?一个标准的第 i 层可以分为线性s^i=z^iW^i+b^i 和非线性 z^{i+1}=f(s^i) 两部分。在深层网络中,激活值的方差会自下而上逐层累积, 链式推导可得:Var[z^i]=Var[x]\prod_{j=0}^{i-1}n_jVar[W^j] \\

对于反向传播,假定输出落在激活函数的线性区,即 f'(s^i_k)\approx1 。同样应用链式法则,各参数的梯度为(d 为总层数):

Var[\frac{\partial{Cost}}{s^i}]=Var[\frac{\partial{Cost}}{s^d}]\prod_{j=i}^{d}n_{j+1}Var[W^j] \\ Var[\frac{\partial{Cost}}{w^i}]=Var[\frac{\partial{Cost}}{s^i}]Var[z^i] \\

现在要回到正题上了。一个稳定的深层网络,要求 Var[z^i]Var[\partial{Cost} /{s^i}] 在各层保持一致,这样可以让激活值(forward)始终远离饱和区,并且梯度(backward)不会消失或爆炸。根据上面的公式,可以得到约束目标:

\forall{i},n_iVar[W^i]=1=n^{i+1}Var[W^i] \\ 注意,由于方差是和输入单元的个数有关的,所以正向和反向分别使用了 n_in_{i+1} 。这两个等式显然不能同时成立,只好折中一下,让 Var[W^i]=2/(n_i+n_{i+1}) 。补充一条基础知识,对于均匀分布 U[a, b] ,其方差为 (b-a)^2/12 。现在开始解方程,让 a=b 并且 (b-a)^2/12=2/(n_i+n_{i+1}) ,得到 Xavier 均匀分布的最终形式:

U[-\sqrt{\frac{6}{n_i+n_{i+1}}},\sqrt{\frac{6}{n_i+n_{i+1}}}] \\

编辑于 04-15