点击下方卡片，关注「集智书童」公众号
语义分割在Transformer架构的适应下取得了显著的进步。与Transformer的进步并行的是基于CNN的U-Net在高质量医学影像和遥感领域的重大进展。这两种双赢激发了作者将它们的优势结合起来，从而诞生了一种专为高效语义编码的U-Net基础视觉Transformer解码器。在这里，作者提出了一种新的Transformer解码器U-MixFormer，它基于U-Net结构，用于高效的语义分割。作者的方法通过利用编码器和解码器阶段之间的横向连接作为注意力模块的特征 Query ，除了依赖跳接的传统方式。此外，作者还创新地将来自不同编码器和解码器阶段的层次特征图混合，形成一个统一的表示形式，用于Key和Value，从而产生作者独特的_mix-attention_模块。
作者的方法在各种配置下都展示了最先进的性能。大量实验表明，U-MixFormer在SegFormer、FeedFormer和SegNet之上取得了巨大的优势。例如，U-MixFormer-B0在mIoU方面分别比SegFormer-B0和FeedFormer-B0提高了3.8%和2.0%，且计算量减少了27.3%，同时比SegNext提高了3.3%的mIoU，在MSCAN-T编码器在ADE20K上的表现。
代码：https://github.com/julian-klitzing/u-mixformer
语义分割是计算机视觉的一个基本下游任务，在工业和学术界一直受到越来越多的关注。语义分割的重要性在于其在实际场景中的广泛应用，例如自动驾驶[23]和医学诊断[1]。尽管取得了这些进展，实现精确的像素级预测仍然是一个挑战，因为需要平衡全局和局部上下文。
全卷积网络（FCN)[11]的引入使得编码器-解码器结构变得流行，其中编码器提取高层语义，解码器将其与空间细节相结合。虽然变体[1, 1, 2]已经改进了这种方法，但传统的卷积神经网络（CNN）在捕捉长程上下文方面存在困难。这一限制激发了人们对基于视觉Transformer的分割方法的研究兴趣。
Transformer[26]，最初是为了自然语言处理而设计的，但由于自注意力机制可以捕获输入序列中的全局关系，因此已经被适应用于视觉任务并取得了显著的成功。受到其成功的启发，[2]也将其适应用于视觉任务，从而导致了视觉Transformer（ViT），该方法将图像解释为嵌入式 Patch 的序列，并使用Transformer编码器来处理它们。这种方法在ImageNet上取得了显著的结果。自从ViT引入以来，已经进行了许多研究，以将其适应于语义分割。主要目标是两个方面：优化编码器并创建能够有效地利用编码器阶段特征的解码器。
在视觉任务中利用Transformer的解码器结构已经引起了显著的转变。DETR首先开创了这种方法，并将Transformer编码器-解码器框架集成到检测和分割中。DETR之后，Segmenter、MaskFormer、Mask2Former引入了具有全局类别标签的掩码预测解码器，强调高级特征。最近，FeedFormer提出了一种解码器设计，它只使用最低级别的编码器特征来解码高级编码器特征。尽管在基于Transformer的解码器方面已经取得了分割方面的进展，但这些方法通常依赖于在其注意力机制内的计算密集型特征配置。此外，这些方法在解码器阶段之间传播特征图方面存在效率问题。
传统上，U-Net架构以其对称的基于卷积神经网络（CNN）的编码器-解码器结构，成为语义分割的 favored 选择，尤其是在医学领域。这种偏好源于U-Net有效地捕获和传播层次特征的特点。此外，其横向连接也起着重要作用，有助于在编码器和解码器之间流动多阶段特征。作者假设，利用U-Net架构固有的这些优势可以导致特征的有效优化，然后可以分层集成到Transformer解码器阶段。
在本文中，作者提出了一种名为_U-MixFormer_的UNet类似Transformer解码器。在U-Net的基础原则之上，U-MixFormer适当地将其专用_mix-attention_模块中的多阶段特征作为Key和Value进行自适应集成。这个模块确保了特征的逐渐传播，并在解码器阶段之间依次混合它们，有效地管理这些阶段之间的依赖关系以捕捉上下文并优化边界。这可以强调传统的CNN层次特征表示，并增强Transformer的全局上下文理解能力。据作者所知，这是首次将U-Net的固有优势与视觉Transformer的转换能力相结合，特别是通过一种新的注意力模块有效地调和 Query 、Key和Value用于语义分割。
作者所做的贡献如下：
基于U-Net的新型解码器架构 作者提出了一种新颖而强大的Transformer解码器架构，该架构旨在实现高效的语义分割。借鉴U-Net在捕获和传播层次特征方面的优势，作者的设计独特地使用Transformer编码器的横向连接作为 Query 特征。这种方法确保了高级语义和低级结构的和谐融合。
优化特征合成以增强上下文理解 为了提高作者UNet类似Transformer架构的效率，作者将多个编码器和解码器输出混合并更新为用于Key和Value的集成特征，从而实现了作者提出的_mix-attention_机制。这种方法不仅为每个解码器阶段提供了丰富的特征表示，而且增强了上下文理解。
兼容性与多种编码器 作者展示了将U-MixFormer与现有的流行的基于Transformer的编码器（如MiT和LVT）和基于CNN的编码器（如MSCAN）相结合的兼容性。
实证基准测试 如图1所示，U-MixFormer在语义分割方法方面，无论是计算成本还是准确性，都实现了新的最先进水平。它始终优于轻量级、中量级和甚至重量级编码器。这种优越性在ADE20K和Cityscapes数据集上得到了证明，特别是在具有挑战性的Cityscape-C数据集上的显著表现。
SETR是第一个将ViT作为语义分割的编码器架构。由于ViT仅将输入图像划分为 Patch ，SETR产生了单 Scale 编码器特征。PVT和Swin Transformer在编码器阶段之间重复将特征图分组到新的非重叠 Patch 中，从而分级生成多 Scale 编码器特征。这两种方法也通过减少Key和Value的空间维度（PVT）或使用位移窗口组（Swin Transformer）来增强自注意力模块的效率。SegFormer重用了PVT的效率策略，同时删除了位置编码和嵌入特征图到重叠 Patch 中。与上述方法不同，SegNeXt和LVT的编码器采用了卷积注意力机制。
DETR是第一个将Transformer解码器用于语义分割的方法。随后的作品改编了DETR，但仍然依赖于可学习的 Query ，这在计算上是昂贵的，特别是与多 Scale 编码器特征相结合时。相比之下，FeedFormer直接使用编码器阶段的特征作为特征 Query ，从而提高了效率。FeedFormer解码使用最低级别的编码器特征（作为 Query 特征）和最高级别的编码器特征。然而，这种设置处理特征图逐个，没有在解码器阶段之间逐步传播特征图，因此错过了改进目标边界检测的机会，以实现更逐步的优化。此外，其他最近的MLP或CNN基解码器也缺乏解码器特征的逐步传播。
在医学和遥感领域，人们已经尝试将UNet架构从基于卷积神经网络（CNN）的框架转变为基于Transformer的框架。TransUNet是首次成功将Transformer引入医学图像分割的方法，它使用ViT与他们的CNN编码器相结合。其他混合方法也在[14, 15, 16]中提出。Cao2023引入了Swin-UNet，这是第一个完全基于Transformer的UNet类似架构。该设计具有用于编码器和解码器的重型Swin Transformer阶段，保留了它们之间的横向连接作为跳接连接。与Swin-UNet相比，作者的方法采用轻量级解码器阶段，使其适合更广泛的下游任务。此外，作者将横向连接解释为 Query 特征而不是跳接连接，并融入了一种独特的注意力机制。
本节介绍了U-MixFormer，一种新颖的UNet类似Transformer解码器架构用于语义分割。一般来说，作者的解码器由与编码器阶段数量相同的个阶段组成。为了清晰起见，图2提供了这种架构的一个视觉概述，示例为一个四阶段的层次编码器，如MiT、LVT或MSCAN。
首先，编码器处理一个大小为的输入图像。四个阶段分别产生层次化、多分辨率的特征，其中。其次，作者的解码器阶段依次通过执行_mix-attention_来生成 refined特征，其中 Query 特征等于相应的横向编码器特征图。Key和Value特征由编码器和解码器阶段的混合给出。值得注意的是，作者的解码器与编码器阶段输出的维度相同。第三，解码器特征使用双线性插值上采样以匹配的高度和宽度。最后，将拼接的特征通过MLP处理以预测具有的分割图。
在Transformer模块中使用的注意力模块计算 Query 、键和值的缩放点积注意力如下：
在这里，表示键和 Query 的嵌入维数，、、 是由选择的特征进行线性投影得到的。
作者方法的核心是选择要投影到Key和Value的特征，这导致了作者提出的mix-attention机制。图3展示了传统的自注意、交叉注意和作者所提出的mix-attention之间的比较。
在自注意力中，用于生成 Query 、Key和Value的特征是相同的（），它们都源自同一个来源，即相同的编码器/解码器阶段。交叉注意力使用两个不同的特征，和，每个特征都源自单一的独特来源。相比之下，作者的mix-attention机制利用来自多个多 Scale 阶段的混合特征来为，从而允许 Query 在不同不同的阶段（即上下文粒度的不同程度）上找到匹配，从而促进特征的增强优化。这种方法的有效性通过作者在Ablation Studies部分中的实验进行验证。
对于解码器阶段的特征集选择，作者采用分段形式进行如下形式化表示：
对于第一个解码器阶段（），所有编码器特征都被选择。对于后续阶段，先前的解码器阶段输出通过替换其横向编码器对应物被传递。
为了使中的特征的空间维度对齐，作者采用了Wang等人（2021年）提出的_空间降维_过程。
在这里，表示特征集的第j个元素，是池化比例，它使特征图的大小与最小的特征图对齐。操作AvgPool和Linear分别配置为AvgPool(, )()和Linear(, )()。
将空间对齐的特征沿通道维度拼接在一起，形成一个混合特征用于Key和Value。
作者改编了传统的Transformer解码器块，如Shim等人（2023年）所建议的那样，丢弃了自注意力模块。此外，作者将交叉注意力模块替换为作者提出的mix-attention模块。得到的结构如图4所示。
使用层规范化（LN）和前馈网络（FFN），第的输出计算如下：
其中表示作者提出的mix-attention。
作者提出了U-MixFormer作为UNet类似架构。然而，作者想要强调作者的方法与其他UNet类似变体的主要区别。
因为作者把横向连接看作是 Query 特征，所以解码器特征图在解码器阶段之间无需显式上采样即可隐式地增加空间分辨率。
作者的方法使用所有解码器阶段来预测分割图，而不仅仅是最后的那个阶段。
最后一个解码器阶段的特征图产生了一个分辨率为，的图像，而其他解码器阶段的特征图恢复了原始的空间分辨率，。
实验是在两个流行的基准数据集上进行的：ADE20K和Cityscapes。ADE20K是一个严格的场景解析基准测试，它突出了150个复杂的语义概念，这些概念被分为20,210张用于训练的图像和2,000张用于验证的图像。Cityscapes包括19个从城市图像中聚合的密集标注物类别，这些类别共有5,000张高分辨率的图像，分辨率为2048 1024。此外，它还引入了19,998张大致标注的图像，用于增强模型的训练。
为了评估U-MixFormer在不同编码器复杂性下的适应性，作者采用了三种不同的编码器 Backbone ：Mix Transformer (MiT)、Light Vision Transformer (LVT)和多尺度卷积注意力编码器（MSCAN）。具体来说，作者使用了MiT-B0、LVT和MSCAN-T作为轻量级模型，使用MiT-B1/2和MSCAN-S作为中量级架构，而更重的变体包括MiT-B3/4/5。最后阶段（MLP）的嵌入维度为轻量级模型的128，其他模型的维度为768。补充材料中的A.1部分提供了关于训练和评估设置的额外信息。
作者将作者的结果与现有的语义分割方法在ADE20K和Cityscapes数据集上进行了比较。表1展示了作者的结果，包括参数数量、浮点运算（FLOPs）和mIoU，这些结果跨越了两个数据集。如图1所示，作者在Cityscapes和ADE20K验证集上绘制了不同方法的表现-计算曲线。
轻量级和中量级模型在表1的顶部部分展示了轻量级模型的性能。如表所示，作者的轻量级U-MixFormer-B0在ADE20K上具有41.2%的mIoU，6.1M的参数和6.1 GFLOPs，在FLOPs和mIoU方面都优于所有轻量级对比模型，展示了性能-计算的更好权衡。值得注意的是，与使用相同编码器（MiT-B0）的SegFormer和Feedformer相比，U-MixFormer在mIoU方面提高了3.8%和2.0%，同时将计算减少了27.3%和21.8%。在Cityscapes上的性能差异更加明显，作者的模型仅使用101.7 GFLOPs就实现了79.0%的mIoU，与SegFormer-B0和FeedFormer-B0相比，分别提高了2.8%和1.1%，并将计算减少了18.9%和5.3%。当使用LVT时，作者的模型在数据集上的性能进一步得到提升，mIoU提高了2.7%和1.3%。此外，作者的U-MixFormer与MSCAN-T，来自SegNeXt的最新编码器，在ADE20K和Cityscapes上分别取得了44.4%和81.0%的mIoU，使用6.7M的参数。
表1的后者部分将重点放在了中量级模型上，作者的方法在中量级模型上继续展示了优越的结果，保持了相对于竞争对手的优势。
表2详细介绍了U-MixFormer在搭配相同的重量级编码器时优于SegFormer的情况，特别是与MiT-B3/4/5的搭配。例如，在ADE20K上，U-MixFormer-B3获得了49.8%的mIoU，仅需56.8 GFLOPs。这显示了与SegFormer-B3相比，mIoU提高了0.4%，并将计算减少了28.1%。此外，作者还推测，通过扩大模型大小（从MiT-B0更改为MiT-B5）可以允许从编码器阶段提取更丰富的上下文信息，可能有助于提高性能。因此，作者训练并评估了重量级模型变体，包括MiT-B4和MiT-B5，并引入了一种从编码器堆叠处提取额外的Key和Value以促进mix-attention的方法。作者称这种增强的变体为U-MixFormer+。对于MiT-B4和MiT-B5配置，作者从编码器的第三阶段中点提取了5和6个Key和Value，以便于mix-attention。结果，当整合更多来自编码器的上下文数据时，作者观察到MiT-B4的性能提高了0.8%，而MiT-B5的性能提高了0.1%，且仅在计算需求上略有增加。
如图5所示，U-MixFormer，FeedFormer和SegFormer在ADE20K和Cityscapes数据集上使用相同的编码器，并呈现了各自的定性结果。与其它方法相比，U-MixFormer在更清楚地分割复杂的物体细节和具有挑战性的区域方面表现出色。它可以显著识别出语义相关的区域和物体细节，这正是它从多阶段编码器中学习上下文特征表示以实现高效分割的能力。
在表3中，作者对不同的设计选择进行了系统性的评估。为了保证比较的公平性，所有模型都在一个统一的随机种子下进行训练和评估。作者设计 FeedFormer，它基于传统的交叉注意力，作为作者的 Baseline 。通过_mix-attention_模块将来自多个编码器阶段的上下文信息集成，可以在mIoU方面提高0.7%，同时降低计算成本。适配一个没有mix-attention的U-Net Transformer 解码器可以提高mIoU 0.9%，但计算成本略有增加（+0.4）。值得注意的是，通过在U-Net架构中应用mix-attention模块，模型的性能增加到41.2%，这意味着在U-Net类似的配置中，传统交叉注意力的性能有了显著的提高。
在像自动驾驶和智能交通系统这样的关键领域中，图像分割的鲁棒性至关重要。在这方面，作者对U-MixFormer的鲁棒性进行了评估，以应对腐蚀和干扰。作者遵循[10]中的方法，引入了Cityscapes-C，这是Cityscapes的_val_的增强版本，它涵盖了16种算法腐蚀，包括噪声、模糊、天气和数字类别。作者将U-MixFormer与共享相同编码器的SegFormer和FeedFormer进行比较。
表4中的结果突显了U-MixFormer的优越鲁棒性。值得注意的是，它在与所有腐蚀类别显著改进，其中与射击噪声的改进达到20.0%和33.3%，在雪天条件下的改进达到21.8%和19.2%。这些结果表明了U-MixFormer的鲁棒性，使其成为在安全性和可靠性至关重要的应用的理想选择。
为了验证基于U-Net的mix-attention模块的有效性，作者进行了一个微观上的ablation study，通过与常见研究（MLP和cross-attention）中的方法（提取相同位置的特征图）进行比较，如图6所示。作者观察到，作者的方法明显优于其他方法。具体来说，作者的特征图更加精确地描绘了物体细节，并清晰地划定了物体之间的边界。这一观察表明，所提出的方法可以在局部和全局上显著分割和捕捉可区分的视觉细节。
为了将作者的方法与最先进的技术进行比较，作者采用了与SegNeXt中相同的一致编码器，具体采用MiT和MSCAN。这些多阶段设计编码器因其效率和创新设计而受到广泛关注。正如SegNeXt中的结果所示，MiT和MSCAN都实现了更高的平均交并集（mIoU）分数和减少了计算开销，其更少的FLOPs表明了这一点。这一分析对于突出作者方法相对于这些 established 编码器的优势至关重要。
如表1和2所示，作者的U-MixFormer在各种模型复杂性（从B0到B5）上始终优于mIoU和FLOPs。表5进一步突显了U-MixFormer的优越性，在ADE20K上比更重的SegNeXt-S实现了3.3%的mIoU增加。此外，MSCAN-S与更强大的MSCAN编码器（SegNeXt-B）的性能非常接近。这些发现表明U-MixFormer是语义分割中一个有前途的解码器架构。
尽管在计算成本和mIoU方面，作者的U-MixFormer具有竞争力的结果，但需要解决一些局限性。作者在使用单个A100 GPU的_mmsegmentation_基准设置下，测试了单个2048 1024图像的推理时间。
如表6所示，与其他轻量级模型相比，U-MixFormer的推理时间较慢。延迟可以归因于U-Net固有的结构，需要通过横向（或残差）连接来保持信息。虽然这对于捕捉层次特征是必需的，但这些连接在推理阶段引入了开销。为了应对这个限制，作者计划在未来的工作中探索模型压缩技术，如剪枝和知识蒸馏，以提高推理速度，同时保留U-MixFormer的准确性优势。这些方法有望潜在地提高推理速度，同时保留U-MixFormer的准确性。
在本文中，作者提出了U-MixFormer，它基于U-Net结构设计用于语义分割。U-MixFormer从最上下文的编码器特征图开始，逐步纳入更细微的细节，这依赖于U-Net捕获和传播层次特征的能力。作者的mix-attention设计强调了合并特征图的组成部分，并将它们与越来越细粒度的横向编码器特征对齐。这确保了高层次上下文信息与复杂低级细节的和谐融合，这对于精确分割至关重要。作者在流行的基准数据集上展示了U-MixFormer相对于不同编码器的优越性。
[1].U-MixFormer: UNet-like Transformer with Mix-Attention for Efficient Semantic Segmentation.
点击上方卡片，关注「集智书童」公众号
内容中包含的图片若涉及版权问题，请及时与我们联系删除
评论
沙发等你来抢