[论文翻译] Visual Saliency Transformer

(2021.8.4)注意：本文仍正在施工中，实验以及RGB-D SOD部分尚未翻译
论文地址：https://arxiv.org/abs/2104.12099
代码：https://github.com/nnizhang/VST
发表于：ICCV 2021(Arxiv 2021.04)

Visual Saliency Transformer

Abstract

最近，大规模的显著性检测方法通过依靠基于CNN的架构取得了可喜的成果。作为一种选择，我们从无卷积的seq2seq角度重新思考这一任务，并通过对长距离依赖关系进行建模来预测显著性，这是卷积所不能实现的。具体来说，我们开发了一个基于纯transformer的新型统一模型，记为Visual Saliency Transformer(VST)，用于RGB和RGB-D显著目标检测(SOD)。其将图像块(image patch)作为输入，利用transformer在图像块之间传播全局上下文。除了在Vision Transformer(ViT)中使用的传统transformer架构，我们利用多级token融合，在transformer框架下提出了一个新的token上采样方法，以获得高分辨率的检测结果。我们还开发了一个基于token的多任务解码器，通过引入任务相关的token和一个新的patch-task注意力机制，以同时进行显著性和边界检测。实验结果表明，我们的模型在RGB和RGB-D SOD基准数据集上都优于现有的SOTA。更重要的是，我们的整个框架不仅为SOD领域提供了一个新的视角，而且还为基于transformer的密集预测模型展示了一个新的范式。

I. Introduction

SOD致力于在观看图像时检测和分割显著的物体或区域。除了在RGB图像上执行SOD之外，RGB-D SOD近年来也获得了越来越多的关注，因为深度数据可以为显著性检测提供有用的空间结构信息。

目前SOD的SOTA以卷积架构为主，在RGB与RGB-D上都是如此。它们通常采用编码器-解码器CNN架构，其中编码器将输入图像编码为多级特征，解码器整合提取到的特征来预测最终的显著图。基于这种简单的架构，大多数研究重心都在建立一个强大的解码器来预测更好的显著性结果。为此，他们引入了各种注意力模型、多尺度特征融合方法和多任务学习框架。对RGB-D SOD的一个额外要求是有效地融合跨模态信息，即外观信息与深度线索。之前的工作提出了各种模态融合方法，如特征融合、知识提炼、动态卷积、注意力模型和图神经网络。因此，基于CNN的SOD方法已经取得了令人印象深刻的结果。

然而，所有先前的方法在学习全局长距离依赖(global long-range dependency)方面都是有限的。长期以来，全局上下文和全局对比已经被证明对于显著性检测至关重要。然而，由于CNN的固有局限性，它们在局部滑动窗口中提取特征，因此先前的方法很难利用关键的全局线索。尽管有些方法利用全连接层、全局池化层和非局部(non-local)模块来纳入全局上下文，但它们只在某些层中进行这种操作，而基于CNN的标准架构仍然没有改变。

最近，Transformer被提出来为机器翻译的词序列之间的全局长距离依赖关系建模。Transformer的核心是自注意力(self attention)机制，它利用查询键的关联性来联系序列中的不同位置。Transformer在编码器和解码器中多次堆叠自注意层，因此可以在每一层中建模长距离依赖关系。因此，将Transformer引入SOD是很自然的，可以一直利用模型中的全局线索。

在本文中，我们首次从一个新的seq2seq角度重新思考SOD，并为RGB和RGB-D SOD开发了一个基于纯transformer的新型统一模型，记为Visual Saliency Transformer。我们遵循最近提出的ViT模型，将每个图像划分为图像块，并在块序列上应用transformer模型。然后，transformer在图像块之间传播长距离依赖关系，而不需要使用卷积。不过，将ViT应用于SOD并不简单。一方面，如何在纯transformer的基础上执行密集预测任务仍然是一个开放的问题。另一方面，ViT通常将在一个非常粗糙的尺度上对图像划分token。如何使ViT适应SOD的高分辨率预测需求也是不清楚的。

