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

arXiv: https://arxiv.org/abs/2511.09559 · HTML: https://mlsanity.com/paper/2511.09559/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 31 references that are in the map.
- First posted on arXiv 31 Oct 2025.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [A Generalization of Transformer Networks to Graphs](https://mlsanity.com/paper/2012.09699/) (2020)
- [Explainable Prediction of Medical Codes from Clinical Text](https://mlsanity.com/paper/1802.05695/) (2018)
- [Code Synonyms Do Matter: Multiple Synonyms Matching Network for Automatic ICD Coding](https://mlsanity.com/paper/2203.01515/) (2022)
- [PLM-ICD: Automatic ICD Coding with Pretrained Language Models](https://mlsanity.com/paper/2207.05289/) (2022)
- [A Label Attention Model for ICD Coding from Clinical Text](https://mlsanity.com/paper/2007.06351/) (2020)
- [Knowledge Injected Prompt Based Fine-tuning for Multi-label Few-shot ICD Coding](https://mlsanity.com/paper/2210.03304/) (2022)
- [Mimic-IV-ICD: A new benchmark for eXtreme MultiLabel Classification](https://mlsanity.com/paper/2304.13998/) (2023)
- [CoRelation: Boosting Automatic ICD Coding Through Contextualized Code Relation Learning](https://mlsanity.com/paper/2402.15700/) (2024)
- [Multi-label Few-shot ICD Coding as Autoregressive Generation with Prompt](https://mlsanity.com/paper/2211.13813/) (2022)
- [Does Head Label Help for Long-Tailed Multi-Label Text Classification](https://mlsanity.com/paper/2101.09704/) (2021)
- [A Two-Stage Decoder for Efficient ICD Coding](https://mlsanity.com/paper/2306.00005/) (2023)
- [Accurate and Well-Calibrated ICD Code Assignment Through Attention Over Diverse Label Embeddings](https://mlsanity.com/paper/2402.03172/) (2024)
- [Multi-stage Retrieve and Re-rank Model for Automatic Medical Coding Recommendation](https://mlsanity.com/paper/2405.19093/) (2024)

## Abstract (arXiv)

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.