3D U-Net

The architecture follows a symmetric contracting and expanding path. The encoder applies successive 3D convolutions (3×3×3 kernels) and 2×2×2 max-pooling to capture hierarchical contextual features while halving spatial resolution at each level. The decoder uses 2×2×2 transposed convolutions to upsample feature maps, doubling resolution until the original input dimensions are restored. This design captures both fine-grained anatomical detail and global spatial context across all three dimensions simultaneously.

Skip connections concatenate high-resolution feature maps from the encoder directly to the corresponding decoder level. This mechanism preserves fine spatial information that would otherwise be lost during downsampling, enabling precise boundary delineation of small structures such as:

• Tumor margins in brain MRI segmentation
• Vessel walls in CT angiography
• Cortical surfaces in neuroimaging

Without skip connections, the decoder would rely solely on coarse, semantically rich features, producing blurred segmentation boundaries.

Unlike 2D U-Nets that process slices independently, the 3D variant uses 3×3×3 convolutional kernels that operate on volumetric patches. Each kernel simultaneously captures spatial relationships in the axial, coronal, and sagittal planes. This is critical for structures with complex 3D morphology, such as branching vasculature or irregular tumor shapes. The receptive field grows cubically with depth, allowing the network to learn true volumetric context rather than inferring it post-hoc from 2D slice stacks.

Due to GPU memory limitations, 3D U-Nets typically process sub-volumes or patches (e.g., 128×128×128 voxels) rather than entire scans. Instance normalization is often preferred over batch normalization because:

• Small batch sizes (often 1-2) make batch statistics unreliable
• Instance normalization normalizes each volume independently, stabilizing training
• It preserves contrast information critical for distinguishing tissue types

Patch-based inference with overlapping tiles and Gaussian weighting at boundaries ensures seamless full-volume reconstruction.

Deep supervision adds auxiliary segmentation outputs at intermediate decoder levels, not just the final layer. Each auxiliary output is compared against a downsampled version of the ground truth segmentation mask, and the combined loss is backpropagated. Benefits include:

• Improved gradient propagation to early encoder layers, combating vanishing gradients in very deep networks
• Faster convergence during training
• Multi-scale regularization, encouraging intermediate features to be semantically meaningful

This technique is especially valuable when training on limited medical datasets where every gradient signal matters.

Medical segmentation datasets often exhibit severe class imbalance, where lesions or small organs occupy a tiny fraction of total voxels. 3D U-Nets commonly employ specialized loss functions:

• Dice Loss: Directly optimizes the Dice Similarity Coefficient, measuring spatial overlap between prediction and ground truth
• Focal Tversky Loss: Weights false positives and false negatives asymmetrically, useful when precision and recall have different clinical costs
• Combined Cross-Entropy + Dice Loss: Balances voxel-wise classification accuracy with region-based overlap

These losses prevent the network from collapsing to an all-background prediction in highly imbalanced scenarios.

The Dice Similarity Coefficient is the primary spatial overlap metric for evaluating 3D U-Net segmentation performance. It quantifies the agreement between predicted and ground-truth masks, defined as 2 × |A ∩ B| / (|A| + |B|). A DSC of 1.0 indicates perfect voxel-wise overlap; values above 0.90 are typically considered clinically acceptable for large organs.

• Sensitive to both false positives and false negatives
• Complements surface-distance metrics like Hausdorff Distance
• Calculated per-class in multi-organ segmentation tasks

A segmentation mask is a label map that classifies every voxel in a 3D volume as belonging to a specific anatomical structure or lesion. 3D U-Nets output dense, voxel-wise probability maps that are thresholded or argmaxed to produce the final mask. Multi-class masks assign unique integer labels to each structure—for example, 1 for liver, 2 for tumor, 0 for background.

• Binary masks: foreground vs. background
• Multi-class masks: multiple anatomical structures
• Instance masks: distinguish individual objects of the same class