RIS 系列 Mask Grounding for Referring Image Segmentation 论文阅读笔记

  一篇 Arxiv 上面的新文章，看看清华大佬们的研究。

• 论文地址：Mask Grounding for Referring Image Segmentation
• 代码地址：原论文说将会开源，静待佳音~
• 预计提交于：CVPR 2024
• Ps：2023 年每周一篇博文阅读笔记，主页 更多干货，欢迎关注呀，期待 6 千粉丝有你的参与呦~

一、Abstract

  Referring Image Segmentation (RIS) 的定义，目前的 SOTA 方法仍然存在像素和词水平上的语言-图像模态鸿沟。主要原因：通常依赖于句子级别的语言特征用于语言-图像对齐；缺乏对细粒度视觉定位的监督。另外，由于弱的视觉和语言特征间的关联，因此需要更有效的推理去理解那些包含多个目标的复杂场景。于是本文引入 Mask Grounding 辅助任务来提升视觉定位的性能，Mask Grounding 直接适用于之前的模型。此外，为全面解决模态鸿沟，设计了一种跨模态对齐损失和一种辅助对齐模块。在 MagNet(Mask-grounded Network) 上达到了 SOTA 的效果。

二、引言

  首先指出图像分割任务与 Referring Image Segmentation 的区别，RIS 的应用。然后就是 RIS 的挑战：如何减少语言和图像特征间的模态鸿沟？需要一个有效性的对齐方法。

  如上图所示，之前的方法主要关注于设计不同的损失函数或者引入网络结构/模块来促进对齐，然而缺陷有俩：往往依赖于句子级别的语言特征对齐；缺乏对细粒度视觉定位的显著训练监督。于是难以处理复杂的目标间关系或者包含很少或混乱上下文的子句。如下图所示：

  于是文本引入一种 Mask Grounding 的辅助任务用于显式地教会模型进行细粒度的对齐。具体来说，在训练过程中，模型随机 mask 掉一些文本词汇，并且让模型来预测这些词汇的实体信息。除了整合文本上下文信息外，还利用了视觉和分割的信息。

  除 Mask Grounding 外，还提出一种跨模态对齐损失和一个对齐模块来全面填补模态鸿沟。整合的模型 MagNet (Mask-grounded Network) 达到了 SOTA 的效果。主要贡献如下：

• 突出了最近 SOTA 的 RIS 方法的缺陷，指出细粒度视觉定位的缺乏；
• 引入 Mask Grounding 辅助任务，旨在增强细粒度的视觉定位算法；
• Mask Grounding + 跨模态对齐损失 + 辅助对齐模块 = MagNet (Mask-grounded Network)，实现了新的 SOTA。

三、相关工作

Architecture Design for RIS

  早期的方法遵循拼接-卷积的操作，后续的工作采用 RNN 或动态卷积的方法。还有一些方法设计出语言-图像融合模块。此外，还有一些工作利用已知的语言结构或目标关系来增强融合。随着注意力结构的成功，当前的工作通常采用无向或双向的交叉注意力模块来执行语言-图像融合。有一些工作使用元学习的方法用于 RIS。受到大语言模型的驱动，一些方法将 RIS 视为自回归向量生成问题。接下来是一些举例：VPD、ReLA、DMMI。这些方法的缺陷在于：总是期待语言-图像对齐发生在 mask 预测过程中。于是本文引入一种辅助任务来显式地对齐语言-图像特征。

Loss Design for RIS

  早期用于训练 RIS 的方法通常采用简单的 binary cross entropy 损失， 接下来是对比学习的损失。与之前使用全局池化的语言特征计算损失相比，本文关注在像素-词水平上学习细粒度的目标联系。

Masked Language Modeling

  Masked language modeling (MLM) 在自然语言处理中很普遍。首先引入的是 BERT，之后就成为预训练模型及视觉-语言模型的标配。最近的 MaskedVLM 采用 MLM 执行 mask vision and language 建模，在一个模态的辅助作用下重建另外一个模态。Mask Grounding 与其不同，通过使用外部 mask 信号直接匹配缺失的单词从而确保重建过程，学习到相应的细粒度视觉定位的信息。

