【深度学习】位置编码

Transformer、DETR、Swin Transformer中的位置编码

beginner_GFL1207

1737人浏览 · 2024-04-28 18:26:19

beginner_GFL1207 · 2024-04-28 18:26:19 发布

一、引言

Self-Attention并行的计算方式未考虑输入特征间的位置关系，这对NLP来说是不可接受的，毕竟一个句子中每个单词都有着明显的顺序关系。Transformer没有RNN、LSTM那样的顺序结构，所以Transformer在提出Self-Attention的同时提出了Positional Encoding。

如图所示，Transformer在Attention模块之前将位置编码加进了待输入特征中。

二、绝对位置编码

首先需要明确NLP中数据的形式，一个批次包含多个句子，每个句子包含多个单词，每个单词被转为长度相同的token向量。由于每个句子中包含的单词数不同，所以会通过padding统一同一批次的embedding。假设一个批次padding后的embedding维度为 $batch\_size,num\_token,dim\_token]$ 。

绝对位置编码是一种将序列中每个token的位置进行编码的方法，它为每个位置分配一个唯一的编码向量。由于绝对位置编码只关注单个位置信息，因此它的实现通常在输入层，可以通过与输入向量相加融入模型。

2.1 直观的位置编码

2.1.1 索引型

将token的索引作为位置编码，第一个token编码为0，第二个token编码为1，以此类推。

其主要问题在于位置编码的值无界。

2.1.2 [0,1]型

为保证值有界，可限制位置编码范围为 $[0, 1]$ ，第一个token编码为0，最后一个token编码为1，其余token等间隔取值。例如共3个token时，位置编码为 $[0, 0.5, 1]$ ，共4个token时，位置编码为 $[0, 0.33, 0.66, 1]$ 。

其主要问题在于两个句子的token个数不同时，两个相同位置间的相对距离不同。共3个token时，第三个与第一个token间距为1，但共4个token时，第三个与第一个token间距为0.66。

2.1.3 二进制型

为保证值有界、句子长度不同时相对距离相同，可通过索引的二进制编码作为位置编码。下图为一个包含8个token，token向量长度为3的句子的位置编码。

如图，因为位置编码与embedding需要相加，所以通常位置编码的维度与embedding的维度相同。直白地说，有几个token就有几个位置编码，token向量的维度是多少位置编码向量的维度就是多少。

其主要问题在于二进制编码的位置向量处于离散空间，与输入相加后进入浮点世界，造成了空间上的浪费。

不过，我们可以观察该类型位置编码的规律。纵向来看，每个维度的编码值变化频率不同，蓝色变化周期为4，绿色变化周期为2，红色变化周期为1。

2.1.4 周期型

为保证值有界、句子长度不同时相对距离相同、节约空间，周期型位置编码包含了类似二进制型位置编码的变化规律，并将离散的二进制转为连续的 $\sin$ 或 $\cos$ 。

以 $\sin$ 为例，我们用 $p os$ 表示embedding中token的索引，用 $i$ 表示token上元素的索引。于是第 $p os$ 个token的位置编码可以表示如下：
$PE_{(pos)}=[\sin(\frac{pos}{2^0}),\sin(\frac{pos}{2^1}),\cdots,\sin(\frac{pos}{2^i}),\cdots,\sin(\frac{pos}{2^{dim\_token-1}})]$

其中， $pos=0,1,\cdots,num\_token-1$ ， $i=0,1,\cdots,dim\_token-1$ 。

可见，每个维度上 $\frac{1}{2^i}$ 被用来控制变化规律，详情如下图。

但是，使用 $\frac{1}{2^i}$ 来控制变化规律会使 $PE_{(pos)}$ 很快形成一个闭环。

如图，当 $dim\_token=3$ 时，以 $0.1$ 的间隔在 $[0, 20]$ 取 $p os$ ，得到200个 $PE_{(pos)}$ ，前100个点为蓝色，后100个点为橙色，可以清晰看到它们的重叠部分。这表明即便 $p os$ 不同， $PE_{(pos)}$ 也有很多点的值是相同的，但我们希望位置编码像地址一样是独一无二的，所以我们使用 $\frac{1}{10000^{i/dim\_token}}$ 替换 $\frac{1}{2^i}$ 来控制变化规律。此时， $PE_{(pos)}$ 如下图，不再有重叠。

于是有：
$PE_{(pos)}=[\sin(w_0pos),\sin(w_1pos),\cdots,\sin(w_ipos),\cdots,\sin(w_{dim\_token-1}pos)]$

其中， $w_i=\frac{1}{10000^{i/dim\_token}}$ 。

但它仍有一个问题，不同位置编码无法相互线性转换。

2.2 Sinusoidal位置编码

