0. 前言

近两年学术界对Transformer在CV上的应用可谓异常青睐，这里重点强调学术界的原因是目前工业界还是比较冷静的（部分公司已经开始考虑Vision Trransformer的落地了），毕竟新方法从学术界到工业界落地一般都会晚几年。在当前背景下，个人对Vision Transformer进行了较浅的调研，本文便是该工作的阶段性总结。本文由三部分组成：

• 第一部分结合代码介绍了Transformer的基本概念和原理，由于我们主要关注的是Transformer在CV上的应用，所以不会涉及过多关于NLP相关的细节。
• 第二部分会介绍几个自己验证过效果较好的Vision Transformer方法。
• 第三部分是一些个人对Vision Transformer的看法，部分观点已有实验验证，但因为涉及工作内容，所以不便贴出来，另外一些观点只是个人的假设，并没有经过验证。

由于时间仓促以及本人能力有限，无法完全保证文中没有错误之处，欢迎指正。

1. Transformer

本部分我们以翻译任务为例介绍Transformer。假设我们有一个法语句子需要翻译为英语，即将"Je suis étudiant"翻译为"I am a student"（为了配合使用别人的图片:))，所以使用这个例子）。如果不关注细节的话，我们可以将翻译模型视为黑箱，翻译过程就如下图所示。

1.1 Transformer之前的Seq2Seq模型

现在我们来探究在Transformer出现之前上图中黑箱的大致细节和原理，从而引入一些常用的概念（用黑体标出）。

一个语言句子其实就是一些词（word）的序列（Sequence），而翻译的过程其实就是将源语言的词序列（ ）翻译为目标语言的词序列（ ），即Sequence to Sequence，简称Seq2Seq。Transformer之前主流Seq2Seq神经网络模型是RNN以及一些变种，主要有LSTM、GRU（GRU是LSTM的变种）等。下面简单了解下RNN和LSTM。

我们人类在翻译一个句子的时候并不是看到一个词就将其立即翻译成目标语言对应的一个单词，并且忘掉之前读过的单词，而是通读一遍句子并理解句子的含义，然后将含义再翻译为目标语言。用计算机语言描述就是先将输入的源词序列编码（Encoding）为一个表示含义的张量（tensor，一组数字），其实主要是对每个词进行编码，然后再将编码的张量解码（Decoding）为目标语言的词序列作为输出。如下图所示，整个模型由两部分组成，编码器（Encoder）和解码器（Decoder）。

上图只是直观的展示，其实在训练和推理（翻译）时是无法直接将原始的语言句子输入给模型进行计算的，因为计算机本质上只支持数字运算，所以需要将序列中的每个词处理成数字才行，常用的方法是word2vec，即将每个单词用唯一的一个张量进行表示，输入序列就是很多个张量的列表，模型的输出也是一个元素为张量的列表，其中每个张量表示目标语言的一个词，下图为源语言句子用word2vec处理后的张量()。

但是，早期的翻译模型仅仅能够做到简单地将源序列的一个或多个单词映射为目标语言的一个或多个单词，句子如果简单的话，这种方法还是比较有效的，但是对于一些大长句就比较吃力了，尤其是无法解决语言翻译中的一个典型问题，指代消除。指代消除是指将句子中表示同一个对象的词关联起来。下面是一个需要进行指代消除的句子:

The animal didn’t cross the street because it was too tired.

句子中的it是一个代名词，具体指的是什么？我们需要结合句子中其他的词才能知道。将it替代为animal便是指代消除。为了更好地编码序列中某个位置的词而需要关注序列中其他位置的词，即考虑上下文信息。RNN（Recurrent Neural Networks）主要就是针对像语言这种长度可变且序列中存在依赖关系的任务设计的。下图左侧是简略的RNN网络结构，结构中具有的循环机制（图中指向A自己的箭头，这也是名称中R的来源）使得模型在编码后续词的时候可以关注之前词的编码信息，好像具有了记忆功能一样。训练和推理的时候都是将序列中的单词逐个输入到模型中进行编码和解码，序列中所有词都是共享同一模型参数。下图右侧是按时间序列展开的结构，为序列中第t个时间步的输入，也就是序列中第t个位置的词张量, A表示模型（可学习参数），表示第t个输出，即目标语言序列中第t个位置的词张量。

下图是标准RNN模型的大致内部结构，这里不进行细致探讨：

RNN存在一个严重问题是无法捕获序列中相隔较远的词之间的依赖关系，即所谓的长程依赖问题。另外，输入序列中的不同词一般具有不同的重要程度，而且输出序列中不同位置的词有时需要考虑输入中的不同位置的词才能更好地进行解码。例如，在翻译任务中，输出的第一个单词一般是基于输入的前几个词确定的，输出的最后几个词可能基于输入的最后几个词。注意力机制（下一节会详细说明）的引入为解码器提供了在每个解码时间步上查看整个输入序列的能力，而且解码器可以在任何时间步决定哪些输入单词是重要的。LSTM（Long Short Term Memory）便是为了解决长程依赖问题而提出的，其后续改进版引入了注意力机制，使得RNN模型的表现不断提升，下图是原始LSTM的大致结构，其中添加了记忆筛选机制，使得模型在后续编解码中有选择地记住重要词的编码信息，并忘掉不重要词的编码信息，这也是其名称的由来。

下图是增加了注意力机制的LSTM模型的改进版本，可以看到预测（图中的“hit”）的解码时不仅依赖于解码器前一个时间步的输出（图中的"<START>"和"he"）和编码器的编码信息，还使用到了由输入序列中每个词的编码信息计算得到的注意力输出（图中的"Attention output"）。"<START>"词为序列开始解码的指示符，因为Decoder的第一次预测时没有来自前一个时间步的输出词作为输入，便用特殊的词代替，类似于占位符。

