CVPR2022 ProtoSeg | 基于不可学习原型的语义分割模型

189****6672

论文标题：Rethinking Semantic Segmentation: A Prototype View

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

摘要

引入主题：尽管流行的语义分割网络设计（基于FCN或基于注意力）和掩码解码策略（基于参数的softmax或基于像素查询）不同，但可以通过将softmax权重或查询向量视为可学习的类原型，故可以都归为一类。
论文算法：鉴于这种原型观点，论文消除了这种参数分割机制的一些局限性，并提出了一种基于不可学习原型的非参数替代方案。模型将每个类表示为一组不可学习的原型，仅依赖于该类中几个训练像素的平均特征，而不是以前的方法以完全参数化的方式学习每个类的单个权重/查询向量。因此，密集预测是通过非参数最近原型检索实现的。这允许模型通过优化嵌入像素和锚定原型之间的排列，直接塑造像素嵌入空间。它能够用恒定数量的可学习参数处理任意数量的类。
实验结果：使用基于FCN和基于注意力的分割模型（即HR-Net、SWN、SegFormer）和主干（即ResNet、HRNet、Swin、MiT），论文的非参数框架在多个数据集（即ADE20K、Cityscapes、COCO等）上产生了令人信服的结果，并在大数据集情况下表现良好。

Introduction

随着连接主义的复兴，语义分割取得了迅速的进展。到目前为止，大多数最先进的分割模型都是建立在完全卷积网络（FCN）上的。尽管其多样化的模型设计和令人印象深刻的结果，现有的基于FCN的方法通常将参数softmax（）应用于像素级特征，以实现密集预测（图1（a））。最近，Transformer的巨大成功刺激了基于注意力分割研究的出现。许多这些“非FCN”模型直接遵循标准掩码解码机制，即估计密集视觉嵌入上的softmax分布（从patch tokens序列中提取）。

有趣的是，其他基于Transformer论文采用像素查询策略（图1（b））：利用一组可学习向量来查询密集嵌入以进行掩码预测。这些论文推测，学习的查询向量可以捕获类属性，但这个观点缺乏深入的分析。

根据上文所述，注意到存在两种不同的掩码解码策略，自然会出现以下问题：

它们之间有什么联系和区别？
如果可学习的查询向量确实隐式地捕获了数据的某些固有属性，那么有没有更好的方法来实现这一点？

基于原型分类的思想既经典又直观（可以追溯到最近邻算法），而且在认知科学中找到依据：数据样本根据其与类的代表性原型的接近程度进行分类。

大多数现有工作都试图直接学习C个类别原型。因此问题2变得更加重要：

这种可学习的基于原型的参数化范式有哪些局限性？
如何解决这些限制？

由问题2驱动，论文发现有三个关键限制：

第一，通常每个类只学习一个原型，不足以描述丰富的类内差异。原型仅以完全参数化的方式学习，而不考虑其代表性能力。

其次，要将H×W×D图像特征张量映射到H×W×C语义掩模，至少需要D×C参数进行原型学习。这损害了其通用性，尤其是在大模型情况下；例如，如果有800个类，D=512，需要0.4M的可学习原型参数。

第三，利用交叉熵损失，仅优化类间距离之间的相对关系；忽略像素和原型之间的实际距离，即类内紧凑性。

Existing Semantic Segmentation Models as Parametric Prototype Learning

首先将现有两种掩模解码策略形式化，然后回答问题1。

Parametric Softmax Projection。几乎所有FCN-like 和许多基于注意力的分割模型都采用这种策略。包括两个可学习的部分:

i)用于密集视觉特征提取的编码器 φ

ii)将像素特征投影到语义标签空间的分类器 ρ (即投影头)。

对于每个像素样本 i，从 φ 中提取的嵌入特征 i∈RD 输入 ρ 进行 C 类别分类：

其中p（c|i）∈[0,1]是i被分配给c类的概率。ρ是线性层，由W=[w1,…,wC]∈RC×D参数化；wc∈RD是第c类的可学习投影向量。

**Parametric Pixel-Query。**一些基于注意力的分割网络以更“‘Transformer- like”的方式工作：给定像素嵌入i∈RD，一组C个查询向量，即，，e=[e1，··，eC]∈RC×D以生成C类上的概率分布：

