Transformer开始往动态路由的方向前进了！厦大&华为提出TRAR，在VQA、REC上性能SOTA！（ICCV 2021）

写在前面

由于具有优越的全局依赖建模能力，Transformer及其变体已成为许多视觉和语言任务的主要结构。然而，在视觉问答（VQA）和指向性表达理解（REC）等任务中，多模态预测通常需要从宏观到微观的视觉信息。因此，如何动态地调度Transformer中的全局和局部依赖建模成为一个新兴的问题 。

在本文中，作者提出了一个依赖于输入样本的路由方案TRAnsformer routing（TRAR），来解决这个问题。在TRAR中，每个视觉Transformer层都配备了具有不同注意力广度的路由模块，模型可以根据前一个step的输出，动态计算相应的Attention权重，从而为每个样本制定最优路由路径。

此外，通过精心的设计，TRAR可以将额外的计算和显存开销减少到几乎可以忽略的程度。为了验证TRAR，作者在VQA和REC的五个基准数据集上进行了大量实验，取得了比标准Transformer和一系列SOTA方法更高的性能。

2、论文和代码地址

TRAR: Routing the Attention Spans in Transformer for Visual Question Answering

论文地址：

https://openaccess.thecvf.com/content/ICCV2021/papers/Zhou_TRAR_Routing_the_Attention_Spans_in_Transformer_for_Visual_Question_ICCV_2021_paper.pdf

代码地址：https://github.com/rentainhe/TRAR-VQA/

3、Motivation

在自然语言处理领域之后，Transformer也成为许多视觉和语言（V&L）任务的首选结构。许多研究者提出各种多模态网络，在各种基准数据集上实现新的SOTA性能。这些模型的巨大成功很大程度上归功于对自注意（SA）的卓越的全局依赖建模，它不仅可以捕获模态内的关系，而且还可以促进视觉和语言的对齐。

然而，在一些V&L任务中，如视觉问答（VQA）和指向性表达理解（REC），多模态推理通常需要来自不同感受野的视觉注意。如上图所示，模型不仅应该理解整体语义，更重要的是，它还需要捕获局部关系，才能回答正确的答案。在这种情况下，仅依赖SA中的全局依赖关系建模仍然不足以满足这样的需求。

为了解决这个问题，帮助Transformer网络探索不同的注意广度已成为一种新兴需求。一个直观的解决方案是建立一个动态的混合网络，其中每一层都有一组不同感受域的注意模块。然后，该模型可以根据给定的样本选择合适的路径。然而，直接应用该解决方案可能会适得其反，因为额外的参数和计算将进一步加剧模型的计算和显存成本，这已经是Transformer的主要缺点。

在本文中，作者提出了一种新的轻量级路由方案，称为Transformer routing（TRAR），它可以在不增加计算和显存开销的情况下实现Attention的自动选择。具体而言，TRAR为每个视觉SA层配备了一个路径控制器，以根据前一步的输出预测下一个注意广度。

为了解决冗余参数和计算问题，TRAR将SA看做全连通图特征更新函数，并为定义的注意广度构造不同的邻接掩码。之后，可以将模块路径选择的任务定义转换为掩码选择的任务，这在很大程度上降低了额外成本。

为了验证TRAR的有效性，作者将其应用于两个多模态任务（VQA和REC），并在五个基准数据集（VQA2.0、CLVER、RefCOCO、RefCOCO+和RefCOCOg）上进行了大量实验。实验结果表明，TRAR能够在VQA2.0和RefCOCOg数据集上实现SOTA性能。

4、方法

Transformer Routing (TRAR) 方案的框架如上图所示。解码层的不同颜色表示不同跨度的视觉注意。通过TRAR，Transformer可以动态选择每个step的视觉注意广度，从而为每个样本实现最佳的推理路径。

4.1. Routing Process

为了实现每个示例的动态路由目标，直观的解决方案是创建一个多分支网络结构，其中每一层都配备了不同设置的模块。具体来说，基于上一个推理步骤的特征和路由空间，下一推理步骤的输出通过以下方式获得：