为保证值有界、句子长度不同时相对距离相同、节约空间、不同位置编码可相互线性转换，Sinusoidal型位置编码交替使用 $\sin$ 和 $\cos$ ，于是第 $p os$ 个token的位置编码可表示如下：
$PE_{(pos)}=[\sin(w_0pos),\cos(w_0pos),\cdots,\sin(w_ipos),\cos(w_ipos),\cdots,\sin(w_{\frac{dim\_token}{2}-1}pos),\cos(w_{\frac{dim\_token}{2}-1}pos)]$

其中， $pos=0,1,\cdots,num\_token-1$ ， $i=0,1,\cdots,\frac{dim\_token}{2}-1$ ， $w_i=\frac{1}{10000^{i/dim\_token}}$ 。

该形式下 $PE_{(pos)}$ 可以线性变换，可由下式证明：
$\begin{split} PE_{(pos+\Delta pos)} &= \left(\begin{array}{c} \sin(w_0(pos+\Delta pos))\\ \cos(w_0(pos+\Delta pos))\\ \cdots\\ \sin(w_{\frac{dim\_token}{2}-1}(pos+\Delta pos))\\ \cos(w_{\frac{dim\_token}{2}-1}(pos+\Delta pos)) \end{array}\right)\\ &= \left(\begin{array}{c} \left[\begin{array}{c} \cos(w_0\Delta pos)&\sin(w_0\Delta pos)\\ -\sin(w_0\Delta pos)&\cos(w_0\Delta pos) \end{array}\right]&\cdots&0\\ \cdots&\cdots&\cdots\\ 0&\cdots&\left[\begin{array}{c} \cos(w_{\frac{dim\_token}{2}-1}\Delta pos)&\sin(w_{\frac{dim\_token}{2}-1}\Delta pos)\\ -\sin(w_{\frac{dim\_token}{2}-1}\Delta pos)&\cos(w_{\frac{dim\_token}{2}-1}\Delta pos) \end{array}\right]\\ \end{array}\right) \left(\begin{array}{c} \sin(w_0pos)\\ \cos(w_0pos)\\ \cdots\\ \sin(w_{\frac{dim\_token}{2}-1}pos)\\ \cos(w_{\frac{dim\_token}{2}-1}pos) \end{array}\right)\\ &= T_{\Delta pos}*PE_{(pos)} \end{split}$

实际上，是用到如下和角公式中的第1项和第3项：

三、相对位置编码

相对位置编码是一种根据位置之间的相对关系来编码序列的方法。相对位置表达的是token两两之间的位置信息，因此它无法直接与输入相加（尺寸不一样），通常是通过为注意力权重增加偏执的方式来实现。

下图为原始注意力和加上相对位置偏执的注意力公式：

其中， $B$ 为相对位置偏执。可见，相对位置编码以偏执形式参与注意力权重的计算过程（没有也不能加在输入上），而绝对位置编码先加在输入上形成 $Q, K, V$ 后再进行注意力计算。

3.1 相对位置偏执

相对位置偏执是从一个可学习的表中查出来的，同一批次下不同token对的相对位置相同，则查出来的偏执也相同。想要获取相对位置偏执，需要确定取偏执时的索引和偏执表的长度及内容。

3.1.1 相对位置索引

取偏执时的索引使用的是相对位置索引，相对位置索引由绝对位置坐标推导得出。我们假设有如下特征图：

如图，特征图大小为 $N * M$ ，特征图的绝对位置坐标由(x,y)表示，即像素点对应的行列索引。在注意力计算时，每个像素点都会与其它像素点计算相似度，相对位置坐标与之类似，但计算的是相对位置。相对位置是以当前像素点绝对位置坐标为基准减去其它像素点的绝对位置坐标。以左下角的绿色点为例，绿色点为基准，绿色点与自身的相对位置坐标为 $(0, 0) = (1, 0) - (1, 0)$ ，绿色点与橙色点的相对位置坐标为 $(1, - 2) = (1, 0) - (0, 2)$ 。

将上述特征图相对位置坐标按行展平，拼接到一起形成如下相对位置坐标矩阵：

相对位置坐标矩阵有 $N * M$ 行、 $N * M$ 列。方便起见，我们将二维的相对位置坐标(x,y)转为一维的相对位置坐标，过程如下：

第一步行坐标加 $N - 1$ 、列坐标加 $M - 1$ 使所有坐标均大于0；第二步行坐标乘上 $2 M - 1$ 使(x,y)不同但x+y相同的坐标变得(x,y)不同x+y也不同，具体来说(0,2)和(2,0)两个相对坐标是不同的，但 $0 + 2 = 2 + 0$ ，此时一维相对坐标不唯一，在行坐标乘上 $5 = 2 * 3 - 1$ 后得到(0,2)和(10,0)，则(0,2) $\neq$ (10,0)且 $0+2\neq10+0$ ，一维相对坐标唯一；第三步行列坐标相加得到一维相对位置坐标。

