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

Summary

This paper addresses the issue of deep networks' sensitivity to domain shifts in medical image analysis, particularly for chest X-rays and skin lesion images. The authors propose Knowledge-enhanced Bottlenecks (KnoBo), a concept bottleneck model that incorporates explicit medical knowledge from resources like textbooks and PubMed to improve generalization. KnoBo leverages retrieval-augmented language models to create a clinically relevant concept space and demonstrates a significant performance improvement—32.4% on average—over fine-tuned models across 20 datasets, highlighting PubMed as a particularly effective knowledge resource for mitigating domain shifts.

Strengths

Innovative Integration of Medical Knowledge: The paper introduces Knowledge-enhanced Bottlenecks (KnoBo), a novel approach that integrates explicit medical knowledge into deep networks. By using retrieval-augmented language models to incorporate clinically relevant factors from medical textbooks and PubMed, the model gains a robust prior that significantly enhances its generalization capabilities. Comprehensive Evaluation Across Diverse Domain Shifts: The authors conduct a thorough evaluation of KnoBo across 20 datasets and two imaging modalities (chest X-rays and skin lesion images). This extensive testing demonstrates the model's ability to handle various domain shifts, such as differences in data from different hospitals and demographic confounders, showcasing its practical applicability and robustness. Significant Performance Improvement: KnoBo shows substantial performance gains, outperforming fine-tuned models by an average of 32.4% on confounded datasets. This impressive improvement highlights the effectiveness of incorporating medical knowledge for reducing sensitivity to domain shifts, with evaluations indicating that PubMed is a particularly valuable resource for enhancing model performance and information diversity.

Weaknesses

Lack of Application to 3D Medical Imaging: Issue: The paper primarily focuses on 2D medical imaging modalities, such as chest X-rays and skin lesion images, without addressing the broader and more widely used domain of 3D medical imaging. Impact: This limitation raises concerns about the generalizability and applicability of the proposed KnoBo model to 3D imaging scenarios, such as CT scans or MRI, which are critical in many clinical contexts. Including evaluations on 3D medical images would significantly enhance the relevance and impact of the research. Insufficient Comparative Analysis with 3D Domain Adaptation Methods: Issue: There is a lack of comparative analysis with existing domain adaptation and generalization methods specifically applied to 3D medical imaging. Impact: The absence of such comparisons makes it difficult to gauge the relative effectiveness of KnoBo in a comprehensive manner, particularly given the significant advancements in 3D domain adaptation techniques. Addressing this gap would provide a more complete evaluation of KnoBo's capabilities and limitations.

Questions

Application to 3D Medical Imaging: Question: Have you considered extending the KnoBo framework to 3D medical imaging modalities such as CT scans or MRIs? If so, what challenges do you anticipate, and how might they be addressed? Suggestion: Expanding your research to include 3D medical images would significantly increase the applicability and impact of your work. Discussing potential adaptations and preliminary results for 3D imaging could provide valuable insights and demonstrate the versatility of KnoBo. Comparative Analysis with State-of-the-Art Methods: Question: How does KnoBo compare with recent state-of-the-art domain adaptation and generalization methods, especially those tailored for 3D medical imaging? Suggestion: Including a detailed comparative analysis with advanced domain adaptation techniques, particularly those used in 3D imaging, would strengthen your paper. This would provide a clearer benchmark and highlight the unique advantages and limitations of KnoBo.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

No

Summary

The authors noticed the domain-shift issue of current medical dataset, and after taking inspiration from medical training, they propose giving deep networks a prior grounded in explicit medical knowledge communicated in natural language. The proposed network can incorporate medical knowledge priors to help model make decisions. Basically, the authors assumed that “medical concepts” are robust enough to against the domain shift. Valid assumption, interesting work.

Strengths

The authors noticed the domain-shift issue of current medical dataset, and after taking inspiration from medical training, they propose giving deep networks a prior grounded in explicit medical knowledge communicated in natural language. The proposed network can incorporate medical knowledge priors to help model make decisions. Basically, the authors assumed that “medical concepts” are robust enough to against the domain shift. Valid assumption, interesting work.

Weaknesses

It is very HARD to read and interpret. Please raise examples for better reading experience. Maybe I have some questions after the clearance. Other points could be seeen under the Questions section.

Questions

1. “In 5 such constructed confounds per modality, covering scenarios of race, sex, age, scan position, and hospital, we find models unable to generalize well, dropping over 63% on average over an in-distribution (ID) evaluation.” What are the tasks? Please specify. 2. “As a baseline, we extract a subset of d pixels directly from the image as the feature without any model-based priors, represented as xp ∈ Rd. We compare the classification performance using x versus xp to probe the efficacy of the vision backbone’s priors.” I didn’t get the point. Why did you want to compare the effectiveness of x and x_p? 3. It is very hard to read the 4.2 section. I cannot fully understand the algorithm 1. Please raise examples for better reading experience. 4. Table 1, why the KnoBo always has not the best, or sometimes the worst performance in in-domain dataset? 5. Is the Structure Prior learnt? Or the questions like “Is there ground-glass opacity?” were pre-defined? How ist he connection between “Is there ground-glass opacity?” and the top-right red box in Figure 1? 6. Is the training end-to-end? Can you please provide an example for 1) how to pretrain using medical books 2) and given an X-ray image, how the workflow is to get the last result? I hold positive thought towards the study, but it is really hard to understand the workflow. 7. Is the parameter prior a part of the bottleneck predictor? In figure 1 it seems like parallel, but in the main text it seems like existing sequentially. Hard to understand. 8. How to determine the best number of concepts? 9. How did you determine the preferred correlation between the label y and concept c?

