Neuroverse3D：首个用于3D神经影像的上下文学习通用模型

在计算神经影像分析领域，深度学习模型正朝着处理多领域、多任务的“通用模型”方向发展。其中，上下文学习（In-Context Learning, ICL）作为一种新兴范式，展现出了卓越的潜力，它允许模型仅通过提供任务相关的图像-标签对作为上下文指导，即可适应新任务和领域，而无需重新训练。本研究提出了首个专为3D神经影像设计的ICL通用模型（Neuroverse3D），其核心创新在于成功解决了3D ICL的内存瓶颈，使模型能够充分利用3D数据的全局信息，同时在多种分割和生成任务上实现卓越性能。

医学影像分析长期受领域偏移（Domain Shift）的困扰，即模型在不同机构、不同扫描设备或参数采集的数据上性能下降。解决方案主要包括领域自适应（需要目标域数据微调）和领域泛化（直接泛化到新领域）。然而，通用模型，尤其是ICL模型，提供了一条更具前景的路径。它们通过在大量多领域数据集上进行训练，学习通用的图像表征，并能通过上下文提示快速适应。

现有的通用模型根据提示类型可分为三类：

1. 符号提示模型：如使用点、边界框的SAM及其3D变体；

2. 自然语言提示模型：利用文本描述指导模型；

3. 上下文学习（ICL）模型：使用图像-标签对作为提示，为模型提供最丰富、最直接的任务指导；

在自然图像处理中，LVM、Painter和SegGPT等ICL模型已展现出强大的多功能性。在神经影像领域，UniverSeg、Neuralizer和One-Prompt等模型也证明了ICL在应对领域偏移方面的有效性。然而，这些模型均局限于2D输入，无法原生处理3D神经影像数据，从而丧失了3D全局信息带来的优势。本研究旨在开发一个能够有效处理3D上下文的ICL模型，在支持低内存占用的同时，提供解剖分割、肿瘤分割、血管分割、图像生成（包括偏置场校正、图像修复、超分辨率、高斯噪声去除、椒盐噪声去除、2d到3d的重建、模态变换）和颅骨去除等功能。

研究团队收集了来自19个数据集、共计43,674个3D多模态扫描，用于模型训练和评估。评估在4个留出数据集上进行，涵盖了14项不同的任务，包括解剖结构分割、肿瘤分割、血管分割以及偏置场校正、去噪、修复、超分辨率等多种生成任务。

Neuroverse3D的设计旨在以低内存开销，实现对3D神经影像数据的有效上下文学习。其核心方法包括三个关键部分：自适应并行-顺序处理（APSP）、U形融合策略以及优化的损失函数。

Neuroverse3D采用双分支3D U-Net架构。一个分支（目标分支）处理需要预测的目标图像，另一个分支（上下文分支）处理提供的上下文图像-标签对：

• 上下文分区：给定一个包含L对上下文的大集合，APSP将其划分为n个较小的、不相交的迷你上下文子集，每个子集包含对图像-标签对。的大小可根据可用内存灵活调整，较大值会减少推理时间但会增加内存使用；

• 并行与顺序处理结合：在每个迷你上下文内部，所有图像-标签对被并行处理（通过共享权重的上下文分支U-Net），并计算该迷你上下文的平均特征表示。然后，模型按顺序处理每一个迷你上下文，并通过一种顺序均值更新算法，逐步融合所有迷你上下文的平均表示，最终得到整个上下文集的全局平均表示。此表示随后输入到目标分支的解码器，生成最终预测；

• 内存优化与梯度等价性：在处理完一个迷你上下文并更新全局表示后，其中间特征会被释放，仅保留最后一个迷你上下文的特征用于训练时的梯度计算。为了确保这种近似处理与处理全部上下文具有相同的梯度期望，论文在梯度计算时对最后一个迷你上下文的表示进行了缩放。这使得内存消耗仅由迷你上下文的大小决定，而非整个上下文集的大小L；

为了实现目标图像与上下文信息的高效交互，模型采用了U形融合策略，它允许模型在顺序处理迷你上下文时，无需存储之前已处理上下文的特征，从而保证了计算效率：

• 双向信息流：信息流动分为两个阶段。首先，在编码器阶段，通过“目标到上下文融合”模块，将目标分支的特征注入到上下文分支，帮助上下文分支更好地理解当前目标图像的结构。然后，在解码器阶段，通过“上下文到目标融合”模块，将APSP计算得到的平均上下文表征注入回目标分支，指导目标图像的重建或分割；

• 融合操作：通过卷积层处理来自两个分支特征的拼接结果，并与原始特征残差连接；

为了应对3D数据中因维度增加而加剧的类别不平衡问题（如海马体等小结构的分割），并增强模型对复杂解剖边界的关注，该研究提出了优化的损失函数（总损失是各任务损失的加权平均，通过任务采样率进行平衡）：

• 分割任务：采用改进的平滑L1损失，通过使用更高的指数项，使模型更关注于具有挑战性的分割区域，而不是简单区域，从而在多样任务间取得更好平衡。该研究有意避开了Dice损失或Focal损失，以简化分割与生成任务间的平衡；

