Artificial Intelligence is rapidly advancing and radically impacting everyday life, driven by the increasing availability of computing power. Despite this trend, the adoption of AI in real-world healthcare is still limited. One of the main reasons is the trustworthiness of AI models and the potential hesitation of domain experts with model predictions. Explainable Artificial Intelligence (XAI) techniques aim to address these issues. However, explainability can mean diferent things to people with diferent backgrounds, expertise, and goals. To address the target audience with diverse needs, we develop storytelling XAI. In this research, we have developed an approach that combines multi-task distillation with interpretability techniques to enable audience-centric explainability. Using multi-task distillation allows the model to exploit the relationships between tasks, potentially improving interpretability as each task supports the other leading to an enhanced interpretability from the perspective of a domain expert. The distillation process allows us to extend this research to large deep models that are highly complex. We focus on both model-agnostic and model-specific methods of interpretability, supported by textual justification of the results in healthcare through our use case. Our methods increase the trust of both the domain experts and the machine learning experts to enable a responsible AI.
The adoption of artificial intelligence (AI) in healthcare has been relatively slow compared to other
industries, despite its immense potential benefits. Acceptance of AI by patients and healthcare
providers requires building trust in these AI-powered systems in hospitals and healthcare systems [
Counterfactual explanations frequently rest on the presumptions that the features are independent
and that the changes in the highlighted features match real-world activities. In actuality, though, the
characteristics are not independent, and actions are likely to have simultaneous efects on several
features. Post-hoc counterfactual explanations need to provide faithful changes that a human can
implement in practice and should conform to the observed correlations in the training data rather
than being anomalies. They may not be faithful to the original data due to problems like overfitting
or excessive generalization, which results in unsatisfactory interpretability. Although ground truth
plausibility of counterfactuals is acknowledged as a crucial attribute, the search results show that
this characteristic is still challenging to measure in practice [
We develop a storytelling XAI framework to provide end-to-end explanations customized for
domain specialists, in our case, healthcare practitioners as well as ML practitioners. The storytelling
XAI uses a knowledge distillation [
Knowledge distillation refers to the process of transferring knowledge from a comparatively large model to a smaller model without compromising on performance. The complex large model is called teacher and the smaller model is called student. The rationale behind model distillation is to train complex large models also called teacher for a specific task and then transfer the knowledge to a smaller model also called student using distillation loss (Equation: 1). The loss function employed in knowledge distillation to instruct the student model to imitate the teacher model’s behavior is called the distillation loss. By using a "temperature" scaling function in the softmax, the logits are softened, so smoothing down the probability distribution and exposing the teacher’s taught inter-class correlations. The probability pi of class i from the logits z is calculated as: = (/ ) ÷ ∑︁ ( / )
In AI, multi-task learning is the process of teaching one model to handle several related tasks at once,
as opposed to training diferent models for every task. The main concept is to increase overall learning
eficiency and generalization performance by utilizing shared representations and commonalities across
tasks Multi-task Learning can be implemented through various methods and the two most common
methods are shared base network extractors with task-specific heads in which the model comprises
task-specific output heads after a shared base network that extracts features. The second method is the
shared decision-making layer in which task-specific layers are connected to the shared decision-making
layer in the model [
Interpretability allows ML practitioners to understand the connections between the features that go
into a model and the results it produces, as well as the relative importance of diferent features in the
decision-making process. Some models, such as linear regression and decision trees, are inherently
interpretable due to their simple structure, while complex models, such as deep neural networks, are
often considered black boxes and require additional techniques to enhance their interpretability [
Attention(, , ) = softmax ︂( )︂ √ where, Q :represents the query matrix.
K :represents the key matrix.
V :represents the value matrix. :represents the dimensionality of the key vectors.
This equation computes the attention weights by taking the dot product of the query and key matrices, scaling it by √, applying the softmax function, and then multiplying it by the value matrix to get the attention output.
Visualization techniques include the analysis of saliency maps, and activation maps of the intermediate layers of a model, the most relevant being GradCAM (Equation:3).
G rad-CAM = ReLU ︃( ∑︁ )︃ where, G rad-CAM : represents the Grad-CAM score for class.
: represents the importance of the th class. : represents the th activation map for the last convolution layer.
ReLU denotes the rectified linear unit function, which clips negative values to zero.
Original Text Cardiac size is normal. Moderate hiatal hernia is again noted. The lungs are hyperinflated. Patient has known emphysema. Small lung nodules, and lingular atelectasis are better seen in prior CT. There is no pneumothorax or pleural effusion.
Augmented Text
['patient has known emphysema. small lung nodules, and lingular atelectasis. there is no
pneumothorax or pleural effusion.']
(3)
(4)
lime() : represents the local interpretable model produced by LIME for input.
ℱ : denotes the set of possible interpretable models.
: approximates the prediction function in the local neighborhood of .
ℒ(, ) : represents the loss function measuring how well the interpretable model.
Ω( ) : represents a regularization term that encourages simplicity or interpretability of the model .
LIME approximates the behavior of the model locally using interpretable models such as linear regression
by perturbing the input and correspondingly reflecting it in the predictions [
Explainable Artificial Intelligence (XAI) is critical in the medical and healthcare domains because
AI-driven systems can have serious and even life-threatening consequences for patients if their
decisions are not transparent and interpretable. Model developers are primarily interested in the
performance, stability, and robustness of the AI model. Their focus is on using explainability to debug
models and improve accuracy [
By using Storytelling XAI, the gap between ML engineers and healthcare professionals is reduced. ML engineers can use Storytelling XAI to transform complex ML models into understandable information that can be shared with healthcare domain experts. ML engineers can more easily communicate the model’s underlying logic and decision-making processes through storytelling, as opposed to providing raw data or technical model outputs, making it easier for domain experts to understand how the model’s predictions fit into actual patient situations and scenarios. Domain experts are more likely to accept the model’s suggestions and incorporate AI into clinical decision-making processes if they can follow the reasoning behind the model’s outputs through a narrative.
For this research, we have used chest X-ray images to detect the abnormalities in the chest, identify the afected regions, and generate a report generation. The report generation task generates a report out of the chest X-image and serves as a text justification for the other task.
Our research contains two diferent tasks which utilize a diferent dataset for each task. The task and the corresponding dataset used are as follows: 1. Chest X-ray report generation: We used a pre-processed dataset provided by the authors of MedVill (Medical Vision Language Learner). The pre-processed dataset was prepared from the original MIMIC-CXR-JPG dataset [23] and Open-I dataset [24]. The MIMIC-CXR-JPG dataset contains 377,110 chest X-ray images and corresponding free-text reports. The Open-I dataset contains 3,851 reports and 7,466 Chest X-ray images, consisting of both the lateral and frontal view of the images [25]. For our ease, the lateral views were removed from the dataset. For our experimentation, we have selected 5000 image-report pairs randomly. The reports were augmented with the help of clinical T5 large which (Fig:1) is a sequence-to-sequence transformer (available on HuggingFace) to obtain the summarized version of the radiology report with the sequence length of 128 words [26]. We performed this step to obtain keywords having higher weightage in the text report through summarization. 2. Chest X-ray abnormality detection: For this task we used a modified version of the VinDr-CXR dataset [27]. The dataset contains 15,000 samples of DICOM images labeled with 14 abnormalities namely aortic enlargement, atelectasis, calcification, cardiomegaly, consolidation, ILD, Infiltration, Lung Opacity, Nodule/Mass, Other lesion, Pleural efusion, Pleural thickening, Pneumothorax, 1. Teacher Abnormality Detection 2. Teacher Lung Segmentation 3. Teacher Report Generation Multi-Task Deep Neural Networks
DISTILL KNOWLEDGE 1. Abnormality Detection 2. Lung Segmentation 3. Report Generation LIME, SHAP, GradCam Attention Map Visualization Shallow Multi-Task Network (Explainer)
and Pulmonary fibrosis along with bounding box for detection and localization of the detected anomalies. The dataset was annotated by a group of seventeen expert radiologists [28]. We selected 5,000 samples randomly without unbalancing the class ratio which could have potentially led to the problem of class imbalance. 3. Lung Segmentation: For this task, we have used the Montgomery dataset which contains images from the Department of Health and Human Services, Montgomery County, Maryland, USA. The dataset consists of 138 CXRs, including 80 normal patients and 58 patients with manifested tuberculosis (TB). The CXR images are 12-bit gray-scale images of dimension 4020 × 4892 or 4892 × 4020. Only the two lung mask annotations are available which were combined into a single image to make it easy for the network to learn the task of segmentation. To make all images of symmetric dimensions we padded the pictures to the maximum dimension in their height or width such that images are 4892 x 4892, this is done to preserve the aspect ratio of CXR while resizing. We scale all images to 512x512 pixels, which retains suficient visual details for vascular structures.
All the images which were not in JPEG format were converted into JPEG from DICOM format.
The images were normalized to the standard range of [
The storytelling XAI framework (Fig:2 is divided into three parts. The priority intent of this framework is not to compete with current state-of-the-art models in terms of evaluation metrics. This framework utilizes knowledge distillation, interpretability, and datasets for a variety of tasks. Using datasets from diferent origins allows the framework to generalize better for real-world scenarios as well. The three parts involved are: 1. The First step involves training the complex deep neural networks for individual tasks. For the task of abnormality detection and localization, a CNN-based model with ResNet 50 backbone is trained using a categorical cross-entropy loss and mean Average Precision(mAP), which also achieves a remarkable performance on this task [29]. For the task of report generation, a CNNRNN-Attention model with ResNet50 backbone is trained for the task using masked loss. This task utilizes the concept of image-captioning, where an image is provided as an input. The hidden layers of the architecture perform the feature extraction process. Then, these features are passed through the RNN layer coupled with the attention layer to generate a radiology report [30]. The image segmentation model uses the same model as the abnormality detection model with the last layers replaced and modified to predict the segmentation task, trained using DICE loss [31]. Algorithm 1 Student Model Training with Knowledge Distillation For Multi-Task Prediction 1: Initialize student model with three heads: , , 2: Initialize teacher models , , 3: Initialize loss functions for each task: , , 4: Initialize hyperparameters: learning rate, temperature and distillation weights 5: Freeze parameters of and 6: for epoch = 1 to do: 7: for ℎ in _ do 8: Forward pass through teacher model for report generation: = (ℎ) 9: Compute loss for report generation: ((ℎ), ) 10: Compute knowledge distillation loss for report generation: ℒ((ℎ), ) 11: Compute total loss for report generation: ℒ = (1 − ) · Lreport + · ℒ distillation 12: Backpropagation: Compute gradients ∇ ℒ 13: Update parameters: ← − · ∇ ℒ 14: end for 15: end for 16: Freeze parameters of and 17: for ℎ = 1 to do 18: for ℎ in _ do: 19: Forward pass through teacher model for abnormality detection: = 20: 21: (ℎ)
Compute loss for abnormality detection: ((ℎ), ) Compute knowledge distillation loss for report generation: ℒ((ℎ), ) 22: Compute total loss for abnormality detection: ℒ = (1 − ) · Labnormality + · ℒ distillation 23: Backpropagation: Compute gradients ∇ ℒ 24: Update parameters: ← − · ∇ ℒ 25: end for 26: end for 27: Freeze parameters of and 28: for ℎ = 1 to do 29: for ℎ in _ do: 30: Forward pass through teacher model for segmentation: = 31: 32: (ℎ)
Compute loss for abnormality detection: ((ℎ), ) Compute knowledge distillation loss for report generation: ℒ((ℎ), ) 33: Compute total loss for segmentation: ℒ = (1 − ) · Lsegmentation + · ℒ distillation 34: Backpropagation: Compute gradients ∇ ℒ 35: Update parameters: ← − · ∇ ℒ 36: end for 37: end for 2. The Second step involves performing knowledge distillation between a shallow CNN neural network which will be the backbone of this model. This neural network consists of diferent prediction heads to perform multiple tasks. The extracted feature representation of the image from the model is passed into the image classification head, text generation head, and image segmentation head to perform abnormality classification, report generation, and image segmentation, simultaneously. Knowledge distillation is performed between the backbone including the specific task head and the individual model trained for the specific task. while the other prediction heads are frozen (C.f., Algorithm: 1). This is performed until all the prediction heads have acquired knowledge from the complex neural network through knowledge distillation. Here the shallow CNN neural network acts like a student and the individual complex models as teachers. 3. The Third and last step involves generating interpretations from the new model which has acquired knowledge from a complex deep neural network model. This step uses model-agnostic and model-specific interpretation techniques to enable interpretability. The LIME interpretations and GradCAM visualizations are incorporated to interpret the abnormality classification model and chest X-ray segmentation model. For report generation, attention visualizer [32] is used to output the results.
The chest X-ray image was given as input to the shallow neural network, which was trained with the knowledge distillation method from three diferent deep neural networks for a diferent task. The original image was resized to a resolution of 512x512 pixels. See figure: 3. The model accurately segmented the two lungs. The abnormality detection model was able to localize and classify the localized afected regions into infiltration, consolidation, pleural efusion, and cardiomegaly. The same model generated the resulting text for radiologists and healthcare professionals. Furthermore, interpretability techniques were applied, which include LIME interpretation, attention map visualization of the generated report, and CAM analysis consisting of GradCAM and GradCAM++. As shown in the ifgure, the text justification asserts the output of interpretability techniques along with the prediction from the student model which will be beneficial for ML practitioners and at the same time it is able to assert the abnormality detection for healthcare professionals as well. Knowledge distillation is a well-capable method to let the shallow neural network to learn representation over a variety of tasks enabling an interpretable neural network. The benefit of the multiple tasks could be inferred from the result (Fig:3), as these multiple tasks when combined, improve the overall understanding of the predictions. The report generation task serves as a textual justification for the other two tasks, namely abnormality detection and chest X-ray segmentation. For interpretation, we have focused on LIME, GradCAM, and visualization techniques. The framework could be extended to incorporate other interpretability techniques including but not limited to SHapley Additive exPlanations (SHAP) [33], Layer-wise Relevance Propagation (LRP) [34], etc.
Our framework is a step toward enabling audience-centric explainability. We have used concept-based explainability to explain the results of the AI models to healthcare professionals. This framework serves as a baseline to demonstrate that the XAI is an amalgamation of interpretability and concept-based approaches to achieve end-to-end explainability. Note that ML practitioners can also interpret the results of the models in technical terms. This framework is also able to explain complex deep neural networks because of the knowledge distillation. The teacher models can each be a very complex deep neural network, but the student model can learn the knowledge. This framework requires multiple datasets for multiple machine tasks. The multiple data sets from diferent sources result in more generalized and robust models in a real-world scenario. Robust and generalized models that perform well in a variety of real-world scenarios are more likely to be trusted and understood by users. Robust and generalized models are less susceptible to adversarial manipulation, improving the reliability of their explanations.
In this research, we have limited our investigation to some specific tasks of abnormality detection, report generation, and chest X-ray segmentation. Other tasks could be added if the datasets for these tasks are available and the inputs are common to all tasks. We have demonstrated a healthcare use case, but this framework could be extended to other domains given the available datasets. Multi-task models are a great way to exploit the relationships between diferent commonalities. This technique allows the model to learn shared feature representations. By leveraging the shared representations across related tasks, multi-task learning can learn more eficiently from limited data compared to Learning a shared representation that performs well across multiple related tasks can lead to more stable and meaningful predictions. Our Storytelling XAI framework can be implemented in critical-stakes domains to enable audience-centric explanations with robust predictions and generalized models, without compromising the performance of even the most complex deep neural networks.
A potential concern is that providing an excessive number of explanations or information may lead to an overload of data, hindering comprehension and user satisfaction. However, the presentation of multiple explanations can enhance transparency and facilitate trust in AI models by addressing the diverse needs of users. By ofering a range of explanations, users can gain a more nuanced understanding of the decision-making process, which can ultimately enhance their satisfaction and comprehension [35]. One might also inquire as to the rationale behind the use of shallow models. The use of shallow models is justified by their reduced susceptibility to overfitting and superior generalization capabilities when processing new data sets, particularly in the context of limited data. Training shallow models for multi-tasking allows for the exploitation of relationships between diferent input modalities, which is not feasible for single-task models. The majority of interpretability methods employ approximation techniques. As the depth of a neural network model increases, the reliance on approximation increases and reliability decreases. Consequently, shallow methods align well with the aforementioned criteria [36].
Our framework currently lacks the ability of a user to interact with the XAI results freely. This may be resolved with a user interface that focuses on integrating Attention Maps, LIME, SHAP, and GradCam visualizations. When extending this framework to other domains, it may be somewhat problematic to collect diferent datasets for diferent tasks. There is limited support for teacher models. The current possibilities are to pick either CNN or RNN-based architectures. For transformer architectures, it may be dificult to perform knowledge distillation. The process of knowledge distillation of transformer models is inherently challenging due to the intricate architectural complexity, the considerable number of parameters, and the dificulty in transferring the intricate patterns learned by self-attention mechanisms. Furthermore, the process is further complicated by the necessity to balance multiple objectives and the risk of overfitting. We have experimented with tasks where the input is solely an image. There is indeed not only a limitation in terms of multimodal support but also textual-only data. Despite the limitations, the work lays down the foundation for future work.
In future work, we would develop a user interface with a human-in-the-loop, preferably domain experts which satisfies the nested model for AI design and validation [ 37]. Then we would focus on conducting a design study and include a new metric to calculate an agreement score between the framework and the domain experts. There would also be an extensive efort to integrate the XAI Question Bank (XAI-QB) into our framework [38]. The XAI-QB is critical to the development of explainable AI systems and is designed to help XAI practitioners understand the explanation needs of diferent end users. It represents common user expectations and covers major categories of questions for standard supervised ML systems. infiltration and pulmonary consolidation. small left-sided pleural effusion moderate cardiomegaly appears tr iton po r
a eR ene
G Textual Justification of the Predictions and the Interpretations infiltration and pulmonary consolidation. small left-sided pleural effusion moderate cardiomegaly appears GradCAM Visualization GradCAM++
Visualization
LIME Interpretability (2021) 794996. [21] N. Prentzas, A. Kakas, C. S. Pattichis, Explainable ai applications in the medical domain: a systematic review, arXiv preprint arXiv:2308.05411 (2023). [22] Y. Rong, T. Leemann, T.-T. Nguyen, L. Fiedler, P. Qian, V. Unhelkar, T. Seidel, G. Kasneci, E. Kasneci, Towards human-centered explainable ai: A survey of user studies for model explanations, IEEE Transactions on Pattern Analysis and Machine Intelligence (2023). [23] A. Johnson, M. Lungren, Y. Peng, Z. Lu, R. Mark, S. Berkowitz, S. Horng, Mimic-cxr-jpg-chest radiographs with structured labels, PhysioNet (2019). [24] D. Demner-Fushman, M. D. Kohli, M. B. Rosenman, S. E. Shooshan, L. Rodriguez, S. Antani, G. R.
Thoma, C. J. McDonald, Preparing a collection of radiology examinations for distribution and retrieval, Journal of the American Medical Informatics Association 23 (2016) 304–310. [25] J. H. Moon, H. Lee, W. Shin, Y.-H. Kim, E. Choi, Multi-modal understanding and generation for medical images and text via vision-language pre-training, IEEE Journal of Biomedical and Health Informatics 26 (2022) 6070–6080. [26] M. Chizhikova, M. Diaz-Galiano, L. A. U. Lopez, M. T. Martín-Valdivia, Sinai at radsum23: Radiology report summarization based on domain-specific sequence-to-sequence transformer model, in: The 22nd Workshop on Biomedical Natural Language Processing and BioNLP Shared Tasks, 2023, pp. 530–534. [27] DungNB, H. Q. Nguyen, J. Elliott, KeepLearning, NguyenThanhNhan, P. Culliton, Vinbigdata chest x-ray abnormalities detection, 2020. URL: https://kaggle.com/competitions/ vinbigdata-chest-xray-abnormalities-detection. [28] H. Nguyen, et al., Vindr-cxr: An open dataset of chest x-rays with radiologist’s annotations.(2022) doi: 10.48550, arXiv (2012). [29] S. Shivadekar, B. Kataria, S. Hundekari, K. Wanjale, V. P. Balpande, R. Suryawanshi, Deep learning based image classification of lungs radiography for detecting covid-19 using a deep cnn and resnet 50, International Journal of Intelligent Systems and Applications in Engineering 11 (2023) 241–250. [30] J. Aneja, A. Deshpande, A. G. Schwing, Convolutional image captioning, in: Proceedings of the
IEEE conference on computer vision and pattern recognition, 2018, pp. 5561–5570. [31] R. Rashid, M. U. Akram, T. Hassan, Fully convolutional neural network for lungs segmentation from chest x-rays, in: Image Analysis and Recognition: 15th International Conference, ICIAR 2018, Póvoa de Varzim, Portugal, June 27–29, 2018, Proceedings 15, Springer, 2018, pp. 71–80. [32] F. Liu, C. Yin, X. Wu, S. Ge, Y. Zou, P. Zhang, X. Sun, Contrastive attention for automatic chest x-ray report generation, arXiv preprint arXiv:2106.06965 (2021). [33] Y. Arslan, B. Lebichot, K. Allix, L. Veiber, C. Lefebvre, A. Boytsov, A. Goujon, T. F. Bissyandé, J. Klein, Towards refined classifications driven by shap explanations, in: International CrossDomain Conference for Machine Learning and Knowledge Extraction, Springer, 2022, pp. 68–81. [34] G. Montavon, A. Binder, S. Lapuschkin, W. Samek, K.-R. Müller, Layer-wise relevance propagation: an overview, Explainable AI: interpreting, explaining and visualizing deep learning (2019) 193–209. [35] W. Yang, Y. Wei, H. Wei, Y. Chen, G. Huang, X. Li, R. Li, N. Yao, X. Wang, X. Gu, et al., Survey on explainable ai: From approaches, limitations and applications aspects, Human-Centric Intelligent Systems 3 (2023) 161–188. [36] S. Xu, Y. Song, X. Hao, A comparative study of shallow machine learning models and deep learning models for landslide susceptibility assessment based on imbalanced data, Forests 13 (2022) 1908. [37] A. Dubey, Z. Yang, G. Hattab, A nested model for ai design and validation, arXiv e-prints (2024) arXiv–2407. [38] Q. V. Liao, D. Gruen, S. Miller, Questioning the ai: informing design practices for explainable ai user experiences, in: Proceedings of the 2020 CHI conference on human factors in computing systems, 2020, pp. 1–15.