Paper: Unified Medical Image Pre-training in Language-Guided Common Semantic Space (ICLR-2024)

Reviewer comment (Novelty):
The novelty appears to be somewhat incremental, as the method may seem to be a direct application of Attentive Mask CLIP proposed by Yang et al. (2023).

Author response:
We re-elaborate more here on the novelty of our method, which is quite different from Attentive Mask CLIP (ACLIP).

First, we target an unified VL framework capable of handling various medical modalities (e.g. 2D X-rays, 3D CT), and Figure 2 shows that merging representations of different medical image modalities into a language-guided common semantic space is challenging. Naively using language guidance and VL contrastive learning (Figure 2a) and simply unifying in one model (Figure 2b) cannot achieve the goal of merging. To this end, we introduce several designs for UniMedI, and *all of the designs aim at better merging 2D and 3D medical images into a language-guided common semantic space*. To best of our knowledge, we are the first one merging 2D and 3D medical images into a unified VL framework (we also provide a table to summarize the difference between our work and existing methods in this field at the end of this answer). Particularly, ACLIP cannot handle multi-modal images, and it can only be used for single modal images. So, *ACLIP cannot be directly applied to our problem*.


The first design is incorporating the guidance from report via VL contrastive learning. Generally, the $[CLS]$ token in the image side is optimized with representation of the paired report. And this image $[CLS]$ token is used in the following parts of UniMedI to ensure that the integrating is guided by language. We also have a novel motivating observation here for our problem (shown in Figure 1 and introduction). We see that despite big differences, medical images from various modalities share a common semantic latent space, which captures the underlying features of an individual's health status, and such status are reflected in medical reports via language. *This motivating observation is insightful in our problem and has not been mentioned in previous works as far as we know.*

The second design is our pipeline. Since no paired 2D and 3D image data exists, so we first select informative 2D slice from 3D volume as the bridge via the proposed attentive selection strategy. And then 2D slices, 2D X-ray and 3D volume are sent to the vision encoder together, enabling better merging. *ACLIP does not have this pipeline and the attentive selection strategy, since bridging multimodal images is not handled in ACLIP.* We have updated our paper by including a subsfigure in Figure 3 to show this pipeline.

The third design is the auxiliary task of mask and recovery implemented by self-distillation. Since 2D and 3D tokens are sent to the visual encoder at the same time and a large proportion of tokens are masked, this auxiliary task can help to enhance interactions among 2D and 3D tokens, facilitating merging representations. *The mask in ACLIP is used for removing redundant information, while the auxiliary task in UniMedI is enhancing dimensional interactions.*

Besides, the vision encoder has *separate tokenizer for 2D and 3D data and a shared backbone*, which is also designed for representation integration, and not used by ACLIP.

In summary, the problem we addressed in this paper is quite novel, and all designs in UniMedI aim at better solving the targeted problem. Maybe some losses are similar with other methods, but they are introduced with different purpose here, and thus brings different effects. And there are lots of difference compared with ACLIP, as we listed above.

| Method | Vision-Language | Unify | Downstream Task | Input Medical Image type |
| :------: | :---------------: | :-----: | :---------------: | :------------------------: |
|COVIRT [1] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|GLoRIA-MIMIC [2] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|MGCA [3] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
| Joint [4]   | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 2D slices  |
| UniMiss [5] | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 3D volumes, 2D slices |
| UniMedI     | $\checkmark$ | $\checkmark$| Classification, Segmentation, Detection(Appendix)          | **2D Images, 3D volumes, 2D slices** |

Reviewer comment (Experiments):
Although the method is capable of handling 3D image data, it treats them as separate 2D images and does not consider the structural aspects of 3D data.

Author response:
Actually, the structural aspects of 3D data is considered in UniMedI. Specifically, after extracting 2D slices, the 3D volume and 2D slices are sent to the vision encoder as inputs at the same time. With separate tokenizers, 2D and 3D tokens exists in one pass. In this way, we can use the mask and recovery task to enhance dimensional interactions. Sorry for causing confusion. We have revised Figure 3 by incorporating a subfigure to illustrate the overall pipeline, and adjust corresponding descriptions, making this point clear.

Reviewer comment (Reproducibility):
As mentioned above, please provide the decoder structure used for the segmentation task.

Author response:
Thanks for the suggestion. We have added more implementation details in Appendix B, including the structures of the segmentation decoder. Generally, we use SETR-PUP [6] for 2D segmentation tasks and UNETR [7] for 3D segmentation tasks.

Reviewer comment (Experiments):
The medical semantic segmentation task on the RSNA and BCV datasets appears to involve predicting bounding boxes that indicate evidence of pneumonia, thus making it more akin to a detection problem rather than a segmentation problem.

Author response:
Generally, predicting bounding boxes is involved for the RSNA dataset, but not for BCV dataset. RSNA is a pneumonia dataset, and detection of bounding boxes can serve as a preliminary step in the segmentation process. So, we add results for segmentation tasks in RSNA in the below table and Appendix C.1 (metric: Dice[%]). Note that comparing with detecting bounding boxes, segmentation can be a more complex task providing more information via precisely segmenting the affected regions within the lungs. Moreover, segmentation task allows for a more detailed comparison with previous methods. So we listed segmentation results in our main paper.

