Medical Vision Generalist: Unifying Medical Imaging Tasks in Context

## Abstract

Abstract This study presents Medical Vision Generalist (MVG), the first foundation model capable of handling various medical imaging tasks—such as cross-modal synthesis, image segmentation, denoising, and inpainting—within a unified image-to-image generation framework.
Specifically, MVG employs an in-context generation strategy that standardizes the handling of inputs and outputs as images. By treating these tasks as an image generation process conditioned on prompt image-label pairs and input images, this approach enables a flexible unification of various tasks, even those spanning different modalities and datasets.
To capitalize on both local and global context, we design a hybrid method combining masked image modeling with autoregressive training for conditional image generation. This hybrid approach yields the most robust performance across all involved medical imaging tasks. To rigorously evaluate MVG’s capabilities, we curated the first comprehensive generalist medical vision benchmark, comprising 13 datasets and spanning four imaging modalities (CT, MRI, X-ray, and micro-ultrasound).
Our results consistently establish MVG’s superior performance, outperforming existing vision generalists, such as Painter and LVM.
Furthermore, MVG exhibits strong scalability, with its performance demonstrably improving when trained on a more diverse set of tasks, and can be effectively adapted to unseen datasets with only minimal task-specific samples.
The code is available at https://github.com/OliverRensu/MVG .

## 1 Introduction

The precise interpretation of medical images is imperative for timely disease detection, diagnosis, and treatment [, ].
Deep-learning based models have emerged as powerful tools in medical image analysis, tackling various challenges spanning from segmenting specific anatomical structures [, , ], localizing single organ diseases [, , , , , , ], to cross-modality image synthesis on brain MRI [, , , ].
However, these models, often referred to as specialist models, are typically customized for specific tasks, modalities, or anatomical regions. While this specialization often results in exceptional performance in certain contexts, it can lead to a severe performance drop when applied to new tasks or when tasked with training multi-domain data.

To address this challenge, recently, there has been a partial shift of research focus in developing generalist medical AI models [, ], which necessitate only a single training phase but are capable of wide application across a diverse array of medical tasks. Specifically, these generalist frameworks unify input and output spaces, allowing for straightforward adaptation to various tasks through the use of user-provided prompts.
While existing generalist medical AI models like MedSAM have demonstrated impressive performance [, , ],
their applicability in medical visual tasks remains limited.
A unified, truly generalist vision model capable of addressing a vast array of medical imaging tasks remains a critical missing piece in the current medical research landscape.

Figure: Figure 1: Medical Vision Generalist enables a single model be capable of performing four types of medical vision tasks on images in four medical imaging modalities of three major body regions.
Refer to caption: x1.png

Motivated by the remarkable success of in-context learning in natural language processing [, ] and computer vision [, , ],
we hereby propose Medical Vision Generalist (MVG), the first generalist model in the medical imaging domain.
Specifically, MVG leverages an in-context learning framework to unify a set of medical imaging tasks, including cross-modal synthesis, denoising, segmentation, and inpainting across modalities like CT, MRI, X-ray, and Micro-ultrasound. In contrast to prior task- and data-specific medical AI models, MVG offers adaptability to new data with minimal labeled samples, eliminating the need for retraining. To achieve this, MVG first standardizes the input/output space using in-context coloring, which maps various tasks into a single-channel coloring scheme. This removes the need for task-specific heads,
thus regulating the model to learn exclusively from prompts. Subsequently, tasks are unified through conditional image generation, where MVG generates output conditioned on both the task prompt and a sample image.

To capture both local and global context, we devise a hybrid strategy that combines masked image modeling and autoregressive training for conditional image generation.
The former involves concatenating prompt images, labels, task inputs, and labels, followed by random masking; the latter constructs prompt image-label pairs, task inputs, and labels as long visual sentences.
During inference, MVG conditions predictions on the prompts selected from locations closely matching the task images, ensuring contextual relevance and guidance that enhances output quality and consistency.

Furthermore, we have curated the first unified medical imaging benchmark, encompassing 13 datasets spanning a range of human anatomies (e.g., abdomen, pelvis, brain, chest) and modalities (e.g., CT, MRI, X-ray, micro-ultrasound). This new benchmark enables a comprehensive assessment of our MVG models.
Experimental results demonstrate the effectiveness of our medical vision generalist in performing various medical vision tasks with only one model. As illustrated in Figure 2, our medical vision generalist outperforms the previous generalist models by a large margin.
For instance, our MVG achieves 0.735 mIoU on all segmentation tasks and outperforms the previous best vision generalist by 0.123 mIoU.
Furthermore, our MVG demonstrates two intriguing properties: 1) it scales well with multiple tasks and datasets, suggesting its potential to excel further as diverse datasets continue to emerge; and 2) it can efficiently generalize to new datasets, with only a few specific examples needed for each task.

## 2 Related Work

#### Medical Image Analysis.

In the field of medical image analysis, there have been key developments in deep-learning models for image segmentation. As the earliest success in this line of work, U-Net [] uses an encoder-decoder architecture with skip-connection, revealing the great potential of deep networks. Following the line, nnUnet [] further improves the model architecture and introduces bags of tricks, building a well-engineered general segmentation model.
TransUnet [] proposed to use pre-trained ViT for better feature extraction.
Recent efforts in medical image analysis have produced remarkable models capable of performing a variety of tasks.
Notable works include "One model to rule them all" [], MedSAM [], and UniverSeg [], which are designed to tackle unified medical segmentation tasks.
UniverSeg adapts UNet [] to intake in-context samples to segment new data and tasks without further training.
Besides, biomedGPT [] proposes a unified generative model for bio-medical vision-language tasks.
In this paper, we propose a novel paradigm to build a generalist model, which is capable of handling various medical vision tasks, including segmentation, inpainting, cross-modal synthesis, and denoising.

