3D U-Net and Unidirectional Transformers

• 3D U-Net and Unidirectional Transformers are distinct architectures that combine localized 3D convolutions with global self-attention to capture detailed and long-range spatial features.
• The UNETR model integrates a transformer encoder with a 3D convolutional decoder using multiresolution skip connections, achieving superior Dice scores on multiple medical segmentation benchmarks.
• Potential advances include leveraging unidirectional (causal) transformer mechanisms to reduce computational complexity while introducing spatial hierarchy, marking an open research avenue.

3D U-Net and Unidirectional Transformers are distinct architectural paradigms for volumetric medical image segmentation, fundamentally differing in their mechanisms for spatial feature learning. The UNETR architecture exemplifies a hybrid approach that integrates a pure transformer encoder—characterized by global, bidirectional self-attention—with the classical U-shaped 3D convolutional decoder, achieving high-resolution segmentation in volumetric data. The dichotomy between bidirectional and unidirectional (causal) transformer attention remains a subject of ongoing research, particularly regarding computational efficiency and inductive bias in 3D context modeling (Hatamizadeh et al., 2021).

1. U-Net Paradigm in Volumetric Segmentation

The U-Net architecture, originally designed for 2D biomedical image segmentation, has been extensively adapted to 3D settings. The 3D U-Net introduces fully convolutional encoder-decoder pathways that contract spatially to extract hierarchical features and then expand to recover fine-grained voxel-wise segmentations. The encoder typically operates through a cascade of spatially down-sampling 3D convolutions, learning both local and contextual cues, while skip connections between corresponding encoder and decoder layers preserve spatial detail lost during down-sampling (Hatamizadeh et al., 2021).

However, a key limitation of the traditional 3D U-Net is the locality of convolutional receptive fields, restricting the model’s efficacy in modeling long-range anatomical dependencies within volumetric data. This constraint motivates the incorporation of architectures capable of capturing global spatial context.

2. Transformer Encoder Integration: UNETR

UNETR reformulates 3D medical image segmentation as a sequence-to-sequence prediction task by embedding volumetric patches and feeding them through a stack of transformer layers. The input $x \in \mathbb{R}^{H \times W \times D \times C}$ is divided into $N = (H \cdot W \cdot D) / P^3$ non-overlapping patches of size $P^3$, each linearly projected into a $K$-dimensional space. Learnable positional embeddings are applied to preserve the spatial ordering of the patches. These embeddings comprise the input to a 12-layer transformer encoder, conforming to the Vision Transformer (ViT-B/16) configuration:

• Embedding dimension $K = 768$
• Number of heads $H = 12$
• MLP hidden dimension $4K = 3072$
• LayerNorm pre-normalization and residual connections after each MSA/MLP sub-layer

Skip connections are extracted at four depths ($i \in \{3, 6, 9, 12\}$), where transformer outputs $z_i$ are recomposed into 3D volumetric features, convolved, and interfaced with the up-sampling stages of the decoder (Hatamizadeh et al., 2021).

3. Self-Attention Directionality and Its Implications

The self-attention mechanism in UNETR is bidirectional: each patch embedding at a given layer attends to all other patches, enabling global context aggregation across the entire volume. This full (non-causal) connectivity is critical for segmenting complex anatomical structures, where spatial relationships may span large regions inaccessible to local convolutional operations.

The UNETR technical report explicitly notes that its attention is non-causal and bidirectional, distinguishing it from unidirectional (causal) transformers that impose sequence ordering via masking. The authors suggest that unidirectional transformer variants—using explicit spatial masking strategies to enforce partial attention—could introduce hierarchical spatial priors or reduce computation, but such approaches remain unexplored in the current UNETR implementation. This suggests potential avenues for structuring long-range dependencies while mitigating the quadratic complexity of standard self-attention (Hatamizadeh et al., 2021).

4. Decoder Design and Multiresolution Skip Connections

UNETR decoders are composed of four cascaded 3D deconvolutional stages. At each up-sampling step, feature maps are concatenated with the corresponding projected encoder features—reshaped from transformer outputs and further refined by convolutional blocks. Each decoder block thus benefits from both the high-level, globally-aware transformer features and the locally-dense convolutional representations.

The process for each decoder stage $k$ can be formalized as:

$U_k = \text{Deconv2×}(F_{k+1}),$

$$F_k = \text{ConvBlock}_{2×}([U_k; \text{Project}(T_i)]),$$

where $T_i$ denotes the projected transformer feature at depth $i$, and $\text{ConvBlock}_{2×}$ is a pair of 3D conv-norm-ReLU layers. The pipeline enables effective spatial alignment between hierarchical contexts, culminating in a 1×1×1 convolutional head for class-specific prediction and softmax activation (Hatamizadeh et al., 2021).

5. Quantitative Evaluation and Ablation Analysis

UNETR establishes state-of-the-art results on prominent 3D medical segmentation benchmarks:

Dataset	nnUNet	TransUNet	CoTr	UNETR
BTCV (multi-organ, CT)	80.2% Dice	83.8% Dice	84.4% Dice	85.6% Dice
BTCV Free Comp.	88.4% Dice	—	—	89.1% Dice
MSD Spleen	—	—	95.4% Dice	96.4% Dice
MSD Brain Tumor (All)	—	—	68.3% Dice	71.1% Dice

Ablation studies demonstrate that alternative decoder variants (NUP, PUP, MLA) with identical transformer encoders underperform the standard UNETR decoder (Brain All Dice: NUP 63.6%, PUP 66.8%, MLA 68.4% vs. UNETR 71.1%). Patch resolution has measurable effects, with $P=16^3$ yielding marginally higher Dice than $P=32^3$ (71.1% vs. 70.3%, respectively) (Hatamizadeh et al., 2021).

6. Computational Constraints and Limitations

UNETR’s multi-head self-attention has $O(N^2)$ computational and memory complexity, where $N$ is the number of 3D patches. This restricts the achievable spatial patch size and, consequently, the minimal granularity of global context modeling, particularly under GPU memory constraints. While bidirectional attention offers maximal spatial context modeling, it does not leverage any causal or hierarchical spatial priors, which a unidirectional or masked transformer formulation could potentially exploit. The authors also identify large-scale volumetric pre-training as an open avenue to improve generalization (Hatamizadeh et al., 2021).

7. Prospects for Unidirectional Transformers in 3D Segmentation

The report notes unidirectional transformer variants as a prospective direction to introduce spatial causality or reduce the model’s computational cost through structured masking. Approaches such as causal attention or spatial orderings may impose beneficial inductive biases for certain volumetric tasks. However, these strategies require further empirical validation in the 3D medical image segmentation context. A plausible implication is that unidirectional transformers could enable more efficient context aggregation or reflect anatomical directionality, but the field currently lacks comparative benchmarks for such variants in volumetric segmentation pipelines (Hatamizadeh et al., 2021).