|         | 1%   | 10%  | 100% |
| :-------: | :----: | :----: | :----: |
| ConVIRT [1] | 8.2 | 5.6 | 17.9 |
| GLoRIA [2] | 9.8 |  14.8  | 18.8 |
|GLoRIA-MIMIC [2] | 11.6 | 16.1 | 24.8 |
| MGCA (ResNet-50) [3] | 12.9 | 16.8 | 24.9 |
| MGCA (ViT-B) [3] | 14.7 | 18.4 | 25.8 |
| UniMedI (ViT-B) | **15.5**  | **19.2**  | **26.6**  |

Besdies, BCV is a organ segmentation dataset, but not a pneumonia dataset. Organ segmentation is a more meaningful task since detailed shape information of organs is useful in clinical problems, rather than bounding boxes of organs. Also, detecting is not involved for organ segmentation in common practice. So, this BCV dataset is only used to evaluate performance of organ segmentation, and no previous work uses it for detection. To further demonstrate the capability of our pretrained model on 3D data, we add a larger 3D segmentation dataset AMOS. Results is listed below and included in Appendix C.2. 

|         | 20%   | 40%  | 100% |
| ------- | ---- | ---- | ---- |
| UniMiss [5]    | 79.5 | 82.3 | 85.8 |
| UniMedI | **78.8** | **82.9** | **86.4** |

Reviewer comment (Novelty):
Please elucidate the distinctions between the proposed method and Attentive Mask CLIP proposed by Yang et al. (2023).

Author response:
We re-elaborate more here on the novelty of our method, which is quite different from Attentive Mask CLIP (ACLIP).

First, we target an unified VL framework capable of handling various medical modalities (e.g. 2D X-rays, 3D CT), and Figure 2 shows that merging representations of different medical image modalities into a language-guided common semantic space is challenging. Naively using language guidance and VL contrastive learning (Figure 2a) and simply unifying in one model (Figure 2b) cannot achieve the goal of merging. To this end, we introduce several designs for UniMedI, and *all of the designs aim at better merging 2D and 3D medical images into a language-guided common semantic space*. To best of our knowledge, we are the first one merging 2D and 3D medical images into a unified VL framework (we also provide a table to summarize the difference between our work and existing methods in this field at the end of this answer). Particularly, ACLIP cannot handle multi-modal images, and it can only be used for single modal images. So, *ACLIP cannot be directly applied to our problem*.


The first design is incorporating the guidance from report via VL contrastive learning. Generally, the $[CLS]$ token in the image side is optimized with representation of the paired report. And this image $[CLS]$ token is used in the following parts of UniMedI to ensure that the integrating is guided by language. We also have a novel motivating observation here for our problem (shown in Figure 1 and introduction). We see that despite big differences, medical images from various modalities share a common semantic latent space, which captures the underlying features of an individual's health status, and such status are reflected in medical reports via language. *This motivating observation is insightful in our problem and has not been mentioned in previous works as far as we know.*

The second design is our pipeline. Since no paired 2D and 3D image data exists, so we first select informative 2D slice from 3D volume as the bridge via the proposed attentive selection strategy. And then 2D slices, 2D X-ray and 3D volume are sent to the vision encoder together, enabling better merging. *ACLIP does not have this pipeline and the attentive selection strategy, since bridging multimodal images is not handled in ACLIP.* We have updated our paper by including a subsfigure in Figure 3 to show this pipeline.

The third design is the auxiliary task of mask and recovery implemented by self-distillation. Since 2D and 3D tokens are sent to the visual encoder at the same time and a large proportion of tokens are masked, this auxiliary task can help to enhance interactions among 2D and 3D tokens, facilitating merging representations. *The mask in ACLIP is used for removing redundant information, while the auxiliary task in UniMedI is enhancing dimensional interactions.*

Besides, the vision encoder has *separate tokenizer for 2D and 3D data and a shared backbone*, which is also designed for representation integration, and not used by ACLIP.

In summary, the problem we addressed in this paper is quite novel, and all designs in UniMedI aim at better solving the targeted problem. Maybe some losses are similar with other methods, but they are introduced with different purpose here, and thus brings different effects. And there are lots of difference compared with ACLIP, as we listed above.

| Method | Vision-Language | Unify | Downstream Task | Input Medical Image type |
| :------: | :---------------: | :-----: | :---------------: | :------------------------: |
|COVIRT [1] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|GLoRIA-MIMIC [2] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|MGCA [3] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
| Joint [4]   | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 2D slices  |
| UniMiss [5] | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 3D volumes, 2D slices |
| UniMedI     | $\checkmark$ | $\checkmark$| Classification, Segmentation, Detection(Appendix)          | **2D Images, 3D volumes, 2D slices** |

Reviewer comment (Presentation):
The description of the framework is not clear. It would be better if you could mark the (1)-(4) in Figure 3.

