【学习日记week6】基于掩蔽图像生成的跨模态学习方法MVLT和金字塔模型PVT

Masked Vision-Language Transformer in Fashion

前言：为什么要看这篇文章？

上周讲的两篇工作中，有一个核心的思路是通过掩蔽学习来实现细粒度对齐，从而帮助预训练模型的性能更好。ViLEM和COTS的区别主要在于，一个掩蔽学习是进行了替换掩蔽，另一个是直接用掩码[MASK]做掩蔽，且ViLEM是仅在文本模态，COTS是在跨模态上做掩蔽（但是这个掩蔽也是模态独立的）。但是COTS在视觉上的掩蔽的思路有个很明显的问题：COTS的掩蔽是在token level的，先将原本的patch进行token化，然后在token上进行复原。但是patch的掩蔽对于细粒度信息的学习帮助可能并没有那么大，而本文在掩蔽学习上的效果是较好的，且其是在像素级别的学习，这可能是最后一篇有关掩蔽学习的论文阅读，之后会更加细化的学习一些其他跨模态学习任务。

Motivation

现有最新的VL模型几乎都是基于transformer实现的，采用ViT/BERT或者其组合结构等。然而一般的vision-language模型在fasion领域里效果不是很好。fashion领域中更注重一些图像的细节信息，而现有的方法一般有两个问题：1. 粒度信息不足；2. 迁移性较差（预训练的特征提取器对于这种特定细化的任务的效果并没有那么好）

Contribution

基于上述问题提出来了MVLT（和上周的MVLM差不多）的一个框架。其中特别的，本文对于掩蔽的学习依靠的是生成的方法，将其作为一个生成任务来进行学习。此外，对于transformer，本文采用了金字塔结构的transformer（pyramid vision transformer）
具体的贡献如下：

1. 提出了masked image reconstruction任务（MIR），实现了像素级别的生成方法来帮助VLP
2. 在MIR任务的基础上，采用了一种端到端的VL框架，在fasion的领域内提高了其迁移性
3. 实验。
   
   一个简单的对比：MIM的任务上，在提取完特征后，会在特征上对齐进行掩蔽，然后基于BERT做出一个prediction。本文的方法是直接将掩蔽的patch喂给模型后，通过生成方法生成一个图像的patch来和原本的进行对比。

背景与相关研究

主要看有关Masked Learning Strategy的部份。
这里对于大部份的Masked Learning方法进行了一个分类：

• regression task：回归任务从一个随机替换中构建token特征
• classification task：分类任务通过预测token的属性/标签来进行学习

具体到工作，Kaleido BERT通过Kaleido策略来进行多粒度语义学习。虽然能够有一定的效果，但是这种方法通过了一些辅助工具来进行token到patch的预对齐（目标检测方法和图像描述方法）; MLIM通过图像重建的方法来增强了图像掩蔽学习（MLIM和本文的思想接近）但是MLIM的目标是对整个图像进行重建（看了一下这篇文章，大概是先将其编码到embedding空间中，然后将embedding中的一部份进行掩蔽，最后重建后计算重建损失）

实验显示，利用重建一整个图像的方法来进行学习是很困难的。近期的工作如BEiT和MAE，采用的是BERT类的预训练模型来当作视觉学习器的一部份，他们发现这么做是有效的。这两个工作让作者认为基于重建的掩蔽学习效果优于基于回归或分类的掩蔽学习。设计一个通用的生成型任务，可以在视觉语言预训练中帮助模型更好地理解和建模多模态数据，同时不需要先验知识。这对于大规模数据集和实际应用非常有用。作者的最初的目标是，希望找到一种方法，可以在没有特定任务上下文的情况下进行预训练，并提高模型的多模态建模性能。

端到端的多模态方法

1. Pixel-BERT：Pixel-BERT是第一个考虑了端到端预训练的方法。它使用2×2的最大池化层来降低图像特征的空间维度，将图像降采样了64倍。尽管它开创了端到端训练的先河，但这种粗糙和刚性的方法在实际应用中效果不佳。它只是与ResNet一起进行联合预训练，没有考虑速度和性能损失。
2. VX2TEXT：VX2TEXT提出将所有模态的数据转化为语言空间，然后使用一种松弛方案进行端到端预训练。尽管将所有模态数据转化为统一的潜在空间很有吸引力，但它忽略了将预训练方法提取的数据作为模型输入并不等同于端到端框架。它没有完全实现端到端的训练。
3. ViLT：ViLT是第一个真正探讨了端到端框架的方法，它通过用基于补丁的投影替代基于区域或网格的特征来实现端到端预训练。然而，除此之外，它没有其他设计，因此在性能上无法与其他方法竞争，因为它只是对ViT的简单扩展。
4. Grid-VLP：Grid-VLP类似于ViLT，但它更进一步，展示了使用预训练的CNN网络作为视觉骨干可以提高下游任务的性能。它认为将CNN网络作为视觉特征提取器是有益的。
5. SOHO：SOHO将整个图像作为输入，并创建了一个视觉词典，以对齐局部区域。然而，这种方法不适用于时尚特定的应用，因为它缺乏可靠的对齐信息，这导致视觉词典可能只学习到背景或前景的位置而不是复杂的语义信息。
6. FashionVLP：FashionVLP使用反馈策略来实现更好的检索性能。在实践中，他们使用从ResNet提取的预训练知识，并对整体、裁剪和地标表示进行建模。此外，他们采用Faster R-CNN作为对象检测器，以弹出RoI（区域兴趣）候选区域。
7. 其他工作：还有一些专门用于端到端预训练的工作，但它们用于特定任务，不直接适用于该研究。

