Rotary Position Embeddings论文阅读

1. 论文

1.1 Transformer中的position embedding

• Transformer中模型不涉及RNN的循环和CNN的卷积，为了学习到sequence中的前后顺序关系，增加了position embedding, embedding的方法采用了sinine和cosine来进行。
  

• 示例如下, 对I am a robot进行编码：
  

• 编码过程中d_model=4, 论文中的10000对应改为100，结果如下：
  

1.2 RoPE: Rotary Position Embeddings

1.2.1 绝对位置编码

• 在原始Transformer中采用的position embedding是通过sine和cosine中计算出来的，加到对应的Q/K/V向量中，在文中称这种情况是绝对位置编码(Absolute position embedding), 如下：
  

1.2.2 相对位置编码

对于相对位置编码(Relative position embedding)来说，有以下几种方式：

1. 给Q/K/V向量加上可训练学习的相对位置编码向量，p_r是可通过训练学习的
   

2. 对QK的计算进行decompose分解，核心是通过p_(m-n)的相对编码替换绝对位置编码p_m
   
   

3. 在2的基础上，使用不同的映射矩阵U来替换W
   

4. 还是使用共同的W映射矩阵，唯一区别在于用p_(m-n)相对距离替换p_m和p_n
   

1.2.3 Rotary Position Embeddings

1. q对应位置是m，k对应位置是n，RoPE编码要解决的问题中计算QK内积只依赖相对位置m-n。
   

2. 先考虑二维情况下的推导

   • 使用复数的表示，Re表示计算结果中的复数部分
     

   • 对应f(x_m, m)写成如下，相对位置编码通过位置m和角度θ的旋转矩阵的乘积来表示。
     

3. 更通用的形式

   • 从二维扩到N维，对应的QK计算中相对位置只依赖R的矩阵
     

   • 计算图更形象的展示如下
     

4. PoRE的特性

   • 距离衰减，QK的内积结果会随着相对距离增加而减少
   • 支持线性Attention

5. 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)

2. 参考

• Roformer: Enhanced transformer with rotary position embedding
• A Gentle Introduction to Positional Encoding in Transformer Models, Part 1
• Neural machine translation with a Transformer and Keras
• Transformer升级之路：2、博采众长的旋转式位置编码
• Rotary Transformer