TAP：对任何事物进行分割、识别和描述

本文链接：https://blog.youkuaiyun.com/sxlsxl119/article/details/144964086

TAP:Tokenize Anything via Prompting，发表于2023年。

一，TAP：

摘要

我们提出了一种统一的、可提示的模型，能够同时对任何事物进行分割、识别和描述。与 SAM 不同，我们旨在通过视觉提示构建一个通用的多功能区域表示。为实现这一目标，我们使用大量的分割掩码（如 SA-1B 掩码）以及来自具有 50 亿参数的预训练 CLIP 模型的语义先验来训练一个可泛化的模型。具体来说，我们通过给每个掩码添加一个语义标记来构建一个可提示的图像解码器。语义标记负责在预定义的概念空间中学习语义先验。通过对掩码标记上的分割和语义标记上的概念预测进行联合优化，我们的模型展现出强大的区域识别和定位能力。例如，一个从头开始训练的38M参数的因果文本解码器在 Visual Genome 区域标记任务上以 164.7 的 CIDEr 分数创下了新纪录。我们相信这个模型可以成为一个通用的区域级图像标记器，能够为广泛的视觉感知任务编码通用区域上下文。代码和模型可在 https://github.com/baaivision/tokenize-anything获取。

1. Introduction

视觉感知的一个关键目标是有效地定位和识别任意感兴趣区域。这需要一个单一的视觉模型，它能够理解区域上下文并同时执行诸如分割、识别和字幕生成等感知任务。然而，现有的模型往往侧重于定位与类别无关的掩码，例如 SAM及其注重效率的后续模型 [50,64,70]，或者仅提取视觉语义，例如 CLIP及其区域级变体 [51,67,74]。具体而言，SAM 开发了一种分割基础模型，可以通过提示分割任何事物，在像素级定位任务中实现了强大的泛化能力。另一方面，CLIP 通过在网络规模的图像 - 文本对上进行对比学习训练了一个识别基础模型，在识别任务中展示出强大的零样本能力。因此，在 SAM 的架构内从 CLIP 模型学习语义先验为实现全面的视觉感知提供了一条有前途的途径。

我们的首要目标是构建一个统一的、可提示的模型，该模型能够一次性对任何事物进行分割、识别和描述（图 1a）。然而，构建这样一个基础模型并非易事，原因如下：

1）目前不存在能够实现通用感知的可提示框架；

2）目前没有公开的网络规模的同时包含掩码、类别和描述的数据集；

3）没有一种训练方法能够有效且高效地将 CLIP 和 SAM 的能力集成到一个单一模型中，以理解任意区域。

这项工作仔细探索了这个方向，并旨在提供一个包含新框架、新数据集和有效学习方法的系统解决方案。

我们首先引入一个可提示的标记化和字幕生成框架（图 3），该框架能够同时实现分割、识别和字幕生成。这需要一个统一的模型，该模型要能够在给出提示任意感兴趣区域的灵活提示时，抽象出通用表示，例如掩码标记和语义标记。我们遵循 SAM 的架构，但将其掩码解码器升级为通用图像解码器，为每个预测的掩码生成一个额外的语义标记。与 SAM 类似，掩码标记有助于像素级分割，而语义标记负责区域级语义预测。通过利用语义标记，该模型可以通过一个 MLP 头同时处理开放词汇分类任务，并通过一个使用自回归过程的轻量级文本解码器处理可提示的字幕生成任务。我们将这个模型称为 TAP，即“Tokenize Anything via Prompting”的缩写，如图 1b 所示。

训练这样一个高性能且可泛化的模型需要一个多样的大规模数据集。然而，目前并没有可用于同时进行分割和识别的网络规模数据源。SA - 1B在 1100 万张图像上构建了 11 亿个高质量掩码标注，用于训练诸如 SAM 之类的分割基础模型。相反，LAION - 2B从网络上收集了 20 亿个图像 - 文本对，有助于训练诸如 CLIP 之类的可泛化识别模型。为了解决缺乏对齐数据带来的挑战，我们引入了 SemanticSA - 1B 数据集（见图 1c）。这个数据集将来自 LAION - 2B 的网络规模语义隐式地集成到 SA - 1B 中。具体来说，对于 SA - 1B 中的每个分割区域，我们提取其在概念词汇表上的概念分布作为其语义先验，该语义先验是由在大量 LAION 图像 - 文本对上训练的强大 CLIP 模型预测得到的。因此，SA - 1B 数据与 LAION - 2B 先验共同构成了我们的预训练数据集。

