UniverSeg: Universal Medical Image Segmentation

In plain words

Usually, to segment a new kind of medical scan, such as a new organ or a new imaging device, someone has to train or fine-tune a neural network for that one job. Clinical researchers often lack the time and skills to do this.

UniverSeg tries a different route. One network is trained once on a large mix of segmentation jobs. At use time you give it the image you want to label, plus a handful of example images with their correct outlines, and it labels the new image by comparing it against those examples. Nothing is retrained. The network passes information back and forth between the query image and the examples at every layer, which is what lets it pick up the new task from the examples alone.

TL;DR

• UniverSeg maps a query image and a support set of labelled examples to a segmentation, with no gradient updates for a new task.
• The core block, CrossBlock, mixes query and support features through a shared convolution and a mean over the support set, so the support set can be any size.
• It is trained on MegaMedical, a collection of 53 medical datasets, plus task-level augmentations and 1,000 synthetic tasks.
• Author-reported: mean Dice 71.8 over six unseen held-out datasets, above the few-shot baselines (41.8 to 50.1) and below per-task nnUNet (84.4) tab. 1.
• Check the paper-vs-code section first: the released code differs from the paper in activation placement, network depth and downsampling. The training code is not released.

Problem

Standard segmentation networks learn one task, so each new anatomy, modality or label set needs new training [abstract]. Few-shot methods from the natural image literature were built for small label sets with similar tasks, and the authors argue they do not transfer to medical tasks that differ widely §5.1, Few-Shot Baselines. Per-task networks (nnUNet) reach high accuracy but need training data and compute for every task §5.1.

Method

UniverSeg is a single function f(x, S) that takes a query image x and a support set S of image-label pairs defining the task, and outputs a label map for x §3. The support set can have any size. Each pair is stacked channel-wise (image and label) before entering the network §3.1.

CrossBlock (the building block). A cross-convolution pairs the query feature map u with each support feature map v_i and runs one convolution on their concatenation. The same weights are reused for every support entry, so the result does not depend on the order of the support set §3.1, eq. 1. A CrossBlock then (a) averages the pairwise outputs over the support set to update the query, and (b) runs a separate convolution on each pairwise output to update each support entry §3.1, eq. 2.

Network. Encoder-decoder with skip connections, UNet-like §3.1, fig. 3. Each encoder level applies a CrossBlock and then halves the resolution of both query and support. Each decoder level doubles the resolution, concatenates the matching encoder skip, then applies a CrossBlock. A final 1x1 convolution maps the query representation to the prediction §3.1. Width is 64 channels per stage with LeakyReLU after each convolution §5.1. Stage counts and resampling differ between paper and released code; see Paper vs code.

Cross-convolution. For a query u and support entries V = {v_i}:

\text{CrossConv}(u, V; \theta_z) = \{z_i\}, \quad z_i = \text{Conv}(u \,\|\, v_i; \theta_z) §3.1, eq. 1

CrossBlock. With A a non-linearity:

z_i = A(\text{CrossConv}(u, v_i; \theta_z)), \quad u' = \tfrac{1}{n} \textstyle\sum_i z_i, \quad v'_i = A(\text{Conv}(z_i; \theta_v)) §3.1, eq. 2

Training objective. Expected segmentation loss over training tasks t, query pairs and support sets, with the query outside its support set §3.2, eq. 3:

\mathcal{L}(\theta; \mathcal{T}) = \mathbb{E}_{t} \mathbb{E}_{(x^t, y^t), S^t} \left[ \mathcal{L}_{\text{seg}}(f_\theta(x^t, S^t), y^t) \right] §3.2, eq. 3

Training loop. Algorithm 1 samples a task, a query and a support set, applies augmentations, then takes one gradient step §3.2, alg. 1.

# UniverSeg training step  [alg. 1] [§3.2]
# Input:  task set T; augmentations Aug_t and Aug_T; learning rate eta
# Output: updated parameters theta

def training_step(theta, tasks):
    # Sample a task; datasets are chosen uniformly first  [app. B]
    task = sample_task_hierarchically(tasks)
    # Sample one query image-label pair from the task  [alg. 1]
    query_image, query_label = sample_pair(task)
    # Sample a support set of 64 pairs, with replacement  [§5.1]
    support = sample_pairs_with_replacement(task, size=64)
    # Augment the query and each support entry independently  [§3.2]
    query_image, query_label = aug_in_task(query_image, query_label)
    support = [aug_in_task(x, y) for x, y in support]
    # Apply the same task-changing transform to all images and labels  [§3.2]
    query_image, query_label, support = aug_task(
        query_image, query_label, support
    )
    # Predict the query label map from the support set  [§3.1]
    prediction = forward(query_image, support)
    # Soft Dice loss; no smoothing constant is stated  [§3.2] [app. B]
    loss = soft_dice_loss(prediction, query_label)
    # Gradient step; Algorithm 1 writes plain SGD, §5.1 says Adam  [alg. 1]
    theta = update(theta, loss, learning_rate=1e-4)
    return theta

