2025-11-24T09:43:19.398688

nnY-Net: Swin-NeXt with Cross-Attention for 3D Medical Images Segmentation

Liu, Tao, Dong et al.
This paper provides a novel 3D medical image segmentation model structure called nnY-Net. This name comes from the fact that our model adds a cross-attention module at the bottom of the U-net structure to form a Y structure. We integrate the advantages of the two latest SOTA models, MedNeXt and SwinUNETR, and use Swin Transformer as the encoder and ConvNeXt as the decoder to innovatively design the Swin-NeXt structure. Our model uses the lowest-level feature map of the encoder as Key and Value and uses patient features such as pathology and treatment information as Query to calculate the attention weights in a Cross Attention module. Moreover, we simplify some pre- and post-processing as well as data enhancement methods in 3D image segmentation based on the dynUnet and nnU-net frameworks. We integrate our proposed Swin-NeXt with Cross-Attention framework into this framework. Last, we construct a DiceFocalCELoss to improve the training efficiency for the uneven data convergence of voxel classification.
academic

nnY-Net: Swin-NeXt with Cross-Attention for 3D Medical Images Segmentation

基本信息

  • 论文ID: 2501.01406
  • 标题: nnY-Net: Swin-NeXt with Cross-Attention for 3D Medical Images Segmentation
  • 作者: Haixu Liu¹, Zerui Tao¹, Wenzhen Dong², Qiuzhuang Sun¹
  • 机构: ¹悉尼大学, ²香港中文大学
  • 分类: cs.CV (计算机视觉)
  • 论文链接: https://arxiv.org/abs/2501.01406

摘要

本文提出了一种名为nnY-Net的新型3D医学图像分割模型结构。该模型通过在U-Net结构底部添加交叉注意力模块形成Y型结构而得名。作者整合了两个最新SOTA模型MedNeXt和SwinUNETR的优势,使用Swin Transformer作为编码器,ConvNeXt作为解码器,创新性地设计了Swin-NeXt结构。模型使用编码器最低层特征图作为Key和Value,将病理和治疗信息等患者特征作为Query来计算交叉注意力权重。此外,基于dynUnet和nnU-net框架简化了3D图像分割的预处理和后处理方法,并构建了DiceFocalCELoss来提高不均衡体素分类的训练效率。

研究背景与动机

问题定义

肝脏肿瘤CT 3D图像分割是医学图像分割中的重要任务。准确的CT图像分割可以帮助医生估计肿瘤体积并制定合理的治疗方案。

现有方法局限性

  1. 传统方法局限: 2016年前主要使用区域生长等无监督算法,精度有限
  2. 单一架构缺陷: 现有方法要么基于纯卷积(如nnU-Net),要么基于纯Transformer(如SwinUNETR),未能充分结合两者优势
  3. 多模态融合不足: 现有方法主要处理图像信息,未能有效利用患者的病理和治疗信息
  4. 类别不平衡问题: 背景标签占据超过90%的像素,导致训练困难

研究动机

本文旨在通过融合Transformer和卷积神经网络的优势,设计一种能够同时处理3D图像和患者临床信息的多模态分割模型,以提高医学图像分割的准确性和实用性。

核心贡献

  1. 提出Swin-NeXt架构: 创新性地将SwinUNETR的编码器与MedNeXt的解码器相结合,充分利用Transformer在特征提取和卷积在像素级解码方面的优势
  2. 设计交叉注意力融合机制: 提出三种特征融合方法,发现交叉注意力机制在多模态融合中表现最佳,能够稳定提升模型性能
  3. 构建DiceFocalCELoss损失函数: 结合DiceLoss、FocalLoss和交叉熵损失,有效解决体素分类中的类别不平衡问题
  4. 简化预处理流程: 基于dynUnet和nnU-Net框架,简化并优化了3D图像分割的预处理和后处理方法

方法详解

任务定义

输入:

  • 3D CT扫描图像 χ ∈ R^(H×W×D×C)
  • 患者临床信息(病理、治疗信息等)

输出: 多类别分割掩码,包括肝脏、肿瘤、血管、主动脉等结构

约束: 处理高分辨率3D图像的内存限制,类别极度不平衡的分割任务

模型架构

1. Swin Transformer编码器

采用3D Swin Transformer作为编码器,具体实现:

  • 补丁分割: 将输入图像分割为大小为M×M×M的3D窗口
  • 窗口注意力机制:
    z^l = W-MSA(LN(z^(l-1))) + z^(l-1)
    z^l = MLP(LN(z^l)) + z^l
    z^(l+1) = SW-MSA(LN(z^l)) + z^l
    z^(l+1) = MLP(LN(z^(l+1))) + z^(l+1)
    
  • 注意力计算:
    Attention(Q,K,V) = Softmax(QK^T/√d + B)V
    

2. ConvNeXt解码器

使用MedNeXt的解码器结构:

  • 转置卷积上采样:
    Y_{i,j,k} = ∑∑∑ K_{p,q,r} · X_{(expanded)}_{i+p,j+q,k+r}
    
  • GELU激活函数:
    GELU(x) = x/2[1 + erf(x/√2)]
    

3. 交叉注意力融合模块

在编码器底层(Bottleneck)集成多模态信息:

  • Key & Value: 编码器最低层特征图
  • Query: 患者临床特征向量
  • 融合策略: 通过全连接层将临床特征映射到合适维度,然后进行交叉注意力计算