其中“∗” 是l2归一化输入之间的内积。

**Prototype-based Classification。**基于原型的分类已经研究了很长时间，可以追溯到机器学习中的最近邻算法和认知科学中的原型理论。它的流行源于其直观的想法：用原型表示类，并引用原型进行分类。设{pm}m=1M是代表类，，{cpm∈{1，·，C}}m的一系列原型。对于数据样本i，通过将i与{pm}m进行比较，并将获胜原型的类别作为响应来进行预测：

其中i和{pm}m是数据样本和原型在特征空间中的嵌入，而〈·〉表示距离度量，通常设置为l2距离（即||i−pm||），但也可以应用其他距离度量。

进一步，等式1-2可以以统一的形式表示：

用等式3-4，准备回答问题1， 2。这两种方法都基于可学习原型；它们是参数模型，因为它们为每个类c学习一个原型gc，即线性权重wc或查询向量ec。因此，可以考虑基于softmax投影的方法“秘密”学习查询向量。至于差异，除了不同的距离度量（即内积与余弦相似度），基于像素查询的方法可以将查询馈送到交叉注意力解码器层，以进行跨类上下文交换，而不是仅利用softmax层内学习的类权重。

在参数化原型学习的统一观点下，该领域中一些固有但长期被忽视的问题没有得到解决。

首先，原型选择是一个重要方面——原型应该适合对应的类别。然而，现有的语义分割算法通常只通过一个原型来描述每个类，没有类内差异。此外，原型是在完全参数化方式来直接学习的，而不考虑其代表性能力。

第二，可学习原型参数的数量，随着类的数量而增长。这可能会阻碍可扩展性，尤其是当存在大量类时。例如，如果有800个类，并且像素特征维数为512，那么仅原型学习就需要至少0.4M的参数，这使得分割成为一项艰巨的任务。此外，如果想用十个原型来表示每个类，而不是仅仅一个，就需要学习4M原型参数。

第三，等式3直观地表明，原型是根据数据进行度量比较的。然而，现有算法通常通过直接优化像素预测精度（如交叉熵损失）来监督密集分割表示，忽略已知的归纳偏差，如关于特征分布的类内紧凑性。这将阻碍学习到的分割特征的识别潜力。

在解决了这个问题3之后, 在下一节中，将详细介绍基于非学习原型的非参数分割方法，该方法是对问题4的可靠回答。

Non-Learnable Prototype based Nonpara-metric Semantic Segmentation

构建了一个非参数分割框架，该框架通过一组不可学习的类原型进行密集预测，如图2所示。

Non-Learnable Prototype based Pixel Classification。

通常，首先采用编码器网络（基于FCN或基于注意力），即φ来映射输入图像i∈Rh×w×3转换为三维特征张量I∈RH×W×D。对于像素级C个类别分类，不是自动学习C类权重{wc∈RD}c=1C（参见等式1）或C个查询向量ec∈RDc=1C（参见等式2），而是使用一组CK个非可学习原型，即，，，{pc，k∈RD}c，K=1C，k，它们仅基于类数据子中心。更具体地说，每个类别c∈{1，·，C}由总共K个原型表示{pc,k}k=1K，原型pc,K被确定为嵌入空间φ中属于C类的训练像素样本的第K个子簇的中心。

通过这种方式，原型可以捕获相应类的特征属性，而不引入额外的可学习参数。类似于等式3，每个像素i的类别预测是通过 winner-take-a的分类实现的。

其中i∈RD代表像素l2归一化嵌入、距离测度〈·〉定义为负余弦相似性。

使用这种基于示例的推理模式，首先对C类上像素i的概率分布进行细化：

其中像素类距离si∈[−1,1]被计算为到c类最近原型的距离。给定每个像素i的gt类，因此，交叉熵损失可用于训练：

等式7可以被视为使像素i更接近其对应类的最近原型，并进一步远离不相关类的其他原型。然而，由于两个原因，仅仅采用这种训练目标是不够的。

首先，等式7仅考虑像素-类距离，而不考虑类内像素-原型关系。例如，对于区分表示学习，像素i被期望进一步推近ci类的某个原型，并远离ci类的其他原型。等式7不能捕捉这种性质。

