突破多模态图像翻译难题:ABS-Mamba借SAM2、CNN、Mamba及LoRA+微调创佳绩 !

大模型向量数据库数据中台

点击下方卡片,关注「AI视界引擎」公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

精确的多模态医学图像翻译需要协调全局解剖语义和局部结构保真度,这一挑战因跨模态信息丢失和结构扭曲而复杂化。作者提出了ABS-Mamba,一种集成Segment Anything Model 2(SAM2)进行器官感知语义表示、专门卷积神经网络(CNN)保留模态特定边缘和纹理细节,以及Mamba选择性状态空间建模以高效处理长距离和短距离特征依赖的新型架构。

在结构上,作者的双分辨率框架利用SAM2的图像编码器从高分辨率输入中捕获器官尺度语义,同时并行CNN分支提取细粒度局部特征。

鲁棒特征融合网络(RFFN)整合这些表示,双向Mamba残差网络(BMRN)通过螺旋扫描和双向状态空间动力学建模空间依赖关系。三阶段跳跃融合解码器增强了边缘和纹理保真度。

作者采用高效低秩适配(LoRA+)微调,以实现精确的域特定化,同时保持预训练组件的基础能力。在SynthRAD2023和BraTS2019数据集上的大量实验验证表明,ABS-Mamba优于现有最佳方法,实现了高保真跨模态合成,保留了解剖语义和结构细节,从而提高临床应用中的诊断准确性。

代码可在https://github.com/gatina-yone/ABS-Mamba获取。

1 引言

多模态医学图像通过不同的组织对比机制提供互补的诊断信息,被广泛应用于临床工作流程中的疾病筛查和治疗计划[1,2]。这些模态特定的对比提供了关键的解剖学信息,并提高了计算分析的准确性[3]。然而,由于时间限制、图像损坏和多样的成像协议,获取完整的多模态数据集具有挑战性,限制了其在临床中的常规应用[4-6]。这推动了计算医学图像翻译框架的发展,这些框架提供了客观、可重复的结果,以支持临床决策和个性化医疗[7,8]。

解剖组织呈现出复杂且空间变化的分布,通常形成复杂的簇集[9]。这种复杂性要求在多模态医学图像翻译中提取细粒度的局部特征和全局依赖关系。近年来,生成对抗网络(GANs)在跨模态合成过程中表现出色,能够保留纹理,但它们经常无法确保高分辨率成像中的全局解剖学一致性[10]。相反,视觉 Transformer (ViTs)擅长捕捉长距离依赖关系,但它们的二次计算复杂度限制了其在高分辨率医学数据集上的效率[11]。

状态空间模型领域的最新进展,如Mamba,通过利用状态空间算子高效地建模像素序列,以捕捉长程和短程的上下文依赖关系[12-17]。这些模型借助微分方程来增强医学图像中纹理与解剖结构之间复杂交互的表征。与此同时,像SAM和SAM2这样的视觉基础模型在多样化的视觉任务中展现出良好的泛化能力和小样本适应性[18-20]。SAM2尤其具有架构上的改进,并在超过100万张图像和5万段视频的大型数据集上进行了训练,其中包括超过110亿个语义 Mask 。结合状态空间模型的高效长程上下文建模与SAM2的先进语义理解,为解决医学图像精确翻译中的挑战开辟了新的可能性。

在这项工作中,作者提出了ABS-Mamba,一种新型图像翻译网络,旨在捕捉全局上下文语义和细节特征,以精确生成目标域图像。其核心创新源于独特的集成策略,结合了SAM2的图像编码器进行全局语义建模、基于CNN的编码器进行局部特征提取,以及Mamba的高效双向序列化交互,以优化跨模态特征互补性和计算效率。所提出的网络采用U型架构,包含三个核心组件:编码器、 Backbone 网络和解码器。

  1. 混合编码器结合了SAM2的预训练Hiera图像编码器(捕捉分层全局上下文)和专门设计的基于CNN的领域特定特征提取器,通过RFFN融合其输出。
  2. Backbone 网络BMRN采用通道混合的Mamba模块,以促进双向交互,通过螺旋扫描轨迹在特征图上进行空间上下文学习,实现平衡的上下文敏感性。
  3. 解码器通过对应编码器-解码器层之间的多尺度 Shortcut 保留Low-Level空间特征,并通过残差连接拼接编码器产生的高频分量,确保在临床级图像翻译中保留解剖细节。