Author response:
Thank you for your valuable suggestion. We apologize if the description of our framework was unclear in the initial submission.  We have revised Figure 3 in our paper to give a clearer presentation without changing the content. The left of Figure 3 demonstrates the overall pipeline of UniMedI, and the right part of Figure 3 and Figure 4 present details on the modules used in the pipeline. Also, we have marked (1)-(4) in Figure 3 upon your request. To further address your concerns, we would like to briefly review UniMedI here.

### Pipeline of UniMedI
As we have described in the first paragraph of Section 3, to overcome the challenges that no paired 2D and 3D image data exists, during the training process, UniMedI first extracts a portion of informative 2D slices from a 3D volume according to language guidance via the proposed **attentive slice selection strategy (shown in Figure 4)**. This creates 2D-3D data pairs to bridge the disparity between the two modalities. Then, selected 2D slices and 2D X-rays, together with original 3D data, are sent to **a unified vision encoder (shown in Figure 3)** to obtain image representations. 

Overall, UniMedI is a vision-language pre-training framework. The paired report of each X-ray/CT is encoded by the text encoder, and VL contrastive learning (shown in Figure 3) are employed to obtain a language-guided common semantic space for both 2D and 3D medical images.

### Key module 1: the unified vision encoder (Figure 3)

For the vision encoder, first, 2D data (selected 2D slices and 2D X-rays) and 3D data (original 3D volumes) are processed by 2D tokenizer $T_{2D}$ and 3D tokenizer $T_{3D}$, respectively. Seperate tokenizers encode input data from different modalities into a series of tokens. 3D CT tokens and corresponding 2D slices tokens are concatenated together for further procesing by a shared backbone. *Note here input contains both 2D and 3D data, so we have 2D and 3D tokens at the same time.* 2D X-rays tokens directly processing by a shared backbone.

Then, to enhance the cross-dimensional communication among 2D and 3D tokens, we adopt **an auxiliary task, i.e., mask and recovery**, and this task is implemented by the **self-distillation method**. Specifically, our visual encoder contains the teacher network $\overline{E}_v$ and the student network $E_v$, where $\overline{E}_v$ is updated by exponential moving averaged (EMA) over $E_v$. Tokens, once encoded by the tokenizer, are input into the teacher network. This process yields global and local embeddings and simultaneously outputs attention scores for each patch of the X-rays or CT scans, as outlined in **Equation 1** in main text. Depending on the scores of each patch, we mask the less significant areas (those with lower scores) for the input into the student network. Since **the student network only inputs important areas of the medical images**, the resulting visual encoding is more closely matched with the semantics included in the report. Finally, loss of the self-distillation method is applied to both the global $[CLS]$ token (i.e., $L_{icl}$) and local patch tokens (i.e., $L_{pcl}$).

### Key module 2: the attentive slices selection strategy (Figure 4)

The aim of attentive slice selection is selecting informative 2D slices from the 3D CT volume. The 3D CT volume is processed by the tokenizer $T_{3D}$ and the teacher network, and thus, attention scores for patches in the CT can be obtained. Then, the scores of all tokens within one slice are averaged to derive the score of the corresponding slice. This process is referred to as the *Inter-Slice Average*, as outlined in **Equation 2** in main text. Slices with top $k$ score are selected as the connection between 2D and 3D data.

Note that these scores are attention weights between the $[CLS]$ token and other patch tokens, and the $[CLS]$ token is used for VL contrastive learning. Therefore, this slice selection process is guided by text information in the report, and 2D slices containing important information in the report can be selected accordingly.

Reviewer comment (Novelty):
Based on the methodology description, it seems that the proposed framework is a combination of the existing methods except for the attentive slice selection. The motivation for introducing the auxiliary task is not clear. Directly saying “inspired by…” is not a good writing style. Please explicitly present the main motivation of all your methodology components.

Author response:
To better explain the motivation of designs in our method, we first elaborate on what we have done in this paper. As we have stated in the introduction, we target expanding the data scale for medical vision language (VL) pre-training. Here, one obvious and important problem is that medical images have different dimensions. Therefore, we propose *UniMedI*, an unified VL framework capable of handling various medical modalities (e.g. 2D X-rays, 3D CT). However, Figure 2 shows that merging representations of different medical image modalities into a language-guided common semantic space is challenging. Naively using language guidance and VL contrastive learning (Figure 2a) and simply unifying in one model (Figure 2b) cannot achieve the goal of merging. To this end, we introduce several designs for UniMedI, and **all of the designs aim at better merging 2D and 3D medical images into a language-guided common semantic space**. To best of our knowledge, we are the first one merging 2D and 3D medical images into a unified VL framework (we also provide a table to summarize the difference between our work and existing methods in this field at the end of this answer).

The first design is incorporating the guidance from report via VL contrastive learning. Generally, the $[CLS]$ token in the image side is optimized with representation of the paired report. And this image $[CLS]$ token is used in the following parts of UniMedI to ensure that the integrating is guided by language.

