# Graph-Constrained Policy Learning for Extreme Clinical Code Prediction

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

## Key facts

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

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

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [BioBERT: a pre-trained biomedical language representation model for biomedical text mining](https://mlsanity.com/paper/1901.08746/) (2019)
- [Publicly Available Clinical BERT Embeddings](https://mlsanity.com/paper/1904.03323/) (2019)
- [Explainable Prediction of Medical Codes from Clinical Text](https://mlsanity.com/paper/1802.05695/) (2018)
- [Efficient Guided Generation for Large Language Models](https://mlsanity.com/paper/2307.09702/) (2023)
- [ICD Coding from Clinical Text Using Multi-Filter Residual Convolutional Neural Network](https://mlsanity.com/paper/1912.00862/) (2019)
- [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)
- [Automated Medical Coding on MIMIC-III and MIMIC-IV: A Critical Review and Replicability Study](https://mlsanity.com/paper/2304.10909/) (2023)
- [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)
- [HiCu: Leveraging Hierarchy for Curriculum Learning in Automated ICD Coding](https://mlsanity.com/paper/2208.02301/) (2022)

## Abstract (arXiv)

Clinical code prediction maps unstructured discharge summaries to ICD-10-CM leaf codes in a large, sparse, and deeply hierarchical label space. Most systems treat the task as flat multi-label classification, scoring codes independently and providing limited training signal for rare labels. We propose a graph-constrained traversal policy that formulates ICD prediction as a finite-horizon decision process over a pruned code hierarchy. A single language model descends the graph level by level, selecting valid child nodes until billable leaf codes are reached. This converts extreme multi-label prediction into sparse, hierarchy-aware subset decisions while guaranteeing structurally valid outputs. On MIMIC-IV discharge summaries, our best supervised policy, SFT-1+, achieves 0.709 micro-F1 on a curated 50-code subset and 0.527 micro-F1 on the full 15,761-code space, outperforming flat baselines including CAML, LAAT, and PLM-ICD. In the full setting, SFT-1+ improves over the strongest flat baseline by 0.044 micro-F1 and 0.157 macro-F1, suggesting that graph-constrained decomposition mitigates the rare-code bottleneck. A controlled factorial study evaluates architecture, training algorithm, and data budget. Across both scales, one shared policy matches a three-specialist cascade while avoiding its context-window overflow on 28-32% of full-space test notes. Increasing supervised trajectory data is the only intervention that consistently improves performance, while GRPO reinforcement learning provides no benefit over supervised continuation with matched data. These results show that simple graph-constrained policy learning can outperform more complex flat, cascaded, and reinforcement-learning alternatives for extreme clinical code prediction.