虽然LSTM已经很优秀了，但是还是有明显的缺陷：无法进行并行训练。由于语言数据本身是不定长的，RNN恰恰就是设计来处理不定长数据任务的，训练时都是单个序列进行训练。有些工作也尝试解决这个问题，但是仍然无法从根本上解决。这个缺陷大大阻碍了将现代GPU用于加速训练大规模语言模型，而这便是Transformer要解决的问题。

1.2 Transformer

如第1.1部分所述，Transformer的提出主要是为了解决LSTM无法进行并行训练的问题，但是仍然沿用Seq2Seq的Encoder-Decoder结构，主要的改进有两点：

1. Encoder和Decoder模块都是由多个Attention（论文中称为Multi-Head Attention）和MLP（论文中称为Feed-Forward Networks）子模块堆叠组成（暂时忽略LayerNorm层）;
2. 类似CNN网络，输入可以一次接受包含多个序列的数据，大大加速了模型的训练。

下面是Transformer的整体网络结构，我们将依次对各个组件内部进行详细探究。

Encoder的输入(Input Embedding)

与LSTM类似，我们需要将输入（语言）序列进行数值化，由于Transformer支持一次输入多个序列，而每个序列的长度不一，为了能够将不同长度的序列组成一个batch（为了进行并行计算），我们需要将所有序列的长度进行对齐，简单的做法便是设置最大序列长度，该最大长度为训练数据集中最长序列的长度加1（加1用于放置序列结束符），其他不足最大长度的序列，多余的部分全部填充为序列结束符，例如符号“<EOS>”表示End of Sentence(实际中可能用另外的符号，如"<blank>")。假设一个数据集中只有两个句子分别为：

Do we really need Attention?
Yes or no?

数据集中最大句子长度为6（包括标点符号），则我们设置最大序列长度为7。所有句子长度对齐后如下。具体实现上，训练时只需要将一个batch内的长度对齐便可,补充的多余部分可以在训练时用mask屏蔽掉，从而不参与训练，mask方式类似后面要介绍的Masked Multi-Head Attention部分。

下一步我们将序列中的每个词进行数值化（又称词嵌入, word embedding），假设整个数据集中最多有10个词（包括标点符号），则我们可以用one-hot方式将每个词表示成唯一的张量，张量的维度为10，例如用“0000000001”表示“do”，用“0000000010”表示“we”，其他依次类推，则对齐的两个句子数值化后如下（为了方便展示，对序列矩阵进行了转置），如果batch size=2，则例子中处理后的输入矩阵大小为[2, 7, 10]，这便是上图中的Input Embedding。当然这个是最简单的策略，存在的主要问题是Embedding的维度随着语料库的词汇数量增加，更常用的是 Word2Vec 方法，这里不再展开。

Encoder和Decoder

Encoder

我们先从整体上了解下Encoder的结构，如下图，从图中可以看出Encoder由N个相同的层堆叠组成（论文中N=6），每个层又由两个子层（sub-layer）组成，第一个子层称为Multi-head Attention，另一个子层称为Feed Forward。每个子层都引入了残差连接，子层的输出与其输入逐元素相加（element-wise addition），再进行LayerNorm（类似于BatchNorm的层归一化），作为本层的输出。Encoder中第一个层(Multi-head Attention+Feed Forward)的输入来自模型的输入，即经过位置编码（图中的Position Encoding）的词嵌入矩阵，其他后续层的输入都来自前一层的输出。

下面我们一一研究Encoder中的每个子层。在研究Multi-Head Attention子层之前，我们需要先弄明白什么是self-attention。

自注意力（Self-Attention）

所谓“自”注意力，个人理解就是指在没有人为先验经验指导的情况下让模型自己学习掌握一种能力，这种能力能够建模所需的注意力。至于注意力机制，论文中给出了非常简洁的定义：为了更好地编码一个序列的表征（representions），将序列中的不同位置关联起来的一种机制，类似于我们在辨别照片中一个不清晰物体时会借助周围其他物体。

常用的关联不同位置信息的方式便是加权求和。具体地，就是先对序列中每个位置进行独立（个人使用的非正式说法）编码，然后将不同位置的编码进行加权求和，得到包含注意力能力的编码信息，这里的权重便是模型需要学习的参数（self的来源）。这种机制其实在CV中也已经有比较成熟的应用，比如2018年提出SENet中使用Squeeze-and-Excitation便是一种通道注意力机制，可以认为一个通道特征就是序列中的一个位置的编码，如下图所示，模型通过学习自注意力函数 ，该函数接受压缩后的通道编码信息，产出每个通道的权重值，权重数量与原始特征通道数量相同，这样便可以将每个权重与原始通道进行相乘，这里相乘后并没有多通道相加，但可以认为是单通上的加权。

我们可以用简单的公式表示这种加权的注意力机制：

其中，表示得是第i个位置词包含注意力的编码信息，是模型学习到的权重，每个位置对应一个，表示所有位置的独立编码信息，表示第j个位置的独立编码信息，为编码信息的长度，为序列的长度，即序列中词的数量。如果要使用矩阵一次计算所有位置的注意力编码信息，则可以使用下面的公式：

其中，，。

在Transformer中这种自注意力机制得到了进一步的发展，但主要的思想是不变的，具体地，要进行加权的各个位置的编码信息矩阵是不变的，变化的是得到权重的方式。如前所述，我们将其他位置的编码信息以加权的方式加入到当前位置的编码中的主要目的是序列中多个位置的词之间存在一定的相关性，这种关系有助于更好地对当前位置的词进行编码（比如解决长程依赖问题）。因此，可以让模型尝试学习对这种相关性进行编码（这是一种1-N的关系，包括当前词自己与自己的相关性），然后将这种相关性直接转换为加权值，相关性越高权值越大，对词最终的编码的贡献越大。

