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

TL;DR

This paper introduces a dictionary learning approach for medical coding with language models, improving interpretability by providing human-understandable explanations beyond traditional label attention methods.

Contribution

It presents a novel dictionary learning method that extracts sparse, interpretable features from language models, enhancing transparency in automated medical coding.

Findings

Dictionary features elucidate 90% of irrelevant tokens

Model behavior can be steered using dictionary features

Enhanced interpretability over label attention mechanisms

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 can steer model behavior, elucidate the hidden meanings of upwards of 90% of medically irrelevant tokens, and are human interpretable.