ChatGLM-6B 主要代码分析 RotaryEmbedding

flyfish

图片链接地址

传统的 Transformer 位置编码（Positional Encoding）被称为绝对位置编码 ，而 Rotary Embedding 被称为相对位置编码 ，主要是因为它们编码位置信息的方式不同，进而影响模型对序列中元素之间位置关系的理解。

1. 传统 Transformer 位置编码：绝对位置编码

在传统的 Transformer 模型中，位置编码使用正弦和余弦函数将每个位置 $t$ 映射到一个固定的向量：$$PE(t,2i)=\sin \left(\frac{t}{10000^{2i/d}}\right)$$

$$PE(t,2i+1)=\cos \left(\frac{t}{10000^{2i/d}}\right)$$
其中，$t$ 是序列中的位置索引，$i$ 是维度索引，$d$ 是嵌入维度。

特点：

• 固定位置编码 ：每个位置 $t$ 的编码是固定的，无论它出现在序列的哪个部分，其编码都是由位置 $t$ 唯一确定的。

• 不变性 ：这种编码方式不会随着序列的变化而变化，意味着同一位置的编码在每次出现时都是相同的。

绝对性：

• 绝对位置感知 ：由于位置编码与序列中的具体位置 $t$ 紧密关联，模型在训练时会将这些编码与特定的序列模式联系起来。这种方式能够让模型感知到序列中每个元素的绝对位置，但对元素之间的相对位置（如相对距离）缺乏直接的建模能力。

• 难以处理相对位置信息 ：在绝对位置编码下，如果需要感知两个元素之间的相对距离或关系，模型必须通过训练学习到这些关系，而不是通过位置编码直接得到。

2. Rotary Embedding：相对位置编码

Rotary Embedding 的核心思想是通过旋转操作，将位置信息嵌入到序列的每个元素中，从而使模型能够自然地感知到序列中元素之间的相对位置关系。

工作原理：

1. 旋转矩阵 ：Rotary Embedding 将位置信息与特征向量通过旋转矩阵结合。假设 $x_1$ 和 $x_2$ 是在位置 $t$ 和 $t+1$ 的特征向量，那么旋转操作后的位置编码变换为：$$R(\theta )\cdot x=\left[\begin{array}{cc}\cos (\theta )& -\sin (\theta )\\ \sin (\theta )& \cos (\theta )\end{array}\right]\cdot \left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]$$
   其中 $\theta$ 是根据位置计算得到的旋转角度。

2. 相对位置感知 ：当两个位置 $t$ 和 $t+1$ 的特征向量进行旋转变换时，模型可以通过旋转角度的差异自然感知到这两个位置之间的相对关系，而无需依赖绝对位置编码。

相对性：

• 相对位置感知 ：Rotary Embedding 通过旋转矩阵直接捕捉相邻元素之间的相对位置信息。例如，元素 $x_1$ 和 $x_2$ 在相邻位置 $t$ 和 $t+1$ 之间的相对关系可以通过旋转角度的差异直接表达。

• 位置编码灵活性 ：由于旋转矩阵使得位置编码可以灵活变化，因此模型能够更自然地处理不同长度的序列和不同的相对位置关系。

3. 绝对 vs. 相对位置编码

• 绝对位置编码 （传统 Transformer）：编码固定，适合处理具体位置相关的任务，但难以直接处理相对位置关系。

• 相对位置编码 （Rotary Embedding）：编码与序列中的相对位置变化相关，更加灵活，适合处理长序列和需要相对位置信息的任务。

Rotary Embedding 与传统位置编码的比较

特点	传统位置编码 (Positional Encoding)	Rotary Embedding
编码方式	正弦和余弦函数的绝对位置编码	旋转矩阵的相对位置编码
位置关系	只能表示绝对位置	更好地表示相对位置
对长序列的处理	长序列时可能失效	能够有效处理长序列
模型适应性	需要在训练期间观察到所有可能位置	更具扩展性，适应超长序列
应用场景	适用于大多数任务	尤其适用于需要处理长序列和复杂依赖关系的任务

