U-Net and V-Net

U-Net passes fine spatial details directly from the encoder to the decoder via skip connections, enabling crisp pixel segmentation even with scarce data.

U-Net encoder-decoder architecture with lateral skip connections alongside 3D volumetric V-Net and Dice loss.

Why Does This Exist?

In biomedical imaging—such as microscopic cell tracking, MRI organ delineation, and CT tumor radiotherapy planning—annotated training datasets are extremely scarce. While ImageNet provides over one million labeled photographs, a hospital clinical trial might provide only 30 patient scans annotated by specialist radiologists.

Standard deep classification networks fail on this task for two reasons:

Severe spatial information loss: As successive pooling and strided convolutions downsample an image by factors of 16 or 32 to extract high-level semantic meaning, high-frequency spatial boundaries (cell membranes, capillary walls, tumor margins) are permanently destroyed. Standard decoders attempting to upsample coarse feature maps back to full resolution produce blurry, inaccurate blobs.
Extreme class imbalance: A small lesion or brain tumor often occupies fewer than 1% of the total voxels in a 3D medical volume. Standard cross-entropy loss gradients are completely dominated by background tissue voxels, causing networks to predict empty masks.

U-Net (Ronneberger et al., 2015) solved the localization bottleneck with an elegant symmetric U-shaped encoder-decoder architecture equipped with lateral skip connections that copy high-resolution spatial features directly from the contracting path to the expanding path. V-Net (Milletari et al., 2016) extended this to full 3D volumetric scans, introducing residual learning stages and the differentiable Soft Dice Loss to eliminate class imbalance issues.

Think of It Like This

A legal translation desk with original source sidebars

Imagine translating a dense ancient manuscript with complex handwriting into a modern printed book.

In a plain autoencoder without skip connections, a researcher reads each ancient paragraph, jots down brief conceptual summaries on note cards, throws away the original manuscript, and hands the cards to a typist. The typist understands the high-level themes, but when tasked with reconstructing the exact archaic punctuation, marginal notations, and ornate initials, they have to guess. The reconstructed book is structurally blurry.

U-Net works like a translation desk where the original illuminated manuscript pages are clipped directly opposite the typist's screen (the lateral skip connections). The typist reads the high-level semantic summary from the encoder, but whenever they need to determine an exact letter shape or boundary stroke, they look straight across at the original high-resolution page and copy the fine visual details directly onto the final draft.

How It Actually Works

The U-Net Encoder-Decoder Blueprint

The U-Net architecture forms a symmetric U-shape consisting of three distinct sections:

Contracting Path (Encoder):
- Repeated blocks of two $3 \times 3$ convolutions followed by ReLU activations.
- A $2 \times 2$ max pooling operation with stride 2 doubles the number of feature channels while halving the spatial height and width:
$(C, H, W) \xrightarrow{\text{Block}} (2C, H/2, W/2)$
Bottleneck:
- The lowest resolution stage (e.g., 512 or 1024 channels at $H/16, W/16$ ), capturing global semantic context.
Expanding Path (Decoder):
- At each stage, a $2 \times 2$ transposed convolution (up-convolution) halves channel depth and doubles spatial resolution:
$(2C, H, W) \xrightarrow{\text{UpConv}} (C, 2H, 2W)$
- Lateral Skip Concatenation: The upsampled feature map is concatenated along the channel dimension with the corresponding high-resolution feature map copied directly from the contracting path:
$X_{\text{decoder}} = \left[ X_{\text{skip\_encoder}} \,\|\, X_{\text{up}} \right] \in \mathbb{R}^{2C \times 2H \times 2W}$
- Two consecutive $3 \times 3$ convolutions fuse the combined channels into clean boundaries.
Final Layer: A $1 \times 1$ convolution projects the final multi-channel feature map to the desired number of output classes ( $K$ ), followed by Softmax or Sigmoid.

V-Net: Volumetric 3D Processing and Residual Blocks

V-Net adapts the U-Net philosophy to 3D volumetric medical images (such as $1 \times D \times H \times W$ MRI tensors):

3D Convolutions: Replaces 2D filters with $3 \times 3 \times 3$ volumetric kernels, processing entire anatomical volumes natively rather than as disjoint 2D slices.
Residual Learning: Adds additive identity shortcuts within each stage: $y = \mathcal{F}(x) + x$ , speeding up convergence.
Volumetric Downsampling: Replaces non-learnable max pooling with strided $2 \times 2 \times 2$ convolutions.

Soft Dice Loss for Extreme Foreground Imbalance

Standard binary cross-entropy treats every voxel equally. In an MRI volume where only 0.5% of voxels represent a prostate tumor, predicting all zeros yields 99.5% accuracy with zero diagnostic value.

V-Net introduced a fully differentiable formulation of the Dice similarity coefficient ( $D \in [0, 1]$ ):

D = \frac{2 \sum_{i=1}^{N} p_i g_i}{\sum_{i=1}^{N} p_i^2 + \sum_{i=1}^{N} g_i^2 + \epsilon}

where $p_i \in [0, 1]$ is the predicted probability at voxel $i$ , $g_i \in \{0, 1\}$ is the ground-truth binary label, and $\epsilon$ is a smoothing constant preventing division by zero.

The Dice Loss to be minimized is:

\mathcal{L}_{\text{Dice}} = 1 - D = 1 - \frac{2 \sum_{i=1}^{N} p_i g_i}{\sum_{i=1}^{N} p_i^2 + \sum_{i=1}^{N} g_i^2 + \epsilon}