# Forward pass f_theta(query, support)  [§3.1]
# query: (B, 1, H, W); support images and labels: (B, S, 1, H, W) each

def forward(query, support_images, support_labels):
    # Stack each support image with its label along channels  [§3.1]
    support = concat([support_images, support_labels], channel_axis)
    # Query gets a group axis of size 1  # (B, 1, 1, H, W)
    query_feats = add_group_axis(query)
    skips = []
    for stage, block in enumerate(encoder_blocks):
        # CrossBlock on query and support  [§3.1] [eq. 2]
        query_feats, support = block(query_feats, support)
        if stage == len(encoder_blocks) - 1:
            break
        # Keep both representations for the skip connection  [§3.1]
        skips.append((query_feats, support))
        # Halve resolution of both; paper says bilinear  [§5.1]
        query_feats = downsample(query_feats)
        support = downsample(support)
    for block in decoder_blocks:
        skip_query, skip_support = skips.pop()
        # Double resolution, then concatenate the matching skip  [§3.1]
        query_feats = concat([upsample(query_feats), skip_query])
        support = concat([upsample(support), skip_support])
        # CrossBlock on the merged representations  [§3.1]
        query_feats, support = block(query_feats, support)
    # A single 1x1 convolution gives the prediction  [§3.1]
    return out_conv(drop_group_axis(query_feats))

# CrossBlock(u, V)  [§3.1] [eq. 2]
# u: query features (B, 1, C, H, W); V: support features (B, S, C, H, W)

def cross_block(query_feats, support_feats):
    # Pairwise cross-conv of the query with each support entry  [eq. 1]
    pairs = []
    for support_entry in support_feats:
        # Non-linearity A on each cross-conv output  [eq. 2]
        pairs.append(activation(cross_conv(query_feats, support_entry)))
    # The query takes the mean over the support set  [eq. 2]
    query_out = mean(pairs)
    # Each support entry gets its own conv after the pairwise step  [eq. 2]
    support_out = [activation(conv(z)) for z in pairs]
    return query_out, support_out

Inference. Prediction depends on which support set is drawn, so the paper averages K predictions, each from an independently drawn support set, pixel by pixel §3.3. The evaluation uses K = 5 draws per test subject §5.1.

# Ensemble inference  [§3.3] [§5.1]
# Paper evaluation: K = 5 draws; support drawn from the held-out support split

def predict(query, support_pool, K=5, N=64):
    predictions = []
    for draw in range(K):
        # Draw N pairs at random; with or without replacement is not stated
        support = sample_pairs(support_pool, size=N)
        predictions.append(forward(query, support))
    # Pixel-wise average of the K predictions  [§3.3]
    return mean(predictions)

Data and augmentation. Each dataset is split per subject into support, development and test parts, 60/20/20 §5.1. Every image is a 2D mid-slice at 128x128, normalized to [0, 1] app. A §4. Augmentation is applied at two levels §3.2, app. C:

• In-task: applied to the query and to each support entry independently (affine, elastic warp, noise, blur, sharpness, brightness and contrast).
• Task: applied to all query and support images and labels together, changing the task itself (intensity or label flips, horizontal or vertical flips, Sobel-edge labels, plus the same geometric and photometric ops).

Synthetic tasks: 1,000 extra tasks. Each has 16 random shapes as labels, 100 smooth deformations of that label map as subjects, and textured intensities app. D. The generator follows SynthMorph app. D.

Dataset. MegaMedical is built from 53 open-access datasets covering 26 domains and 16 modalities §4. Held-out datasets: ACDC, SCD and STARE (anatomies seen in training) and PanDental, SpineWeb and WBC (unseen anatomies) §5.1. Labels that come as several classes are split into one binary task per label, so the network has a single output channel §4.

Hyperparameters

Augmentation parameters app. C; p = probability per iteration:

Results

Author-reported. Setup: six held-out datasets, 2D mid-slices, support size 64 during training, five random support sets per test subject averaged, Dice score 0 to 100 averaged over labels and then over tasks §5.1. Hardware for the runtime column is not stated.

Main comparison tab. 1:

Per dataset, mean Dice app. E:

Training-strategy ablation tab. 2:

Support size and ensembling (text and figures only; no table):

• Average held-out Dice rises from 53.7 at N = 1 to 69.9 at N = 16, then flattens to 71 at N = 64 §5.3, fig. 7.
• Ensembling helps at every N, with gains of 2.4 to 3.1 Dice for N < 16 §5.3, fig. 7.
• Non-ensembled N = 64 beats heavily ensembled N = 2, 4, 8 with K = 64 §5.3, fig. 9.
• For WBC and PanDental, variance across random support sets shrinks as N grows §5.3, fig. 8.

Number of training tasks (fig. 6): more tasks improve held-out Dice. At 10 % of datasets, the best and worst subsets differ by 17.3 Dice §5.3.

Paper vs code

Compared against the official repository at commit 833a0c34c65e38d675e21bd48ddec6797cc03259 (JJGO/UniverSeg, Apache-2.0, owned by an author). The repository contains the model, the CrossConv layer and example data loaders. It does not contain the training loop, loss, augmentations, dataset construction or ensembling.

