3D TransUNet: Volumetric Segmentation with Transformers

• 3D TransUNet is a neural architecture that merges 3D U-Net with Vision Transformers to capture both local spatial details and global context in volumetric image segmentation.
• It offers architectural variants (encoder-only, decoder-only, and encoder+decoder) to optimally tackle challenges in multi-organ, small-structure, and tumor segmentation.
• Empirical studies demonstrate state-of-the-art Dice scores and efficiency improvements through MAE pre-training and adaptive training strategies.

3D TransUNet is a neural architecture designed for volumetric medical image segmentation that integrates Vision Transformer modules within the U-Net framework. By embedding transformer-based global self-attention into 3D convolutional encoder–decoder structures, 3D TransUNet addresses both the localized spatial feature extraction typical of U-Nets and the long-range dependency modeling characteristic of transformers. The approach offers multiple architectural variants targeting the challenges of multi-organ, small-structure, and tumor segmentation, leveraging both the volumetric capabilities of nnU-Net and the non-local attention of the Transformer. Extensive empirical evidence establishes state-of-the-art performance in diverse clinical tasks and highlights the impact of strategically assigning Transformer modules to encoder, decoder, or both.

1. Architectural Foundations and Variants

3D TransUNet extends the original 2D TransUNet to fully volumetric segmentation, building upon the adaptive 3D nnU-Net design (Chen et al., 2023). The backbone is a U-shaped encoder–decoder model, with each component realized in three dimensions:

• Encoder: Four or five 3D convolutional down-sampling stages extract hierarchical feature maps, each stage halving the spatial dimensions and increasing the channel count, producing multi-scale features (e.g., [D/2, H/2, W/2, C₁], …, [D/8, H/8, W/8, C₃]).
• Transformer Encoder: CNN feature maps are partitioned into non-overlapping 3D patches (P×P×P), each flattened and linearly projected into tokens $z_0$. Positional encoding is added. $L_{enc}$ stacked Transformer layers process these tokens, implementing multi-head self-attention and MLP blocks. Output tokens are reshaped into volumetric feature maps and reintegrated into the decoder via up-sampling and concatenation.
• Transformer Decoder: Candidate object/organ queries $P^0$ are refined over $T$ iterative cross-attention updates using multi-scale CNN decoder features $\mathcal{F}_t$. Query updates are masked via coarse foreground predictions to focus on relevant spatial regions. Hungarian matching after mask-classification queries aligns proposals to ground truth.
• Variants: Encoder-only (Transformer in encoder only), Decoder-only (Transformer in decoder only), Encoder+Decoder (both modules). Encoder-only excels in multi-organ segmentation; Decoder-only is superior for small/complex lesions (Yang et al., 2024).

Implementation follows the canonical nnU-Net framework for convolutional modules, data normalization, and augmentation. Skip-connections are maintained between encoder and decoder stages.

2. Mathematical Formulation

Key transformations in 3D TransUNet rely on patch-embedding, attention, and mask refinement:

• Patch Embedding: For an input $x \in \mathbb{R}^{D\times H\times W\times C}$,

$z_0 = [x_1^pE; x_2^pE; \ldots; x_N^pE] + E^{pos},$

where $x_i^p$ is the flattened $i$-th patch, $E$ is a linear projection, and $E^{pos}$ is positional encoding.

• Transformer Encoder Layer:

$$z'_l = \mathrm{MSA}(\mathrm{LN}(z_{l-1})) + z_{l-1}, \ z_l = \mathrm{MLP}(\mathrm{LN}(z'_l)) + z'_l.$$

• Self-attention: For queries $Q$, keys $K$, values $V$,

$$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}(QK^\top/\sqrt{d_k})V.$$

• Cross-attention in Decoder: For queries $P^t$, features $\mathcal{F}$,

$$P^{t+1} = P^t + \mathrm{Softmax}\left[(P^t W_q)(\mathcal{F} W_k)^\top\right](\mathcal{F} W_v).$$

Foreground mask $Z^t$ constrains attention spatially:

$$h(Z^t) = 0 \text{ if } Z^t(i,j,s)=1; \text{ else } -\infty.$$

• Loss Functions: Encoder-only uses

$L = L_{ce} + L_{dice},$

while Decoder-only utilizes Hungarian-matched mask/classification loss:

$$\mathcal{L} = \lambda_0 (L_{ce} + L_{dice}) + \lambda_1 L_{cls},\quad \lambda_0=0.7,\ \lambda_1=0.3.$$

