正/余弦位置编码 Sinusoidal Encoding

2401_828890642026-02-17 11:30

1 公式

其中，pos 为词汇在句子中的位置索引；i 为特征维度索引，2i 代表偶数位置，2i + 1 代表奇数位置；d 为特征维度，在传统 Transformer 模型中默认为 512。

2 代码实现

$"你", "今天", "好", "漂亮", "啊"$ ，句子 token 长度为 5，若特征维度为 512，则位置编码矩阵的形状为 (5, 512)。

python 复制代码

# 正/余弦位置编码
class SinusoidalEncoding(nn.Module):
    def __init__(self,
                 d_model: int = 512,
                 max_len: int = 5000,
                 p: float = 0.1) -> None:
        """
        PE(pos, 2i) = sin(pos / 10000 ** (2i / d_model))
        PE(pos, 2i + 1) = cos(pos / 10000 ** (2i / d_model))

        Args:
            d_model: 特征维度，默认为 512
            max_len: 最大句子长度，默认为 5000
            p: 丢弃率，默认为 0.1
        """
        super(SinusoidalEncoding, self).__init__()
        # 初始化位置向量
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)
        # 1 / 10000 ** (2 * k / d)
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000) / d_model))
        # 偶数位置
        pe[:, 0::2] = torch.sin(position * div_term)
        # 奇数位置
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        # 缓存
        self.register_buffer("pe", pe)
        # 丢弃层
        self.dropout = nn.Dropout(p)

    def forward(self, x: Tensor) -> Tensor:
        # 截取与词向量相同形状的位置向量，并相加，但位置向量不参与学习，(N, S, E) -> (N, S, E)
        x += Variable(self.pe[:, :x.size(1)], requires_grad=False)

        return self.dropout(x)