为了解决第一个问题，我们设计了一个基于token的transformer解码器，通过引入任务相关的token来学习决策embedding。然后，我们提出了一个新的patch-task-attention机制，以生成密集预测结果，这为在密集预测任务中使用transformer提供了一个新的范例。受之前利用边缘检测提高SOD性能的模型启发，我们建立了一个多任务解码器，通过引入一个显著性token和一个边界token来同时进行显著性和边界检测。这一策略通过简单地学习与任务相关的token简化了多任务预测的工作流程，从而在很大程度上降低了计算成本，同时获得了更好的结果。为了解决第二个问题，受减少token长度的Token-to-Token(T2T)变换的启发，我们提出了一个新的反向T2T变换，通过将每个token扩展为多个子toekn来对token进行上采样。然后，我们逐步对块token进行上采样，并将其与低级token融合，以获得最终的全分辨率显著图。此外，我们还使用了一个跨模态transformer来深入探索RGB-D SOD的多模态信息之间的互动。最终，我们的VST在参数量和计算成本相当的情况下，在RGB和RGB-D数据上都优于现有SOTA。

本文的主要贡献总结如下：

• 我们首次从seq2seq建模的新角度，为RGB和RGB-D SOD设计了一个基于纯transformer架构的新型统一模型
• 我们设计了一个多任务transformer解码器，通过引入任务相关的token以及patch-task-attention，共同进行显著性和边缘检测
• 我们为基于transformer的框架提出了一种新的token上采样方法
• 我们提出的VST模型在RGB和RGB-D SOD基准数据集上取得了SOTA，这证明了其有效性和基于transformer的SOD模型的潜力

II. Related Work

Deep Learning Based SOD

基于CNN的方法已经成为RGB和RGB-D SOD的主流趋势，并取得了良好的性能。大多数方法通过使用UNet或HED风格的网络结构来利用多级特征融合策略。一些工作引入了注意力机制来学习更多的鉴别性特征，包括空间和通道注意力或像素级的上下文注意力。其他工作试图设计递归网络来逐步完善显著图。此外，一些作品引入了多任务学习，例如固定点预测(fixation prediction)、图像字幕(image caption)和边缘检测以提高SOD的性能。

对于RGB-D SOD，许多方法都设计了各种模型来融合RGB和深度特征，并取得了显著的效果。一些模型采用了简单的特征融合方法，即concatenation、summation或multiplication。其他模型则利用深度线索来生成空间注意力或通道注意力来增强RGB特征。Pang等人引入了动态卷积，Li等人引入了仿射变换参数来增强RGB特征。Piao等人提出了一个深度蒸馏器(depth distiller)，将深度知识转移到RGB流中。Luo等人引入了图神经网络来建模多模态数据之间的关系。此外，一些工作还利用交叉注意力机制来传播RGB和深度线索之间的长距离跨模态互动。

与以往基于CNN的方法不同，我们首次从seq2seq的角度重新思考SOD，并为RGB和RGB-D SOD提出了一个基于纯transformer的统一模型。在我们的模型中，我们遵循一些CNN方法，利用边缘检测来提高SOD的性能。然而，与这些基于CNN的模型不同，我们设计了一个新颖的基于token的多任务解码器，以在transformer框架下实现这一目标。

Transformers in Computer Vision

Vaswani等人首次提出了用于机器翻译的transformer编码器-解码器架构，其采用了多次叠加的multi-head自注意力机制以及和point-wise前向传播层。最近，越来越多的工作将transformer模型引入到各种计算机视觉任务中，并取得了优异的成绩。一些工作将CNN和transformer结合到混合架构中，用于目标检测、全景分割(panoptic segmentation)、车道形状预测、多目标追踪等。通常情况下，他们首先使用CNN来提取图像特征，然后利用transformer来引入长距离的依赖关系。

其他作品设计了纯transformer模型，从seq2seq的角度处理图像。ViT将每张图像划分为一串展平的二维patch，然后采用transformer进行图像分类。Touvron等人引入了teacher-student策略来提高ViT的数据效率。Yuan等人提出了T2T模块对局部结构进行建模，并提出了一种高效的deep-narrow ViT结构。在本文中，我们采用T2T-ViT作为backbone，并提出了一个新颖的多任务解码器和一个反向的T2T token上采样方法。值得注意的是，我们对任务相关token的使用与以前的模型不同。在以往的工作中，通过将token嵌入多层感知器，直接将类token用于图像分类。然而，我们不能直接从单个任务token中获得密集预测结果。因此，我们提出在patch token和task token之间进行patch-task-attention，以预测显著图和边界图。我们相信该策略也会给未来其他密集预测任务的transformer模型带来启发。

