Integrating Clinical Knowledge into Concept Bottleneck Models

TL;DR

Concept Bottleneck Models (CBMs) often learn from data in ways that can propagate biases and fail under domain shifts. The paper introduces a clinical-knowledge guided CBM framework that aligns the model's concept importance with clinician priorities by perturbing each concept and using alignment losses on a measure $\Delta Y_{k,l} = |\hat{y}_k - \hat{y}_{(\hat{c}_l\to0)_k}|$, balanced by high/low importance constraints. This approach is evaluated on white blood cell and skin image datasets, demonstrating improved out-of-domain performance and better alignment with expert knowledge while maintaining interpretability. The work provides a practical method to enhance robustness and transferability of interpretable medical imaging models by incorporating domain-specific clinical insights into the CBM training objective.

Abstract

Concept bottleneck models (CBMs), which predict human-interpretable concepts (e.g., nucleus shapes in cell images) before predicting the final output (e.g., cell type), provide insights into the decision-making processes of the model. However, training CBMs solely in a data-driven manner can introduce undesirable biases, which may compromise prediction performance, especially when the trained models are evaluated on out-of-domain images (e.g., those acquired using different devices). To mitigate this challenge, we propose integrating clinical knowledge to refine CBMs, better aligning them with clinicians' decision-making processes. Specifically, we guide the model to prioritize the concepts that clinicians also prioritize. We validate our approach on two datasets of medical images: white blood cell and skin images. Empirical validation demonstrates that incorporating medical guidance enhances the model's classification performance on unseen datasets with varying preparation methods, thereby increasing its real-world applicability.

Figures (6)

• Figure 1: We propose a method to guide concept bottleneck models (CBMs) using knowledge aligned with clinicians' perspectives. (a): CBMs predict interpretable concepts (e.g., granule color, cell shape, etc.) and then make a final prediction (e.g., eosinophil) based on them. During training, models usually do not consider the clinical importance of the concepts. Therefore, granule color and cell shape are treated equally despite granule color being a much more important factor for predicting eosinophils. (b): To incorporate clinical knowledge, we enforce the CBM to exhibit a significant drop in cell type prediction probabilities when a clinically important concept is removed from the prediction. For instance, the predicted eosinophil probability should be lower when granule color, a key factor in recognizing eosinophil, is missing. (c): Conversely, the cell type prediction probabilities should experience a negligible drop when a less clinically important concept is removed from the prediction. For instance, the eosinophil probability should not be affected much when cell shape, which is irrelevant to recognizing eosinophil, is missing.
• Figure 2: Integration of Clinical Knowledge into CBM: The lower path illustrates the original CBM, where the final class prediction ($\hat{\mathbf{y}}$) is based on the prediction of concepts ($\hat{\mathbf{c}}$). The concept importance learned from the model ($\Delta Y$) is obtained through computing the absolute difference between the prediction probabilities when all concepts are used for prediction ($\hat{\mathbf{y}}$) and when the concept is removed from prediction ($\hat{y}_{\hat{c}_{l}\to0}$). We align the $\Delta Y$ with the concept importance from the clinician's perspective through $Loss_{align}$. Specifically, we enforce the model to maximize the $\Delta Y_{:,l}$ for the concept that is considered as High important by the clinicians, while we constrain the model to minimize the $\Delta Y_{:,l}$ for the concept that is considered Low important by the clinicians.
• Figure 3: (a) Examples for each cell type from the in-domain dataset (PBC), as well as the out-of-domain datasets (Scirep and RaabinWBC). (b) Skin images for Benign and Malignant classes from Fitzpatrick 17k (in-domain) and DDI (out-of-domain) datasets.
• Figure 4: The effect of $\lambda$ in loss function for WBC classification, using VGG + MLP(128).
• Figure 5: Qualitative results on the out-of-domain WBC and skin datasets demonstrate that our method improves concept predictions, leading to correct class predictions.