Motivation:
由于卷积核的局部特性,传统的基于cnn的分割模型缺乏建模长期相关性的能力。为了解决这样的问题,已经开发了各种强大的关系建模方法。
(a)基于空间金字塔的方法, 采用不同大小的卷积核将不同范围的上下文信息聚集在一层中
(b)基于UNet[16]的编解码器网络[16,24,12]采用跳跃连接的方式将相同尺度的粗粒度深度特征和细粒度浅层特征合并
虽然这些方法在密集像素预测方面取得了很大的成功,但全局关系的建模仍然是低效的,使得进一步提高分割精度变得困难
Transformer是可以解决这种问题,之前的transformer©在最高级CNN特征上计算全局上下文,在特定尺度下保证各种范围依赖关系
这样的设计对于医学图像分割来说仍然是不完美,首先,它只使用自注意机制在单个尺度上进行上下文建模,而忽略了跨尺度的依赖性和一致性。后者通常在大小变化较大的病变的分割中起关键作用。其次,除了上下文建模之外,如何学习不同语义类别之间的相关性以及如何保证同一类别区域的特征一致性仍然没有考虑到。但它们都成为了基于cnn的分割方案设计的关键。
Contribution:
(1)我们提出MCTrans,它可以构建跨尺度的上下文依赖关系和适当的语义关系。
(2)引入一种新的可学习的proxy embedding方法,分别构建类别依赖关系、通过自注意和交叉注意增强特征表示。
(3)我们将设计的MCTrans插入到一个类似UNet的网络中,并在分割数据集上评估其性能。实验结果表明,MCTrans在所有任务的计算量上都有轻微的提高,显著优于先进的方法,证明了所有提出的网络组件的有效性。
Method:
A、Unet
采用深度CNN提取不同尺度的多层次特征𝑋_𝑖。将i = 2,3,4级的特征串联起来,形成整体token T。
在这篇文章中Patch被设置成1,意思就是将第i个feature map的每个位置视为patch。将不同层次的特征图通过1*1的卷积形成相同通道数,得到嵌入标记𝑇_𝑖∈ℝ^(𝐿_𝑖×𝐶_𝑒 )。然后将i = 2;3;4级的特征串联起来,形成整体token𝑇∈𝑅^(𝐿×𝐶),为了弥补缺失的位置信息,对标记进行了位置嵌入Epos的补充,以提供关于该特征在序列中的相对或绝对位置的信息。
B、Transformer-Self-Attention (TSA)
对多个特征之间的上下文信息进行编码,产生丰富而一致的像素级上下文。
多头自注意网络(MSA)和前馈网络(FFN)组成,但是由于token具有非常长的序列,我们使用Deformable Self - Attention (DSA)机制来替代SA。
DSA是一个稀疏attention,并不是全部key被查询的,只关心参考点周围的一组关键采样点,而不考虑特征图的空间大小。通过为每个查询分配少量固定数量的键,可以缓解收敛性和特征空间分辨率问题。
C、Transformer-Cross-Attention (TCA)
引入了语义关系建模的可学习嵌入,进一步增强了特征表示。
我们将不同分辨率上的token输入TSA模块以进行建模。输出增强的tokens将与proxy嵌入通过TCA模块进行交互。最后,我们将编码后的token折叠回金字塔特征,以获得用于预测的最终特征映射。
经过Ks层的自注意力,然后把增强后的tokens 𝑇^(𝐾_s )输入交叉注意中,同时还有一个可学习的proxy嵌入𝐸_𝑝𝑟𝑜学习全局语义关系(类内和类间的)。
每个TCA模块包含两个多头自注意块。具体来说,对于第j层,代理嵌入𝐸_𝑝𝑟𝑜^(𝑗−1)通过各种线性变换,得到第一个MSA块的输入(查询、键、值)。MSA块的自注意机制与每一类别相互连接和交互,从而建模各种标签的语义相关性。
接下来,学习到的代理嵌入通过另一个MSA块中的交叉注意提取输入tokens 𝑇^(𝐾_s )的特征并与之交互,query输入是代理嵌入,key和value输入是tokens 𝑇。通过交叉关注,token的特征与学习到的全局语义关系进行通信,全面提高了特征表示的类内一致性和类间可判别性,生成更新的代理嵌入E j pro。
Experiment:
消融实验
在UNet的最高级别特性上使用了Non-local操作和ViT-Encoder来实现单尺度上下文建模。
进一步评估了TCA、TSA模块和辅助损失的影响,说明了语义关系学习对特征表示的有效性。
TCA、TSA模块数
随着TSA尺寸的增大,DSC评分先增大后减小。固定Ns后,在Nc = 4时达到顶点。对于小数据集,并不是Transformer的容量越大越好。
细胞核分割
病灶分割