Model Distillation for Faithful Explanations of Medical Code Predictions (2022)

Abstract: Machine learning models that offer excellent predictive performance often lack the interpretability necessary to support integrated human machine decision-making. In clinical medicine and other high-risk settings, domain experts may be unwilling to trust model predictions without explanations. Work in explainable AI must balance competing objectives along two different axes: 1) Models should ideally be both accurate and simple. 2) Explanations must balance faithfulness to the model’s decision-making with their plausibility to a domain expert. We propose to use knowledge distillation, or training a student model that mimics the behavior of a trained teacher model, as a technique to generate faithful and plausible explanations. We evaluate our approach on the task of assigning ICD codes to clinical notes to demonstrate that the student model is faithful to the teacher model’s behavior and produces quality natural language explanations.

## introduction machine learning (ml) methods have demonstrated predictive success in medical settings, leading to discussions of how ml systems can augment clinical decision-making @xcite . however, a prerequisite to clinical integration is the ability for healthcare professionals to understand the justifications for model decisions. clinicians often disagree on the proper course of care, and share their justifications as a means of agreeing on a treatment plan. explainable artificial intelligence (ai) can enable models to provide the explanations needed for them to be integrated into this process @xcite . however, modern ai models that often rely on complex deep neural networks with millions or billions of parameters pose challenges to creating explanations that satisfy clinician's demands @xcite . ## methods our proposed method is post-hoc and seeks to balance faithfulness and plausibility. we assume that we have a trained teacher model with good predictive performance but low interpretability. we train a student model that takes the same input from the dataset, but uses the teacher model's predictions as its labels. figure 1 gives a visual representation of our model distillation setup. the mimic-iii dataset contains anonymized english-language icu patient records, including physiological measurements and clinical notes @xcite . following mullenbach et al. (2018) , we focus on discharge summaries which describe a patient's visit and are annotated with icd-9 codes. there are 8,922 different icd-9 codes that describe procedures and diagnoses that occurred during a patient's stay. the manual assignment of these codes to patient records are required by most u.s. healthcare payers @xcite . to train the teacher models, we duplicate the experimental setup of @xcite and @xcite , which use the text of the discharge summaries as input to the dr-caml and han models, respectively, which then are trained to predict all icd-9 codes associated with that document. after applying their pre-processing code to tokenize the text, the dataset contains 47,724 discharge summaries divided into training, dev, and test splits. we also duplicate the experimental setup of @xcite , which has a similar experimental setup but only predicts the top 50 most common icd-9 codes. we apply dr-caml and han to the texts in mimic-iii and save its continuous-valued probabilities as the labels for our student model. we similarly apply transicd to mimic-iii-50, which contains the top 50 most frequent labels in mimic-iii and save the continuous-valued probabilities as the labels. 3 for all three models, we use the code released by the authors. 4 training the student model on predictions from the existing teacher model optimizes for faithfulness. we want the student model to produce plausible explanations and fulfill the criteria from lipton (2018): simulatibility, decomposibility, and algorithmic transparency. to fulfill these desiderata, each student is a linear regression trained on bag-ofwords representation of the clinical text. the fundamental trade-off here is that if we overly restrict our model class, the student will be unfaithful and unable to mimic the behavior of the teacher model. but if we allow for a student model that is too complex, it may not provide plausible or otherwise desirable explanations. these trade-offs may be domain-specific based, for example, on the target audience of the explanations. if the student model demonstrates sufficient empirical performance, a domain expert may even prefer to use it in place of the teacher model, an option unsupported by lime or shap models. we train a student model for each medical code independently, and we refer to the student model trained on x model's predictions as "student-x" (e.g. student-drcaml). each student uses only 50k parameters, allowing us to train each model on a single cpu in a matter of minutes. we implement our method using the linear sgdregressor model from sklearn @xcite , and apply a log transform to the model's probability outputs and train the student to minimize squared loss. after a brief foot_3 grid search on the validation set, we use l1 regularization with α = 0.0001 for the dr-caml student and α = 0.01 for han and transicd proxies. to extract a rationale, we take the feature importance weights of the student model and average over a sliding window of n-grams from the discharge summary. we extract the n-gram with the highest average feature importance weight. future work could use extracted rationales to train a student model that remains faithful to a black-box model. in the next two sections, we introduce our evaluation for the student model's faithfulness to each model and the plausibility of its explanations. ## faithfulness evaluation to establish that this collection of linear regressions is faithful to the trained models, we want to show that it makes similar predictions across all icd-9 codes on held-out data. recall from figure 1 that the student is trained not to predict the true icd-9 codes but to output the same label probabilities as the teacher model. in fact, the student model never sees the true icd-9 codes. we evaluate faithfulness by comparing the outputs of the student and teacher models on the held-out test set. if the two systems produced identical outputs on held-out data, we would say that the student was perfectly faithful. we make this comparison in three different ways -first with regression metrics for the continuous outputs of the two models, then using classification metrics with binarized teacher predictions, and finally comparing student outputs as predictions for the true icd-9 codes. for all these comparisons, we use a logistic regression baseline that is trained to directly predict the icd-9 codes, independent of any black-box model. while we would expect the logistic baseline's predictions to roughly correlate with those of other models, we would not expect it to be faithful. similar to @xcite , our first evaluation uses regression metrics that assess the correlation between the student's predictions and the original teacher model's predicted probabilities. we use spearman and pearson correlation coefficients and the non-parametric kendall tau rank correlation. these metrics range from -1 to 1 with 1 indicating perfect faithfulness. regression results are on the left side of our second evaluation treats the teacher model's predictions as binary labels to compute f1, auc, and precision scores. we then evaluate the faithfulness of our student model by treating its outputs as probabilities and using classification metrics such as f1 score. prec @ n is the fraction of the n highest scored labels that are present in the ground truth. these metrics range from 0 to 1, where perfectly faithful predictions would have 1.0 auc and f1 scores. the student model is considered faithful if it correctly predicts whether the teacher model will make a binary prediction. we again use the logistic regression baseline. classification results are on the right side of finally, we use the student model's predictions to predict the ground-truth icd code labels and compare its predictive performance against that of the teacher model's in our results show that our proxies are quite faithful to the teacher models. mullenbach et al. (2018) caml (hi) ... line placed bronchoscopy performed showing large mucus plug on the left on transfer to... cosine ... also needed medication to help your body maintain your blood pressure after receiving iv... cnn ... found to have a large lll lingular pneumonia on chest x ray he was... logistic ... impression confluent consolidation involving nearly the entire left lung with either bronchocentric or vascular ... ours dr-caml 0.38 ... line placed bronchoscopy performed showing large mucus plug on the left on transfer to... logistic 0.28 ... tube down your throat to help you breathe you also needed medication to help... student-drcaml 0.38 ... a line placed bronchoscopy performed showing large mucus plug on the left on transfer ... student-han 0.39 ... line and r radial a line placed bronchoscopy performed showing large mucus plug on... rudin (2019) critiques post-hoc methods in general, arguing that "if we cannot know for certain whether our [post-hoc] explanation is [faithful], we cannot know whether to trust either the explanation or the original model." because no post-hoc method can ever be perfectly faithful to an original model, our explicit measurement of faithfulness provides a useful approach for understanding whether the student is "faithful enough" for a given application. it also allows for a prediction-specific analysis -if we wish to use the student model to explain a high-stakes prediction made by a blackbox model, we can first check whether both agree upon that specific prediction. in applications where explainability is essential, our student model could be used as a more interpretable replacement for a high-performing blackbox model. in such a case, a domain expert might care less about the evaluation of faithfulness in ## plausibility evaluation explanations are considered plausible if they can be reasoned about by a person @xcite . evaluating plausibility is thus typically more difficult than faithfulness, because it requires input from annotators @xcite . furthermore, an explanation that is plausible to a domain expert may not be plausible to a layperson. @xcite evaluated the plausibility of their models' explanations by collecting annotations from a clinical expert. for 100 notes, each of four models produced an explanation in the form of a 14-token subsequence taken from the discharge summary. the clinician read the four (anonymized) explanations and the corresponding icd code and subjectively rated each explanation as "informative" foot_4 . across the 100 examples, the clinician rated caml as slightly more informative than the logistic regression and cnn baselines. to conduct our plausibility evaluation, we first use or reproduce the baseline methods from @xcite and @xcite . each model, including the student, produces a 14-token explanation from the discharge summary by first finding the 4-gram with the largest average feature importance and then including five tokens on either side of the 4-gram. the logistic regression baseline is the same as in § 4, where feature importance is computed using the coefficients of the logistic model. the student model's explanations are computed in the same manner, finding the 4-gram with the largest average coefficient weights. for caml, dr-caml, and the cnn models, we use the code released by @xcite to extract explanations. the cnn baseline primarily differs from caml in that it does not use an attention mechanism. finally, we reimplement their cosine baseline which picks the 4-gram with the highest cosine similarity to the icd-9 code description text. we extract the model's explanations for the same foot_5 discharge summaries as were evaluated by @xcite . for each explanation, we use the annotation classifier described above to predict the probability that each explanation would have been labeled as informative. if we set the classifier threshold such that 45% of explanations are rated as informative (matching the proportion from the original annotations), we get the results in the score column of ## discussion we have introduced a method that uses knowledge distillation to generate post-hoc explanations and is designed to be interpretable and plausible while maintaining faithfulness to the trained model. by constraining the student to a class of models that is decomposible, simulatible, and algorithmically transparent, our optimization for faithfulness gives us a clear way to evaluate several dimensions of interpretability. we evaluated our method on the task of clinical code prediction. a key benefit of our method is its simplicity and wide applicability. even for a proprietary trained model for which the learned parameters are unknown, a student can be trained as long as we have a dataset that includes the trained model's predictions. our approach has the additional benefit of producing a standalone student model that can provide global feature explanations. if the student has sufficient predictive performance, a skeptic of post-hoc methods (e.g. @xcite ) might prefer to use the inherentlyinterpretable student. the present work has several limitations that are left for future work. though the task of medical code prediction has important implications and has been widely studied in interpretability research, we only consider this single task on a single englishlanguage dataset. while we have shown our student approach works for three different black-box models, it requires additional study in new domains and tasks. there may be black-box models for which no linear student is faithful. our evaluation is also limited to only a single form of explanation: n-grams extracted via importance or attention weights. counterfactual explanations (i.e., an alternative input that would have been classified differently) might be harder or easier for our student method to generate @xcite . our plausibility evaluations rely on a small set of annotations from which we extrapolate. future work should collect new annotations that consider metrics such as sufficiency and simulatibility that require human evaluations @xcite @xcite . as the ml community continues to explore new directions for interpretable methods, new desiderata may arise based on the domain experts who turn to ml methods for decision support. interpretable ml methods should clearly define how they expect to satisfy criteria such as faithfulness or plausibility; by designing for plausibility and transparency and optimizing for faithfulness, our proposed method is broadly applicable. we release our code to enable future work. ## ethics and broader impacts this paper is situated in a broader field of clinical applications of machine learning. while our work does not raise new ethical issues within this domain, there are general concerns that also apply to this work. ml methods should not be deployed in real-world settings without extensive validation @xcite . in the clinical domain, particular attention must be paid to the possibility of perpetuating disparities that have been encoded in the training data @xcite . while mimic-iii provides a useful benchmark for developing and evaluating methods, it is not representative of the the enormous variety of clinical and linguistic data. domain experts and those most likely to be affected by new ml systems should be given oversight of potential deployments.