利用 SemanticSA-1B 数据集，我们从一开始就使用真实掩码和相关语义对模型进行预训练，有效地将 CLIP 的能力整合到 SAM 的架构中。这只需通过同时针对通用分割和概念预测训练一个可提示的分词器便可实现。为了预测每张掩码图像的语义概念，我们进一步提出最小化预测概念分布与目标分布之间的 KL 散度损失，旨在最大限度地转移 CLIP 的知识。这种联合训练目标使得模型在定位和识别方面都能实现强大的泛化能力，从而有助于完成通用视觉任务。

我们对 TAP 模型及其组件进行了广泛的评估。TAP 在实例分类中展示出强大的零样本性能，例如在具有挑战性的 LVIS 基准测试中达到 59.1 的 AP，同时保持了具有竞争力的零样本分割性能，例如 TAP 和 SAM 的 AP 分别为 43.0 和 43.1。值得注意的是，在 Visual Genome 的区域字幕任务中，我们以 164.7 的 CIDEr 分数创下新纪录，且与之前的工作相比使用的参数显著更少。我们的研究结果表明，标记化的区域级特征对于分割和分类任务都是可泛化的，甚至可以直接提示因果文本生成。最重要的是，我们相信 TAP 模型可以成为一个通用的区域级图像标记器，能够为广泛的视觉感知任务编码区域上下文（见图 2）。

2 相关工作

2.1 视觉基础模型

视觉基础模型是当下计算机视觉领域极为关键的研究方向。这类模型旨在构建一种通用的、具有强大泛化能力的架构，以此处理各式各样的视觉任务，像是图像分类、目标检测、语义分割以及图像字幕生成等等。举例来说，SAM（Segment Anything Model）算是其中的佼佼者，它聚焦于分割任务，能借助提示技术，对几乎任何给定图像中的任意物体进行分割，构建起分割基础模型，在像素级别的定位任务里通用性极强。而 CLIP（Contrastive Language-Image Pretraining）走的是另一条路线，通过海量网络规模的图像 - 文本对展开对比学习，训练出识别基础模型，拥有出色的零样本识别本领，让模型在没见过的类别上也能做出识别判断。如今，不少研究都致力于融合这类各有所长的基础模型的优势，打造更为全面、功能更丰富的视觉感知体系，来应对现实场景中复杂多变的视觉需求。

2.2 开放词汇分割

与先前那些在有限词汇范围内工作的实例分割和语义分割模型不同，开放词汇分割（OVS）旨在对超出训练所用封闭词汇集的区域进行分类。大量研究工作专注于利用像CLIP这样预训练的视觉 - 语言模型（VLMs），并着重设计特定的对齐技术，以便有效地将VLM知识整合到现有分割模型中。例如，LSeg将文本和像素嵌入到一个公共特征空间中，为每个像素分配标签。MaskCLIP构建了一个两阶段模型，以便与CLIP视觉编码器无缝集成。ZegFormer将问题分解为一个类别无关的分组任务和一个区域级分类任务，以此利用VLM。通过利用字幕数据，一些研究以弱监督的方式将视觉特征与文本对齐。例如，GroupViT在没有像素级标注的图像 - 字幕对上进行训练，基于文本监督直接对掩码进行分组。OVSeg使用从图像字幕中的名词生成的伪标签，在带掩码的图像上对CLIP进行微调。另一方面，CGG结合了定位和生成损失，以深入挖掘图像字幕中的知识。此外，其他研究在单个网络中联合学习多个任务，或者研究文本到图像的扩散模型。我们的工作与基于CLIP的方法一致，但与两阶段模型有所不同，两阶段模型通常依赖图像级的CLIP来对掩码进行分类。相反，我们的方法侧重于开发一个具有区域级语义感知能力的单一模型。

2.3 零样本区域理解

