# Multi-stage Retrieve and Re-rank Model for Automatic Medical Coding Recommendation

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 25 references that are in the map.
- First posted on arXiv 29 May 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Explainable Prediction of Medical Codes from Clinical Text](https://mlsanity.com/paper/1802.05695/) (2018)
- [ICD Coding from Clinical Text Using Multi-Filter Residual Convolutional Neural Network](https://mlsanity.com/paper/1912.00862/) (2019)
- [Incorporating Hierarchy into Text Encoder: a Contrastive Learning Approach for Hierarchical Text Classification](https://mlsanity.com/paper/2203.03825/) (2022)
- [Knowledge Injected Prompt Based Fine-tuning for Multi-label Few-shot ICD Coding](https://mlsanity.com/paper/2210.03304/) (2022)
- [Multi-Label Classification of Patient Notes a Case Study on ICD Code Assignment](https://mlsanity.com/paper/1709.09587/) (2017)
- [Towards Automated ICD Coding Using Deep Learning](https://mlsanity.com/paper/1711.04075/) (2017)
- [A Comparative Study of Pretrained Language Models for Long Clinical Text](https://mlsanity.com/paper/2301.11847/) (2023)
- [Dilated Convolutional Attention Network for Medical Code Assignment from Clinical Text](https://mlsanity.com/paper/2009.14578/) (2020)

## Built on by (most cited first, 8 of 8)

- [A General Knowledge Injection Framework for ICD Coding](https://mlsanity.com/paper/2505.18708/) (2025)
- [Similarity-Dissimilarity Loss for Multi-label Supervised Contrastive Learning](https://mlsanity.com/paper/2410.13439/) (2024)
- [Don't Pay Attention, PLANT It: Pretraining Attention via Learning-to-Rank](https://mlsanity.com/paper/2410.23066/) (2024)
- [Unlocking Historical Clinical Trial Data with ALIGN: A Compositional Large Language Model System for Medical Coding](https://mlsanity.com/paper/2411.13163/) (2024)
- [LTR-ICD: A Ranking-Aware Framework for Automatic ICD Coding](https://mlsanity.com/paper/2510.13922/) (2025)
- [TraceCoder: Towards Traceable ICD Coding via Multi-Source Knowledge Integration](https://mlsanity.com/paper/2510.15267/) (2025)
- [Enhancing Rare Codes via Probability-Biased Directed Graph Attention for Long-Tail ICD Coding](https://mlsanity.com/paper/2511.09559/) (2025)
- [Multimodal Dual-Encoder Retrieval for Automated ICD Coding](https://mlsanity.com/paper/2610.04263/) (2026)

## Abstract (arXiv)

The International Classification of Diseases (ICD) serves as a definitive medical classification system encompassing a wide range of diseases and conditions. The primary objective of ICD indexing is to allocate a subset of ICD codes to a medical record, which facilitates standardized documentation and management of various health conditions. Most existing approaches have suffered from selecting the proper label subsets from an extremely large ICD collection with a heavy long-tailed label distribution. In this paper, we leverage a multi-stage ``retrieve and re-rank'' framework as a novel solution to ICD indexing, via a hybrid discrete retrieval method, and re-rank retrieved candidates with contrastive learning that allows the model to make more accurate predictions from a simplified label space. The retrieval model is a hybrid of auxiliary knowledge of the electronic health records (EHR) and a discrete retrieval method (BM25), which efficiently collects high-quality candidates. In the last stage, we propose a label co-occurrence guided contrastive re-ranking model, which re-ranks the candidate labels by pulling together the clinical notes with positive ICD codes. Experimental results show the proposed method achieves state-of-the-art performance on a number of measures on the MIMIC-III benchmark.