我们希望每个相对位置坐标都能取到对应的相对位置偏执，这个一维的相对位置坐标即可作为相对位置索引，按照索引即可在相对位置偏执表中取出相对位置偏执（完全相同的相对位置坐标产生的相对位置索引相同，所以相对位置偏执也相同）。

3.1.2 偏执表

偏执表的长度由相对位置坐标的数量决定。将相对位置坐标矩阵中所有坐标点绘制在坐标系上，如下图所示：

可见，行坐标x的最大值为 $N - 1$ 、最小值为 $- (N - 1)$ ，所以共有 $2 N - 1 = (N - 1) - [- (N - 1)] + 1$ 个点；同样地，列坐标y的最大值为 $M - 1$ 、最小值为 $- (M - 1)$ ，所以共有 $2 M - 1 = (M - 1) - [- (M - 1)] + 1$ 个点。所以，偏执表的长度，即相对位置坐标的总个数为 $(2 N - 1) * (2 M - 1)$ 。

偏执表中的内容是学习得到的。通过relative_position_bias_table=nn.Parameter(torch.zeros((2N-1)*(2M-1)))初始化为0，接着通过nn.init.trunc_normal_(relative_position_bias_table, std=.02)初始化为正态分布，之后会随训练过程不断更新。

四、应用

4.1 Transformer中的位置编码

4.1.1 简介

上述Sinusoidal位置编码是在Transformer中针对NLP问题提出的。一个 $num\_token=50,dim\_token=128$ 的句子的位置编码如下图：

纵向来看，不同列的变化频率不同，从左到右频率依次下降。即使后58个维度无变化，为保证位置编码与embedding维度相同，仍然保留完整的128个维度。

位置编码在Transformer模型中的作用主要体现在以下几个方面：

(1) 捕捉词序信息：由于位置编码与词的位置相关，因此它们可以帮助模型理解输入序列中词的顺序。这对于依赖关系分析、句法分析等任务非常重要。
(2) 防止重复使用相同输入：由于位置编码是随机的，相同的输入序列会产生不同的位置编码。这有助于防止模型重复使用相同的输入来生成输出，从而提高模型的多样性和泛化能力。
(3) 增加模型的鲁棒性：位置编码的随机性可以帮助模型更好地处理噪声和异常值，从而提高其鲁棒性。

在实际应用中，位置编码通常在自注意力机制之前添加到输入序列中。这样，自注意力机制可以同时考虑词的语义信息和位置信息，从而更好地捕捉输入序列中的依赖关系。

4.1.2 实现

import torch.nn as nn
import torch


class PositionalEncoding(nn.Module):
    def __init__(self, dim_token, max_num_token=5000):
        super(PositionalEncoding, self).__init__()

        self.encoding = torch.zeros(max_num_token, dim_token)

        pos = torch.arange(0, max_num_token).unsqueeze(dim=1)  # 不是每次实时计算，而是预估一个pos上限
        _2i = torch.arange(0, dim_token, step=2)  # 共计算dim_token/2-1次，每次计算两个值sin和cos

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / dim_token)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / dim_token)))

    def forward(self, x):
        batch_size, num_token, dim_token = x.size()
        return self.encoding[:num_token, :]  # num_token是padding前单词的数量


if __name__ == '__main__':
    x = torch.randn((2, 3, 6))  # [batch_size,num_token,dim_token]
    pe = PositionalEncoding(6)  # dim_token必须是偶数
    y = pe(x)

4.2 DETR中的位置编码

4.2.1 简介

DETR将Transformer用在了CV的目标检测任务中，目标检测要求预测目标框，因此位置信息也很重要，所以也需要引入位置编码。DETR引入的位置编码也是Sinusoidal形式的。

不过，图像的维度与句子的维度不同。句子维度为 $batch\_size,num\_token,dim\_token]$ ，Transformer中位置编码与句子维度相同，一个位置编码向量表示一个句子（位置编码向量长度 $dim\_token$ ）。图像维度为 $batch\_size,num\_channel,height,width]$ ，DETR中一个位置编码向量表示一个像素（位置编码向量长度 $num\_channel$ ，一半的向量表示横坐标，另一半表示纵坐标）。此外，DETR还考虑了padding问题，仅针对非padding区域计算位置编码。

不仅如此，DETR中位置编码仅在Attention的 $Q$ 和 $K$ 中，而Transformer在 $Q$ 、 $K$ 、 $V$ 上都有。

4.2.2 实现

import torch.nn as nn
import torch


