Mask-Free Neuron Concept Annotation for Interpreting Neural Networks in Medical Domain

TL;DR

The paper addresses the need for transparent AI in medicine while reducing the costly requirement for pixel‑level concept masks. It introduces MAMMI, a mask‑free medical neuron concept annotation framework that leverages vision–language models to map neuron activations to domain‑relevant medical concepts. Key innovations include constructing a medical concept set from domain medical reports, adaptive neuron representative image selection to mitigate class imbalance, and adaptive cosine‑based concept matching in CLIP space with a threshold $\theta_{concept}$ and parameter $\beta=0.95$. On NIH Chest X‑ray14 with DenseNet121, MAMMI outperforms mask‑based methods like TSI and other CLIP‑based baselines, with clear gains when using domain‑specific concept sets such as MIMIC_Nouns; this demonstrates improved interpretability and potential for transparent clinical decision‑making. The work also provides evidence that reducing data annotation costs via mask‑free methods can retain, or even improve, explanatory quality in medical imaging tasks, with code available for replication.

Abstract

Recent advancements in deep neural networks have shown promise in aiding disease diagnosis and medical decision-making. However, ensuring transparent decision-making processes of AI models in compliance with regulations requires a comprehensive understanding of the model's internal workings. However, previous methods heavily rely on expensive pixel-wise annotated datasets for interpreting the model, presenting a significant drawback in medical domains. In this paper, we propose a novel medical neuron concept annotation method, named Mask-free Medical Model Interpretation (MAMMI), addresses these challenges. By using a vision-language model, our method relaxes the need for pixel-level masks for neuron concept annotation. MAMMI achieves superior performance compared to other interpretation methods, demonstrating its efficacy in providing rich representations for neurons in medical image analysis. Our experiments on a model trained on NIH chest X-rays validate the effectiveness of MAMMI, showcasing its potential for transparent clinical decision-making in the medical domain. The code is available at https://github.com/ailab-kyunghee/MAMMI.

Figures (3)

• Figure 1: Overview of neuron concept annotation method of MAMMI. The representative images of neurons for concept annotation are selected by images with activation values over the neuron-wise computed adaptive threshold from the probing data. The concept set is constructed by extracting nouns from the medical report. The representative images and the text concept set are projected into the CLIP model embedding space for concept identification by the adaptive concept matching module.
• Figure 2: The data distribution of adaptive neuron representative image selection compared to the probing set. (a) Ablation study on the value of parameter $\alpha$. (b) Distribution regarding the number of selected representative images of neurons in the final layer and train set class distribution.
• Figure 3: Qualitative result on the penultimate layer. We display two penultimate layer neurons selected by Shapley Value khakzar2021neuralwww24 as the most important penultimate neurons for identifying the respective ground truth label in the image.