The second design is our pipeline. As we have explained in the last question, since no paired 2D and 3D image data exists, so we first select informative 2D slice from 3D volume as the bridge via the attentive selection strategy. And then 2D slices, 2D X-ray and 3D volume are sent to the vision encoder together, enabling better merging.

The third design is the auxiliary task of mask and recovery implemented by self-distillation. Since 2D and 3D tokens are sent to the visual encoder at the same time and a large proportion of tokens are masked, this auxiliary task can help to enhance interactions among 2D and 3D tokens, facilitating merging representations.

Besides, the vision encoder has separate tokenizer for 2D and 3D data and a shared backbone, which is also designed for representation integration.

In summary, the problem we addressed in this paper is quite novel, and all designs in UniMedI aim at better solving the targeted problem. Maybe some losses is similar with other methods, but they are introduced with different purpose here, and thus brings different effects. Moreover, we have changed the writing upon your suggestion in our paper.





| Method | Vision-Language | Unify | Downstream Task | Input Medical Image type |
| :------: | :---------------: | :-----: | :---------------: | :------------------------: |
|COVIRT [1] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|GLoRIA-MIMIC [2] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
|MGCA [3] | $\checkmark$ | --- | Classification, Segmentation, Detection | 2D Images |
| Joint [4]   | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 2D slices  |
| UniMiss [5] | ---| $\checkmark$| Classification, Segmentation  | 2D Images, 3D volumes, 2D slices |
| UniMedI     | $\checkmark$ | $\checkmark$| Classification, Segmentation, Detection(Appendix)          | **2D Images, 3D volumes, 2D slices** |

Reviewer comment (Experiments):
The medical classification experiments are weak since most selected datasets are old and many other dedicated methods already achieved great performance. Please use the latest RSNA datasets for validation as well.

Author response:
Thank you for your feedback. Actually, we **have already use the latest RSNA datasets** (specifically, stage 2 data but not stage 1) for evaluation. The stage 2 dataset is larger than stage 1, and is adopted by latest previous works [2,3]. We apologize for the confusion and have revised the manuscript by explicitly saying that we use the stage 2 data of RSNA in Section 4.2.

Besides the classification experiments, we also add the **segmentation** (Dice [%]) and **detection** tasks (mAP [%]) and compared them with the most advanced method MGCA, as shown in the following tables.


Segmentation:
|         | 1%   | 10%  | 100% |
| ------- | ---- | ---- | ---- |
| MGCA [3]    | 66.2 | 71.3 | 73.6 |
| UniMedI | **67.8** | **73.1** | **75.3** |

Detection:
|         | 1%   | 10%  | 100% |
| :-------: | :----: | :----: | :----: |
| ConVIRT [1] | 8.2 | 5.6 | 17.9 |
| GLoRIA [2] | 9.8 |  14.8  | 18.8 |
|GLoRIA-MIMIC [2] | 11.6 | 16.1 | 24.8 |
| MGCA (ResNet-50) [3] | 12.9 | 16.8 | 24.9 |
| MGCA (ViT-B) [3] | 14.7 | 18.4 | 25.8 |
| UniMedI (ViT-B) | **15.5**  | **19.2**  | **26.6**  |

Above results are added to Appendix C.1 to further demonstrate the effectiveness of our approach.

Reviewer comment (Experiments):
For the 3D segmentation tasks, the small BCV dataset cannot provide statistically significant results (other previous works also used this BCV dataset is not a good excuse). Please use larger datasets like MICCAI FLARE and AMOS.

Author response:
In response to your suggestion, we add an experiment using a larger dataset AMOS. Results are listed below and this experiment is added to Apppendix C.2. We believe that this result can demonstrate the superiority of UniMedI with statistical significance.

|         | 20%   | 40%  | 100% |
| ------- | ---- | ---- | ---- |
| UniMiss [5]    | 79.5 | 82.3 | 85.8 |
| UniMedI | **78.8** | **82.9** | **86.4** |



As for MICCAI FLARE, this dataset contains images of many other body parts besides the chest, but chest is the main focus of our pre-trained model due to that publicly available paired image-report datasets are mainly about chest. Adapting a model pre-trained with chest data to images of other body parts is non-trivial and beyond the scope of this paper. We may leave it as our future work.

Reviewer comment (Experiments):
The authors designed an attentive slice selection method to select the disease relevant 2D slices from 3D volume. What’s the accuracy of the method? It is not validated yet but it could be easily done by testing it on some lung/abdomen tumor CT datasets that have tumor segmentation masks, e.g., check whether the selected slices cover all the tumors and compute the accuracy.

Author response:
Thanks for this constructive feedback, we have conducted an additional experiment following your suggestion. Specifically, we utilize the subset0 of LUNA16 dataset, a standard dataset in LUNA Nodule Detection, to verify whether the selected slices encompass the nodules. We filter the original CTs from 89 to 67 by selecting nodules larger than 3 mm. The input size of CT is 32 $\times$ 224 $\times$ 224. We see that if the nodules are included in the slices UniMedI selected. A continuous set of 32 slices that contain nodules are selected for this performance validation. We follow the LUNA challenge guidelines that a nodule is deemed to be within three slices near the central coordinate. We selected the top 20% of slices with the highest attentive scores.

