超越UNet与ViT:混合Transformer模块MTM如何革新医学图像分割
1. 医学图像分割的“老难题”与“新挑战”
如果你接触过医学图像分析,比如从CT、MRI影像里把肿瘤、器官的轮廓给“抠”出来,那你对UNet这个名字一定不陌生。它就像这个领域的“老黄牛”,凭借经典的U型编码器-解码器结构,加上跳跃连接,在过去几年里立下了汗马功劳。我最早用它做细胞分割的时候,感觉真是顺手,效果也够用。但用久了,尤其是在处理一些大尺寸、结构复杂的图像时,比如要分割一整幅腹部CT里相互挨着的多个器官,问题就来了。
你会发现,有些分割出来的边界总是不太准,或者一些明明离得很远但属于同一类别的区域,模型就是“看”不到它们之间的联系。这背后的根本原因,就是卷积神经网络(CNN)的“天生缺陷”:局部感受野。卷积核每次只能看到图像上的一小块区域,它擅长捕捉纹理、边缘这些局部特征,但对于理解整幅图像中所有部分之间的全局关系,就显得力不从心了。这就是所谓的“长程依赖”建模能力不足。医生看片时,一眼就能根据器官的整体形态和相对位置做出判断,但传统的UNet做不到这一点。
那怎么办呢?近两年,Transformer这个在自然语言处理领域大杀四方的模型被引入了计算机视觉,诞生了Vision Transformer(ViT)。它的核心——自注意力机制,理论上可以计算图像中任意两个像素(或图像块)之间的关系,完美解决了长程依赖的问题。我当时也兴奋地尝试过,但很快就遇到了新的“坑”。
首先,ViT对数据“胃口”极大,需要在大规模数据集(比如ImageNet)上进行预训练,才能学到好的特征表示。但医学图像数据标注成本极高,公开数据集规模通常有限,直接拿过来用,效果往往不尽如人意。其次,自注意力机制的计算复杂度是像素数量的平方级。一张224x224的图片,拆成 patch 后,计算量已经不小了,对于更高分辨率的医学图像,这个计算成本是很多研究团队和临床机构难以承受的。最后,还有一个容易被忽视但很重要的问题:标准的自注意力只在单个样本内部计算关系,它忽略了不同样本之间可能存在的共性模式。比如,所有健康人的左心室形状都有一定的相似性,这种跨样本的统计规律,单样本自注意力是学不到的。
所以,我们似乎陷入了一个两难境地:UNet轻量、易用,但“眼光”不够长远;ViT“视野”开阔,但“饭量”大、“体力”消耗也大。有没有一种方法,能取两者之长,补两者之短呢?这就是今天要深入聊的MT-UNet及其核心——混合Transformer模块(Mixed Transformer Module, MTM) 要解决的问题。它不依赖大规模预训练,计算更高效,还能同时捕捉样本内和样本间的信息,可以说是为医学图像分割这个特定场景“量身定制”的解决方案。
2. MT-UNet的整体设计哲学:把好钢用在刀刃上
MT-UNet的整体架构一眼看去,还是我们熟悉的U型编码器-解码器结构,有跳跃连接来融合浅层细节和深层语义信息。这很好,意味着我们之前基于UNet的开发经验和直觉大部分都能沿用,学习成本低。但它的精髓,也是与普通UNet最大的不同,在于它并非在所有层级都使用Transformer。
仔细看它的网络设计图,你会发现那个创新的混合Transformer模块(MTM) 只被放在了编码器路径的较深层次。这是什么道理呢?我来打个比方。处理图像就像我们看一幅画,一开始我们凑得很近,关注的是笔触、颜料颗粒这些极其细微的局部细节(对应网络浅层的高分辨率特征图)。在这个阶段,使用卷积操作就非常合适,因为它天生就是捕捉局部模式的专家,而且能为模型引入一些关于图像结构(如边缘、角落)的先验知识,这对于数据量不大的医学图像任务至关重要,相当于给了模型一个“好起点”。
随着我们一步步深入网络(特征图空间尺寸变小,通道数变多),我们的“视野”逐渐拉远。这时,我们不再只看笔触,而是开始看物体的形状、各部分之间的相对位置关系(对应深层的低分辨率、高维语义特征)。在这个阶段,建立长程依赖、理解全局上下文就变成了首要任务。此时,再使用局部操作的卷积就会显得捉襟见肘,而Transformer的自注意力机制正好大显身手。
MT-UNet的设计哲学正是“把好钢用在刀刃上”:
- 浅层(早期):使用卷积。专注于提取局部细节特征,并注入归纳偏置,稳定训练,减少对海量数据的依赖。
- 深层(后期):引入MTM模块。在特征图尺寸已经通过下采样变得较小的时候,计算全局注意力的成本已经大大降低。在这里,MTM开始发挥其强大的全局建模和跨样本关联能力。
这种混合架构是一种非常务实的工程思维。它既避免了ViT从一开始就处理所有全局关系带来的巨大计算开销和“数据饥渴”,又弥补了UNet在深层语义理解上的不足。整个网络的下采样和上采样依然使用步长卷积和转置卷积来实现,保持了端到端训练的简洁性。接下来,我们就来拆解这个核心的MTM模块,看看它内部到底有什么魔法。
3. 混合Transformer模块(MTM)的核心:两把利剑
MTM模块是MT-UNet性能提升的关键,它不是一个单一的注意力机制,而是一个精心组装的“组合拳”,主要由两部分构成:局部-全局高斯权重自注意力(LGG-SA) 和 外部注意力(EA)。你可以把它理解为一个拥有“双重思维”的处理器。
3.1 第一把剑:局部-全局高斯权重自注意力(LGG-SA)
标准的自注意力(SA)虽然强大,但它是一种“平均主义”的全局计算:每个像素(或token)都要和图像中所有其他像素计算关联度。这在视觉任务中其实是有些“浪费”的。想象一下,你要识别一只猫的耳朵,肯定是它附近的眼睛、胡须区域提供的信息最相关,而远处背景的沙发纹理,关联度就低很多。为所有像素对分配同等的计算优先级,不仅计算量大,还可能引入噪声。
LGG-SA的聪明之处在于,它采用了“先局部,后全局,且局部优先”的两阶段策略,并加入了一个可学习的“聚焦镜”——高斯权重。
第一步:局部注意力(Local Attention) 首先,特征图被划分成一个个不重叠的局部窗口(比如7x7的小块)。注意力计算只在每个窗口内部进行。这就像你先用放大镜仔细审视画面的每一个小区域,理清区域内各部分的关系。这一步的计算复杂度只和窗口大小有关,与整个图像大小呈线性关系,大大降低了计算量。完成局部信息聚合后,每个窗口会被总结(聚合)成一个代表该窗口主要信息的“全局token”。这个聚合函数作者尝试了多种(最大池化、平均池化、卷积等),发现使用轻量的动态卷积效果最好。
第二步:全局注意力(Global Attention) 现在,我们手上有了一组“全局token”,每个token代表原图的一个局部区域。这些token的数量远少于原始像素数。在这个缩编的“代表委员会”层面,我们再执行一次自注意力计算,来建立各个区域(窗口)之间的全局依赖关系。这就像你听完各个部门的汇报后,再统筹思考部门间的协作关系。这一步的计算成本因为token数量减少而变得可以接受。
高斯加权:让网络学会“聚焦” 在第二步的全局注意力计算中,LGG-SA还引入了一个关键创新:高斯加权轴向注意力。它不再使用普通的注意力计算方式,而是引入了一个可学习的高斯权重矩阵。这个矩阵的作用是,在计算两个区域(token)的关联度时,会根据它们之间的相对距离进行加权。距离越近的区域,权重可能越高(取决于网络学到的方差参数)。这相当于给了网络一个可调节的“焦距”,让它能够自适应地强调或弱化不同距离的依赖关系,更加符合视觉直觉——附近的上下文通常更重要。同时,为了进一步节省计算,全局注意力采用了轴向注意力(分别计算行和列方向上的注意力),将复杂度从平方级降到了1.5次方级。
所以,LGG-SA可以总结为:用线性复杂度处理局部细节,用降低后的复杂度处理全局结构,并用一个可学习的高斯核引导模型更关注邻近上下文。它是对原始ViT中“笨重”的自注意力的一次非常针对性的、高效的改造。
3.2 第二把剑:外部注意力(EA)
LGG-SA解决了样本内部的依赖问题,但MTM的野心不止于此。它通过外部注意力(External Attention) 模块,将目光投向了样本之间。
标准自注意力(SA)的查询(Q)、键(K)、值(V)都来自输入样本自身的线性变换。这意味着每个样本都是“自闭”地计算注意力,样本与样本之间是孤立的。然而,在医学图像中,不同病人的同类器官、同类病灶,在图像形态上必然共享着某种潜在的、共同的分布规律。EA的目标就是捕捉这种跨样本的统计规律。
EA的实现思想非常巧妙。它不再为每个样本单独生成K和V,而是引入了两个共享的、可学习的记忆单元,通常记为 M_k 和 M_v。你可以把它们想象成两个存储在模型中的“公共知识库”或“原型字典”。
- 对于任何一个输入样本,它的查询向量(Q)会去和共享记忆单元
M_k计算相似度,得到注意力权重。 - 然后,这个权重被用于对另一个共享记忆单元
M_v进行加权求和,得到输出。
这个过程的好处是:
- 建模样本间关系:所有样本都通过同一套
M_k和M_v来进行注意力计算,因此它们的计算过程被“关联”了起来。网络在训练过程中,会优化这两个记忆单元,使其能够编码整个训练数据集中最有鉴别力的共性模式。 - 极致降低计算量:
M_k和M_v的尺寸是固定的(比如,[S, C],其中S是预设的记忆槽数量,通常远小于图像token数N)。因此,EA的计算复杂度从SA的 O(N²) 骤降到 O(NS),而S是一个常数,所以复杂度是线性的 O(N)。这对于处理高分辨率图像至关重要。
在MTM的设计中,EA模块被放置在LGG-SA之后。这样,特征首先经过LGG-SA的“内部梳理”(理清本样本内的局部与全局关系),然后再经过EA的“外部洗礼”(融入数据集的共性知识)。两者协同工作,使得模型获得的信息既全面又具有泛化性。
4. 实战效果:数据不说谎
理论设计得再巧妙,最终还是要靠实验效果说话。MT-UNet的作者在两个公认的医学图像分割公开数据集上进行了验证:Synapse多器官CT分割数据集和ACDC心脏MRI电影序列分割数据集。这两个数据集场景不同、模态不同,很有代表性。
为了证明MTM中每个组件的有效性,作者首先进行了详细的消融实验。结果非常清晰:
- 只使用局部注意力(Local)或只使用全局注意力(Global),模型性能都会下降,证明两者是互补且必需的。
- 移除高斯加权机制,指标也会下滑,说明这个可学习的“聚焦”功能确实帮助模型更好地利用了局部上下文。
- 在LGG-SA的基础上加入EA模块,在两个数据集的DSC(戴斯相似系数,越高越好)和HD95(豪斯多夫距离,越低越好)指标上均有显著提升(例如DSC提升约1%)。这直接证明了引入样本间关联知识对提升分割精度有实实在在的帮助。
在与其他先进方法的对比实验中,MT-UNet展现出了强大的竞争力。它不仅大幅超越了UNet、Attention UNet等基于CNN的经典模型,也超过了需要ImageNet预训练的TransUNet、Swin-Unet等早期视觉Transformer分割模型。特别是在一些解剖结构复杂、边界模糊的器官(如胰腺、胃、心脏右心室)分割上,MT-UNet的优势更加明显。从可视化结果看,它的分割轮廓更平滑、更完整,对细小结构的捕捉也更准确。
最让我个人觉得有吸引力的一点是,所有这些成绩是在单张RTX 1080Ti显卡上训练得到的。这意味着它的计算需求相对亲民,很多高校实验室和小型研究团队都有条件复现和尝试,降低了前沿技术的应用门槛。它不需要大规模的预训练,从零开始训练就能达到很好的效果,这对于数据稀缺的医学影像领域来说,是一个巨大的实用优势。
5. 总结与启发:一种务实的创新路径
回顾MT-UNet和它的MTM模块,我觉得它给医学图像分割乃至更广泛的视觉任务带来了一种非常宝贵的思路:在经典结构与前沿技术之间寻找最佳平衡点,针对领域痛点进行定制化改进。
它没有全盘抛弃经过时间检验的CNN和UNet架构,而是将其作为稳健的基石。它也没有盲目套用原始的ViT,而是深刻分析了其在高分辨率图像处理上的计算瓶颈和在有限数据下的泛化难题,从而创造了LGG-SA和EA这两个针对性极强的组件。
LGG-SA通过局部-全局分级处理和高斯加权,实现了对长程依赖的高效、精准建模。EA则通过共享记忆单元,以一种计算高效的方式挖掘了数据集的群体智慧。两者的结合,使得模型同时具备了强大的样本内理解能力和样本间知识迁移能力。
在实际项目里尝试这类模型时,我的体会是,首先要吃透你的数据特点:图像分辨率多大?目标物体是局部的还是全局结构化的?数据量有多少?标注质量如何?MT-UNet这种混合架构尤其适合数据量中等、目标结构需要全局上下文理解、且计算资源有限的场景。比如,我们在处理一些病理切片图像的分割时,就借鉴了这种“浅层卷积+深层Transformer”的思想,在保证精度的同时,确实加快了训练速度。
当然,MT-UNet也不是终点。比如,记忆单元 M_k/M_v 的大小如何设定?如何更好地与3D医学图像结合?这些都是可以继续探索的方向。但无论如何,它清晰地展示了一条路径:通过精妙的模块设计,我们完全可以在不依赖巨量数据和算力的前提下,让模型“看得更远、想得更深”。这对于推动AI在医疗等垂直领域的真正落地,意义可能比单纯刷高几个Benchmark分数更为重大。
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐



所有评论(0)