Because the denominator is bounded by the volume of predicted and ground-truth foreground regions rather than the vast empty background, Dice loss maintains strong gradient magnitudes even for microscopic structures.

Worked Example

Let us compute the Soft Dice Loss on a small $4$ -pixel region with ground-truth binary mask $G$ and predicted probabilities $P$ :

Ground truth $G$ : $[1, 1, 0, 0]$ (2 positive pixels, 2 negative pixels)
Predicted probabilities $P$ : $[0.8, 0.7, 0.1, 0.2]$

Compute numerator term:

2 \sum_{i=1}^{4} p_i g_i = 2 \cdot (0.8 \times 1 + 0.7 \times 1 + 0.1 \times 0 + 0.2 \times 0) = 2 \cdot (0.8 + 0.7) = 2 \cdot 1.5 = 3.0

Compute denominator terms:

\sum_{i=1}^{4} p_i^2 = 0.8^2 + 0.7^2 + 0.1^2 + 0.2^2 = 0.64 + 0.49 + 0.01 + 0.04 = 1.18

\sum_{i=1}^{4} g_i^2 = 1^2 + 1^2 + 0^2 + 0^2 = 1 + 1 + 0 + 0 = 2.0

\text{Denominator} = \sum p_i^2 + \sum g_i^2 = 1.18 + 2.0 = 3.18

Compute Soft Dice coefficient and loss ( $\epsilon = 0$ ):

D = \frac{3.0}{3.18} \approx 0.9434

\mathcal{L}_{\text{Dice}} = 1 - D = 1 - 0.9434 = 0.0566

If the model had predicted near-zero probabilities for the foreground ( $[0.05, 0.05, 0.0, 0.0]$ ), the numerator would collapse to $2(0.1) = 0.2$ while the denominator remained $\approx 2.0$ , yielding $D \approx 0.1$ and $\mathcal{L}_{\text{Dice}} \approx 0.90$ —heavily penalizing the omission regardless of background size.

Code

1import torch2import torch.nn as nn3
4class DiceLoss(nn.Module):5    """Differentiable Soft Dice Loss for binary semantic segmentation."""6    def __init__(self, smooth: float = 1e-5) -> None:7        super().__init__()8        self.smooth = smooth9
10    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:11        probs = torch.sigmoid(logits)12        probs_flat = probs.view(-1)13        targets_flat = targets.view(-1)14
15        intersection = 2.0 * (probs_flat * targets_flat).sum()16        denominator = (probs_flat ** 2).sum() + (targets_flat ** 2).sum()17
18        dice_score = (intersection + self.smooth) / (denominator + self.smooth)19        return 1.0 - dice_score20
21class SimpleUNetStage(nn.Module):22    """A minimal 1-level U-Net decoder block verifying skip concatenation."""23    def __init__(self, in_ch: int, skip_ch: int, out_ch: int) -> None:24        super().__init__()25        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)26        self.conv = nn.Sequential(27            nn.Conv2d(in_ch // 2 + skip_ch, out_ch, kernel_size=3, padding=1),28            nn.ReLU(inplace=True),29            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),30            nn.ReLU(inplace=True)31        )32
33    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:34        x_up = self.up(x)35        concat = torch.cat([skip, x_up], dim=1)36        return self.conv(concat)37
38# Test verification39x = torch.randn(2, 256, 16, 16)      # Bottleneck features40skip = torch.randn(2, 128, 32, 32)   # High-resolution skip from encoder41decoder_block = SimpleUNetStage(in_ch=256, skip_ch=128, out_ch=128)42
43out = decoder_block(x, skip)44print(out.shape)45# -> torch.Size([2, 128, 32, 32])46
47# Verify Dice Loss with our worked example values48loss_fn = DiceLoss(smooth=0.0)49logits = torch.tensor([1.386, 0.847, -2.197, -1.386]) # sigmoid -> [0.8, 0.7, 0.1, 0.2]50targets = torch.tensor([1.0, 1.0, 0.0, 0.0])51print(round(loss_fn(logits, targets).item(), 4))52# -> 0.0566

Watch Out For

Spatial dimension mismatches during skip concatenation and unpadded cropping

In the original 2015 U-Net paper, convolutions used unpadded "valid" boundaries ( $P=0$ ). As a result, the contracting feature maps were slightly larger than the upsampled decoder maps (for instance, $64 \times 64$ in the encoder versus $56 \times 56$ in the decoder), necessitating central cropping of the skip features before concatenation.

In modern PyTorch implementations, failing to account for this either causes a shape crash during torch.cat or requires complex cropping. Always use padded convolutions (padding=1 on $3 \times 3$ filters) and input image dimensions that are clean powers of 2 (divisible by $2^4 = 16$ or $2^5 = 32$ ). This guarantees that each upsampling layer reproduces the exact spatial dimensions of its corresponding encoder stage without cropping.

The Quick Version

U-Net uses a symmetric encoder-decoder structure with lateral skip connections to copy fine spatial features directly to the decoder.
Skip connections resolve the fundamental trade-off between semantic abstraction and pixel-precise boundary localization.
V-Net expands U-Net to 3D volumetric images using volumetric $3 \times 3 \times 3$ kernels and internal residual skip connections.
Soft Dice Loss replaces cross-entropy to prevent gradient collapse when segmenting tiny structures against overwhelming background tissue.
Using padded convolutions (padding=1) preserves identical spatial dimensions across corresponding encoder and decoder levels.