谷歌提出「卷积+注意力」新模型,超越ResNet最强变体!

Transformer 跨界计算机视觉虽然已取得了一些不错的成绩,但大部分情况下,它还是落后于最先进的卷积网络。

现在,谷歌提出了一个叫做 CoAtNets 的模型,看名字你也发现了,这是一个 Convolution + Attention 的组合模型。


图片

该模型实现了 ImageNet 数据集 86.0% 的 top-1 精度,而在使用 JFT 数据集的情况下实现了 89.77% 的精度,性能优于现有的所有卷积网络和 Transformer!


图片

卷积结合自注意,更强的泛化能力和更高的模型容量

他们是如何决定将卷积网络和 Transformer 结合起来造一个新模型呢?

首先,研究人员发现,卷积网络和 Transformer 在机器学习的两个基本方面——泛化和模型容量上各具优势。

由于卷积层有较强的归纳偏置(inductive bias),所以卷积网络模型具有更好的泛化能力和更快的收敛速度,而拥有注意机制的 Transformer 则有更高的模型容量,可以从大数据集中受益。

那将卷积层和注意层相结合起来,不就可以同时获得更好的泛化能力和更大的模型容量吗!

那好,关键问题来了:如何有效地将它们结合起来,并在准确性和效率之间实现更好的平衡?


图片

研究人员又发现,常见的深度卷积(depthwise convolution)只需简单的相对注意就可以合并进注意力层!

因此,他们将 CoAtNets 的实现分为两步:

1、将卷积和自注意结合在一个基本计算块中;

2、将不同类型的计算块垂直堆叠在一起(垂直布局设计),形成完整的网络。

具体实现

首先,由于 Transformer 和 MBConv 中的 FFN 模块都采用了“反向瓶颈”的设计,加上深度卷积和自我注意都可以用预定义的感受野中的加权和来表示,CoAtNets 主要采用 MBConv 卷积块。

具体来说,卷积依赖于一个固定核从局部感受野收集信息:


图片

相比之下,自注意允许感受野成为整个空间位置,并基于对(xi,xj)之间的重归一化成对相似性来计算权重:


图片

将它们以最佳形式结合之前,研究人员比较了一下两者各自的理想特性。


图片

一个结合深度卷积和自注意的理想模型需要具备以上 3 个特性,而研究人员的实现方法很简单:在 Softmax 归一化前或后,将全局静态卷积核与自适应注意矩阵求和。


图片

将卷积和注意力结合起来之后,就可以开始堆叠整个网络。

全局上下文在空间大小(spatial size)方面具有二次复杂性,如果直接将上述公式中的相对注意应用到原始图像输入,会因为普通尺寸图像中过多的像素,导致计算速度过慢。

因此,要构建一个在实际操作中可行的模型,研究人员选择在 feature map 达到可管理水平后,进行一些下采样以减小空间大小,再使用全局相对注意。

其中下采样可以像 ViT 那样用具有积极步幅(stride,例如 16×16)的 convolution stem 或像 ConvNets 里具有渐进池的多级网络来实现。

研究人员按照上面两种方法堆叠出 5 个变体:C-C-C-C、C-C-C-T、C-C-T-T 和C-T-T-T 以及 ViT,其中C和T分别表示卷积和 Transformer。

为了作出选择,他们进行了对照实验,选出了泛化能力和模型容量最好的两个:C-C-T-T 和C-T-T-T,优中选优,最终选择了迁移性能更好的C-C-T-T。


图片

至此,完整的 CoAtNets 模型就实现了。

NO.1 的 top-1 精度

其实,将卷积和自注意相结合用于计算机视觉的想法并不新鲜,但此前的一些方法通常会带来额外的计算成本等问题。

而他们的相对注意实例化是深度卷积和基于内容的注意力的自然结合,加上优中选优的垂直布局设计,新模型 CoAtNets 造成的额外成本最小。

下面就来看具体的实验数据:


图片

仅使用 ImageNet-1K 数据集的实验结果表明:CoAtNet 不仅优于 ViT 变体,还能匹敌最佳的仅卷积结构(EfficientNet-V2 和 nfnet)。

从上表和下图还可以看出:使用 ImageNet-21K 进行预训练,CoAtNet 的优势变得更加明显,大大优于以前的所有模型。

且最好的 CoAtNet 变体实现了 88.56% 的 top-1 精度,而与之匹敌的 ViT-H/14 则需要在 23 倍大的 JFT 数据集上预先训练 2.3 倍规模的 ViT 才能达到 88.55% 的成绩。

这标志着 CoAtNet 模型数据效率和计算效率的显著提高。


图片

△ImageNet-21K ImageNet-1K 的参数精度

最后,研究人员用 JFT 进一步评估了大规模数据集下的 CoAtNet,发现它达到 89.77% 精度,超过 ResNet 最强变体 NFNet-F4+ 排名第一,同时在 TPU 训练时间和参数计数方面提高了 2 倍效率。


图片

作者介绍

论文的四位作者全部来自谷歌大脑。


图片

Dai Zihang,本科清华大学工商管理专业,硕士卡内基梅隆大学计算机专业,博士卡内基梅隆大学计算机和哲学双学位。


图片

刘寒骁,同为清华本科,卡内基梅隆大学计算机博士,谷歌大脑高级研究科学家,曾在 DeepMind 研究神经架构搜索。


图片

Quoc Le,这位大佬应该很多人都知道,斯坦福大学计算机科学博士毕业,吴恩达的学生,谷歌大脑的创始成员和 AutoML 的缔造者之一。


图片

Tan Mingxing ,北京大学博士毕业、康奈尔大学博士后。

本文链接

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注