技术创新点

  1. 混合架构设计: 首次将Swin Transformer编码器与ConvNeXt解码器相结合,发挥各自优势
  2. 多模态交叉注意力: 创新性地使用交叉注意力机制融合图像特征和临床信息
  3. 组合损失函数: 设计DiceFocalCELoss解决医学图像分割中的类别不平衡问题

实验设置

数据集

  • 图像数据: 110张肝脏CT扫描图像,来自98名患者
  • 临床数据: 56个变量的患者病理和治疗信息
  • 数据清理: 删除了4组有问题的数据(HCC 017、008、025、009)
  • 缺失值处理: 使用机器学习模型进行缺失值填补

评价指标

  1. Dice系数: Dice = 2×|X∩Y|/(|X|+|Y|)
  2. 平均IoU: MIoU = (1/N)∑|X_i∩Y_i|/|X_i∪Y_i|
  3. Hausdorff距离(HD95): 95百分位数的Hausdorff距离
  4. 准确率、召回率、精确率: 标准分类指标

对比方法

  • U-Net
  • UNETR
  • SwinUNETR
  • MedNeXt
  • 提出的Swin-NeXt及其变体

实现细节

  • 框架: 基于MONAI开发,兼容Jupyter Notebook
  • 预处理: 图像裁剪、归一化、重采样、数据增强
  • 训练策略: 滑动窗口策略,分块处理避免内存溢出

实验结果

主要结果

在肝脏肿瘤分割任务上的性能对比:

模型类别DiceMIoUHD95准确率召回率精确率
U-Net总体0.7090.61416.8470.9910.7700.704
SwinUNETR总体0.6560.5537.30.9830.7330.64
MedNeXt总体0.6830.58321.60.990.7520.674
Swin-NeXt总体0.6620.56914.6140.9920.7140.684

消融实验

特征融合方法对比:

融合方法模型DiceMIoUHD95
AddSwinUNETR0.6620.56226.956
ConcatSwinUNETR0.6410.53739.197
Cross AttentionSwinUNETR0.6660.56432.883
Cross AttentionMedNeXt0.6830.58926.428
Cross AttentionSwin-NeXt0.6570.56511.28

关键发现:

  1. 交叉注意力机制在所有模型上都能提供稳定的性能提升
  2. Add和Concat方法在使用转置卷积的模型(MedNeXt、Swin-NeXt)上无法收敛
  3. 交叉注意力机制与转置卷积结构兼容性更好

案例分析

通过患者HCC066的CT切片可视化结果显示,虽然Swin-NeXt在数值指标上未达到最优,但在分割形状和大小方面最接近真实标签,具有更好的临床实用价值。

相关工作

3D医学图像分割发展

  1. 传统方法: 区域生长等无监督方法
  2. 深度学习时代: 3D U-Net奠定基础
  3. 自动化框架: nnU-Net集成自动化预处理和参数选择
  4. Transformer应用: UNETR引入ViT,SwinUNETR使用Swin Transformer
  5. 现代卷积: MedNeXt基于ConvNeXt设计

本文贡献定位

本文是首个系统性结合Transformer编码器和现代卷积解码器的3D医学图像分割方法,并引入了多模态交叉注意力机制。

结论与讨论

主要结论

  1. 架构融合有效: Swin Transformer编码器与ConvNeXt解码器的结合是有效的
  2. 交叉注意力优越: 在多模态融合中,交叉注意力机制显著优于简单的加法和拼接方法
  3. 组合损失函数有效: DiceFocalCELoss能够改善类别不平衡问题的训练收敛

局限性

  1. 参数设置保守: 为公平比较减少了参数量,未使用最优的1:1:3:1特征提取块比例
  2. 计算资源限制: 由于服务器排队等因素,未能完成所有预定实验
  3. 数据集规模: 相对较小的数据集可能限制了模型的泛化能力

未来方向

  1. 探索更优的网络架构比例设计
  2. 在更大规模数据集上验证方法的有效性
  3. 研究其他医学图像分割任务的适用性

深度评价

优点

  1. 创新性强: 首次系统性结合Swin Transformer和ConvNeXt的优势
  2. 多模态融合: 有效利用临床信息辅助图像分割
  3. 实验充分: 包含详细的消融实验和对比分析
  4. 实用价值高: 针对医学图像分割的实际需求设计

不足

  1. 性能提升有限: 相比基线方法,数值提升不够显著
  2. 实验条件不一致: 不同模型的训练时间和条件存在差异
  3. 理论分析不足: 缺乏对为什么交叉注意力更有效的深入理论分析
  4. 泛化性未知: 仅在肝脏肿瘤分割上验证,其他任务的表现未知

影响力

  1. 方法论贡献: 为3D医学图像分割提供了新的架构设计思路
  2. 多模态融合: 为医学图像分析中的多模态信息融合提供了有效方案
  3. 实用价值: 考虑临床信息的分割方法更符合实际医疗需求

适用场景

  1. 3D医学图像分割: 特别是需要结合临床信息的场景
  2. 多模态医学分析: 图像与结构化数据结合的任务
  3. 类别不平衡分割: 背景占主导的分割任务

参考文献

论文引用了该领域的重要工作,包括:

  • 3D U-Net: 3D医学图像分割的奠基工作
  • nnU-Net: 自动化医学图像分割框架
  • SwinUNETR: Swin Transformer在医学图像分割中的应用
  • MedNeXt: 基于ConvNeXt的医学图像分割方法

总体评价: 这是一篇在3D医学图像分割领域具有一定创新性的工作,特别是在多模态融合和架构设计方面。虽然性能提升相对有限,但方法的实用性和对领域的贡献值得认可。