多模态论文笔记——NaViT

大家好，这里是好评笔记，公主号：Goodnote，专栏文章私信限时Free。本文详细解读多模态论文NaViT（Native Resolution ViT），将来自不同图像的多个patches打包成一个单一序列——称为Patch n’ Pack——从而实现可变分辨率并保持长宽比。在训练过程中采用序列封装的方式处理任意分辨率和纵横比的输入，除了具备灵活性的模型应用外，还展示了通过大规模监督和contrastive image-text pretraining来提高训练效率。

论文

论文名：Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
论文链接：https://arxiv.org/pdf/2307.06304

摘要

1. 现有问题与研究背景：将图像调整为固定分辨率，再用计算机视觉模型（如训练视觉生成模型）处理的普遍做法存在缺陷，但尚未被有效解决。视觉Transformer（ViT）能进行灵活的序列建模，可处理不同长度输入序列。
2. NaViT模型：作者利用ViT特性开发了NaViT（原生分辨率ViT），训练时采用序列打包技术（Patch n’ Pack），可处理任意分辨率和宽高比的输入图像。
3. 模型优势
   • 训练效率：在大规模监督和对比图像文本预训练中，NaViT训练效率显著提升。
   • 任务迁移：能高效迁移到图像和视频分类、目标检测、语义分割等标准任务。
   • 性能表现：在鲁棒性和公平性基准测试中结果更优；推理时可通过输入分辨率灵活性权衡成本和性能。
4. 研究意义：NaViT突破了传统卷积神经网络标准输入和建模管道的限制，为视觉Transformer发展开辟新方向。

ViT 参考历史文章：多模态论文笔记——ViT、ViLT

1. 引言

视觉Transformer（ViT）简单、灵活且可扩展的特性，使其几乎普遍地取代了基于卷积的神经网络。该模型的基础是一个简单操作：将图像分割成小块（patch），每个patch被线性投影成一个标记（token）。通常情况下，输入图像会被调整为固定的正方形宽高比，然后分割成固定数量的patch。最近的一些研究探索了这一范式的替代方案：

• FlexiViT（Beyer等人，2023）架构：通过支持多种块尺寸和随机采样技术实现序列长度与计算成本的调整。
• Pix2Struct（Lee等人，2022）：引入了另一种补丁方法，该方法保留了图像的宽高比，这在图表和文档理解等任务中特别有用。

作者在本文提出NaViT，采用Patch n’ Pack技术，把不同图像的多个patch打包到一个序列，能在保留宽高比的同时支持可变分辨率，灵感来自自然语言处理中的示例打包技术。

作者证明：

1. 在训练时随机采样分辨率可显著降低训练成本。
2. NaViT在广泛的分辨率范围内都能实现高性能，在推理时可实现平滑的成本 - 性能权衡，并且可以以较低成本适配新任务。
3. 示例打包实现的固定批次形状带来了新的研究思路，如保持宽高比的分辨率采样、可变的标记丢弃率和自适应计算。

NaViT性能对比如下：

• 性能比较：在固定计算预算下，NaViT性能始终优于ViT，用四分之一计算量的NaViT能达到最佳ViT的表现（见图1左）。
• 性能优势原因：NaViT性能更优主要是因为训练示例数量大幅增加。示例打包、可变分辨率输入和可变标记丢弃，使NaViT - L/16训练时处理图像数量是ViT的五倍（见表2）。
• 微调效率：NaViT在微调过程中也展现出效率提升（见图1中）。
• 推理优势：预训练和微调时接触多种分辨率，让NaViT在不同分辨率评估时性能出色，在推理成本上优势显著（见图1右）。

NaViT的训练和适应效率，以及灵活的推理能力，为视觉Transformer开辟了一条充满前景的道路。Patch n’ Pack技术使计算机视觉系统能够超越当前数据和建模流程所带来的限制，让那些此前受固定批次形状约束而无法实现的想法成为可能，为创新和发展开启了新的可能性。

2. 方法

图像尺寸问题

1. 深度神经网络训练与图像尺寸的关系：深度神经网络训练和运行依赖批量输入，为了在当前硬件上高效处理，需固定批量形状，这使得计算机视觉应用中图像大小也被固定。
2. 现有图像尺寸处理方法的问题：受卷积神经网络架构限制，常采用图像调整大小或填充为固定尺寸的方法，但这两种方法存在缺陷，调整大小会损害性能，填充则效率低下。
3. 图像宽高比分析结果：对ImageNet、LVIS和WebLI等代表性数据集的图像宽高比分析（如下图3）显示，多数图像并非正方形。
   

