Beyond Label Attention: Transparency in Language Models for Automated Medical Coding via Dictionary Learning (2024)

Abstract: Medical coding, the translation of unstructured clinical text into standardized medical codes, is a crucial but time-consuming healthcare practice. Though large language models (LLM) could automate the coding process and improve the efficiency of such tasks, interpretability remains paramount for maintaining patient trust. Current efforts in interpretability of medical coding applications rely heavily on label attention mechanisms, which often leads to the highlighting of extraneous tokens irrelevant to the ICD code. To facilitate accurate interpretability in medical language models, this paper leverages dictionary learning that can efficiently extract sparsely activated representations from dense language model embeddings in superposition. Compared with common label attention mechanisms, our model goes beyond token-level representations by building an interpretable dictionary which enhances the mechanistic-based explanations for each ICD code prediction, even when the highlighted tokens are medically irrelevant. We show that dictionary features are human interpretable, can elucidate the hidden meanings of upwards of 90% of medically irrelevant tokens, and steer model behavior.

## introduction transparency is a vital factor in healthcare to gain patients' trust, especially when ai models make critical decisions in clinical practice @xcite . one of the essential applications of ai models is to assign international classification of diseases (icd) codes automatically based on the clinical text (we name this task as medical coding). these icd codes categorize patient diagnoses, conditions, and treatments for billing, reporting, and treatment purposes @xcite . however, assigning icd codes is complex and requires expertise and time @xcite . recent advancements in medical pre-trained language models (plms) have made it possible to treat medical coding as a high-dimensional multilabel classification challenge @xcite . these ai models led to significant success in efficient icd coding @xcite . however, their transparency remains of great concern @xcite . therefore, developing automated interpretability methods is crucial to upholding transparency in medical coding processes. significant progress has occurred in the field of black-box interpretability, particularly concerning feature attribution, with the emergence of perturbation-based methods such as shap @xcite and its approximate counterpart lime @xcite . these techniques, rooted in information and game theory, are recognized for assessing feature relevance in detail by intelligently perturbing and ablating input features @xcite . while approximation methods have greatly improved the speed of calculating shapley values, exact computations remain expensive @xcite @xcite . the huge computational cost makes their application impractical towards automated medical icd coding since clinical notes usually contain thousands of high dimensional token embeddings in a vast multilabel prediction space @xcite . as such, we seek a human-interpretable approach that scales efficiently with large datasets, highlights essential features, and offers more comprehensive explanations of plm predictions. recent advancements in mechanistic interpretability methods @xcite have demonstrated the potential to surpass the computational challenges posed by traditional black-box approaches by elucidating the roles of specific neuron subsets within a network. this level of mechanistic understanding is precious in the medical field, where explaining the significance of a feature is as crucial as its identification. in recent years, the attention mechanism @xcite @xcite has been heavily used to interpret and explain the behavior of large transformer models. within the realm of automated medical icd coding, label attention (laat) variants are most prevalent due to their efficiency in handling extensive sequence lengths. they calculate an attention score for each token relative to each label, thereby identifying tokens crucial for icd code predictions @xcite . nevertheless, studies such as @xcite have questioned the interpretability and validity of explanations provided by attention mechanisms. for example, the laat mechanism may highlight incoherent or irrelevant tokens, such as stop words for highly medically specific icd codes. for instance, laat attributes the stop token "and" to the medically specific icd code "998.59 postoperative wound infection" despite conjunctions being irrelevant to medical prognosis, thus undermining the interpretability of laat as shown in figure 1 . we attribute this issue to be the result of neuron polysemanticity-where a single neuron responds to diverse, unrelated inputs-complicates direct interpretation @xcite . @xcite theorizes polysemanticity to be a form of superposition, occurring when the count of independent data features surpasses layer dimensions, leading to data features being represented by linear combinations of neurons. as a result, individual neurons are often directly uninterpretable @xcite . addressing this, sparse autoencoders @xcite have been applied to distill these complex dense representations into interpretable, sparse linear combinations, performing what is known as dictionary learning (dl). this strategy has effectively decomposed different language model layers, such as neural word embeddings @xcite , mlps @xcite , and residual connections @xcite , proving scalable and monosemantic. our paper generalizes these sparse coding concepts to better understand the attention mechanism and pre-trained language model (plm) embeddings by constructing effective dictionaries with laat to improve the interpretability of the medical icd coding task. we summarize our main contributions: ## methodology as depicted in figure 2 , our focus within dictionary learning involves explicitly building dictionaries where relevant tokens and icd codes are mapped to dictionary features. we examine two sparse autoencoder approaches aimed at creating interpretable representations from dense language model embeddings: one via l 1 minimization @xcite and another via spine's loss function @xcite . while our discussion primarily centers on the l 1 minimization technique for its widespread application and illustrative clarity regarding dictionary learning's objectives in section 3.1, further details on spine are provided in the appendix a.6. then, using our trained sparse autoencoder, we perform ablation studies to understand the downstream effects of dictionary features in section 4.1 and map the relevant icd codes to each dictionary feature as discussed in section 3.2. finally, we leverage sparse encoding and its ablation techniques in constructing our final dictionary, mapping both relevant tokens and icd codes to each dictionary feature in section 3.3, which is used in our new proposed method autocodedl in figure 3 . ## interpretability evaluations while there does not exist a unified singular definition of neural model interpretability, there is a general consensus that model interpretability methods should be both model explainable and human understandable @xcite . we define explainability as how much do the dictionary features learned predict the pre-trained language model's icd predictions. we define human understandability within the lens of monosemanticity: a dictionary feature is only human interpretable when the tokens that highly activate said dictionary feature are related and its underlying concept can be easily identified. sparse autoencoders. we explore two major dictionary learning (dl) sparse autoencoder approaches, spine by @xcite and l 1 described in section 3. baselines. following @xcite , we explore four unsupervised baseline encoders that we compare to a true dictionary encoding: independent component analysis (ica), capable of decomposing word embeddings into semantically meaningful independent components (ics) @xcite ; principal component analysis (pca), which has been used to analyze the structure of word embeddings @xcite ; an identity relu encoder, which effectively treats each element in the plm embedding as its own dictionary feature; and a random encoder. dataset and model. we train on plm embeddings from a 110m medical roberta model leveraged by the state of the art plm-icd coding model @xcite and evaluate our method using the cleaned mimic-iii dataset, using 38,427 clinical notes for training and 8,750 for evaluation, as detailed by @xcite . additional details on training are in appendix a.1. ## conclusion this study introduces a novel method that combines dictionary learning with label attention mechanisms to improve the interpretability of medical coding language models. by uncovering interpretable dictionary features from dense language embeddings, the proposed approach offers a dictionary-based rationale for icd code predictions, addressing the growing demand for transparency in automated healthcare decisions. this work lays the groundwork for future explorations in dictionary learning to enhance healthcare interpretability.