Transformer使用类似查字典的方式来显式地对这种相关性进行建模。在实际使用字典的时候，我们首先得有个查询的字（这里称为query，简称为q）,然后使用一定的标准逐一与字典中的关键字（这里称为key，简称为k）进行比较，符合比较标准（比如拼音或笔画相同）的我们就认为找到了要查找的内容，或者说两个字匹配上了。在Transformer中，也是显示地为序列中的每个词指定一个查询的字，并且为序列中每个词指定一个关键字，它们都是与独立编码信息（这里称为value）类似的张量。到这里，有两个问题需要解决，一是如何得到具体的、和张量，另一个是如何计算和的相似性。

第一个问题，由于我们希望模型自己学习到序列中不同的相似性，最直接的方法就是通过不断地迭代训练，用监督的方式让模型自己学习每个词（或称为token，对于更深的层来说，不能再称为词）对应的，和，这其实类似于CNN中让模型自己学习特征，而不是人工设计特征。简单的做法就是各使用一个可学习的矩阵将序列的每个词线性映射为，和，假设这三个矩阵分别为，和，输入为（行向量，表示一个词嵌入），则每个词对应的，和可通过如下计算得到：

过程如下图示例，图中假设输入的维度是4，,，，：

如果使用矩阵一次计算序列中所有位置的、和，则可以使用下面的公式：

其中假设输入序列矩阵，每个词张量为维，一共N个词，,，，、和是模型超参，后续我们还会提到他们，计算的,, 。

过程如下图示例：

Tranformer论文中作者将分别计算的输入直接用表示，如下公式，但其实是同一输入的副本。

对于第二个问题，相对比较简单了，两个张量相似的话，可以用它们的欧式距离衡量，距离越小越相关，或者先计算两个张量的点积，然后使用sigmoid函数映射到区间，越相关则值越接近1。如前所述，每个都要与序列中所有的计算相似性，而且序列中的这种相似性可能不是独立的（互斥性），那么可以让与每个先计算点积，然后在产生的点积张量（一个点积一个标量，多次点积就是一个多维张量）上执行softmax函数，点积张量的每个值都会映射到区间，即可得到第i个词与序列中其他词之间的相关性权重。计算的方式如下：

每个张量独立计算点积：

与所有的key的点积计算如下：

使用矩阵一次计算序列中所有词之间的点积：

表示第个词的query 与第j个词的key 的点积。

权重计算如下：

其中

过程如下图示例，该过程展示了计算与和的权重，结果分别为0.88和0.12：

使用矩阵一次计算序列中所有词的权重张量：

该权重矩阵便可以用于前面公式1中计算注意力编码了：

可以看到，这个公式与Transformer论文中的注意力计算公式已经很接近了：

过程如下图示例：

作者增加缩放因子的原因，是当和元素的维度太大时会导致softmax函数处于梯度极小的区域，这样会导致模型的最终效果较差，所以增加一个缩放因子。作者也将这样的注意力机制称为带缩放的点积注意力(Scaled Dot-Product Attention)，模型的示意图如下（后面会介绍Mask的作用）：

Pytorch代码实现如下：

import torch

class ScaledDotProductAttention(nn.Module):
    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature
        self.dropout = torch.nn.Dropout(attn_dropout)

    def forward(self, q, k, v, mask=None):
        # self.temperature就是缩放因子
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)
        # dropout是正则化方法
        attn = self.dropout(torch.nn.functional.softmax(attn, dim=-1))
        output = torch.matmul(attn, v)
        return output, attn

Multi-Head Attention

Transformer作者发现，相比使用单一的注意力函数(即公式2)，对输入序列并行执行多次线性映射（即计算不同的）并使用多个注意力函数，模型效果会更好。个人猜测原因是一个序列中一个词可能与多个词的相关性同样重要，但是使用了softmax函数就是显式地约束相关性是互斥的。通过并行计算多个注意力函数，可以大大缓解这个问题。但是，每增加一个注意力函数，就是显著增加模型的参数量和计算量。作者做法是将原来的张量维度分为h份，每一部分在一个独立的注意力函数中进行计算，这样的话，参数量和计算量并未增加，但将单注意力（Single—Head Attention）变成了多注意力（Multi-Head Attention）。例如，在使用Multi-Head Attention之前的维度为，使用后每个Attention中的维度为，总体上的张量维度是不变的，还是维。每个独立的Attention分别计算出一部分注意力编码后再将它们拼接起来（Concat），然后使用一个矩阵进行线性变换（融合），具体公式如下：

其中，,，，计算的,, 。论文中，, ，整体的示意图如下。

代码实现如下：

import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from transformer.Modules import ScaledDotProductAttention

class MultiHeadAttention(nn.Module):
    # n_head就是上述的h超参
    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()
        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        # 为了实现高效，不同头的Q、K和V的计算矩阵可以是同一个，这其实应用了矩阵分块计算
        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

    def forward(self, q, k, v, mask=None):
        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        # len_q、len_k、len_v分别是序列的长度，即词的数量
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = 
        # 同时计算所有head中的q,k,v，然后再分拆，便于在多个head中分别计算注意力
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1) 

        # 实现上看起来是Single-Head, 但是q,k,v的shape比Single-Head的多一维
        # 使用了矩阵运算的高效实现
        q, attn = self.attention(q, k, v, mask=mask)
        
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        # 对多个head的注意力编码进行融合
        q = self.dropout(self.fc(q))
        q += residual

        q = self.layer_norm(q)

        return q, attn

Position-wise Feed-Forward Networks

Multi-Head Attention子层输出的编码张量与残差连接的原始张量相加并进行LayerNorm后，作为Position-wise Feed-Forward Networks的输入，一个由二层全连接网络和位于两者之间的ReLU激活函数组成的子层，加“Position-wise”的原因是每个全连接线性变换是应用到序列中的每个词的，即所有词共用一组权重参数，公式表达如下：

表示序列中某个位置“词”的注意力编码，，，。论文中，。

和Multi-Head Attention子层一样，两个全连接网络产生的编码张量与残差连接的原始张量相加，并进行LayerNorm，代码实现如下：

import torch.nn as nn