有一些与我们类似的工作，它将transformer引入了语义分割任务。作者采用了一个vision transformer作为backbone，然后将token序列reshape为二维图像特征。然后，他们使用卷积和双线性上采样来预测全分辨率的分割图。他们的模型仍然属于混合架构的范畴。相比之下，我们的模型是一个纯transformer架构，不依赖任何卷积操作和双线性上采样。

III. Visual Saliency Transformer

图1. 我们所提出的用于RGB&RGB-D SOD的VST模型的整体架构。虚线框住的表示用于RGB-SOD的模块。

图1展示了我们提出的VST模型的整体架构。主要组件包括一个基于T2T-ViT的transformer encoder，一个将patch token从编码器空间转换到解码器空间的变换器transformer convertor，以及一个多任务transformer decoder。

III.A Transformer Encoder

与其他基于CNN的SOD方法类似，这些方法通常利用预训练的图像分类模型，如VGG和ResNet作为其编码器backbone来提取图像特征。在本文中，我们采用预训练的T2T-ViT模型作为我们的backbone，具体内容如下。

☆ Tokens to Token

给定一个来自上一层的长度为$l$的patch token序列${\mathbf{T}}^{\prime }$，T2T-ViT迭代地使用T2T模块，该模块由re-structurization和soft split这两步组成，对${\mathbf{T}}^{\prime }$中的局部结构信息进行建模，并得到一个新token序列。

图2. T2T模块以及我们所提出的RT2T模块的细节

Re-structurization 如图2(a)所示，token ${\mathbf{T}}^{\prime }$首先用一个transformer层进行转换，得到新的token $\mathbf{T}\in {\mathbb{R}}^{l\times c}$：$$\mathbf{T}=\mathrm{M}\mathrm{L}\mathrm{P}\left(\mathrm{MSA}\left({\mathbf{T}}^{\prime }\right)\right)$$ 其中$MSA$和$MLP$分别表示vanilla Transformer中的multi-head self-attention和multilayer perceptron。需要注意的是，在每个block前都要进行layer normalization处理。然后，$\mathbf{T}$被reshape为二维图像$\mathbf{I}\in {\mathbb{R}}^{h\times w\times c}$，其中$l=h\times w$，以恢复空间结构，如图2(a)所示。

Soft split 在re-structurization之后，$\mathbf{I}$首先被分割成$k\times k$的patch，此时使用了$s$的overlapping以及$p$的zero padding来填充图像边界。然后，图像patch被unfold为一个token序列${\mathbf{T}}_o\in {\mathbb{R}}^{l_o\times c{k}^2}$，其中序列长度$l_o$计算如下：$$l_o=h_o\times w_o=\lfloor \frac{h+2p-k}{k-s}+1\rfloor \times \lfloor \frac{w+2p-k}{k-s}+1\rfloor$$ 与vanilla ViT不同，T2T-ViT中采用的overlapped patch splitting引入了相邻patch内的局部对应关系，从而带来了空间先验知识。

T2T转换可以多次迭代进行。在每次迭代中，步骤re-structurization首先将之前的token embedding转化为新的embedding，同时整合所有token内的长距离依赖关系。然后，soft split操作将每个$k\times k$范围内的token融合成一个新的token，准备用于下一层。此外，当参数设置为$s<k-1$时，token的长度可以逐步减少。

我们按照T2T-ViT的做法，首先将输入图像软分割成patch，然后两次采用T2T模块。在这三个软分割步骤中，patch大小被设定为$k=[7，3，3]$，overlapping被设定为$s=[3，1，1]$，padding大小被设定为$p=[2，1，1]$。因此，我们可以得到多级token ${\mathbf{T}}_1\in {\mathbb{R}}^{l_1\times c}$，${\mathbf{T}}_2\in {\mathbb{R}}^{l_2\times c}$，${\mathbf{T}}_3\in {\mathbb{R}}^{l_3\times c}$。给定输入图像的高和宽分别为$H$和$W$，有$l_1=\frac{H}{4}\times \frac{W}{4}$，$l_2=\frac{H}{8}\times \frac{W}{8}$，$l_3=\frac{H}{16}\times \frac{W}{16}$。我们遵循T2T-ViT的原始实现，设置$c=64$，并在T3上使用线性投影层，将其embedding dimension从$c$转化为$d=384$。

☆ Encoder with T2T-ViT Backbone

最后的token序列T3被加入sinusoidal position embedding以编码二维的位置信息。然后，$L^{\mathcal{E}}$ transformer layer被用来建模T3内部的长距离依赖关系，以提取强大的patch token embedding ${\mathbf{T}}^{\mathcal{E}}\in {\mathbb{R}}^{l_3\times d}$。