The final accuracy of UniMedI is 76.1%. For reference, accuracy of random choosing is 53.73%. Although the absolute accuracy value is not pretty high, the improvement compared with the random model is significant. Note that (1) no explicitly direct supervision regarding nodules and its position is used in UniMedI, and we only use the report and no extra processing for reports is involved; (2) no specific and direct learning strategy is used in UniMedI for extracting slices with nodules (attentive slice selection is an indirect one because it emphasize report information but not exact nodule slices). The selection is based on attention weights, which is automatically learned in an end-to-end way. In other words, this selection ability is a kind of "free lunch" brought by UniMedI, so it might not be fair if it is compared with some specifically designed methods.

Reviewer comment (Writing):
In Sec 3., “selected 2D input is fed into the network along with the 3D tokens…”. You already convert the 3D images into 2D key slices. What are the 3D tokens here?

Author response:
The 3D tokens are generated by feeding 3D volume into the vision encoder and being processed by the 3D tokenizer. After 2D slices are selected from 3D volume, we have 2D and 3D data as input at the same time, so with separate tokenizers, 2D and 3D tokens exists in one pass. In this way, we can use the mask and recovery task to enhance dimensional interactions. Sorry for causing confusion. We have revised Figure 3 by incorporating a subfigure to illustrate the overall pipeline, and adjust corresponding descriptions, making this point clear.

Reviewer comment (Novelty):
Please explicitly clarify the unique motivations of EMA and auxiliary tasks rather than just saying "inspired by".

Author response:
EMA is a technique used in the self-distillation method, to accomplish the mask and recovery auxiliary task. And the motivation to this task is better merging presentations of 2D and 3D medical images, as we have explained in details in the response to **W2**. The way to use EMA has been explained in the response to **W1**, "Key module 1: the unified vision encoder (Figure 3)" part. We have revised all "inspired by" sentences to make our statement precise.

Reviewer comment (Writing):
Minors: Please polish the writing: Page 5. Sec. 3.3 “In Section 3.1…. Then in Section 3.1…”

Author response:
Fixed!

Reviewer comment (Experiments):
The motivation of the paper, bridging features in 2D and 3D, is not validated (see Q1).

Author response:
First, in the revised version, we have highlighted the differences among each subfigures of Figure 2 as much as possible. In 2(b), overall, the 2D X-ray representations are distributed in the upper part of the figure, while the 3D CT representations are distributed in the lower part of the figure. In 2\(c\), representations of 2D and 3D images are generally merged together and closer compared with 2(b). In summary, Figure 2 shows that naively using language guidance (2a) and simply unifying in one model (2b) cannot achieve the our challenging goal, which is using text information in reports as the guidance to map data from various medical image modalities into the shared semantic space. In contrast, UniMedI can better integrate different medical image modalities.

Also, to further address your concerns, **we have visualized data representations of two different classes (i.e., pneumonia and cardiomegaly) in Figure 6 in Appendix C.4**. As can be seen, *even data coming from different classes, UniMedI can bring together representations for medical data of different modalities, while not making features of all classes similar, well preserving semantic information from reports in the representation space*.

We really agree with the reviewer on the viewpoint that discrepancy between 2D and 3D modalities could be much larger than that among classes. Therefore, Figure 6 in Appendix address this point and nicely demonstrate the effectiveness of UniMedI on integrating medical multi-modal images into a language-guided common semantic space, since Figure 6 shows semantic meaning is preserved while modalities are merged in the representation space, as we expected.

Reviewer comment (Writing):
The details of the method are confusing (see Q2).

Author response:
We appreciate the reviewer's suggestion to depict our method more clearly. Firstly, we would like to elaborate more on the pipeline of UniMedI to make the whole process clear. Based on the pipeline, we further explain contents in Figure 3 and Figure 4 clearly. After that, we address your question point-to-point in order to clarify any confusion. Lastly, **we have adjusted Figures 3 and 4 without changing the content in the updated version of our paper, including clearer annotations and some accompanying explanations**. We believe these enhancements will greatly improve the clarity and comprehensibility of our framework.

### Pipeline of UniMedI
As we have described in the first paragraph of Section 3, to overcome the challenges that no paired 2D and 3D image data exists, during the training process, UniMedI first extracts a portion of informative 2D slices from a 3D volume according to language guidance via the proposed **attentive slice selection strategy (shown in Figure 4)**. This creates 2D-3D data pairs to bridge the disparity between the two modalities. Then, selected 2D slices and 2D X-rays, together with original 3D data, are sent to **a unified vision encoder (shown in Figure 3)** to obtain image representations. 

Overall, UniMedI is a vision-language pre-training framework. The paired report of each X-ray/CT is encoded by the text encoder, and VL contrastive learning (shown in Figure 3) are employed to obtain a language-guided common semantic space for both 2D and 3D medical images. To make this pipeline clear, we add a subfigure to Figure 3 for illustration.

