A Textbook Remedy for Domain Shifts: Knowledge Priors for Medical Image Analysis

Thank you for your positive feedback on our work! Here, we address your concern about applications in 3D images.

We agree with you that 3D modalities are important. However, we scoped our study to 2D modalities as these are cheaper in practice. Part of the motivation is to transfer to less-resourced hospitals where X-rays are generally available but more expensive imaging modalities are not (i.e., in developing countries). That being said, we are very concerned with establishing the broad applicability of our work, and so we evaluated 20 datasets within the modalities we considered.

The KnoBo framework is flexible enough to adapt to 3D modalities because the key contribution is the factorization. Many components of KnoBo can be readily reused, such as the concept generation from medical documents. We only need an appropriate 3D encoder (e.g., 3D CNN) and a pretraining dataset. Fortunately, large-scale public 3D medical datasets are available, e.g., DeepLesion [1] has 32,000 annotated CT scans, CT-RATE [2] has 50,188 CT volumes with associated clinical reports, and the recently released CLIP-CT [2] could serve as a backbone. We agree that 3D would enrich the work, but we feel we have provided significant evidence already that KnoBo is effective and that there is a clear path on how one could apply it to 3D domains in the future.

[1] Yan et al. DeepLesion: Automated Deep Mining, Categorization and Detection of Significant Radiology Image Findings using Large-Scale Clinical Lesion Annotations. 2017.

[2] Hamamci et al. A foundation model utilizing chest CT volumes and radiology reports for supervised-level zero-shot detection of abnormalities. 2024.

Hi,

Do you feel we have provided enough detail about possible extensions of our work to 3D? We will add some of this discussion to the limitations and future work section given extra space. Would you consider updating your rating or are there other weaknesses that we can discuss now during the discussion period?

Thanks!

Thanks for the feedback! We will try to include more examples in future versions. We hope the answers below clarify:

Q1. What are the tasks?

We studied the medical image classification tasks in a confounded setting (the medical class names are found in Tables 5 and 6). As explained in the second paragraph of the introduction (Lines 30-34), the samples are confounded with different factors during training/validation (ID) and testing (OOD) time. For example, in an X-ray classification task to classify COVID and normal patients, if we confound the data with sex, the dataset is constructed as follows:

Train/Validation (ID): (COVID, male) (Normal, female)

Test (OOD): (COVID, female) (Normal, male)

A robust model should predict based on COVID-related features rather than spurious correlations (sex in this example).

Q2. Why compare the effectiveness of $x$ and $x_p$?

This experiment aims to validate our motivation that deep neural networks lack priors for medical domains. $x$ is the feature extracted by a deep network at initialization (i.e., a ViT). $x_p$ is a feature vector constructed by extracting pixel values in the image. If a network has good priors for a domain, we expect the network output at initialization ($x$) to be a more useful representation of the image content than the pixels themselves ($x_p$). For natural images, this is known from prior work on Deep Image Priors, but, as we show in Figure 2, this is not true for medical images. Without good priors, models can easily adopt incorrect hypotheses and rely on spurious features. This leads to catastrophic failures when the spurious features are unavailable in out-of-domain tests and motivates the need to incorporate priors from elsewhere.

Q3. explain algorithm 1

Algorithm 1 aims to collect a set of concepts from medical documents. Given the class name (e.g., COVID), we want to extract useful knowledge from medical documents as concepts (e.g., ground glass opacity, a highly indicative feature of COVID). The whole process is a retrieval-augmented-generation task. We first use class names as queries (e.g., COVID) to retrieve relevant documents from a corpus (e.g., PubMed). The retrieved document may contain useful information, e.g., in the upper left of Figure 3, the document says, “the most frequent radiological patterns found were ground-glass opacities.” We feed these documents as contexts for an LLM, which extracts the information as a concept (ground glass opacity). We repeat this process using generated concepts as queries until a predefined number of concepts is reached.

Q4. Why doesn’t KnoBo always have the best ID performance?