四、方法

4.1 结构

  MagNet (Mask-grounded Network) 由三个模块组成，首先 Mask Grounding 旨在提升细粒度视觉定位性能，使用视觉线索、语言上下文、分割信息来教会模型预测 masked 文本 tokens。之后 Cross-modal Alignment Module (CAM) 用于微调语言和图像特征间的双向交互。最后 Crossmodal Alignment Loss (CAL) 监督像素-像素和像素-文本的对齐。

4.2 Mask Grounding

  如图 3 所示，给定输入图像及其对应的指代表达式和分割 mask。首先利用一些特定词汇来代替指代表达式中的一些 tokens，从而训练模型来预测这些 tokens（类似于 MLM）。具体来说，首先取得 mask 区域的中心坐标，然后通过一个 2 层的 MLP，将其 mask 编码到一个 mask embedding 中。同时采用一个线性层将语言 embedding 投影到与图像 embedding 相同的维度。然后应用提出的 Masked Token 预测器在注意力机制的作用下来处理所有拼接的 embedding 用于 masked token 预测。最后，使用一个 cross-entropy 损失 ${\mathcal{L}}_{\text{grounding}}$ 来比较最终的预测分布与目标分布。数学公式描述如下：令 $\mathbf{T},\mathbf{I},\mathbf{M}$ 分别表示语言编码器、图像编码器、mask 编码器的输入：