和之前的工作架构类似的，本文在进行不同模态编码器的编码后，用的是一个金字塔形编码器的结构来进行学习，能够为fashion任务学习到不同层级的特征

模型结构

这篇文章的结构整体参考了PVT的结构分成了四个阶段，因为这个架构几乎就是在PVT基础上而做的，所以这里直接先跳转到后文PVT部份

在PVT基础上继续进行学习。整个MVLT也用了4stage的结构，其中有两个关键点，一是多模态的编码器，二是基于此设计的预训练目标函数。

1. 多模态的encoder

如结构图，能够接受语言和视觉的两种输入。

对于语言，首先对于图像的caption进行tokenize，然后用[MASK]随机掩蔽（比率为$r_l$）。像掩蔽过程一样，首先先获得一个词token序列，然后添加一个[CLS] token到序列前，此外，如果序列长度小于128，采用[PAD]对序列进行填充。最后生成输入序列$\text{T}\in {\mathbb{R}}^L=⟨t_1;\dots ;t_L⟩$

对于视觉，将整个$\text{I}\in {\mathbb{R}}^{H\times W\times 3}$作为输入。然后将其进行划分成patches：$\text{V}\in {\mathbb{R}}^{N\times P\times P\times 3}$。关于视觉信息的掩蔽，会放在下个部份讲，这里提一下同样也有一个比率$r_v$

上面的两个模态的信息会先编码为各自的embedding ${\mathbf{T}}^1,{\mathbf{V}}^1$ 忽略其他的stages，仅从k-stage来进行讨论。首先将文本的嵌入${\text{T}}^k\in {\mathbb{R}}^{L\times D_k}$编码到一个隐藏特征$m^k\in {\mathbb{R}}^{L\times D_{k+1}}$，有
其中后面两项是线性投影层和一个position embedding。对于视觉来讲，保持着一个PVT块内的操作，先进行一个卷积下采样，然后将其拉直（PVT内也是先reshape到一维嵌入做attention再reshape回到特征图），即

（这里原论文的卷积核好像写错了，应该是${\text{W}}_v^k\in {\mathbb{R}}^{D_k\times R_k\times R_k\times D_{k+1}}$，P应该也不是$N\times D_{k+1}$维的）

然后对其进行一个连接操作，经过多个transformer块（一般是两个）进行多个MSA操作后，再给他解除连接后转化为特征图或者language embedding。

在对齐的时候，会同时生成四个阶段的text embeddings和vision embeddings。对于具体的参数，基本参考了PVT的参数，如下

2.预训练目标（MIR）

主要分三个部份，对于语言，采用MLM，对于视觉，采用MIR，对于跨模态，采用ITM。
MIR：这里采用的是一个重构损失，直接将图像的细粒度信息匹配拉到像素级别，在fasion任务上非常重要。为了更好的用MIR学习，mask策略利用了PVT的结构，进行灵活的masking strategy，如下图所示：

在ViT为backbone的方法中，mask的策略因为特征图尺度不会进行变化，所以只能说固定大小掩蔽，PVT的结构允许了多尺度掩蔽。(如上图右部分所示)掩蔽的公式可以写成如下形式：

这个M函数在$q\in \Phi$时取得1，其余时间取0。而$\Phi$是一组在[1,Q]中以比率$r_v$选取的随机整数。而$Q=\frac{H\times W}{(\alpha \times P{)}^2}$。其中$\alpha$的取值范围为1-8，P是patch的大小。（注意一下，P应该是会改变的？）

重建损失采用的是了l1和l2结合的smooth-l1范数，公式如下。

这里的${\text{I,I}}^{\prime }$都是像素级别的，前者是原始图像后者是重建图像，是通过一个Unet和之前不同阶段的输出

就是用了一个Encoder-Decoder的结构和一个smooth-l1来做一个粗略的重建，这真的有效吗？

3. 图文匹配损失（ITM）

