作者：HeptaAI
链接：https://www.zhihu.com/question/488561011/answer/3131570354
来源：知乎
 

最近回顾了一下ViT的结构，发现从文本Transformer出发解释会好理解很多。ViT其实就是图像版的BERT，除了一开始从输入到embedding与BERT不同之外，其他的技巧基本都是相同的。例如，ViT的[CLS] token，后期的Encoder block，MLP head，与BERT都是完全一致的：

动图版本：

从图中可以看出，两者的主要差别就是Embedding层。因此在这篇文章中，我们将从一个NLPer的视角出发，看看这份工作是如何调整输入数据的处理方法，从而将BERT迁移到CV领域的。文中代码来自博客transformer is all you need!。

Semantic Embedding

文本Embedding

有关文本的Embedding，我写过一篇详细的分析：Tokenizer与Embedding串讲。对于文本，Embedding矩阵的本质就是一个查找表。由于输入向量是one-hot的，embedding矩阵中有且仅有一行被激活。行间互不干扰。这是什么意思呢？如下图所示，假设词汇表一共有6个词，则one-hot表示的长度为6。现在我们有三个单词组成一个句子，则输入矩阵的形状为 (3,6)(3, 6) 。然后我们学出来一个embedding矩阵，根据上面的推导，如果我们的embedding size为4，则embedding矩阵的形状应该为 (6,4)(6, 4) 。这样乘出来的输出矩阵的形状应为 (3,4)(3, 4) 。

我在图中用不同颜色标明了三个subword embedding分别的计算过程。对于第一个单词'I'，假设其one-hot编码为 [0,0,1,0,0,0][0, 0, 1, 0, 0, 0] ，将其与embedding矩阵相乘，相当于取出embedding矩阵的第3行（index为2）。同理，对于单词'love'，相当于取出embedding矩阵的第二行（index为1）。这样一来大家就理解了，embedding矩阵的本质是一个查找表，每个单词会定位这个表中的某一行，而这一行就是这个单词学习到的在嵌入空间的语义。

图像Embedding

图像的输入不是一个一个字符，而是一个一个像素。设每个像素有C个通道，图片有宽W和高H，因此一张图片的所有数据可以用一张大小为 H×W×CH\times W\times C 的张量来无损地表示。例如，CIFAR-10数据集里面，数据的大小就是 224×224×3224\times224\times3 。但是一个一个像素输入transformer粒度太细了，一张最小的图片也要 224⋅224224\cdot 224 个token，所以一般把图片切成一些小块（patch）当作token输入。因此，patch的大小 Ph×PwP_h\times P_w 必须是能够被图片的宽和高整除的。例如对于CIFAR-10，一般的设定是大小为 16×16×316\times16\times3 的patch，这样就能获得 14×1414\times 14 个patch。注意，patch只是几何上的切分，通道数 CC 与原图是保持不变的。

由于文本一直以来用的Transformer都只支持一维输入，所以很容易想到可以直接将patch拍平成一维，变成 196196 张patch。由于patch的分割和拍平都会让图像失去部分几何联系，这 14⋅14=19614\cdot14=196 张大小为 16×16×316\times16\times3 的patch就构成一串有损表示图像的tokens。这些图像的token意义上等价于文本的token，都是原来信息的序列表示。不同的是，文本的token是通过分词算法分到的subword，这些subword会被映射到字典的index；也就是说，文本的token是一个数字。而图像的一个token（patch）是一个 16×16×316\times16\times3 的矩阵。那么现在问题来了，如果文本的embedding矩阵可以用查找表来建模，那应该怎样构建图像的embedding矩阵呢？

其实想法非常简单，文本的每个token都等价于一个one-hot向量，因此embedding的转换是一个 V→DV\rightarrow D 的过程，其中 DD 是embedding的长度。那么由于transformer接受的输入只能是 DD ，所以我们只需要一个线性的图像转换过程 Ph×Pw×C→DP_h\times P_w\times C\rightarrow D 。也就是说，对于每个形状为 Ph×Pw×CP_h\times P_w\times C 的三维token（patch）张量，我们寻找一个线性变换 L\mathcal{L} ，使得token变成一个长度为 DD 的一维向量。

