DILA: Dictionary Label Attention for Mechanistic Interpretability in High-dimensional Multi-label Medical Coding Prediction
This addresses the need for interpretable predictions in medical coding, where existing methods often fail to explain the overall mechanism behind each label, though it is incremental as it builds on attention-based interpretability approaches.
The paper tackles the problem of providing comprehensive explanations for high-dimensional multi-label medical coding predictions by proposing a mechanistic interpretability module called DILA, which disentangles dense embeddings into sparse embeddings representing globally learned medical concepts and achieves at least 50% better human understandability while maintaining competitive performance.
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.