点击下方卡片，关注「集智书童」公众号
导读
虽然卷积神经网络长期以来被认为是图像处理领域的最佳技术，但Transformer架构的引入对其地位提出了挑战。虽然Transformer在图像分类和分割方面取得了优异的成绩，但它们本质上仍然依赖于大量的训练数据集，且计算成本较高。
一种新引入的Transformer衍生品——KV Transformer，在合成、自然语言处理和图像分类任务中显示出有希望的结果，同时降低了复杂性和内存使用。这特别有利于需要局部推理的应用场景，例如医学筛查应用。作者努力进一步评估KV Transformer在语义分割任务上的优势，特别是在医学影像领域。通过直接比较相同基础架构的传统和KV变体，作者提供了关于降低模型复杂性的实际权衡的进一步见解。作者观察到参数数量和乘加操作显著减少，而在与QKV实现直接比较时，大多数KV变体模型实现了相似的性能。
1 引言
1.1 动机
卷积神经网络（CNNs）长期以来一直是图像分类和分割的顶尖技术，尤其是在医学影像领域。然而，由于CNNs依赖于局部感受野和卷积核，它们在捕捉长距离依赖关系和全局上下文方面存在局限性，这可能导致全局信息的丢失。视觉Transformer（ViTs）[1]，源自于在自然语言处理（NLP）中使用的Transformer [2]，通过自注意力机制有效地捕捉长距离依赖关系和空间关系来解决这个问题。然而，由于自注意力机制的二次复杂性和缺乏权重共享，与使用卷积核的CNNs不同，Transformers和ViTs在计算上仍然比传统的CNNs更昂贵。为此，Borji [3] 提出了KV Transformer，以研究在保留性能的同时减少模型复杂性的可能性，通过省略注意力机制中使用的其中一个向量。这在资源受限的环境，如医学筛查中特别有益，在这些环境中，高效的推理更受欢迎。因此，作者致力于评估KV注意力在纯和混合Transformer模型中的性能，以及在医学图像分割任务中的应用。
1.2 相关工作
首次由 Vaswani、Shazeer、Parmar 等人于 2017 年提出[2]，Transformer 通过其新颖的自注意力机制的使用，彻底改变了自然语言处理（NLP）领域，并随后通过图像分类任务对计算机视觉产生了影响。Dosovitskiy、Beyer、Kolesnikov 等人[1]将 Transformer 架构应用于图像分类任务，通过将输入图像分割成固定大小的块来实现。每个块被扩展成一个向量，然后通过一个线性变换投影到一个更高维度的表示中，该变换定义为：
，其中
是扩展后的块向量，
是权重矩阵，
是偏置向量，
是得到的高维表示。为了解决与 CNN 相比，Transformer 内在缺乏空间归纳偏差的问题，添加了位置嵌入以保留块的位置。一旦块被嵌入，它们就被输入到一个标准的 Transformer 编码器中，该编码器应用自注意力机制。自注意力层的输出可以通过以下方式计算：
注
意
力
机
制
（
）
其中，
、
和
分别是 Query 、 Key和Value 矩阵，
是键的维度。最后， Patch 表示通过多层感知器（MLP） Head 处理，将高维Transformer输出转换回特定任务的预测。
对于图像分类，附加的类别 Token 通过MLP Head 处理，以产生单个类别的预测。2020年，郑、卢、赵等人[4]介绍了SETR（SEgmentation TRansformer），它使用基于Transformer（ViT）的编码器进行语义分割（图1a）。通过使用具有渐进上采样（PUP）的卷积解码器来实现语义分割。在此，从Transformer编码器输出的高维特征图被重建为输入图像的维度。编码器的输出是一个形状为
的 Token 序列，其中
是特征维度。该序列首先被 Reshape 为形状为
的特征图，然后通过四个卷积层，每个层执行2
上采样操作。结果特征图的形状与输入图像相同，深度等于要分割的类别数
（图2）。
在传统的Transformer中，自注意力机制为序列中的每个 Token 动态计算三个组件， Query （Q）、键（K）和值（V），从而确定每个 Token 相对于其他 Token 的重要性。Borji [3] 通过研究是否需要包含 Query 、 Key和Value 来实现有竞争力的性能，扩展了原始Transformer架构及其各种衍生版本。标准Transformer中的自注意力机制由公式1给出。权重矩阵
和
是负责将输入
投影到维度为
的输出张量的参数，其中
，和
。这种公式依赖于计算 Query （
）中的每个 Token 与键（
）中相应 Token 的点积。Borji [3] SETR-PUP通过将
替换为
来修改这一点，使得
KV自注意力机制最初创建一个对称结构。为了引入非对称性，Borji [3] 建议通过添加维度为
的2D位置编码来增强
的注意力矩阵。这种修改将矩阵转换为
的维度。然后，使用一个由
个神经元组成的线性层将这个扩展矩阵投影回其原始的
大小（图1b）。表1说明了与KV注意力相关的复杂性和计算成本预期的优势。KV
Pos注意力导致参数和FLOPS计数增加，这取决于参数
。需要注意的是，
可以设置得比
低得多。然而，KV+Pos注意力的一个显著缺点是其计算复杂度依赖于
，这可能导致高计算成本。
TransUNet [5] 提出了一种混合（CNN和Transformer）模型。编码器利用ResNet-50架构 [6] 从输入图像中提取High-Level特征。TransUNet不是将图像分割成 Patch ，而是将CNN生成的特征图应用于 Patch 嵌入，创建特征 Patch 作为Transformer的输入。这种方法允许更好地保留空间信息并捕捉到细粒度细节。然后，Transformer层处理这些 Patch 以建模长距离依赖关系，增强网络整合全局上下文与局部特征的能力。Wu，Xiao，Codella等人 [7] 引入了卷积视觉Transformer（CvT），这是一种混合模型，通过集成两个关键的卷积操作（ Token Embedding 和投影）来增强视觉Transformer（ViT）。
在CvT中，卷积 Token Embedding 通过应用卷积层来替换ViT的 Patch 分割。这种方法允许模型通过应用多个在整个图像上滑动的过滤器来有效地利用局部信息。结果的特征图保持了图像的空间结构，同时减少了其维度。卷积投影在将特征图传递到自注意力机制之前对其应用进一步的卷积。与原始ViT相比，这种混合方法提高了准确性，减少了参数数量，并提高了计算效率，尤其是在图像分类任务中。
2 材料与方法
2.1 数据集
所使用的医学图像数据集为UW-Madison GI Tract Image Segmentation [8]，该数据集包含腹部MRI切片。三个类别（大肠、小肠、胃）的标注以运行长度编码的器官分割形式提供。在预处理过程中，作者将RLEGT数据转换为2D灰度多类别 Mask 。此外，数据集按照80:16:4的比例分为训练集、验证集和测试集。
2.2 模型架构
为了探索不同架构下的KV注意力机制，作者实现了具有KV和QKV注意力变体的多个模型，具体如下。所有模型都共享一个从SETR（第1.2节）改编的卷积解码器。作者修改了解码器，在上采样过程中将特征维度减半，从而减少了整体参数数量（图2）。作者首先按照[4]中所述实现了SETR编码器，使用ViT-B/16主干网络，特征维度
，头数
，层数
，并使用QKV和KV注意力机制。作者将这些架构分别称为SETR-QKV和SETR-KV。此外，作者还探索了SETR-KV-pos，作者在KV注意力块中引入位置编码以创建不对称性，遵循[3]中SETR-PUP。2D位置编码维度
设置为50。
此外，作者构建了两个具有混合编码器的模型。从TransUNet[5]中汲取灵感，作者将ResNet-50架构[6]的前四个卷积层集成到作者的编码器中，以在 Patch 嵌入阶段之前捕获更高维度的特征。在第四层，作者将块的数量从六个增加到九个，以提高特征提取能力，同时保持特征维度为1024。与[5]中的方法不同，作者没有使用 Shortcut 。作者将这些模型分别称为SETR-QKVCE和SETR-KV-CE。最后，作者开发了一个使用卷积视觉Transformer（CvT）[7]作为编码器的混合模型。SETR-QKV-CVT和SETR-KV-CVT模型使用CvT-13编码器，在卷积Transformer块中分别使用QKV和KV注意力实现多头注意力（MHA）。
2.3 实验
第2.2节中描述的所有模型均训练了100个Epoch，未使用提前停止策略以确保结果可比。输入分辨率设置为
，并选择了固定的块大小
。使用了学习率为
的AdamW优化器以及因子为0.9的多项式学习率调度。
此外，训练时选择了32个样本的批量大小。在训练过程中，应用了实时数据增强，包括水平翻转、垂直翻转、平移缩放旋转、粗略dropout和随机亮度对比度，每种操作有
的概率被应用。所有模型均在第2.1节中描述的数据集上训练，未使用任何预训练的 Backbone 网络。
3 结果
测试架构的性能指标包括Jaccard指数和加权Jaccard指数。模型复杂度由可学习参数的数量表示，而计算效率通过乘加操作（MACs）的数量来评估（通过torchinfo和ptflops Python模块收集）。作者的结果表明，所有KV变体在性能上与相应的QKV实现相当或略好，同时参数数量和MACs的数量也减少了大约10%（见表2）。
4 讨论
作者的结果表明，与它们的QKV注意力对应物相比，KV Transformer在分割任务上表现出具有竞争力的性能。编码器中卷积元素的应用导致参数数量的预期增加；然而，通过利用KV注意力可以减轻这一点。表1中概述的计算成本和复杂性在作者的实验中并未实现，因为解码器中卷积元素的应用增加了这两个指标。
尽管如此，这些发现激励作者进一步调查KV注意力在未来的项目中的应用。由于KV注意力在 Transformer 中的主要架构变化发生在MHA块内，因此理论上它可以适应各种ViT-like架构，这为研究该设计的鲁棒性提供了额外的机会。总之，作者研究了KV Transformer架构在分割任务中的应用，特别是在医学成像的背景下。作者的发现表明，所评估的KV变体与它们的QKV对应物具有相似的性能，同时也在参数数量和计算成本方面实现了降低。
参考
[1]. Exploring the Integration of Key-Value Attention Into Pure and Hybrid Transformers for Semantic Segmentation*
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）