This is because, in the confounded settings (explained in Q1), models can achieve high ID performance by taking shortcuts using spurious correlations, e.g., learning to classify males and females instead of focusing on diseases. Models that achieve high ID performance in this way will have catastrophic failures on the OOD test (see OOD numbers of ViT-L/14 and DenseNet in Table 1). In this case, a robust model must have a high average performance across ID and OOD. There is a natural tradeoff between ID and OOD, but our models strike the best balance.

Q5. Is the Structure Prior learnt? The connection between “Is there ground-glass opacity?” and the top-right red box.

As explained in Q3, the structure prior is not learned from training instances but instead constructed from medical documents.

The top-right red box illustrates the concept grounding module (Sec 4.3, Lines 180-188). Given a concept (e.g., is there ground-glass opacity?) and an X-ray image, we aim to estimate the probability of this concept existing on this X-ray.

To achieve this goal, We train a binary classifier for each concept. We obtain the positive and negative examples for each concept by leveraging a pretraining dataset with (X-ray, clinical report) pairs. We can estimate the presence of a concept in each example by prompting an LLM to generate a response indicating whether the clinical report implies the concept. For example, if the clinical report mentioned related terms about “ground-glass opacity,” it is a positive example of this concept. Otherwise, it will be a negative example. With positive and negative examples for a concept, we can train a binary classifier to predict the concept given an X-ray image.

Q6. Is the training end-to-end?

No. To be computationally feasible, the training has stages: (1) constructing the prior from documents (explained in Q3), (2) concept grounding (explained in Q5), and (3) linear layer for final label prediction. Given an X-ray image, the workflow is: (1) execute the concept grounding functions to get the probabilities on all concepts in the structure prior, (2) use the concept probabilities as the input to a linear layer to get the final label prediction, e.g., predict COVID or Normal. Please also refer to Lines 142-144 for Preliminary on Concept Bottleneck Model and Sec 4.3 for details about the workflow.

Q7. Is the parameter prior a part of the bottleneck predictor?

No. As explained in Sec 4.4, the parameter prior is used to regularize the learning of the linear layer. We will compute the L1 distance between the linear layer and the parameter prior as the loss to align the learned parameters with priors.

Q8. How to determine the best number of concepts?

We did not tune this hyperparameter much and used 150 concepts. In Figure 4, we examine the impact of the number of concepts on ID and OOD performance. More concepts lead to better ID but worse OOD so that an appropriate trade-off can be selected for a problem.

Q9. How did you determine the preferred correlation between the label y and concept c?

A language model (explained on Line 198) determines the preferred sign of the correlation between y and c.

Please do not hesitate to ask more questions, and we will be happy to answer them. We look forward to your response!

Hi,

Has our rebuttal clarified sufficiently that you feel you can judge our work? Are there other aspects that we can clarify for you to make an assessment now that we have provided answers to your question?

Thanks!

Thank you for your constructive feedback. Here we address your comments:

Q1. The concept grounding seems pretty data-hungry.

This is a valid concern, but while we sample from a large dataset like MIMIC, in reality, we only use 22k total examples. Training the concept grounding component is just learning a linear classifier, so we sampled a small subset of the data.

On average, each X-ray grounding function uses 1,523 examples, and each skin lesion grounding function uses 1,342 examples. This can be reduced significantly with little loss in performance.

The table below shows the performance of KnoBo on X-ray datasets using varied sizes of samples for training each grounding function. It highlights that KnoBo achieves much of its performance with a small number of examples, suggesting that it is not data-hungry.

Q2. For skin lesion datasets KnoBo only brings marginal overall improvement.

Empirical challenges within the skin lesions datasets are not because of a lack of data but because the datasets have very strong non-generalizable cues (skin colors, hairs, etc). For this reason, it is easy to fine-tune models to high in-domain performance and learn non-transferable predictors.

Our overall metric averages both in-domain performance and out-of-distribution performance. KnoBo is much better out of distribution (see OOD and Avg in Table 2, shows differences of +22.9 and +6.7, respectively), but when averaging in domain performance on datasets with no explicitly identified confounds (unconfd), helps more moderately (+.3). We believe that OOD measures more accurately capture if models are able to learn the correct hypothesis and therefore the more important measure of if a model is trustworthy in real applications.