class PositionwiseFeedForward(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid) 
        self.w_2 = nn.Linear(d_hid, d_in)
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x

        x = self.w_2(F.relu(self.w_1(x)))
        x = self.dropout(x)
        x += residual
        x = self.layer_norm(x)

        return x

LayerNorm

CV中不同Norm的计算方式比较：

为什么Transformer使用LayerNorm，而不是BatchNorm?

Layer normalization is used in the transformer because the statistics of language data exhibit large fluctuations across the batch dimension, and this leads to instability in batch normalization.

image.png

具体实现上，NLP和CV中LayerNorm的计算方式不一样，NLP中的计算方式如下图左侧所示，其只在在每个token的embedding维度上计算均值和方式，然后进行归一化：

image.png

Pytorch中的示例代码：

>>> # NLP Example
>>> batch, sentence_length, embedding_dim = 20, 5, 10
>>> embedding = torch.randn(batch, sentence_length, embedding_dim)
>>> layer_norm = nn.LayerNorm(embedding_dim)
>>> # Activate module
>>> layer_norm(embedding)
>>>
>>> # Image Example
>>> N, C, H, W = 20, 5, 10, 10
>>> input = torch.randn(N, C, H, W)
>>> # Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
>>> # as shown in the image below
>>> layer_norm = nn.LayerNorm([C, H, W])
>>> output = layer_norm(input)

位置编码(Positional Encoding)

从前面的注意力机制我们可以看到，Transformer最大的好处就是让序列中每个词的编码是同时进行的，且对序列中每个词进行编码时都会注意到序列中的其他词，有利于对每个词进行更好地编码。但是，序列中词的相对位置关系也是特别重要的，在LSTM模型中序列的每个词是以先后顺序输入到模型中的，所以天然地具有利用词之间的位置信息。Transformer的解决方法是显式地将位置信息直接附加到词的embedding张量中。所谓直接附加，就是将每个词的位置信息直接编码成与每个词相同维度（即）的张量，与词的embedding相加，如下图所示。

这里还有个问题，那就是每个词的位置编码怎么得到。Transformer的作者指出直接编码位置信息或者让模型学习每个词的位置信息都是可行的，且效果相当，但是作者用了sine和cosine函数来对每个词的位置信息进行编码，公式如下：

其中pos表示词在序列中的位置，i表示词张量中的位置。代码的实现如下：

import torch
import torch.nn as nn
import numpy as np

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()

        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)

    def forward(self, x):
        return x + self.pos_table[:, :x.size(1)].clone().detach()

这样的位置编码方式，使模型很容易学习到序列中每个词的相对位置，因为对于任何一个固定的位置偏移k，可以用位置的编码信息进行线性表示，如下公式所示，而且这样的编码信息可以使模型的编码位置扩展到训练时未见过的长度。

以上是Encoder模块中一个层的组成，较详细的组成如下图，为了便于理解，图中将一个序列的和分开输入给Encoder，串联重复N次就组成了Encoder，论文中N=6。总体上，Encoder的编码输出维度为 ，B为Batch size。

其实，如果只是想要了解Transformer的Attention机制，并应用到CV中，那么了解整个Encoder就可以了，因为绝大部分的CV任务（除了类似Image Captrain这样的任务）从本质上来说并不需要用Encoder-Decoder网络结构进行建模，只需要对输入进行encoding就行。所以，目前几乎所有的（个人的片面经验）ViT模型都只利用了Transformer的Encoder部分的思想。

Decoder

Decoder也是由个（论文中是6）相同的层串联组成，每个层除了由分别带有残差连接的Multi-Head Attention和Feed Forward子层组成外（与Encoder一样），在两个子层之间还加入了一个Multi-Head Attention子层，该子层也带有残差连接和LayerNorm。其主要作用是融合Encoder的输出和Docoder的中间编码信息，便于更好地进行解码。其中的Feed Forward子层与Encoder中的完全一样。整体的结构如下图所示。

Masked Multi-Head Attention

细心的读者肯定已经发现，在前面的Multi-Head Attention源码的forward函数有个mask参数，这其实就是为Masked Multi-Head Attention实现的。所以，Masked Multi-Head Attention与Multi-Head Attention几乎是一样的，只是多了一个mask操作。

总体上来说，这mask是为了屏蔽掉训练时位于序列前面的词与序列中位于该词后面的其他词计算注意力，并将它们的编码信息用于计算该词的编码。屏蔽的原因主要是Decoder在训练时与推理时的使用方式不同造成的，在推理的时候，Decoder与RNN模型一样，每个时间步只能预测一个单词，不断循环直到模型预测序列结束符或者达到预定义的最大长度，即只能根据前面所有已经预测的单词预测当前时间步的单词。例如，将"Je suis étudiant"翻译为"I am a student"，在推理时，要先给解码器输入开始解码标志符，这是一个特殊的单词，例如“<START>”，如果没有这个预先输入的单词，Decoder就无法对句子的第一个单词“I”进行解码，因为在序列上它是第一个，没有可以依赖的之前时间步的输入。当输入“<START>”的embedding张量（包含位置编码信息），Decoder结合embedding和Encoder的编码信息预测输出“I”，然后再将“<START>”和“I”的embedding作为Decoder的输入，结合Encoder的编码信息预测输出“am”,以此类推，不断循环，直到模型预测出句子结束符“<EOS>”或者达到提前设定的最大预测序列长度。下面的动图展示了推理的过程，但是其中没有从“<START>”作为输出开始解码。