其次，由于像素-类距离在所有类之间都进行了归一化（参见等式6），等式7仅优化类内和类间距离之间的相对关系，而不是直接正则化像素和类之间的余弦距离。例如，当像素i的类内距离相对小于其他类间距离时，等式7的惩罚较小，但类内距离可能仍然较大。

接下来，首先阐述了类内在线聚类策略，然后详细介绍了两个额外训练目标，这两个目标依赖于原型任务（即聚类结果），并分别解决了上述两个问题。

Within-Class Online Clustering

采用在线聚类方法进行原型选择和分配：将同一类中的像素样本分配给属于该类的原型，然后根据分配更新原型。聚类带来了一个自然瓶颈，迫使模型发现类内区分模式，但丢弃了特定于实例的细节。因此，被选为子集群中心的原型是相应类的代表。在线进行聚类使方法可扩展到大量数据，而不是离线聚类，需要在整个数据集上进行多次传递以进行特征计算。

形式上，给定属于c类的一个训练batch中的像素Ic={in}n=1N，目标是将像素Ic映射到c类的K个原型，{pc，k}k=1K。将像素到原型的映射表示为Lc=[lin]n=1N∈0,1K×N，其中，lin=[lin，K]k=1K∈{0,1}K是K个原型上的一个one-hot向量。Lc的优化通过最大化像素嵌入之间的相似性来实现，即Xc=[in]n=1N∈RD×N和原型，Pc=[pc，k]k=1K∈RD×K:

为了求解等式8，可以将Lc作为transportation polytope的一个元素:

其中，，，h(Lc)=∑n，k−lin，kloglin，k 是一个熵，κ > 0是一个控制分布光滑度的参数。利用软赋值松弛和额外的正则化项 h (Lc) ，得到了方程的求解器：

其中u∈RK和v∈RN是重正则化向量，通过Sinkhorn-Knopp迭代的几个步骤计算。

在线聚类在GPU上是高效的，因为它只涉及两个矩阵乘法；实际上，将10万像素聚类成10个原型只需2.5毫秒。

Pixel-Prototype Contrastive Learning

分配概率矩阵Lc=[lin]n=1N∈[0,1]K×N，在线将训练像素I=inn=1N分为c类中的K个原型{pc,k}k=1K。当前批次中的所有样本处理后，每个像素i被分配给ci类的第K个原型，其中ki˙=arg⁡maxk{li,k}k=1K。为原型分配预测导出一个训练目标是很自然的，即最大化原型分配后验概率。这可以看作是一种像素原型对比学习策略，并解决了等式7的第一个限制：

等式11强制每个像素嵌入i与其指定的（“正”）原型、pci、ki相似，并与其他CK−1（‘负’）原型P−不相关. 与先前基于像素度量学习的分割模型相比，该模型消耗大量负像素样本，该论文方法仅需要CK原型进行像素-原型对比度计算，既不会产生较大的记忆成本，也不需要大量的像素对比较。

Pixel-Prototype Distance Optimization

基于像素类/原型距离的相对比较，等式7和等式11激发了类间/簇间的差异性，但较少考虑减少簇内变化，即使同一原型的像素特征紧凑。因此，通过直接最小化每个嵌入像素与其指定原型之间的距离，使用紧致性感知损失进一步正则化表示

请注意，i和，pci，ki均进行l2归一化”。该训练目标最大限度地减少簇内变化，同时保持具有不同原型符号的特征之间的分离，使模型更具鲁棒性。

Network Learning and Prototype Update

该模型是一种非参数方法，通过直接优化像素嵌入空间φ来学习语义分割。它被称为非参数，因为它直接从训练像素样本本身构建原型假设。因此，通过在所有训练像素上最小化组合损失，通过随机梯度下降来学习特征提取器φ的参数：

同时，不可学习原型，，，{pc，k}c，k=1C，k不是通过随机梯度下降学习的，而是作为相应嵌入像素样本的中心计算的。为此，通过考虑在线聚类结果，让原型不断演化。特别是，在每次训练迭代之后，每个原型都会更新一次:

μ∈ [0,1]是动量系数，，ic，k表示嵌入训练像素的l2归一化平均向量，通过在线聚类将其分配给原型，pc，k。

有了原型的明确含义，分割过程可以直观地理解为检索最相似的原型（子类中心）。图3提供了人和车的原型检索结果，每个人和车具有K＝3个原型