[Paper Review] Characterizing Renal Structures with 3D Block Aggregate Transformers

Xin Yu, Yucheng Tang+11 more|arXiv (Cornell University)|Mar 4, 2022

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes UNesT, a 3D hierarchical transformer model with 3D block aggregation to improve data efficiency and segmentation accuracy for renal substructures (cortex, medulla, pelvicalyceal system) in contrast-enhanced CT scans. The method achieves state-of-the-art performance with a Dice score of 0.8467 and high correlation (Pearson R = 0.9891) with manual segmentation, demonstrating strong reproducibility and generalizability on public datasets.

ABSTRACT

Efficiently quantifying renal structures can provide distinct spatial context and facilitate biomarker discovery for kidney morphology. However, the development and evaluation of the transformer model to segment the renal cortex, medulla, and collecting system remains challenging due to data inefficiency. Inspired by the hierarchical structures in vision transformer, we propose a novel method using a 3D block aggregation transformer for segmenting kidney components on contrast-enhanced CT scans. We construct the first cohort of renal substructures segmentation dataset with 116 subjects under institutional review board (IRB) approval. Our method yields the state-of-the-art performance (Dice of 0.8467) against the baseline approach of 0.8308 with the data-efficient design. The Pearson R achieves 0.9891 between the proposed method and manual standards and indicates the strong correlation and reproducibility for volumetric analysis. We extend the proposed method to the public KiTS dataset, the method leads to improved accuracy compared to transformer-based approaches. We show that the 3D block aggregation transformer can achieve local communication between sequence representations without modifying self-attention, and it can serve as an accurate and efficient quantification tool for characterizing renal structures.

Motivation & Objective

• To address data inefficiency in transformer-based models for small-structure segmentation in 3D medical imaging.
• To develop a hierarchical 3D transformer architecture that improves local representation learning without modifying self-attention.
• To create the first clinically annotated dataset of renal substructures (116 subjects) for accurate volumetric analysis.
• To enable reliable, automated quantification of renal cortex, medulla, and collecting system volumes for biomarker discovery.
• To validate the method’s generalizability and clinical utility on external datasets like KiTS19.

Proposed method

• Proposes UNesT, a 3D U-shape medical segmentation model using hierarchical transformer encoders with 3D block aggregation to enable communication between sequence representations.
• Introduces a block aggregation function that merges patches at a factor of 8 every two transformer layers, enhancing local inductive bias without altering the global self-attention mechanism.
• Employs a simplified design that retains standard self-attention while improving data efficiency and representation learning for small structures.
• Uses a nested transformer architecture inspired by nested ViT, enabling hierarchical modeling of high-resolution 3D medical volumes.
• Trains the model on a newly collected dataset of 116 contrast-enhanced CT scans with expert-annotated renal substructures.
• Validates performance using Dice score, Pearson correlation with manual segmentation, and convergence analysis on varying data regimes.

Experimental results

Research questions

• RQ1 Can a 3D block aggregation mechanism improve data efficiency in transformer-based 3D medical image segmentation without modifying self-attention?
• RQ2 Does the proposed hierarchical transformer design yield better performance for small renal substructures compared to standard transformers and local self-attention variants?
• RQ3 To what extent does the model generalize to external datasets like KiTS19, and how does it compare to existing state-of-the-art methods?
• RQ4 How does the model’s performance and convergence rate vary under low-data regimes, and what is its reproducibility relative to manual segmentation?
• RQ5 Can the method achieve high correlation with manual volumetric measurements, indicating clinical utility for biomarker discovery?

Key findings

• The proposed UNesT model achieves a Dice score of 0.8467 on the renal substructure dataset, outperforming the baseline (0.8308) and demonstrating state-of-the-art performance.
• The method shows a Pearson correlation of 0.9891 between automated and manual volumetric measurements, indicating strong reproducibility and reliability.
• On low-data regimes, the block aggregation design improves Dice score by over 4% compared to the second-best method, highlighting enhanced data efficiency.
• The model converges 15% faster at 2K iterations and 4% faster at 30K iterations when using hierarchical modules, indicating improved training dynamics.
• When evaluated on the public KiTS19 dataset, UNesT achieves a Dice score of 0.9778 for kidneys and 0.8398 for tumors, showing strong generalizability.
• Bland-Altman analysis confirms that automated segmentation agrees more closely with manual annotations (mean difference: 0.01) than inter-rater agreement (0.29), indicating high reproducibility.