点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
医学图像分割(MIS)涵盖了从骨骼到器官分割等多样化任务,每个任务在寻找最佳分割模型时都面临独特的挑战。当前最先进的与AutoML相关的MIS框架nnU-Net自动化了模型配置的许多方面,但仍受限于固定的超参数和启发式设计选择。作为MIS的完整AutoML框架,作者提出了Auto-nnU-Net,这是一种新型nnU-Net变体,支持超参数优化(HPO)、神经架构搜索(NAS)和分层NAS(HNAS)。
此外,作者提出了正则化先验带(Regularized PriorBand),以平衡模型精度与训练所需的计算资源,解决现实世界医疗环境中常见的资源限制问题,从而限制大规模训练程序的可行性。
作者在医学分割十项全能赛(Medical Segmentation Decathlon)提供的多样化MIS数据集上评估了yinqing-Auto-nnU-Net_2505,分析了AutoML技术对分割性能、计算效率和模型设计选择的影响。结果表明,作者的AutoML方法在10个数据集中的6个上显著提升了nnUNet的分割性能,在其他数据集上表现相当,同时保持了实用的资源需求。
作者的代码可在https: //github. com/automl/AutoNNUnet获取。
1 引言
机器学习(ML)在现代医疗领域发挥着关键作用,能够实现精准诊断、早期癌症检测(Cao等人,2023)以及科学发现(Falk等人,2018)。医学图像分割(MIS)旨在识别医学扫描中的解剖结构,但由于数据集差异、类别不平衡和任务特定约束(Litjens等人,2017;Isensee等人,2020a;Ali等人,2024),这一任务具有挑战性。自配置方法通过针对特定数据集调整模型来减少人工调优的需求(Ali等人,2024)。nnU-Net(Isensee等人,2020a)已成为一种先进的框架,能够自动配置基于U-Net的架构以实现强大的分割性能。然而,nnU-Net出人意料地依赖于一些固定的超参数和人工设计的启发式方法,这限制了其灵活性,且在不同数据集上可能无法始终获得最优结果(Bergstra等人,2012;Quinton等人,2024)。
在本工作中,作者利用自动化机器学习(AutoML)(Hutter et al., 2019) 来应对这些挑战,并对 AutoML 对医学图像分割(MIS)的影响进行大规模研究。作者引入了 Auto-nnU-Net,这是一种集成了 AutoML 的 nnU-Net 新变体,能够实现超参数优化(HPO)和神经架构搜索(NAS)。通过将 PriorBand (Mallik et al., 2023) 与多目标优化 (Karl et al., 2023) 相结合,作者提出了正则化 PriorBand 用于联合架构和超参数搜索(JAHS)(Awad et al., 2023),以应对近期研究(Rayed et al., 2024)中强调的 MIS 资源效率日益增长的关注。作者的研究在医学分割十项全能(MSD)数据集(Simpson et al., 2019; Antonelli et al., 2022)上评估了 Auto-nnU-Net,为优化策略、超参数重要性和数据集特征的影响提供了见解。值得注意的是,与大多数关于 MIS 的 AutoML 研究(Ali et al., 2024)不同,作者报告了所有十个 MSD 数据集的结果,从而对跨不同医学分割挑战的泛化性和鲁棒性进行了更全面的评估。
在这项工作中,作者做出了以下贡献:
- Auto-nnU-Net用于AutoML驱动的MIS。作者提出了一种新型框架,该框架自动化了nnU-Net中的关键设计决策,以实现灵活且结构化的超参数优化(HPO)和神经网络架构搜索(NAS)。
- 基于正则化先验带宽的高效优化。作者引入了正则化先验带宽,将训练时间作为优化目标,以反映现实世界的约束条件,其中有限的资源和频繁的再训练使得高效训练至关重要。该方法仅在模型能够提升准确率时选择较慢的模型,并内在地提供权衡解决方案。
- 在全部十个MSD数据集上进行广泛评估。作者分析了AutoML对分割精度的影响,包括超参数重要性和数据集可迁移性,从而更深入地理解泛化行为,并指导设计更鲁棒、高效的模型,以应对多样化的医学图像任务。
2 图像分割背景
参考Szeliski(2022)的定义,语义分割是指将图像划分为与特定类别相关的区域。作者使用图像分割这一术语来描述该任务,其中每个像素被 Token 以实现对视觉数据的结构化分析。
医学图像分割(MIS)涉及对医学图像,例如磁共振成像(MRI)或计算机断层扫描(CT)扫描进行划分,以识别感兴趣区域,包括器官或潜在的恶性结构如肿瘤(Antonelli等人,2022)。在实际应用中,自动分割通过准确识别患者治疗的关键区域来辅助临床医生(Liang等人,2019)。最近的MIS数据集主要关注前景类,将背景视为一个单独的排除类(Menze等人,2015;Heller等人,2019;Simpson等人,2019;Antonelli等人,2022)。与自然图像分割不同,MIS面临诸如训练数据有限、类别不平衡、小或分支的解剖结构、边界模糊以及强度分布变化等挑战。来自MRI和CT扫描的由 Voxel (体积像素)组成的三维图像的分割进一步增加了MIS的分割复杂性和计算需求(Isensee等人,2020a;Ali等人,2024)。
3 相关工作
在本节中,作者回顾了Auto-nnU-Net核心组件的相关研究:自配置分割框架、超参数优化(HPO)、神经网络架构搜索(NAS)以及面向多目标优化的模型集成(MIS)。作者的工作旨在将这些技术统一到一个针对当前任务定制的自动化机器学习(AutoML)框架中。
nnU-Net。自配置框架解决了为特定任务和数据集设计和调整MIS模型的设计挑战。基于U-Net的成功(Ronneberger等人,2015年),Isensee等人(2020a)引入了nnU-Net,该框架优化了针对特定任务的U-Net。与CASH(Thornton等人,2013年)类似,它联合选择训练超参数和用于推理的最终模型或集成模型。nnU-Net通过利用AutoML中常见的元数据特征来实现这一点(Vanschoren,2019年)。作者关注nnU-Net的自配置机制,省略了预处理和后处理步骤。该流程包括三个阶段:(i)实验规划,其中基于规则的超参数选择利用数据集属性;(ii)训练,其中使用5折交叉验证训练2D U-Net、3D U-Net,如有必要,则训练3D U-Net级联;(iii)推理,根据验证分数选择表现最佳的模型或集成模型。nnU-Net依赖于三种类型的超参数:(i)固定(例如,学习率、优化器、损失函数);(ii)基于规则(例如,预处理、网络拓扑);(iii)启发式(例如,集成选择、后处理)。
MIS的超参数优化(HPO)。一般来说,针对MIS的HPO研究较少。Yang等人(2019)提出使用强化学习来优化3D分割模型的数据增强和学习率。Quinton等人(2024)将HPO应用于不同模型,包括nnUNet,通过依次对三组超参数进行贝叶斯优化(BO):(i) 图块大小,(ii) 数据预处理和增强,以及(iii) 损失函数和优化器。
面向医学图像分割的神经网络架构搜索(NAS)。多种NAS方法已被应用于医学图像分割。一些方法基于DARTS(H. Liu等人,2019)并应用于基于编码器-解码器的医学图像分割模型(Weng等人,2019;Zhu等人,2019;Y. He等人,2021)。另一种方法采用粗粒度到细粒度的策略用于U-Net结构的网络,首先优化整体拓扑结构,然后细化单元级操作(Q. Yu等人,2020)。针对医学图像分割,也已提出了基于进化和基于图的NAS方法,使用遗传算法(Hassanzadeh等人,2020;Khouy等人,2023;C. Yu等人,2023)以及架构的图表示,这些架构在训练过程中被优化或扩展以减少搜索时间并提高灵活性(R. Liu等人,2023;Qin等人,2023)。
多目标优化与MIS的联合超参数优化和神经网络架构搜索。已有工作将多目标神经网络架构搜索应用于MIS,以平衡性能和资源约束(Baldeon-Calisto等人,2020;Lu等人,2022),但未进行超参数调整。Yang等人(2021)通过 Agent 模型结合超参数优化和神经网络架构搜索来优化U-Net配置,但未考虑资源效率。
4 Auto-nnU-Net用于MIS
在这项工作中,作者提出了Auto-nnU-Net,这是一种将AutoML方法集成到nnU-Net中的新方法。此外,作者引入了Regularized PriorBand,以在Auto-nnU-Net中实现高效的联合架构和超参数搜索(JAHS)(Awad等人,2023)。
4.1 将AutoML方法集成到nnU-Net中
nnU-Net提供了稳健的分割流程,包括数据预处理、实验规划、训练和推理。然而,其固定且基于规则的超参数限制了可配置性。为解决这些限制,作者提出了Auto-nnU-Net,通过灵活的实验规划和训练增强nnU-Net。图1展示了yinqing-Auto-nnU-Net_2505概述。与nnU-Net不同,Auto-nnU-Net将超参数和架构配置作为输入,实现JAHS。它返回泛化误差和训练运行时间,使优化过程能够同时考虑分割性能和计算效率。
4.2 正则化先验Band用于高效的联合超参数优化和神经架构搜索
基于灵活的Auto-nnU-Net框架,作者进一步利用Regularized PriorBand增强优化过程。在本节中,作者描述了如何将PriorBand(Mallik等人,2023)从超参数优化(HPO)扩展到联合自适应超参数优化(JAHS)。鉴于nnU-Net需要较高的训练成本并提供强大的先验配置,作者旨在将这种知识融入优化过程以提高其效率。为此,作者采用了PriorBand(Mallik等人,2023),这是一种多保真度超参数优化方法,专门设计用于将先验知识整合到计算成本高昂的深度学习模型的优化中。它通过结合随机采样、基于先验的采样和基于当前最优配置的采样策略来增强探索,并根据优化进程动态调整。随机采样探索搜索空间,基于先验的采样利用专家知识,而基于当前最优配置的采样则优化当前表现最佳的配置。
作者通过将架构编码在统一的配置空间(Zela等人,2018)中扩展了PriorBand的超参数搜索空间。然而,探索更大规模的模型会带来计算挑战。虽然增加模型规模可以提高准确性,但它也会提高优化成本并延长训练时间。作者考虑将训练运行时间作为优化目标,以更好地反映医疗环境中的实际约束,在这些环境中计算资源通常有限,大规模训练可能不可行(Rayed等人,2024)。数据集异质性——由于技术因素(例如,扫描仪、协议)和解剖学变异(例如,器官形状、结构数量)——通常需要重复微调或模型适配。在这样的持续学习场景中,重新训练是反复且昂贵的(Isensee等人,2020a;Wagner等人,2024),因此高效训练至关重要。隐私约束通常阻止对患者数据的集中访问,当新数据可用时需要局部或联邦重新训练,这进一步强调了最小化训练成本的重要性(Wagner等人,2024)。
Regularized PriorBand的核心思想是只有当大模型能够提升准确率时,才应考虑使用它们。最终目标仍然是优化准确率,确保性能最佳的配置不会被资源受限的选择所取代。Regularized PriorBand的概述见附录C中的算法2。
连续减半中的选择策略。Regularized PriorBand的一个关键改进在于修改了连续减半(Successive Halving, SH)子程序中的配置选择策略(Jamieson等人,2016)。在标准的SH方法中,下一更高预算的配置是根据其成本进行选择的,其中成本最低的配置被优先用于评估。然而,在同时优化准确性和训练运行时间时,作者必须考虑一个成本向量而非单个标量,以同时考虑这两个目标。
为了同时兼顾准确性和运行时间,作者通过采用Deb等人(2002年)提出的非支配排序和拥挤距离排序,对选择过程进行了改进,这与近期研究(Izquierdo等人,2021年;Schmucker等人,2021年;Awad等人,2023年)的方法类似。在当前预算下评估配置后,作者应用非支配排序将配置分组到不同的前沿。为了促进解的多样性,作者在每个前沿内根据拥挤距离对配置进行排序。从这些排序好的前沿中,作者选择排名前
的配置,在下一个更高的预算下进行评估,从第一个前沿开始,直到选出
个配置。如果两个配置的拥挤距离相同,选择过程将优先考虑准确性。这确保了具有最高准确性的配置总是被优先推进。
当前解选择。在正则化先验Band中,最终当前解配置的选择基于准确率,不考虑运行时间。然而,为了在整个优化过程中实现基于当前解的采样,作者结合了准确率和运行时间。选择仅限于近似Pareto前沿上的配置,确保在两个目标之间取得平衡。为了选择用于局部搜索的当前解,作者计算归一化目标成本的覆盖区域,并选择最大化该区域的配置,从而促进准确率和运行时间之间权衡的探索。
5 实验设置
基于作者的Auto-nnU-Net框架,作者开展了迄今为止最全面的MIS领域AutoML研究,相当于约60千GPU小时和10^964千克CO2当量(详见附录A)。所有实验、结果和可视化的复现说明可在作者的GitHub仓库https://github.com/automl/AutoNNUnet中找到。所有实验均采用5折交叉验证进行。更多详情请参见附录D.3。
5.1 数据集
为确保对方法的全面评估,作者使用了医学分割十项全能赛(MSD)(Simpson等人,2019年;Antonelli等人,2022年),这是一个包含十个MIS数据集的基准,旨在捕捉临床任务、成像模式和数据特征方面的多样性(参见附录D.1)。MSD使用Dice相似性系数(DSC)(Dice,1945年),这是一个有效的评估指标,用于评估Mid neo
5.2 Baseline 模型
在作者的实验中,作者旨在研究AutoML方法如何提升当前MIS方法的分割性能。作者的首个 Baseline 是nnU-Net框架中的3D U-Net,特别是其默认配置(ii)Conv,(ii)ResM和(ii)ResL。此外,作者还评估了MedSAM2(Ma等人,2024a),这是一种基于基础模型的MIS方法。与nnU-Net不同,MedSAM2利用大规模预训练,是yinqing-Auto-nnU-Net_2505最先进的竞争者。作者采用作者提出的流程,在每一个单独的MSD数据集上对MedSAM2进行100个epoch的微调,这大致相当于在训练运行时间最长的nnU-Net配置D01数据集上的训练运行时间。
5.3 Auto-nnU-Net评估
本节概述了评估作者Auto-nnU-Net方法的实验设置。Auto-nnUNet使用Regularized PriorBand将Training Runtime目标与
一同纳入优化,以在同等性能下优先选择更高效的模型。对于PriorBand优化器,作者采用Mallik等人(2023年)提出的设置(参见附录D.2.2),仅使用一个随机种子,因为否则需要大量的计算资源。
Auto-nnU-Net搜索空间包括常规超参数,这些参数定义了训练和配置设置(例如学习率和数据增强),以及控制网络结构的架构超参数(例如编码器类型和dropout率)。这种JAHS搜索空间公式化(Bansal等人,2022年)能够同时调整训练动态和模型容量。完整的Auto-nnU-Net搜索空间如表4所示,以及作为HPO和NAS空间的组合。超参数的详细信息在附录D.2.3中说明。
5.4 消融实验
为评估Auto-nnU-Net框架中不同组件的贡献,作者定义了两个消融变体,它们隔离或修改了JAHS搜索空间的部分内容:
使用PriorBand进行超参数优化。在此变体中,作者禁用Auto-nnU-Net的架构搜索,并仅优化常规超参数,使用PriorBand(Mallik等人,2023年),而不添加来自Auto-nnU-Net的正则化,以最小化
。通过排除架构超参数,这种消融实验隔离了调整配置选择的影响,并有助于量化当网络架构固定为nnU-Net默认值时,仅归因于超参数优化的性能提升。在表4(顶部),定义了每个超参数的可能值范围和集合。
用于U-Net的分层神经网络架构搜索。虽然Regularized PriorBand使得JAHS成为可能,如同在Auto-nnU-Net中,但其搜索空间仅限于预定义的修改。为了探索更广泛的U-Net架构,同时保持效率和可行性,作者引入了分层神经网络架构搜索(HNAS)搜索空间,利用上下文无关文法(CFG)(Schrodi等人,2023)系统性地优化和细化U-Net结构(参见附录C.3.1)。yinqing-Auto-nnU-Net_2505保留了nnU-Net的默认配置,同时引入了灵活的拓扑和单元 Level 的设计选择。作者应用基于先验的采样方法,使用Regularized PriorBand,将CFG生成规则建模为分类和整数超参数,以平滑地与现有的nnU-Net组件集成(参见附录C.3.2)。
6 结果
在本节中,作者通过MSD数据集展示了Auto-nnU-Net的有效性。作者评估了分割精度、效率和配置可迁移性,并与 Baseline 和消融实验进行了比较。为了更深入地理解其优化行为,作者进一步分析了各个超参数的重要性。
6.1 Auto-nnU-Net 结果
首先,作者讨论Auto-nnU-Net在MSD数据集上的优化进展。图2展示了Auto-nnU-Nets当前的1-DS C值随时间的变化,并与默认nnU-Net Baseline 进行了比较。除D08外(在该数据集上DSC值相同),Auto-nnU-Net的表现优于nnU-Nets卷积默认模型。值得注意的是,对于D04,Auto-nnU-Net比训练nnU-Net(ResL)更快地识别出当前最优配置,这突显了其在计算成本较高的模型上的效率优势。
最终验证结果,包括MedSAM2,在附录E的表6中列出。作者在成本随时间变化的比较中排除了MedSAM2,因为它在十个数据集中的九个上表现不如nnU-Net,仅在D10上表现更优。作者假设MedSAM2在D10上的表现是由于它需要从真实 Mask 中获取边界框 Prompt ,这有助于检测小目标区域。虽然这种方法有利于检测任务,但它需要额外的标注,而yinqing-Auto-nnU-Net_2505无需此类监督。
为了评估在未见过数据上的性能,作者评估了MSD测试集的结果。表1展示了所有方法的最终测试集DSC[%],不包括MedSAM2,因为MSD的ground truth分割结果不可用。与验证结果一致,Auto-nnUNet实现了最高的平均DSC(77.04%)。yinqing-Auto-nnU-Net_2505超越了所有nnU-Net Baseline ,并展现出强大的泛化能力,在十个数据集中排名最佳的有五个。
正则化先验Band在优化过程中天生平衡了准确性和训练运行时间(第4.2节)。图4比较了Auto-nnU-Net、其消融实验(HPO、HPO
HNAS)以及 Baseline 模型(nnU-Net、MedSAM2)在D03和D04上的Pareto前沿,展示了目标权衡。在D03上,Auto-nnU-Net和HPO
HNAS显示出明显的准确率-运行时间权衡,而nnU-Net(ResM)以低运行时间实现了高DSC。在D04上,Auto-nnU-Net和
在准确率上优于所有nnUNet变体,并显著减少了训练时间——HPO
HNAS将运行时间缩短了26倍。MedSAM2在这两个数据集上都表现不佳。这些结果表明Auto-nnU-Net能够联合优化准确性和效率。
消融实验结果。作者在图2和表1中讨论了Auto-nnU-Net的HPO和HPO+HNAS消融实验结果。HPO在除D08和D10以外的所有数据集上都优于nnU-Net(Conv),而Auto-nnU-Net显示出相似的改进,但在除D08以外的所有数据集上都表现优异。HPO+HNAS在六个数据集上优于nnU-Net(Conv),并且在某些数据集(如D01)上通常比Auto-nnU-Net和HPO优化得更加高效。然而,与Auto-nnU-Net相比,HPO和HPO+HNAS在测试集上的DSC值较低,这表明它们对未见数据的鲁棒性有所下降。对于D01,尽管验证DSC值相似,HPO+HNAS相对于Auto-nnU-Net表现较差,这表明其对未见数据的敏感性更高,以及将神经网络架构编码为超参数的有效性。
6.2 自动机器学习网络超参数重要性分析
为评估单个超参数对准确性的影响,作者采用功能ANOVA(fANOVA)(Hutter等人,2014年)通过将性能方差分解为每个超参数及其交互作用的贡献来估计它们在配置空间中的全局重要性。图14显示了它们在不同数据集上的重要性。像前景过采样、初始学习率和动量(SGD)这样的关键超参数始终具有影响力,尽管它们的重要性在不同数据集之间存在显著差异。这突出了AutoML相对于原始nnU-Net中使用的固定设置的价值。相比之下,像编码器类型和归一化这样的超参数重要性较低,可能不需要进一步优化。
6.3 跨数据集迁移现有配置
评估数据集对AutoML的影响的一个关键问题是针对一个数据集优化的配置是否能够泛化到其他数据集。作者分析了Auto-nnU-Nets的迁移性,排除了D08(因为它未优于默认配置),结果得到一个
矩阵(图13)。在半数数据集中,定制化的 incumbent 未达到最高的 DSC,其中D02配置在D05上显示出最大增益(+2.78%),而D03 incumbent在四个数据集上达到了最高的 DSC。然而,它在D05上的较低性能导致其平均表现未能超越nnU-Net默认配置。这些结果表明配置可以在数据集间迁移——例如D03在D03、D08-D10上表现最佳——但其他数据集如D04和D05泛化能力较差,尤其是在D06-D10上。这突出了元学习HPO在提高MIS迁移性方面的潜力(Feurer等,2015;Wistuba等,2015;Schilling等,2016)。
7 结论与未来工作
在这项工作中,作者提出了Auto-nnU-Net,一个结合nnU-Net与结构化HPO和NAS的医学图像分割自动化框架。通过集成Regularized PriorBand,作者联合优化了分割性能和训练运行时间,解决了医学环境中的实际约束。作者对全部十个Medical Segmentation Decathlon数据集进行了综合评估,结果表明Auto-nnU-Net始终优于或匹配强 Baseline ,同时保持了实用的资源需求。作者通过消融研究进一步分析了HPO和NAS的贡献,考察了优化配置在数据集间的可迁移性,并评估了超参数的重要性。这些见解有助于更深入地理解作者的分割方法在不同临床环境中的设计和优化。总体而言,Auto-nnU-Net为自动化医学图像分割提供了一个灵活且资源感知的基础,能够在现实世界约束下实现鲁棒的模型设计。
局限性。本研究包括nnU-Net和Auto-nnU-Net的结果,基于3D U-Net架构且不进行后处理或集成,可能无法完全反映原始nnU-Net的性能(Isensee等人,2020a)。然而,引入集成——在AutoML中常见(Erickson等人,2020)——可能会提升Auto-nnU-Net的结果。关于作者的帕累托分析,低预算配置近似全预算性能,并揭示了运行时与准确性的权衡。最后,虽然DeepCAVE中的 Agent 模型可能引入轻微的近似误差,但研究结果为推进MIS中的AutoML提供了坚实基础。
未来研究。未来的工作可以将评估扩展到完整的nnU-Net流程,并进一步研究数据集属性如何影响AutoML结果。使用多个默认配置(Pfisterer等人,2018年)和元学习(Feurer等人,2015年;Vanschoren,2019年;Aguiar等人,2019年)来启动AutoML可以提高MIS的AutoML效率。最后,使用预训练模型的零样本AutoML(Ozturk等人,2022年)可以增强适应性,同时降低成本。
参考
[1]. Auto-nnU-Net: Towards Automated Medical Image Segmentation
点击上方卡片,关注「AI视界引擎」公众号