SED: A Simple Encoder-Decoder for Open-Vocabulary Semantic Segmentation

Abstract

开放词汇语义分割旨在将像素划分为来自开放类别集合的不同语义组。现有的大多数方法依赖于预训练的视觉–语言模型，其中关键在于如何将图像级模型适配到像素级分割任务中。在本文中，我们提出了一种简单的编码器–解码器框架，称为SED，用于开放词汇语义分割。SED 由两部分组成：基于层级编码器的代价图生成模块以及带有类别早期剔除的渐进式融合解码器。其中，层级编码器代价图生成模块采用层级式骨干网络（而非普通的 Transformer）来预测像素级的图像–文本代价图。与普通 Transformer 相比，层级骨干网络能够更好地捕获局部空间信息，并且其计算复杂度相对于输入尺寸是线性的。渐进式融合解码器采用自顶向下的结构，将代价图与骨干网络不同层级的特征图逐步融合以完成分割。为加速推理速度，我们在解码器中引入了一种类别早期剔除机制，能够在解码的早期阶段剔除大量不存在的类别，在不降低精度的情况下，推理速度最多可提升4.7 倍。我们在多个开放词汇语义分割数据集上进行了实验，结果验证了 SED 方法的有效性。在使用ConvNeXt-B作为骨干网络时，SED 在ADE20K（150 类别）数据集上取得了31.6% 的 mIoU，并且在单张A6000 GPU上的推理速度为82 毫秒/图像。我们的代码已开源，地址为：
https://github.com/xb534/SED

1. Introduction

语义分割旨在解析图像中每一个像素所属的语义类别。传统方法 [6, 31, 52] 通常假设语义类别是封闭集合，因此在推理阶段难以识别未见过的类别。为了解决这一问题，近年来的研究开始探索开放词汇语义分割[3, 51, 59]，以实现对任意类别像素的分割。

近年来，诸如CLIP[38] 和ALIGN[25] 等视觉–语言模型通过数百万规模的图像–文本配对数据学习到了对齐的图像–文本特征表示。预训练的视觉–语言模型在识别开放词汇类别方面展现出卓越的泛化能力。这也激发了大量研究工作，探索将视觉–语言模型用于开放词汇语义分割任务 [14, 29]。

早期的研究工作主要采用两阶段框架[15, 29, 55]，将视觉–语言模型直接适配到开放词汇分割任务中。具体而言，这类方法首先生成与类别无关的掩码候选，随后利用预训练的视觉–语言模型将这些候选掩码分类为不同的语义类别。然而，这种两阶段框架在掩码生成和类别分类阶段使用了两个相互独立的网络，从而限制了计算效率；此外，该框架也未能充分利用上下文信息。

不同于上述两阶段方法，基于单阶段框架的方法直接扩展单一的视觉–语言模型以实现开放词汇语义分割。一些方法通过移除图像编码器最后一层的池化操作，生成用于分割的像素级特征图。例如，MaskCLIP[62] 去除了 CLIP 图像编码器最后一层的全局池化，并利用value embedding与文本 embedding直接预测像素级分割结果。CAT-Seg[12] 则首先生成像素级的图像–文本代价图，随后通过空间聚合和类别聚合对该代价图进行细化。

尽管这些方法取得了较好的性能，但我们注意到它们仍存在以下局限性。首先，MaskCLIP 和 CAT-Seg 均采用普通 Transformer ViT[16] 作为骨干网络，该结构在局部空间信息建模能力较弱，且输入分辨率较低。为了解决这些问题，CAT-Seg 引入了一个额外的网络来补充空间信息，但这也带来了额外的计算开销。其次，随着类别数量的增加，CAT-Seg 的计算成本会显著上升。

为了解决上述问题，我们提出了一种简单而高效的编码器–解码器方法，称为SED。SED 由两部分组成：基于层级编码器的代价图生成模块以及带有类别早期剔除的渐进式融合解码器。其中，基于层级编码器的代价图生成模块采用层级式骨干网络（而非普通 Transformer）来预测像素级图像–文本代价图。与普通 Transformer 相比，层级骨干网络能够在不同层级上更好地保留空间信息，并且其计算复杂度相对于输入尺寸呈线性增长。渐进式融合解码器则逐步融合来自层级骨干网络不同层级的特征图以及代价图，以生成分割预测结果。为了提升推理速度，我们在解码器中设计了一种类别早期剔除机制，能够在解码的早期阶段有效预测真实存在的类别并剔除不存在的类别。我们在多个开放词汇语义分割数据集上进行了全面实验，结果表明所提出方法在精度和效率方面均具有明显优势。我们的主要贡献总结如下：

• 我们提出了一种用于开放词汇语义分割的编码器–解码器框架，该框架包含基于层级编码器的代价图生成模块和渐进式融合解码器。

• 我们引入了一种类别早期剔除机制，能够在解码早期剔除不存在的类别，在几乎不损失分割性能的前提下显著提升推理速度。例如，在PC-459数据集上可实现4.7 倍的加速。

