RIS 系列 Synchronizing Vision and Language: Bidirectional Token-Masking AutoEncoder for RIS 论文阅读笔记

写在前面

  又是一周周末，可惜今天运气不咋好，上午被小汽车撞自行车后座上了，手臂皮外伤，所幸人没事。下午继续淦论文吧。

• 论文地址：Synchronizing Vision and Language: Bidirectional Token-Masking AutoEncoder for Referring Image Segmentation
• 代码地址：原论文未提供
• 预计投稿于：CVPR 2024
• Ps：2023 年每周一篇博文阅读笔记，主页 更多干货，欢迎关注呀，期待 6 千粉丝有你的参与呦~

一、Abstract

  Referring Image Segmentation 指代图像分割旨在在像素水平上分割出自然语言表达式所指的特定目标。最近一些基于 Transformer 的方法凭借着注意力机制生成上下文 query，虽然很是成功，但是未能理解复杂表达式中的上下文。于是本文受到 masked autoencoder (MAE) 的启发提出 bidirectional token-masking autoencoder (BTMAE) 双向 token-masking 自编码器。通过在 token 水平上利用图像-文本以及文本-图像构建出图像和文本上缺失的特征。换句话说，这一方法涉及到两个模态的相互补全。实验效果很好。

二、引言

  RIS 的定义，最近的一些进展，应用。一些工作基于 Transformer 编码器利用 Transformer 的编码器 tokensize 视觉和文本特征，使用交叉注意力建模多模态间的关系。其缺点在于可能会产生混乱的视觉-语言关系，一些工作尝试利用丰富的多模态上下文信息提取 tokens、使用类似对比学习的方式、视觉-语言 Transformer 微调、多模态 embedding 空间来增强联系，然而这些方法仍然在建模高维度视觉-语言上下文时有难度。

  为解决这些问题，本文关注于 Masked Autoencoder (MAE) 的自特征建模能力。然而很难将其应用在 RIS 任务上。原因：基于 MAE 的预训练需要大量数据；传统的单模态 MAE 不能学习多模态的知识。

  于是本文提出 bidirectional token masking token autoencoder (BTMAE) 双向 token masking 自回归编码器来解决 RIS 及 MAE 的缺点。

  如上图所示，将 BTMAE 划分为双流来建模模态间的高维度关联信息：其中一流用于从完整的文本中重构缺失的图像，另外一流从完整的图像中重构缺失的文本。进一步，为维持模型的鲁棒性，提出一个解码器层，名为 Impact Token Attention (ITA)，ITA 从 BEMAEs 生成的视觉语言 tokens 中采样关键的 tokens。于是提出的模型会尽可能消除混乱句子的噪声干扰，并生成一个正确的预测 mask。

  在 RefCOCO、RefCOCO+、GRef 数据集上进行评估，提出的方法达到了 SOTA。本文贡献总结如下：

• 提出 BTMAE，在无大数据集的情况下能够捕捉多模态上下文信息，有效增强模型的性能；
• 引入 IAT 模块消除复杂句子中无关信息的影响。ITA 模块通过 BTMAE 采样多模态 tokens，来建立起关键的 tokens；
• 方法达到了 SOTA 的性能。

三、相关工作

Referring Image Segmentation

  RIS 的定义，传统的方法将图像特征图映射到语言特征空间，最近的一些方法应用 Transformer 中的注意力机制实现更高维度的视觉-语言特征融合，但问题在于未能定位到那些包含太多无关描述以及复杂背景的句子。于是有一些方法单独 tokenizing 目标和文本，然后融合这两个模态的 embedding，然而仍然依赖于图像和语言特征间的关联。而本文提出的 BTMAE 以无监督学习的方式能够跨模态建模上下文信息。

Masked Autoencoder

  MAE 旨在重构输入的数据，同时选择性地忽略或 mask 掉输入的某些部分，通常用于特征学习和维度降低。提出的方法利用 MAE 特征表示建模的能力学习模态间复杂的上下文关系。

四、提出的方法

