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

TL;DR

This paper introduces DILA, a novel interpretability module for high-dimensional multi-label medical coding that disentangles dense embeddings into a sparse, concept-based space, improving global interpretability and maintaining performance.

Contribution

The paper proposes DILA, a mechanism that transforms dense embeddings into a sparse, interpretable space with learned medical concepts, enhancing global understanding in medical coding models.

Findings

Sparse embeddings are at least 50% more human understandable.

Automated pipeline uncovers thousands of medical concepts.

Maintains competitive accuracy and scalability.

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.