Three-Stage U-Net Architecture

• Three-stage U-Net-based architectures are advanced CNN systems that integrate parallel, cascaded, or modality-specific U-Nets to enhance segmentation precision.
• They leverage multi-resolution inputs and hierarchical fusion to improve feature learning, yielding faster convergence and higher Dice scores across benchmarks.
• These architectures are applied in biomedical imaging tasks such as skin lesion, brain tumor, and histology segmentation, exemplified by mrU-Net and Triple U-Net.

A three-stage U-Net-based architecture refers to a class of convolutional neural network (CNN) segmentation systems composed of either (i) @@@@1@@@@ multi-resolution streams fused within a single network, or (ii) serial cascades of multiple U-Net or closely related sub-networks, each segmenting different features or regions. These systems exploit either spatial scale, semantic priors, or task-specific feature extraction at each stage, and are designed to improve accuracy, convergence, and handling of hierarchical or instance-differentiating structures in biomedical images. Representative implementations include the multi-resolution U-Net (mrU-Net), sequential U-Nets for multi-label segmentation, and multi-branch ("triple U-Net") pipelines for modality-specific feature fusion. Below, key instantiations, architectural variants, methodological details, and their experimental impact are reviewed.

1. Multi-Resolution and Multi-Stream Three-Stage U-Net

The multi-resolution U-Net architecture, as embodied by mrU-Net, augments the base U-Net design by inserting three parallel input streams at different spatial resolutions. The core objective is the simultaneous extraction and integration of low- and high-frequency features, allowing faster and more robust convergence for segmentation tasks involving scale-heterogeneous objects (Jahangard et al., 2020).

Formal Construction

Let $X \in \mathbb{R}^{H \times W \times C}$ denote the input image. Three down-sampled variants are defined:

• $X_1 = D_1(X) \equiv X$ (original scale, $H \times W$),
• $X_2 = D_2(X)$ (downsampled by $2 \times$, $H/2 \times W/2$),
• $X_4 = D_4(X)$ (downsampled by $4 \times$, $H/4 \times W/4$).

Each $X_k$ passes through a standalone convolutional block:

$$f_k(X_k;\theta_k) = \text{ReLU}(\text{Conv}_{3 \times 3}(\text{ReLU}(\text{Conv}_{3 \times 3}(X_k))))$$

and the resulting feature maps $f_k$ are upsampled as necessary and concatenated with the main contracting path at matching spatial depths. The fused feature tensor before the bottleneck is

$$F = \text{concat}\left[f_1,\; \text{up}_2(f_2),\; \text{up}_4(f_4)\right]$$

where $\text{up}_2$ and $\text{up}_4$ denote upsampling by $2 \times$ and $4 \times$, respectively.

Encoder-Decoder Details

The encoder comprises classic U-Net levels (4 total), with "hashed" blocks for embedding low-resolution features. Decoding mirrors the encoder, maintaining skip connections and dimension consistency; the final 1×1 convolution reduces channels to output labels, followed by a softmax. All convolutions are stride-1, zero-padded; pooling and upsampling are executed with stride-2 operators.

Training and Loss

A soft Dice loss is minimized:

$$L = 1 - \frac{2\sum_i p_i g_i}{\sum_i p_i + \sum_i g_i}$$

where $p_i$ is the predicted probability and $g_i$ the ground-truth label for pixel $i$. The model is optimized with Adadelta (learning rate 1.0), batch size 16, and heavy geometric augmentation for some datasets.

Quantitative Performance

Evaluations on four medical segmentation benchmarks (skin lesions, lung nodules, retinal vessels, prostate MRI) reveal modest but statistically significant Dice improvements over standard U-Net for most tasks, with mrU-Net achieving, e.g., 97.9% Dice on LUNA (vs. 97.3%) and 73.6% on DRIVE (vs. 73.1%). Convergence is faster (peak Dice achieved in fewer epochs), consistent with the hypothesis that direct multi-resolution input facilitates feature learning (Jahangard et al., 2020).

Limitations

The three-stream design increases computational and memory demands, and may require further architectural adaptation for high-dimensional or 3D tasks. The method was validated solely on 2D datasets.

2. Cascaded and Sequential U-Nets for Hierarchical and Multi-Label Segmentation

A second paradigm leverages sequential three-stage U-Nets, where each stage segments a specific spatial or semantic subregion with conditioning on previous output masks. This is prominent in glioma and brain tumor segmentation, exploiting known anatomical or pathological hierarchies (Beers et al., 2017, Ghaffari et al., 2020).

Clinical Motivation and Design

In brain tumor imaging (BraTS dataset), lesions are organized as "whole tumor" (WT), "tumor core" (TC), and "enhancing tumor" (ET), with strict spatial containment: ET ⊆ TC ⊆ WT. The pipeline consists of three networks:

1. Stage 1: Segment WT using multimodal MRI. For (Beers et al., 2017), a cascade of two U-Nets at different resolutions (2 mm then 1 mm) generates a high-quality mask.
2. Stage 2: Segment ET, using input MRI plus the binary WT mask as a fifth channel.
3. Stage 3: Segment TC, again using the WT mask for spatial restriction.

This pipeline enforces hierarchical consistency, restricts later stages to smaller regions of interest, and improves training stability for difficult sublabels (Beers et al., 2017, Ghaffari et al., 2020).

Implementation and Loss

Each 3D U-Net comprises four levels of down- and up-sampling, with patch-based training. The soft Dice loss is applied independently for each output:

$$L_{\text{Dice}} = 1 - \frac{2\sum_i p_i g_i + \epsilon}{\sum_i p_i + \sum_i g_i + \epsilon}$$

where $p_i$ are predictions, $g_i$ ground-truth, $\epsilon$ for stability.

Performance and Ablation

On BraTS, the sequential approach yields strong region-wise Dice coefficients: WT 0.882, ET 0.732, TC 0.730 (Beers et al., 2017). Ablation studies indicate that providing the WT mask as input to stages 2 and 3 increases ET and TC Dice by 3–5 points, confirming its utility as an anatomical prior. The cascaded Dense U-Net variant (Ghaffari et al., 2020) with self-ensembling achieves higher Dice for WT (0.90), TC (0.82), ET (0.78) and benefits from focused patch sampling and post-hoc connected component filtering.

3. Multi-Branch Triple U-Net Architecture for Modality-Specific Fusion

A third design, exemplified by the Triple U-Net, constructs three parallel U-Net branches with modality- or feature-specific inputs, fusing them at multiple decoder scales to enhance fine-grained, instance-level accuracy (Ahmed et al., 2024).

Architectural Composition

The Triple U-Net consists of:

• RGB branch: standard U-Net for learning semantic and appearance features from raw RGB images, supervised with pixel-wise cross-entropy loss.
• Hematoxylin branch ("H branch"): U-Net for contour and edge features using the Hematoxylin channel (obtained by Beer–Lambert color deconvolution), supervised with a Soft-Dice loss.
• Segmentation branch: A third U-Net fusing features from both preceding branches using Progressive Dense Feature Aggregation (PDFA) modules at each scale. The PDFA progressively concatenates current and prior outputs from all branches, enabling dense, hierarchical feature integration.

Mathematically, for scale $l$, PDFA:

$$y_1^l = \text{Conv}\left([F_{RGB}^l, F_H^l, S^{l+1}]\right) \ y_2^l = \text{Conv}\left([F_{RGB}^l, F_H^l, S^{l+1}, y_1^l]\right) \ \ldots$$

where $F_{RGB}^l$, $F_H^l$ are encoder feature maps, $S^{l+1}$ is the deeper scale's fused feature.

Preprocessing and Postprocessing

Color deconvolution isolates the Hematoxylin channel from RGB using optical density and a stain matrix. Final instance segmentation is achieved via marker-based watershed, leveraging the nuclei center distance transform to resolve overlapping instances.

Training and Metrics

Total loss is a weighted sum of binary cross-entropy and Dice for each branch. On the CryoNuSeg dataset, Triple U-Net outperforms the single-branch baseline substantially in Average Jaccard Index (AJI; 0.6741 vs. 0.525) and Panoptic Quality (PQ; 0.5056 vs. 0.477). AJI gains, in particular, are attributed to improved separation of overlapping nuclei facilitated by the H branch and dense decoder fusion (Ahmed et al., 2024).

4. Comparative Summary of Three-Stage U-Net Architectures

Approach	Staging Principle	Fusion/Conditioning	Representative Task	Reported Dice/AJI (%)
mrU-Net (Jahangard et al., 2020)	Multi-scale, parallel streams	Concatenation at encoder	2D medical segmentation	Skin 70.6, LUNA 97.9, DRIVE 73.6
Cascaded U-Net (Beers et al., 2017, Ghaffari et al., 2020)	Sequential, subregion hierarchy	Mask as input at each stage	Brain tumor (3D MRI)	WT 88–90, TC 73–82, ET 73–78
Triple U-Net (Ahmed et al., 2024)	Parallel, modality-specific branches	PDFA fusion in decoder	Histology instance segmentation	AJI 67.4 (vs. U-Net 52.5), PQ 50.6

The mrU-Net emphasizes improved multiresolution feature extraction and faster convergence for 2D tasks. Cascaded/Sequential U-Nets focus on anatomical or semantic hierarchy enforcement in 3D, providing spatial priors and subregion containment. The Triple U-Net exploits feature complementarities between image modalities and specialized loss functions, primarily targeting high-precision, instance-level instance segmentation under challenging image conditions.

5. Limitations and Future Research Directions

Three-stage U-Net-based architectures—multi-resolution, sequential, or multi-branch—require significantly more parameters, GPU memory, and careful optimization relative to basic U-Net. In 3D or very high-resolution contexts, these demands scale rapidly and may necessitate more parameter-efficient fusions, e.g., 1×1 convolutions or channel bottlenecks. None of the cited works performed comprehensive hyperparameter searches, and specialized loss formulations (e.g., for fine vessel topology or boundary adherence) remain open for further study (Jahangard et al., 2020, Ahmed et al., 2024).

Extensions to fully 3D input, integration with attention mechanisms, or incorporation of additional biological priors are plausible improvements. A plausible implication is that as problem complexity increases—e.g., for highly nested or overlapping structures—the added architectural granularity and fusion of three-stage U-Nets is likely to yield greater improvements in segmentation quality than on simpler, globally homogeneous tasks. Empirical evaluation on non-biomedical domains remains limited.

6. Conclusion

Three-stage U-Net-based architectures generalize the canonical U-Net by introducing explicit multi-scale, hierarchical, or modality-diverse feature pipelines, yielding consistent improvements in precision, convergence, and robustness across a range of biomedical image segmentation benchmarks. Their tight fusion of feature maps or sequential use of predicted masks enables faithful instance delineation and explicit topology enforcement, albeit at higher computational cost. Continued investigation into parameter-efficient fusions, loss tailoring, and broader application domains is warranted.