在训练的时候，我们是可以一次看到要预测序列中的所有词的，如果还像推理时那样，那么整个模型并没有完全实现真正的多个序列并行训练，Decoder将是整个模型的训练效率瓶颈。所以，当然是希望训练时像Encoder一样，一次将多个序列同时作为Decoder的输入，直接预测后续的单词。但是，如果使用不带mask的Multi-Head Attention，那么模型就可以很开心地作弊了，因为已经全部知道要预测的答案了，为什么还要费力气学习解码呢，直接将输入作为输出不要太简单，而且训练Loss直接就是为0。解决这个问题的方式就是在训练的时，将序列中后续的词的编码信息遮盖住，不让其参与到前面词的编码中，这便有效阻止了模型作弊，而且可以提升并行训练。这里说的遮盖其实很容易在Multi-Head Attention基础上实现的，在Multi-Head Attention中会计算一个词与序列中每个词的权重，这样会得到一个权重矩阵，矩阵的每一行表示改行对应的单词与序列其他词的相关性，那么在计算这个权重前，我们只需要将该词位置之后的点积值重置为一个特别小的值（例如），在进softmax计算后，那些位于该词位置之后的词的权重都特别小，即与重置的权重对应的词的编码信息就几乎不会参与到当前词的解码信息中。例如下面是对“"I am a student”解码时，在mask之前计算的点积矩阵：

经过mask处理后如下所示，可以看到权重矩阵变成了一个下三角矩阵。

具体的代码实现可以参考Multi-Head Attention小结。

Decoder的输入

从前面的Decoder整体结构图中我们可以看到Decoder的输入包含两部分，一部分是Encoder对原始序列（源语言序列）的编码信息，一部分是目标序列（目标语言句子）。其中，训练时的目标序列是ground truth序列，而推理时每次的输入是前面所有时间步预测的词组成的序列。与Encoder一样，也需要对输入的序列进行位置编码，编码方式与Encoder相同，这里不再赘述。

从图中可以看到，Masked Multi-Head Attetion模块的输入都是来自输入的目标序列，为了利用原序列中的编码信息，每个Encoder层（一共N个）的第二个Multi-Head Attetion子层会接受Encoder中的编码信息作为注意力计算中的K和V，将前一层的输出作为Q，目的是用找到与目标序列词相关性较大的源序列词的编码信息，从而更好地预测当前时间步的词。

这里有个细节需要注意，那就是Encoder本身的输出只是一个的矩阵，那又是如何转换为两个K和V矩阵的？Transformer的做法是直接将outputs既作为K，也作为V直接使用。

具体可以参考下面这个较详细的结构图。

整个Decoder的输出与Encoder一样，都是。

Linear、Softmax和Loss

Transformer将翻译视为分类任务进行目标语言输出，具体地，首先获取数据集中目标语言的词数量，即词库的大小，假设为M，然后每个时间步就是进行M分类，将对应预测概率最高的词作为本次时间步的输出。

所以，就像大部分分类模型一样，Transformer会在Decoder后面跟一个线性分类层和Softmax层。这里详细说明下，Linear是如何将Decoder的输出映射为分类的logits张量的。类似于Feed Forward子层一样，这个Linear层其实也是Position-wise的，即序列中所有位置的词都是共用同一个线性映射权重的。如前，假设词库的大小为M，则Linear的权重大小为，则映射后的logits张量的大小为，然后在logits的最后一个维度（M）上执行Softmax操作，选择概率最大的作为预测输出，就可以得到个词的预测，即序列中每个位置词的预测，大致流程如下图所示。训练时，使用分类任务常用的CrossEntropy函数就可以了。

到此，Transformer相关的内容基本上就介绍完了，下面是Vision Transformer相关的内容。

2. Vision Transformer

2.1 ViT

在ViT之前的很多工作都或多或少尝试在CNN网络中添加各种各样的attention机制。但是ViT的提出完全颠覆了之前的观点，它让我们看到，对于图像识别任务，也可以完全使用attention的深度网络来解决。总体上，理解了Transformer的原理，就会发现ViT的实现还是很简单的。

下面是整个模型的结构图，可以看出，核心就使用了Transformer的Encoder模块（下图右侧），从图中看出该模块并非是原始的Transformer论文的Encoder，而是后续的改进版，主要是将LayerNorm放在每个子层的前面执行。另外，MLP子层中的全连接隐藏层后非线性激活函数使用的是GELU。

为了使用原生的Encoder模块，作者对输入的2D图像做了以下处理，从而适应Encoder的一维序列形式的输入要求。

1. 将图片按固定分辨率大小切成N个切片（patches），然后将切片内的所有像素值扁平化，即延通道方向拼接起来。假设切片大小为，图片的大小为，则经过切片和扁平化处理后的张量为， 其中N为序列的长度且；
2. 对每个切片张量进行线性映射，整个Encoder的每个子层中使用的序列词的embedding维度为D，所以使用一个全连接层对其进行线性映射，则维度变成；
3. 类似于BERT，在序列的头部添加一个可学习的类别embedding，维度与切片的维度相同（至于为什么这么做，笔者没有深究，感兴趣的读者可以进一步阅读BERT文章），这样序列的长度就变成N+1，输入变成；
4. 与Transformer一样，需要对切片的嵌入添加位置编码，ViT使用的一维可学习位置编码（维度与切片embedding的维度相同），并未使用二维位置编码，因为作者发现二者在效果上并没有差异。

经过以上处理后，就可以作为Encoder的输入，如前所述，Encoder最后一个层的输出仍然是。ViT并没有使用整个嵌入进行类别计算，而是只使用序列的第0个位置，即可学习类别embedding作为类别推理时使用的图片表征，这也是BERT的做法。最后，使用一个MLP模块进行分类预测。该MLP模块在大规模数据集上预训练时使用的是两层结构，在小数据集上微调时只使用单层，输出层参数维度为，K为类别概率。

整体的推理计算过程如下面公式所示：

以上就是ViT所有的相关细节，下面是github上ViT的Pytorch实现主页提供的动态推理效果图。具体的代码实现也很简单，可参考Pytorch实现版本。

评价