4.1 整体框架

  如上图所示，提出的模型包含视觉和语言编码器用于特征提取，两个 BTMAEs 模块用于增强上下文信息。此外，ITA 的解码器层生成多模态紧凑型的 tokens，用于生成最终预测的 masks。具体来说，首先输入的 RGB 图像通过视觉编码器生成多尺度的融合特征 ${\mathbf{X}}_{\mathbf{v}}\in {\mathbb{R}}^{C\times H\times W}$，其中 $C$ 是 embedding 维度，$H$、$W$ 分别为特征的高、宽。

  对于文本输入，其通过一个预训练的语言编码器来建立语言特征 ${\mathbf{X}}_{\mathbf{t}}\in {\mathbb{R}}^{C\times H\times W}$，其中 $L$ 为特征的长度。这两种特征分别提炼为 tokens ${\mathrm{T}}_{\mathbf{v}}\in {\mathbb{R}}^{C^{\prime }\times (H\times W)}$ 和 ${\mathbf{T}}_{\mathbf{l}}\in {\mathbb{R}}^{C^{\prime }\times L}$。而 ${\mathrm{T}}_{\mathbf{v}}$ 和 ${\mathrm{T}}_1$ 则分别用于两种类型的输入：$BTMA{E}_{l\to v}$、$BTMA{E}_{v\to l}$。在 $BTMA{E}_{l\to v}$ 中，其生成一个重构的视觉 tokens ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{r}}\in {\mathbb{R}}^{C^{\prime }\times (H\times W)}$，具有精炼的空间上下文信息。而对于 $BTMA{E}_{v\to l}$，其生成一个重构的语言 tokens ${\mathrm{T}}_{\mathbf{l}}^{\mathbf{r}}\in {\mathbb{R}}^{C^{\prime }\times l}$，具有增强的语言上下文信息。此外，ITA 在目标上生成一个关键的凝练 tokens ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{i}}\in {\mathbb{R}}^{C^{\prime }\times K}$，接下来穿过 impact token transformer (ITT) 层的得到预测的 mask。

4.2 双向 token-masking 自回归编码器

  如上图所示，将 BTMAE 划分为两步，首先 BTMAE 编码器和解码器生成自编码图像，然后仅在 BTMAE 编码器上执行 token 提炼。上图 (a) 和 (b) 预测 $BTMA{E}_{l\to v}$，而 $BTMA{E}_{v\to l}$ 则是相同的流程。

  具体来说，在视觉 tokens ${\mathrm{T}}_{\mathbf{v}}$ 上执行位置编码和随机 masking，生成 masked token ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{m}}$ 。如果 masking 比例为 $\alpha$，则 ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{m}}$ 的尺寸为 $C^{\prime }\times (\lfloor H\times W\times (1-\alpha )\rfloor )$。其次，${\mathrm{T}}_{\mathbf{v}}^{\mathbf{m}}$ 作为 $E$ Transformer 的编码器输入特征，需要注意的是一个不同的模态 tokens ${\mathrm{T}}_1$ query 用作第二个多头自注意力机制中的 key 和 value。再者，类似于 MAE，应用一个随机的 masking tokens 和位置编码到编码器的输出 tokens。而这些 masking tokens 分别插入到 ${\mathrm{T}}_{\mathbf{v}}$ 的 masked 位置上。于是 ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{m}}$ 的尺寸和 ${\mathrm{T}}_{\mathbf{v}}$ 相同了。最后将这些 tokens 穿过 $D$ 个 Transformer 解码器块，并将其 reshape，投影到一个自编码图像 ${\mathrm{I}}_{\mathbf{pred}}$。

  解码器采用 ${\mathrm{T}}_1$ 作为 key 和 value，其与 BTMAE 编码器一致。${\mathrm{T}}_{\mathbf{v}}^{\mathbf{r}}$ 仅仅采用未 masked 的 ${\mathrm{T}}_{\mathbf{v}}$，${\mathrm{T}}_{\mathbf{l}}$ 和 $E$ 个 Transformer 编码器。此外 ，$BTMA{E}_{(v\to l)}$ 有着与 $BTMA{E}_{(l\to v)}$ 相同的结构以及自回归编码器语言特征。

4.3 凝缩的 Token 注意力

  ITA 模块从增强的多模态特征 ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{r}}$、${\mathrm{T}}_{\mathbf{l}}^{\mathbf{r}}$ 中生成浓缩的 tokens ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{i}}$，用于 mask 的生成。具体来说，通过在 ${\mathrm{T}}_{\mathbf{v}}^{\mathbf{r}}$ 上执行空间全局平均池化，生成 ${\mathrm{T}}_{\mathbf{l}}^{\mathbf{r}}$。${\mathrm{T}}_{\mathbf{v}}^{\mathbf{r}}$ 和 ${\mathrm{T}}_{\mathbf{l}}^{\mathbf{r}}$ 穿过 MHA 和 FFN 层后，其输出 shape 改为 ${\mathrm{X}}_{\mathbf{v}}$ 相同的形状。接下来采用 $1\times 1$ 卷积和逐元素 softmax 操作生成空间采样 masks ${\mathbf{S}}_{\mathbf{v}}\in {\mathbb{R}}^{K\times H\times W}$，其中 $K$ 为视觉浓缩 tokens 的数量。因此，定义 ${\mathbf{S}}_{\mathbf{v}}$ 的第 $k$ 个通道为 ${\mathrm{S}}_{{\mathbf{v}}^{\mathbf{k}}}\in {\mathbb{R}}^{1\times H\times W}$，这一过程表示如下：
