NaViT

一、总览

题目: Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
机构：Google DeepMind
论文: https://arxiv.org/pdf/2307.06304.pdf
代码：
任务:
特点:
前置相关工作：Pix2Struct
同期类似工作：
后续衍生工作：

二、摘要

当我们在用CV模型处理图像时，会把图像调整为固定分辨率，但是往往这是一个次优的选择。
然而诸如Vision Transformer（ViT）之类的模型能够灵活地进行序列建模，因此可以改变输入序列的长度。本文提出NaViT（Native Resolution ViT）来利用这一点，它在训练期间使用序列打包（sequence packing） 来处理任意分辨率和宽高比的输入。
除了让模型更加灵活之外，NaViT能够有效地迁移到图像/视频分类、目标检测和语义分割等标准任务，并且在基准评测上提高鲁棒性。
在推理时，这样一种输入分辨率的灵活性，能够在测试时对性能和效率之间做一个平稳的平衡。作者相信，NaViT不同于标准的，CNN型的输入建模设计，能够为ViTs指明一个有前途的方向。

三、引言

ViT（Dosovitskiy 等人，2021）的简单、灵活和可扩展的性质使其成为基于卷积的神经网络的几乎无处不在的替代品。该模型的基础是一个简单的操作：将图像分割成块，每个块都线性投影到一个token。通常，输入图像被调整为固定的正方形长宽比，然后分割成固定数量的补丁。

最近的工作探索了这种范例的替代方案：

1. FlexiViT（Beyer 等人，CVPR 2023）在一个架构内支持多种patch size，从而实现序列长度/计算成本的平滑变化。这是通过在每个训练步骤中对patch size进行随机采样以及调整resize算法以允许初始卷积特征支持多个patch size来实现的。
2. Pix2Struct（Lee 等人，2022）引入了一种保留宽高比的替代patching方法，这对于图表和文档理解等任务特别有用。

我们提出了另一个备选方案，NaViT。许多来自不同图像的patches被打包在一个序列里面，因此也被叫做patch n’pack （补丁包），这能够允许可变的分辨率而且能够保持宽高比。方法也是受自然语言处理领域的example packing启发，将多个样例打包在一个序列里面，来在变长输入上进行更有效率的训练。

本文发现：

1. 训练时，对分辨率进行随机采样，极大程度减少了训练代价。
2. NaViT由于能够处理不同的范围的分辨率，因此能够在推理的时候，很好地在成本-效果之间做一个平衡，并且以一种较小的代价迁移到新的任务上。
3. 由example packing产生的固定batch形状，能够激发一些新的研究点，比如保持长宽比，可变token dropping rate，自适应计算。

这些发现有重大的实际影响，在一个固定计算代价下，NaViT能够稳定地比ViT效果更优。比如，我们仅用了少于1 / 4的计算成本，就比ViT的最好表现更优。如下图中的左图。

我们认为在分配的计算预算内处理的训练示例数量的大幅增加是比 ViT 性能提高的主要贡献者 - 示例打包与可变分辨率输入和可变令牌丢弃相结合，使 NaViT-L/16 的处理能力增加了五倍训练期间的图像（表2）。

我们认为在分配的计算预算内处理的训练样本数量的大幅增加是比 ViT 性能提高的主要原因 - 示例打包与可变分辨率输入和可变token丢弃相结合，使 NaViT-L/16 的处理能力增加了五倍多的训练图像（表2）。

这种效率上的提高，在finetuning阶段也能有效地体现，比如上面图1中的（middle），在更小的预训练成本下，就能取得更高的finetune acc。

除此之外，由于在预训练和微调阶段都处理可变分辨率，因此只有一个模型也能在多种分辨率上取得很好的效果，很好地体现了NaViT在推理代价上的优势，如图1中的（right）。

NaViT在训练和迁移上的效率，灵活的推理，为 Vision Transformers 提供了一条充满希望的途径。Patch n’ Pack使计算机视觉系统能够超越当前数据和建模pipeline所施加的限制，实现以前受固定batch形状限制的想法，从而释放创新和进步的新可能性。

四、方法

深度神经网络通常使用批量输入进行训练和运行。为了在当前硬件上进行高效处理，这意味着固定的批次形状，这又意味着计算机视觉应用程序的图像大小固定。再加上历史上与卷积神经网络相关的架构限制，导致了调整图像大小或将图像填充到固定大小的做法。这两种方法都已被证明存在缺陷：前者会损害性能，后者效率低下（Lee et al., 2022）。对 ImageNet (Deng et al., 2009)、LVIS (Gupta et al., 2019) 和 WebLI (Chen et al., 2022c) 中长宽比的分析，结果表明大多数图像通常不是方形的（图 3）。

在语言建模中，通常通过示例打包来绕过固定序列长度的限制：来自多个不同样例的token被组合在一个序列中，这可以显着加速语言模型的训练（Krell 等人，2021）。通过将图像视为补丁（令牌）序列，我们表明 Vision Transformers（Dosovitskiy 等人，2021）可以从相同的范例中受益，我们称之为 Patch n’ Pack。使用这种技术，ViT 可以在“原始”分辨率的图像上进行训练，我们将这种方法命名为 NaViT。

4.1、架构改动

相对于原始的ViT，在模型的架构上做了如下的一些改动：