• 我们提出的SED方法在多个开放词汇语义分割数据集上取得了领先性能。具体而言，SED 在分割性能与推理速度之间实现了良好的权衡（见图 1）。在使用ConvNeXt-L作为骨干网络时，SED 在A-150数据集上取得35.2% 的 mIoU，在PC-459数据集上取得22.6% 的 mIoU。

3. Method

在本节中，我们详细介绍所提出的用于开放词汇语义分割的编码器–解码器框架 SED。如图 2 所示，SED 的整体架构由两个主要组成部分构成：基于层级编码器的代价图生成模块以及带有类别早期剔除的渐进式融合解码器。在基于层级编码器的代价图生成模块中，我们采用层级式图像编码器和文本编码器来生成像素级图像–文本代价图(F_{cv})，以及供解码器使用的多层级特征图(F_2, F_3, F_4)。渐进式融合解码器通过特征聚合模块（FAM）和跨层融合模块（SFM），逐步融合像素级代价图 (F_{cv}) 与层级特征图 (F_2, F_3, F_4)，从而生成高分辨率特征图(F_h)。在获得 (F_h) 之后，我们通过输出层预测不同类别的语义分割结果。此外，在解码器中引入了类别早期剔除（CER）策略，用于在推理的早期阶段剔除不存在的类别，以进一步提升推理速度。

3.1. Hierarchical Encoder-based Cost Map

基于层级编码器的代价图生成（HECG）采用视觉–语言模型CLIP[11, 38, 41] 来生成像素级图像–文本代价图。具体而言，我们首先分别使用层级式图像编码器和文本编码器来提取视觉特征和文本嵌入；随后，计算这两种特征之间的像素级代价图。现有方法（如MaskCLIP[62] 和CAT-Seg[12]）通常采用普通 Transformer作为图像编码器来生成像素级代价图。然而，如前文所述，普通 Transformer 在局部空间信息建模能力方面相对较弱，并且其计算复杂度相对于输入尺寸为二次增长。为了解决这些问题，我们提出使用层级式骨干网络作为代价图生成的图像编码器。层级编码器能够更好地捕获局部信息，并且其计算复杂度相对于输入尺寸呈线性增长。代价图生成的具体过程如下所述。

3.2. Gradual Fusion Decoder

语义分割任务在很大程度上依赖于高分辨率特征图。然而，由编码器生成的代价图 (F_{cv}) 分辨率较低且噪声较大，直接利用该代价图进行预测并不利于生成高质量的分割结果。为了解决这一问题，我们提出了一种渐进式融合解码器（Gradual Fusion Decoder，GFD）。GFD 通过级联两个模块，在多层结构中逐步生成高分辨率特征图 (F_h)。这两个模块分别为特征聚合模块（Feature Aggregation Module，FAM）和跳层融合模块（Skip-layer Fusion Module，SFM）。其中，FAM 旨在建模局部区域与不同类别之间的关系；而 SFM 则利用分层编码器中的浅层特征来增强特征图的局部细节，从而提升分割结果的精细度。

Feature Aggregation Module:

Skip-layer Fusion Module:

与普通 Transformer 相比，采用带跳层融合的分层编码器能够显著提升模型性能。这主要归因于两个方面：一方面，分层编码器能够为语义分割提供更加丰富的局部信息；另一方面，阻断梯度向图像编码器的反向传播，避免了对开放词汇语义分割能力产生的负面影响。

3.3. Category Early Rejection

渐进式融合解码器的计算开销与类别数量成正比。当类别数量非常大时，推理时间会显著增加。事实上，大多数图像只包含少数几个类别，因此解码器在推理阶段的大部分时间被用于计算并不存在的类别对应的特征。为提升推理速度，我们提出了一种类别早期剔除机制（category early rejection），用于在解码器的早期层识别图像中真实存在的类别，并提前剔除不存在的类别。被剔除类别对应的特征图将从当前解码层中移除，后续解码层仅对保留下来的类别进行计算。

在训练阶段，如图 4(a) 所示，我们在每一层后加入一个辅助卷积分支，分别预测分割结果，并使用真实标注进行监督。为避免对主干解码器训练产生负面影响，我们阻断了这些辅助分支向解码器的梯度反向传播。

在推理阶段，我们在分割预测图上采用top-k 策略来判断图像中存在的语义类别。具体而言，对于每个像素，选取响应值最大的 top-k 个类别，并对所有像素的结果取并集，得到当前图像中预测存在的类别集合，该集合将作为输入传递给下一解码层。实验表明，当 (k = 8) 时，基本可以保证识别出绝大多数真实存在的类别。图 4(b) 展示了推理阶段的类别早期剔除过程。

具体来说，我们首先基于 (F^{l_1}{dec}) 预测分割结果，并通过 top-k 策略选取 (N{l_1}) 个类别；随后，移除未被选中的类别对应的特征图，生成输出特征图
[
F^{clear}{l_1} \in \mathbb{R}^{H_v \times W_v \times N{l_1} \times D}.
]
该特征图被送入下一解码层。对后续各层同样执行该过程，逐层生成包含更少类别的特征图。

因此，大量不存在的类别在解码器的早期层即被剔除，从而显著提升了解码器的整体推理速度。