$${\mathbf{S}}_{{\mathbf{v}}^{\mathbf{k}}(\mathbf{x},\mathbf{y})}=\frac{e^{{\mathbf{X}}_{{\mathbf{v}}^{\mathbf{k}}(\mathbf{x},\mathbf{y})}^{\mathbf{r}}}}{{\sum }_{k=1}^K{e}^{{\mathbf{X}}_{{\mathbf{v}}^{\mathbf{k}}(\mathbf{x},\mathbf{y})}^{\mathbf{r}}}}$$其中 $(x,y)$ 为像素坐标，$k=1,2,\dots ,k$，${\mathrm{X}}_{\mathbf{v}}^{\mathbf{r}}\in {\mathbb{R}}^{K\times H\times W}$ 为标准自注意力序列的输出，其计算如下：
X v r = f 1 × 1 ( f F F N ( X a t t ) ) X a t t = ψ ( w Q T v r ( w K T l r ) ⊤ C ′ ) ( w V T l r ) + T v r

$$\begin{aligned} &\mathbf{X_{v}^{r}}=f_{1\times1}\left(f_{FFN}\left(\mathbf{X_{att}}\right)\right)\\ &\mathbf{X_{att}}=\psi\left(\frac{\mathbf{w_{Q}T_{v}^{r}}\left(\mathbf{w_{K}T_{l}^{r}}\right)^{\top}}{\sqrt{C^{\prime}}}\right)(\mathbf{w_{V}T_{l}^{r}})+\mathbf{T_{v}^{r}} \end{aligned}$$ ​Xvr​=f1×1​(fFFN​(Xatt​))Xatt​=ψ(C′ ​wQ​Tvr​(wK​Tlr​)⊤​)(wV​Tlr​)+Tvr​​其中 ${\mathrm{w}}_{\mathbf{K}}\in {\mathbb{R}}^{C^{\prime }\times C^{\prime }}$、${\mathrm{w}}_{\mathbf{Q}}\in {\mathbb{R}}^{C^{\prime }\times C^{\prime }}$，${\mathrm{w}}_{\mathbf{V}}\in {\mathbb{R}}^{C^{\prime }\times C^{\prime }}$ 分别为可学习的投影矩阵。此外，$\psi (\cdot )$ 和 $f(\cdot )$ 分别表示 softmax 和 $1\times 1$ 卷积操作。因此通过下列过程生成 $K$ 个浓缩的 tokens ${\mathrm{T}}_{\mathbf{V}}^{\mathbf{i}}={\mathrm{T}}_{{\mathrm{v}}^1}^{\mathrm{i}},{\mathrm{T}}_{{\mathrm{v}}^2}^{\mathrm{i}},...,{\mathrm{T}}_{{\mathrm{v}}^{\mathrm{K}}}^{\mathrm{i}}$：
$${\mathrm{T}}_{\mathbf{vk}}^{\mathrm{i}}=\frac{{\sum }_{x=1}^H{\sum }_{y=1}^W({\mathrm{S}}_{\mathbf{vk}(\mathbf{x},\mathbf{y})}\cdot {\mathrm{X}}_{(\mathbf{x},\mathbf{y})})}{{\sum }_{x=1}^H{\sum }_{y=1}^W{\mathrm{S}}_{\mathbf{vk}(\mathbf{x},\mathbf{y})}}$$其中 ${\mathrm{T}}_{{\mathrm{v}}^{\mathrm{K}}}^{\mathrm{i}}$ 为第 $k$ 个视觉浓缩 token。

  ITA 最后一次的处理涉及到多模态浓缩 tokens 的融合，通过 impact token transformer (ITT) 和编码器的特征来表示 IoU 区域。ITT 的结构类似于标准的 Transformer 解码器。

4.4 目标函数

  首先，在 $BTMA{E}_{(l\to v)}$ 中使用均方误差损失来重建输入的图像，这一过程如下：
