ViT的学习笔记

1. 致谢

感谢paddle开源的Vit系列教程《飞桨——从零开始学视觉Transformer》；

2. Vit经典模型

I	II	III	IV	V
DETR 2020	ViT 2020	DeiT 2021	Swin 2021

3. 基础知识

FFN: Feed Forward Network，前馈神经网络，= MLP。
MLP: Multilayer Perceptron，多层感知机。

4. Transformer理论介绍

4.1 注意力张量：[B, nhead, N, N]

张量中的N*N矩阵表示了不同单词之间的相关性；

4. Vit参考实现

Library	Description
xFormers	Meta推出的开源库，注重效率，内部使用了自定义的CUDA核进行实现
PaddleViT	百度研究院推出的Vit相关的模型库
vit-pytorch	使用PyTorch实现的Vit系列模型库

4 Vit组成模块

4.1 Attention模块（Multi-head Attention）

Attention模块是一种seq2seq的变换模块，其输入输出的特征维数相同；
参数：

• num_heads: attention head的数量，默认为8；将向量分成8份的话，计算效率也比较高。
• qkv_bias: 在qkv映射时使用偏置参数，默认为False，但是qkv_bias默认需要开启。
• attn_drop: self.attn_drop的丢弃率，默认为0。
• proj_drop: self.proj_drop的丢弃率，默认为0。

Note:

• qkv_bias在一般情况下必须开启，这样才符合逻辑回归的数学原理，我们在timm的Vit实现中也可以看到此参数默认是开启的，这里我们引用timm中Vit模型的参数解释来说明：

  qkv_bias: bool, If True, enable qkv(nn.Linear) layer with bias, default: True

  可知qkv_bias默认需要开启。

# timm库中的Attention实现
# 在torch中声明模块需要继承nn.Module
class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        # dim需要是num_heads的整数倍，可以参考HuggingFace的实现;
        # TODO: assert dim % num_heads == 0
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
    	# B为batch-size，
    	# N为patch-num，
    	# C为通道数，
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # self.qkv() --> x: (B, N, 3C)
        # complex: B*3C^2
        # reshape() --> x1: (B, N, 3, nhead, dhead)
        # complex: 0
        # permute() --> qkv: (3, B, nhead, N, dhead)
        # complex: 0
        q, k, v = qkv.unbind(0)   # make torchscript happy (cannot use tensor as tuple)

		# q: (B, nhead, N, dhead)
		# k: (B, nhead, N, dhead)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        # q@k' --> attn: (B, nhead, N, N)
		# complex: B*nhead*N*dhead*N = BN^2*C
		# *self.scale --> attn
		# complex: B*nhead*N*N = BN^2*nhead
        attn = attn.softmax(dim=-1)
        # complex: BN^2*nhead
        attn = self.attn_drop(attn)
        # complex: 0
		
		# attn: (B, nhead, N, N)
		# v: (B, nhead, N, dhead)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        # attn @ v --> x
        # complex: B*nhead*N*N*dhead = BN^2C
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

复杂度分析

可以看到，Attention模块的计算复杂度为
$$3B{C}^2+B{N}^{2}C+B{N}^2\cdot nheads+B{N}^2$$
我们将复杂度按照算子分开列出：

• FC算子self.qkv: $3B{C}^2$
• 乘法注意力q@k': $B{N}^{2}C$
• Softmax算子: $B{N}^2\cdot nheads$
• 注意力加权运算: $B{N}^2$

3.2 Vit模型为什么要使用Multi-head-Self-Attention?

这里我们针对仅原始Vit的Multi-head-Self-Attention进行讨论，并没有扩展到所有SOTA的vit类型算法；
由Attention模块的复杂度分析可知，仅有Softmax算子的复杂度$B{N}^2\cdot nheads$与注意力head的数量有关；于是我们可以知道使用Multi-head-Self-Attention会使得Self-Attention算法的复杂度增加，且其增加的复杂度主要是来源于softmax算子的运算量增加了，且其运算量增加了$nheads$倍；

3.3 关于Self-Attention和Bilinear-Pooling

最近在学习PyTorch代码的时候，看到了一个有趣的方法—— Bilinear Pooling，我第一次看到这个方法时，感觉它跟Self-Attention有点相似，（关于 Bilinear Pooling的解读，请参考《双线性池化（Bilinear Pooling）详解、改进及应用》，以下简称为“BiPool详解”）
首先，我们来看看 Bilinear Pooling的公式，（此公式改进自原论文公式），
对于图像$\mathcal{I}$的位置$l$，（这里使用张量符号“$\mathcal{I}$”表示图像，因为图像维度一般是NCHW），有两种特征$f_A(\mathcal{I},l)\in {\mathbb{R}}^{T\times M}$和$f_B(\mathcal{I},l)\in {\mathbb{R}}^{T\times N}$，则 Bilinear Pooling的计算为
$$\begin{array}{rll}\text{opr}(\mathcal{I},l,f_A,f_B)& ={f_A}^T(\mathcal{I},l){f_B}^T(\mathcal{I},l)& \in {\mathbb{R}}^{M\times N}\\ \sigma \left(\mathcal{I}\right)& ={\sum }_l\text{opr}(\mathcal{I},l,f_A,f_B)& \in {\mathbb{R}}^{M\times N}\\ x& =\text{vec}\left(\sigma \left(\mathcal{I}\right)\right)& \in {\mathbb{R}}^{MN}\\ y& =\text{sign}\left(x\right)\sqrt{\left|x\right|}& \in {\mathbb{R}}^{MN}\\ y& =y/{\Vert y\Vert }_2& \in {\mathbb{R}}^{MN}\end{array}$$
直观上来说， Bilinear Pooling就是首先将计算不同位置上两种特征的自相关外积矩阵；

4 学习笔记

4.1 使用einsum()实现Vit中的矩阵乘法

请参考《矩阵视角下的Transformer详解（附代码）（by 孙裕道）》；

4.2 Vit中使用的LayerNorm是一般认为的那个Layer Norm吗？

这个问题的来源于这篇文章《BERT用的LayerNorm可能不是你认为的那个Layer Norm？》；这篇文章的核心观点是：作者经过探究发现，BERT里面使用的“layer-norm”实际上都是在做“instance-norm”。
南溪认为，这个观点是不能同意的，作者出现这种混淆的原因是将CV中的norm的范式套用在NLP中norm的范式上，于是觉得似乎可以相提并论，

也就是认为：$[H,W]⟺L,C⟺D,N⟺B$；
实际上这样的“等价”是“行不通”的，将“$[H,W]$等价于$L$”，会产生“layer-norm在做“instance-norm”的误解；而对于序列模型而言，$L$和$D$分别是两种不同层级的特征通道，在$L$或$D$维内部特征点之间需要是尽可能“互斥”的关系，也就是它们并不是相互独立的；
在南溪看来，$[H,W]$之所以被合成到一个维度中，是因为CV中会有“图像上的特征点常常具有位置无关的特性”，也就是可以认为特征图上的像素点是“独立同分布”的，于是恰好可以套用IN的假设来对每个通道的特征图进行归一化，如下图展示的IN的原理，

而BN则是将范围扩大至$[N,H,W]$三个维度，（在实践中使用较多）。