• Non-linearity on the cross-conv output. Paper: z_i = A(CrossConv(...)) §3.1, eq. 2. Code sets the cross-op non-linearity to none model.py, and CrossOp only applies one when set model.py model.py. In code, the query update is a plain mean of linear cross-convolutions model.py. The LeakyReLU appears only after the support convolution model.py.
• Encoder and decoder depth. Paper: 5 encoder stages and 4 decoder stages §5.1. Code: encoder_blocks=[64, 64, 64, 64], so 4 encoder stages, and the decoder is encoder_blocks[-2::-1], so 3 stages model.py model.py.
• Resampling. Paper: bilinear for both down- and up-sampling §5.1. Code: max-pool 2x2 for downsampling model.py; bilinear for upsampling model.py.
• Output activation. Paper: a 1x1 convolution to the prediction §3.1; the loss is soft Dice §3.2. Code: linear output with no sigmoid or softmax model.py. Whether a sigmoid is applied before the soft Dice loss in training is not stated, and the training code is absent.
• Update rule. Algorithm 1 writes plain SGD alg. 1; §5.1 and app. B say Adam app. B. The repository has no optimizer.
• Ensembling. Paper: K-draw averaging is part of inference §3.3. Repository: a single forward pass only README.md. Averaging must be written by the user.
• Consistent parts. Support-set concatenation of image and label §3.1 matches the code model.py. The query update is the mean over the support set §3.1 and the code takes mean(dim=1) model.py. Input is 128x128 with min-max normalization to [0, 1] README.md and app. A.

Pitfalls

• Pick the code or the paper, and say which. The code is a 4/3-stage network with max-pool downsampling and no activation on the cross path. The paper describes a 5/4-stage network with bilinear resampling and an activation on the cross path. Which of the two the released weights were trained with is not verified here: the weights file is named universeg_v1_nf64_ss64_STA.pt model.py, and the card did not load it. Check the weights against the code before relying on either description.
• The soft Dice loss is written as a score. Eq. app. B is 2·sum(y·ŷ) / (sum y² + sum ŷ²), with no minus sign and no smoothing term. Minimizing it as written would reward low overlap. Implement the loss as 1 minus this ratio, and add a small constant if an empty target can occur. Neither choice is stated in the paper.
• The Dice metric formula uses squared sums. app. B writes |y|² + |ŷ|² in the denominator. This equals the usual sum for binary masks, which is what the evaluation uses, but not for soft predictions.
• Text and table disagree on the Dice gains. The text says UniverSeg beats the best baseline by 7.3 to 34.9 Dice §5.2. The per-dataset table app. E gives gaps of 6.4 (PanDental) to 33.0 (SpineWeb) over the best baseline. The table is the one the card uses; the text range could not be reproduced from it.
• Support sampling at test time is ambiguous. Training samples support pairs with replacement §5.1. For evaluation the paper says only "randomly drawn" support sets §5.1. Replacement, and whether the query is excluded from the support pool at test time, are not stated.
• In-task augmentation probability is ambiguous. App. C says p controls whether all support entries are augmented, and applies to them individually. Whether one draw applies to all entries or each entry gets its own draw is unclear from the text.
• Cost grows with support size. Cross-convolution runs for every query-support pair and tiles the inputs along the batch dimension cross_conv.py. The code's own docstring says runtime and memory are O(Sx · Sy) cross_conv.py. Memory grows with N, so larger support sets cost more per prediction. Table 1 runtime is at one support size; the support size used is not stated there.
• Each run predicts one binary label. The head has one output channel model.py. Multi-label tasks need one run per label §4.
• 2D only. Every task is a mid-slice from a 3D volume §4 app. A. The reported numbers do not cover full-volume segmentation.
• Training budget is not reported. Batch size 1 app. B, learning rate 1e-4 app. B, training until development Dice stops improving §5.1. Total steps, hardware and wall-clock time are not stated.
• Baselines were tuned by the authors. Baselines were trained with support size 64 and the UniverSeg augmentation pipeline, and the best variant was chosen §5.1. The comparison is therefore against a tuned variant of each baseline, not the original methods.
• Ablation gains are on one held-out average. The main ablation is the mean over six datasets tab. 2. Per-dataset ablations are in the appendix app. E and are not summarized here.
• Weights have a different license. The code is Apache-2.0. The pretrained weights are under the OpenRAIL++-M license, which the repository README says limits use to research LICENSE-model README.md. This card's CC BY 4.0 applies to the paper text only. Do not redistribute the weights from this card.
• Data is not released. The authors say they cannot re-release the processed MegaMedical data because of dataset agreements. They offer data-processing code instead §4. Reproducing the full training run needs the source datasets.

Lineage

• Builds on: UNet-style encoder-decoder with skip connections §3.1; synthetic shapes adapted from SynthMorph app. D. No arXiv ids are linked here.
• Compared against: SE-net, ALPNet and PANet as few-shot baselines §5.1; nnUNet as the supervised per-task upper bound §5.1.
• Superseded by: not stated in this wiki's content.