#### Universal Models and In-Context Learning.

The advent of the universal Transformer architecture and its success in generative pretraining has inspired the development of universal models that tackle a wide range of computer vision tasks [, , , , , ].
In-context learning is a novel few-shot learning paradigm that emerged in large language models and was first proposed by GPT-3 [].

Figure: Figure 2: Comparison with other generalists. Our model achieves state-of-the-art performance on all involved medical vision tasks of five types.
Refer to caption: x2.png

Specifically, in-context learning enables one model to perform different tasks with only in-context examples as prompts.
While the prompts for language models are mostly defined as a few sentences, in-context learning in other domains is still in an early exploration stage.
As one of the earliest works, Flamingo [] extends the modality of in-context learning with language instructions and sequences of images and videos. Perceiver-IO [] uses the Transformer architecture for a general-purpose model that handles data from arbitrary settings like natural language, visual understanding, multi-modal reasoning, and StarCraft II. AD [] introduces in-context learning to reinforcement learning with algorithm distillation.
DPT [] provides a sample-efficient RL algorithm with strong in-context decision-making. In this paper, we use a sequence of paired medical images to build a vision model with in-context learning ability, unifying 13 medical tasks as a generation task.

## 3 Method

Unlike previous medical AI models, which are specific to one or a few predefined imaging tasks and produce a predetermined set of outputs, the proposed MVG aims to offer unprecedented flexibility across tasks, modalities, and datasets.
The key idea is to unify medical imaging tasks, such as cross-modal synthesis, image segmentation, denoising, and inpainting, within an image-to-image generation framework.

### 3.1 Tasks

Our Medical Vision Generalist (MVG) is designed to address various medical imaging tasks, with a particular emphasis in this study on segmentation, cross-modal synthesis, inpainting, and denoising, tasks for which well-represented public datasets are available.
However, it is crucial to note that its design should be widely applicable to any image-to-image generation task.

Segmentation.
Medical image segmentation, including CT, MRI, X-ray, and Micro-ultrasound segmentation, involves dividing an image obtained from these modalities into distinct segments to isolate regions of interest, such as organs or abnormalities. The input space for these tasks typically consists of images from CT, MRI, X-ray, or Micro-ultrasound scans. The output space is represented by a mask, where each value (excluding the background) in the mask corresponds to a different class or type of object, such as a liver or kidney.

Cross-modal synthesis.
Cross-modal synthesis aims to generate images in one modality from images of another modality for the same subject, aiding in visualization and facilitating multi-modal medical image analysis. The input space and output space are different medical imaging modalities.

Brain image inpainting. In the context of brain image processing, inpainting refers to the process of synthesizing healthy brain tissue in regions affected by glioma, a type of brain tumor []. Inpainting allows professionals to effectively utilize non-standard imaging protocols and directly apply brain parcellation tools to facilitate treatment planning.
The input space is the corrupted brain MRI and the output space is the corresponding brain MRI restoring the affected regions to a normal state.

Denoising. Denoising aims to reconstruct full-dose CT images from low-dose CT images, allowing for reduced radiation doses during CT scans while preserving diagnostic image quality. The input space is the scanned CT image with low-dose radiation, while the output space is the corresponding image with full-dose radiation.

### 3.2 Unifying the Input/Output Space

Assume an input image is denoted as $\mathbf{x}\in\mathcal{R}^{H\times W}$,
the output could be a segmentation map, a synthesized brain image in the target modality, a restored normal brain MRI or a full-dose CT image of the same size.
To unify the output space of images across tasks, our MVG adopts a strategy beyond task-specific heads: mapping different tasks into a single-channel coloring scheme, inspired by [, ].
We explore three different in-context coloring methods that circumvent reliance on label values, including binary, pre-defined, and random colorization, as detailed below.

Binary colorization.
We break down the problem of segmenting multiple classes into individual binary segmentation tasks, each focusing on separating one class from the background. Specifically, if a segmentation mask contains $N_{k}$ foreground classes, we simply split it to $N_{k}$ binary masks. However, this requires multiple inferences when an image contains more than one foreground class.

Pre-defined colorization.
In this approach, we allocate a predetermined unique color to each segmentation mask derived from diverse datasets. Suppose there are $K$ segmentation datasets, with each dataset containing $N_{k}$ classes. Consequently, the $n_{th}$ class of the $k_{th}$ dataset is assigned the value of $\sum\limits_{i=1}^{k-1}i*N_{i}+n$. Note that different tasks may involve classes with identical semantics; for example, both the AMOS segmentation dataset [] and the Synapse dataset [] include the class "Liver". However, distinct colors are assigned to the same class across different tasks.

Random colorization.
The use of pre-defined colors may restrict the adaptability and efficacy of medical vision generalists, as they can cause the model to focus on learning tasks based on the color of the prompt rather than the contextual information [].
To address this limitation, we build a set of colors and randomly sample colors for different semantics in one iteration but the same semantic in the prompt label and task label share the same color.

