[论文解读] UNesT: Local Spatial Representation Learning with Hierarchical Transformer for Efficient Medical Segmentation

Xin Yu, Qi Yang+13 more|arXiv (Cornell University)|Sep 28, 2022

Advanced Neural Network Applications被引用 4

一句话总结

UNesT 提出了一种基于分层3D Transformer的医学图像分割模型，通过分块特征聚合增强局部空间表征，在全脑（133类）和肾亚结构分割任务上实现了最先进性能。其在Colin数据集上实现了0.8505的平均Dice相似系数（DSC），在CANDI数据集上实现了0.7906的平均DSC，优于先前方法（包括27个网络集成模型），同时在小样本数据集上展现出高数据效率和鲁棒性。

ABSTRACT

Transformer-based models, capable of learning better global dependencies, have recently demonstrated exceptional representation learning capabilities in computer vision and medical image analysis. Transformer reformats the image into separate patches and realizes global communication via the self-attention mechanism. However, positional information between patches is hard to preserve in such 1D sequences, and loss of it can lead to sub-optimal performance when dealing with large amounts of heterogeneous tissues of various sizes in 3D medical image segmentation. Additionally, current methods are not robust and efficient for heavy-duty medical segmentation tasks such as predicting a large number of tissue classes or modeling globally inter-connected tissue structures. To address such challenges and inspired by the nested hierarchical structures in vision transformer, we proposed a novel 3D medical image segmentation method (UNesT), employing a simplified and faster-converging transformer encoder design that achieves local communication among spatially adjacent patch sequences by aggregating them hierarchically. We extensively validate our method on multiple challenging datasets, consisting of multiple modalities, anatomies, and a wide range of tissue classes, including 133 structures in the brain, 14 organs in the abdomen, 4 hierarchical components in the kidneys, inter-connected kidney tumors and brain tumors. We show that UNesT consistently achieves state-of-the-art performance and evaluate its generalizability and data efficiency. Particularly, the model achieves whole brain segmentation task complete ROI with 133 tissue classes in a single network, outperforming prior state-of-the-art method SLANT27 ensembled with 27 networks.

研究动机与目标

• 解决视觉Transformer在保留局部空间信息以及处理3D医学影像中大规模、异质性组织结构方面的局限性。
• 提升在医学影像中常见低数据场景下的数据效率和模型鲁棒性。
• 实现对数百种组织类别及层级连接结构（如脑部和肾脏亚结构）的精确分割。
• 开发一种单模型解决方案，其性能和效率超越集成模型。
• 创建一个由放射科医生标注的新型院内肾亚结构CT数据集，用于验证和泛化能力评估。

提出的方法

• UNesT采用分层3D块聚合模块，将空间相邻的图像块分组为块，实现在分层级别上局部通信的同时保持全局自注意力机制。
• 模型采用类似U-Net的架构，包含嵌套的Transformer编码器和卷积解码器，通过多分辨率的跳跃连接相连。
• 通过分层堆叠Transformer编码器增强局部空间表征，在保持全局上下文的同时提升对局部结构的归纳偏差。
• 该方法整合了一种简化、收敛更快的Transformer设计，在保持高分辨率3D体数据性能的同时降低计算成本。
• 编码器与解码器之间的跳跃连接实现多尺度特征融合，提升对小尺寸或复杂结构的定位精度。
• 该框架在多模态数据集（包括T1加权MRI和CT）上端到端训练，结合大量数据增强和正则化技术。

实验结果

研究问题

• RQ1 分层Transformer架构是否能有效学习3D医学图像中的局部空间表征，同时保留全局上下文？
• RQ2 与现有模型（包括集成方法）相比，UNesT在133类全脑分割任务上是否能实现更优性能？
• RQ3 在低数据场景下，UNesT对罕见或复杂亚结构（如肾组件）的性能如何？
• RQ4 该模型在不同解剖区域、模态和组织复杂度水平上的泛化能力如何？
• RQ5 所提出的方法是否能在无需大量人工标注的情况下，实现与放射科医生相当的肾亚结构分割可重复性？

主要发现

• 在使用Colin数据集的全脑分割任务中，UNesT实现了0.8505的平均DSC，显著优于先前最先进水平的0.7444。
• 在CANDI数据集中，UNesT将平均DSC从0.6968提升至0.7025，表明其在具有挑战性的公开基准上性能增强。
• UNesT优于SLANT27（一个27个网络集成模型），以单一模型实现更高精度，降低推理成本与复杂度。
• 在肾亚结构分割任务中，UNesT实现了0.7906的DSC，表明其在复杂且层级连接的肾脏组件上表现优异。
• 该模型表现出高可重复性，其体积测量的平均差异为0.03，优于放射科医生之间的组间一致性（0.29）。
• 在异常值分析中，90%的分割结果可靠，10%存在轻微过分割或欠分割，可通过视觉检查识别，并可能通过后处理修正。