先前的工作侧重于将视觉语言模型（VLMs）的开放词汇能力扩展到目标检测任务中。近期的研究则致力于融合CLIP在开放词汇分类方面的优势与SAM在分割方面的能力。例如，SAM-CLIP通过使用从两位“教师”（指SAM和CLIP）那里采样的数据子集重新训练视觉编码器，从SAM和CLIP中提炼知识，保留了CLIP和SAM各自的优势。RegionSpot通过添加一个在检测数据集上训练的适配器统一了提示方式，使得SAM的掩码标记能够与从带掩码图像片段导出的CLIP特征相互作用。一些工作试图构建能够识别任意区域中物体的统一模型。SEEM基于X-Decode构建而成，擅长处理各类提示，包括点击、边界框、涂鸦、文本以及参照图像片段。继SAM之后，ASM为SA-1B创建了一个新数据集（AS-1B），构建了语义标签、问答对和详细字幕的丰富标注。利用这个数据集，他们开发了一个用于全景视觉识别的新模型ASM。与这些依赖手工制作的多模态数据集的模型不同，我们充分利用来自SA-1B的大量分割掩码以及来自高性能CLIP模型的语义先验知识，旨在开发一个可提示的分词器，使其能够理解任意给定区域的语义上下文。

3 方法

我们引入了一种新颖的可提示框架，该框架能够有效地对任意感兴趣区域进行分割、识别和字幕生成。这是通过预训练一个可提示的标记器来实现的，该标记器利用带有 CLIP 先验的大量分割掩码（第 3.1 节），随后扩展模型的能力以包含可提示字幕生成的生成能力（第 3.2 节）。

3.1 可提示标记化

我们的主要目标是在一个可提示的分割模型中协调视觉与语言，使模型具备区域级语义感知能力。为实现这一目标，在本小节中，我们将介绍模型架构、预训练数据集，以及涉及概念预测与可提示分割的学习方法，还有预训练损失。

模型架构：我们的模型包含三个关键模块：图像编码器、提示编码器和图像解码器（见图3）。我们保留了SAM的架构，但将其掩码解码器升级为通用图像解码器。此外，为了更高效且有效地达成目标，我们对SAM的架构做了几处修改。具体来说，图像编码器采用标准的视觉Transformer（ViT），使用16×16的非重叠窗口。为了减轻计算强度，我们用卷积跨窗口模块替代图像编码器中的全局注意力机制，并用基于索引的相对位置偏差替换基于查询的相对位置嵌入。对于提示编码器，我们不会将前一阶段的掩码预测添加到图像嵌入中，因为这会在前提示（例如草图点）和高级提示（例如交互点）之间引入差异。因此，提示编码器中的所有掩码嵌入层都被移除了。在图像解码器中，我们给每个预测出的掩码额外添加一个语义标记。其中，掩码标记用于像素级别的分割，而语义标记则有助于区域级别的识别。所以，我们的图像解码器一共会生成4个掩码和9个标记：4个掩码标记、4个语义标记，以及1个IoU标记。

预训练数据集：传统的视觉 - 语言对齐方法依赖于图像 - 文本对，这限制了细粒度的区域理解。与先前依赖精心收集或近似区域 - 文本数据的方法不同，我们仅使用分割数据和 CLIP 先验将图像片段与语言对齐。由于 SA - 1B 是一个类别无关的数据集，我们利用高性能开源 CLIP 模型 EVA - CLIP，计算概念分布，将其作为 SA - 1B 内每个图像片段的语义先验。我们首先创建一个标签列表，该列表由从各种热门图像数据集收集的 2560 个类别组成。然后，我们使用一个简单的提示模板：“a {}” 或 “a photo of a {}”，利用 CLIP 生成文本嵌入。同时，对于来自 SA - 1B 的每个带掩码的图像片段，我们获取由 CLIP 生成的其视觉嵌入。概念分布可定义如下：

在这里，

表示温度参数。因此，分割数据及其现成的 CLIP 先验）

被本地存储，共同构成了我们的预训练数据集 SemanticSA - 1B。

概念预测：为了增强我们模型的语义感知能力，我们提议使用语义标记来预测区域概念。具体而言，我们利用语义标记获取预测的视觉嵌入

，它会进一步被投影到概念分布

上。

我们进一步提议对模型预测与 CLIP 目标之间的概念分布进行对齐。概念对齐损失可定义为

与

之间的 KL 散度损失，其表达式为：

