PaperlayerView exact revision on arXiv

Source-linked AI summary

UniverSeg: Universal Medical Image Segmentation

Victor Ion Butoi, Jose Javier Gonzalez Ortiz, Tianyu Ma, Mert R. Sabuncu, John Guttag, Adrian V. Dalca

arXiv:2304.06131v1cs.CVcs.LG

TL;DR

Medical segmentation models often require retraining for unseen anatomies, modalities, or labels, creating burdens for researchers without machine-learning resources. UniverSeg learns a single task-agnostic model that uses labeled support examples and CrossBlocks to segment new tasks without retraining, substantially outperforming few-shot methods across held-out datasets. Its demonstrated scope is 2D data with single labels.

  • Problem

    Domain shift and task diversity make unseen medical segmentation difficult, while adapting models requires expertise, computational resources, and labor.

  • Method

    UniverSeg is a single task-agnostic segmentation model that uses labeled examples and CrossBlock interactions to specify and solve new tasks without retraining.

  • Results

    UniverSeg substantially outperforms existing few-shot methods across all held-out datasets, with Dice improvements of 7.3 to 34.9 over the highest-performing baseline.

  • Takeaways & Limitations

    Performance depends strongly on task diversity during training and support-set diversity during inference, while variably sized support sets provide user flexibility.

  • Takeaways & Limitations

    The study focuses on 2D data and single labels, leaving 3D volumes and multi-label maps for future extensions.

Abstract

from arXiv · show

While deep learning models have become the predominant method for medical image segmentation, they are typically not capable of generalizing to unseen segmentation tasks involving new anatomies, image modalities, or labels. Given a new segmentation task, researchers generally have to train or fine-tune models, which is time-consuming and poses a substantial barrier for clinical researchers, who often lack the resources and expertise to train neural networks. We present UniverSeg, a method for solving unseen medical segmentation tasks without additional training. Given a query image and example set of image-label pairs that define a new segmentation task, UniverSeg employs a new Cross-Block mechanism to produce accurate segmentation maps without the need for additional training. To achieve generalization to new tasks, we have gathered and standardized a collection of 53 open-access medical segmentation datasets with over 22,000 scans, which we refer to as MegaMedical. We used this collection to train UniverSeg on a diverse set of anatomies and imaging modalities. We demonstrate that UniverSeg substantially outperforms several related methods on unseen tasks, and thoroughly analyze and draw insights about important aspects of the proposed system. The UniverSeg source code and model weights are freely available at https://universeg.csail.mit.edu

Adrian V. Dalca MIT CSAIL & MGH, HMS

The supplied section includes a contribution marker and a fragment referring to anatomies and imaging modalities.

  • The text refers to anatomies and imaging modalities.
  • An asterisk denotes equal contribution.

1. Introduction

Medical image segmentation spans diverse modalities, domains, and labels, while domain shift makes unseen tasks difficult and training task-specific models burdensome. UniverSeg addresses this by using labeled examples to specify new tasks without retraining, and reports strong generalization across held-out tasks.

  • Medical segmentation spans diverse imaging modalities, biomedical domains, and target labels.
  • Domain shift can make models perform poorly on out-of-distribution examples, while new-task training requires expertise, compute, and labor.
  • UniverSeg uses one general model to segment new tasks from labeled examples without retraining.
  • The framework transfers information from an example set to a new image through a flexible CrossBlock mechanism.
  • UniverSeg substantially outperforms several models across held-out tasks involving unseen anatomies and approaches fully supervised networks.

2. Related Works

Prior medical segmentation methods commonly train task-specific or predetermined models, while transfer and few-shot approaches often retain domain or retraining limitations. UniverSeg instead dynamically specifies broad new tasks during inference without fine-tuning.

  • Supervised medical segmentation methods typically train models from scratch for each new problem.
  • nnUNet automates some design decisions but still requires substantial training overhead.
  • Multi-task learning uses tasks predetermined during design, limiting networks to tasks presented during training.
  • Transfer learning still involves substantial training for each new task and can be unhelpful across natural and medical image domains.
  • Few-shot segmentation methods often target new classes within particular input domains rather than generalizing across anatomies or imaging modalities.
  • UniverSeg avoids fine-tuning while targeting broad generalization across anatomies, modalities, and previously unseen datasets.

3. UniverSeg Method