本文解决方案

• 语言建模常见做法：在语言建模里，通过示例打包把多个不同示例的标记组合到一个序列，以此绕过固定序列长度限制，能显著加速语言模型训练（Krell等人，2021 ）。
• 视觉建模处理方案：通过将图像视为小块（patches）序列，作者发现ViT 也可以从相同的范式中受益，这种技术被称为Patch n’ Pack。借助 Patch n’ Pack 技术，ViT 能够直接在图像的 “原生” 分辨率下进行训练，该训练方法被命名为 NaViT。

2.1 架构变化

图2：示例打包技术支持可变分辨率且保持宽高比的图像输入，能减少训练时间、提升性能并增加灵活性。此处展示了为支持Patch n’ Pack技术，在数据预处理和建模方面需要进行的修改。网络中的逐位置操作，如多层感知器（MLPs）、残差连接和层归一化等，无需改变。

NaViT基于原始的视觉Transformer（ViT）构建，但原则上可以使用任何对补丁序列进行操作的ViT变体。为了实现Patch n’ Pack技术，作者进行了以下架构修改：

• 掩码自注意力和掩码池化：为了防止不同示例之间相互关注，作者引入了额外的自注意力掩码。类似地，在编码器之上进行掩码池化，目的是对每个示例内的标记表示进行池化，从而在序列中为每个示例生成单个向量表示。图2展示了如何通过掩码控制注意力的感受野。
• 分解式和分数式位置嵌入：为了处理任意分辨率和宽高比，需要重新审视位置嵌入。对于分辨率为$R\times R$的正方形图像，具有patch尺寸$P$的普通ViT的位置编码为可学习的1D位置嵌入，长度为$(R/P{)}^2$。在更高分辨率$R$下进行训练或评估时，需要对这些嵌入进行线性插值。

Pix2struct（Lee等人，2022年）引入了可学习的二维绝对位置嵌入。通过这种方式，可以学习尺寸为[maxLen, maxLen) 的位置嵌入，并使用每个补丁的(x, y) 坐标进行索引。这使得宽高比可变，分辨率最高可达R = P · maxLen。然而，在训练过程中必须看到(x, y) 坐标的每一种组合。

为了支持可变宽高比并能轻松外推到未见过的分辨率，作者引入了分解式位置嵌入，将其分解为$x$和$y$坐标的单独嵌入${\varphi }_x$和${\varphi }_y$。然后将它们相加（在3.4节中探索了其他组合策略）。考虑了两种模式：

• 绝对嵌入，其中$\varphi (p):[0,maxLen]\to {\mathbb{R}}^D$是绝对patch索引的函数；
• 分数嵌入，其中$\varphi (r):[0,1]\to {\mathbb{R}}^D$是$r=P/side-length$的函数，即沿着图像的相对距离。

后者提供了与图像大小无关的位置嵌入参数，但会部分模糊原始宽高比，而原始宽高比仅隐含在patch数量中。作者考虑简单的可学习嵌入$\varphi$、正弦嵌入以及NeRF（Tancik等人，2020）中使用的可学习傅里叶位置嵌入。

2.2 训练变化

Patch n’ Pack使得在NaViT的训练过程中可以使用一些新技术。

• 连续标记丢弃：标记丢弃（在训练期间随机省略输入图像块）已被用于加速训练。然而，通常情况下，所有样本丢弃的标记比例相同；而打包技术实现了连续标记丢弃，即每张图像的标记丢弃率可以不同。这既能通过丢弃标记实现更快的处理速度，又能保留部分完整图像，从而减少训练与推理之间的差异。此外，借助打包技术，丢弃分布可以在整个训练过程中按照预先设定的计划发生变化。3.3节将探究不同的计划安排以及灵活标记丢弃的益处。
• 分辨率采样：NaViT可以使用每个图像的原始分辨率进行训练。或者，也可以在保持宽高比的同时对总像素数进行重采样。NaViT更加灵活，它允许通过从图像大小分布中采样来进行混合分辨率训练，同时保留每个图像的原始宽高比。这既提高了处理速度，又能接触到更大的图像，与同等的ViT相比（在模型大小和训练时长方面），性能有了显著提升。3.2节探索了不同的采样策略，以及预训练和微调中的可变分辨率训练。

  普通ViT中，更高的处理速度（在较小图像上训练）和更高的性能（在较大图像上训练，以便在评估时实现高分辨率）之间存在矛盾。通常情况下，模型在较低分辨率下进行预训练，在较高分辨率下进行微调（Touvron等人，2019）

2.3 NaViT的效率