对于RGB SOD，我们采用单个transformer encoder，从每个输入的RGB图像中获得RGB编码器patch token ${\mathbf{T}}_r^{\mathcal{E}}\in {\mathbb{R}}^{l_3\times d}$。对于RGB-D SOD，我们遵循two-stream架构，进一步使用另一个transformer encoder，以类似的方式从输入深度图中提取深度编码器patch token ${\mathbf{T}}_d^{\mathcal{E}}$，如图1所示。

III.B Transformer Convertor

我们在transformer encoder和decoder之间插入一个convertor模块，将编码器的patch token ${\mathbf{T}}_{\ast }^{\mathcal{E}}$从编码器空间转换到解码器空间，得到转换后的patch token ${\mathbf{T}}^{\mathcal{C}}\in {\mathbb{R}}^{l_3\times d}$。

☆ RGB-D Convertor

Cross-modality-attention

☆ RGB Convertor

为了与我们的RGB-D SOD模型保持一致，对于RGB SOD，我们只需在${\mathbf{T}}_r^{\mathcal{E}}$上使用${\mathbf{L}}^{\mathcal{C}}$标准transformer层来获得转换后的patch token sequence ${\mathbf{T}}^{\mathcal{C}}$。

III.C Multi-task Transformer Decoder

我们的解码器旨在将patch token ${\mathbf{T}}^{\mathcal{C}}$解码为显著图。因此，我们提出了一种新型的token上采样方法，并采用多级token融合以及一个基于token的多任务解码器。

☆ Token Upsampling and Multi-level Token Fusion

我们认为，直接从${\mathbf{T}}^{\mathcal{C}}$预测显著图不能获得高质量的结果，因为${\mathbf{T}}^{\mathcal{C}}$的长度相对较小($l_3=\frac{H}{16}\times \frac{W}{16}$)，这对密集预测来说是不足的。因此，我们认为应该先对patch token进行上采样，然后再进行密集预测。大多数基于CNN的方法采用双线性上采样来恢复大尺度特征图。作为一种替代方案，我们在transformer框架下提出了一种新的token上采样方法。受T2T模块逐步融合相邻token以减少token长度的启发，我们提出了一种反向T2T(RT2T)变换，通过将每个token扩展为多个子token来对token进行上采样，如图2(b)所示。

具体来说，我们首先对输入的patch token进行投影，将其embedding dimension从$d=384$减少到$c=64$。然后，我们使用另一个线性投影，将embedding dimension从$c$扩大到$c{k}^2$。接下来，类似于T2T中的软分割，每个token可以被看作是一个$k\times k$的image patch，相邻patch的有$s$的overlapping。然后，我们可以使用$p$的zero-padding将token串reshape成一个图像。输出图像的大小可以用反过来的相应公式来计算，即给定输入patch token的长度为$h_o\times w_o$，输出图像的空间大小为$h\times w$。最后，我们将图像reshape至大小为$l_o\times c$的上采样token，其中$l_o=h\times w$。通过设置$s<k-1$，RT2T变换可以逐渐增大token的长度。受T2T-ViT的启发，我们使用了三次RT2T，并设置$k=[3，3，7]$，$s=[1，1，3]$，$p=[1，1，3]$。因此，patch token的长度可以逐渐上采样为$H\times W$，相当于输入图像的原始尺寸。