class PositionEmbeddingSine(nn.Module):
    def __init__(self, num_channel=64, temperature=10000):
        super().__init__()
        self.num_channel = num_channel
        self.temperature = temperature

    def forward(self, mask):
        assert mask is not None
        not_mask = ~mask  # mask中True表示padding区域，False表示非padding区域
        pos_y = not_mask.cumsum(1)  # 如果是padding区域，pos不增加
        pos_x = not_mask.cumsum(2)  # 横、纵坐标均计算pos

        i = torch.arange(self.num_channel)
        wi = self.temperature ** (2 * (i // 2) / self.num_channel)  # 2i = i // 2

        pos_x = pos_x[:, :, :, None] / wi  # 所有像素都有num_channel/2个横坐标
        pos_y = pos_y[:, :, :, None] / wi  # 所有像素都有num_channel/2个纵坐标

        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)  # 原本dim只到3，在第4个维度上stack然后flatten能使sin和cos交替出现
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)

        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)  # 前半部分为纵坐标，后半部分为横坐标
        return pos


if __name__ == '__main__':
    x = torch.randn((8, 4, 5, 6))  # [batch_size,num_channel,height,width]
    mask = torch.zeros((8, 5, 6))  # 同一图像上mask在每个通道上都一样
    mask = mask.bool()  # 默认没有padding
    pes = PositionEmbeddingSine(2)  # num_channel必须是偶数，这里输入的是num_channel/2，一半用于横坐标，另一半用于纵坐标
    y = pes(mask)

4.3 Swin Transformer中的位置编码

4.3.1 简介

Swin Transformer将Transformer用在了CV的各个任务中，也包括目标检测，所以同样需要引入位置编码。Swin Transformer引入的位置编码是相对位置编码形式的。

正如相对位置编码部分的描述，Swin Transformer的位置编码以偏执形式存在，参与注意力权重的计算。图像维度为 $batch\_size,num\_channel,height,width]$ 时，token的数量为 $h e i g h t * w i d t h$ ，token的长度为 $num\_channel$ ，所以其分别作为 $Q$ 和 $K$ 时 $QK^T$ 的维度应该是 $batch\_size,height*width,height*width]$ ，表示token两两之间的相似度。此时，相对位置坐标的索引矩阵维度是 $[h e i g h t * w i d t h, h e i g h t * w i d t h]$ ，根据索引取的相对位置偏执 $B$ 的维度也是 $[h e i g h t * w i d t h, h e i g h t * w i d t h]$ ，同一批次不同图片的相对位置偏执相同。

4.3.2 实现

import torch.nn as nn
import torch
import math


class RelativePositionBias(nn.Module):
    def __init__(self, N, M):
        super(RelativePositionBias, self).__init__()
        # 查询表要能包含所有可能的相对位置
        self.relative_position_bias_table = nn.Parameter(torch.zeros((2 * N - 1) * (2 * M - 1)))
		nn.init.trunc_normal_(self.relative_position_bias_table, std=.02)  # 初始化偏执表中的值
		
        # 获取绝对位置坐标矩阵
        coords_h = torch.arange(N)
        coords_w = torch.arange(M)
        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))  # [2, N, M]
        coords_flatten = torch.flatten(coords, 1)  # [2, N*M] 第一行为行坐标，第二行为列坐标

        # 根据绝对位置坐标导出相对位置坐标矩阵
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # [2, N*M, N*M]
        relative_coords = relative_coords.permute(1, 2, 0)  # [N*M, N*M, 2]

        # 将二维相对位置坐标矩阵转为一维索引矩阵
        relative_coords[:, :, 0] += N - 1
        relative_coords[:, :, 1] += M - 1
        relative_coords[:, :, 0] *= 2 * M - 1
        self.relative_position_index = relative_coords.sum(-1)

    def forward(self, attn):
        batch_size, NM, NM = attn.size()
        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(NM, NM)
        return relative_position_bias


if __name__ == '__main__':
    x = torch.randn((8, 6, 256))  # [batch_size,num_token,dim_token] num_token = height * width = N * M = 2 * 3
    q = k = v = x  # 没做线性映射的单头自注意力
    attn = torch.bmm(q, k.transpose(1, 2)) * math.sqrt(k.size(-1))  # (QK^T)/sqrt(dim_k)
    rpb = RelativePositionBias(2, 3)
    pos = rpb(attn)
    attn = attn + pos.unsqueeze(0)  # (QK^T)/sqrt(dim_k)+B 同批次不同图像相对位置编码相同
    attn = torch.softmax(attn, dim=-1)  # SoftMax((QK^T)/sqrt(dim_k)+B)
    attn = torch.bmm(attn, v)  # SoftMax((QK^T)/sqrt(dim_k)+B)V