作者对ABS-Mamba在两个公共数据集上进行了广泛评估:一个包含大脑和骨盆的CT和MRI图像,另一个由大脑的T1和T2加权MRI序列组成。定量结果表明,ABS-Mamba在合成目标域图像方面优于当前最先进的方法,并在多机构临床环境中表现出强大的泛化能力。

2 方法论

2.1 ABS-Mamba概述

ABS-Mamba是一种新型架构,它协同集成了SAM2的视觉Transformer与基于CNN的局部特征提取器,并通过最新的Mamba模块增强纵向上下文建模能力。如图1所示,混合编码器采用冻结的SAM2图像编码器,其16×16的块嵌入用于处理分辨率为512×512的高分辨率输入图像,而CNN分支则从256×256的低分辨率图像中提取层次化特征。RFFN动态融合多尺度特征,使模型能够学习到统一表示,从而协调双分辨率医学成像协议中互补信息的融合。在训练过程中,作者仅对SAM2解码器参数的19.47%应用低秩适配(LoRA+),以保持预训练模型的诊断泛化能力,并实现特定领域的精细化调整[22]。

picture.image

Backbone 架构集成了双向通道混合Mamba模块,以有效捕获医学图像序列中的长程上下文依赖关系,同时保持高局部保真度,解决了医学图像数据固有的非因果特性。BMRN模块采用螺旋扫描轨迹对特征图中的各向同性空间上下文进行建模,并结合自适应通道混合层,分层聚合跨通道上下文信息,特别针对医学图像翻译中的解剖一致性进行设计。

ABS-Mamba的解码器通过三阶段过程实现高保真度医学图像翻译。在解码过程中,它通过 Shortcut 逐层整合编码器特征,并包含两个混合混淆特征层和一个CNN特征层。

2.2 混合编码器

混合编码器采用跨分辨率分层特征融合设计,包含两个并行路径:1) CNN编码器处理低分辨率输入

;2) 预训练的SAM2-Hiera图像编码器处理高分辨率输入

,两者共同提取多尺度全局上下文特征。通过RFFN实现跨模态特征融合,该模块对每个通道应用注意力机制,对齐并整合两个路径的异构特征表示,确保下游任务的多尺度特征聚合具有鲁棒性。

鲁棒特征融合网络(RFFN):如图1c所示,作者基于交错分割策略设计了一种多尺度深度可分离卷积架构,通过特征通道重组和多核并行处理实现高效特征融合。混合编码器第i层的输出特征张量定义为F,E RCinHW,其中

表示输入通道维度。首先,根据模2余数准则,将输入特征沿通道维度解耦为两个互补子空间:

其中

。基于特征金字塔理论,通过深度可分离卷积在解耦的子特征上分别实现不同尺度的卷积核(

),增强跨尺度特征完整性 [23]。多核处理之后,分解的特征会经历:

其中

控制负信号保留率,增强了对对抗样本的鲁棒性。

表示 1×1 点卷积。

跨分辨率层次特征融合:令

表示CNN编码器第

层的输出特征。SAM2-Hiera编码器输出四个层次化的

。其中

[24]。令

表示通过

融合的第

层和第

层的特征。编码器的第一层处理单通道输入

,在保持分辨率的同时将通道维度扩展至3。反射填充(nn.ReflectionPad2d(3))将其扩展至

,通过镜像边界减少伪影。一个

卷积(步长为2)输出

的特征图,随后进行BatchNorm和ReLU激活,得到

。对于第二层,一个

的步长为2的卷积将

下采样至

。经过归一化和激活处理后,RFFN通过残差连接整合

,生成融合特征

。第三层遵循相同范式,进一步下采样至

,并与

进行二次融合,最终生成

2.3 双向Mamba残差网络

作者提出的BMRN如图1b所示,该模型采用螺旋扫描机制将二维特征表示转换为渐进式螺旋序列。形式上,给定输入批次

,该模块首先应用维度 Reshape 生成归一化特征

。通过预定的螺旋变换矩阵

,执行空间映射以获得

如图1a所示,双向螺旋轨迹设计能够有效捕获长程上下文信息。该足迹展示了每个学习算子部署的焦点点的空间分布,这些焦点点在寻求中心像素的上下文交互时被使用。对于正向和反向扫描序列,作者采用独立的Mamba进行特征提取。输入张量

(记为

)在输入到两个独立的状态空间模块Mamba

和Mamba

A之前会进行层归一化。

其中

分别表示Leaky ReLU激活函数和逐元素乘法。在将双路径特征进行拼接后,通过残差卷积模块进行特征融合:

其中