这里采用的是一个交叉熵损失，计算“是否匹配”的01标签和匹配概率的交叉熵。匹配概率有${\text{p}}_{\text{ITM}}={\mathcal{F}}_{\text{ITM}}(⟨\text{T},\text{V}⟩;{\text{W}}_{\text{ITM}})$这是一个线性分类层，只有一个二维的输出（直接用一个分类头来监督）

4. MLM损失

文本就是直接掩蔽为[MASK]，采取的loss是负对数似然函数（交叉熵）。

同样的，p是一个从分类头中得到的一个结果。
综合上面三个损失加权和得到最终的损失函数。

具体实验（代码阅读）

首先是一个整体的框架，本次实验基本上是在PVT的基础上做的，可以在libs.pvlt中查看：

首先是PVT的代码块

class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
            attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, H, W, T_num):
        x = x + self.drop_path(self.attn(self.norm1(x), H, W, T_num))   # # bs, 7744+40, 64
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x

每个stage下采样所用的patch embedding：
首先将输入的维度提取后将其用一个Conv2D进行下采样（注意kernel的大小和stride都是patch的大小），然后将其进行flatten操作，然后将其前两个维度互换

为什么要互换？
首先，在进行了卷积后，通道数变为embeddim，总维度为：(B, embed_dim, H, W)
注意H和W是进行块化之后的宽和高，每个H或者W代表的是多个像素。即H*W=patch_size
相当于有了embed_dim维的特征图，然后将其拉直后变为 (B, embed_dim, patch_size)
但是因为要变成embedding，所以后两个维度需要进行互换。固有了下面的方法实现（PVT中也有）

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)

        self.img_size = img_size
        self.patch_size = patch_size
        assert img_size[0] % patch_size[0] == 0 and img_size[1] % patch_size[1] == 0, \
            f"img_size {img_size} should be divided by patch_size {patch_size}."
        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
        self.num_patches = self.H * self.W
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x):
        B, C, H, W = x.shape # torch.Size([5, 3, 352, 352])

        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm(x)
        H, W = H // self.patch_size[0], W // self.patch_size[1]

        return x, (H, W)

然后是加入了语言嵌入的MLVT的前馈块的实现：

        for i in range(num_stages):
            patch_embed = PatchEmbed(img_size=img_size if i == 0 else img_size // (2 ** (i + 1)),
                                     patch_size=patch_size if i == 0 else 2,
                                     in_chans=in_chans if i == 0 else embed_dims[i - 1],
                                     embed_dim=embed_dims[i])
            text_embed = nn.Sequential(
                nn.Linear(in_features=token_hidden_size if i==0 else embed_dims[i-1], 
                          out_features=embed_dims[i]),
                nn.LayerNorm(embed_dims[i]))
            num_patches = patch_embed.num_patches if i != num_stages - 1 else patch_embed.num_patches + 1
            pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dims[i]))
            text_pos_embed = nn.Parameter(torch.zeros(1, num_text_tokens, embed_dims[i]))
            pos_drop = nn.Dropout(p=drop_rate)
            
            block = nn.ModuleList([Block(
                dim=embed_dims[i], num_heads=num_heads[i], mlp_ratio=mlp_ratios[i], qkv_bias=qkv_bias,
                qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + j],
                norm_layer=norm_layer, sr_ratio=sr_ratios[i])
                for j in range(depths[i])])
            cur += depths[i]

            setattr(self, f"patch_embed{i + 1}", patch_embed)
            setattr(self, f"text_embed{i + 1}", text_embed)
            setattr(self, f"pos_embed{i + 1}", pos_embed)
            setattr(self, f"text_pos_embed{i + 1}", text_pos_embed)
            setattr(self, f"pos_drop{i + 1}", pos_drop)
            setattr(self, f"block{i + 1}", block)

            trunc_normal_(pos_embed, std=.02)
            trunc_normal_(text_pos_embed, std=.02)
        
        # define PVLT-TextEmbedding

