FlexiViT

一个标准的 ViT 模型，训练时候的 Patch Size 是多大，它只能够在那个 Patch Size 的情况下取得良好的性能。当 Patch Size 改变时，一般模型就需要重新训练。为了克服这一限制，本文提出了一个可以适应所有 Patch Size 的 ViT 模型FlexiViT。

一个适应所有 Patch 大小的 ViT 模型

论文名称：FlexiViT: One Model for All Patch Sizes

论文地址： https://arxiv.org/pdf/2212.08013.pdf

视觉 Transformer (ViT) 和 CNN 模型不同的一点是对输入图片处理方式的不同：ViT 的处理方式是分块化 (Patchification) 将图像切割成不重叠的 Patches，并对这些 Patches 通过线性映射得到的 tokens 进行后续的运算。而 CNN 通过一些密集的，重叠的卷积来处理输入图片，使之转化成特征图。那么我们可以发现，ViT 这种与 CNN 完全不同的处理图片的范式解锁了新的功能。比如，我们可以随机丢弃一些不重要的 tokens 来提升模型的效率，额外增加一些 tokens 来适应新的任务，或者混合一些不同模态的 tokens。

尽管分块化对 ViT 模型很重要，但块的大小的作用却在之前的工作中很少受到关注。最初的 ViT 使用三种 Patch Size (32×32, 16×16 和 14×14 像素)，但许多后续工作将 Patch Size 固定在 16×16。在这项工作中，作者提出 Patch Size 在不改变模型参数量的前提下，为模型的计算复杂度和预测性能提供了一个简单而有效的平衡。

比如，一个 ViT-B/8 模型在 156 GFLOPs 和 85 M 参数的情况下，在 ImageNet-1K 上的 Top-1 准确率为 85.6%，而一个 ViT-B/32模型在 8.6 GFLOPs 和 87 M 参数的情况下，在 ImageNet-1K 上的 Top-1 准确率仅仅为 79.1%。尽管在性能和计算复杂度方面存在巨大的差异，但是这些模型在本质上具有相同的参数。一个标准的 ViT 模型，训练时候的 Patch Size 是多大，它只能够在那个 Patch Size 的情况下取得良好的性能。当 Patch Size 改变时，一般模型就需要重新训练。

为了克服这一限制，本文提出 FlexiViT，希望在不增加成本的前提下，训练一个可以适应所有 Patch Size 的 ViT 模型。为了训练 FlexiViT，作者在训练过程中随机化 Patch Size，并对位置编码参数和 patch embedding 参数做 Resize 操作。这些简单的修改已经足以获得强大的性能，此外，作者借助 KD 获得更好的结果。

标准 ViT 对于 Patch Size 灵活吗

 

图1：FlexiViT 方法框架

可以看到，标准 ViT 模型的性能随着 Patch Size 偏离预训练 Patch Size 的大小程度越大，ViT 性能的下降也越大。因此，标准 ViT 对于 Patch Size 是不灵活的，作者也因此开发了 FlexiViT，希望解决这个问题。 

图2：一个预训练好的标准 ViT 模型，在不同的 Patch Size 下会产生较差的性能

对于 Patch Size 更灵活的 FlexiViT 模型

如图2中的实线所示是 FlexiViT-B 模型的性能，可以看到性能与 ViT-B/16 和 ViT-B/30 接近，且当 Patch Size 是16和30除外的其他数时，性能均优于 ViT-B/16 和 ViT-B/30。FlexiViT-B 模型的训练方法与 ViT-B/16 和 ViT-B/30 一致，但是不同的是，在训练的每一步，Patch Size 的大小都是从一组预定义的 Patch Size 中均匀随机选择的。只需要对模型和训练的代码做以下2处修改即可。

注意，改变 Patch Size 大小不等于改变图片分辨率的大小，这二者是完全不一样的两回事。Patch Size 大小纯粹是对模型的改变，而改变图像大小可能会大大减少可用信息。以上2处的改进，如下列算法所示。

登录后复制

model = ViT(...)
for batch in data:
  ps = np.random.choice([8, 10, ..., 40, 48])
  logits = model(batch["images"], (ps, ps))
  # [...] backprop and optimize as usual

class ViT(nn.Module):
  def __call__(self, image, patchhw):
    # Patchify, flexibly:
    w = self.param("w_emb", (32, 32, 3, d))
    b = self.param("b_emb", d)
    w = resize(w, (*patchhw, 3, d))
    x = conv(image, w, strides=patchhw) + b
    # Add flexible position embeddings:
    pe = self.param("posemb", (7, 7, d))
    pe = resize(pe, (*x.shape[1:3], d))
    return TransformerEncoder(...)(x + pe)

