“一把梭ViT”来了，谷歌提出可以灵活应对各种图像块尺寸的FlexiViT

大数据文摘授权转载自将门创投

作者：seven_

视觉Transformer（ViT）目前已成为学界和工业界常用的视觉任务backbone，其通过将一整幅图像切割成图像块（patch）的方式来将其转换为序列送入到Transformer编解码结构中进行特征建模。很多工作已经证明，切割出来的图像块的大小往往决定了整体模型的性能和速度，如果输入模型较小的图像块，会得到更高的识别准确率，但是计算成本会增加，反之亦然。图像块的大小直接控制了模型速度与准确率之间的权衡。那我们究竟是该选择大图像块还是小图像块呢，是否需要根据下游任务来决定呢？来自谷歌研究院的研究者告诉我们，“小孩子才做选择，成熟的ViT模型可以灵活应对各种尺寸”。

论文链接：

代码链接：

https://github.com/google-research/big_vision

在本文中，作者提出了一种ViT训练策略，即在训练时随机确定图像块的大小来训练模型，这样训练得到的模型称为FlexiViT，其权重可以适应多种输入尺寸，这使得模型在部署阶段根据不同任务不同数据定制模型成为可能。作者在大量下游任务上进行了验证实验，包括图像分类、图像-文本检索、开放世界目标检测、全景分割和语义分割等，结果表明FlexiViT可以超过在单一尺寸下训练的诸多模型，可以灵活的应对各种现实场景，下图展示了FlexiViT与其他baseline模型在ImageNet-1K上的速度-准确性曲线对比情况。

视觉Transformer的火热标志着学界对之前占主导地位的卷积神经网络（CNN）模式的重大革新，其通过“图像序列化”的方式来进行底层的像素处理，这种“序列化”也带来了一些新的处理模式，例如mask图像块、为特定任务添加专门的token，以及混合不同模态的图像块来处理多模态任务。

作者发现，ViT模型中的图像块作为最基本的视觉处理单元，但是对于图像块尺寸的研究却很少，相比CNN时代的多尺度几何分析、金字塔特征建模等工作，ViT还缺乏在这一方面的探索。虽然初代版本的ViT中使用了三种图像块尺寸（32×32、16×16和14×14），但在后续的演变中，为了训练方便，大家更愿意将图像块固定为16x16大小。在这项工作中，作者对图像块尺寸对于模型预测精度和计算速度的影响进行了研究，期间模型的参数容量并不发生变化。实验结果证明固定图像块训练得到的ViT只能在单一的输入尺度上表现良好，如果输入发生变化，模型性能会严重受损，但是如果想要适应新输入，就只能重新训练模型。

为了克服这个限制，本文作者提出了FlexiViT，顾名思义，这是一种可以灵活处理输入尺度的ViT模型。作者在训练期间随机生成不同尺寸的图像块，并针对每种尺度自适应的调整位置和嵌入参数的大小，如下图所示。

作者发现，这种自适应的训练pipeline可以有效提升模型的泛化能力，作者在广泛的下游任务上进行了实验，作者还发现，以这种方式训练得到的backbone在跨尺度的情况下性能也会得到保留，根据这一观察我们可以使用其来进行资源高效的迁移学习。

2.1 对标准的ViT进行分析

标准的ViT首先将一幅输入图像  切割成一系列图像块  ，可以将此过程称为“序列化”，接下来对每个图像块进行嵌入编码  ，最后将学习到的位置编码   添加到嵌入编码  中，以便送入到Transformer编码器完成运算。在这一过程中，对于给定的图像大小  ，图像块大小  决定了Transformer模型输入序列的长度。我们可以将ViT模型表示为  ，其中  是模型规模（小、中、基础、大、…）， 是图像块大小。

作者发现，标准的ViT结构并不灵活，其在不同的图像块情况下，性能会断崖式下降。为了使标准ViT能够强行输入不同图像块，作者简单地通过双线性插值调整嵌入权重  和位置编码  的大小，结果如下图所示，其中虚线和点线分别代表标准ViT的性能随着图像块尺度的变化而下降。

2.2 如何训练更加灵活的ViT模型

在上图中，作者展示了基础FlexiViT-B模型的性能，FlexiViT-B是在与ViT-B/16和ViT-B/30模型相同的设置下进行训练的，除了输入图像块大小进行了动态调整，其输入的图像块大小是从一组预定义的尺寸集合中随机均匀采样得到的，在具体实现时，仅需要对训练代码进行两处小改动。

首先需要为嵌入权重  和位置编码  定义一个基础参数shape，保证其是可学习参数，将其作为网络前向传播的一部分。为了使FlexiViT能够应对几乎任何尺寸，作者使用240x240的输入图像，切割图像块大小为{240, 120, 60, 48, 40, 30, 24, 20 , 16, 15, 12, 10, 8, 6, 5, 4, 2, 1}，在均匀采样时，采样范围为8-48。

下图展示了将标准ViT转换为FlexiViT的伪代码，需要注意的是，调整图像块大小不等同于调整输入图像大小，调整图像块尺寸不会影响模型参数，而调整输入图像尺寸则会丢失图像信息。

2.3 如何调整图像块嵌入的大小

在调整完输入图像块  之后，还应该考虑对块嵌入向量以及块嵌入权重  进行相应的变换。如果简单的使用双线性插值来缩放，则生成的token大小会很难控制，例如  。理想情况下，我们希望调整过程中不会丢失图像中的关键信息。一种较为直接的实现方法是在对图像块进行嵌入编码后，立即使用LayerNorm进行规范化处理，但是这种方法需要改变模型架构，并且与现有的预训练ViT不兼容。因此作者提出了一种更加巧妙的调整方法，即设计了一系列矩阵线性变化来实现：

其中  是任意输入，这种方法避免了对原有架构的修改，对本文的pipeline更加合适。

2.4 FlexiViT与知识蒸馏能否碰撞出火花

鉴于FlexiViT的灵活性，作者突发奇想，能否将其与知识蒸馏技术相结合来提高性能。一般情况下，知识蒸馏框架的优化难度要远远大于传统的标签监督方法，一些工作会将学生模型与教师模型架构设置的非常接近来解决这一问题，但是这样做又会影响蒸馏性能，因为教师模型通常要比学生模型具有更大的架构和规模，并且构成结构上的差异性。对于这种两难的情况，使用FlexiViT在合适不过，由于FlexiViT天然的架构灵活性，可以使用更加庞大复杂的ViT作为教师模型来优化学生FlexiViT，FlexiViT的随机块采样可以完美覆盖多种大型模型，并显著提高蒸馏性能。作者使用大型ViT-B/8模型作为教师模型进行实验，通过最小化教师和学生 FlexiViT的预测与随机块嵌入之间的KL距离进行优化：

其中  是FlexiViT模型在图像块大小为  的类分布， 是教师对完全相同输入的预测分布。下图比较了使用教师初始化与随机初始化和标签监督训练的蒸馏效果。其中淡绿色曲线展示了蒸馏效果，可以看到，蒸馏FlexiViT在图像块尺寸较小时与教师ViT的性能较为接近，当尺寸增加时，蒸馏FlexiViT的性能有很大提升。

参考

[1] Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Algorithms for learning kernels based on centered alignment. The Journal of Machine Learning Research, 13:795–828, 2012. 5

点「在看」的人都变好看了哦！

关注公众号：拾黑（shiheibook）了解更多

[广告]赞助链接：

四季很好，只要有你，文娱排行榜：https://www.yaopaiming.com/
让资讯触达的更精准有趣：https://www.0xu.cn/