Figure: Figure 3: Method overview. Left: Four types of medical tasks (i.e., segmentation, cross-modal synthesis, inpainting, and denoising) are unified as a universal image-to-image generation task with in-context learning. Right: We adopt mask image modeling and auto-regressive training for in-context generation.
Refer to caption: x3.png

Except for medical image segmentation, the outputs of all other tasks in this study do not involve categorical values that need to be predicted. Therefore, we do not apply coloring for these tasks.

### 3.3 Task Unification via Conditional Image Generation

After standardizing the input and output space for all tasks as images of identical sizes,
we construct the training input, including 1) the task prompt consisting of paired prompt images and prompt labels, and 2) the task input and its associated label. We then unify various medical imaging tasks within a conditional image-to-image generation framework using the task prompt as task specification.
All the tasks are unified to generate the task label $Y\in\mathcal{R}^{H\times W}$ based on the condition including the task image $X\in\mathcal{R}^{H\times W}$, the prompt image $P_{x}\in\mathcal{R}^{H\times W}$, and the prompt label $P_{y}\in\mathcal{R}^{H\times W}$.
Specifically, we use two conditional image generation frameworks - mask image modeling [] and autoregressive training [].

#### Architecture Selection.

Following the same setting in [, ], we take vanilla ViT [] as an encoder including a patch embedding layer and several Transformer blocks. The decoder is a simple prediction head with two convolution layers and takes four feature maps [] from ViT as input.

#### Mask Image Modeling.

During training, we form a square image by concatenating the prompt image (upper left) with its corresponding label (upper right), as well as the task image (lower left) with its associated label (lower right), as illustrated in Figure 3(a). We perform random masking on the square image and train ViT to reconstruct the masked region []:

$$ $\begin{split}p(x)&=\prod\limits_{i=1}^{M}p(x_{i}|x_{x\not\in x_{M}},\theta).\\ \end{split}$ (1) $$

where $x_{M}$ is the mask region, $x_{x\not\in x_{M}}$ is the visible region, and $\theta$ denotes the model parameters.
However, in practice, we observed that mask image modeling yields unsatisfactory results for medical image segmentation.
We hypothesize that this may be attributed to the masking strategy’s potential to compromise the preservation of global contextual information and to ignore small regions-of-interest.
To ensure the efficacy on medical image segmentation, we introduce an additional auto-regressive training, which preserves the global context within individual images, as shown below.

#### Auto-Regressive Training.

In auto-regressive training, each image, including paired prompt images, prompt labels, task inputs, and associated labels, is treated as a single element in a sequential data structure. The model is fed with a partial sequence and trained to predict the next image in the sequence conditioning on the preceding ones.

Mathematically, let $P_{x_{1}},P_{y_{1}},...,P_{x_{n}},P_{y_{n}},X,Y$ denote $n+1$ pairs of images and labels. The first $n$ pairs serve as the task prompt, and the model learns to predict the task output $Y$ given the task input $X$ and the prompt. This process iterates through each pair in the sequence. For each iteration, auto-regressive training is conducted with supervision solely on prompt labels and the task label:

$$ $\begin{split}S=[S_{1},S_{2},...,S_{2n-1},&S_{2n},S_{2n+1},S_{2n+2}]=[P_{x_{1}} ,P_{y_{1}},...,P_{x_{n}},P_{y_{n}},X,Y],\\ p(x)&=\prod\limits_{i=1}^{n+1}p(S_{2i}|S_{1},...,S_{2i-1},\theta).\\ \end{split}$ (2) $$

#### Loss Function.

Any regression loss function like $\boldsymbol{l}_{1}$ or $\boldsymbol{l}_{2}$ can serve as the loss function of our MVG. Different from the $\boldsymbol{l}_{2}$ loss function in MAE [], we find the smooth $\boldsymbol{l}_{1}$ performs best for MVG.

#### Inference.

We first construct a sequence $S=[P_{x},P_{y},X,\hat{Y}]$, where prompts, the task image, and the desired output are concatenated together. MVG leverages the task prompt, composed of the prompt image $P_{x}$ and label $P_{y}$, for task specification, subsequently generating predictions by conditioning on both the task input $X$ and the task prompt.
Since the task prompt is formulated as images, MVG demonstrates versatility in defining imaging tasks, capable of handling data sourced from diverse scanning machines, procedures, settings, or populations. For instance, if $P_{x}$ and $P_{y}$ represent an image and label extracted from the AMOS CT training set, respectively, MVG performs multi-organ segmentation on the image $X$ derived from the AMOS CT testing set, maintaining consistency within the dataset setting and guided by the provided context.

Different task prompts can yield varying results. In this study, we address this variability by selecting the prompt image from a location that closely matches the task image.
Given the instance $X_{TE}$ which has $N_{TE}$ slices from the testing set, we randomly choose an instance $X_{TR}$ which has $N_{TR}$ slices and the corresponding label $Y_{TR}$ from the training set. For the $n_{th}$ slice of $X_{TE}$, we also choose the the $\mathrm{floor}(\frac{n_{th}*N_{TR}}{N_{TE}})$ slice as the prompt.

## 4 Experiment