与特征对齐不同，特征对齐通常是最小化预测视觉嵌入与 CLIP 视觉嵌入之间的负余弦相似度，其形式为

而概念对齐则是最小化两个分布之间的

。它衡量了

与

之间的相似度，使

更靠近正例的

（也就是相关概念），同时让它远离负例的

（即无关概念）。这一做法促使

趋向正交，从而最大程度地转移 CLIP 的开放世界知识。

可提示分割：SAM内的掩码解码器会针对通用分割任务对输入提示做出响应。因此，我们将可提示分割视为开启语义能力的必要前奏。与SAM类似，我们的模型默认会为每个提示预测四个掩码，不过会通过路由策略选择其中一个来解决歧义问题。为了提高在大规模SA - 1B数据集上的训练效率，我们实施了一个两阶段采样策略，最多使用9个提示点，而在原始SAM中这一过程是在11个交互阶段内完成的。在第一阶段，我们以等概率从真实掩码中采样一个框或点。在后续阶段，我们从预测掩码与真实掩码之间的误差区域均匀采样1到8个点。为了启用掩码作为先验提示（这是SAM尚未探索的一个方面），我们在第二阶段引入了一种非交互式采样方法，概率为50%。这种采样方式从真实掩码中均匀获取1到9个点，提供了更广阔的提示空间。关于分割损失

，我们采用焦点损失、骰子损失和IoU预测损失的线性组合，权重比例为20:1:1。IoU预测头使用均方误差损失进行训练，依据预测掩码与真实掩码之间的实际IoU进行监督，这与SAM的做法一致。

预训练损失：我们最终的预训练损失是结合概念预测与可提示分割的联合损失：

，其中平衡权重

、

经过搜索，并根据经验设置为 (1, 1)，以便充分学习丰富的 CLIP 语义。利用这个联合损失，我们在 SemanticSA - 1B 数据集上训练一个可提示的标记器。我们的方法概览如图 3 所示。

3.2 可提示字幕生成

为了评估可提示语义标记的有效性，在 SemanticSA - 1B 上完成预训练后，我们在模型顶部附加了一个额外的轻量级文本解码器，并在 Visual Genome（VG）数据集上对其进行微调。我们的文本生成架构概览如图 4 所示。

区域字幕任务：许多先前的工作使用CLIP视觉特征以及大语言模型（LLMs）来生成区域字幕。近期的方法利用 SAM 解码器的特征，同样依赖大语言模型（LLMs）来增强弱语义上下文。然而，简单地附加LLMs不仅会带来计算负担，而且对于区域级别的视觉理解而言可能也并非必要，因为区域字幕通常由不到15个单词组成。在我们构建紧凑型视觉模型的过程中，我们开发了一个生成式标记器。这是通过扩展我们的分词器（见3.1节）的能力来实现的，具体方式是借助因果语言建模融入文本生成功能。具体而言，我们训练一个轻量级文本解码器，用来自我们标记器的语义标记进行提示，以生成区域字幕。通过利用这种具有语义感知的视觉标记器，我们的模型可以高效地对该任务进行端到端训练，无需使用LLMs。

因果文本解码器：我们使用一个嵌入维度为512的标准Transformer来生成简短的区域描述。如果有语义上下文作为提示，这个轻量级文本解码器足以执行从掩码到文本的转换。给定由可提示分词器生成的语义标记（参考图3），我们仅对这些语义标记应用线性投影，使其维度与文本嵌入对齐（见图4）。随后，我们将语义标记置于序列的首位，接着是一个[开始]（[BOS]）标记和单词标记。旋转嵌入用于整合多模态序列的位置编码。我们采用字节对编码，其词表包含32k个标记。最后，我们通过因果语言建模执行下一个标记预测，并使用交叉熵损失。

3.3 推理

在SemanticSA - 1B数据集上完成视觉感知训练，并在Visual Genome数据集上完成文本生成训练之后，我们的模型能够同时进行分类、分割和字幕生成任务。以下概述了推理流程：

掩码选择：给定一个视觉提示，我们的图像解码器会生成4个掩码和9个标记。最终的掩码及相关语义标记通过启发式策略进行选择。具体来说，如果是用边界框进行提示，我们选择第一个掩码；如果是用松散的点进行提示，则选择排名靠前的剩余掩码，这类似于混合专家（MoE）的简化实现。

