【论文阅读】Patch n‘ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution_patch n' pack

Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution

引用： Dehghani M, Mustafa B, Djolonga J, et al. Patch n’pack: Navit, a vision transformer for any aspect ratio and resolution[J]. Advances in Neural Information Processing Systems, 2024, 36.

论文链接： （NeurIPS 2023）https://proceedings.neurips.cc/paper_files/paper/2023/hash/06ea400b9b7cfce6428ec27a371632eb-Abstract-Conference.html

代码链接： https://github.com/kyegomez/NaViT

论文概述

此前的模型将图像大小调整为固定分辨率明显是一个次优的选择。然而，诸如 Vision Transformer (ViT) [1]之类的模型提供了灵活的基于序列的建模，因此可以改变输入序列长度。支撑Vit的操作很简单：将图像拆分为多个补丁，每个补丁都线性投影到一个token上。通常，输入图像的大小会调整为固定的正方形纵横比，然后拆分为固定数量的patches。FlexiViT[2]在一个架构中支持多种补丁大小，从而实现序列长度的平滑变化，从而计算成本。这是通过在每个训练步骤中随机抽取一个补丁大小和调整大小的算法来实现的，以允许初始卷积嵌入支持多个补丁大小。Pix2Struct[3]引入了一种保留纵横比的替代修补方法，这对于图表和文档理解等任务特别有用。

论文提出了 NaViT（Native Resolution ViT）来实现这一点，NaViT在训练期间使用sequence packing来处理任意分辨率和宽高比的输入。NaViT将来自不同图像的多个补丁打包在一个序列中，称为 Patch n’ Pack，这样可以在保留宽高比的同时实现可变分辨率。这是受自然语言处理中示例打包的启发，其中多个示例被打包到一个序列中，以适应对可变长度输入的有效训练。NaViT 可以有效地转移到图像和视频分类、对象检测和语义分割等标准任务，并提高鲁棒性结果和基准。在推理时，输入分辨率的灵活性可用于平稳地进行测试时的成本性能权衡。

论文证明了：

1. 在训练时随机抽样分辨率显着降低了训练成本。
2. NaViT 在各种分辨率下都能实现高性能，从而在推理时实现平滑的性价比权衡，并且可以以更低的成本适应新任务，
3. 通过示例打包实现的固定批次形状带来了新的研究思路，例如纵横比保留分辨率采样、可变令牌丢弃率和自适应计算。
4. 在分配的计算预算内处理的训练样本数量的大幅增加是提高 ViT 性能的主要因素——example packing加上可变分辨率输入和可变token drop使 NaViT-L/16 能够在训练期间处理五倍以上的图像。这种效率的提高延伸到微调过程。此外，通过在预训练和微调期间将 NaViT 暴露于多种分辨率，单个模型在评估各种分辨率时表现出出色的性能，在推理成本方面显着提高了 NaViT。

Method

NaViT 建立在原始 ViT 之上，但原则上可以使用任何在一系列patches上运行的 ViT 变体。为了启用 Patch n’ Pack，我们进行了以下架构修改：

• Masked self attention and masked pooling：为了防止示例相互关注，引入了额外的self attention mask。同样，编码器之上的masked pooling旨在将每个示例中的标记表示形式集中起来，从而在序列中为每个示例提供单个向量表示。图 2 显示了如何通过掩蔽来控制注意力的接受领域。
• Factorized & fractional positional embeddings：为了处理任意分辨率和纵横比，论文重新审视了位置嵌入。给定分辨率 R×R 的方形图像，具有补丁大小 P 的vanilla ViT 学习长度的 1-D 位置嵌入。线性插值这些嵌入对于以更高分辨率 R 进行训练或评估是必要的。为了支持可变纵横比并易于推断到看不见的分辨率，引入了因式分解位置嵌入，其中将 x 和 y 坐标分解为单独的嵌入 φx 和 φy。
• Continuous Token dropping：packing支持连续tokens丢弃，因此tokens丢弃率可以因图像而异。这样一来，在看到一些完整图像的同时，通过丢弃实现更快的吞吐量，从而减少了训练/推理差异。
• Resolution sampling：NaViT可以使用每个图像的原始分辨率进行训练。或者，可以在保留纵横比的同时对像素总数进行重新采样。在原版ViT中，更高的吞吐量（在较小的图像上进行训练）和更高的性能（在较大的图像上进行训练，以便在评估时实现高分辨率）之间存在紧张关系。通常，模型以较小的分辨率进行预训练，并以较高的分辨率进行微调。NaViT 允许通过从图像大小分布中采样来进行混合分辨率训练，同时保留每个图像的原始纵横比。这既可以提高通量，又可以曝光到大图像，与同等的 ViT 相比，性能显著提高（在模型大小和训练持续时间方面）。

实验

在两种设置中预训练NaViT：JFT-4B上的分类训练和WebLI上的对比语言图像训练。通常，对于JFT，初始裁剪是在预训练中应用的，在这两种情况下，图像的大小都会调整为正方形（扭曲的纵横比）。NaViT在JAX中使用FLAX库实现，并内置在Scenic中。

图 1 显示了不同 NaViT 模型与计算匹配的 ViT 基线相比的 JFT 预训练性能。NaViT 在性能上始终优于 ViT，同时在不同的计算和参数规模上使用相同的计算预算;例如，性能最佳的 ViT 的性能可以与计算量少四倍的 NaViT 相匹配。相反，图 1 中计算最轻的 NaViT 的成本效益是其等效 ViT 对应物的 5 倍。NaViT 模型受益于保留的纵横比和对许多分辨率进行评估的能力，但这里的主要贡献者是 NaViT 在分配的计算预算内处理的训练示例数量显着增加。这是通过对多个可变分辨率示例进行采样和标记丢弃相结合来实现的，从而产生可变大小的图像，这些图像被有效地打包成与原始模型相似的序列长度。Patch n’ Pack（将序列打包简单应用于视觉转换器）可以显著提高训练效率。由此产生的 NaViT 模型可以在推理时应用于许多分辨率，并且可以廉价地适应新任务。总体而言，Patch n’ Pack 支持以前因需要固定批次形状而受到阻碍的各种研究，包括自适应计算和用于提高训练和推理效率的新算法。

参考文献

[1] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas
Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.

[2] Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, and Filip Pavetic. Flexivit: One model for all patch sizes. In CVPR, 2023.

[3] Kenton Lee, Mandar Joshi, Iulia Turc, Hexiang Hu, Fangyu Liu, Julian Eisenschlos, Urvashi Khandelwal, Peter Shaw, Ming-Wei Chang, and Kristina Toutanova. Pix2struct: Screenshot parsing as pretraining for visual language understanding, 2022.