常用位置编码理解

目录

一、正弦曲线位置编码

二、可学习位置编码

三、相对位置编码

以下部分都未解析

 四、旋转位置编码

---

原文：https://zhuanlan.zhihu.com/p/674133494

一、正弦曲线位置编码

实现方式如下：

class SinPositionEncoding(nn.Module):
    def __init__(self, max_sequence_length, d_model, base=10000):
        super().__init__()
        self.max_sequence_length = max_sequence_length #序列的最大长度
        self.d_model = d_model #模型维度，与词嵌入维度相同
        self.base = base #用于计算编码的缩放因子

    def forward(self):
        pe = torch.zeros(self.max_sequence_length, self.d_model, dtype=torch.float)  
        # 初始化size(max_sequence_length, d_model)的零矩阵，用于存储位置编码

        exp_1 = torch.arange(self.d_model // 2, dtype=torch.float)  
        #1.创建一个从0到 d_model/2 - 1 的浮点数序列，初始化一半维度，sin位置编码的维度被分为了两部分 2.tensor([0., 1., 2., 3., 4.])

        exp_value = exp_1 / (self.d_model / 2)
        #tensor([0.0, 0.2, 0.4, 0.6, 0.8])，目的是生成一个在0到1之间的序列，其中每个元素代表了一个缩放因子，当 d_model 较大时，exp_1 中较小的索引（接近0）将对应较大的缩放因子，而较大的索引（接近 d_model // 2）将对应较小的缩放因子

        alpha = 1 / (self.base ** exp_value)  # size(dmodel/2)
        out = torch.arange(self.max_sequence_length, dtype=torch.float)[:, None] @ alpha[None, :]  
        # size(max_sequence_length, d_model/2)，创建一个从0~self.max_sequence_length - 1 的浮点数序列，表示序列中的位置索引。然后将一维数组转换为二维的列向量，None用于增加维度，@矩阵乘法

        embedding_sin = torch.sin(out)#计算正弦编码
        embedding_cos = torch.cos(out)

        pe[:, 0::2] = embedding_sin  
        # 将正弦编码分配到位置编码矩阵， pe 中所有行的第0列、第2列、第4列等（即所有奇数列）

        pe[:, 1::2] = embedding_cos 
        return pe

SinPositionEncoding(d_model=10, max_sequence_length=4, base=10000).forward()

exp_value实现了指数部分，变量out实现了三角函数内的值d-model表示每个token向量化后的维度。对于out部分的理解：

d_model=8, max_sequence_length=4, base=10000

1.生成位置索引：positions = torch.arange(4, dtype=torch.float)  # [0., 1., 2., 3.]。

2.将位置索引增加一个维度，使其成为列向量：positions = positions[:, None]  # [[0.], [1.], [2.], [3.]]

3.假设 alpha 已经根据之前的步骤计算得到，并且是一个形状为 (8 // 2,) 的张量，即 [4,]。为了简化，我们假设 alpha 如下：

alpha = 1 / (10000 ** torch.arange(4, dtype=torch.float))  
# 简化的例子，实际中 alpha 会根据 d_model 计算
alpha = alpha[None, :]  # [[1., 0.25, 0.0625, 0.015625]]

4.使用矩阵乘法 @ 计算 out：

out = positions @ alpha  # [[0., 0.25, 0.0625, 0.015625],
                         #  [1., 0.5 , 0.125 , 0.078125],
                         #  [2., 1.  , 0.25  , 0.125  ],
                         #  [3., 1.5 , 0.375 , 0.234375]]

5.然后我们将这些值分配到位置编码矩阵pe：

pe = torch.zeros(4, 8)
pe[:, 0::2] = embedding_sin
pe[:, 1::2] = embedding_cos

最终，pe将如下所示：

pe = tensor([[sin(0), cos(0), sin(0.25), cos(0.25), sin(0.0625), cos(0.0625), sin(0.15625), cos(0.15625)],
             [sin(1), cos(1), sin(0.5 ), cos(0.5 ), sin(0.125 ), cos(0.125 ), sin(0.078125), cos(0.078125)],
             [sin(2), cos(2), sin(1   ), cos(1   ), sin(0.25  ), cos(0.25  ), sin(0.125   ), cos(0.125)],
             [sin(3), cos(3), sin(1.5 ), cos(1.5 ), sin(0.375 ), cos(0.375 ), sin(0.234375), cos(0.234375)]])

二、可学习位置编码

#直接将位置编码当作可训练参数，让它随着训练过程更新，bert\gpt一些架构的实现方式
class TrainablePositionEncoding(nn.Module):
    def __init__(self, max_sequence_length, d_model):
        super().__init__()
        self.max_sequence_length = max_sequence_length
        self.d_model = d_model

    def forward(self):
        pe = nn.Embedding(self.max_sequence_length, self.d_model)
        nn.init.constant(pe.weight, 0.)#将 pe.weight 中的所有元素初始化为0
        return pe

举例说明，最大序列长度 max_sequence_length 为 10，模型维度 d_model 为 4，创建形状为 (10, 4)嵌入层pe饼使用0初始化权重。

三、相对位置编码

论文《Self-Attention with Relative Position Representations》，通常用于Transformer模型的变种中，例如Transformer-XL，它在处理长序列时能够提供更好的性能，因为它允许模型在计算注意力时考虑元素之间的相对位置关系。

class RelativePosition(nn.Module):
    def __init__(self, num_units, max_relative_position):
        super().__init__()
        self.num_units = num_units #嵌入维度，每个位置编码的维度
        self.max_relative_position = max_relative_position #考虑最大相对位置距离
        self.embeddings_table = nn.Parameter(torch.Tensor(max_relative_position * 2 + 1, num_units))
        #创建一个可学习的参数，用于存储相对位置编码。
        #size为 (max_relative_position * 2 + 1, num_units)

        nn.init.xavier_uniform_(self.embeddings_table)
        #使用Xavier均匀初始化方法初始化 embeddings_table

    def forward(self, length_q, length_k):
        range_vec_q = torch.arange(length_q)
        range_vec_k = torch.arange(length_k)
        distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
        # 计算查询序列和键序列之间的距离矩阵。

        distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
        #将距离矩阵中的值限制在 -self.max_relative_position 到 self.max_relative_position 之间。

        final_mat = distance_mat_clipped + self.max_relative_position
        #将限制后的矩阵的值平移，使得中心位置编码为0。

        final_mat = torch.LongTensor(final_mat).cuda()
        #将最终的矩阵转换为长整型张量，并移动到GPU上。
        embeddings = self.embeddings_table[final_mat].cuda()

        return embeddings

假设有一个查询序列（query）和键序列（key），它们的长度分别为 length_q = 2 和 length_k = 3。使用 RelativePosition 类来生成相对位置编码，其中 num_units = 2 表示每个位置编码的维度，max_relative_position = 1表示最大相对位距离。

#先初始化，在调用forward方法
length_q = 2
length_k = 3
num_units = 2 #表示每个位置编码的维度。
max_relative_position = 1 #表示我们关心的最大相对位置距离

1. 初始化位置索引：

   • range_vec_q = torch.arange(length_q) 生成查询序列的位置索引：[0, 1]
   • range_vec_k = torch.arange(length_k) 生成键序列的位置索引：[0, 1, 2]

2. 计算距离矩阵：

   • distance_mat = range_vec_k[None, :] - range_vec_q[:, None] 计算键序列中每个位置与查询序列中每个位置之间的距离，结果为：

     tensor([[0, 1, 2],
             [-1, 0, 1]])

3. 限制距离值：

   • distance_mat_clipped = torch.clamp(distance_mat, -1, 1) 将距离矩阵中的值限制在 -21到 1 之间，结果为：

     tensor([[0, 1, 1],
             [1, 0, 0]])

4. 平移距离值：

   • final_mat = distance_mat_clipped + 1 将限制后的矩阵的值平移，使得中心位置编码为0，结果为：

     tensor([[1, 2, 2],
             [2, 1, 1]])

5. 结果：

   • embeddings 将是一个形状为 (2, 3, 2) 的张量，其中第一个维度对应查询序列的长度，第二个维度对应键序列的长度，第三个维度对应位置编码的维度。

假设 embeddings_table 初始化如下：

embeddings_table = torch.Tensor([
    [0.1, 0.2],  # 索引0
    [0.3, 0.4],  # 索引1
    [0.5, 0.6],  # 索引2
])

• final_mat 的形状是 (2, 3)，其中：
  • 第一个维度（行）对应于查询序列（query）的位置。
  • 第二个维度（列）对应于键序列（key）的位置。
• final_mat 张量理解：
  • final_mat[0, 0] = 1 表示查询序列的第一个位置与键序列的第一个位置之间的相对位置编码索引是1。
  • final_mat[0, 1] = 2 表示查询序列的第一个位置与键序列的第二个位置之间的相对位置编码索引是2。
  • final_mat[0, 2] = 2 表示查询序列的第一个位置与键序列的第三个位置之间的相对位置编码索引也是2。
  • final_mat[1, 0] = 2 表示查询序列的第二个位置与键序列的第一个位置之间的相对位置编码索引是2。
  • final_mat[1, 1] = 1 表示查询序列的第二个位置与键序列的第二个位置之间的相对位置编码索引是1。
  • final_mat[1, 2] = 1 表示查询序列的第二个位置与键序列的第三个位置之间的相对位置编码索引也是1。

因此，使用 final_mat 作为索引，从 embeddings_table 中检索对应的编码，embedding张量将如下所示：

tensor([[[0.3, 0.4],
         [0.5, 0.6],
         [0.5, 0.6]],
        [[0.5, 0.6],
         [0.3, 0.4],
         [0.3, 0.4]]])

以下部分都未解析

class RelativeMultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1, batch_size=6):
        "Take in model size and number of heads."
        super(RelativeMultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.batch_size = batch_size

        assert d_model % n_heads == 0
        self.head_dim = d_model // n_heads

        self.linears = _get_clones(nn.Linear(d_model, d_model), 4)
        self.dropout = nn.Dropout(p=dropout)
        self.relative_position_k = RelativePosition(self.head_dim, max_relative_position=16)
        self.relative_position_v = RelativePosition(self.head_dim, max_relative_position=16)

        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).cuda()

    def forward(self, query, key, value):
        # embedding
        # query, key, value = [batch_size, len, hid_dim]
        query, key, value = [l(x).view(self.batch_size, -1, self.d_model) for l, x in
                             zip(self.linears, (query, key, value))]

        len_k = query.shape[1]
        len_q = query.shape[1]
        len_v = value.shape[1]

        # Self-Attention
        # r_q1, r_k1 = [batch_size, len, n_heads, head_dim]
        r_q1 = query.view(self.batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        r_k1 = key.view(self.batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        attn1 = torch.matmul(r_q1, r_k1.permute(0, 1, 3, 2))

        r_q2 = query.permute(1, 0, 2).contiguous().view(len_q, self.batch_size * self.n_heads, self.head_dim)
        r_k2 = self.relative_position_k(len_q, len_k)
        attn2 = torch.matmul(r_q2, r_k2.transpose(1, 2)).transpose(0, 1)
        attn2 = attn2.contiguous().view(self.batch_size, self.n_heads, len_q, len_k)
        attn = (attn1 + attn2) / self.scale

        attn = self.dropout(torch.softmax(attn, dim=-1))
        # attn = [batch_size, n_heads, len, len]
        r_v1 = value.view(self.batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        weight1 = torch.matmul(attn, r_v1)
        r_v2 = self.relative_position_v(len_q, len_v)
        weight2 = attn.permute(2, 0, 1, 3).contiguous().view(len_q, self.batch_size * self.n_heads, len_k)
        weight2 = torch.matmul(weight2, r_v2)
        weight2 = weight2.transpose(0, 1).contiguous().view(self.batch_size, self.n_heads, len_q, self.head_dim)

        x = weight1 + weight2
        # x = [batch size, n heads, query len, head dim]

        x = x.permute(0, 2, 1, 3).contiguous()
        # x = [batch size, query len, n heads, head dim]

        x = x.view(self.batch_size * len_q, self.d_model)
        # x = [batch size * query len, hid dim]

        return self.linears[-1](x)

 四、旋转位置编码

来源于论文《RoFormer: Enhanced Transformer with Rotary Position Embedding》，在在Llama及Llama2，QWen等模型中使用，Rope是将绝对位置编码与相对位置编码进行结合，通过绝对位置编码的方式实现相对位置编码。

# 生成旋转矩阵
def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
    # 计算词向量元素两两分组之后，每组元素对应的旋转角度\theta_i
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成 token 序列索引 t = [0, 1,..., seq_len-1]
    t = torch.arange(seq_len, device=freqs.device)
    # freqs.shape = [seq_len, dim // 2] 
    freqs = torch.outer(t, freqs).float()  # 计算m * \theta

    # 计算结果是个复数向量
    # 假设 freqs = [x, y]
    # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs) 
    return freqs_cis

# 旋转位置编码计算
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # xq.shape = [batch_size, seq_len, dim]
    # xq_.shape = [batch_size, seq_len, dim // 2, 2]
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)

    # 转为复数域
    xq_ = torch.view_as_complex(xq_)
    xk_ = torch.view_as_complex(xk_)

    # 应用旋转操作，然后将结果转回实数域
    # xq_out.shape = [batch_size, seq_len, dim]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)
    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.wq = Linear(...)
        self.wk = Linear(...)
        self.wv = Linear(...)

        self.freqs_cis = precompute_freqs_cis(dim, max_seq_len * 2)

    def forward(self, x: torch.Tensor):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(batch_size, seq_len, dim)
        xk = xk.view(batch_size, seq_len, dim)
        xv = xv.view(batch_size, seq_len, dim)

        # attention 操作之前，应用旋转位置编码
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        # scores.shape = (bs, seqlen, seqlen)
        scores = torch.matmul(xq, xk.transpose(1, 2)) / math.sqrt(dim)
        scores = F.softmax(scores.float(), dim=-1)
        output = torch.matmul(scores, xv)  # (batch_size, seq_len, dim)