• 生成任务：采用标准平滑L1损失，并额外增加了一个基于图像强度梯度差的损失项，有助于模型在生成任务中更好地保留脑图像的复杂解剖边界细节；

为进一步提升模型的泛化能力并模拟真实世界中的数据变异，训练中应用了广泛的数据增强（随机仿射变换、翻转、旋转等）和任务增强技术。模型在8块NVIDIA V100 GPU上进行训练，每块GPU的批次大小（batch size）为1。优化器选用ADAM，初始学习率设置为1×10⁻⁴，整个训练过程持续120K steps。验证损失每1.2K步评估一次，如果连续20次评估未见改善，学习率将减半。总训练时间约为8天，最终选择在验证集上损失最低的模型作为最终模型。

在上下文大小为8的设置下，Neuroverse3D与现有的ICL模型（如UniverSeg, Neuralizer, Painter, SegGPT）以及任务专用的3D U-Net模型（分少样本和全监督两种设置）进行了全面比较（分割采用Dice评估，其他任务采用非PSNR评估）：

• 分割任务：Neuroverse3D显著优于所有其他ICL模型，在海马体、丘脑等目标上的Dice分数提升超过20个百分点。其性能也全面超越了少样本任务专用模型，并与全监督U-Net模型的性能非常接近。这证明了在数据有限的医疗中心，采用Neuroverse3D这类ICL模型是更经济且准确的方案；

• 生成任务：Neuroverse3D同样优于其他ICL模型，在椒盐噪声去除和颅骨去除等任务上，性能接近全监督水平；

定性结果清晰地展示了Neuroverse3D的3D处理优势。2D ICL模型由于对轴向切片进行独立处理，在矢状面或冠状面上经常出现不连续和错误。例如，在分割任务中，2D模型在矢状面视图下表现出较低准确性和较多假阳性/阴性。而在生成任务中，Neuralizer会产生过多背景误差，Painter则难以处理如偏置场校正等医疗专用任务。相比之下，Neuroverse3D在所有任务和不同视图下均能保持高精度和一致性。

随着上下文数量的增加，Neuroverse3D和对比模型的性能均有所提升，这印证了上下文数量对ICL模型的重要性。Neuroverse3D在所有上下文大小下均保持领先。性能提升在上下文大小达到约16后趋于平缓，但预测结果的稳定性（标准差）持续改善。APSP机制使得Neuroverse3D在有限资源下能够支持更大甚至无限的上下文大小，而SegGPT等模型则因内存限制，在实验中最多只能支持16个上下文对。

此外，研究者还训练了一组在特定任务上未经过训练的Neuroverse3D-unseen模型，以评估其跨任务泛化能力。如图所示，对于分割任务，Neuroverse3D-unseen仅比Neuroverse3D略有下降（Dice系数约低5点），但仍显著优于其他ICL模型。对于生成任务，性能下降更为明显。研究者认为，这是因为训练集中包含了更多样化的分割任务（超过30种脑结构及其随机组合），使分割学习更接近元学习，而生成任务的多样性相对较少，模型更容易记忆特定任务。这提示未来增加任务的多样性是提升ICL模型跨任务泛化能力的关键方向。

当迷你上下文大小时，模型以纯顺序模式运行，内存消耗最低（7.35GB），但计算时间较长。增大会提高内存占用，但能减少计算时间。关键在于，所有设置下模型都能产生一致的结果，这为模型在不同资源约束下的部署提供了极大的灵活性。

此外，尽管Neuroverse3D处理的3D上下文信息量更大（8对3D图像），但其推理速度却快于需要逐片处理2D上下文的对比模型，因为它避免了为每个切片重复计算上下文特征。

消融研究验证了所提出损失函数的有效性。移除改进的平滑L1损失会导致分割性能下降，尤其是在海马体等小而复杂的区域。移除梯度损失则显著降低了生成任务的性能。研究发现，即使梯度损失未直接用于分割任务，其存在也能提升分割性能，可能是因为它增强了模型对边界的敏感性。

本文提出的Neuroverse3D成功解决了在3D神经影像中应用ICL范式所面临的核心挑战——高内存消耗。通过创新的自适应并行-顺序处理（APSP）和U形融合策略，模型能够高效处理大规模3D上下文。优化的损失函数则有效平衡了多任务学习并增强了对解剖细节的关注。在大规模多中心数据上的实验证明，Neuroverse3D在多种3D神经影像任务上显著超越了现有的ICL模型，其分割性能可与全监督模型相媲美，且无需针对新中心或新任务进行重新训练，展现出强大的临床应用潜力。

汇智灵曦推出的训练推理一体化平台，集成了多种经典AI模型，内置AI模型商城，为用户提供高效便捷的训练与推理解决方案。用户只需设置输入输出路径并调整参数，即可快速完成模型训练与推理，无论是科研探索还是临床应用，都能轻松上手。平台具备简洁的界面和操作流程，免去复杂的技术要求，帮助用户节省时间、提高效率，是实现AI技术快速落地的理想选择。

点击下方“阅读原文”

立即申请“训练推理一体化平台”试用