### Figure 3 shows designs in the unified visual encoder and VL contrastive learning
For the vision encoder, first, 2D data (selected 2D slices and 2D X-rays) and 3D data (original 3D volumes) are processed by 2D tokenizer $T_{2D}$ and 3D tokenizer $T_{3D}$, respectively. Separate tokenizers encode input data from different modalities into a series of tokens. 3D CT tokens and corresponding 2D slices tokens are concatenated together for further processing by a shared backbone. 2D X-rays tokens directly processing by a shared backbone.

Then, to enhance the cross-dimensional communication among 2D and 3D tokens, we adopt **an auxiliary task, i.e., mask and recovery**, and this task is implemented by the **self-distillation method**. Specifically, our visual encoder contains the teacher network $\overline{E}_v$ and the student network $E_v$, where $\overline{E}_v$ is updated by exponential moving averaged (EMA) over $E_v$. Tokens, once encoded by the tokenizer, are input into the teacher network. This process yields global and local embeddings and simultaneously outputs attention scores for each patch of the X-rays or CT scans, as outlined in **Equation 1** in main text. Depending on the scores of each patch, we mask the less significant areas (those with lower scores) for the input into the student network. Since **the student network only inputs important areas of the medical images**, the resulting visual encoding is more closely matched with the semantics included in the report. Finally, loss of the self-distillation method is applied to both the global $[CLS]$ token (i.e., $L_{icl}$) and local patch tokens (i.e., $L_{pcl}$).

### Figure 4 shows the attentive slice selection strategy
The aim of attentive slice selection is selecting informative 2D slices from the 3D CT volume. The 3D CT volume is processed by the tokenizer $T_{3D}$ and the teacher network, and thus, attention scores for patches in the CT can be obtained. Then, the scores of all tokens within one slice are averaged to derive the score of the corresponding slice. This process is referred to as the *Inter-Slice Average*, as outlined in **Equation 2** in main text. Slices with top $k$ score are selected as the connection between 2D and 3D data.

Note that these scores are attention weights between the $[CLS]$ token and other patch tokens, and the $[CLS]$ token is used for VL contrastive learning. Therefore, this slice selection process is guided by text information in the report, and 2D slices containing important information in the report can be selected accordingly.

>>2.a. What are the light orange boxes used for?

In Figure 3, the orange boxes denote tokens of medical texts encoded by the language encoder $E_l$. The dark orange box represents the $[CLS]$ token, and the remaining light orange boxes represent all other tokens.

In Figure 4, the orange box highlights a slice and the Inter-Slice Average process (explained in the above descriptions for Figure 4). In essence, the function of the orange box is to provide a clear understanding that this process occurs within a slice.

>>2.b. What is the operation after T_3D and T_2D?

Details are explained in the above descriptions for Figure 3. The merged line after $T_{2D}$ and $T_{3D}$ represents that 2D and 3D tokens are concatenated for further procesing by a shared backbone.


>>2.c. The main text in the method section is not corresponding to Figure 3 or Figure 4. It is unclear how Equations (1-2) were integrated into the framework.

Equation 1 is utilized to produce attention scores for each patch in the teacher network. Equation 2 calculates the average attention score for patches in one slice, and this average score is used for slice selection. Refer to the above descriptions for Figure 3 and 4 for details.

Reviewer comment (Experiments):
Lack of baseline methods for 2D and 3D pre-training (see Q3).

Author response:
Thanks for listing these works for us. We now provide comparisons and discussions to these work.

Most listed models [1, 2, 4, 5] are trained via self-supervised learning (SSL). Comparing SSL models with the model by vision-language (VL) pre-training is not very reasonable, since VL pre-training is dealing with some multi-modality problem, which is quite different with the SSL problem. Moreover, UniMedI also handles multimodal medical images at the same time. Therefore, the comparison might not be as fair as what we listed in the paper. Nevertheless, we provide compared results in the below table (CheXpert$^\ast$ denoting full backbone fine-tuning and CheXpert$^\diamond$ representing linear classification, only fine-tuning the classification head). Here, we additionally include some latest SOTA SSL methods in the general domain, further demonstrating the superiority of our method.


| Method | Perfomance | Description |
| :--------: | :--------: | :--------: |
| SimMIM [1]    | 88.07 (CheXpert$^\ast$, 100%) | Benchmark of SSL in Medical Image Classification |
| SAM [2] | --- | Designed specific backbones for 2D and 3D data, only focus on landmark detection, Lesion matching |
| Swin UNETR[4] | --- | Designed for specific task,  semantic segmentation of brain tumors in MRI Images |
| DiRA [5] | 87.59 (CheXpert$^\ast$, 100%) | Unites discriminative, restorative, and adversarial learning in a unified manner |
| Moco v2 [7]     |  84.90 (CheXpert$^\diamond$, 100%)    | General domain SSL method     |
| Moco v3 [8]     |  82.02 (BCV, 100%) | General domain SSL method |
| DINO [9]        |    82.61(BCV, 100%) | General domain SSL method |
| UniMedI | **90.50** (CheXpert$^\diamond$, 100%), **85.40** (BCV, 100%) |---|


