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

TL;DR

KnoBo, a class of concept bottleneck models that incorporates knowledge priors that constrain it to reason with clinically relevant factors found in medical textbooks or PubMed, is introduced, and evaluations reveal that PubMed is a promising resource for making medical models less sensitive to domain shift.

Abstract

While deep networks have achieved broad success in analyzing natural images, when applied to medical scans, they often fail in unexcepted situations. We investigate this challenge and focus on model sensitivity to domain shifts, such as data sampled from different hospitals or data confounded by demographic variables such as sex, race, etc, in the context of chest X-rays and skin lesion images. A key finding we show empirically is that existing visual backbones lack an appropriate prior from the architecture for reliable generalization in these settings. Taking inspiration from medical training, we propose giving deep networks a prior grounded in explicit medical knowledge communicated in natural language. To this end, we introduce Knowledge-enhanced Bottlenecks (KnoBo), a class of concept bottleneck models that incorporates knowledge priors that constrain it to reason with clinically relevant factors found in medical textbooks or PubMed. KnoBo uses retrieval-augmented language models to design an appropriate concept space paired with an automatic training procedure for recognizing the concept. We evaluate different resources of knowledge and recognition architectures on a broad range of domain shifts across 20 datasets. In our comprehensive evaluation with two imaging modalities, KnoBo outperforms fine-tuned models on confounded datasets by 32.4% on average. Finally, evaluations reveal that PubMed is a promising resource for making medical models less sensitive to domain shift, outperforming other resources on both diversity of information and final prediction performance.

Figures (10)

• Figure 1: In-domain (ID), out-of-domain (OOD), and average of ID and OOD (Avg) performance on confounded medical image datasets. Our interpretable Knowledge-enhanced Bottlenecks (KnoBo) are more robust to domain shifts (e.g., race, hospital, etc) than fine-tuned vision transformers dosovitskiy2020image.
• Figure 2: Classification performance on natural and medical images through linear probing using features extracted from untrained and frozen models versus pixels features (See Sec \ref{['sec: prior']} for details).
• Figure 3: Overview of Knowledge-enhanced Bottlenecks (KnoBo) for medical image classification, comprising three main components: (1) Structure Prior (Sec \ref{['sec: structure_prior']}) constructs the trustworthy knowledge bottleneck by leveraging medical documents; (2) Bottleneck Predictor (Sec \ref{['sec: bottleneck_predictor']}) grounds the images onto concepts which are used as input for the linear layer and ; (3) Parameter Prior (Sec \ref{['sec: parameter_prior']}) constrains the learning of linear layer with parameters predefined by doctors or LLMs.
• Figure 4: Ablation of bottleneck sizes on X-ray datasets. The x-axis is the number of randomly selected concepts (KnoBo) or visual features (Linear Probe). $+$prior means adding parameter prior.
• Figure 5: Prompt template for retrieval-augmented concept bottleneck generation. The text in the square brackets is words that need to be changed when using this prompt for skin lesion images.