NaViT：训练任意分辨率和长宽比的 ViT
TL; DR：修改 ViT 的位置编码，使得模型布局支持任意分辨率、任意长宽比的输入。同时使用 sequence packing 和 masked self-attention 来将多个样本打包到一个序列内并制止彼此干扰，实现高效的边长序列训练。
弁言
固定分辨率输入图像在 CV 范畴一直是一个问题。ViT 的出现，将图片进行图块化，得到 1D token 序列进行同一处理的这种情势，使得我们有望降服这一问题。我们可以参考 NLP 在 Transformer 的序列建模方面的一些履历，应用到 ViT 上，实现任意分辨率、任意长宽比的高效训练上。
起首，在模型布局上，原始 ViT 无法支持任意分辨率输入图像的根本原因是在于其位置编码是可学习的 1D 位置嵌入，只能处理固定分辨率、固定长宽比的输入。作者起首对 ViT 的位编码进行了修改。但是今天来看更好地的 ViT 位置编码方式或许是 2D RoPE。
其次，在支持任意分辨率输入后，不同图片在同一个 batch 内，为了进行并行计算需要 padding 以保持序列长度雷同，假如 batch 内图片分辨率差异过大，会造成大量的计算资源浪费。NaViT 鉴戒自然语言处理中的处理方式，提出训练时将多张图片打包到同一个序列中，从而制止浪费，提高可变长序列的训练效率。同时为了制止同 sequence 内的不同图片相互干扰，NaViT 还在 self attention 和终极的输出 pooling 上为每张独立图片添加了对应的 mask。
最后，作者提出了 continuous token dropping 和分辨率采样的训练策略，进一步提高模型的训练效率和在可变分辨率上的终极性能。
模型布局适配
1 位置编码
起首，在模型布局上，ViT 需要支持变长的输入 token 序列。原始的 ViT 显然是无法做到这一点的，它只能处理方形、固定分辨率 ( R , R ) (R,R) (R,R) 的输入图片，由于其位置编码接纳的是可学习 1D 位置嵌入，如许在训练完成之后，原始 ViT 只能处理特定分辨率的输入图像，假如想要拓展到其他分辨率，需要对学习到的位置嵌入进行插值。
Pix2Struct 在训练时学习尺寸为 [ maxLen , maxLen ] [\text{maxLen},\text{maxLen}] [maxLen,maxLen] 的 2D 位置嵌入，对每个图块位置通过 ( x , y ) (x,y) (x,y) 进行二维索引，如许就能在推理时处理分辨率最大为 ( P ⋅ maxLen , P ⋅ maxLen ) (P\cdot\text{maxLen},P\cdot\text{maxLen}) (P⋅maxLen,P⋅maxLen) 的输入图片。然而这需要在训练时见到过全部 ( x , y ) (x ,y) (x,y) 组合，也比较麻烦。
NaViT 考虑了拆分的位置编码（factorized positional embedding），即将图片的长、宽两个位置维度分开来进行编码，使用两个分离的嵌入 ϕ x , ϕ y \phi_x,\phi_y ϕx,ϕy 来表现位置 ( x , y ) (x,y) (x,y) 处的位置编码，然后将它们加和起来。作者考虑了绝对索引嵌入 ϕ ( p ) : → R D , p ∈ [ 0 , maxLen ] \phi(p):\rightarrow\mathbb{R}^D,\ p\in[0,\text{maxLen}] ϕ(p):→RD, p∈[0,maxLen] 和相对位置嵌入 ϕ ( r ) : → R D , r = p / sideLen ∈ [ 0 , 1 ] \phi(r):\rightarrow\mathbb{R}^D,\ r=p/\text{sideLen}\in[0,1] ϕ(r):→RD, r=p/sideLen∈[0,1]。后者可以给出与图片绝对尺寸无关的位置编码，但是会对原始长宽比造成肯定程度的混淆。对于位置编码的详细情势，作者也考虑了可学习、正弦和傅里叶三种情势。
2 掩码自留意力和掩码池化
在模型布局支持了任意分辨率输入图像后，我们在训练时需要对不同序列长度的输入图片进行 padding，以包管 batch 内序列长度相当。这会造成很多的计算资源浪费。作者鉴戒 NLP 中的做法，将多张图片打包到同一个系列中，以只管减少 padding。
但是我们将多个样本的 token 打包进了同一个序列中，需要包管不同样本彼此之间没有影响。在 ViT 中，会导致 token 之间彼此影响的地方就是自留意力层，因此我们在 ViT 的自留意力计算时，根据 token 所来自的不同样本设置掩码，制止来自不同样本的 token 相互 attend 到。此外，ViT 会采取对各 token 进行 pooling 的方式来得到终极的单个输出 embedding，在这里我们也需要进行 pooling，对来自同一个样本的输出 token 进行 pooling，每个样本终极得到单个 embedding。
训练策略
在模型架构适配后，利用支持可变分辨率的特性，NaViT 又提出了两个训练策略，进一步提高模型的训练效率和在可变分辨率上的终极性能。
1 continuous token dropping
token dropping，即训练时随机丢弃一些输入图块，可以用来加快训练。但之前的方法都需要对全部样本丢弃雷同的比例，NaViT 接纳 sequence packing 之后我们可以进行 continuous token dropping，对不同的样本的丢弃比例可以不同。如许就可以在使用 token dropping 来提高训练效率的同时，让一些样本保持是完备的图片，从而减小训练-推理时的差异。
2 分辨率采样
在传统的ViT中，我们需要在训练/推理效率和模型性能之间进行权衡，要么接纳高分辨率图像，终极性能更强但是训练和推理开销更大，要么接纳低分辨率图像，开销降低但是性能稍差。一样平常我们会分阶段训练，先在较小的分辨率下预训练，然后在更高的分辨率下微调。
NaViT 更加灵活，我们可以随时通过从图像大小分布中采样来进行混合分辨率训练，同时保存每张图像的原始长宽比。这使得 NaViT 团体吞吐更高，且也能见到高分辨率图像，团体在训练推理耗时和模型性能上相比原始 ViT 有较大的提高。而且模型可以顺应大范围内的可变分辨率图像。
效率分析
self attention 开销
我们将多张图片 packing 到一个序列中，会导致序列长度大幅增长。如许一个最直接的担心就是 self attention 中 O ( n 2 ) \mathcal{O}(n^2) O(n2) 的复杂度。但是作者指出，随着隐层维度的增长，NaViT packing 带来的额外开销是不绝减小的。下图展示了这种趋势。除了速度外，长序列另一个需要担心的是空间复杂度，作者说有 Flash Attention 这也不成问题。
Packing, and sequence-level padding
我们需要包管 packing 了多张图片的终极序列在 batch 内的长度是雷同的，作者接纳了一种贪心的策略来进行 packing，不外一样平常没办法完美地 packing 到固定的序列长度，因此或多或少还是得 padding 一点 token。假如想进一步做得更细致一点，可以设计方案动态地选择分辨率和 token dropping 的个数，来实现完美 sequence packing，不外作者说目前的实现 padding token 只占 2%，已经是一个简单且不错的方案了。
Padding examples and the contrastive loss.
做了 sequence packing 之后，计算 token-level 的损失是很直接的，但是计算 example-level 的损失比较麻烦，而很多视觉任务都是基于 example-level 的损失。起首，我们要进行上面提到的 masked pooling，来对可变长度的 token 序列进行 pooling。每个 sequence 中的 example 个数不同，假如我们要固定 batch size，需要设置一个 E max E_\text{max} Emax，每个序列中的样本凌驾该值则丢弃，少于该值则用 padding token 的表征。
这类似对比损失中的一个问题，对比损失计算在时间和内存上的规模约为 O ( n 2 ) O(n^2) O(n2)。为了制止这种情况，可以使用 chunked constrastive loss（这里引的 paper 没搜到，我估计就是 局部 softmax + block online softmax 全局合并？），各个设备先在本地计算，然后累积全局 softmax 归一化所需的统计数据，这就制止了全局 softmax 需要网络全部数据。这使得我们可以设置很高的 E max E_\text{max} Emax（从而有效使用模型编码器），而损失计算不会成为瓶颈。
总结
仅从布局上来说，ViT 要支持任意分辨率、任意长宽比输入图像，只要位置编码是可外推的绝对位置编码就行了，现在来看比较常用的是苏神的 2D RoPE，Qwen 2/2.5 VL 中就接纳了这种方式。NaViT 紧张是将 NLP 训练中的 sequence packing 引入了过来，将多个长度不一的 token 序列放到一个 sequence 里，并对 self attention 进行 masking 制止不同图片之间的干扰。（不外我看大部分训练库都没实现这个？不知道那里能找到参考实现
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