总体上，ViT还是属于挖坑的方法，存在很多的问题需要解决，比如很弱的归纳偏置，这需要相比CNN更大的训练数据集，当然，可以认为这也是它的优点。还有就是特别耗GPU内存，难以直接迁移到检测、分割等下游任务，计算复杂度与输入尺寸是平方关系，缺乏像CNN这样的局部注意力机制，推理速度慢等问题。后续的工作，基本都是针对这些问题展开的。下面我们介绍的Swin Transformer和PvtV2主要解决ViT的高计算复杂度和难以直接迁移到下游任务的问题。

2.2 Swin Transformer

概述

如前所述，Swin Transformer(后面简称Swin)主要解决难以直接迁移到下游任务和ViT的高计算复杂度的问题。下面是Swin网络的结构图。

为了解决ViT难以直接迁移到下游任务的问题，Swin实现了非常类似ResNet系列的网络结构，整个网络仍然由4个独立的stage串联而成，第2~4个stage输出的feature map的高和宽都是前一个stage的1/2，通道数都是前一个stage的2倍。这样的网络结构非常适合与后续的FPN结合并应用于检测和分割任务上。对feature map空间进行降采样并对通道进行升维的操作主要由网络中的Patch Merging模块实现。

为了解决计算复杂度问题，作者将ViT中（global）计算全局注意力的方式（feature map中每个patch都会与其他所有的patch计算注意力）改为局部窗口（local window）计算，每个窗口由固定数量的相邻patch划分得到，窗口之间没有重叠，如下图左侧部分所示（红色框为窗口，灰色框为patch），这样便可以将计算复杂度与输入分辨率的二次方关系（）降低为线性关系（）。该功能通过将ViT中的MSA（Multi-head Self-Attention）子层替换为基于局部非重叠窗口计算注意力的W-MSA（Window MSA）来实现。然而，局部窗口注意力的引入限制了feature map中patch之间的信息交互，这会影响模型的表达能力。基于此，作者将W-MSA改进为带偏移的（shifted）SW-MSA模块，简单理解就是将前一层W-MSA的所有窗口平移一定的位置，这样之前不在一个窗口内的patch经过窗口重新划分后处于同一窗口（窗口大小不变）内，这样便能达到更大范围的patch进行信息交互的目的。如下图右侧所示，图中标示蓝点的patch处于不同层的不同窗口内，便能够与更大范围的patch计算注意力。为了实现以上效果并兼容较低的计算复杂度，作者将配置有W-MSA的注意力子模块与配置有SW-MS的注意力子模块串联起来，作为一个整体应用到网络结构的设计中，文中称为Swin Transformer Block，如上图b所示。根据网络规模的设计，不同stage会配置不同数量的Swin Transformer Block，非常类似ResNet的Residule Block和BottleNeck Block。

以上是Swin整体的概述，下面我们再分别探究下网络中几个关键组件的具体细节。

Patch Partition & Linear Embedding

类似于ViT，Swin先将原图片按照固定数量的像素值（patch size）切分成不重叠的patch（或称为token），作为注意力计算的最小单元，Swin的patch size为4，假设原图为，则经过Patch Partition后的尺寸为。Swin中第一个stage的通道输入假设为C（论文中），则需要将输入线性映射为。总体上，这个功能是通过一个kernel_size=4，输入通道数为3，输出通道数为C，步长=4的卷积操作一步完成。

Patch Merging

如前所述，第2~4个stage输出的feature map的高和宽都是前一个stage的1/2，通道数都是前一个stage的2倍。这由网络中的Patch Merging模块实现。这个过程其实也可以用一个kernel_size=2，输入通道数为C（假设是第2个stage），输出通道数为，步长=2的卷积操作完成。但是Swin的实现中并未这么做，而是像ViT中划分patch的方式一样，先对输入矩阵进行切分，然后对属于同一patch的像素进行通道维度的合并（假设为），最后使用线性映射，将每个patch的通道数映射为指定大小（如）。个人觉得这两种实现基本等价，可能是为了降低模型的参数量，所以作者选择了后者。

Swin Transformer Block

如前所述，具体的Swin Transformer Block分为两种，分别是使用W-MSA的Block和SW-MSA block，二者的区别仅仅就是计算注意力时窗口的划分方式有些区别，下面分别介绍这两种注意力。

W-MSA

W-MSA的原理相对比较简单，ViT的注意力计算是在feature map空间内的全部patch上进行的，为了在不重叠的局部窗口内计算注意力，Swin将原特征图矩阵重新排列成维度更高尺寸更小的特征图，然后计算注意力就和在原特征图上没什么区别了。例如，原特征图尺寸为，假设窗口大小为M，即一个窗口由相邻的个patch组成，那么重新排列后的特征图维度为，这样注意力计算就在每个的子特征图上进行。代码实现如下：

def window_partition(x, window_size):
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    return windows

论文中也给出了W-MSA将注意力的计算复杂度由与输出尺寸的平方关系，降到了线性关系，如下图公式（2）所示，公式中忽略了softmax计算。

现给出以上公式具体的推理过程。假设原图已经通过Patch Partition划分为的特征图。为了便于阅读，我们先给出Transformer的注意力计算公式：

假设Q、K、V的维度都是，则、、和的维度都是。在MSA中，计算Q、K和V，需要的计算量，然后计算（维度为）需要的计算量，最后进行加权平均，也需要，计算过程中使用的是多头注意力，多头特征融合（）还需要，所以总体上就是公式(1)的复杂度。但是由于W-MSA中的注意力是在的窗口内进行的，所以计算仅需要个的计算量，即。同样的，计算加权平均也需要计算量，其他部分计算量不变，则最终计算量如公式(2)所示。推理过程参考了这篇博客。

SW-MSA

基于SW-MSA的Block是位于W-MSA Block之后，换句话说SW-MSA Block就是为了扩展W-MSA Block的注意力视野的，所以都是在前一个W-MSA Block配置上将所有窗口的位置延H和W方向平移的距离，从而产生新的窗口，如上面图2右侧所示。然而，这样会带来两个问题，第一个是窗口的数量增加了，从原来的变成，如果窗口的数量较小时，如，则会增加到，增加了2.25倍。另外一个问题是这样划分后，每个窗口内的patch数量不相等，这样就很难进行矩阵计算，从而影响计算效率，简单的方式是使用padding的方式，将每个窗口的尺寸补全，但是这样会增加计算量。作者通过循环偏移的方式高效解决了这两个问题。整体的过程如下图所示。