**Table 1: Datasets overview. Our MVG is trained and evaluated on 13 different datasets covering four major human body regions (i.e., Abdomen, Pelvis, Brain, Chest). #Training/Testing refers to the number of samples for training and testing.**
| Region | Dataset | Modality | #Training | #Testing | Task |
| --- | --- | --- | --- | --- | --- |
| Abdomen | AMOS [] | CT | 240 | 120 | Segmentation |
| Abdomen | WORD [] | CT | 100 | 20 | Segmentation |
| Abdomen | BTCV [] | CT | 21 | 9 | Segmentation |
| Abdomen | AMOS [] | MRI | 60 | 50 | Segmentation |
| Pelvis | MicroSegNet [] | Micro-US | 55 | 20 | Segmentation |
| Pelvis | PROMISE [] | MRI | 50 | 30 | Segmentation |
| Brain | BraTS-GLI [] | MRI | 1251 | 219 | Cross-modal synthesis |
| Brain | BraTS-Local [] | MRI | 1000 | 251 | Inpainting |
| Chest | Low dose  [] | CT | 200 | 59 | Denoising |
| Chest | Defect Detection [] | Xray | 15 | 6 | Segmentation |
| Chest | ACDC [] | MRI | 100 | 50 | Segmentation |
| Chest | LA [] | MRI | 81 | 20 | Segmentation |

### 4.1 Implementation Details

#### Data.

As shown in Table 1, our model is developed on 13 different datasets covering four major human body regions (i.e., Abdomen, Pelvis, Brain, Chest). More details can be found in the appendix. Following the standard preprocessing strategy, we apply a windowing range of [-100, 200] to all involved CT scans for better contrast. Input images are firstly resized to 512 $\times$512 and then randomly cropped with a size of 448 $\times$ 448.
To evaluate the generalization of MVG, we choose MSD [], a multi-organ segmentation dataset as an out-of-distribution dataset.

#### Training details.

AdamW optimizer is used with a weight decay of 0.05. The peak learning is set to $1e^{-3}$ with a cosine learning rate scheduler. We train our model 100 epochs with 5 warm-up epochs. We only adopt the random crop as the data augmentation. The sampling weight of segmentation tasks is 0.5 while the rest of the tasks share 0.5.
We use 8 A5000 GPUs to train our models.
We use 1 in-context sample for both training and inference.

#### Training Objective.

In practice, for all tasks except segmentation, we perform 90% training iteration with mask image modeling and 10% training iteration with auto-regressive training. For segmentation tasks, we perform 100% training iterations with auto-regressive training.

#### Evaluation.

We use mean IoU (mIoU) as the evaluation metric for segmentation. For cross-modal synthesis, inpainting, and denoising, we use mean absolute error (MAE), peak signal-to-noise ratio (PSNR), and structural similarity index measure (SSIM) as the evaluation metric.

### 4.2 A Generalist to 13 Medical Tasks

#### Baselines.

Our generalist baselines includes LVM [] and Painter [], which are trained on our benchmark.
UniverSeg [] is used as a segmentation generalist baseline.
While our specialist baselines includes ResNet-18 [] with a two-layer MLP decoder, UNet [], VNet [], TransUNet [], and nnUNet [].
For synthesis tasks, we involve Pix2Pix [] as an additional baseline.

#### Quantitative evaluation.

In Table 2, we compare our method with the latest vision generalist models[, ] across a range of segmentation tasks. Our medical vision generalist achieves the best performance of 0.79 mIoU among all the generalists.
Specifically, Our medical vision generalist outperforms Painter [] by 0.24 mIoU, LVM [] by 0.62 mIoU on average, and UniverSeg [] by 0.35 mIoU.
All the generalists only require one model to perform these different tasks.
UniverSeg is trained with up to 64 in-context samples, yet still yields inferior performance to our method which only relies on 1 in-context sample.
At the same time, specialist models like UNet [], TranUNet [], and nnUNet [], which need to train different models for different tasks, still hold the edge in performance.

We report the image synthesis results in Table 3. we present a detailed quantitative comparison of our method with the latest generalist and specialist models across various tasks including cross-modal synthesis, inpainting, and denoising. Our Medical Vision Generalist (MVG) model demonstrates strong capabilities and achieves competitive performance, particularly in the generalist category. For instance, in the task of cross-modal synthesis, MVG shows an improvement over Painter in all metrics: a lower Mean Absolute Error (MAE) by 0.002, a higher Peak Signal-to-Noise Ratio (PSNR) by 0.69, and a better Structural Similarity Index (SSIM) by 0.009 over the best vision generalist.

**Table 2: Quantitative evaluation in segmentation tasks. Compared to other generalists, our method achieves state-of-the-art performance with solid improvements. ^∗: We inference the official weights with the 64 in-context samples from training set.**
| Method | AMOS CT | WORD | BTCV | AMOS MRI | MicroSegNet | PROMISE | Chest Defect | ACDC | LA |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| *Specialists* |  |  |  |  |  |  |  |  |  |
| ResNet-18 | 0.55 | 0.50 | 0.51 | 0.53 | 0.67 | 0.75 | 0.62 | 0.69 | 0.68 |
| UNet | 0.81 | 0.83 | 0.82 | 0.81 | 0.90 | 0.91 | 0.89 | 0.86 | 0.83 |
| VNet | 0.70 | 0.75 | 0.72 | 0.73 | 0.90 | 0.89 | 0.86 | 0.87 | 0.84 |
| TranUNet | 0.80 | 0.82 | 0.84 | 0.82 | 0.94 | 0.90 | 0.88 | 0.88 | 0.84 |
| nnUNet | 0.87 | 0.90 | 0.91 | 0.88 | 0.97 | 0.93 | 0.90 | 0.90 | 0.89 |
| *Generalists* |  |  |  |  |  |  |  |  |  |
| UniverSeg^∗ | 0.20 | 0.29 | 0.37 | 0.25 | 0.71 | 0.55 | 0.55 | 0.54 | 0.57 |
| Painter | 0.52 | 0.48 | 0.45 | 0.51 | 0.69 | 0.68 | 0.50 | 0.52 | 0.55 |
| LVM | 0.12 | 0.14 | 0.10 | 0.15 | 0.36 | 0.30 | 0.10 | 0.12 | 0.13 |
| MVG | 0.73 | 0.74 | 0.73 | 0.74 | 0.91 | 0.85 | 0.79 | 0.85 | 0.81 |

