文章目录
前言
随着自然语言处理(NLP)领域的快速发展,预训练的语言模型如BERT、GPT系列、PaLM、Qwen等取得了显著的成功。这些模型能够有效地捕捉文本中的语义信息,并在各种下游任务中表现出色。然而,在处理长文本序列时,准确地建模词项之间的相对位置关系对于提高模型性能至关重要。传统的绝对位置编码方法存在一定的局限性,尤其是在处理非常长的序列时。因此,研究者们提出了多种改进方案,其中旋转位置编码(RoPE)因其独特的优势而受到了广泛关注。本文旨在详细介绍旋转位置编码(Rotary Position Embedding, RoPE),一种用于增强深度学习模型特别是大型语言模型(LLMs)捕捉序列数据相对位置信息的方法。文章简单解释RoPE的基本原理,重点介绍RoPE代码实现与应用。
一、 旋转位置编码原理
1、RoPE概述
旋转位置编码是一种基于正弦和余弦函数的位置编码方法,它允许模型直接操作位置信息,而不是简单地将其添加到输入向量中。RoPE的核心思想是利用复数域内的旋转矩阵来表示位置信息,从而使得不同位置的信息可以通过简单的线性变换进行交互。这种方法不仅提高了模型对长序列的理解能力,而且有助于减少参数数量,加快训练速度。
2、 复数域内的旋转
RoPE的关键在于引入了复数域的概念,即每个位置都由一对实部和虚部组成。具体来说,对于一个维度为d的隐藏状态,我们定义前d/2个元素作为实部,后d/2个元素作为虚部。然后,通过构造特定形式的旋转矩阵,可以实现任意两个位置之间信息的有效传递。
1、位置编码生成
给定一个最大序列长度max_seq_len
,我们可以预先计算出所有可能位置对应的旋转矩阵。这一步骤通常只需要执行一次,并且可以被保存下来供后续使用。位置编码的具体生成过程如下:
- 对于每个位置i,计算其对应的频率
inv_freq = 1 / (base ** (torch.arange(0, dim, 2) / dim))
,其中base
是一个超参数,默认值为10000。 - 使用
inv_freq
构造旋转矩阵freqs = torch.outer(seq_idx, inv_freq)
,这里seq_idx
是位置索引的向量。 - 最终得到的位置编码是由cosine和sine组成的四元组
cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
。
2、 应用位置编码
一旦获得了位置编码缓存rope_cache
,就可以很容易地将其应用于模型的隐藏状态中。具体而言,对于每个时间步t,我们取出相应的旋转矩阵,并将其作用于当前时刻的隐藏状态上。这样做的好处是可以直接改变位置信息,而不需要额外的学习参数。
二、RoPE的实现细节
1、RotaryEmbedding类设计
为了方便地生成和管理位置编码,我们可以创建一个名为RotaryEmbedding
的PyTorch模块。该类包含初始化方法__init__()
、生成位置编码的方法forward()
以及辅助函数impl()
和forward_impl()
。以下是RotaryEmbedding
类的主要组件:
__init__(self, dim, rope_ratio=1, original_impl=False, device=None, dtype=None)
:初始化成员变量并设置默认参数。impl(self, seq_length: int, dim: int, device: torch.device, dtype: torch.dtype)
:原始实现版本的位置编码生成逻辑。forward_impl(self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000)
:优化后的实现版本。forward(self, max_seq_len, offset=0)
:调用上述两种实现之一来生成最终的位置编码。
2、apply_rotary_pos_emb函数
有了RotaryEmbedding
类之后,下一步就是编写apply_rotary_pos_emb
函数,用于将生成的位置编码应用于给定的输入张量x
。此函数接收两个参数:一个是待处理的张量x
,另一个是之前生成的位置编码缓存rope_cache
。它的主要工作包括:
- 分离需要应用位置编码的部分和不需要的部分。
- 调整
rope_cache
以匹配输入张量的形状。 - 执行旋转操作并将结果与未处理部分合并。
def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
b, np, sq, hn = x.size(0), x.size(1), x.size(2), x.size(3)
rot_dim = rope_cache.shape[-2] * 2
x, x_pass = x[..., :rot_dim], x[..., rot_dim:] # 最后维度的hn分成两份
rope_cache = rope_cache[:, :sq]
xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)
rope_cache = rope_cache.view(-1, 1, sq, xshaped.size(3), 2)
x_out2 = torch.stack([
xshaped[..., 0] * rope_cache[..., 0] - xshaped[...,