Rating

6

Confidence

5

Soundness

2

Presentation

2

Contribution

3

Limitations

N/A

Summary

The presented paper addresses the challenge of domain shifts in medical image classification, where conventional neural networks often lack effective priors for medical datasets. The authors introduce KnoBo, a novel class of concept bottleneck networks that integrate medical knowledge priors to enhance neural network performance in medical image classification tasks. KnoBo has three primary components: 1) Structure Prior: leverages medical documents to construct a knowledge bottleneck; 2) Bottleneck Predictor: maps images onto concepts; 3) Parameter Prior: imposes constraints on learning the linear layer's parameters, ensuring that the network's predictions remain consistent with the established medical knowledge. The proposed method of KnoBo is evaluated using two medical datasets, ISIC and X-ray, under both confounded and unconfounded settings.

Strengths

The paper is well-organized and easy to follow. The authors propose a novel approach to incorporating medical knowledge priors into existing neural networks by optimizing three distinct terms. The theory behind each term is solid and clearly explained, making the methodology easy to understand. One notable aspect is the parameter prior (Sec. 4.4), which adds interpretability by aligning the estimations with medical concepts. They also use human evaluation on the learned concepts. The authors conduct a comprehensive set of experiments to evaluate their proposed method, covering several scenarios, including confounded and unconfounded data. This broad range of experiments supports the validity of their findings. The ablation studies are thorough, encompassing five knowledge sources, appropriate baselines, and various bottleneck sizes. These ablations provide valuable insights into the proposed method's performance and sensitivity.

Weaknesses

The paper's primary weakness is its reliance on content generated by LLMs to create new medical concepts. Although GPT-4 is one of the most advanced models for generating concepts, it can still produce hallucinations, especially when the generated concepts are loosely aligned with common terms such as hair, skin color/tone, gel bubble, and rules (in the context of skin lesions).

Questions

- Whats is the main reason to define $ \Delta = | ID - OOD |$ as robustness measure (row 257)? In my opinion, this metric is more related to some "fair predictions" concept, since it measures how ID and OOD performances differ. - Do the authors know why the diversity values for skin lesion datasets (Table 3, last column) are much lower than those for X-rays?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Not a limitation, but I recommend incorporating a filtering scheme into the concept generation process to enhance the paper. While the authors have conducted human evaluations, relying solely on such assessments may not be feasible due to the limited availability of human resources. Refining the concept generation could be achieved by improving the prompts or employing a medical foundation model to assess the quality of the generated concepts about the target image.

Summary

This paper proposed Knowledge-enhanced Bottlenecks (KnoBo) to leverage the concept bottleneck model (CBM) to improve model robustness to various domain shifts and confounders. KnoBo explores using retrieval augmented generation to generate concept space with a large number of concepts using a large language model (i.e., GPT-4). KnoBo brings significant overall improvement when large-scale image-text pairs are available for concept grounding. A comprehensive evaluation across various scenarios and various baselines confirms the effectiveness of KnoBo in reducing the gap between in-distribution and out-of-distribution performance.

Strengths

1. This paper explores exploiting LLM and retrieval augmentation generation, which enables building large concept space from the large-scale knowledge base. 2. Overall, the presentation is clear. Technical details are clearly presented in the supplemental material. However, some results are a bit confusing and lack discussion of limitation and failure cases. 3. Evaluation is comprehensive.

Weaknesses

1. The concept grounding seems pretty data-hungry, requiring large-scale multi-modal datasets for pretraining (e.g., MIMIC-CXR with over 300k image-text pairs). For smaller datasets (ISIC with about 60k images with generated captions from LLM), the proposed KnoBo only brings marginal overall improvement. This requirement significantly limits the broad application of the proposed framework in the data-scarcity domain, such as cancer/tumor classification. 2. Lack of a detailed discussion of limitations (one one-line limitation in Section 7) and presentation of failure cases. No failure cases (wrong classification / wrong concept) were included in either the main text or supplemental material. 3. Some results are confusing to the reviewer (See No. 2 in the Questions section).

Questions

1. In Table 1, the best performance of ISIC-age should be ViT-L/14 and LSL, not the proposed KnoBo (which is bolded). 2. The KnoBo with 150 concepts has an average OOD Acc of 58.8 across Chest X-ray Datasets, as stated in Tables 2 and 4. Why does Figure 4 indicate that KnoBo with 135 concepts has an accuracy below 50 for OOD X-ray datasets? As the performance of OOD datasets decreases with the increased number of concepts, these two results seem contradicted. 3. Typo in the title of Algorithm 1 in Section 4.2. Should be 'retrieval'

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

As explained in the Weaknesses and Questions sections.

Response to the rebuttal

I thank the authors for their time and effort in preparing the rebuttal. The rebuttal resolved some of my questions, and I appreciate the explanation and revision for Figure 4 and the supplement of failure concepts in the attachment. I have some remaining questions: 1. About the examples per grounding functions, does this part correspond to lines **177-179**? Could the author point out or supply how many grounding functions are used for each dataset, how they are selected, and how is the # of grounding functions confirmed? I tried but failed to find these details in the manuscript. If you use 22k examples in x-ray, you should have (22k/2k=11) grounding functions? Is that different from the skin dataset? And why does each x-ray grounding function use 1,523 examples when you actually supply 2,000 examples per grounding function? 2. For the failure case, would it be possible to supply some failure examples of KnoBo in Fig. 7 (i.e., examples without the green checkmarks)? Thank you very much, and I am looking forward to your reply. Great work!