窝牛号

嵌套分层转换器:ViT实现准确、数据高效和可解释的视觉分类

在视觉理解方面,Vision Transformer (ViT) 及其变体最近受到了极大的关注,因为它们在许多核心视觉应用(例如图像分类、对象检测和视频理解)上具有卓越的性能。ViT 的核心思想是利用自注意力层的力量来学习小块图像之间的全局关系。然而,补丁之间的连接数量随图像大小呈二次方增加。这种设计导致数据效率低下, 尽管原始的 ViT 可以比具有数亿张图像进行预训练的卷积网络表现得更好,但这样的数据要求并不总是可行的,而且在给定较少数据的情况下,它的性能仍然低于卷积网络。许多人正在探索寻找更合适的架构重新设计,以有效地学习视觉表示,例如通过添加卷积层和构建具有局部自注意力的 分层结构。

层次结构原理是视觉模型的核心思想之一,底层在高维像素空间学习更多局部对象结构,而顶层在低维特征空间学习更多抽象和高级知识。现有的基于 ViT 的方法专注于在自注意力层内设计各种修改以实现这样的层次结构,但是虽然这些方法提供了有希望的性能改进,但它们通常需要大量的架构重新设计。此外,这些方法缺乏可解释的设计,因此很难解释训练模型的内部工作原理。

为了应对这些挑战,在“嵌套分层转换器:迈向准确、数据高效和可解释的视觉理解”中,我们提出了对现有分层结构驱动设计的重新思考,并提供了一种新颖的正交方法来显着简化它们。这项工作的中心思想是将特征学习和特征抽象(池化)组件解耦:嵌套的转换器层分别编码图像块的视觉知识,然后聚合处理后的信息。这个过程以分层方式重复,形成金字塔的网络结构。由此产生的架构在ImageNet上取得了具有竞争力的结果并且在数据效率基准上优于结果。我们已经证明,这样的设计可以通过更快的收敛显着提高数据效率,并提供宝贵的可解释性优势。此外,Google介绍了 GradCAT,这是一种在推理时解释训练模型的决策过程的新技术。

架构设计

整个架构很容易实现,只需在原始 ViT 的源代码中添加几行 Python 代码即可。原始的ViT 架构将输入图像划分为小块,将每个块的像素投影到具有预定义维度的向量,然后将所有向量的序列馈送到包含多个堆叠相同变压器层的整体 ViT 架构。虽然 ViT 中的每一层都处理整个图像的信息,但使用这种新方法,堆叠的变换层仅用于处理包含几个空间相邻图像块的图像区域(即块)。这一步对于每个块都是独立的,也是进行特征学习的地方。最后,一个称为块聚合的新计算层将所有空间相邻的块组合在一起。在每个块聚合之后,对应于四个空间相邻块的特征被馈送到具有一堆变换层的另一个模块,然后联合处理这四个块。这种设计自然地构建了网络的金字塔层次结构,由于块聚合,底层可以专注于局部特征(如纹理),而上层专注于降维的全局特征(如对象形状)。

网络处理图像的可视化:给定输入图像,网络首先将图像划分为块,其中每个块包含 4 个图像块。每个块中的图像块被线性投影为向量,并由一堆相同的转换器层处理。然后提出的块聚合层聚合来自每个块的信息,并将其空间大小减少 4 倍。最高层级的块数减少为 1,并在其输出后进行分类。

可解释性

该架构具有非重叠的信息处理机制,在每个节点上都是独立的。这种设计类似于决策树结构,表现出独特的可解释性能力,因为每个树节点都包含其父节点正在接收的图像块的独立信息。我们可以通过节点追踪信息流来了解每个特征的重要性。此外,我们的层次结构保留了整个网络中图像的空间结构,从而获得了对解释有效的空间特征图。下面我们展示两种视觉可解释性。

首先,我们提出了一种在测试图像上解释训练模型的方法,称为基于梯度的类感知树遍历 (GradCAT)。GradCAT 从层次结构的顶部到底部跟踪每个块(树节点)的特征重要性。主要思想是找到从顶层的根节点到底层的子节点的最有价值的遍历,该遍历对分类结果贡献最大。由于每个节点都处理来自图像某个区域的信息,因此这种遍历可以很容易地映射到图像空间进行解释(如下图叠加的点和线所示)。

