Enhancing Rare Codes via Probability-Biased Directed Graph Attention for Long-Tail ICD Coding

TL;DR

This paper introduces ProBias, a novel graph attention model that improves rare ICD code prediction by leveraging co-occurrence probabilities and enriched textual descriptions, achieving state-of-the-art results on benchmark datasets.

Contribution

ProBias is the first model to incorporate probability-biased directed graph attention and large language model-based code descriptions for long-tail ICD coding.

Findings

Significant improvement in macro F1 score for rare codes

Effective use of co-occurrence probabilities to guide attention

Enhanced code representations with external clinical context

Abstract

Automated international classification of diseases (ICD) coding aims to assign multiple disease codes to clinical documents and plays a critical role in healthcare informatics. However, its performance is hindered by the extreme long-tail distribution of the ICD ontology, where a few common codes dominate while thousands of rare codes have very few examples. To address this issue, we propose a Probability-Biased Directed Graph Attention model (ProBias) that partitions codes into common and rare sets and allows information to flow only from common to rare codes. Edge weights are determined by conditional co-occurrence probabilities, which guide the attention mechanism to enrich rare-code representations with clinically related signals. To provide higher-quality semantic representations as model inputs, we further employ large language models to generate enriched textual descriptions for ICD codes, offering external clinical context that complements statistical co-occurrence signals. Applied to automated ICD coding, our approach significantly improves the representation and prediction of rare codes, achieving state-of-the-art performance on three benchmark datasets. In particular, we observe substantial gains in macro-averaged F1 score, a key metric for long-tail classification.