大语言模型中的位置编码

1. 为什么需要位置编码

词向量本身没有不包含位置信息。
自注意力机制对词向量的排布不敏感，无法区分词序。

知乎: Transformers 中，词向量的排列本身就包含了位置信息，为什么还需要位置编码？

Arxiv: Transformer Language Models without Positional Encodings Still Learn Positional Information

2. 位置编码

位置编码大致可以分为两类： Absolute Positional Encoding 和 Relative Positional Encoding。代表性方法包括：Sinusoidal PE 和 RoPE。

位置编码分类，[图源](https://arxiv.org/pdf/2312.17044)

2.1 Preliminary

在transformer架构中，编码器部分的数据流可以如下表示：

输入词向量序列 , 其中为序列长度，为词向量维度。
, 其中为通过线性变换得到的 Query, Key, Value 矩阵。计算为注意力权重。
, 计算得到注意力分数.
, 计算得到输出。
, 计算得到最终输出。

其中，为全连接网络。

2.2 Absolute Positional Encoding

Sinusoidal PE使用正弦和余弦函数来编码位置信息，其公式如下：

其中，为位置，为维度。在维的词向量中,每两个一组,分别对应正弦和余弦函数, 并且每个位置编码可由其它位置编码线性组合得到, 可以使模型学习到相对位置关系.

代码实现如下：

import numpy as np

def get_positional_encoding(n_position, d_pos_vec):
    pos_m = np.array([[pos / np.power(10000, 2 * (i // 2) / d_pos_vec) for i in range(d_pos_vec)] for pos in range(n_position)])
    pos_m[:, 0::2] = np.sin(pos_m[:, 0::2])  # dim 2i
    pos_m[:, 1::2] = np.cos(pos_m[:, 1::2])  # dim 2i+1
    return pos_m

使用相对位置编码不利于模型长度外推, 存在一些变体如下:

添加平移不变性, 在位置编码中随机添加全局偏移量和局部偏移量,使得模型对于相对距离的敏感度降低。如CAPE: Encoding Relative Positions with Continuous Augmented Positional Embeddings - 2021
增强平滑性, 如利用复数融合词向量表示与位置信息. Encoding word order in complex embeddings - 2019

2.3 Relative Positional Encoding

相对位置编码关注于词向量之间的相对位置. 在后续的自注意力分数计算中, 采用相对位置编码的通用形式如下:

自注意力权重为词向量和相对位置的函数, 其中为相对位置信息, 为添加位置信息的任意操作。

旋转位置编码(Rotary Position Embedding RoPE)基于旋转矩阵实现相对位置编码，以往的大多数位置编码方式将位置编码信息直接作用上词向量上，破坏了向量的独立性，不适合线性自注意力。

设词向量序列为 $\bold{E}N={x_i}^N{i=1} $其中$ x_i\in R^d $为词向量，$ N$ 为序列长度。则自注意力机制中的QKV向量可表示为：

其中，为线性变换函数，为位置索引。

则对位置的注意力权重可表示为：

自注意力计算

则相对位置编码问题可建模为寻找函数，满足：

论文中给出的一个解为：

相对位置编码的一个解

其中，为旋转矩阵，为旋转角度。

旋转矩阵

原始注意力计算公式如下，复杂度为：

original self attention

线性自注意力机制通过优化相似度计算，将复杂度降低至：

linear attention

其中。

添加RoPE后，线性注意力计算公式如下：

RoPE attention

RoPE的代码参考llama3的模型可按如下方式实现：

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    # 1/(theta ^ [0, 2, 4, ..., dim] / dim)
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # [0, 1, 2, ..., end]
    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
    # [end x dim // 2]
    freqs = torch.outer(t, freqs)
    # torch.polar(幅度, 相位角)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
    return freqs_cis

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    # 只保留第2维（序列长度）和最后一维（头维度/2）的实际大小，其它维度为1，用于广播
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(
    xq: torch.Tensor, # [bsz, seqlen, n_local_heads, head_dim]
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # 每两个相邻值表示一个复数的实部和虚部
    # [bsz, seqlen, n_local_heads, head_dim] -> [bsz, seqlen, n_local_heads, head_dim/2, 2]
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    # (R_Θ(m)·q)ᵀ(R_Θ(n)·k) = qᵀ·R_Θ(n-m)·k
    # 两个不同位置的向量内积会自然体现它们的相对位置
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)


class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        # 并行训练的设备，GPU数量
        model_parallel_size = fs_init.get_model_parallel_world_size()
        # 当前设备的Query头数量
        self.n_local_heads = args.n_heads // model_parallel_size
        # 当前设备的KV头数量
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        # 分组查询注意力，每个KV头需要复制的次数
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = args.dim // args.n_heads

        # 特殊线性层，每个 GPU 只存储并计算权重矩阵的一部分列
        self.wq = ColumnParallelLinear(
            args.dim,
            args.n_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wk = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wv = ColumnParallelLinear(
            args.dim,
            self.n_kv_heads * self.head_dim,
            bias=False,
            gather_output=False,
            init_method=lambda x: x,
        )
        self.wo = RowParallelLinear(
            args.n_heads * self.head_dim,
            args.dim,
            bias=False,
            input_is_parallel=True,
            init_method=lambda x: x,
        )

        # KV缓存张量
        self.cache_k = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,
                args.max_seq_len,
                self.n_local_kv_heads,
                self.head_dim,
            )
        ).cuda()

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int,
        freqs_cis: torch.Tensor,
        mask: Optional[torch.Tensor],
    ):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        # 对Q,K分别进行旋转，使得不同位置的向量内积体现相对位置
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)

        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv

        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

        # repeat k/v heads if n_kv_heads < n_heads
        keys = repeat_kv(
            keys, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)
        values = repeat_kv(
            values, self.n_rep
        )  # (bs, cache_len + seqlen, n_local_heads, head_dim)

        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
        keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        values = values.transpose(
            1, 2
        )  # (bs, n_local_heads, cache_len + seqlen, head_dim)
        scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
        if mask is not None:
            scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(output)

RoPE相比于Sinusoidal PE的优势在于：

RoPE是相对位置编码，两个位置向量内积自然体现相对距离，Sinusoidal PE是绝对位置编码。
RoPE：通过旋转矩阵旋转向量，是乘性操作，Sinusoidal PE：直接将位置向量加到词向量上，是加性操作
RoPE：保留了词向量的独立性，不直接修改原始词向量，Sinusoidal PE：直接修改词向量，破坏了词向量的独立性。
RoPE：线性自注意力机制，复杂度为，Sinusoidal PE：原始自注意力机制，复杂度为。
RoPE：可以在不同设备上并行计算，Sinusoidal PE：需要在同一设备上计算。
RoPE具有更强的长度外推能力。

2.4 可学习的位置编码

将位置编码作为模型参数，通过模型训练得到。

如T5模型考虑两个位置之间的相对位置关系，在计算注意力权重乘积之后，加上基于对数线性模型的位置编码。即

T5模型的位置编码

3. 大语言模型的上下文长度外推

大语言模型的上下文长度外推是指模型在训练时只使用了有限长度的上下文，但在推理时需要处理更长的上下文。一般的解决方式为插值。

基于RoPE的方法包括

线性插值：对位置索引进行缩放，
NTK-aware缩放：对旋转频率进行缩放，

大模型

#LLM #PE

大语言模型中的位置编码

https://wenzhaoabc.github.io/llm/postion_encoding/

作者

wenzhaoabc

发布于

2025年2月12日

许可协议

分词和词嵌入上一篇

大语言模型中的采样参数下一篇