**Table 3: Quantitative comparison with other tasks. Our model shows strong capabilities in the tasks of cross-modal synthesis, impainting, and denoising.**
| Method | Cross-modal synthesis | Inpaiting | Denoise |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| MAE | PSNR | SSIM | MAE | PSNR | SSIM | MAE | PSNR | SSIM |  |
| *Specialists* |  |  |  |  |  |  |  |  |  |
| ResNet-18 | 0.026 | 20.984 | 0.860 | 0.008 | 30.981 | 0.959 | 0.022 | 30.519 | 0.709 |
| Pix2Pix | 0.018 | 24.311 | 0.899 | 0.008 | 34.891 | 0.982 | 0.020 | 33.011 | 0.730 |
| TranUNet | 0.016 | 25.541 | 0.938 | 0.005 | 35.561 | 0.989 | 0.016 | 33.999 | 0.761 |
| *Generalists* |  |  |  |  |  |  |  |  |  |
| Painter | 0.021 | 24.031 | 0.920 | 0.006 | 33.595 | 0.978 | 0.020 | 33.104 | 0.721 |
| MVG | 0.019 | 24.721 | 0.929 | 0.006 | 34.521 | 0.981 | 0.018 | 33.521 | 0.731 |

#### Qualitative evaluation.

To provide a more intuitive observation of our medical vision generalist, we provide the visualization of different tasks in Figure 4.

|  |  |  |  |
| --- | --- | --- | --- |
|  |  |  |  |
|  |  |  |  |
|  |  |  |  |
|  | Image | Prediction | label |

Figure: Figure 4: Qualitative evaluation of four tasks: Segmentation (1st row), denoising (2nd row),cross-modal synthesis (3th row), and inpainting (4th row).

|  |  |
| --- | --- |
|  |  |
|  |  |
|  |  |

Figure: Figure 5: Impact of training data scale. We ablate on various scales of the training data (randomly sampled from each dataset), ranging from 1% to 100%.

#### Generalize to unseen datasets.

The advantage of in-context learning is that it allows models to adapt to new datasets quickly. We report the result in Table 4, MVG achieves 0.84 mIoU on MSD-Liver with only new prompts without any fine-tuning. After fine-tuning MVG with one instance on MSD-Spleen and MSD-Lung, MVG achieves 0.87 mIoU and 0.48 mIoU.

| Dataset | mIoU |
| --- | --- |
| MSD-Liver | 0.84 |
| MSD-Spleen | 0.87 |
| MSD-Lung | 0.48 |

Table: Table 4: Generalization to the unseen MSD dataset.

#### Generalize to unseen datasets.

The advantage of in-context learning is that it allows models to adapt to new datasets quickly. We report the result in Table 4, MVG achieves 0.84 mIoU on MSD-Liver with only new prompts without any fine-tuning. After fine-tuning MVG with one instance on MSD-Spleen and MSD-Lung, MVG achieves 0.87 mIoU and 0.48 mIoU.

### 4.3 Ablation Study

#### Data scalability.

Nowadays, more and more datasets are available, which motivates us to study whether a scale-up dataset can train a stronger medical vision generalist.
As shown in Table 5, we randomly choose 1%, 10%, 50% as comparison with the full data. The performance of our medical vision generalist consistently improves with the growth of dataset size. These results show the strong dataset scalability of our medical vision generalist.
We also compare it with the specialists, which are trained on each dataset separately, with the same data sampling proportion (i.e., 1%, 10%, 50%).
We can interpret that in the low data regime, the specialists suffer from overfitting.

#### Color space.

To unify the output space of different segmentation tasks in which the same value from different datasets may have different semantics, we propose to unify the output space with a pre-defined color for each class or random color that we keep the same semantics have the colors.
As shown in Table 5, the random color performs much better than the pre-defined color on abdominal segmentation while having a similar performance on prostate segmentation.
In particular, our medical vision generalist with random color space gains the average result of 0.735 mIoU on abdominal segmentation and improves 0.123 mIoU over that with pre-defined color space.
In contrast, our medical vision generalist achieves inferior performance with pre-defined or random colors.
The random color makes medical vision generalists learn more from the context instead of the color itself and avoid the model being limited by the number of colors.

#### Isolated and Unified training.

To prove our medical vision generalist can benefit from large-scale datasets across different tasks. We compare two settings: 1) isolated training: we train different models on different datasets in isolation. Namely, we train 13 models for the 13 datasets. 2) unified training: we train our medical vision generalist on all datasets together. Note that both settings have the same model architecture. We report the results in Table 5. The unified model makes significant improvements over the isolated model in all tasks and the improvements reach 0.14 mIoU. Such results indicate that the medical vision generalist can benefit from large-scale datasets even if this dataset has different annotation semantics which motivates the medical image analysis community to further expand the datasets.

#### Auto-regressive v.s. MIM.