这段代码是分为了四个阶段的，每个阶段需要设置一个相应的函数方法，在运行的时候进行调用。
这里解释一下加入了language的前向过程：首先先将图像传入图像嵌入，文本传入文本嵌入（图像嵌入用降采样的嵌入方法计算）

    def forward_pyramid_features_vl(self, x, y):    # x = vision, y = language
        img_feats, text_feats = [], []

        B = x.shape[0]
        y = self.text_embeddings(y) # torch.Size([5, 128]) -> torch.Size([5, 128, 768])

        for i in range(self.num_stages):    # four stages
            # init functions
            patch_embed = getattr(self, f"patch_embed{i + 1}")
            text_embed = getattr(self, f"text_embed{i + 1}")
            pos_embed = getattr(self, f"pos_embed{i + 1}")
            text_pos_embed = getattr(self, f"text_pos_embed{i + 1}")
            pos_drop = getattr(self, f"pos_drop{i + 1}")
            block = getattr(self, f"block{i + 1}")

            # process
            x, (H, W) = patch_embed(x)  # 5, 7744, 64
            y = text_embed(y)   # torch.Size([5, 128, 768]) -> torch.Size([5, 128, 64]) # TODO

            if i == self.num_stages - 1:
                pos_embed = self._get_pos_embed(pos_embed[:, 1:], patch_embed, H, W)
            else:
                pos_embed = self._get_pos_embed(pos_embed, patch_embed, H, W)
            # print('debug376', i, x.shape, y.shape)
            x = pos_drop(torch.cat((x + pos_embed, y + text_pos_embed), dim=1)) # [bs, 88*88=7744, 64] [bs, 44*44=1936, 128], [bs, 22*22=484, 320], [bs, 11*11=121, 512]
            
            for blk in block:
                x = blk(x, H, W, self.T_num)
            x, y = torch.split(x, [H*W, self.T_num], dim=1)

            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
            img_feats.append(x)
            text_feats.append(y)

        return img_feats, text_feats

这里会将每一层学到的特征都存储在队列中，因为会用多层的特征进行损失计算。
损失函数的计算看具体代码，这里主要讲核心的模型结构。
下面讲一下视觉mask的实现:

        # load image and transform it
        image = self.rgb_loader(self.images[index])
        image = self.img_transform(image)

        # get image mask for masking strategy
        if self.dataset_type == 'train':
            if self.mask_strategy == 'square':
                img_mask = self.generate_square_mask(im_size=self.trainsize, mask_size=self.trainsize//self.mask_ratio)
            elif self.mask_strategy == 'stroke':
                img_mask = self.generate_stroke_mask(im_size=self.trainsize)
            elif self.mask_strategy == 'random_grid':
                # use this in our final version
                img_mask = self.generate_grid_mask(input_size=(self.trainsize, self.trainsize), mask_ratio=self.mask_ratio, patch_size=16)
            else:
                raise NameError('>>> invalid parameter: {}'.format(self.mask_strategy))
        elif self.dataset_type == 'valid':
            # print(self.grid_masking_images[index])
            img_mask = self.pkl_loader(self.grid_masking_images[index])
        else:
            raise Exception('No type named {}'.format(self.dataset_type))
        
        masked_images = image.clone().masked_fill_(torch.Tensor(img_mask).byte().bool(), value=torch.tensor(1e-6))
        t2i_labels = torch.Tensor(img_mask)

具体的，采用了grid_mask的方法

    def generate_grid_mask(self, input_size=(352, 352), mask_ratio=0.75, patch_size=16):
        # ensure input_width and input_height are divisible by patch_size
        assert input_size[0] % patch_size == 0
        assert input_size[1] % patch_size == 0

        num_width = input_size[0] // patch_size
        num_height = input_size[1] // patch_size

        num_patches = num_width * num_height
        num_mask = int(mask_ratio * num_patches)

        mask = np.concatenate([
            np.zeros((num_patches - num_mask, patch_size, patch_size)),     # the number of unmasked patches
            np.ones((num_mask, patch_size, patch_size)),    # the number of masked patches
        ], axis=0)

        mask_split = np.split(mask, num_patches, axis=0)
        np.random.shuffle(mask_split)   # random

        h_list = list()
        for i in range(num_height):
            cur_list = mask_split[0 + i: i + num_width]
            np.random.shuffle(cur_list)
            h_list.append(np.transpose(np.hstack(cur_list), (2, 1, 0)))

        final_mask = np.vstack(h_list)

        # final_mask = np.uint8(final_mask.squeeze()) * 255
        # cv2.imwrite(save_name, final_mask)
        return np.transpose(final_mask, (2, 0, 1))

Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

真的想理解这个结构让我头疼了很久，很难理解这个。。。主要是motivation很让人头痛，大概的意思是ViT的总体结构是直接一次下采样，主要是为了做分类任务而设计的，而金字塔结构的设计可以让原本的transformer可以和resnet一样用到下游任务中。
原作者的见解更为透彻，他在知乎上写了专栏：大白话VPT
问题重点在于分辨率和中间输出上。原始transformer是第一次输入进行分patch，之后通过加深层数来加深感受野。但本文的方法在原始patch的基础上，在后续的transformer前再对图像进行划分。

可以看到基本上MVLT的结构就是这个结构，只是扩展到了多模态上。这也是
整个结构分为四个阶段，分别对应四个分辨率上的transformer。每个transformer block前会进行一次patch embedding来进行下采样，然后在transformer block内，在进行一次SR的操作，以适应相应的分辨率。原作者的解释是这样的：

为了在保证feature map分辨率和全局感受野的同时降低计算量，我们把key（K）和value（V）的长和宽分别缩小到以前的1/R_i。通过这种方法，我们就可以以一个较小的代价处理4-stride，和8-stride的feature map了。