It is evident that UniMedI outperforms all other methods, even without fine-tuning for the evaluation on CheXpert. Other methods not included in the above table are specially designed for some downstream tasks. SAM [2] is for landmark detection and Lesion matching with designed backbone, and it use separate models for 2D and 3D data, respectively. Swin UNETR [4] is designed for semantic segmentation of brain tumors in MRI Images, which cannot be applied to the same downstream tasks as our method. KAD [3] is designed for classification tasks, incorporating external knowledge and a classification head related to the disease. Med3D [6] is for transfer learning with a coarse-to-fine structure. Comparing with these specially designed methods is not a fair setting. Besides, all these designs for specific tasks are orthogonal with our methods, and can be involved in our framework for further performance improvement. Unifying all these designs can be our future work.

Reviewer comment (Evaluation):
From Figure 1, I fail to observe (c) is particularly better than (b) in terms of feature quality. I understood that the authors try to present the features of pneumonia get closer between 2D and 3D images, but this needs features of another class (e.g., nodule) as a comparison. It is possible that the proposed UniMedI makes features of all classes similar. I think the discrepancy between 2D and 3D modalities could be much larger than that among classes.

Author response:
First, in the revised version, we have highlighted the differences among each subfigures of Figure 2 as much as possible. In 2(b), overall, the 2D X-ray representations are distributed in the upper part of the figure, while the 3D CT representations are distributed in the lower part of the figure. In 2\(c\), representations of 2D and 3D images are generally merged together and closer compared with 2(b). In summary, Figure 2 shows that naively using language guidance (2a) and simply unifying in one model (2b) cannot achieve the our challenging goal, which is using text information in reports as the guidance to map data from various medical image modalities into the shared semantic space. In contrast, UniMedI can better integrate different medical image modalities.

Also, to further address your concerns, **we have visualized data representations of two different classes (i.e., pneumonia and cardiomegaly) in Figure 6 in Appendix C.4**. As can be seen, *even data coming from different classes, UniMedI can bring together representations for medical data of different modalities, while not making features of all classes similar, well preserving semantic information from reports in the representation space*.

We really agree with the reviewer on the viewpoint that discrepancy between 2D and 3D modalities could be much larger than that among classes. Therefore, Figure 6 in Appendix address this point and nicely demonstrate the effectiveness of UniMedI on integrating medical multi-modal images into a language-guided common semantic space, since Figure 6 shows semantic meaning is preserved while modalities are merged in the representation space, as we expected.

Reviewer comment (Presentation):
The illustration of Figure 3 and Figure 4 is unclear to me. What are the light orange boxes used for? What is the operation after T_3D and T_2D? The main text in the method section is not corresponding to Figure 3 or Figure 4. It is unclear how Equations (1-2) were integrated into the framework.

Author response:
We appreciate the reviewer's suggestion to depict our method more clearly. Firstly, we would like to elaborate more on the pipeline of UniMedI to make the whole process clear. Based on the pipeline, we further explain contents in Figure 3 and Figure 4 clearly. After that, we address your question point-to-point in order to clarify any confusion. Lastly, **we have adjusted Figures 3 and 4 without changing the content in the updated version of our paper, including clearer annotations and some accompanying explanations**. We believe these enhancements will greatly improve the clarity and comprehensibility of our framework.

### Pipeline of UniMedI
As we have described in the first paragraph of Section 3, to overcome the challenges that no paired 2D and 3D image data exists, during the training process, UniMedI first extracts a portion of informative 2D slices from a 3D volume according to language guidance via the proposed **attentive slice selection strategy (shown in Figure 4)**. This creates 2D-3D data pairs to bridge the disparity between the two modalities. Then, selected 2D slices and 2D X-rays, together with original 3D data, are sent to **a unified vision encoder (shown in Figure 3)** to obtain image representations. 

Overall, UniMedI is a vision-language pre-training framework. The paired report of each X-ray/CT is encoded by the text encoder, and VL contrastive learning (shown in Figure 3) are employed to obtain a language-guided common semantic space for both 2D and 3D medical images. To make this pipeline clear, we add a subfigure to Figure 3 for illustration.

### Figure 3 shows designs in the unified visual encoder and VL contrastive learning
For the vision encoder, first, 2D data (selected 2D slices and 2D X-rays) and 3D data (original 3D volumes) are processed by 2D tokenizer $T_{2D}$ and 3D tokenizer $T_{3D}$, respectively. Separate tokenizers encode input data from different modalities into a series of tokens. 3D CT tokens and corresponding 2D slices tokens are concatenated together for further processing by a shared backbone. 2D X-rays tokens directly processing by a shared backbone.

Then, to enhance the cross-dimensional communication among 2D and 3D tokens, we adopt **an auxiliary task, i.e., mask and recovery**, and this task is implemented by the **self-distillation method**. Specifically, our visual encoder contains the teacher network $\overline{E}_v$ and the student network $E_v$, where $\overline{E}_v$ is updated by exponential moving averaged (EMA) over $E_v$. Tokens, once encoded by the tokenizer, are input into the teacher network. This process yields global and local embeddings and simultaneously outputs attention scores for each patch of the X-rays or CT scans, as outlined in **Equation 1** in main text. Depending on the scores of each patch, we mask the less significant areas (those with lower scores) for the input into the student network. Since **the student network only inputs important areas of the medical images**, the resulting visual encoding is more closely matched with the semantics included in the report. Finally, loss of the self-distillation method is applied to both the global $[CLS]$ token (i.e., $L_{icl}$) and local patch tokens (i.e., $L_{pcl}$).