概念预测：随后，所选的语义标记用于在特定数据集的概念词汇表（例如COCO和LVIS）上预测概念。具体而言，我们利用语义标记，通过一个3层的多层感知机（256→1024→1024）获取一个1024维的视觉embedding。该视觉嵌入会进一步投影到概念分布对数几率（即Ppred）以进行分类。

字幕生成：最后，我们以所选语义标记为提示，采用贪婪采样策略生成最多40个单词标记。为了加快注意力计算，我们遵循自回归解码的标准做法，缓存序列中前一步生成的键值对。

4. 实验

4.1 实验设置

预训练：我们在SemanticSA - 1B上对TAP模型进行预训练，该数据集包含SA - 1B数据及其相关的CLIP先验知识。完整的SA - 1B由1100万张高分辨率图像组成，每张图像大约有100个区域，总共产生11亿个分割掩码。为了获取SA - 1B数据的CLIP先验知识，受[39, 68, 74]启发，我们使用EVA - CLIP [48]在一个经过整理的标签空间上生成文本嵌入，这个标签空间是从COCO [31]、ADE20K [75]、LVIS [12]、Objects365 [44]、Visual Genome [21]和OpenImagesV4 [23]数据集合并而来的。这就产生了一个涵盖2560个类别的概念列表，覆盖了用于分割的物体和物质。

评估：我们在COCO和LVIS数据集上评估零样本实例分割性能。对于零样本实例分类，由于LVIS涵盖1203个类别，相比之下，COCO仅覆盖80个常见类别，与开放世界假设有所偏离，所以我们优先选用LVIS。在区域级字幕任务中，考虑到SA - 1B与Visual Genome（VG）之间的域差距，我们采用两阶段微调方法。首先，我们冻结图像编码器 - 解码器，仅使用VG v1.0训练集微调文本解码器，将其标记为“部分微调（partial - FT）”。随后，我们解冻图像编码器 - 解码器，对模型进行端到端的微调，把这种两阶段微调策略标记为“完全微调（full - FT）”。由于不存在具备对齐的掩码、类别和字幕的网络规模数据集，所有的消融研究（4.3节）均采用“部分微调”方式进行。我们在VG测试集和RefCOCOg [37]验证集上报告以下四个指标：BLEU@4、METEOR、ROUGE和CIDEr。

实现细节：在所有实验中，我们使用AdamW优化器（β1 = 0.9，β2 = 0.999），基础学习率设为1e-3 。采用余弦退火学习率策略。在SemanticSA - 1B上进行预训练期间，应用尺度抖动，范围是[0.5, 2.0]，共进行180k次迭代（约4个轮次），使用256个GPU，批量大小为256。我们在不进行数据增强的情况下对VG数据集进行微调，共60k次迭代（约50个轮次），使用64个GPU，批量大小为64 。其他超参数包括：权重衰减设为0.1，ViT - B/ViT - L的随机失活路径率分别设为0.1/0.2，图像/文本解码器的随机失活率分别设为0.1/0.4。图像编码器使用MAE预训练权重进行初始化，而其他所有层则从零开始训练。对于所有实验，在每个采样阶段，每个GPU最多采用64个采样提示。

4.2 主要结果

零样本实例分类：我们用真实标注（GT）框提示我们的模型，以此评估其在LVIS数据集上的纯识别能力。使用GT框作为视觉提示时，我们的模型显著超越了RegionCLIP和RegionSpot，后两者是在有限的图像区域上训练的。这些乐观的结果表明，在详尽的图像区域上进行概念预测，能够有效地赋予SAM语义感知能力。如表1所示，性能卓越的EVA-CLIP在零样本评估中胜过了所有其他方法，取得了令人瞩目的稀有类别平均精度（rare AP）。尽管如此，对于实时视觉系统而言，部署独立的CLIP（5B）模型来计算大量图像块是不切实际的。我们证明了大型CLIP模型的知识能够整合到一个紧凑的分词器（0.1B）中，且性能尚可接受。