Dice loss evaluates spatial overlap; binary cross-entropy penalizes voxel-wise misclassification.

3. Integration with the nnU-Net Framework

3D TransUNet inherits the adaptive design, cropping, channel width, augmentation strategy, and normalization protocol of nnU-Net (Chen et al., 2023). Specifically:

• Early-stage convolutional outputs provide input space for transformer tokenization.
• LayerNorm is adopted in transformer modules; instance normalization in convolutional layers.
• Skip connections and decoder upsamplings exploit nnU-Net's multi-scale framework for fine spatial localization.
• Sliding window inference aggregates predictions, with softmax applied over cross-attention queries.
• Low-resolution simulation, random geometric perturbations, and contrast/brightness normalization are standard augmentations.

This integration ensures both scalability to large volumes and compatibility with existing medical segmentation benchmarks.

4. Empirical Results and Task-Specific Performance

Performance was benchmarked on multi-organ CT (Synapse), hepatic vessel segmentation (MSD), brain metastases (BraTS2023), and pancreatic tumor datasets.

Variant/Task	Avg Dice (%)	HD95 (mm)	Dataset
Encoder-only	88.11	6.02	Synapse
Decoder-only	59.3	96.4	BraTS2023 METS
nnU-Net Baseline	87.33	6.46	Synapse
Decoder-only	69.69	–	Pancreas Tumor

On Synapse, Encoder-only yields +0.8% Dice over nnU-Net. On BraTS METS, Decoder-only surpasses Encoder-only by 2.7 percentage points in lesion-wise Dice (Yang et al., 2024). On hepatic vessels, Decoder-only improves average Dice by 1.6%.

Ablation studies confirm that multi-scale cross-attention and masked queries are critical for localizing fine structures; combining both encoder and decoder transformers did not always yield additive gains, with task-dependent optimal configurations. Query count (Nₑ) exhibited negligible impact on segmentation accuracy.

5. Training Strategies, Pre-training, and Computational Considerations

Training protocols leverage AdamW and SGD optimizers, cosine/polynomial learning rate schedules, and deep supervision via per-stage losses. For Encoder-only, Masked Autoencoder (MAE) pre-training on image tokens accelerates convergence and stabilizes optimization, halving the necessary epochs to reach competitive Dice (Yang et al., 2024). Decoder-only configuration demands longer epochs due to increased computational complexity from high-resolution cross-attention and Hungarian matching.

Input volumes and patch sizes are tuned per dataset, with typical patch embeddings of P=16. During training, batch sizes and data augmentations are adjusted for hardware constraints (NVIDIA RTX 8000) and task specificity.

Computational cost is elevated for Decoder-only variants (≈1.5× slower per epoch) and memory footprint increases with multi-scale cross-attention. Encoder-only is more efficient and easier to deploy at scale.

6. Extensions, Limitations, and Future Directions

3D TransUNet demonstrates superior segmentation for brain metastases, multi-organ CT, and small tumor delineation. Limitations arise in combining both transformer modules, where synergy is non-universal and requires task-specific tuning (Chen et al., 2023). Hungarian matching, essential for query-centric mask classification, constitutes a training bottleneck; potential acceleration via approximate methods is identified.

MAE pre-training is essential for encoder configuration and may be extended to deeper decoder layers. Incorporation of sparsity via dynamic attention or token selection is a promising avenue for reducing resource demands. The architecture is well positioned for further extension into multi-class segmentation and uncertainty quantification.

7. Comparative Models and Related Research

Comparison studies incorporate nnU-Net, CoTr, nnFormer, Swin UNETR, TransBTS, and Attention U-Net baselines (Chen et al., 2023, Wang et al., 2022):

• 3D TransUNet outperforms nnU-Net and other transformer hybrids by margins of 0.8–3.0% in Dice on key datasets.
• MISSU embeds a small transformer on the deepest 3D feature map and employs self-distillation to refine shallow features, achieving 89.98/85.77/80.14 Dice on BraTS 2019 (WT/TC/ET), besting previous benchmarks (Wang et al., 2022).
• Dynamic Linear Transformer leverages ROI-based dynamic token sampling and linear self-attention (O(n)) for 3D medical segmentation (Zhang et al., 2022).

This landscape highlights the value of transformer-based global attention, adaptive query refinement, and hybrid architectures for volumetric medical image analysis.