参数量从 1400 万到 14 亿的可扩展 U-Net 系列，基于 TotalSegmentator 大规模预训练，面向通用医学图像分割
大规模预训练模型已深刻改变了自然语言处理和计算机视觉领域。然而，医学图像分割领域的主流模型参数量仍停留在数千万量级，规模扩展潜力尚未得到充分探索。能否将模型规模提升一个数量级，并验证更大模型在临床任务间具有更强的迁移能力，是 STU-Net 工作探索的核心问题。
我们设计了一系列可扩展可迁移 U-Net（STU-Net）模型，参数量从 1460 万（STU-Net-S）到 14.57 亿（STU-Net-H）不等。STU-Net-H 是迄今为止最大的医学图像分割模型。所有变体均基于 nnU-Net 框架，并进行了关键架构改进：引入残差连接以支持深度扩展，采用无权重插值上采样以消除跨任务迁移时的权重不匹配问题。
在包含 104 个解剖结构标注的大规模数据集 TotalSegmentator 上预训练后，STU-Net 验证了规模扩展能够持续提升分割精度。在 TotalSegmentator 基准上，STU-Net-H 达到 90.06% 平均 DSC，超越所有 CNN 和 Transformer 竞争方法。其迁移能力延伸至 14 个下游数据集的直接推理及 3 个数据集的微调任务，覆盖 CT、MRI、PET 等多种模态和多样化分割目标。
STU-Net 提供四种规格——S（1460 万）、B（5826 万）、L（4.40 亿）和 H（14.57 亿参数）。联合扩展网络深度与宽度的策略优于单独扩展任一维度。STU-Net-B 已在 TotalSegmentator 上超越 nnU-Net 0.36%、超越 SwinUNETR-B 4.48%。STU-Net-H 以 90.06% 平均 DSC 创下该基准的最优纪录。
架构改进使规模扩展成为可能：各模块中的残差连接有效缓解了超深网络的梯度消散问题；固定 6 阶段、各向同性卷积核的设计确保预训练权重可跨任务复用而无需担心形状不匹配。
在 TotalSegmentator 上预训练后，STU-Net 可直接在 14 个包含预训练类别子集的下游 CT 数据集上进行推理，无需额外训练。在这 14 个数据集（共 2,494 个样本）上，STU-Net-H 以 84.02% 平均 DSC 超越 nnU-Net 的 76.37%，提升幅度达 7.65%。
在 FLARE22、AMOS22（CT+MRI）和 AutoPET22（CT+PET）三个挑战性下游数据集上进行微调时，STU-Net-H-ft 达到 80.69% 平均 DSC，优于 nnU-Net 的 77.06%。值得注意的是，基于 CT 预训练后在 MRI、PET 等非 CT 模态上进行微调同样受益显著，表明模型习得了超越模态特异性的基础解剖结构表征。
在拥有 104 个解剖结构标注的 TotalSegmentator 验证集上，STU-Net-H 以 90.06% 平均 DSC 刷新纪录，超越最优 CNN 模型 nnU-Net（86.76%）达 3.3 个百分点，超越最优 Transformer 模型 SwinUNETR-B（82.64%）达 7.4 个百分点。
性能提升在全部五个解剖子类上均保持一致，其中脊椎（86.97%→90.43%）和肋骨（86.11%→90.29%）提升最为显著，表明模型扩展带来的是全面覆盖能力的提升，而非仅体现在整体平均值上。
医学图像分割领域长期存在一种假设：针对单一类别训练的专科模型必然优于同时处理所有类别的通用模型。STU-Net 的实验结果对此提出了挑战。
我们分别训练了五个专科模型（器官、脊椎、心脏、肌肉、肋骨）并与单一通用 STU-Net 进行对比。在 STU-Net-H 规模（14 亿参数）下，通用模型以 90.06% 整体平均 DSC 超越所有专科模型中的最优组合（89.07%）。这表明在足够大的规模下，单一统一模型确实能够同时精通所有分割目标——迈向真正意义上的医学分割基础模型的关键一步。
四种 STU-Net 变体通过系统性地联合扩展编码器深度与通道宽度来定义：S（1460 万参数，12.8B FLOPs）、B（5826 万，60.9B）、L（4.40 亿，416B）和 H（14.57 亿，1,623B）。消融实验表明仅扩展深度或宽度的收益均不如两者均衡扩展。尽管 S 与 H 之间存在百倍参数差距，所有变体均共享完全相同的 6 阶段编解码拓扑和各向同性卷积核配置——正是这一设计约束实现了真正的权重可迁移性，无需任何形状适配器。
在三个挑战性微调基准——FLARE22（13 个腹部器官）、AMOS22（CT + MRI，15 个器官）和 AutoPET22（CT + PET 病灶分割）——上，从预训练权重微调的 STU-Net-H 在所有数据集上均超越从随机初始化微调的 nnU-Net。跨模态迁移结果尤为引人注目：仅在 CT 上预训练的 STU-Net-H，在 AMOS-MRI 和 AutoPET-PET 上微调后仍优于从头训练的 nnU-Net——表明预训练权重编码了超越模态特异性的基础解剖先验知识。
STU-Net 证明了在自然语言和计算机视觉领域发现的扩展律同样适用于三维医学图像分割。凭借 14 亿参数及在涵盖 CT、MRI、PET 模态的 17 个数据集上的强大迁移能力，STU-Net-H 代表了通用医学分割的当前前沿，是迈向医学人工通用智能（MedAGI）的重要基础模型构建模块。
Ziyan Huang, Haoyu Wang, Zhongying Deng, Jin Ye, Yanzhou Su, Hui Sun, Junjun He, Yun Gu, Lixu Gu, Shaoting Zhang, Yu Qiao
arXiv 2023