突破经典网格特征？AutoFocusFormer: Image Segmentation off the Grid 论文阅读笔记

写在前面

  这一周赶上五一五天假了，朋友们出去 happy 了吗？有没有赶上人山人海的热闹？反正我只是在 5.1 那天出去走走，哈哈。
  这是一篇关于实例分割的文章，所解决的问题在于实例分割中需要的小目标像素分辨率太低，于是本文提出一种自适应下采样的方法来解决这一之前没有人想到要做的隐藏问题。

• 论文地址：AutoFocusFormer: Image Segmentation off the Grid
• 代码地址：https://github.com/apple/ml-autofocusformer/tree/main
• 收录于：CVPR 2023
• 温馨提示：这篇论文写作废话太多，抓不住重点，此篇写到一半结束了，但可作为反面教材学习避免一些糟糕的表达。
• TiP：2023 每周一篇博文，主页更多干货输出，欢迎浏览、关注。粉丝量破 3K，期待您的参与~

一、Abstract

  真实世界的图像内容总是密度不平衡分布，即大块的蓝天白云，小块的目标物体，如上图所示。经典的卷积网络对所有图像内容一视同仁。但由于小目标占据的空间位置非常小，以至分割等任务比较糟糕。于是，在下采样中保留更多的表示小目标的像素能够有助于留存重要的信息。本文提出一种局部注意力 Transformer 图像识别 Backbone，即 AutoFocusFormer (AFF) 方法，执行自适应下采样来学习保留最重要的像素信息。由于自适应下采样会产生一些不规则的像素分布，因此放弃了经典的grid结构，而提出了一种新的基于点的局部注意力 block，由一个平衡聚类模块和一个可学习的聚类融合模块组成。实验表明 AFF 的性能很棒。

二、引言

  首先指出现实世界图像内容分布通常不平衡，经典的卷积网络平等地对待图像上的内容，而忽视了不同目标的重要性。最近的 Transformer 方法虽然并不是严格依赖于 grid 结构，但是仍采用基于 grid 的技术，例如步长为 16 的卷积和 $7\times 7$ 的局部注意力滑动窗口。
  尽管卷积网络被普遍使用，但均匀下采样对于需要像素水平的细节信息例如分割任务，就力不从心了。于是首先想到的解决方法是提高输入图像分辨率，这样下采样后的特征图分辨率也随之提高，但是计算内存和计算量也随之增加。还有一些方法在分割解码器中采用不规则的采样点来尝试解决这一问题，但是仍然依赖于均匀下采样编码器，留下了一些均匀采样的隐患。
  为解决这一问题，需要模型能够在每张图像上非均匀采样，即 保留重要的细节信息同时压缩那些不太重要的区域。有一些方法简单的应用全局注意力机制进行自适应下采样，但全局注意力不会放缩到远高于 ImageNet 的分辨率，从而无法适应于类似图像分割的任务。
  为满足上述要求，本文提出提出第一个端到端的、连续自适应下采样阶段的网络，AutoFocusFormer (AFF)。为缩放高分辨率图到分割的要求，AFF 采用局部注意力 Block。为定义不规则采样 tokens 的局部注意力区域，提出一种平衡聚类算法，利用空间填充曲线将这些不规则的位置分组到区域，提出一种自适应下采样模块通过可微分的区域融合过程来学习不同图像位置的重要性。最终，修改分割头以使得能够适应 backbone 生成的不规则空间表示。
  本文提出的方法通过保留任意数量的token从而能拥有大量的计算预算，总体而言，贡献如下：

• 第一个引入端到端的、有着弹性下采样率的连续自适应下采样；
• 提出一组平衡聚类算法将 tokens 分组到区域内；提出一组聚类融合模块来启用端到端的自适应下采样学习；
• 应用本文提出的算法到一些 SOTA 的框架内，验证了方法的泛化性；
• 实验结果表明本文提出的方法能够显著地识别实例分割任务中的小目标。

三、相关工作

视觉 Transformer Backbones

  ViTs 模型缺乏稠密预测任务所需要的级联特征图，于是 MViT、PVT 出世了。但仍然是在全局注意力图上进行操作，因此所需要的计算成本随着输入图像的尺寸增大而增加。解决办法之一是应用低分辨率的特征图，例如 BoT、LeViT 等。另外一些方法修改注意力操作：Seg-Former、PoolFormer、AFT、Swin Transformer、SASA。

