[Paper Review] 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 Applications4 citations

TL;DR

UNesT proposes a hierarchical 3D transformer-based medical image segmentation model that enhances local spatial representation through block-wise feature aggregation, achieving state-of-the-art performance on challenging whole-brain (133 classes) and renal substructure segmentation tasks. It outperforms prior methods, including a 27-network ensemble, by achieving a mean DSC of 0.8505 on the Colin dataset and 0.7906 on CANDI, while demonstrating high data efficiency and robustness on small datasets.

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.

Motivation & Objective

• To address the limitations of vision transformers in preserving local spatial information and handling large-scale, heterogeneous tissue structures in 3D medical imaging.
• To improve data efficiency and model robustness in low-data regimes common in medical imaging.
• To enable accurate segmentation of hundreds of tissue classes and hierarchically connected structures, such as brain and kidney substructures.
• To develop a single-model solution that surpasses ensembled models in performance and efficiency.
• To create a novel in-house renal substructure CT dataset annotated by radiologists for validation and generalization.

Proposed method

• UNesT employs a hierarchical 3D block aggregation module that groups spatially adjacent patches into blocks, enabling localized communication while preserving global self-attention across hierarchical levels.
• The model uses a U-Net-like architecture with a nested transformer encoder and a convolutional decoder, connected via skip connections at multiple resolutions.
• Local spatial representation is enhanced by stacking transformer encoders hierarchically, maintaining global context while improving inductive bias for local structures.
• The method integrates a simplified, faster-converging transformer design that reduces computational cost while maintaining performance on high-resolution 3D volumes.
• Skip connections between encoder and decoder allow multi-scale feature fusion, improving localization accuracy for small or complex structures.
• The framework is trained end-to-end on multi-modality datasets including T1-weighted MRI and CT, with extensive data augmentation and regularization.

Experimental results

Research questions

• RQ1 Can a hierarchical transformer architecture effectively learn local spatial representations while preserving global context in 3D medical images?
• RQ2 Can UNesT achieve superior performance on whole-brain segmentation with 133 tissue classes compared to existing models, including ensembled methods?
• RQ3 How does UNesT perform in low-data regimes, particularly for rare or complex substructures like renal components?
• RQ4 To what extent does the model generalize across different anatomical regions, modalities, and tissue complexity levels?
• RQ5 Can the proposed method achieve radiologist-level reproducibility in renal substructure segmentation without requiring extensive manual annotation?

Key findings

• UNesT achieves a mean DSC of 0.8505 on the whole-brain segmentation task using the Colin dataset, a significant improvement from the prior state-of-the-art of 0.7444.
• On the CANDI dataset, UNesT improves mean DSC from 0.6968 to 0.7025, demonstrating enhanced performance on a challenging, publicly available benchmark.
• The model outperforms SLANT27, a 27-network ensemble, by achieving higher accuracy with a single model, reducing inference cost and complexity.
• UNesT achieves a DSC of 0.7906 on the renal substructure segmentation task, indicating strong performance on complex, hierarchically connected kidney components.
• The model shows high reproducibility, with a mean difference of 0.03 in volumetric measurements compared to radiologists, outperforming inter-rater agreement (0.29).
• In outlier analysis, 90% of segmentations were reliable, while 10% showed under- or over-segmentation, which were detectable via visual inspection and potentially correctable with post-processing.