1. 工作简介
这篇文章是苏剑林的一篇关于Transformer当中的位置编码的优化考察。
众所周知，transformer的attention机制本身是不带有位置信息的，因此对于文本序列，attention机制本身就会丢失掉原文当中的序列信息，造成信息缺失，影响到模型的效果表达，这个应该已经算是面试中常见的八股文问题了。
但是，对于位置编码的具体实现，大概会去考虑这个问题的人就大幅减少了。而苏剑林这里就是对这部分内容进行了详细的考察和优化尝试，这点真心还是很佩服他的。
2. 常见位置编码方式
首先，我们来考察一下现有的一些位置编码的方法。
这部分的内容我们主要参考苏剑林的博客《让研究人员绞尽脑汁的Transformer位置编码》进行一些自己的整理。
1. 绝对位置编码
1. Bert
绝对位置编码的一个典型的例子就是Bert模型。
他的思路非常的简单粗暴，既然attention层本身无法识别位置信息，那么我就在输入当中显式地多加一个序列位置信号，然后让模型自己去学习这个序列位置信号的分布。
具体来说，就是讲原本的输入xi 变成了 ，其中， 表示的就是第i个位置的绝对位置编码。
这种编码方式的好处在于说简单直接，而且模型直接参数拟合的方式也不会产生人为的信号偏差，但是缺点在于说推理阶段只能使用训练中预先定义好的位置编码，比如bert，最大输入句长就是512，超过了512模型就失去了编码能力。
此外，位置编码的训练充分度也会是一个需要考虑的问题。
2. Attention Is All You Need
绝对位置编码的另一个典型例子就是vanilla transformer。
不过，不同于Bert那样直接交给模型去训练position embedding，经典Transformer的位置Embedding的定义是直接通过三角函数的方式进行预先定义好的。
他的出发点在于说参考了三角函数的性质：
=cosα⋅cosβ−sinα⋅sinβ
=sinα⋅cosβ−cosα⋅sinβ
这个性质刚好可以用于描述两个位置之间的相对距离关系。
因此，文中先验地给出一个人工预设的位置编码如下：
其中，
表示第
个位置，
，表示position embedding当中某一个具体维度上的值。
可以看到，这样的情况下位置表征就是先验确定的了，我们无需交由模型进行额外地拟合，而且长度使用上也更加自由，基本没有长度限制，但是代价就是给模型增加了人工的先验限制，学习到的embedding特征不但要满足token本身的信息表征，还要满足三角函数形式的数据分布表达，收缩了解空间的表达域。
2. 相对位置编码
在考察相对位置编码之前，我们首先来看一下位置编码到底做的是一个什么样的事。
它本质上就是在原本的词向量上面额外地加上一个用于标定位置信息的位置向量 ，从而使得attention层的输入带有位置信息。
我们将attention层的运算进行具体的展开如下：
我们将位置向量p 写入之后可以得到:
因此，事实上位置向量的加入本质上也可以通过加入偏置矩阵的方式直接作用到Attention矩阵A 以及V上面。
1. 经典相对位置编码
相对位置编码的思路来自于文献Self-Attention with Relative Position Representations。
如前所述，位置编码的核心就是给每一个位置添加一个具体的position embedding从而令attention层的输入可以识别出其具体的位置，但是由于句长的无限性所以限制了绝对位置编码的使用方法。
而相对位置编码的核心思路就是说通过一个滑动窗口，换句话来说，就是对于每一个位置i , j i,ji,j，如果它们的相对距离相同，那么他们共享同一个position向量。
而对于句长特别长的输入，如果两个位置i , j i,ji,j的距离特别大，那么我们就对其做一个截断，从而规避掉句长带来的限制。
而关于这部分内容的具体实现，文中首先去除了Query当中包含的位置信息，从而将Attention矩阵和输出Output变成了如下的形式：
然后，文中将两个位置相关的矩阵直接替换成了两个相对位置的偏移矩阵，具体而言：
{ A = s o f t m a x ( Q ⋅ ( K + R K ) T d ) O = A ⋅ ( V + R V ) \left\{ \begin{aligned} A &= softmax(\frac{Q \cdot (K + R_K)^T}{\sqrt{d}}) \\ O &= A \cdot (V + R_V) \end{aligned} \right.
⎩
⎪
⎨
⎪
⎧
A
O
=softmax(
d
Q⋅(K+R
K
)
T
)
=A⋅(V+R
V
)
其中，R i , j R_{i,j}R
i,j
表示第j jj个位置相对于第i ii个位置的位置偏移向量，具体而言：
R i , j = P [ c l i p ( i − j , p m i n , p m a x ) ] R_{i, j} = P[clip(i-j, p_{min}, p_{max})]
R
i,j
=P[clip(i−j,p
min
,p
max
)]
因此，R是一个三维矩阵，即R K , R V ∈ R n × n × d R_K, R_V \in \mathbb{R}^{n \times n \times d}R
K
,R
V
∈R
n×n×d
。
故在计算中会略带一点特殊，具体而言可以表达如下：
Attn = tf.math.softmax(
(tf.einsum("bik,bjk->bij", Q, K) + tf.einsum("bik,ijk->bij", Q, R))/tf.math.sqrt(d)
)
O = tf.einsum("bij,bjk->bik", Attn, V) + tf.einsum("bij,ijk->bik", Attn, R)
1
2
3
4
2. XLNet
XLNet较之经典的相对位置编码进一步移除了V上面的位置信息，所有的位置信息全部都只发生在attention矩阵上面。
因此，我们这里只需要考察attention矩阵在这里的具体形式。
如前所述，Attention权重矩阵事实上就是Query（Q QQ）和Key（K KK）的内积，调整权重之后求一个softmax，因此，我们只需要看Q QQ和K KK的内积部分即可。
事实上，在后续的各种相对位置编码当中，似乎都已经只在Attention矩阵当中加入位置信息来影响权重分布，而不会对value加入权重信息。
言归正传，XLNet的位置编码加入的方式具体如下：
Q ⋅ K T = x W Q ⋅ W K T x T + x W Q ⋅ R T + u W Q ⋅ W K T x T + v W Q ⋅ R T Q \cdot K^T = xW_Q \cdot W_K^T x^T + xW_Q \cdot R^T + uW_Q \cdot W_K^T x^T + vW_Q \cdot R^T
Q⋅K
T
=xW
Q
⋅W
K
T
x
T
+xW
Q
⋅R
T
+uW
Q
⋅W
K
T
x
T
+vW
Q
⋅R
T
其中，R RR的定义和经典相对位置编码中的定义相同，而u , v u,vu,v则是两个可训练的向量。
换成伪代码即：
Attn = tf.math.softmax(
(
tf.einsum("bik,bjk->bij", Q, K) \
+ tf.einsum("bik,ijk->bij", Q, R) \
+ tf.einsum("bik,bjk->bij", tf.matmul(u, W_Q), K) \
+ tf.einsum("bik,ijk->bij", tf.matmul(v, W_Q), R)
)/tf.math.sqrt(d)
)
O = tf.einsum("bij,bjk->bik", Attn, V)
1
2
3
4
5
6
7
8
9
3. T5
T5的位置信息编码则更加暴力一些，都不是使用相对位置的偏置向量，而是直接给出一个偏置矩阵，然后对这个矩阵进行训练。
具体而言：
Q ⋅ K T = x W Q ⋅ W K T x T + B Q \cdot K^T = xW_Q \cdot W_K^T x^T + B
Q⋅K
T
=xW
Q
⋅W
K
T
x
T
+B
其中，B ∈ R n × n B \in \mathbb{R}^{n \times n}B∈R
n×n
就是一个权重矩阵。
Attn = tf.math.softmax(
(tf.einsum("bik,bjk->bij", Q, K) + B)/tf.math.sqrt(d)
)
O = tf.einsum("bij,bjk->bik", Attn, V)
1
2
3
4
4. DeBerta
DeBerta矩阵的定义与T5相反，T5是去除了位置与token的交叉项，只保留相对位置产生的偏移矩阵，而DeBerta与之相反，它去除掉了相对位置产生的偏置矩阵，但是留下了两个位置与token之间的交叉项，并将之用相对位置矩阵的方式进行保留。
具体而言：
Q ⋅ K T = x W Q ⋅ W K T x T + x W Q ⋅ W K T R T + R W Q ⋅ W K T x T Q \cdot K^T = xW_Q \cdot W_K^T x^T + xW_Q \cdot W_K^T R^T + R W_Q \cdot W_K^T x^T
Q⋅K
T
=xW
Q
⋅W
K
T
x
T
+xW
Q
⋅W
K
T
R
T
+RW
Q
⋅W
K
T
x
T
用伪代码表述就是：
Attn = tf.math.softmax(
(
tf.einsum("bik,bjk->bij", tf.matmul(x, W_Q), tf.matmul(x, W_K)) \
+ tf.einsum("bik,ijk->bij", tf.einsum("bil,lk->bik", x, W_Q), tf.einsum("kl,ijk->ijl", W_K, R)) \
+ tf.einsum("ijk,bjk->bij", tf.einsum("ijk,kl->ijl", R, W_Q), tf.einsum("bil,lk->bik", x, W_K))
)/tf.math.sqrt(d)
)
O = tf.einsum("bij,bjk->bik", Attn, V)
1
2
3
4
5
6
7
8
3. RoPE方法介绍
RoPE（Rotary Position Embedding）位置编码是Rofermer这篇文献的核心贡献点。
他的核心想法就是，借用苏剑林自己的话说：
通过绝对位置编码的方式实现相对位置编码。
如前，我们已经注意到了，相对位置编码在实现上是直接作用于attention权重矩阵的，这样的实现方式使得类似Linformer这种直接先对Attention矩阵中间结果进行投影的方式无法实现（关于Linformer相关的内容也可以参考我的博客文献阅读：Linformer: Self-Attention with Linear Complexity）。
因此，苏剑林基于复数相乘的特性设计了如下的位置编码函数：
{ f q ( x m , m ) = ( W Q x m ) ⋅ e i m θ f k ( x n , n ) = ( W K x n ) ⋅ e i n θ g ( x m , x n , m , n ) = R e [ ( W Q x m ) ( W K x n ) e i ( m − n ) θ ] \left\{ \begin{aligned} f_q(x_m, m) &= (W_Q x_m) \cdot e^{im\theta} \\ f_k(x_n, n) &= (W_K x_n) \cdot e^{in\theta} \\ g(x_m, x_n, m, n) &= Re[(W_Q x_m) (W_K x_n) e^{i(m-n) \theta}] \end{aligned} \right.
⎩
⎪
⎪
⎨
⎪
⎪
⎧
f
q
(x
m
,m)
f
k
(x
n
,n)
g(x
m
,x
n
,m,n)
=(W
Q
x
m
)⋅e
imθ
=(W
K
x
n
)⋅e
inθ
=Re[(W
Q
x
m
)(W
K
x
n
)e
i(m−n)θ
]
对于二维情况，我们可以给出一组可行解：
f ( q , m ) = ( c o s m θ − s i n m θ s i n m θ c o s m θ ) ( q 0 q 1 ) f(q, m) = \begin{pmatrix} cos m\theta & -sin m\theta \\ sin m\theta & cos m\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix}
f(q,m)=(
cosmθ
sinmθ
−sinmθ
cosmθ
)(
q
0
q
1
)
我们很快可以仿照上述方式给出一组高维情况下的可行解：
此时，由R矩阵的稀疏性，我们可以直接用下述变换来进行替换：
f ( q , m ) = ( c o s m θ 0 c o s m θ 0 c o s m θ d / 2 − 1 c o s m θ d / 2 − 1 ) ⊗ ( q 0 q 1 . . . q d − 2 q d − 1 ) + ( s i n m θ 0 s i n m θ 0 s i n m θ d / 2 − 1 s i n m θ d / 2 − 1 ) ⊗ ( − q 1 q 0 . . . − q d − 1 q d − 2 ) f(q, m) = \begin{pmatrix} cos\ m\theta_0 \\ cos\ m\theta_0 \\ cos\ m\theta_{d/2-1} \\ cos\ m\theta_{d/2-1} \end{pmatrix} \otimes \begin{pmatrix} q_0 \\ q_1 \\ ... \\ q_{d-2} \\ q_{d-1} \end{pmatrix} + \begin{pmatrix} sin\ m\theta_0 \\ sin\ m\theta_0 \\ sin\ m\theta_{d/2-1} \\ sin\ m\theta_{d/2-1} \end{pmatrix} \otimes \begin{pmatrix} -q_1 \\ q_0 \\ ... \\ -q_{d-1} \\ q_{d-2} \end{pmatrix}
f(q,m)=
⎝
⎜
⎜
⎛
cos mθ
0
cos mθ
0
cos mθ
d/2−1
cos mθ
d/2−1
⎠
⎟
⎟
⎞
⊗
⎝
⎜
⎜
⎜
⎜
⎛
q
0
q
1
...
q
d−2
q
d−1
⎠
⎟
⎟
⎟
⎟
⎞
+
⎝
⎜
⎜
⎛
sin mθ
0
sin mθ
0
sin mθ
d/2−1
sin mθ
d/2−1
⎠
⎟
⎟
⎞
⊗
⎝
⎜
⎜
⎜
⎜
⎛
−q
1
q
0
...
−q
d−1
q
d−2
⎠
⎟
⎟
⎟
⎟
⎞
而关于其具体实现，我们摘录苏剑林在他们自己在GitHub上面的伪代码实现如下：
sinusoidal_pos.shape = [1, seq_len, hidden_size] # Sinusoidal position embeddings
qw.shape = [batch_size, seq_len, num_heads, hidden_size] # query hiddens
kw.shape = [batch_size, seq_len, num_heads, hidden_size] # key hiddens
cos_pos = repeat_elements(sinusoidal_pos[..., None, 1::2], rep=2, axis=-1)
sin_pos = repeat_elements(sinusoidal_pos[..., None, ::2], rep=2, axis=-1)
qw2 = stack([-qw[..., 1::2], qw[..., ::2]], 4)
qw2 = reshape(qw2, shape(qw))
qw = qw * cos_pos + qw2 * sin_pos
kw2 = K.stack([-kw[..., 1::2], kw[..., ::2]], 4)
kw2 = K.reshape(kw2, K.shape(kw))
kw = kw * cos_pos + kw2 * sin_pos
# Attention
a = tf.einsum('bjhd,bkhd->bhjk', qw, kw)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
4. 实验效果考察
1. 翻译任务上的表现
RoPE编码方式最直接的一个检验方法就是直接在翻译任务上测试一下，文中也是这么干的，直接在transformer当中将位置编码进行了一下替换，得到结果如下：
可以看到：
RoPE位置编码是可以正常work的。
2. 预训练任务上的表现
然后，文中考察了一下RoPE在预训练任务当中的效果，得到结果如下：
可以看到：
RoPE编码在预训练任务当中也是有效的。
3. GLUE下游Finetune效果
然后，文中还考察了一下与训练得到的模型在GLUE任务当中finetune的效果，得到结果如下：
可以看到：
在STS-B，QQP以及MRPC任务上Roformer效果优于BERT，而在剩余的三个任务当中效果不及BERT。
因此可以认为RoPE编码的效果和经典的方法各有优劣吧。
4. 中文任务中效果表达
最后，文中还在中文语料上面进行了预训练以及下游finetune任务的效果考察，得到的结果和前面基本一致。
5. 总结 & 思考
结论而言，RoPE编码较之传统的几种位置编码方式其实效果上可能难分伯仲，但是其优点在于说既不受句长的限制，又可以被应用于Linformer，这点较之于他的前辈们还是有很大的优势的。
因此，如果是选择长文本处理任务的位置编码，可能RoPE会是一个更好的选择方案。
6. 参考链接
https://kexue.fm/archives/8265
https://kexue.fm/archives/8130