这里讨论Patch n’ Pack对NaViT计算效率的一些影响。

图4：假设每张图像有256个标记，由于打包而产生的额外注意力负担；这种负担会随着模型规模的增大而减小。

• 自注意力成本：当将多个图像打包成长序列时，注意力的$O(n^2)$成本是一个自然的问题。尽管许多工作旨在消除这种二次缩放（Tay等人，2022，2020），但作者在此证明，随着Transformer隐藏维度的增加，注意力在总成本（包括MLP的计算成本）中所占的比例越来越小。图4展示了这一趋势，表明与打包示例相关的开销相应减少。除了速度方面的考虑，自注意力的内存成本对于极长序列也是一个挑战。然而，这个问题也可以通过采用内存高效的方法（Rabe和Staats，2021；Dao等人，2022）来解决。
• 打包和序列级填充：包含多个示例的最终序列长度必须固定。作者使用附录A.3中讨论的贪心打包方法；通常情况下，不存在恰好能组合成固定长度的示例，因此必须使用填充标记。例如，可以动态地选择序列中最后一个示例的分辨率或标记丢弃率，以恰好填充剩余的标记；作者发现通常填充标记的比例不到2%，因此这种简单的方法就足够了。
• 填充示例和对比损失：对于打包后的序列，基于每个标记的损失很容易实现。然而，许多计算机视觉模型是通过示例级损失进行训练的，这种损失通常应用于池化表示。首先，这需要对典型的池化层进行修改，以适应打包操作。其次，必须从每个序列中提取多个池化表示。固定的批量形状要求我们做出这样的假设：从B个序列的批量中，我们最多提取$B\times E_{max}$个池化表示（即每个序列最多$E_{max}$个示例）。如果一个序列包含超过$E_{max}$张图像，多余的图像将被丢弃，这会浪费模型编码器的计算资源。如果一个序列的示例少于$E_{max}$个，那么损失计算将处理大量虚假的填充表示。

后者（如果一个序列的示例少于$E_{max}$个）在对比学习中是一个问题，因为损失计算在时间和内存上的复杂度为$O(n^2)$。为了避免这个问题，作者使用了分块对比损失（Mustafa等人，2023），它通过在本地设备子集上进行计算，并有效地累积全局softmax归一化所需的统计信息，避免了收集所有数据点进行softmax计算的需求。这使得可以使用较高的$E_{max}$值（从而有效利用模型编码器），而不会受到损失计算的瓶颈限制。

3. 实验

• NaViT架构：基于普通ViT，融入2.1节支持打包的修改，还纳入查询 - 键归一化、省略偏差及注意力池化等改进。
• 预训练设置：在JFT - 4B上分类训练，在WebLI上进行对比语言 - 图像训练；预训练时通常对JFT数据集应用Inception裁剪，而NaViT模型预训练不做此操作并保留图像宽高比。
• 实现框架：使用JAX和FLAX库，基于Scenic构建。

3.1 提高训练效率和性能

• 性能对比：图1展示不同NaViT模型在JFT预训练性能，并与计算量匹配的ViT基线模型对比。在不同计算量和参数规模下，NaViT在相同计算预算时性能优于ViT，如性能最佳的ViT性能，NaViT可用四分之一计算量达到，最小计算量的NaViT比对应ViT成本效益高五倍。附录A.1有实验设置详情。
• 性能提升因素：NaViT因保留宽高比、多分辨率评估受益，主要得益于在计算预算内显著增加处理的训练示例数量，通过可变分辨率示例采样与标记丢弃结合，得到可变大小图像，高效打包成类似原始模型的序列长度，后续将对这些因素进行消融实验。

3.2 可变分辨率的好处

• NaViT “有效分辨率”：因保留原生宽高比，NaViT的“分辨率”指与给定分辨率正方形图像面积相同的图像，即“有效分辨率”，推理成本与同面积正方形图像上的普通ViT相同。

• 可变分辨率预训练：

  • 低分辨率图像处理FLOPs少，固定分辨率训练存在处理速度与处理高分辨率图像细节能力的权衡。
  • NaViT可混合低、高分辨率图像，发挥两者优势。
  • 图5对比两种NaViT变体在不同分辨率下的训练情况（同FLOPs），一种是固定分辨率$R=R_{max}$ ，另一种是可变分辨率$R\sim U(64,R_{max})$，结果显示可变分辨率模型性能优于单一分辨率训练模型，在固定分辨率最佳情况（训练与评估分辨率相同）下，可变分辨率模型性能相当或更优。