Q3. Confusing results in Tables 2 and 4.

In Figure 4, we deactivated the prior loss for KnoBo, which resulted in lower OOD performance. In this ablation, we focus on comparing different inputs (black-box features vs. concept scores) and their sizes while removing prior loss forms a more comparable setting. We provide the updated Figure 4 in the PDF of the global response, which includes the curves of KnoBo with prior loss. We will clarify this experimental setup in the final draft of our paper.

Q4. Failure cases and limitation section.

Failure cases are difficult to examine and generally not very accessible to non-medical professionals (in the main text, we tried to avoid adding examples that might make some readers uncomfortable if they are squeamish about medical images). Sections C.4 (Table 11) and Section D (Figures 8 and 9) of the appendix contain some of this, where we recruited medical students to evaluate limitations.

The medical students annotated 300 concepts and often found the information in our bottlenecks highly relevant, but some were not groundable (see Table A in the PDF for examples). This did not vary much by information source, but skin lesion information was judged to be easier to ground than X-ray information. In a possible camera-ready version, given the extra page, we can try to integrate some more of this content in the main body of text, including example failures, and expand the limitations section to more explicitly discuss the data efficiency of our methods.

Q5.Typos.

We have carefully proofread the manuscript again to correct the issues.

Hi,

Has our rebuttal addressed the weakness you felt the paper had? In particular, we strongly feel our method is not data-hungry and we have a fairly extensive discussion of problems (currently in the appendix) that could be moved up given extra space. Would you consider updating your rating or are there other weaknesses you feel we can address during this discussion period?

Thank you for your positive feedback on our work! Here is our reply to your questions:

Q1. Hallucinations of GPT-4

This is a valid concern, but the way we use GPT-4 highly encourages it to directly copy information from documents instead of inventing it. Our concept generation is conditioned on medical documents, which differs from previous work [1] that solely relies on LLM. We employ LLM as a text extraction tool, which basically copy-paste or paraphrases. In addition, each concept is attributed to a document that medical professionals can fact-check. Table 12 of the Appendix and supplementary materials (HTML files in the example_concepts/ folder) show example concepts with corresponding reference documents.

We examined concepts from the X-ray bottleneck. 80% of the concepts were word-for-word copies of a feature in a sentence that the authors of a document said they examined. 14% included some paraphrases but were still grounded in the document and never completely invented (sometimes a combination of multiple factors in different places in the document). 6% were of confusing origin and potentially hallucinated but still mentioned relevant factors. This aligns with the analysis of the bottlenecks we did with medical students in section C.4 of the appendix.

Q2. Robustness measure ($\Delta$)

Our analysis is focused on multiple metrics at once. The delta measure $\Delta$ needs to be combined with the other metrics (ID, OOD) when measuring models’ robustness (it is trivially minimized by creating classifiers with chance performance). Our perspective is that a robust model should have high performance across ID and OOD while maintaining a small drop to demonstrate its robustness to domain shifts.

Q3. Diversity values for skin lesion datasets are much lower than those for X-rays

This is because X-ray is a more information-dense modality as it shows multiple organs (e.g., lung, heart, stomach). Skin lesion diagnosis requires fewer medical factors in practice. Moreover, X-ray is the most accessible medical imaging modality, resulting in more available medical documents.

Q4. Filtering scheme of concept generation

Thank you for the suggestion! To clarify a bit, we don’t have a phase in our method that requires medical professionals, although we used them later to verify the validity of our approach. We designed three filtering criteria to ensure the concepts are diverse and visually identifiable. (Appendix B.3, Line 712-717) As you suggest, more filtering of the concepts using auxiliary signals would likely substantially improve the results and is something we are exploring.

[1] Yang et al. Language in a bottle: Language model guided concept bottlenecks for interpretable image classification. CVPR 2023.

I just want to add on here as I believe RAG is quite an effective and reasonable measure to prevent hallucinations by far. Plus, the number/stats they reported in the rebuttal seem pretty convincing.