以下是模型在左侧输入图像(包含 4 只动物)上的 top-4 预测和相应的可解释性结果的示例。如下图所示,GradCAT 突出显示了沿层次结构的决策路径以及图像上局部图像区域中的相应视觉线索。

给定左侧输入图像(包含四个对象),该图可视化了前 4 个预测类的可解释性结果。遍历沿树定位模型决策路径,同时定位对预测目标类影响最大的相应图像块(图像上的虚线所示)

此外,下图可视化了 ImageNet 验证集的结果,并展示了这种方法如何实现一些直观的观察。例如,下面的打火机示例(左上图)特别有趣,因为地面实况类——打火机/火柴棒——实际上定义了右下角的火柴棒对象,而最显着的视觉特征(具有最高节点值)是实际上来自左上角的红灯,它在概念上与打火机共享视觉线索。这也可以从覆盖的红线中看出,它表示对预测影响最大的图像块。因此,虽然视觉提示是错误的,但输出预测是正确的。此外,下面木勺的四个子节点具有相似的特征重要性值(参见节点中可视化的数字;

GradCAT 获得的结果的可视化。图像来自ImageNet验证数据集

其次,与原始 ViT 不同,我们的分层架构在学习表示中保留了空间关系。顶层输出输入图像的低分辨率特征图,使模型能够通过在顶层学习表示上应用类注意力图(CAM) 轻松地执行基于注意力的解释。这使得仅使用图像级标签即可实现高质量的弱监督对象定位。有关示例,请参见下图。

基于 CAM 的注意力结果的可视化。较暖的颜色表示更高的注意力。图像来自ImageNet验证数据集

通过这种设计,特征学习只在局部区域独立发生,特征抽象发生在聚合函数内部。这种设计和简单的实现对于分类以外的其他类型的视觉理解任务来说已经足够通用了。它还大大提高了模型的收敛速度,显着减少了训练时间以达到所需的最大精度。

我们通过两种方式验证了这一优势。首先,我们将 ImageNet 精度上的 ViT 结构与不同数量的总训练 epoch 进行比较。结果显示在下图的左侧,展示了比原始 ViT 更快的收敛速度,例如,与总共 30 个训练 epoch 的 ViT 相比,准确度提高了约 20%。

其次,我们修改架构以执行无条件的图像生成任务,因为由于收敛和速度问题,为图像生成任务训练基于 ViT 的模型具有挑战性。通过转置建议的架构来创建这样的生成器很简单:输入是嵌入向量,输出是 RGB 通道中的完整图像,块聚合被Pixel Shuffle支持的块解聚合组件替换。令人惊讶的是,我们发现我们的生成器比容量可比的SAGAN更容易训练,并且表现出更快的收敛速度,以及更好的FID 分数(衡量生成的图像与真实图像的相似程度)。

左:与标准 ViT 架构相比,给定不同总训练时期数的 ImageNet 准确度。右图:ImageNet 64x64 图像生成 FID 分数(越低越好),单次 1000 次训练。在这两个任务上,我们的方法都显示出更好的收敛速度。

在这项工作中,Google展示了一个简单的想法,即在这种嵌套层次结构设计中解耦特征学习和特征信息提取,通过一种新的基于梯度的类感知树遍历方法来提高特征的可解释性。此外,该架构不仅提高了分类任务的收敛性,还提高了图像生成任务的收敛性。所提出的想法专注于聚合功能,因此与用于自注意力的高级架构设计正交。我们希望这项新研究能够鼓励未来的架构设计师探索更多可解释和数据效率更高的基于 ViT 的视觉理解模型,例如采用这项工作来生成高分辨率图像。我们也发布了源代码对于这项工作的图像分类部分。

本站所发布的文字与图片素材为非商业目的改编或整理,版权归原作者所有,如侵权或涉及违法,请联系我们删除

窝牛号 wwww.93ysy.com   沪ICP备2021036305号-1