• 1.
• 2.
• 3.
• 4.
• 5.
• 6.
• 7.
• 8.
• 9.
• 10.
• 11.
• 12.
• 13.
• 14.
• 15.
• 16.
• 17.

如何改变 Patch Embedding 的尺寸

图3：PI-resize 在模型输入和 Patch size 变化时可以维持性能

与知识蒸馏的关系

知识蒸馏中，一个通常较小的学生模型被训练来模仿一个通常较大的教师模型的预测。与标准标签监督训练相比，这可以显著提高学生模型的性能。使用 FlexiViT，可以用一个强大的 ViT 教师的权重初始化一个学生 FlexiViT，并显著提高蒸馏性能。本文从 ViT-B/8 模型初始化并且蒸馏得到一个 FlexiViT-B 模型。在初始化时，作者使用 PI-resize 调整教师模型的 Patch Embedding 权重为 32×32，并双线性地重新采样其位置编码为 7×7。然后按照 FunMatch 方法训练学生模型，最小化教师和学生 FlexiViT 之间的 KL Divergence。

FlexiViT 的内部表征

FlexiViT 有没有以与 ViT 相似的方式处理不同 Patch Size 的输入？

作者通过分析模型的内部表征来研究这个问题。作者应用了 Centered Kernel Alignment (CKA) 方法，这是一种广泛用于比较神经网络内部和跨模型表征的方法。可视化结果如下图3所示。从 Block1 到 Block 6 的 MLP 层，特征映射的表征相似。在 Block 6 的 MLP 层，开始发散，然后在最后一个 Block 再次收敛。

图4：FlexiViT 的内部表征的 t-SNE 可视化结果。不同大小的点代表不同的 grid size，不同颜色代表不同的层   

实验：使用预训练的 FlexiViT 模型

作者使用预训练的 FlexiViT 模型在下列迁移学习任务中验证 FlexiViT 的性能。实验结果如下图4所示。在不同的任务集中，FlexiViT 模型大致与两个固定 Patch Size 的 ViT 模型的性能相匹配，在大 Patch Size 时几乎不落后，在小 Patch Size 时有改进。这些结果证实，与针对不同 Patch Size 预训练多个 ViT 模型相比，使用预训练的 FlexiViT 没有明显的缺点。FlexiViT 提供了一种新的方式，使迁移学习更有效地利用资源，节省加速器内存和计算量。我们可以使用大的输入 Patch Size 进行廉价的迁移学习，但稍后使用小的 Patch Size 部署得到的模型，以获得不错的性能。

图5：迁移学习实验结果

如下图5左侧所示，是微调 grid size 为8的 FlexiViT 模型，然后在其他 grid size 下评估的结果。标准 ViT 模型的精度很快就会下降，而 FlexiViT 随着网格尺寸的增加表现出极大的性能提升。

如下图5右侧所示，是一个 FlexiViT-B 模型在三种不同的 grid size 下进行了微调，并在不同的 grid size 下进行了评估。可以看到，无论FlexiViT-B 模型在哪种 grid size 下做微调，都能在更大的 grid size 下正常工作。比如，可以在 8×8 的 grid size 小下执行相对低计算代价的微调。当在 8×8 grid size 下评估时，模型达到 81.8% 的精度，但当在 24×24 grid size 下评估时，它达到 85.3% 的精度。

图6：FlexiViT 更灵活的迁移学习

作者进一步表明，FlexiViT 的预训练模型在迁移到下游任务时也更加灵活。如下图6所示是 SUN397 数据集的实验结果。在迁移学习实验中，作者随机化 Patch Size 的大小，并在不同的 Patch Size 大小下评估单个模型。FlexiViT 的迁移效果最好。当 Patch Size 的大小发生变化时，FlexiViT 的灵活性最好。

 

图7：FlexiViT 的预训练模型在迁移到下游任务时也更加灵活

如下图7所示，作者对 FlexiViT 的 attention 做了一些分析。可以观察到，attention relevance 在不同的 Patch Size 大小下会发生显著变化。随着 Patch Size 的减小，注意力会集中到更多的小区域。

图8：attention relevance 可视化结果

总结

一个标准的 ViT 模型，训练时候的 Patch Size 是多大，它只能够在那个 Patch Size 的情况下取得良好的性能。当 Patch Size 改变时，一般模型就需要重新训练。为了克服这一限制，本文提出 FlexiViT，希望在不增加成本的前提下，训练一个可以适应所有 Patch Size 的 ViT 模型。为了训练 FlexiViT，作者在训练过程中随机化 Patch Size，并对位置编码参数和 patch embedding 参数做 Resize 操作。这些简单的修改已经足以获得强大的性能，此外，作者借助 KD 获得更好的结果。

 

原创作者: whaosoft 转载于: https://blog.51cto.com/whaosoft/11611752