Graph-Constrained Policy Learning for Extreme Clinical Code Prediction

2026-07-11

AI Summary

This work addresses the extreme multi-label classification problem of mapping clinical text to leaf-level ICD-10-CM codes, where the label space is vast, sparse, and highly hierarchical, rendering conventional flat approaches ineffective for predicting rare codes. To overcome this challenge, the authors propose a language model–based graph-constrained traversal strategy that formulates code assignment as a sequential decision process over a pruned ICD hierarchy. A single model traverses the taxonomy top-down, selecting only valid child nodes at each level until reaching a leaf code, thereby guaranteeing structurally consistent and valid predictions. Evaluated on the MIMIC-IV dataset, the SFT-1+ model achieves a micro F1 score of 0.527 across the full label space of 15,761 codes, outperforming the strongest flat baseline by 0.044 in micro F1 and 0.157 in macro F1, substantially alleviating the bottleneck in rare code prediction.

Abstract

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.

Problem

• clinical code prediction
• ICD-10-CM
• hierarchical label space
• extreme multi-label classification
• rare labels

Innovation

• graph-constrained policy learning
• hierarchical code prediction
• extreme multi-label classification
• ICD-10-CM
• structured output generation