• 可变分辨率微调：

  • 过往方法：先前研究在预训练后期或微调时提高分辨率，虽提升质量但成本增加。
  • 实验操作：对NaViT和ViT在不同固定分辨率下微调，同时对NaViT进行可变分辨率微调。
  • 实验结果：在ImageNet - 1k数据集微调预训练模型，预训练性能提升在各分辨率下转移良好，且有两个关键发现：一是可变分辨率微调的NaViT（“NaViT 64:512”）与单一分辨率微调的NaViT性能相当且优于单一分辨率的ViT，无需选定单个下游微调分辨率；二是低分辨率（64）微调的NaViT在高分辨率评估时性能良好，实现经济适配，证实相关研究成果。

• 分辨率采样策略：

  • 采样方式：示例打包使多种分辨率采样策略可行，可选择对目标边长（平均高/宽，$R$）或目标面积（序列长度$R^2$）采样，均匀分布采样边长偏向低序列长度，采样面积偏向高边长。
  • 具体操作：从支持范围[-1, 1]的分布$D$为每个图像采样$u$，线性缩放$u$到[64, 384]采样边长，或到$[64^2,384^2]$采样面积。考虑四种分布$D$：均匀分布$U(-1,1)$、截断标准正态分布$N_t(0,1)$、偏向低分辨率的$N_t(-0.5,1)$和偏向高分辨率的$N_t(0.5,1)$。
  • 结果结论：图7表明最佳分辨率重采样策略优于默认分辨率，直接采样边长优于采样面积，偏向较低值的截断正态分布最佳，二者都通过优先采样小序列提高处理速度。

3.3 可变标记丢弃的好处

1. 标记丢弃策略：
   • 实验了连续采样的标记丢弃率和与分辨率相关的标记丢弃率（具体解释见附录A.4）。
   • 图9a显示，从Beta分布采样的可变丢弃率比恒定丢弃率性能更好。
   • 图9b表明，对于训练分辨率 $R\sim U(64,384)$ 且丢弃率在 $[0.5-\delta ,0.5+\delta ]\propto R$ 范围内缩放的模型，与分辨率相关的标记丢弃率相比Beta分布能进一步提升性能。
2. 调度标记丢弃率：
   • 打包操作便于在训练中改变标记丢弃率。
   • 通过调整丢弃率可平衡处理图像数量与单张图像信息利用，在不改变总训练成本下最大化最终准确率。
   • 依据已处理图像数量改变丢弃率（细节见附录A.5），图8显示在NaViT - B/16的JFT预训练中，降低标记丢弃率能提升性能。

3.4 位置嵌入

图10：分解式位置嵌入提升了对新分辨率和宽高比的泛化能力。(a) 各分辨率下的最佳准确率（浅色）和平均准确率（深色）。(b) 相对于分辨率256归一化后的准确率。

1. 评估对象与关注点：对2.1节引入的分解式嵌入及其设计选择进行评估，既关注绝对性能，也关注模型对训练分辨率范围外分辨率的外推能力。
2. 实验设置：
   • 在JFT上对NaViT - B/16模型训练200,000步，训练分辨率为 $R\sim U(160,352)$ 。
   • 将其与在固定分辨率256下训练相同数量图像的ViT - B/16模型对比，后者在新分辨率评估时采用标准位置嵌入插值方法。
3. 实验结果（图10）：
   • 分解式方法性能优于基线ViT模型和Pix2struct中的学习型2D嵌入。
   • Pix2struct中的学习型2D嵌入在泛化到更高分辨率时存在困难，可能因需大量未见过的 $(x,y)$ 对。
   • 分解式嵌入以相加方式组合效果最佳。

3.5 NaViT性能的其他方面

1. 分布外泛化
   • 评估方式：用JFT预训练的NaViT模型直接在下游数据集评估，包括ImageNet及其鲁棒性变体数据集，与计算量匹配的ViT基线模型对比。
   • 结果：图11显示NaViT在各数据集表现出色，在ImageNet - A上远超ViT，但在ObjectNet上优势缩小，可能因保持宽高比的中心裁剪对两数据集作用不同，不裁剪时差异更显著（附录图20）。
2. 校准
   • 评估内容：对ImageNet1K微调的模型，量化计算不确定性质量，计算改变每示例补丁数量时顶级预测的预期校准误差。
   • 结果：当每图像补丁数量在[128, 1024]变化时，校准误差在（0.045, 0.047）稳定，无需事后校准（附录E详细）。