光看上图有些抽象，我将该过程拆解为如下图5步进行解释，首先我们假设原特征图的尺寸为，窗口大小为，则窗口数量为4，下图步骤1为W-MSA的窗口划分方式，称为正常窗口划分（regular window partition）, 步骤2为偏移过的窗口划分，可以看到窗口数量增加为9个（9个不同颜色表示），其中只有中间的黑色窗口的尺寸是4，其他都不正常，为了保持黑色窗口内的4个patch关系，作者将特征图矩阵沿着（0,0）位置进行滚动，从而将黑色窗口patch张量组移到特征图空间的左上角，具体的做法如下图中的步骤2和步骤3，首先将其延H方向滚动1个像素，然后延W方向滚动1个像素，就得到了步骤4中的特征图，其中步骤2和步骤3可以用类似torch.roll的函数一次实现。可以看到步骤4既保持了黑色窗口内的patch组合关系没有被打乱，而且将窗口数量减少到4，且每个窗口的尺寸都是一样的，如图5所示。另外，计算完注意力后还要将特征矩阵还原回去，即按照相反的方式将步骤5中每个位置的patch还原回到步骤1中的位置。

虽然目前的SW-MSA解决了窗口数量和尺寸的问题，但是又引入了一个新的问题，有些窗口内的patch并不是相邻的，以上图中步骤5个红色窗口为例，里面的4个patch在原特征图中刚好位于四个角上，互相之间的距离太远，作者并不希望这些patch之间进行信息交互（这样会不会带来什么负面影响，论文中并没有提及，可能是因为它会破坏作者引入的局部窗口的归纳偏置吧）。解决这个问题的方法，可以使用Transformer中Decoder模块用到的mask方法屏蔽掉这些patch之间计算注意力。具体的图文介绍可以参考这篇博客。

由于每次平移的尺寸都是，那么无论特征图和窗口的尺寸是多少（窗口尺寸必须小于等于特征图尺寸），总体上只有三种mask生成情况，这一点在这篇博客中并未详细说明，分别是：窗口内的所有patch之前就位于同一窗口（称为正常窗口)，即上图中的黑色窗口；第二种是正常窗口的右侧和下方，窗口内的patch由之前的两个窗口的patch组成，即上图中的绿色和黄色窗口，最后一种就是由之前至少四个窗口的patch组成的红色窗口。下面是特征尺寸为, 窗口尺寸为，平移尺寸为2的窗口划分矩阵，数字相同的表示在roll之前位于同一窗口内。

tensor([[0., 0., 0., 0., 1., 1., 2., 2.],
        [0., 0., 0., 0., 1., 1., 2., 2.],
        [0., 0., 0., 0., 1., 1., 2., 2.],
        [0., 0., 0., 0., 1., 1., 2., 2.],
        [3., 3., 3., 3., 4., 4., 5., 5.],
        [3., 3., 3., 3., 4., 4., 5., 5.],
        [6., 6., 6., 6., 7., 7., 8., 8.],
        [6., 6., 6., 6., 7., 7., 8., 8.]])

具体的mask生成过程就不赘述了，这里截取了作者开源代码生成mask的部分，可以进一步验证：

def window_partition(x, window_size):
    H, W = x.shape
    x = x.view( H // window_size, window_size, W // window_size, window_size)
    windows = x.permute(0, 2, 1, 3).contiguous().view(-1, window_size, window_size)
    return windows

def get_mask(res, window_size, shift_size):
    H, W = res
    img_mask = torch.zeros(H, W)
    h_slices = (slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None))
    w_slices = (slice(0, -window_size),
                slice(-window_size, -shift_size),
                slice(-shift_size, None))
    cnt = 0
    for h in h_slices:
        for w in w_slices:
            img_mask[h, w] = cnt
            cnt += 1
            
    print(img_mask)

    mask_windows = window_partition(img_mask, window_size)
    print(mask_windows)
    mask_windows = mask_windows.view(-1, window_size * window_size)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
    
    return attn_mask

mask = get_mask((8, 8), 4, 2)

另外，值得一提的是作者这里使用mask的方式与Transformer中的不太一样，Transformer中的方式是直接用极小值填充Q、K计算的注意力矩阵：

attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
if mask is not None:
    attn = attn.masked_fill(mask == 0, -1e9)

而Swin的实现是将一个包含负值（-100）的mask（参考上面的mask生成代码）与原attention进行element-wise加和：

attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)

这样使用的原因个人不太清楚（也许这样是等价的），有了解的读者，还望不吝赐教。

相对位置编码

与ViT中类似，Swin中也是用的是可学习的位置编码，而且是2维相对位置编码：

然而，可学习位置编码参数的初始化大小为窗口尺寸M的大致2倍，称为参数表，每次使用的位置参数B都是从中预定义的位置获取，论文并未详细说明这样实现的原因，可能是引用了某篇论文中相对位置编码方法，这里先记录下，回头研究了再补充。

评价

Swin虽然在效果和性能上取得了显著的提升，但是还有个明显的缺点，那就是在224分辨率上预训练的模型无法直接用于任意分辨率的finetune或下游任务，而且推理时并不支持像CNN一样的任意分辨率输入。

2.3 PVTv2

概述