We ablate the conditional image generation method: mask image modeling and auto-regressive training. As shown in Table 5, Auto-regressive training emerges as the significantly superior method, outperforming MIM across all datasets. The results underscore the fundamental weakness of random masking strategies in segmentation tasks, especially when dealing with small anatomical organs. By leveraging medical images’ inherent spatial and contextual information, auto-regressive training offers a powerful alternative that significantly enhances in-context learning and segmentation performance.

**Table 5: Ablation study: *a) Color space for segmentation.* Using semantic masks in a random color space as prompts significantly improves the segmentation performance of our generalist model.**
| Method | AMOS CT | WORD | BTCV | AMOS MRI | MicroSegNet | PROMISE | Chest Defect | ACDC | LA |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| *a) Color Space* |  |  |  |  |  |  |  |  |  |
| Binary | 0.46 | 0.48 | 0.48 | 0.49 | 0.78 | 0.78 | 0.45 | 0.49 | 0.52 |
| Pre-defined | 0.60 | 0.62 | 0.62 | 0.61 | 0.89 | 0.86 | 0.71 | 0.74 | 0.73 |
| Random | 0.73 | 0.74 | 0.73 | 0.74 | 0.91 | 0.85 | 0.79 | 0.85 | 0.81 |
| *b) Isolated vs. Unified Training* |  |  |  |  |  |  |  |  |  |
| Isolated | 0.55 | 0.57 | 0.57 | 0.58 | 0.80 | 0.77 | 0.70 | 0.69 | 0.68 |
| Unified | 0.73 | 0.74 | 0.73 | 0.74 | 0.91 | 0.85 | 0.79 | 0.85 | 0.81 |
| *c) Auto-regressive vs. MIM in Segmentation Training* |  |  |  |  |  |  |  |  |  |
| MIM (mask 50%) | 0.56 | 0.48 | 0.46 | 0.54 | 0.70 | 0.66 | 0.50 | 0.50 | 0.52 |
| MIM (mask 75%) | 0.53 | 0.42 | 0.44 | 0.52 | 0.70 | 0.63 | 0.48 | 0.50 | 0.51 |
| Auto-regressive | 0.73 | 0.74 | 0.73 | 0.74 | 0.91 | 0.85 | 0.79 | 0.85 | 0.81 |

## 5 Conclusion and Discussion

#### Conclusion.

In this work, we present MVG, a versatile model capable of handling various medical imaging tasks, including cross-modal synthesis, segmentation, denoising, and inpainting, within a unified image-to-image generation framework. MVG employs an in-context generation strategy to standardize inputs and outputs as images, allowing flexible task unification across various modalities and datasets.
A hybrid approach combining masked image modeling and autoregressive training proved the most effective.
To thoroughly assess MVG’s potential and limitations, we also curate the first comprehensive generalist medical vision benchmark consisting of 13 datasets across 4 imaging modalities, including CT, MRI, X-ray, and Micro-ultrasound.
Experiment results demonstrate that MVG consistently outperforms existing vision generalists.
Benefiting from the in-context learning scheme, MVG demonstrates exceptional flexibility, scalability and potential for generalization to unseen datasets with minimal samples. We will make our code and benchmark publicly available to encourage future research in medical AI generalists.

#### Limitations.

Despite the encouraging results above, our MVG still falls behind the specialist models on different benchmarks. Additionally,
the quality of MVG’s final prediction is sensitive to the choice of the in-context sample. We conjecture this is because, currently, we only use one in-context sample for both training and inference (due to computing resource constraints), but for the next step, we plan to explore techniques such as strategically selecting in-context samples to alleviate the problem.
Our current framework is in 2D but in the future we plan to extend our method to 2.5D or 3D models.

### Acknowledge

This work is partially supported TPU Research Cloud (TRC) program, Google Cloud Research Credits program, and AWS Cloud Credit for Research program.

## Appendix A Visualization

Figure: Figure 6: Qualitative evaluation of segmentation. MVG shows strong capabilities on various segmentation tasks covering multiple modalities and body regions.
Refer to caption: x16.png

In the paper, we only visualize one segmentation dataset. Here, we provide visualization results on more segmentation datasets as shown in Figure 6.

**Table 6: Pre-training on medical data is important. We initialize the encoder with natural image supervised pretrained model, natural image mask image modeling pretrained model, and medical image mask image modeling pretrained model.**
| Method | AMOS CT | WORD | BTCV | AMOS MRI | MicroSegNet | PROMISE | Chest Defect | ACDC | LA |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Sup. | 0.55 | 0.57 | 0.57 | 0.58 | 0.80 | 0.70 | 0.68 | 0.69 | 0.68 |
| Natural MIM | 0.69 | 0.70 | 0.68 | 0.69 | 0.89 | 0.83 | 0.76 | 0.81 | 0.78 |
| Medical MIM | 0.73 | 0.74 | 0.73 | 0.74 | 0.91 | 0.85 | 0.79 | 0.85 | 0.81 |

## Appendix B Pretraining

Most parameters of our medical vision generalist are from the ViT encoder. Recent works prove that pretraining makes ViT learn more robust and transferable representations. Traditionally, the supervised pretrained model [] on ImageNet is loaded as the model initialization for various medical segmentation tasks (Sup.). Recent progress in self-supervised learning demonstrates that the model pretrained by mask image modeling achieves better generalization. Therefore, for comparison, we load the MAE, which trains the model on a natural image for comparison (Natural MIM). However, the domain gap between natural images and medical images limits the power of medical vision tasks. Therefore, we train a model with mask image modeling on medical images for comparison. As shown in Table 6, the superiority of the Medical MIM approach is evident, with the model achieving the highest performance metrics across all evaluated tasks