3. 推理权衡
   • 评估内容：鉴于模型灵活性，探讨在给定计算预算下最大化总体准确率的方法，通过Cloud TPUv3芯片延迟量化。在线推理需确定为每个示例分配固定数量标记的算法。
   • 结果：图1显示NaViT比ViT权衡更好，收益递减明显，少量补丁也有竞争力。附录C研究级联方法提供帕累托最优模型和更多权衡机会。
4. 公平性信号标注
   • 研究目的：探究NaViT能否克服公平性信号标注中因标注不准影响组校准等指标的问题。
   • 实验方法：在FairFace和CelebA数据集训练标注员作为线性探针，用NaViT或ViT冻结特征，比较准确率。
   • 实验结果：
     • NaViT提供更高质量表示，处理正方形图像时提高公平性信号标注准确率，Wilcoxon符号秩检验验证显著（$p=3\times 10^{-4}$ ）。
     • 应用Inception风格裁剪后调整图像面积，NaViT保持原始宽高比提高准确率，95%置信水平下显著（$p=0.02$ ）。附录G有完整图表。

3.6 其他下游任务

1. 语义分割
   • 实验设置：在ADE20k数据集对NaViT进行语义分割微调，遵循Segmenter线性解码器协议，以ViT - L/16为基线，两模型在JFT - 4B预训练且计算预算可比。微调时，ViT输入随机正方形裁剪（$R_{max}\times R_{max}$），NaViT输入随机调整大小（保持宽高比，总像素数$R_{max}^2$），评估在原始分辨率进行。
   • 实验结果：图13显示，相同最大微调分辨率$R_{max}$时，NaViT性能优于ViT，如$R_{384}$的NaViT速度比$R_{512}$的ViT快两倍且性能更优。NaViT优势在于训练分辨率灵活，可处理不同形状图像。
2. 目标检测
   • 实验设置：使用计算量匹配的NaViT和ViT作为OWL - ViT - L/14目标检测器骨干，遵循其训练和评估协议。
   • 实验结果：表1表明，基于NaViT的检测器在常见和未见的LVIS “罕见” 类别上表现更优。虽预训练时间短导致绝对性能低，但相对差异显示NaViT为细粒度视觉任务提供强大表示。
3. 视频分类
   • 实验设置：因Transformer处理视频有挑战，NaViT可不同分辨率和时间长度训练以缓解。通过提取三个时空补丁 “tubelets” 对JFT训练的NaViT进行Kinetics400分类微调，扩展位置嵌入含时间维度，用 “中心帧嵌入” 初始化嵌入内核。
   • 实验结果：NaViT - L在约少6倍训练轮数下，与ViViT - L（80.4%）性能相当，且无需多裁剪评估。使用的Kinetics400数据集数据量少于先前工作。

4. 相关工作

1. 灵活的视觉Transformer
   • FlexiViT开发新颖内核调整大小方法，借支持多种补丁大小的模型实现可变“分辨率”，并研究其与蒸馏和神经架构搜索的关系。
   • Pix2struct通过新位置嵌入模式支持可变宽高比，在非自然图像任务中展现显著效率和性能提升。
2. 多尺度视觉Transformer
   • 在分割、检测等定位任务中，使用多空间尺度特征图常见，不少工作开发类似视觉Transformer，但也有人对简单定位任务使用该方法的必要性存争议。
   • NaViT未构建层次化多尺度表示，认为结合自身独特灵活性与该模型家族优势有研究潜力。
3. 混合分辨率加速训练
   • 多数考虑分辨率的图像建模工作聚焦用固定低分辨率加速预训练，如FixRes在预训练后期提高分辨率，后续很多工作沿用，但该方法不可逆，调高分辨率后无法降分辨率减少计算量。
   • 多网格训练通过在训练早期用“粗糙”时空分辨率大批量数据，后期用“精细”分辨率加速视频建模，需分层网格采样调度和学习率缩放。相比之下，Patch n’ Pack整合混合分辨率更轻松，无需复杂调度或训练管道。
4. 标记丢弃提高效率
   • 早期研究探索随机标记丢弃，后续研究表明结构化或基于“重要性”的策略有优势。
   • Patch n’ Pack避开固定小批量形状限制，而此前相关研究受此制约。

5. 结论和未来工作

作者通过本文证明，Patch n’ Pack，即简单地将序列打包应用于视觉Transformer，能够显著提高训练效率。由此产生的NaViT模型在推理时可以应用于多种分辨率，并且可以低成本地适配新任务。Patch n’ Pack开启了许多之前因需要固定批次形状而受阻的研究方向，包括自适应计算以及提高训练和推理效率的新算法。

热门专栏

机器学习

机器学习笔记合集

深度学习

深度学习笔记合集