由于PVT的v1和v2两篇论文很接近，而且v2显著优于v1，所以这里就不分别介绍v1和v2了，直接总结v2版本。PVT要解决的问题和Swin一样，构建直接可以用于下游检测和分割任务的backbone网络，并且降低ViT的计算量。下图是从v1论文中截取的网络结构图，v2和v1在总体结构上没有差异。类似于ResNet系列模型结构，PVT和Swin都是由多个stage组成，每个stage内部由若干个统一的注意力模块组成，每个stage都会对输入的feature map的空间尺度进行降采样并增加通道维数，最后输出呈金字塔形的多层feature map。而且二者在降低空间尺寸和增加通道维度的方法上基本类似，最大的区别是feature map的注意力计算方式上。

具体地，PVT的每个stage由一个Patch Embedding子层和多个串联Attention模块（下图中的Transformer Encoder）组成。Patch Embedding便是用来降低空间尺寸和增加通道维度的，类似于Swin中的Patch Merging子层。每个stage会将输入feature map的宽和高减少为原来的1/2（第一个stage为1/4），通道数也会增加，但不是增加2倍，每个stage的输出通道数分别为64、128、320和512。

下面分别探究PVT的部分细节。

Patch Embedding

PVTv2的Patch Embedding实现与Swin类似，使用一层卷积操作实现patch划分和线性映射。不同的是，PVTv2使用了重叠的大尺寸卷积操作，这样使得相邻的patch之间有一定比例的信息重叠，作者认为这样能够保持图像数据的局部连续性（并未解释这样的好处，文中也没有消融实验），如下图（a）所示。具体地，第i个stage中Patch Embedding的卷积核步长为，卷积核大小为，padding尺寸为，输入通道数为，输出通道数为（即卷积核数量），输出为（flatten之前的shape），即为每个stage的空间降采样比例，然后接一个LayerNorm。

Transformer Encoder

SRA和Linear SRA

PVTv2对Transformer Encoder（已成为Transformer Block）的改进也是为了减少Attention模块的计算量，具体的做法是直接减少K和V的数量，这样便能明显较少矩阵计算。假设原始的Q、K、V的维度都是，长宽减小的比例为，则K、V的维度都是，则计算的复杂度由原来的：

减小到：

虽然不像Swin能够直接将计算复杂度降低为的线性关系。作者将这样的Attention改进版称为SRA（Spatial Reduction Attention）。既然是空间的降采样（K、V重排列为），同样有多个实现方式，Patch划分+线性映射的方式，或者使用步长大于1的卷积操作，在PVTv2中使用的是卷积操作，PVTv1中使用的是前者。是超参，不同stage具体的值不同，文中使用的分别为8、4、2和1。

另外作者还提出了一种称为Linear SRA的注意力模块，可以将计算复杂度降低为的线性关系。具体的做法是暴力地将K和V的Patch数量（或空间尺寸）降采样为常数量，具体地可以使用AdaptiveAvgPool2d实现，一般池化的窗口大小设置地比较大，文中使用的是7。

两种SRA的比较如下图。

Position Encoding

最后，为了使得PVT像CNN网络一样可以接受任意尺寸的输入，作者将ViT中使用的与patch数量必须一致的可学习位置编码改为带0值padding的位置编码方法（Zero Padding Postion Encoding）,该方法首次是在”Conditional Positional Encodings for Vision Transformers"这篇论文中提出，具体的实现是在MLP子模块的第一个Linear层与非线性激活函数（例如GELU）之间增加一个3 * 3的带zero padding的深度可分离卷积，如上面的图1(b)所示。具体的原理这里就不展开了，关于Vision Transformer的位置编码方式已经可以另开一篇文章了。

评价

整体上PVT实现简单且有效，尤其是用SRA的方式降低注意力计算量，应该比较契合Kaiming He最新论文中的观点：

因为本身图像数据存在冗余性，均匀间隔地对Q、K的尺寸（flatten后指数量）做一定的降采样，也不会影响模型的能力。

缺点是论文中缺乏消融实验，不清楚每个改进点带来多少的效果提升。

3. 个人的看法

1. 理性看待paper里的效果，尤其是在COCO和ImageNet上的效果。当然，纯粹发论文另当别论。
2. Vision Transformer对细粒度分类可能有效。
3. 背景单一的检测任务，可能并不需要注意力机制，例如工业数据集。
4. 可迁移性还是较弱，可能配合自监督预训练会有所改善。
5. 非常吃显存仍然是Vision Transformer的弊端，这将严重导致Vision Transformer在检测任务中的发挥，尤其对于小目标检测任务（大分辨率才能让模型看清楚目标特征）。
6. Vision Transformer依赖大规模数据集，但是无监督方法可能会缓解这个问题。
7. 模型参数量容易达到饱和，更多的参数，迁移到其他任务并没有提升，也可能和数据量有关。
8. Vision Transformer在通用检测任务上提升很多。

4. 参考

文中绝大部分图片来自参考文献。

1. Attention Is All You Need
2. https://jalammar.github.io/illustrated-transformer
3. https://zhuanlan.zhihu.com/p/54356280
4. https://mp.weixin.qq.com/s/S89kak4El3hPZJc0EzRszA
5. https://zhuanlan.zhihu.com/p/308301901
6. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
7. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
8. CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows
9. PVTv1: Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions.
10. PVTv2: Improved Baselines with Pyramid Vision Transformer
11. Conditional Positional Encodings for Vision Transformers.
12. https://colah.github.io/posts/2015-08-Understanding-LSTMs/
13. https://looperxx.github.io/CS224n-2019-08-Machine%20Translation,%20Sequence-to-sequence%20and%20Attention/
14. https://medium.com/deeper-learning/glossary-of-deep-learning-word-embedding-f90c3cec34ca
15. https://github.com/jadore801120/attention-is-all-you-need-pytorch
16. https://github.com/lucidrains/vit-pytorch
17. https://zhuanlan.zhihu.com/p/360513527
18. Masked Autoencoders Are Scalable Vision Learners
19. https://www.borealisai.com/en/blog/tutorial-17-transformers-iii-training/
20. PowerNorm: Rethinking Batch Normalization in Transformers
    20.https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html?highlight=layernorm#torch.nn.LayerNorm