$$\begin{gathered} \mathbf{O}=\text{Language}\mathrm{Encoder}(\mathrm{Mask}(\mathbf{T})) \\ \mathbf{P}=\mathrm{Image}\text{Encoder}(\mathbf{I}) \\ \mathbf{C}=\text{MaskEncoder}(\mathbf{M}) \\ {\mathcal{L}}_{\mathrm{grounding}}={\mathcal{L}}_{\mathrm{CE}}({\mathbf{y}}_{\mathrm{gt}},\text{Predictor}(\mathrm{Concat}([\mathbf{O},\mathbf{P},\mathbf{C}]) \end{gathered}$$其中预测器为类似 BERT 编码器的结构，$M$ 为 GT masks 的中心坐标，${\mathbf{y}}_{\mathrm{gt}}$ 为 masked token 的标签，${\mathcal{L}}_{\mathrm{CE}}$ 为交叉熵损失。实验中设置 Swin-B 为图像编码器，BERT-base 为语言编码器，但方法不限于此。

讨论

  如上表所示，Mask Grounding 超越了标准的 masked language modeling (MLM) 和 masked-vision language modeling (MaskedVLM)。原因在于：模型整合：传统的 MLM 为单模态，缺乏了指代表达式及其匹配的视觉目标的联系，而 MaskedVLM 为多模态，Mask Grounding 能够超越的目的在于引入了额外的 masking 信号来对齐 masked words 和匹配的视觉目标。这一结果表明 词-目标联系和细粒度的视觉定位很重要；任务属性：MLM 和 MaksedVLM 作为一般的预训练任务，需要在下游任务上进行微调，而 Mask Grounding 设计于一个 RIS 的辅助任务，在训练阶段增强了细粒度的视觉定位性能，且不需要额外的微调；预测上下文：MLM 和 MaskedVLM 采用文本或文本-视觉上下文预测，而 Mask Grounding 整合了外部的分割信息，于是性能更好。

4.3 跨模态对齐模块

  如上图所示，提出的跨模态对齐、 cross-modal alignment module (CAM) 将全局上下文先验注入到图像特征中，再进行跨模态融合。CAM 首先采用不同窗口尺寸的池化操作生成 $K$ 个不同尺度的特征图构成特征金字塔。然后，每个特征图将会通过一个 3 层的 MLP 用于提取全局特征。之后所有的输出特征将通过双线性插值上采样到原始特征图，然后沿着特征维度拼接。同样采用一个门控单元来调制最后的输出。最终，输出后的特征返回到输入特征上用于下一阶段图像或语言编码器的输入。将语言编码器划分为 4 个阶段，并在每个阶段的末尾添加 CAM 模块。

  用数学公式表示如下：令 ${\mathbf{T}}_i$ 和 ${\mathbf{I}}_i$ 分别表示语言和图像编码器每个阶段的输入，于是每个阶段有：
$$\begin{array}{c}{\mathbf{O}}_i=\text{LanguageStage}({\mathbf{T}}_i),{\mathbf{P}}_i=\mathrm{Image}\text{Stage}({\mathbf{I}}_i)\\ {\mathbf{P}}_i^k={\mathrm{MLP}}_k({\mathrm{Pool}}_k({\mathbf{P}}_i))\\ {}_{,p2t}^K,{\mathbf{P}}_{i,t2p}^k={\mathrm{X}-\mathrm{MHA}}_k({\mathbf{O}}_i,{\mathbf{P}}_n^k)\\ {\mathbf{O}}_{i,p2t}=\mathrm{Concat}([{\mathbf{O}}_{i,p2t}^i,...,{\mathbf{O}}_{i,p2t}^N]\\ {}_{2p}=\mathrm{Concat}([\mathrm{Up}({\mathbf{P}}_{i,t2p}^1,...,\mathrm{Up}({\mathbf{P}}_{i,t2p}^N)]\\ {\mathbf{O}}_{i+1}={\mathbf{O}}_i+\tanh (\mathsf{MLP}({\mathbf{O}}_{i,p2t}))\\ {\mathbf{P}}_{i+1}={\mathbf{P}}_i+\tanh (\mathsf{MLP}({\mathbf{P}}_{i,t2p}))\end{array}$$其中 $Up$ 表示上采样，X-MHA 表示双向跨模态多头注意力。

4.4 跨模态对齐损失

  采用跨模态对齐损失来对齐语言和图像特征，其中 cross-modal alignment loss (CAL) 全面总结了像素-像素级别的损失 ${\mathcal{L}}_{\text{P2P}}$ 和像素-文本损失 ${\mathcal{L}}_{\text{P2T}}$。用数学公式表示如下：给定语言编码器产生的语言特征 $\mathbf{T}\in {\mathbb{R}}^{M\times D}$，包含 $|\mathcal{P}|$ 个正样本像素特征的最终像素解码器 mask 特征 $\mathbf{I}\in {\mathbb{R}}^{C_l\times H_l\times W_l}$，$|\mathcal{N}|$ 个负样本像素特征。${\mathbf{I}}_i^{+}$ 表示正样本集合 $\mathcal{P}$ 中的第 $i^{th}$ 个像素特征，$I_j^{-}$ 表示负样本集合 $\mathbb{N}$ 中的第 $j^{th}$ 个像素特征，${\mathbf{T}}_k$ 表示第 $k^{th}$ 个语言特征，然后有：
$$\begin{array}{r}{\mathcal{L}}_{\mathrm{CAL}}={\mathcal{L}}_{P2P}+{\mathcal{L}}_{P2T}\\ {\mathcal{L}}_{P2P}=-\frac{1}{|\mathcal{P}|}\sum _i^{|\mathcal{P}|}\frac{e^{{\mathbf{I}}_i^{+}\cdot {\mathbf{I}}_{\mathbf{avg}}^{+}/{\tau }_1}}{e^{{\mathbf{I}}_i^{+}\cdot {\mathbf{I}}_{\mathbf{avg}}^{+}/{\tau }_1}+{\sum }_j^{|\mathcal{N}|}{e}^{{\mathbf{I}}_i^{+}\cdot {\mathbf{I}}_j^{-}/{\tau }_1}}+-\frac{1}{|\mathcal{N}|}\sum _j^{|\mathcal{N}|}\frac{e^{{\mathbf{I}}_j^{-}\cdot {\mathbf{I}}_{\mathrm{avg}}^{-}/{\tau }_1}}{e^{{\mathbf{I}}_j^{-}\cdot {\mathbf{I}}_{\mathrm{avg}}^{-}/{\tau }_1}+{\sum }_i^{|\mathcal{P}|}{e}^{{\mathbf{I}}_j^{-}\cdot {\mathbf{I}}_i^{+}/{\tau }_1}}\\ {\mathcal{L}}_{P2T}=-\frac{1}{|\mathcal{P}|}\sum _i^{|\mathcal{P}|}\frac{e^{{\mathbf{I}}_i^{+}\cdot {\mathbf{T}}_{\mathrm{avg}}/{\tau }_2}}{e^{{\mathbf{I}}_i^{+}\cdot {\mathbf{T}}_{\mathrm{avg}}/{\tau }_2}+{\sum }_j^{|\mathcal{N}|}{e}^{{\mathbf{I}}_i^{+}\cdot {\mathbf{I}}_j^{-}/{\tau }_2}}\end{array}$$其中 ${\mathbf{I}}_{\mathrm{avg}}^{+}=\frac{1}{|\mathcal{P}|}{\sum }_i^{|\mathcal{P}|}{\mathbf{I}}_i^{+}$ 和 ${\mathbf{I}}_{\mathrm{avg}}^{-}=\frac{1}{|\mathcal{N}|}{\sum }_j^{|\mathcal{N}|}{\mathbf{I}}_j^{-}$ 分别表示平均池化后正样本像素特征和负样本像素特征。${\mathbf{T}}_{\mathrm{avg}}=\mathrm{proj}(\frac{1}{M}{\sum }_m^M{\mathbf{T}}_k)$ 为平均池化和线性投影后的词特征，${\tau }_1$ 和 ${\tau }_2$ 为超参数。需要注意的是所有的特征在进行点乘前均经过 L2 归一化处理，但并未在上式中体现。

4.5 损失函数

  损失函数为下列 4 个不同损失的加权求和：
$$\begin{array}{c}\mathcal{L}={\lambda }_{\mathrm{BCE}}{\mathcal{L}}_{\mathrm{BCE}}+{\lambda }_{\mathrm{Dice}}{\mathcal{L}}_{\mathrm{Dice}}+{\lambda }_{\mathrm{CAL}}{\mathcal{L}}_{\mathrm{CAL}}+{\lambda }_{\mathrm{grounding}}{\mathcal{L}}_{\mathrm{grounding}},\end{array}$$实验中 ${\lambda }_{\mathrm{BCE}}=2.0$，${\lambda }_{\mathrm{Dice}}=2.0$，${\lambda }_{\mathrm{grounding}}=1.0$ 。

五、实验

5.1 数据集及评估指标

• 数据集：RefCOCO、RefCOCO+、GRef
• 评估指标：overall intersection-over-union (oIoU)、mean intersection-overunion (mIoU)

5.2 主要结果

5.3 可视化

5.4 消融研究

  训练 10 个 epoch，输入图像尺寸 $224\times 224$。所有的消融实验执行在 RefCOCO 和 RefCOCO+ 数据集上。

• RIS 性能的影响;
• Mask 编码器的设计；
• Mask Token 预测器的设计；
• Mask Grounding 的统一效果。
• CAM 的有效性
• CAL 的有效性

语言-图像对齐的影响

MagNet 组件的兼容性

六、结论

  本文提出 Mask Grounding，基于周围的文本、视觉和分割信息，通过教导模型预测随机 mask 掉的文本 tokens，实验效果很好。为全面解决模态鸿沟，设计了一种跨模态对齐损失和一种辅助对齐模块。当一齐作用时，提出的 MagNet 实现了 SOTA 的性能。

写在后面

  这篇论文咋说呢，感觉就是那种顶会的边缘，创新点属于可拒可不拒的那种。最大的败笔还是论文的写作确实不咋地，没有一种连贯之感。另外，实验缺少了实验细节的介绍，放在补充材料中吗？