UniverSeg learns a universal segmentation function that uses a query image and task-defining support image-label pairs, rather than training a separate model for each task. Its CrossBlock modules exchange information between query and support representations within an encoder-decoder network, while augmentation and support-set ensembling improve generalization and robustness.

  • The universal function predicts a query label map from the query image and a support set of example image-label pairs.
  • CrossBlock: Cross-convolution combines query and support feature maps, and shared weights make the operation permutation invariant to support-set ordering.
  • CrossBlock: CrossBlock updates query and support representations through their average representation, enabling interactions and variably sized support sets.
  • Network: The network uses CrossBlocks in a UNet-like encoder-decoder with residual connections, spatial downsampling and upsampling, skip concatenations, and a final 1x1 convolution.
  • Training: Training averages a segmentation loss over tasks and query-support examples, with the query excluded from the support set.
  • Augmentation and inference: Task and in-task augmentations diversify training tasks and examples, while averaging predictions from independently sampled support sets reduces dependence on support-set choice.

4. MegaMedical Dataset

MegaMedical is a standardized, diverse training collection assembled to support generalization across medical segmentation tasks. It combines 53 datasets across domains and modalities, transforms heterogeneous data into binary 2D tasks, and adds synthetic tasks for greater diversity.

  • MegaMedical contains 53 open-access datasets spanning 26 medical domains and 16 imaging modalities.
  • The collection standardizes heterogeneous images and label maps and expands task diversity with synthetic segmentation tasks.
  • Processed MegaMedical data cannot be re-released because of dataset agreements, so the authors provide code to reconstruct it from source datasets.
  • Datasets: The datasets cover varied biomedical domains, including eyes, lungs, spine vertebrae, white blood cells, abdominal structures, and brain.
  • Medical Image Task Creation: For 3D datasets, the pipeline extracts mid-slices along major axes; each modality becomes a task, and each available label becomes a binary segmentation task.

5. Experiments

Experiments evaluate UniverSeg on six held-out medical datasets spanning seen and unseen anatomies, comparing it with few-shot baselines and task-specific nnUNets. Results show stronger few-shot performance, competitiveness with supervised models on some datasets, and benefits from task diversity, larger support sets, and ensembling.

  • Experimental setup: Evaluation used six held-out datasets, including three anatomies represented during training and three anatomies absent from the rest of MegaMedical.The datasets cover heart, retinal blood vessels, mandible, vertebrae, and white blood cells.
  • Method comparison: UniverSeg significantly outperformed all few-shot methods across every held-out dataset.The comparison used the best-performing variant of each few-shot baseline, including additions to their training and inference procedures.
  • Method comparison: Dice improvements over the highest-performing baseline ranged from 7.3 to 34.9 across datasets.Qualitative predictions also showed clearer improvements in the example segmentations.
  • Method comparison: UniverSeg was competitive with task-specific supervised networks on datasets such as PanDental and WBC, despite those networks being extensively trained for each task.Its segmentations more closely matched supervised predictions than other few-shot methods, especially on SpineWeb and STARE.
  • Task generalization analysis: More training tasks improved held-out performance, while dataset selection also mattered substantially: 10%-dataset subsets differed by 17.3 Dice points.The strongest subset covered a broad range of anatomies, including heart, abdomen, brain, and eyes.
  • Inference analysis: Average held-out Dice increased from 53.7 to 69.9 as support size grew from 1 to 16, reaching a maximum of 71 at support size 64.Ensembling consistently improved performance, with gains of 2.4-3.1 Dice points for support sizes below 16.

6. Discussion and Conclusion

UniverSeg learns a single task-agnostic medical-image segmentation model that uses diverse training tasks and support examples to generalize without retraining. The work focuses on 2D, single-label, single-modality segmentation and identifies support diversity and size as important factors.

  • Core contribution: UniverSeg uses a single task-agnostic model to segment unseen medical tasks through cross-convolution between query and support representations.The model interacts query and support representations at different scales.
  • Scope and limitations: The experiments focus on single-label, single-modality, 2D segmentation, while future work targets 2.5D or 3D models and multi-label maps.The stated limitation is an explicit scope boundary rather than a claim about those extensions.
  • Task construction: Medical tasks are represented by dataset, modality, label, and axis tuples after splitting multi-class datasets and 3D datasets into single-label and axis-specific tasks.This construction defines the task units used by the framework.
  • Training design: During training, hierarchical sampling and augmentation are used to increase task and data heterogeneity across datasets, modalities, axes, and labels.The sampling procedure is designed to avoid overrepresenting domains with many tasks.

C. Data Augmentation