这显然是一个降维操作，直接用矩阵是做不了的，因此显然不是查找表。那可以直接展开吗？例如CIFAR-10中，每个patch是 16×16×316\times16\times3 的张量，那我可以直接拍平成一条 16⋅16⋅3=76816\cdot16\cdot3=768 的一维向量吗？这不符合embedding的本质：在文本embedding中，这个转换过程（embedding矩阵）是学出来的，也就是说embedding的本质是一种结合了整个句子意义的、对某个token的隐层表示，而非单纯来自输入图像的一种显式的表示。因此，我们更进一步，我们需要的线性变换 L\mathcal{L} 需要带着参数，我们用 Lθ\mathcal{L}_\theta 来表示。

到此为止，熟悉图像算法的同学应该已经猜到答案了。最naive的能用来对多通道数据进行信息的降维融合的算子，一定是卷积。由于图像是二维的，很容易想到Conv2D。具体怎么卷积呢？假设embedding的维度 D=768D=768 ，对于每一个patch，我们都做如下的Conv2D卷积：

这个最后得到的长度为 768768 的一维向量，就是我们要的这个patch的embedding。然后我们将这个过程对每一个patch都重复一遍，就能得到 196196 个长度为 768768 的一维embedding，相当于一个长度为 196196 的句子。将这些embedding向量拼成一个embedding矩阵，形状就是 196⋅768196\cdot 768 。现在这个矩阵就可以扔到transformer里面做运算了。

如果你仔细思考，可以发现，这个过程其实是可以进一步并行的，因为各个patch的embedding计算并不相互影响。因此，只要设置stride为patch的长宽 1616 ，就可以用Conv2D一步到位：

这里我们只给出一步到位的代码，因为比一个一个算的优雅多了：

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x

Positional Embedding

文本Embedding

在之前的讲Transformer的文章Self-Attention & Transformer完全指南中，我们已经分析了原版Transformer是如何编码token的位置信息的。简单地说，原版Transformer根据token在句子中出现的位置，为token赋予一个个位置值，然后加到semantic embedding上即可。但是，这种方法的encoding是handcraft的。根据各位炼丹师的经验，这往往也可以被参数化，通过模型自动学习来获得更加准确的encoding。于是，在后期的Transformer实践中（例如BERT），作者就直接将positional embedding也直接参数化，用学习的方法来获得。

具体的实现方法和semantic embedding是完全一样的，也是查找表。假设BERT最长支持的输入序列长度为 NN ，embedding维度为 DD ，则我们可以创造一个形状为 N×DN\times D 的embedding矩阵。这样，由于每个token都有一个index，就可以直接去embedding矩阵里面找对应index的那个embedding向量。

图像Embedding

ViT采用了和BERT一致的方法，具体实现如下：

# 这里多1是为了后面要说的class token，embed_dim即patch embed_dim
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) 

# patch emded + pos_embed
x = x + self.pos_embed

但是要注意一个问题，文本transformer规定了最大的输入长度 NN ，但是图像一般不会给出最大图像规模，而是一个推荐的图像规模，例如 ≈224⋅224\approx 224\cdot 224 。因此，图像的positional embedding一般都是给出一个推荐的patch数 N0N_0 （就是代码里的num_patches），例如 14⋅1414\cdot14 。而由于patch的大小是确定的，所以patch的数量 NN 是不确定的。这种情况下，就不能直接通过形状为 N0×DN_0\times D 的embedding矩阵来解决。当 N≶N0N\lessgtr N_0 的时候，需要使用插值法来将positional embedding的个数调整到 N0′=NN_0'=N 。具体代码一般不会用到，在此略过。