与自然语言处理类似，预训练视觉主干的迁移提高了模型在各种视觉任务上的性能。更大的数据集、可扩展的架构和新的训练方法都推动了模型性能的提高。然而，视觉模型仍然远远落后于语言模型。具体来说，迄今为止最大的视觉模型ViT只有4B个参数，而入门级语言模型通常都超过10B个参数，更不用说540B个参数的大型语言模型了。为了探索AI模型的性能极限，GoogleResearch最近在CV领域的一项研究率先将VisionTransformer参数量扩展到22B，并提出了ViT-22B。与之前类似模型参数量的4B相比，可以说这是迄今为止最大的denseViT模型。论文地址：https://arxiv.org/pdf/2302.05442.pdf对比之前最大的ViT-G和ViT-e，表1为对比结果。从下表来看，ViT-22B主要扩展了模型的宽度，使参数量变大，深度与ViT-G相同。现在的ViT模型就像这位知乎网友说的一样。难道谷歌在ChatGPT上输了一场比赛，势必要在CV领域展开竞争？如何？在最初的研究阶段，他们发现在扩展ViT的过程中，会出现训练不稳定的情况，并可能带来架构的变化。研究人员随后精心设计了模型，并以前所未有的效率对其进行了并行训练。ViT-22B的质量在一组全面的任务中进行评估，从（小样本）分类到密集输出任务，它达到或超过了最先进的水平。例如，即使用作冻结视觉特征提取器，ViT-22B在ImageNet上也能达到89.5%的准确率。通过训练一个文本塔来匹配这些视觉特征，它在ImageNet上的零样本设置中达到了85.9%的准确率。此外，该模型可以被视为教师，用作蒸馏目标，研究人员在ImageNet上训练了一个准确率为88.6%的ViT-B学生模型，这对于这种规模的模型来说是最先进的。模型架构ViT-22B是一个基于Transformer的编码器模型，类似于原始的VisionTransformer架构，但包含以下三个主要修改以提高大规模训练的效率和稳定性：并行层、query/key(QK)Normalization和省略偏见。平行层。正如Wang和Komatsuzaki的研究中所述，该研究设计了一种注意力和MLP并行结构：这通过结合MLP和注意力块的线性投影来实现额外的并行化。值得注意的是，用于查询/键/值投影的矩阵乘法和MLP的第一线性层被融合到一个操作中，MLP的非注意力投影和第二线性层也是如此。QK归一化。训练大型模型的难点之一是模型的稳定性。在扩展ViT的过程中，研究人员发现训练损失在数千轮步骤后出现了分歧。这种现象在8B参数的模型中尤为突出。为了稳定模型训练，研究人员采用了Gilmer等人的方法。在点积注意力计算之前对查询和键应用LayerNorm归一化操作，以提高训练的稳定性。具体来说，注意力权重计算为：省略偏差。在PaLM之后，偏差项从QKV投影中移除，并且所有Layernorms都没有偏差地应用，从而提高了加速器利用率(3%)而没有质量损失。然而，与PaLM不同的是，研究人员对MLP的致密层使用了一个偏差项，即便如此，在不影响质量的情况下也没有速度损失。图2显示了一个ViT-22B编码器模块。embedding层在原有ViT的基础上进行patch提取、线性投影、添加位置embedding等操作。研究人员使用多头注意力池来聚合头部中的每个标记表示。ViT-22B使用图像分辨率为224×224的14×14块。ViT-22B使用学习的1D位置嵌入。在对高分辨率图像进行微调时，研究人员根据原始图像中预训练位置嵌入的位置进行二维插值。培训基础设施和效率ViT-22B使用FLAX库，在JAX中实施，并在Scenic中构建。它利用模型和数据并行性。值得一提的是，研究人员使用了jax.xmapAPI，它提供了对所有中间体（例如权重和激活）和芯片间通信的分片的明确控制。研究人员将芯片组织成大小为t×k的二维逻辑网格，其中t是数据平行轴的大小，k是模型轴的大小。然后，对于t组中的每一个，k个设备获得同一批图像，每个设备只保留1/k激活并负责计算所有线性层输出的1/k（详见下文）。图3：异步并行线性运算(y=Ax)：为跨设备的重叠通信和计算建模并行矩阵乘法。异步并行线性操作。为了最大化吞吐量，必须考虑计算和通信。也就是说，如果你希望这些操作在分析上等同于未分片的情况，你必须尽可能少地通信，理想情况下让它们重叠，这样你就可以保持矩阵乘法单元（其中大部分容量在FLOPs）总是很忙.参数切片。该模型是在第一轴上平行的数据。每个参数都可以在这个轴上完全复制，或者可以让每个设备保存它的一个块。研究人员选择从模型参数中拆分出一些大张量，以适应??更大的模型和批量大小。使用这些技术，ViT-22B在TPUv4训练期间每核每秒处理1.15k个令牌。ViT-22B的模型触发器利用率(MFU)为54.9%，表明硬件使用效率很高。请注意，PaLM报告的MFU为46.2%，而研究人员在同一硬件上测量ViT-e（仅数据并行）的MFU为44.0%。实验结果实验探索了ViT-22B对图像分类的评估结果。表2的结果表明，ViT-22B在各项指标上仍有明显提升。此外，研究表明，像ViT-22B这样的大型模型的线性探测可以接近或超过具有高分辨率的小型模型的完整微调性能，通常成本更低且更容易实现。该研究进一步测试了细粒度分类数据集iNaturalist2017的线性可分性，将ViT-22B与其他ViT变体进行了比较。该研究测试了224像素和384像素的输入分辨率。结果如图4所示。该研究观察到ViT-22B明显优于其他ViT变体，尤其是在标准224px输入分辨率下。这表明ViT-22B中的大量参数对于从图像中提取详细信息很有用。表3显示了ViT-22B对于CLIP、ALIGN、BASIC、CoCa、LiT模型的零样本传输结果。表3的底部比较了三种ViT模型的性能。在所有ImageNet测试集上，ViT-22B都取得了相当或更好的结果。值得注意的是，ObjectNet测试集上的零样本结果与ViT模型大小高度相关。最大的ViT-22B在具有挑战性的ObjectNet测试集上设置了新的SOTA设置。分布外(OOD)。我们构建了一个从JFT到ImageNet，以及从ImageNet到不同分布外数据集的标签图，即ObjectNet、ImageNet-v2、ImageNet-R和ImageNet-A。目前可以确认的结果是扩展模型增加了分布外性能，与ImageNet上的改进一致。这适用于只看过JFT图像的模型，以及在ImageNet上微调的模型。在这两种情况下，ViT-22B都延续了在较大模型上具有更好OOD性能的趋势（图5，表11）。此外，研究人员调查了ViT-22B模型在语义分割和单目深度估计任务中捕获的几何和空间信息的质量。语义分割。研究人员在三个基准测试中评估ViT-22B作为语义分割主干：ADE20K、PascalContext和PascalVOC。从表4可以看出，当只看到少量分割掩码时，ViT-22B主干传输效果更好。单目深度估计。表5总结了该研究的主要发现。从顶行（DPT解码器）可以看出，与不同的主干相比，使用ViT-22B功能可产生最佳性能（在所有指标上）。通过将ViT-22B主干与ViT-e（一个较小的模型，但使用与ViT-22B相同的数据进行训练）进行比较，研究发现扩展架构可以提高性能。此外，将ViT-e主干与ViT-L（与ViT-e类似的架构，但训练的数据较少）进行比较，研究发现这些改进也来自于扩展预训练数据。这些发现表明，更大的模型和更大的数据集都有助于提高性能。该研究还探索了视频数据集。表6显示了Kinetics400和MomentsinTime数据集上的视频分类结果，表明使用冻结的主干可以获得有竞争力的结果。该研究首先与ViT-e进行了比较，ViT-e拥有最大的先验视觉主干模型，由40亿个参数组成，并且也在JFT数据集上进行了训练。我们观察到更大的ViT-22B模型在Kinetics400上提高了1.5分，在MomentsinTime上提高了1.3分。最后，该研究指出，通过全面的端到端微调还有进一步改进的空间。有关更多技术细节，请参阅原始论文。