UniverSeg uses task-level and in-task augmentations to expose the model to more task variation and more examples within each task. The augmentation set includes geometric, intensity, label, blur, noise, and sharpness transformations, including a Sobel-based operation for thin structures.

  • Augmentation strategy: Task and in-task augmentation increase effective task diversity and the number of examples seen for each task during training.The two augmentation types are separated according to whether they vary tasks or individual support/query examples.
  • Task transformations: Task augmentations can also flip image intensities or reverse foreground and background labels, altering the task specification without changing the underlying image geometry.Intensity inversion excludes label maps, whereas label flipping acts on segmentation maps.
  • Task augmentation: Sobel-Edge Label applies a Sobel filter to label maps to create additional tasks containing thin segmentation structures.The x- and y-direction responses are combined through their squared norm.
  • Geometric transformations: Geometric augmentations include elastic warps, affine shifts, and horizontal or vertical flips applied consistently across relevant support entries and queries.Affine parameters control rotation, translation, and scale, while elastic warps control deformation strength and smoothness.
  • Image transformations: Appearance augmentations include Gaussian blur, Gaussian noise, brightness and contrast changes, and sharpness changes applied to images.Their parameters control kernel size, variance, brightness, contrast, or sharpness, depending on the operation.

D. Synthetic Tasks

The training procedure introduces synthetic segmentation tasks to increase task diversity. Each task is generated from random label shapes, spatial deformations, and synthesized intensity images.

  • Synthetic task creation: 1,000 synthetic tasks are generated to increase training task diversity and improve held-out performance.The paper reports improvement in held-out performance after introducing synthetic tasks.
  • Synthetic task creation: Each synthetic task begins with a label map containing 16 random shapes representing regions of interest.The label map is the basis for subsequent spatial and image synthesis.
  • Generation process: For each task, the label map is given 100 random smooth spatial variations before intensity images are synthesized.The generation process proceeds from labels to spatial variations and then to images.
  • Synthetic diversity: The generated tasks vary in shapes, textures, and label shapes, providing heterogeneous synthetic segmentation problems.Examples visualize 10 of the 1,000 generated tasks.

E.1. Main Results

UniverSeg is compared with few-shot baselines and supervised upper bounds across held-out datasets, with analyses of augmentation, support size, support-pool size, ensembling, and training-data scale. Results consistently favor broader support information, although task-diversity gains are not universal.

  • Method comparison: UniverSeg is evaluated against few-shot baselines and supervised nnUNet upper bounds using held-out test Dice scores.The method comparison reports per-dataset test performance for the baselines, UniverSeg, and nnUNet.
  • Training strategies: Increasing task diversity generally helps, but it decreases performance on the STARE eye-vessel segmentation task.The overall trend holds for most datasets, with STARE as the reported exception.
  • Available inference data: Using more available support examples during inference consistently improves performance and reduces variance across random support subsets.The trend is reported across all evaluated tasks.
  • Support-set ensembling: For support-set ensembling, N > 1 yields substantial improvements and lower variance, whereas N = 1 produces only small gains that decline as K increases.The comparison separates the effect of support-set size from the number of ensembled predictions.
  • Train–inference support sizes: Using inference support sizes larger than those seen in training improves results for models trained with N ≥2.This pattern is interpreted as evidence that the model benefits from interacting with larger support sets.

F. Additional Visualizations

Additional visualizations show that UniverSeg adapts to diverse support-set labels and image transformations, while performance improves with larger support sets. Soft predictions can preserve thin-structure segmentation that thresholding may obscure.

  • Soft Predictions: UniverSeg’s soft predictions can capture thin structures that are obscured when predictions are thresholded into binary masks.The visualization uses STARE as an example of thin-structure performance that is lost after thresholding.
  • WBC Visualizations: For WBC cytoplasm segmentation, UniverSeg closely matches ground-truth masks despite substantial variability among support-set examples.The WBC visualizations evaluate adaptation using held-out predictions and cytoplasm labels.
  • WBC Visualizations: Changing the support-set labels from cytoplasm to nucleus causes UniverSeg to adapt its predictions to the nucleus label.The query images remain the same while the supplied labels change.
  • WBC Visualizations: UniverSeg correctly segments labels after inverting both query and support images, demonstrating invariance to this image transformation.The inversion is implemented as 1 − x for the image data.
  • Support Sets: Segmentation mask quality substantially improves as the support set grows, while ensembling reduces prediction variability across support sets.The multi-label visualization combines independently predicted labels using a softmax operation.
Loading 2304.06131v1…