1. Masked self attention and masked pooling 为了防止example之间彼此attend，引入了额外的self-attention掩模。类似地，masked pooling旨在池化每个示例中的token表示，从而导致序列中每个示例都有一个向量表示。图 2 展示了如何通过掩码来控制注意力的感受野。
2. Factorized & fractional positional embeddings. 为了处理任意分辨率和纵横比，我们重新审视position embedding。给定分辨率为 $R\times R$ 的方形图像，补丁大小为 P 的普通 ViT 可以学习长度为$(R/P{)}^2$ 的一维位置嵌入（Dosovitskiy 等人，2021）。线性插值这些嵌入对于以更高分辨率 R 进行训练或测试是必要的。

Pix2struct（Lee 等人，2022）引入了可学习的 2D 绝对位置编码，从而学习大小为 [maxLen, maxLen] 的位置嵌入，并用每个补丁的 (x, y) 坐标进行索引。这使得宽高比可变，分辨率高达 $R=P\cdot maxLen$。然而，每一种组合(x, y) 坐标在训练过程中必须可见。

为了支持可变的纵横比并轻松外推（extrapolate ）到没有见过的分辨率，我们引入了因式分解的位置嵌入，其中我们分解为 x 和 y 坐标的单独嵌入 ${\varphi }_x$ 和 ${\varphi }_y$。然后将它们组合在一起（第 3.4 节探讨了可选的组合策略）。我们考虑两种模式：绝对嵌入，其中 ${\varphi }_{(}p)$ : [0, maxLen] → $R^D$ 是绝对补丁索引的函数，以及分数嵌入，其中 ${\varphi }_{(}r)$ : [0, 1] → $R^D$ 是一个函数$r=p/side\_length$，即沿图像的相对距离。后者提供独立于图像大小的位置嵌入参数，但部分地混淆了原始长宽比，从而仅隐含在补丁数量中。我们使用了简单学习的embedding $\varphi$，sinusoidal embeddings，以及用NeRF里面用到的傅立叶位置编码。

4.2、训练改变

Patch n’ pack技术使得在训练NaViT的时候，能够有一些新的训练技术。

1. Continuous Token dropping. token丢弃（训练期间随机省略输入补丁）（Akbari 等人，2021；Li 等人，2023 flip）已被开发来加速训练。然而，它们通常从所有示例中删除相同比例的令牌；打包可以实现连续的令牌丢弃，从而可以根据图像改变令牌丢弃率。这样可以通过丢弃同时仍然看到一些完整图像来实现更快的吞吐量，从而减少训练/推理差异。此外，通过打包，在整个训练过程中，drop分布可能会根据一些预定义的schedule而变化。在第 3.3 节中，我们探讨了不同的schedules以及灵活token丢弃的好处。
2. Resolution sampling. NaViT 可以使用每个图像的原始分辨率进行训练。或者，可以在保留纵横比的同时对像素总数进行重新采样。在普通 ViT 中，更大的吞吐量（在较小的图像上进行训练）和更高的性能（在更大的图像上进行训练，以在评估时实现高分辨率）之间存在着紧张关系。通常，模型会在较小的分辨率下进行预训练，并在较高的分辨率下进行微调。 （Touvron等人，2019）。 NaViT 更加灵活；它允许通过从图像尺寸分布中采样来进行混合分辨率训练，同时保留每个图像的原始长宽比。这允许更高的吞吐量和大图像的曝光，与同等 ViT 相比，性能得到显着提高（在模型大小和训练持续时间方面）。 3.2 节探讨了不同的采样策略以及用于预训练和微调的可变分辨率训练。

4.3、NaViT的效率

在这里，我们讨论 Patch n’ Pack 对 NaViT 计算效率的一些影响。
self attention cost. 当将多个图像打包成更长的序列时，$O(n^2)$ 的注意力成本自然是一个值得关注的问题。尽管许多工作旨在消除这种二次缩放（Tay et al., 2022, 2020），但我们在这里证明，随着transformer隐藏维度的缩放，注意力在总体成本中所占的比例越来越小，其中包括计算成本MLP 也是如此。图 4 说明了这一趋势，表明与打包示例相关的开销相应减少。除了速度考虑之外，自注意力的内存成本也可能对极长的序列构成挑战。然而，这一挑战也可以通过采用内存高效的方法来解决（Rabe 和 Staats，2021；Dao 等人，2022）。

Packing, and sequence-level padding. 最终序列包含多个示例的长度必须是固定的。我们使用附录 A.3 中讨论的贪婪打包方法；通常不存在完全加起来等于固定长度的示例的完美组合，并且必须使用填充标记。例如，可以动态选择分辨率或令牌丢弃序列中最后一个示例与剩余标记完全匹配的比率；然而，我们发现通常只有不到 2% 的标记是填充标记，因此简单的方法就足够了。

Padding examples and the contrastive loss. 用打包序列可以很容易地实现每个令牌的损失。然而，许多计算机视觉模型都是通过示例级损失进行训练的，通常应用于池表示。首先，这需要对典型的池头进行修改以考虑堆积。其次，必须从每个序列中提取多个合并表示。固定批次形状需要假设，从一批 B 序列中，我们最多提取 B × Emax 池表示（即每个序列的 Emax 个示例）。如果序列包含超过 Emax 图像，则多余的图像将被丢弃，从而浪费模型编码器的计算。如果序列的示例少于 Emax，则损失将处理大量假填充表示。
后者是对比学习的一个问题，其中损失计算在时间和内存上扩展 ~ O(n2)。为了避免这种情况，我们使用了分块对比损失（Mustafa et al., 2023），它通过在本地设备子集上执行计算并有效地积累全局 softmax 归一化所需的统计数据来避免收集 softmax 的所有数据点的需要。这使得 Emax 值较高（从而有效地使用模型编码器），而不会受到损失的瓶颈。