基于聚类的注意力

  基于稀疏注意力的聚类方法已经应用在视觉语言领域了，但是这些方法并未涉及自适应下采样。大多数方法尝试聚类学习到的特征，而本文提出基于 token 位置的聚类。有一些方法将聚类应用在解码器中，不同于这种，本文将聚类应用在编码器中。

自适应下采样

  有很多尝试结合自适应下采样与 ViTs 的方法，例如 AdaViT、DynamicViT、A-ViT、Dynamic Grained Encoder、EViT、PS-ViT。但这些方法仍然基于全局注意力图，因此不能缩放到高分辨率的分割任务，而是仅关注与加速 ImageNet 上的分类。有一些方法关注于自适应下采样策略的设计，总体趋势则是发现 tokens 的融合比 tokens 的减少似乎更自然。相应的，本文提出一种区域融合模块，能够自适应地选择融合位置，直接从损失中获得重要 tokens 的得分。

点云网络

  有一些 3D 点云方法，直接在一系列不规则的点上进行操作，例如：PointNet++、PointConv、Point Transformer。这些方法总是选择 K 个最近邻的或者为每个点找到若干个最近聚类中心。本文利用 PointConv 作为解码器模型来代替 $3\times 3$ 卷积。

四、方法

  如上图所示：本文提出的 Backbone 模型主要由 patch embedding 模块和 4 个阶段以及一个特定任务头组成。其中每个阶段由一个聚类算法、一些局部注意力 Transformer 块、自适应下采样模块组成。

4.1 聚类和区域

  原始的聚类方法有着平方复杂度，于是一些方法通过引入对给定的点进行初次 k-means 聚类来加速 KNN。本文同样采用聚类来定义局部区域，即将 tokens 划分为类，将区域定义为包含多个邻近的聚类。但是传统的聚类方法不能直接完成这一目标，原因在于这类方法通常需要迭代或者多轮采样导致速度太慢，并且聚类中心有着不同数量的赋值点，于是在计算时只能用 0 填充到同一维度从而进行批量运算，这就占据了大量的内存和时间上的浪费。
  于是本文提出一种新颖的平衡聚类方法。

4.1.1 平衡聚类

  通过空间填充曲线将 2D 坐标转化为 1D 坐标，同时保留 2D 距离信息，然后将这些坐标划分到等大小的组内从而得到一个平衡的聚类，于是在线上靠近的点在 2D 空间中也相当接近。然而从 2D 转为为 1D 的过程中，很难完整地保存 2D 的信息。为减轻这一丢失问题，本文采用一种两阶段的处理过程。这一想法实质上是尽在粗糙的层次上利用空间填充曲线来获得稀疏且规则的 2D 样本位置，然后基于 2D 距离来排列 tokens 的位置。
  确切来说，首先将图像划分为粗糙且规则的 Patch 块，其数量与打算聚类的类别数量相同。每一个 patch 的中心记为空间填充锚点，空间填充曲线在锚点之间创建序列。于是在这序列上位置为 $p\in {\mathbb{R}}^2$ 处，且属于锚点 $a_i\in {\mathbb{R}}^2$ 的 token，同时定义前面的锚点 $a_{i-1}$ 和 $a_{i+1}$。计算 token p 到这两个锚点的距离比 $r$ 如下：
$$r(p)=\frac{d_{i-1}(p)}{d_{i+1}(p)}=\frac{\Vert p-a_{i-1}{\Vert }_2}{\Vert p-a_{i+1}{\Vert }_2}$$在每个 pathc 上，通过对 $r$ 进行降序排列，就能得到一组有序的 tokens。如下图所示：

注意这一聚类算法在每个阶段开始时仅需要执行一次，而所有注意力 bolock 和结尾使用的下采样模块在这一阶段中共享聚类信息，但这一平衡聚类算法仅适用于低维度的点。

4.1.2 聚类的区域

  为鼓励信息在整个图像上流动，那么注意力就不能仅仅局限于相同的聚类上，例如 Swin-Transformers 通过滑动窗口使得注意力能够关注不同层上的不同区域。但是在本文中，每一层的重聚类会增加计算量。于是采用更小的聚类，并使得每个 token 关注周围 $R$ 个邻近的聚类。

结语

  这可能是第一篇写到最后没太完成的博文，论文中的写作一言难尽，废话太多，抓不住重点，遂到此为止，及时止损。这篇论文的写作可以算得上反面教材了。