Tyche: Stochastic In-Context Learning for Medical Image Segmentation

TL;DR

Tyche addresses the dual challenges of needing task-specific retraining for medical image segmentation and the lack of multiple plausible segmentations by introducing stochastic in-context segmentation. It comprises two variants, Tyche-TS (train-time stochasticity) and Tyche-IS (inference-time stochasticity), and employs SetBlock-CrossBlock interactions plus a best-candidate Dice loss to produce diverse segmentation candidates without retraining. Across twenty unseen tasks and multiple data modalities, Tyche outperforms in-context and interactive baselines and matches or approaches specialized stochastic methods, while offering faster inference than some baselines. The framework enables robust uncertainty quantification and practical deployment in clinical settings by modeling annotator disagreement without requiring retraining for new tasks.

Abstract

Existing learning-based solutions to medical image segmentation have two important shortcomings. First, for most new segmentation task, a new model has to be trained or fine-tuned. This requires extensive resources and machine learning expertise, and is therefore often infeasible for medical researchers and clinicians. Second, most existing segmentation methods produce a single deterministic segmentation mask for a given image. In practice however, there is often considerable uncertainty about what constitutes the correct segmentation, and different expert annotators will often segment the same image differently. We tackle both of these problems with Tyche, a model that uses a context set to generate stochastic predictions for previously unseen tasks without the need to retrain. Tyche differs from other in-context segmentation methods in two important ways. (1) We introduce a novel convolution block architecture that enables interactions among predictions. (2) We introduce in-context test-time augmentation, a new mechanism to provide prediction stochasticity. When combined with appropriate model design and loss functions, Tyche can predict a set of plausible diverse segmentation candidates for new or unseen medical images and segmentation tasks without the need to retrain.

Figures (43)

• Figure 1: Tyche: the first in-context stochastic segmentation framework. Human annotators (top) can handle a wide variety of tasks, and different annotators often produce differing segmentations. Existing automated methods (middle) are typically task-specific and provide only one segmentation per image. Tyche (bottom) can capture the disagreement among annotators across many modalities and anatomies without retraining or fine-tuning.
• Figure 2: Tyche Model Schematic. The target $x^t$, context set ${(x^t_j, y^t_j)}_{j=1}^S$, and noise images $\{z_k\}_{k=1}^K$ are inputs to the network. The architecture employs UNet-like levels, but uses SetBlocks that enable interactions between the context set and the target segmentation candidates.
• Figure 3: CrossBlock Mechanism The CrossBlock involves interactions between a single feature and a set of features and outputs new feature for the target and new features for each.
• Figure 4: SetBlock Mechanism. SetBlock enables interactions between the set of features from the context set and the set of features from the prediction candidates. It outputs two sets of features, one for the context and one for the prediction candidates.
• Figure 5: Visualization of predictions for three different samples, 1 per row. Left: LIDC-IDRI. Right: Hippocampus dataset. The leftmost columns are raters' annotations. The 4 last columns are model predictions. Tyche provides a set of prediction that is diverse and matches the raters, for tasks unseen at training time.