论文链接: 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还有以下三个明显优势:
NesT拥有更快的收敛速度,所需的训练数据量也很小,可以同时在ImageNet和小型数据集(如CIFAR)上得到良好的泛化性能。
将NesT扩展到图像生成任务上,在生成速度上相比之前的方法快了近8倍。
在NesT中,通过嵌套的层次结构可以实现特征学习和抽象过程的解藕,同时这种结构可以天然的类比于决策树过程,本文基于此提出了一种全新的模型可解释方法CradCAT,可以对图像的显著性区域进行定位。