其中α是路径控制器预测的路径概率，是一组模块。在测试过程中，可以对α进行二值化以进行hard selection ，也可以保持连续以获得soft routing 。然而，从上式中，我们可以看出，这样的路由方案不可避免地会使网络变得非常复杂，并大大加剧训练开销。

在这种情况下，关键是优化路径路由的定义，以减轻实验负担。标准自注意力的定义为：

我们可以看到，SA可以看作是一个全连通图的特征更新函数，可以看做为加权邻接矩阵。因此，为了获得不同注意广度的特征，只需要限制每个输入元素的图连接。这可以通过在dot-product的结果乘上一个邻接的mask来实现，表示如下：

其中，为二进制的0/1矩阵，如果在目标元素的注意范围内，则设置为1。

基于上面的等式，SA的路由层可以定义为：

其中可以在不同SA层之间共享，从而减小参数量。

然而，上面的公式在计算上仍然是昂贵的。因此，作者进一步将模块选择问题简化为邻接mask 的选择，表示为：

从上式中，我们可以看出，额外的计算和显存占用可以减少到几乎为0，并且从不同的注意广度中选择特征的目标仍然可以实现。

4.2. Path Controller

在TRAR中，每个SA层都具有路径控制器，用于预测路由的概率，即模块选择。具体来说，给定输入特征，路径概率定义如下：

这里，MLP指的是多层感知机，AttentionPool指的是基于注意力的池化方法。

4.3. Optimization

在本文中，作者为TRAR提供了两种类型的推理方法，即软路由（soft routing）和硬路由（hard routing）。

Soft routing

通过应用Softmax函数，可以将路由路径的分类选择作为连续可微的操作。然后可以通过特定任务的目标函数，将Router和整个网络进行端到端的联合优化。在测试过程中，不同注意广度的特征被动态组合。由于软路由不需要额外的调参，因此训练相对容易。

Hard routing

硬路由是实现二进制路径选择，它可以进一步引入特定的CUDA内核，以加速模型推理。然而，分类路由使得Router的权重不可微，直接对软路由的结果进行二值化可能导致训练和测试之间的特征Gap。为了解决这个问题，作者引入了Gumbel-Softmax Trick来实现可微的路径路由：

4.4. Network Structure

作者在MCAN上构建了动态路由网络。具体而言，与标准Transformer类似，MCAN有六个编码层，用于对LSTM提取的语言特征进行建模，还有六个解码层，用于同时处理视觉特征和跨模态对齐。作者用本文提出的路由模块替换视觉SA模块，如上图所示。

5、实验

5.1. Ablations

上表展示了REC和VQA任务上本文方法的消融实验结果，可以看出，权重共享的Router会带来一定程度的性能下降，Soft Routing性能比Hard Routing性能好一点。

上表展示了TRAR路由空间的影响，从上表可以看出，加入局部注意力掩码，可以显著提高性能，加入更高阶的掩码，性能提升不大。

5.2. Comparison with SOTA

上表展示了，基于VQA2.0数据集，本文方法在VQA任务上和SOTA方法的性能对比。

上表展示了，基于CLEVR数据集，本文方法和其他方法的性能对比。

上表展示了在三个数据集上，基于REC任务，本文方法和SOTA方法的性能对比。

5.3. Qualitative Analysis

上图展示了用硬路由（）和软路由（）时，TRAR路由熵的变化。

上图展示了本文方法和baseline的Attention对比，以及不同任务上Attention的可视化结果。

6、总结

在本文中，作者研究了Transformer在两个视觉-语言任务（即VQA和REC）中的依赖关系建模。这两项任务通常需要来自不同感受野的视觉注意，而标准Transformer无法完全处理这两项任务。

为此，作者提出了一种轻量级的路由方案Transformer routing（TRAR），以帮助模型为每个样本动态选择注意广度。TRAR将模块选择问题转化为选择注意掩码的问题，从而使额外的计算和显存开销可以忽略不计。为了验证TRAR的有效性，作者在5个基准数据集上进行了大量实验，实验结果证实了TRAR的优越性。