区域级字幕：我们在Visual Genome和RefCOCOg数据集上评估我们的模型。首先，我们使用真实标注（GT）框提示图像解码器，随后，再用生成的语义标记提示文本解码器。评估结果如表2所示。令人惊讶的是，在Visual Genome数据集上，即便图像编码器 - 解码器是在SA - 1B上预训练的、之前从未见过VG图像且处于冻结状态（“部分 - 微调”），我们的模型也能取得154.7的CIDEr分数。通过采用两阶段微调策略（“完全 - 微调”），仅使用一个轻量级文本解码器，我们就创造了164.7的CIDEr分数新纪录。值得注意的是，同期工作ASM是在一个多模态数据集上训练的，该数据集包含大量区域 - 文本对。我们模型的语义知识是从CLIP模型中学到的。另一项同期工作SCA，额外训练了一个12层的图像解码器，以学习用于文本提示的字幕标记。这些结果表明，我们的语义标记在预训练期间有效地编码了足够的区域级信息用于字幕生成，这支持了我们之前的说法，即TAP可以充当一个位置感知的图像分词器。

零样本实例分割：我们在零样本实例分割任务中评估我们的模型，这是原始SAM擅长的任务。按照常见做法[20, 67]，我们首先从ViTDet-H模型获取检测边界框。随后，我们利用这些框来提示图像解码器，并比较在COCO和LVIS数据集上的纯分割性能（即，使用框的类别）。为了进行公平比较，我们报告原始SAM以及我们复现版本（记为我们的实现）的结果。如表3所示，我们的模型在不同模型规模下都取得了与原始SAM相当的分割结果。这表明额外的概念预测和区域字幕任务并不会损害SAM原本的能力。此外，这还说明通用分割作为一项基本的几何任务，可能无法充分利用视觉基础模型中的语义表示。

4.3 消融实验

预训练损失：关于预训练损失的消融研究展示在表 4 和表 5 中，其中、和分别代表使用分割、特征以及概念进行预训练。特征预测和概念预测（见 3.1 节）。从表 4 可以看出，仅用进行预训练时（模型 A），字幕指标非常低。当与语义预测相结合（模型 B/C）时，字幕性能有了显著提升。尽管特征预测表现出了语义感知能力，但在分类和字幕任务中，它都不如概念预测。这些发现表明，概念空间对于获取 CLIP 先验知识至关重要。我们推测，这个空间有效地帮助模型从 CLIP 中学习到了负文本嵌入（即）。此外，表 4 中的分割结果表明，增加语义预测的预训练既不会提高也不会降低 COCO 和 LVIS 上的掩码平均精度（mask AP）。这一观察结果还表明，SAM 架构可以在分割掩码之外融入更多的任务监督。

语义标记：为了评估语义标记的有效性，我们进行了四项实验。首先，我们使用 “预训练” 列中列出的损失对模型进行预训练。随后，我们使用从冻结的预训练模型生成的、在 “文本提示” 中概述的项目微调文本解码器。模型 A 作为我们的基线，仅使用进行预训练。在这里，掩码标记直接用于区域级字幕任务，类似于使用原始 SAM 的输出训练文本解码器。模型 D 是我们的默认模型，通过可提示的分割和概念预测联合优化，使用语义标记来提示文本解码器。如表 4 所示，在字幕任务中，语义标记始终优于掩码标记，同时在分割任务中实现了相当的平均精度（AP）。最终，事实证明语义标记是最有效的。这表明语义标记化极大地释放了基础模型的潜力，有助于开展更多感知任务。

缩放文本解码器：我们沿着深度和嵌入维度扩大文本解码器，来消除字幕生成的瓶颈问题。如表 6 所示，在 VG 数据集上，增加模型规模并没有带来显著的提升。这表明，除非能够进一步增加文本的长度与数量，否则在区域字幕任务中使用更大的解码器可能并非必要。

4.4 定性结果

我们使用基于点的提示对 TAP 进行定性评估。只需简单点击，或是用密集的点网格自动进行提示，我们的模型就能同时生成分割掩码、类别名称以及文本描述。

开放世界知识：图 5 展示了在开放世界场景中具有挑战性的示例实例。由于词汇设计具有主观性，像 “百事可乐”“可口可乐”“龙”“蜘蛛侠” 以及 “威士忌” 这类精心设定的概念，很难通过检索（即分类）被选中。然而，我们的模型在处理这些与特定概念相关的实例时表现熟练，这表明它具备处理开放世界知识的能力。