## Appendix C Societal Impact

MVG promises to adapt to new medical imaging tasks easily, given in-context samples that represent the new tasks.
We believe MVG can serve as a stepstone to make medical imaging tools more accessible to clinical researchers, other scientists, or beyond, by lowering the bar of machine learning expertise, computational resources, and human labor.
Since our method is data-driven, the bias in the data would be inherited by models.
The users should be aware of the bias and its potential detriment.

## Peer Reviews

confidence: 5
contribution: 4
initial_rating: 6
presentation: 3
questions: Line 146 and subsection 4.2 :  “13 tasks…”. 
I suppose it is a miswriting, since there are 13 datasets, but only 4 tasks. 

Line 246 and Figure 3:
The sub-numeration of blocks in the figure would make the reference to its blocks and understanding easier. Instead of referring to blocks as "upper right" or "lower left".

Line 248: There is a reference to “Fig 3(a)”, but, again, no numeration in the image.

Table 1. I think it is important to note in the beginning of the article, that in the benchmark the proportion of datasets for 4 task in non-equal, with Segmentation task comprising the majority.

Line 255: “We hypothesize that this may be attributed to the masking strategy..." - Indeed, medical images data have fundamental and well known differences from the natural domain data; here, you could cite some papers that discuss this topic in more detail. For example (Huang et al., Self-supervised learning for medical image classification: a systematic review and implementation guidelines)

Line 465 and Table 4. The difference between Liver and Spleen datasets and Lung dataset is almost two times. Can you elaborate on the causes of such a huge difference?

Line 470: Data scalability section - there is no reference to Fig. 6.

The reliance on prompts to condition predictions introduces variability in outputs and performance, because different prompts may yield different results. Did authors address it somehow in their experiments? At the very least, how is the performance affected by the different prompts?
soundness: 4
strengths: The comprehensiveness of imaging modalities and tasks. The paper addresses segmentation, cross-modal synthesis, inpainting, and denoising tasks across various modalities, like CT, MRI, X-ray, and micro-ultrasound.

High performance and generalization ability is remarkable. MVG outperforms SOTA models such as Painter and LVM in most metrics. It also demonstrates scalability with different datasets and adaptability to unseen datasets with minimal samples.

Methodology of training - hybrid approach with masked image modeling and autoregressive training is a novelty that provides strong performance across diverse tasks. Also this approach partly solves the problem with medical imaging data augmentation (like cropping operations).
summary: The work is about a new foundation model for medical image analysis that aims to unify multiple imaging tasks (segmentation, cross-modal synthesis, inpainting, and denoising), called Medical Vision Generalist (MVG). The implementation is within a single model using a standardized image-to-image generation framework. MVG distinguishes itself by employing in-context learning strategies, which eliminate the need for retraining on new datasets and enable quick adaptation to unseen tasks with minimal labeled samples. Also authors introduce a benchmark that covers 13 datasets across 4 medical tasks for various imaging modalities (CT, MRI, X-ray, and micro-ultrasound). The latter makes this study not only about the methodology but also an important contribution to the research community.
weaknesses: As shown by authors, the hybrid use of autoregressive training boosts performance. However, it may impose higher computational costs during inference. The authors could provide a more detailed analysis of the trade-offs between performance and inference efficiency, especially when using MVG with heavy medical files, like the high resolution MRI.

The ablation study lacks detailed insights into the contribution of individual components. For example, how do masked image modeling and autoregressive training individually affect the performance? For example, as the authors mentioned, usual augmentation techniques for SSL, like image cropping, can not be applied to medical imaging. The cropping procedure might mislead the model training and plumet the performance. I wonder if the authors addressed that and conducted some experiments with (specific for medical domain) data augmentation?

The paper compares the authors model with generalist models, however the comparative analysis with specialist models such as U-Net and nnUNet could be more detailed. A deeper examination of the performance trade-offs between generalist and specialist models would help to position MVG’s contribution more clearly and to understand the limits of its applicability. Also, I suggest addressing not only the segmentation task in a great detail; additional qualitative examples for inpainting and denoising would improve the clarity and show how MVG compares visually to specialists.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Can the authors explain the discrepancy in ACDC Dice scores for UniverSeg between Table 2 in this paper (0.54) and Table 6 in the UniverSeg paper appendix (0.70)?

2. Is the failure of mask modeling in segmentation tasks due to the use of L1 loss? Have alternative segmentation losses, such as Dice loss, been tested?

3. Could the authors elaborate on the motivation and practical benefits of MVG in medical imaging? Clinically, how does a foundational model outperform specialist models in practice, such as in speed, ease of deployment, or cost-effectiveness? Technically, what key insights from MVG could guide future researchers in developing clinically and economically viable foundational models?
soundness: 2
strengths: 1. The authors insightfully recognize that a pure masking strategy is insufficient for medical image segmentation, leading them to incorporate an autoregressive training pipeline. Experimental results confirm the effectiveness of this approach.