此外，受现有SOD方法中被广泛证明的多级特征融合的成功经验的启发，我们利用来自T2T-ViT编码器中长度较大的低级token(即${\mathbf{T}}_1$和${\mathbf{T}}_2$)，来提供准确的局部结构信息。对于RGB和RGB-D SOD，我们均只使用来自RGB变换器transformer的低级token。具体来说，我们通过concatenation和linear projection将${\mathbf{T}}_2$和${\mathbf{T}}_1$与上采样的patch token逐步融合。然后，我们采用一个transformer层来获得第i层的解码器token ${\mathbf{T}}_i^{\mathcal{D}}$，其中$i=2,1$。 整个过程可被表述为：$${\mathbf{T}}_i^{\mathcal{D}}=\mathrm{MLP}(\mathrm{MSA}\left(\mathrm{Linear}\left(\left[\mathrm{R}\mathrm{T}2\mathrm{T}\left({\mathbf{T}}_{i+1}^{\mathcal{D}}\right),{\mathbf{T}}_i\right]\right)\right)$$ 其中$[,]$表示沿着token embedding dimension进行concatenation。"Linear"是指linear projection，将concatenation后的embedding dimension减少到$c$。最后，我们使用另一个linear projection，将${\mathbf{T}}_i^{\mathcal{D}}$的embedding dimension恢复到$d$。

☆ Token Based Multi-task Prediction

受现有的纯transformer方法的启发，这些方法在patch token sequence上添加了一个class token，用于图像分类。因此，我们也利用任务相关的token来预测结果。然而，我们不能像T2T-ViT那样，直接利用MLP对task token embedding进行处理来获得密集预测结果。因此，我们提出在patch token和task related token之间进行patch-task-attention，以进行SOD。

此外，受SOD模型中广泛使用的边缘检测的启发，我们还采用了多任务学习策略来联合进行显著性和边缘检测，以利用后者来提升前者的性能。

为此，我们设计了两个任务相关的token，即一个显著性token ${\mathbf{t}}_s\in {\mathbb{R}}^{1\times d}$和一个边缘token ${\mathbf{t}}_b\in {\mathbb{R}}^{1\times d}$。在每个解码器层$i$，我们在patch token sequence ${\mathbf{T}}_i^{\mathcal{D}}$上添加显著性和边缘token ${\mathbf{t}}_s$和${\mathbf{t}}_b$，然后用${\mathbf{L}}^{\mathcal{D}}$ transformer层处理它们。因此，这两个task token可以从与patch token的互动中学习与图像相关的task-related embedding。之后，我们将更新的patch token作为输入，并执行token上采样和多级融合过程，得到上采样的patch token ${\mathbf{T}}_{i-1}^{\mathcal{D}}$。接下来，我们在下一层$i-1$中重新使用更新的${\mathbf{t}}_s$和${\mathbf{t}}_b$，以进一步更新它们和${\mathbf{T}}_{i-1}^{\mathcal{D}}$。我们重复这个过程，直到我们达到最后一个解码器，此时尺寸为$\frac{1}{4}$。

对于显著性和边缘预测，我们在最终解码器的patch token ${\mathbf{T}}_1^{\mathcal{D}}$与显著性和边缘标记${\mathbf{t}}_s$和${\mathbf{t}}_b$之间进行patch-task-attention。对于显著性预测，我们首先将${\mathbf{T}}_1^{\mathcal{D}}$ embedding至Query $Q_s^{\mathcal{D}}\in {\mathbb{R}}^{1\times d}$，将${\mathbf{t}}_s$ embedding至Key ${\mathbf{K}}_s\in {\mathbb{R}}^{1\times d}$以及Value ${\mathbf{V}}_s\in {\mathbb{R}}^{1\times d}$。同样，对于边缘预测，我们将${\mathbf{T}}_1^{\mathcal{D}}$ embedding至$Q_b^{\mathcal{D}}$，将${\mathbf{t}}_b$ embedding至${\mathbf{K}}_b$和${\mathbf{V}}_b$。然后，我们采用patch task attention来获得任务相关的patch token：$$\begin{array}{rl}{\mathbf{T}}_s^{\mathcal{D}}& =\mathrm{sigmoid}\left({\mathbf{Q}}_s^{\mathcal{D}}{\mathbf{K}}_s^{\top }/\sqrt{d}\right){\mathbf{V}}_s+{\mathbf{T}}_1^{\mathcal{D}}\\ {\mathbf{T}}_b^{\mathcal{D}}& =\mathrm{sigmoid}\left({\mathbf{Q}}_b^{\mathcal{D}}{\mathbf{K}}_b^{\top }/\sqrt{d}\right){\mathbf{V}}_b+{\mathbf{T}}_1^{\mathcal{D}}\end{array}$$ 在这里，我们使用sigmoid激活函数来计算attention，因为在每个等式中我们只有一个key。

由于${\mathbf{T}}_s^{\mathcal{D}}$和\boldsymbol{T}_{b}^{\mathcal{D}}是$\frac{1}{4}$尺寸的，我们采用第三个RT2T变换将它们上采样到全分辨率。最后，我们应用两个带有sigmoid激活函数的线性变换，将它们投影到$[0,1]$中的标量，然后将它们分别reshape为二维显著图和二维边缘图。整个过程在图1中给出。

IV. Experiments