### Figure 4 shows the attentive slice selection strategy
The aim of attentive slice selection is selecting informative 2D slices from the 3D CT volume. The 3D CT volume is processed by the tokenizer $T_{3D}$ and the teacher network, and thus, attention scores for patches in the CT can be obtained. Then, the scores of all tokens within one slice are averaged to derive the score of the corresponding slice. This process is referred to as the *Inter-Slice Average*, as outlined in **Equation 2** in main text. Slices with top $k$ score are selected as the connection between 2D and 3D data.

Note that these scores are attention weights between the $[CLS]$ token and other patch tokens, and the $[CLS]$ token is used for VL contrastive learning. Therefore, this slice selection process is guided by text information in the report, and 2D slices containing important information in the report can be selected accordingly.

>>2.a. What are the light orange boxes used for?

In Figure 3, the orange boxes denote tokens of medical texts encoded by the language encoder $E_l$. The dark orange box represents the $[CLS]$ token, and the remaining light orange boxes represent all other tokens.

In Figure 4, the orange box highlights a slice and the Inter-Slice Average process (explained in the above descriptions for Figure 4). In essence, the function of the orange box is to provide a clear understanding that this process occurs within a slice.

>>2.b. What is the operation after T_3D and T_2D?

Details are explained in the above descriptions for Figure 3. The merged line after $T_{2D}$ and $T_{3D}$ represents that 2D and 3D tokens are concatenated for further procesing by a shared backbone.


>>2.c. The main text in the method section is not corresponding to Figure 3 or Figure 4. It is unclear how Equations (1-2) were integrated into the framework.

Equation 1 is utilized to produce attention scores for each patch in the teacher network. Equation 2 calculates the average attention score for patches in one slice, and this average score is used for slice selection. Refer to the above descriptions for Figure 3 and 4 for details.

Reviewer comment (Experiments):
As a paper for pre-training, the authors did not compare with the many existing pre-trained models, neither 3D or 2D. For 2D X-ray imaging, there are many pre-trained models publicly available [e.g., 1-3]. For 3D CT imaging, a variety of public models are also missing [e.g., 4-6].

Author response:
Thanks for listing these works for us. We now provide comparisons and discussions to these work.

Most listed models [1, 2, 4, 5] are trained via self-supervised learning (SSL). Comparing SSL models with the model by vision-language (VL) pre-training is not very reasonable, since VL pre-training is dealing with some multi-modality problem, which is quite different with the SSL problem. Moreover, UniMedI also handles multimodal medical images at the same time. Therefore, the comparison might not be as fair as what we listed in the paper. Nevertheless, we provide compared results in the below table (CheXpert$^\ast$ denoting full backbone fine-tuning and CheXpert$^\diamond$ representing linear classification, only fine-tuning the classification head). Here, we additionally include some latest SOTA SSL methods in the general domain, further demonstrating the superiority of our method.


| Method | Perfomance | Description |
| :--------: | :--------: | :--------: |
| SimMIM [1]    | 88.07 (CheXpert$^\ast$, 100%) | Benchmark of SSL in Medical Image Classification |
| SAM [2] | --- | Designed specific backbones for 2D and 3D data, only focus on landmark detection, Lesion matching |
| Swin UNETR[4] | --- | Designed for specific task,  semantic segmentation of brain tumors in MRI Images |
| DiRA [5] | 87.59 (CheXpert$^\ast$, 100%) | Unites discriminative, restorative, and adversarial learning in a unified manner |
| Moco v2 [7]     |  84.90 (CheXpert$^\diamond$, 100%)    | General domain SSL method     |
| Moco v3 [8]     |  82.02 (BCV, 100%) | General domain SSL method |
| DINO [9]        |    82.61(BCV, 100%) | General domain SSL method |
| UniMedI | **90.50** (CheXpert$^\diamond$, 100%), **85.40** (BCV, 100%) |---|


It is evident that UniMedI outperforms all other methods, even without fine-tuning for the evaluation on CheXpert. Other methods not included in the above table are specially designed for some downstream tasks. SAM [2] is for landmark detection and Lesion matching with designed backbone, and it use separate models for 2D and 3D data, respectively. Swin UNETR [4] is designed for semantic segmentation of brain tumors in MRI Images, which cannot be applied to the same downstream tasks as our method. KAD [3] is designed for classification tasks, incorporating external knowledge and a classification head related to the disease. Med3D [6] is for transfer learning with a coarse-to-fine structure. Comparing with these specially designed methods is not a fair setting. Besides, all these designs for specific tasks are orthogonal with our methods, and can be involved in our framework for further performance improvement. Unifying all these designs can be our future work.