2. Unlike other foundational models focused primarily on segmentation, this paper addresses a broader range of tasks, including cross-modal synthesis, image denoising, and inpainting, opening potential new research directions.
summary: This paper introduces Medical Vision Generalist (MVG), a foundational model designed to handle diverse medical imaging tasks within a unified framework. MVG standardizes input and output as images through an in-context generation approach, treating tasks as an image generation process conditioned on prompt image-label pairs. To leverage both local and global context, MVG combines masked image modeling with autoregressive training for conditional image generation. Experimental results indicate that MVG outperforms existing vision generalists, such as Painter and LVM, demonstrating scalability and adaptability across modalities and datasets.
weaknesses: 1. While the addition of multiple tasks is beneficial, the paper overlooks essential medical imaging tasks, such as image registration and inverse reconstruction, making MVG appear more like an expanded segmentation model than a comprehensive foundation model. The reviewer suggests that MVG’s learned feature representations could potentially support image registration by integrating a flow estimation head, and inverse reconstruction by using denoising as a regularizer, unfolding the inverse optimization problem with forward consistency [1] within the network.

2. Restricting training to 2D images raises concerns about MVG’s utility as a foundational model for medical imaging. Effective 3D analysis is crucial, as many anatomical structures (e.g., brain cortex, lung vessels, heart) span significant volumes, where 2D slices may miss critical contextual information.

3. The authors claim that MVG “scales well with multiple tasks and datasets,” yet the evidence provided in Figure 6 and Table 6 only demonstrates that more data improves performance, a known property of deep learning models, and that unified training is preferable to isolated training, which reiterates established insights for vision transformers requiring large-scale data.

4. The paper lacks runtime and complexity analysis, particularly GPU resources for training. Comparisons with resource-efficient models, such as nnU-Net trained on individual datasets, would offer a clearer picture. Specifically, what GPU/hours are required for training MVG on all datasets, and how does GPU memory usage compare to nnU-Net or other specialist models trained on individual datasets?
[1] MoDL: Model-Based Deep Learning Architecture for Inverse Problems, TMI 2019.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: To better demonstrate the advantage of unifying multiple medical vision tasks through the MVG model, the authors could make the following efforts to address the concerns,
1. compare the segmentation performance while isolating the effects of the colorization strategy from task unification.
2. discuss potential reasons for the marginal improvements in these tasks and propose strategies to close the gap with specialist models
3. conduct scalability experiments on the unified datasets.
soundness: 2
strengths: The paper first uses in-context learning to unify multiple medical vision tasks, which is original.
The proposed output space unification strategy is useful when training with multiple segmentation tasks.
summary: The paper introduces a medical vision generalist (MVG) model, which unifies various medical imaging tasks including segmentation, cross-modal synthesis, denoising, and inpainting within a single image-to-image generation framework. MVG utilizes in-context learning, treating tasks as image generation processes conditioned on prompt image-label pairs, allowing for flexible adaptation to different modalities and datasets. The authors have also curated a comprehensive benchmark with 13 datasets across four imaging modalities to evaluate MVG, which consistently outperforms existing generalist models.
weaknesses: The advantage of unifying multiple medical vision tasks through the MVG model could not be verified based on the evidence provided in this paper.  
According to Table 2 and Table 5, the performance gain in segmentation tasks could be due to the colorization strategy instead of unifying other vision tasks. 
Regarding cross-modal synthesis, inpainting, and Denoising tasks, the improvement by the MVG model is marginal compared to the previous generalist model and all generalist models performs worse than specialist models in each task according to Table 3.
Besides, the experiments on the scalability were conducted on several small datasets, which could not demonstrate the potential to increase the performance when unifying all datasets.

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: - Could the authors explain the potential clinical benefits of your proposed framework in detail? 
- Could the authors add more unseen-dataset evaluation to test the generalization of the proposed framework, for example cell data whose semantic is not available in the training data?
- Could the authors explain architecture of your ViT, such as the embedding size or number of layers or the total number of parameters, since the network requires 8 A5000 gpus to train?
- Could the authors explain the training time of auto-aggression modeling?
- Could the authors explain the inference resources needed of your model? Can it run on CPU under an acceptable time? How much time does it require to infer a given test image under a given test hardware?
soundness: 3
strengths: - The paper is well written and easy to follow. 
- The paper proposes a solution to address the generalization of medical imaging analysis, such as cross-domain problem.
- The paper proposes a unified colorization formulation to unify the different output types of medical imaging analysis tasks. 
- The paper treats the different learning tasks as a prompt-based learning task. 
- The paper introduces a new benchmark for generalist medical imaging analysis.
summary: In this paper, the authors introduce a unified framework, named Medical Vision Generalist, for medical imaging analysis tasks, including segmentation, cross-modality synthesis, inpainting and denoising. The authors formulate the learning task as prompt-based learning task and prompt include task-image and task-label pairs. The authors adopt a single-channel colorization method to unify the output space of images across tasks. The authors combine the masked imaging modeling and auto-aggressive training method to train the model. To demonstrate the framework's capabilities, the authors curate a comprehensive generalist medical vision benchmark.
weaknesses: - The proposed framework is limited to 2D scenario.
- The paper does not explain the clinical motivation why a generalized medical imaging analysis model is needed. To me, since there are many pre-trained specialized models, medical researchers can just pick one of the state-of-the-art models and get better performance than the generalized models. 
- The motivation of combining the masked imaging modeling and auto-aggressive is not clear. In the experiment, auto-aggressive training is more superior than masked imaging modeling. Then why combining them in the first place?
- The generalization of proposed framework is not very convincing, since the organs in the unseen-dataset also appears in training set. What about other types of medical imaging dataset, such as pathology data?