定义为两个连续的3x3卷积层,每个卷积层后接ReLU激活函数和批量归一化(BN)。

3 实验与结果

3.1 数据集和实现细节

为验证所提出的ABS-Mamba在医学图像翻译任务中的有效性和性能,使用SynthRAD2023 Grand Challenge数据集[25]和BraTS 2019公共数据集[26]进行实验。实验具体针对两个关键挑战:脑部和盆腔CT到MRI翻译中的金属伪影抑制,以及T1加权图像和T2加权图像之间的脉冲序列参数映射。数据集根据患者ID随机分为70%的训练集和30%的测试集,SynthRAD数据采用轴向层面级分层抽样,BraTS 2019遵循官方划分策略[26]。

作者将ABS-Mamba与ResViT [27]、VMamba [28]、TransUNet [29]、DiffMa [30]、I2I-Mamba [31]和UMamba [32]等几种最先进的方法进行比较。性能评估采用两种成熟指标:结构相似性指数(SSIM)和峰值信噪比(PSNR)。所有模型均在配备24GB内存的NVIDIA RTX 4090 GPU上进行训练,批处理大小设置为4。优化过程采用Adam优化器(

),基础学习率为2e-4,并使用余弦退火调度衰减至1e-7。混合损失结合了对抗损失(

)和像素级L1损失(

)。

自适应层(

)被集成到SAM2Encoder模块中,使整个架构中有19.47%的可训练参数, Adapter 参数比例为3.2e-3。

3.2 结果

为使用SynthRAD2023 Grand Challenge数据集(脑部和骨盆)评估ABS-Mamba在CT到MRI图像转换任务中的性能,并与现有最优方法进行比较,表1中的全面基准测试结果表明,ABS-Mamba在多个指标上均优于现有方法。

picture.image

它在脑部数据集上实现了最高的SSIM和PSNR,在骨盆数据集上实现了最高的PSNR。与次优方法相比,它在骨盆数据集上的SSIM差异微乎其微,展示了ABS-Mamba的总体优越性。

作者进一步在BraTS 2019数据集上验证了ABS-Mamba,以评估图像翻译质量。定量结果如表2所示,ABSMamba实现了0.9350±0.009的SSIM和31.70±0.56 dB的PSNR,在SSIM和PSNR方面分别优于次优方法0.65%和4.03%。这些结果表明ABS-Mamba在处理不同临床站点和扫描仪之间的变化方面的有效性。

picture.image

图2展示了SynthRAD2023大挑战数据集和BraTS 2019数据集中代表性样本通过不同方法获得的图像翻译结果的视觉比较。这些定性比较突出了ABS-Mamba在实现准确翻译结果方面的优越性能。

picture.image

3.3 消融实验

为探究所提模块的有效性,作者在BraTS2019数据集上进行了消融实验。作者评估了整合SAM2图像编码器(SIE)和RFFN,以及改变BRMN模块数量对图像翻译性能的影响。结果汇总于表3。当同时整合SIE和RFFN模块与BRMN时,性能最佳,默认配置的2个BRMN模块展现出0.934的平均SSIM。这些结果验证了融合SAM2编码器提供的全局语义表征、CNN编码器提取的医学图像局部细节特征,以及Mamba高效序列双向交互机制,实现互补跨模态特征计算的联合优化的有效性。

picture.image

4 结论

本文提出了一种名为ABS-Mamba的新型医学图像翻译框架,该框架巧妙地融合了SAM2视觉基础模型的全局解剖学先验知识与Mamba的高效状态空间建模。其架构创新通过三个关键组件实现:

1)结合SAM2-Hiera解剖学先验知识保留与基于CNN的局部特征提取的双流编码器;

2)螺旋扫描的双向Mamba模块,实现多尺度上下文依赖的线性复杂度建模;

3)具有不确定性感知融合的层级 Shortcut ,用于精确的高频细节重建。

在多中心脑部和盆腔数据集上的实验验证表明,ABS-Mamba在合成高质量目标图像时具有精确的解剖学准确性。这一创新有望通过减少多次成像程序的需求,从而降低辐射暴露并降低医疗成本,进而转变患者护理模式。这一进展提升了患者安全性,并提高了资源配置效率。

点击上方卡片,关注「AI视界引擎」公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
火山引擎大规模机器学习平台架构设计与应用实践
围绕数据加速、模型分布式训练框架建设、大规模异构集群调度、模型开发过程标准化等AI工程化实践,全面分享如何以开发者的极致体验为核心,进行机器学习平台的设计与实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论