【AAAI2022】可解释性ViT登场，谷歌AI提出层次嵌套Transformer模型

论文链接： https://www.zhuanzhi.ai/paper/05c2c9c5620a7ac93e61c1a76370e548 代码链接： https://github.com/google-research/nested-transformer

目前的ViT模型首先将输入的图像划分为图像块，然后使用类似于NLP模型中的方式对图像块提取特征。随后，使用多个自注意力层（self-attention）进行全局的信息交互，以实现长程的特征提取。最近的一些Transformer工作表明，ViT模型在大规模的标记数据集上可以超过卷积神经网络的精度，但是其在小数据集上进行训练时，优势并不明显。这可能是因为Transformer本身缺乏类似于卷积网络中局部性和平移等效性之类的归纳偏置。

目前已有一些工作对此进行改进，其中具有代表性的例如Transformer原作者Ashish Vaswani发表在CVPR2021上的HaloNets[1]，以及微软亚研院获得ICCV2021 Best Paper的Swin Transformer[2]。这些工作都提出了一些特殊操作对自注意力机制进行了修改，先通过对局部块进行关注，随后进行块与块的信息交互。本文提出的NesT采用另一种方式进行改进，即保证原自注意力机制不变，引入聚合函数（block aggregation function）来提高模型的整体性能，同时可以为模型带来一定的可解释性。

本文对这种层次化的设计进行了分析和改进，作者认为在Transformer模型中通过在非重叠图像块上提取特征，并以分层次的形式将其进行聚合是这种层次化设计的精髓，其中对多个块特征融合的函数（block aggregation function）可以促进跨区域非局部信息的交互。基于这一理念，本文提出了一种层次化嵌套Transformer（NesT），其仅需要在原始ViT代码上修改几行就可以实现。

相比普通的ViT模型，NesT还有以下三个明显优势：

1. NesT拥有更快的收敛速度，所需的训练数据量也很小，可以同时在ImageNet和小型数据集（如CIFAR）上得到良好的泛化性能。

2. 将NesT扩展到图像生成任务上，在生成速度上相比之前的方法快了近8倍。

3. 在NesT中，通过嵌套的层次结构可以实现特征学习和抽象过程的解藕，同时这种结构可以天然的类比于决策树过程，本文基于此提出了一种全新的模型可解释方法CradCAT，可以对图像的显著性区域进行定位。