import torch
class RotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, base=10000, precision=torch.half, learnable=False):
        super().__init__()
        inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
        inv_freq = inv_freq.half()
        self.learnable = learnable
        if learnable:
            self.inv_freq = torch.nn.Parameter(inv_freq)
            self.max_seq_len_cached = None
        else:
            self.register_buffer('inv_freq', inv_freq)
            self.max_seq_len_cached = None
            self.cos_cached = None
            self.sin_cached = None
        self.precision = precision

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys,
                              error_msgs):
        pass

    def forward(self, x, seq_dim=1, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[seq_dim]
        if self.max_seq_len_cached is None or (seq_len > self.max_seq_len_cached):
            self.max_seq_len_cached = None if self.learnable else seq_len
            t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
            # Different from paper, but it uses a different permutation in order to obtain the same calculation
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            if self.precision == torch.bfloat16:
                emb = emb.float()

            # [sx, 1 (b * np), hn]
            cos_cached = emb.cos()[:, None, :]
            sin_cached = emb.sin()[:, None, :]
            if self.precision == torch.bfloat16:
                cos_cached = cos_cached.bfloat16()
                sin_cached = sin_cached.bfloat16()
            if self.learnable:
                return cos_cached, sin_cached
            self.cos_cached, self.sin_cached = cos_cached, sin_cached
        return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]

    def _apply(self, fn):
        if self.cos_cached is not None:
            self.cos_cached = fn(self.cos_cached)
        if self.sin_cached is not None:
            self.sin_cached = fn(self.sin_cached)
        return super()._apply(fn)
    
    


# 初始化 RotaryEmbedding 模块
dim = 64  # 嵌入维度
rotary_emb = RotaryEmbedding(dim=dim)

# 模拟输入张量
batch_size = 2
seq_len = 10
embedding_dim = dim
x = torch.randn(batch_size, seq_len, embedding_dim)

# 调用 forward 方法
cos, sin = rotary_emb(x)

# 输出 cos 和 sin 的形状
print("Cosine Embedding Shape:", cos.shape)
print("Sine Embedding Shape:", sin.shape)

输出

Cosine Embedding Shape: torch.Size([10, 1, 64])
Sine Embedding Shape: torch.Size([10, 1, 64])

Rotary Embedding 的设计思想是将位置编码嵌入到一个旋转的向量空间中，从而为序列建模提供更强的相对位置感知能力。

1. 三角函数基础

三角函数 cos 和 sin 描述了一个角度在单位圆上的投影，定义如下：$$\cos (\theta )=\frac{\text{邻边}}{\text{斜边}},\sin (\theta )=\frac{\text{对边}}{\text{斜边}}$$
这些函数具有周期性，对于任何角度 $\theta$，都有以下性质：$$\cos (\theta +2\pi )=\cos (\theta ),\sin (\theta +2\pi )=\sin (\theta )$$

2. 位置编码（Positional Encoding）

在传统的 Transformer 模型中，位置编码通过 sin 和 cos 函数来表示输入序列中的位置信息。对于一个给定的位置 $t$，对应的编码可以表示为：$$PE(t,2i)=\sin \left(\frac{t}{10000^{2i/d}}\right)$$

$$PE(t,2i+1)=\cos \left(\frac{t}{10000^{2i/d}}\right)$$
其中，$t$ 是序列中的位置，$i$ 是维度索引，$d$ 是嵌入维度。这个编码方式保证了不同维度具有不同的频率，以便模型能够感知到位置的不同。

3. 旋转嵌入（Rotary Embedding）

Rotary Embedding 是一种改进的相对位置编码方法，其核心思想是将位置信息通过旋转矩阵嵌入到序列中的每个特征向量中。它通过以下步骤实现：

1）. 逆频率生成 ：
首先，生成一个逆频率向量 inv_freq：$${\text{inv\_freq}}_j=\frac{1}{{\text{base}}^{\frac{2j}{d}}}$$
其中 base 通常取 10000，j 是维度索引，d 是嵌入维度。

2）. 频率矩阵生成 ：
接下来，计算频率矩阵 freqs，将逆频率与时间步长（即序列位置）相乘：$${\text{freqs}}_{i,j}=t_i\times {\text{inv\_freq}}_j$$
其中 $t_i$ 是序列位置。

3）. 三角函数编码 ：
频率矩阵的每个元素通过 cos 和 sin 进行编码，并合并为一个编码矩阵：$$\text{emb}=[\cos (\text{freqs}),\sin (\text{freqs})]$$

4）. 旋转变换 ：
在旋转嵌入中，编码后的 cos 和 sin 矢量与输入向量进行旋转变换。给定一个输入向量 $x$ 及其旋转矩阵 $R(\theta )$：$$R(\theta )\cdot x=\left[\begin{array}{cc}\cos (\theta )& -\sin (\theta )\\ \sin (\theta )& \cos (\theta )\end{array}\right]\cdot \left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]$$