Branched U-Net for Multitask Segmentation

• Branched U-Net is a convolutional encoder-decoder architecture with a shared encoder and multiple decoder branches for simultaneous, expert-specific segmentation tasks.
• It maintains skip-connection design for high spatial fidelity while incorporating dynamic, branch-specific losses to handle multimodal inputs and inter-annotator variability.
• Empirical evaluations of variants like U-Net-and-a-half and MBDRes-U-Net demonstrate improved generalization, reduced computational overhead, and enhanced segmentation accuracy in 3D medical imaging.

A branched U-Net is a family of convolutional encoder–decoder networks in which multiple output streams (branches) emerge from a shared encoder, enabling simultaneous learning of distinct yet related segmentation tasks or the integration of multiple supervision signals. While preserving the fundamental skip-connection architecture of the vanilla U-Net, branched designs are motivated by the need to model inter-annotator variability, enhance multitask learning, exploit multimodal input, or reduce computational burden in 3D medical image processing.

1. Architectural Principles and Design Variants

Branched U-Nets are defined by a single, shared encoder that produces a low-dimensional feature space, followed by two or more parallel decoder branches. Each branch specializes in a particular supervision signal, expert annotator, or data modality, with skip-connections preserved independently for each decoder to maintain high spatial fidelity. Key instantiations include:

• U-Net-and-a-half (Zhang et al., 2021): Employs a ResNet-50 encoder and $M$ parallel decoders (typically $M=2$), each receiving the same shared feature map at the deepest layer (Conv5_x output). Each decoder has a standard upsampling path with skip-connections. All decoders are jointly supervised, typically with expert-specific annotation masks (Zhang et al., 2021).
• MBDRes-U-Net (Chen et al., 2024): Implements a multi-branch residual block strategy in a 3D U-Net, where the encoder and each decoder stage process multi-channel inputs by distributing channel groups to parallel streams, followed by adaptive weighted fusion, dilated convolutions, and spatial–channel attention. Decoders are not expert-specific but leverage branching within each residual block to exploit multimodal sources and spatial context (Shen et al., 4 Nov 2024).
• Supervised Bottleneck Branch (Rahman et al., 2020): Applies a fully-connected branch at the bottleneck, with a parallel pixel-wise cross-entropy supervision signal, thereby enforcing semantic content even at the network's deepest representation. The decoder remains single-branched, but the bottleneck itself is branched for targeted supervision (Zahra et al., 2020).

2. Detailed Layer Structures and Branching Mechanisms

U-Net-and-a-half (Zhang et al., 2021)

Encoder:

• ResNet-50 backbone with initial 7×7 convolution, batch normalization, ReLU activation, max pool, followed by four residual blocks (Conv2_x to Conv5_x). The Conv5_x output at 2048×7×7 forms the shared low-dimensional space.

Branching Point:

• The feature map at 2048×7×7 is copied to each decoder, ensuring gradients from all branches propagate to the same encoder.

Decoder (for each branch $k$, $k=1,\ldots,M$):

• Sequence of up-convolutions, each concatenated with skip-connections from the corresponding encoder stage.
• Final 1×1 convolution with sigmoid activation for binary segmentation ($C$ output channels).
• Each decoder is supervised by its corresponding ground-truth mask.

Losses:

• Expert-specific branches supervised by hybrid cross-entropy/focal losses, with outputs summed before evaluation.
• Dynamic weighting parameter $w$ reflects inter-annotator Dice/IoU agreement.
• See Section 4 below for explicit loss equations.

MBDRes-U-Net (Shen et al., 4 Nov 2024)

Encoder/Decoder Path:

• Six-level 3D U-Net, input size 128³, FLAIR/T1/T1c/T2 channels.
• Each encoder block: MBDRes block (multi-branch, group convolutions), stride-2 downsampling.
• Multi-branch mechanism: Input is projected by 1×1×1 conv, split into $g$ channel groups ($g=8$). Each branch processes its group by (possibly adaptive) 3D dilated convolution (dilation rates 1,2,3), and the fused sum is projected back and added to the input (residual connection).

Attention Mechanisms:

• Pre-encoder 3D SACA (spatial–channel attention) block splits channels, applying channel and spatial attention separately, then shuffles and concatenates.

Branching Strategy:

• The "branch" in MBDRes-U-Net refers to per-block parallel processing, not multiple output decoders. This design enables the model to separately encode features from different modalities or receptive fields within each block.

Supervised Bottleneck Branch (Zahra et al., 2020)

Bottleneck FC Branch:

• After the deepest convolutional block, activations are flattened; two FC layers produce a tensor with pixel-wise output, reshaped to match the bottleneck's spatial size.
• This auxiliary output is supervised with cross-entropy loss.
• Output is element-wise combined with the bottleneck feature map before upsampling in the decoder.

3. Supervision Strategies and Loss Formulations

Multi-Expert Supervision (Zhang et al., 2021)

Let $O_1, O_2$ be decoder outputs, $A_1, A_2$ the respective annotation masks. The key loss formulations are:

Hybrid Cross-Entropy Loss:

$$\mathcal{L}_{\mathrm{hybrid}}(O_1, O_2; A_1, A_2) = -w \sum_i [A_{1,i} \log(O_{1,i} + O_{2,i})] - (1-w) \sum_i [A_{2,i} \log(O_{1,i} + O_{2,i})]$$

with $w \in [0,1]$ dynamically updated per epoch.

Hybrid Focal Loss:

$$\mathcal{L}_{\mathrm{focal}} = w c_w (1 - e^{-\mathrm{CE}_1})^\gamma \mathrm{CE}_1 + (1-w) c_w (1 - e^{-\mathrm{CE}_2})^\gamma \mathrm{CE}_2$$

with parameter definitions as in the provided equations.

At inference, each branch predicts an output, and ensemble strategies (e.g., sum, average) produce the final segmentation mask.

Bottleneck Supervision (Zahra et al., 2020)

• Bottleneck FC branch output is supervised with pixel-wise cross-entropy loss.
• Decoder output is supervised with $L_1$ loss against ground truth.
• Total loss: $L_\mathrm{total} = L_\mathrm{ce} + \lambda L_1$.

Multimodal and Parameter Efficiency (Shen et al., 4 Nov 2024)

• Standard voxel-wise cross-entropy across four MRI-based output classes.
• Optionally, soft Dice loss may be included.

4. Implementation Considerations

Data and Augmentation:

Study	Input Data	Patch Size	Augmentation
(Zhang et al., 2021)	10 WSI kidney (512×512), 10 IVUS (full frame)	WSI: 224×224	Rotation, flip, crop, color
(Shen et al., 4 Nov 2024)	BraTS18/19, four MRI channels	128×128×128	Z-score norm, flip, rotate
(Zahra et al., 2020)	MRI, CT (2D)	Not specified	Not specified

• U-Net-and-a-half uses PyTorch, large GPU memory (RTX 3090, 24 GB), batch sizes 8–16, and convergence times ranging from 10 minutes (IVUS) to 4 hours (WSI).
• MBDRes-U-Net: 3× NVIDIA A30 GPUs, batch size 16, epochs = 500.

Resource Efficiency:

• MBDRes-U-Net achieves $\sim4\times$ reduction in parameters and $>50\times$ in FLOPs relative to conventional 3D U-Net, with improved Dice (+1.8–13.6%) on BraTS datasets (Shen et al., 4 Nov 2024).
• U-Net-and-a-half introduces little overhead relative to a single large U-Net, as decoder duplication is efficient in the context of modern GPUs (Zhang et al., 2021).

5. Empirical Performance and Generalization

Results on Public Datasets

Architecture	Domain	Dice (mean ± SD)	IoU (mean ± SD)	Notes
U-Net-and-a-half	WSI (A₁)	0.9874 ± 0.005	0.7573 ± 0.02	Matches inter-annotator Dice
Single U-Net crossval	WSI (1→2)	0.9834 ± 0.006	0.7604 ± 0.01	Degrades when cross-tested
MBDRes-U-Net	BraTS18	79.9/90.5/84.6	–	3.2–13.6% Dice increase; ≈¼ params

The branched U-Net outperforms single-expert U-Nets in generalization to held-out annotations, mitigating overfitting to any one expert's viewpoint. In multimodal, 3D contexts, branching reduces computation without degrading segmentation accuracy.

6. Practical Applications and Limitations

Branched U-Nets are particularly suitable when:

• Multiple experts annotate the same images, and no consensus ground truth exists (e.g., digital pathology, medical imaging with variable context).
• Segmentation must integrate multi-modality data, requiring explicit modeling of feature dependencies (as in MBDRes-U-Net for MRI sequences).
• Parameter or compute constraints demand lightweight segmentation architectures.

Current limitations include:

• Additional complexity during training, especially in managing branch-specific losses and dynamic loss weights.
• Increased parameter count relative to single-output U-Net, though mitigated by group convolutions or in-block branching as in MBDRes-U-Net.
• For some branched designs (e.g., supervised bottleneck), insufficient evidence is presented regarding quantitative improvement or optimal hyperparameters (Zahra et al., 2020).

7. Relationship to Related Architectures and Future Perspectives

Branched U-Nets generalize standard encoder-decoder architectures (U-Net, V-Net) by enabling parallelization at crucial layers. They share design philosophy with multitask learning networks, mixture-of-experts models, and architectures that address annotator or modality heterogeneity.

Potential future directions include:

• Expanding to more than two decoders for applications requiring the modeling of more extensive expert variability.
• Incorporation of semi-supervised or unsupervised branches to leverage unlabeled data.
• Further architectural innovations combining explicit attention, adaptive loss weighting, and multi-resolution aggregation to maximize annotation-signal efficiency and domain generalization.

Branched U-Net architectures thereby offer a principled extension of the U-Net paradigm for multi-source, multi-expert, and multimodal segmentation tasks, with reproducible design specifications and strong empirical performance in biomedical imaging benchmarks.