$${\mathcal{L}}_{BTMA{E}_{(l\to v)}}=\sum _{x,y}{\left({\mathrm{I}}_{\mathrm{Pred}}\left(x,y\right)-{\mathrm{I}}_{\mathrm{GT}}\left(x,y\right)\right)}^2$$其中 ${\mathrm{I}}_{\mathrm{Pred}}$ 和 ${\mathrm{I}}_{\mathrm{GT}}$ 指的是预测的 RGB 图像和原始的 RGB 图像。

  对于 $BTMA{E}_{(v\to l)}$，采用交叉熵损失来性训练预测的语言 tokens ${\mathrm{T}}_{\mathrm{l}}^{\mathrm{r}}$ 和原始的语言 tokens ${\mathrm{X}}_1$。

  最后采用 IoU 损失求和以及加权的交叉熵损失来赋值更多的权重给难的像素，定义损失函数如下：
L I O U = 1 − ∑ ( x , y ) min ⁡ ( M P , M G ) ∑ ( x , y ) max ⁡ ( M P , M G ) L b c e w = − ∑ ( x , y ) w [ M G ln ⁡ ( M P ) + ( 1 − M G T ) ln ⁡ ( 1 − M P ) ]

$$\begin{aligned} &\mathcal{L}_{IOU}=1-\frac{\sum_{(x,y)}\operatorname{min}\left(\mathbf{M_P},\mathbf{M_G}\right)}{\sum_{(x,y)}\operatorname{max}\left(\mathbf{M_P},\mathbf{M_G}\right)}\\ &\mathcal{L}_{bce}^{w}=-\sum_{(x,y)}w\left[\mathbf{M_{G}}\ln\left(\mathbf{M_{P}}\right)+\left(1-\mathbf{M_{GT}}\right)\ln\left(1-\mathbf{M_{P}}\right)\right] \end{aligned}$$ ​LIOU​=1−∑(x,y)​max(MP​,MG​)∑(x,y)​min(MP​,MG​)​Lbcew​=−(x,y)∑​w[MG​ln(MP​)+(1−MGT​)ln(1−MP​)]​其中 $w=\sigma |{\mathrm{M}}_{\mathrm{P}}-{\mathrm{M}}_{\mathrm{G}}|$，${\mathbf{M}}_{\mathbf{G}}$ 和 ${\mathbf{M}}_{\mathbf{P}}$ 为 GT maps 以及预测的 maps。最后，总体的损失表达如下：
$${\mathcal{L}}_{\mathrm{total}}={\mathcal{L}}_{BTMA{E}_{(l\to v)}}+{\mathcal{L}}_{BTMA{E}_{(v\to l)}}+{\mathcal{L}}_{IOU}+{\mathcal{L}}_{bce}^w$$

五、实验

5.1 数据集 & 评估指标

  Ref-COCO、RefCOCO+、GRef；
  mean intersection-overunion (mIoU)、overall intersection-over-union (oIoU)

5.2 实施细节

  使用 ResNet-101、Darknet-53、Swin-B 预训练在 ImageNet 上作为视觉 Backbone，BERT 12 层的 Transformer，768 维度。输入图像尺寸 $480\times 480$，$BTMA{E}_{(l\to v)}$ 中编码器层数 $E=4$，$BTMA{E}_{(v\to l)}$ 中解码器层数 $D=4$。设置 BTMAE 中所有的 masking ration $\alpha =0.5$。每个 ITA 用于设定浓缩 tokens ${\mathrm{T}}_{\mathrm{v}}^{\mathrm{r}}$ 的数量分别为 16，8，2。于是，在 ITA 最后一层仅有两个 token 保留下来，即对应前景和背景。Adam 优化器，用于训练和微调的超参数为 ${\beta }_1=0.9$，${\beta }_2=0.999$，$ϵ=10^{-8}$。学习率从 $10^{-4}$ 衰减为 $10^{-5}$，余弦学习率计划。总体 epochs 设为 $200$，Batch_size 12，两块 NVIDIA RTX A6000 GPUs。

5.3 结果

定性结果

定量结果

视觉 & 语言重构结果

5.4 消融分析

  Swin-B 视觉 Backbone + BERT 语言 Backbone

BTMAE 的效果

ITA 的效果

  上表 2。

Masking 比例的影响

六、结论

  提出的 BTMAE 解决了现有 RIS 模型的限制，增强了对复杂和混乱的上下文信息的理解 ，大量实验表明 BTMAE 达到了 SOTA 的效果，提升了鲁棒性。

写在后面

  总算是把这篇博文写完了 ，怎么说，思路想法还是可以的，就是这个写作好像有点不太行，需要注意下。