DILA: Dictionary Label Attention for Mechanistic Interpretability in High-dimensional Multi-label Medical Coding Prediction

TL;DR

This work tackles interpretability in extreme multilabel medical coding by introducing DILA, a framework that disentangles dense PLM embeddings into sparse dictionary features and maps them to ICD codes via a sparse global matrix $\mathbf{A_{f_{icd}}}$. An automated interpretability pipeline using medical LLMs identifies and summarizes these dictionary concepts, enabling global mechanistic explanations while maintaining competitive predictive performance on the MIMIC-III dataset. Across extensive experiments, DILA yields thousands of human-interpretable features, enables precise code-level explanations through feature ablations, and highlights the tradeoffs between interpretability and accuracy in a scalable, auditable system. Overall, DILA advances trustworthy AI in healthcare by coupling intrinsic sparsity with automated concept summarization and global explanation capabilities.

Abstract

Predicting high-dimensional or extreme multilabels, such as in medical coding, requires both accuracy and interpretability. Existing works often rely on local interpretability methods, failing to provide comprehensive explanations of the overall mechanism behind each label prediction within a multilabel set. We propose a mechanistic interpretability module called DIctionary Label Attention (\method) that disentangles uninterpretable dense embeddings into a sparse embedding space, where each nonzero element (a dictionary feature) represents a globally learned medical concept. Through human evaluations, we show that our sparse embeddings are more human understandable than its dense counterparts by at least 50 percent. Our automated dictionary feature identification pipeline, leveraging large language models (LLMs), uncovers thousands of learned medical concepts by examining and summarizing the highest activating tokens for each dictionary feature. We represent the relationships between dictionary features and medical codes through a sparse interpretable matrix, enhancing the mechanistic and global understanding of the model's predictions while maintaining competitive performance and scalability without extensive human annotation.

Figures (20)

• Figure 1: Medical coding as a high dimensional multilabel prediction task.
• Figure 2: DILA composes of three steps: First, we disentangle each token embedding into its dictionary features. Then, we project each set of dictionary features with our globally interpretable $\mathbf{A_{f_\text{icd}}}$ to generate our local explanation $\mathbf{A_{\text{laat}}}$ for downstream multilabel prediction. Finally, medical LLMs identify the learned dictionary feature to understand the learned relationships in $\mathbf{A_{f_\text{icd}}}$.
• Figure 3: Top 5 Dictionary Features for Diabetes-related ICD Codes.
• Figure 4: UMAP of $\mathbf{A_{f_\text{icd}}}$ with respect to all medical codes. We observe clusters of medical codes with relative distances that are intuitive. For instance, neuropathy is a common condition associated with diabetes, and vehicle accidents are more closely linked to bone and spinal fractures.
• Figure 5: Example of token sorting to acquire the necessary token contexts for $f_i$ for dictionary feature 1,871, relating to subdural hemotomas. The far left column indicates the position of the token in the text corpus.