Graph-Constrained Policy Learning for Extreme Clinical Code Prediction
Summary
Proposes a graph-constrained traversal policy that reformulates ICD-10-CM code prediction as a finite-horizon decision process over a pruned code hierarchy, outperforming flat baselines on MIMIC-IV discharge summaries.
View Cached Full Text
Cached at: 07/15/26, 04:17 AM
# Graph-Constrained Policy Learning for Extreme Clinical Code Prediction
Source: [https://arxiv.org/html/2607.11954](https://arxiv.org/html/2607.11954)
###### Abstract
Clinical code prediction maps unstructured discharge summaries to ICD\-10\-CM leaf codes from a large, sparse, and deeply hierarchical label space\. Most automated coding systems treat this problem as flat multi\-label classification, scoring each code independently and leaving rare labels with limited training signal\. We propose a graph\-constrained traversal policy that reformulates 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 from the current frontier until billable leaf codes are reached\. This constraint converts extreme multi\-label prediction into a sequence of sparse, hierarchy\-aware subset decisions while guaranteeing structurally valid outputs\.
On MIMIC\-IV discharge summaries, our best supervised traversal 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 label space, outperforming flat baselines including CAML, LAAT, and PLM\-ICD\. The gains are largest in the full setting, where SFT\-1\+improves over the strongest flat baseline by\+0\.044\+0\.044micro\-F1 and\+0\.157\+0\.157macro\-F1, suggesting that graph\-constrained decomposition mitigates the rare\-code bottleneck\. We further conduct a controlled factorial study over architecture, training algorithm, and data budget\. Across both scales, a single shared policy matches a three\-specialist cascade while avoiding its context\-window overflow failure on 28–32% of full\-space test notes; increasing supervised trajectory data is the only intervention that consistently improves performance; and GRPO reinforcement learning provides no benefit over supervised continuation under 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\.
## IIntroduction
Clinical code prediction maps unstructured documentation to standardized diagnosis codes used for reimbursement, quality measurement, cohort discovery, epidemiology, and risk adjustment\. In U\.S\. inpatient care, diagnoses are recorded using ICD\-10\-CM, a deeply hierarchical system that refines broad chapters into sections, subsections, categories, subcategories, and billable leaf codes\.
Automated ICD coding from discharge summaries is challenging because it combines long\-document understanding with extreme multi\-label prediction\. Notes may contain several thousand tokens, while the label space is large, sparse, and heavily imbalanced: common conditions appear frequently, but many clinically important leaf codes occur rarely\. Prior work has explored linear classifiers, CNNs, RNNs, label\-wise attention, transformer encoders, graph\-enhanced architectures, and retrieval\-augmented systems\[[1](https://arxiv.org/html/2607.11954#bib.bib1),[2](https://arxiv.org/html/2607.11954#bib.bib2),[3](https://arxiv.org/html/2607.11954#bib.bib3),[4](https://arxiv.org/html/2607.11954#bib.bib4),[5](https://arxiv.org/html/2607.11954#bib.bib5)\]; however, most methods still treat coding as flat multi\-label classification over leaf codes\.
This paper casts ICD\-10\-CM code prediction as graph\-constrained traversal, formalized as a finite\-horizon MDP over the pruned ICD\-10\-CM hierarchy \(Section[IV\-A](https://arxiv.org/html/2607.11954#S4.SS1)\)\. A single decoder\-only language model descends the hierarchy level by level\. At each depth, a lightweight environment exposes the valid children of the current frontier through tool calls; the model selects the children supported by the note, and the descent continues until billable leaves are reached\. Hard validity constraints restrict each selection to the candidate setCtC\_\{t\}, so structurally invalid codes receive zero probability by construction\. Instead of scoring the full label space,O\(\|L\|\)O\(\|L\|\), for each note, the traversal policy reduces inference toO\(∑ℓ\|Cℓ\|\)O\(\\sum\_\{\\ell\}\|C\_\{\\ell\}\|\), where\|Cℓ\|≪\|L\|\|C\_\{\\ell\}\|\\ll\|L\|after pruning to relevant branches\. Extreme multi\-label classification therefore becomes a short sequence of sparse, graph\-constrained subset decisions\. As Table[I](https://arxiv.org/html/2607.11954#S2.T1)details, our system combines hierarchy\-guided inference, hard structural constraints, a trained policy, multi\-turn context across depths, and formal MDP grounding; prior systems omit at least one of these properties\.
On MIMIC\-IV discharge summaries, the trained traversal policy significantly outperforms three strong flat baselines, CAML\[[1](https://arxiv.org/html/2607.11954#bib.bib1)\], LAAT\[[3](https://arxiv.org/html/2607.11954#bib.bib3)\], and PLM\-ICD\[[14](https://arxiv.org/html/2607.11954#bib.bib14)\], on both a curated 50\-code subset and the full 15,761\-code ICD\-10\-CM space\. On the full label space, SFT\-1\+exceeds the strongest flat baseline, LAAT, by\+0\.044\+0\.044micro\-F1 and\+0\.157\+0\.157macro\-F1\. CAML and PLM\-ICD collapse to 0\.002 macro\-F1, indicating limited ability to discriminate among 15,761 candidates without structural guidance\. These results show that graph\-constrained decomposition partially insulates against the rare\-code bottleneck\.
We further ask whether additional architectural or algorithmic complexity is warranted once graph\-constrained traversal is introduced\. We conduct a controlled factorial study on the same MIMIC\-IV data, isolating architecture, training algorithm, and annotation budget: one shared traversal policy versus a three\-specialist cascade, supervised fine\-tuning versus GRPO reinforcement learning, and 2,000 versus 5,000 gold trajectories\. Three findings hold at both label\-space scales\. First, the shared policy matches the cascade on accuracy while avoiding the cascade’s context\-window overflow failure on 28–32% of full\-space test notes\. Second, additional supervised trajectories are the only intervention that reliably improves performance, with gains roughly doubling from the 50\-code to the full\-space setting\. Third, GRPO does not improve over supervised continuation under matched data\. Together, these findings support a simple practical conclusion: use one shared traversal policy and invest in supervised trajectories rather than cascaded architectures or reinforcement learning\.
## IIRelated Work
### II\-AHierarchical and Graph\-Aware Coding
Because ICD is hierarchical, several models incorporate code structure\. HyperCore represents ICD codes with hyperbolic and co\-graph embeddings to exploit hierarchical relations and co\-occurrence patterns\[[4](https://arxiv.org/html/2607.11954#bib.bib4)\]\. Later graph\-based and knowledge\-guided models\[[15](https://arxiv.org/html/2607.11954#bib.bib15),[17](https://arxiv.org/html/2607.11954#bib.bib17),[18](https://arxiv.org/html/2607.11954#bib.bib18)\]use code descriptions, ancestor relations, and label graphs to improve rare\-code generalization\. These methods show that ICD structure carries predictive signal, but most use the hierarchy as an auxiliary representation or regularizer rather than traversing it directly\. In Section[V](https://arxiv.org/html/2607.11954#S5), we show that direct graph\-constrained traversal yields larger gains as the label space grows\.
### II\-BClinical Foundation Models and LLM Coding
Clinical language models such as ClinicalBERT\[[12](https://arxiv.org/html/2607.11954#bib.bib12)\]and BioBERT\[[13](https://arxiv.org/html/2607.11954#bib.bib13)\]showed that biomedical pretraining improves representation learning for healthcare NLP\. For ICD coding, transformer encoders\[[14](https://arxiv.org/html/2607.11954#bib.bib14)\]and generative prompting have been proposed as alternatives to CNN/RNN pipelines, but long discharge notes and large code vocabularies remain limiting factors\. Prompt\-only LLM coding also raises reliability concerns\[[21](https://arxiv.org/html/2607.11954#bib.bib21)\], including nonexistent codes, mixed ICD revisions, or plausible but invalid descendants\. Closest to our setting, Boyle et al\.\[[22](https://arxiv.org/html/2607.11954#bib.bib22)\]use the ICD hierarchy to search sparsely rather than score the full vocabulary, but rely on prompting alone\. We instead bind generation to an explicit candidate set from the current ICD\-10\-CM graph and train the traversal policy on oracle demonstrations\.
### II\-CPositioning
Table[I](https://arxiv.org/html/2607.11954#S2.T1)situates our approach against representative prior systems along five axes that together define the design space for hierarchical ICD coding\.
TABLE I:Design\-space comparison of representative ICD coding systems\. ✓ = yes, ✗ = no, ~ = partial\.Our formulation combines all five properties: the ICD\-10\-CM graph is traversed at inference time with hard validity constraints, the traversal policy is trained directly on oracle demonstrations, decisions at each depth condition on the full accumulated history, and the procedure is grounded in a formal MDP \(Section[IV\-A](https://arxiv.org/html/2607.11954#S4.SS1)\)\.
## IIIData and Task Construction
### III\-ASource and Cohort Construction
Following the MIMIC\-IV coding setup of Boyle et al\.\[[22](https://arxiv.org/html/2607.11954#bib.bib22)\], we construct our dataset from MIMIC\-IV\[[8](https://arxiv.org/html/2607.11954#bib.bib8)\]and MIMIC\-IV\-Note\[[9](https://arxiv.org/html/2607.11954#bib.bib9)\], joining admissions, diagnosis codes, and discharge notes by admission identifier\. We retain admissions with non\-null discharge notes coded exclusively under ICD\-10, discarding any admission containing at least one ICD\-9 code to avoid mixing code schemas\. This yields 122,281 admissions, each represented as a discharge note and gold diagnosis code set\.
We then validate diagnosis codes against the official active ICD\-10\-CM code list \(April 2026 release, 74,719 codes\), retaining only admissions whose gold codes are present in this reference set\. This removes admissions with deprecated, malformed, or non\-billable codes, leaving 122,197 notes\. This validated cohort serves as the base population for both label spaces and downstream splits\. All data access, processing, and model training were conducted within an AWS SageMaker instance operating under a Business Associate Agreement \(BAA\) with Amazon, in compliance with MIMIC\-IV data use requirements\.
### III\-BLabel Space and Code Graph
Similar to Boyle et al\.\[[22](https://arxiv.org/html/2607.11954#bib.bib22)\], we define two label spaces over the validated cohort:*all*, containing every observed code with 15,761 leaves, and*frequent*, containing the 50 most common codes\. These provide controlled settings at different points on the difficulty–realism trade\-off\. For each setting, we build a pruned ICD\-10\-CM graph retaining gold codes and their ancestors up to the chapter level, while discarding unreferenced branches \(Table[II](https://arxiv.org/html/2607.11954#S3.T2)\)\.
TABLE II:Label space, code graph, and split statistics\.
### III\-CSplits and Statistics
For each label space, we sample 6,500 notes from the validated cohort and partition them into four disjoint splits: 2,000 for supervised fine\-tuning, 3,000 for continuation training, 500 for validation, and 1,000 for testing\. The continuation split is used for both additional supervised fine\-tuning and GRPO reinforcement learning, supporting the factorial design in Section[IV\-E](https://arxiv.org/html/2607.11954#S4.SS5)\. Splitting is performed at the note level with seed 42 so that no admission appears in more than one split\. Test and validation sets are held fixed across both label spaces and all training arms\.
Table[II](https://arxiv.org/html/2607.11954#S3.T2)reports summary statistics\. Notes in the*all*setting carry more gold codes on average than in the*frequent*setting because*frequent*includes only diagnoses falling within the top\-50 codes\. Because ICD coding comparisons are sensitive to cohort construction, split strategy, and code\-version handling\[[10](https://arxiv.org/html/2607.11954#bib.bib10),[20](https://arxiv.org/html/2607.11954#bib.bib20)\], we fix note\-level splits, evaluate both label spaces, and hold test and validation sets constant across all arms\.
### III\-DGold Trajectory Construction
Each training example is a multi\-turn gold trajectory demonstrating the correct graph traversal for a note\. We expand each gold ICD\-10\-CM code to its full ancestor path, from chapter root to billable leaf, and replay the resulting paths through the traversal environment\. At each turn, the environment presents candidate nodes and the oracle selects exactly those on a gold ancestor path\. Leaf codes are recorded automatically upon expansion, and the trajectory ends with anANSWERturn listing all recorded leaf codes\.
This yields one deterministic sequence per note, with assistant tool calls and environment responses tracing every gold code from root to leaf\. During tokenization, environment\-response tokens are masked from the training loss, so the model is supervised only on its own decisions: which nodes to select and when to answer\.
## IVMethodology
### IV\-ATraversal as a Markov Decision Process
We formalize ICD\-10\-CM code prediction as a finite\-horizon MDP over the pruned ICD\-10\-CM graph\.𝒢=\(𝒱,ℰ\)\\mathcal\{G\}=\(\\mathcal\{V\},\\mathcal\{E\}\), where𝒱\\mathcal\{V\}is the set of code\-graph nodes andℰ\\mathcal\{E\}the directed parent→\\rightarrowchild edges\.
State\.A statest=\(x,Ft,Rt\)s\_\{t\}=\(x,F\_\{t\},R\_\{t\}\)comprises the discharge notexx\(fixed throughout the episode\), the*frontier*Ft⊆𝒱F\_\{t\}\\subseteq\\mathcal\{V\}of nodes pending expansion, and the setRtR\_\{t\}of billable leaf codes recorded so far\. The initial state iss0=\(x,\{root\},∅\)s\_\{0\}=\(x,\\\{\\textsc\{root\}\\\},\\emptyset\)\.
Action\.At each step the policy first expands every node inFtF\_\{t\}, revealing the candidate set
Ct=⋃v∈Ftchildren𝒢\(v\)\.C\_\{t\}=\\bigcup\_\{v\\in F\_\{t\}\}\\mathrm\{children\}\_\{\\mathcal\{G\}\}\(v\)\.The actionat⊆Cta\_\{t\}\\subseteq C\_\{t\}is the subset of children the note supports\. The constraintat⊆Cta\_\{t\}\\subseteq C\_\{t\}is enforced by the environment, so structurally invalid selections receive zero probability by construction rather than by penalty\.
Transition\.The transition is deterministic\. Any selected node that is a billable leaf is added toRtR\_\{t\}; the remainder form the next frontier:
Ft\+1\\displaystyle F\_\{t\+1\}=\{v∈at∣vis not a leaf\},\\displaystyle=\\\{v\\in a\_\{t\}\\mid v\\text\{ is not a leaf\}\\\},Rt\+1\\displaystyle R\_\{t\+1\}=Rt∪\{v∈at∣vis a leaf\}\.\\displaystyle=R\_\{t\}\\cup\\\{v\\in a\_\{t\}\\mid v\\text\{ is a leaf\}\\\}\.The episode terminates whenFt\+1=∅F\_\{t\+1\}=\\emptyset\.
Reward\.For supervised fine\-tuning the reward is implicit: the policy is trained to imitate oracle trajectories constructed from the gold code sets \(Section[III\-D](https://arxiv.org/html/2607.11954#S3.SS4)\)\. For reinforcement learning \(GRPO\-1, Section[IV\-E](https://arxiv.org/html/2607.11954#S4.SS5)\), we define an explicit terminal reward
r\(sT\)=12\[set\-F1\(RT,C∗\)\+path\-F1\(RT,C∗\)\],r\(s\_\{T\}\)=\\tfrac\{1\}\{2\}\\left\[\\mathrm\{set\\mbox\{\-\}F1\}\(R\_\{T\},C^\{\*\}\)\+\\mathrm\{path\\mbox\{\-\}F1\}\(R\_\{T\},C^\{\*\}\)\\right\],whereC∗C^\{\*\}is the gold code set and path\-F1 awards partial credit for correct ancestor paths \(Equation[1](https://arxiv.org/html/2607.11954#S4.E1)\)\. Rewards are zero for all non\-terminal transitions; only the terminal state receives the reward defined above\.
Policy\.The policyπθ\(at∣st\)\\pi\_\{\\theta\}\(a\_\{t\}\\mid s\_\{t\}\)is parameterized by a single QLoRA adapter over a decoder\-only backbone\. Because the full traversal history\(a0,C0,…,at−1,Ct−1\)\(a\_\{0\},C\_\{0\},\\ldots,a\_\{t\-1\},C\_\{t\-1\}\)is retained in the context window, the Markov property holds with respect to the model’s input: each action conditions on the note, the complete descent so far, and the current candidate set\.
The horizon is bounded by the depth of the deepest selected path in𝒢\\mathcal\{G\}\(typically 5–7 expand\-and\-select rounds\)\. The effective action space at each step satisfies\|Ct\|≪\|𝒱\|\|C\_\{t\}\|\\ll\|\\mathcal\{V\}\|, reducing the combinatorial multi\-label problem to a short sequence of sparse, graph\-constrained subset selections\. Section[IV\-B](https://arxiv.org/html/2607.11954#S4.SS2)describes the concrete implementation of this MDP as a tool\-calling environment\.
### IV\-BTraversal Environment
We implement the MDP defined in Section[IV\-A](https://arxiv.org/html/2607.11954#S4.SS1)as a guided descent through the pruned ICD\-10\-CM graph, mediated by a lightweight environment that maintains a traversal agenda and exposes two tools and a terminal action:
- •expand\_level\(\)pops all currently pending nodes from the agenda and reveals their children, grouped by parent\. If a pending node is a billable leaf \(a code with no descendants in the pruned ICD\-10\-CM graph\) it is recorded automatically as a final answer rather than expanded, since there is nothing below it to drill into\. When the agenda is empty, the environment returnsDONEand lists all recorded codes\.
- •select\(children=\[\.\.\.\]\)accepts a list of child IDs drawn exclusively from the most recentexpand\_level\(\)output and queues them as pending for the next expansion\. The environment validates every ID against the set of children just presented; invalid IDs are rejected, which realizes the constrained\-decoding principle of restricting generation to a valid candidate set\[[19](https://arxiv.org/html/2607.11954#bib.bib19)\]at the level of tool arguments\. An empty list is permitted, allowing the model to skip an entire level when no candidates are supported by the note\.
- •ANSWER: code1, code2, \.\.\.terminates the trajectory and emits all recorded leaf codes as the model’s final prediction\.
The environment enforces breadth\-first traversal: eachexpand\_level\(\)call processes every pending node at the same hierarchy depth in a single turn, so a complete traversal spans 5\-7 depth levels \(approximately 12–14 model turns\) regardless of how many codes the note carries \(Figure[1](https://arxiv.org/html/2607.11954#S4.F1)\)\.
Discharge note \(excerpt\)\.“ … M with a h/oHTNwho presents with 5 days of headache, fatigue, fever …Secondary Diagnoses: Hypertension\.”
*No lipid disorder documented anywhere in the note\.*ROOTCh9Ch4I10 – I1AE78I10E785CirculatoryHypertensive dis\.Ess\. hypertensionEndocrine & metabolicLipoprotein metab\.Hyperlipidemiarecorded \(hit\)missedChapterSectionBillable codeFigure 1:A real traversal from the top\-50 test set \(gold=\{=\\\{I10, E785\}\\\}\)\. Fromrootthe policy selects Chapter 9 \(circulatory\), enters*Hypertensive diseases*, and records the billable leafI10\(essential hypertension\), which is correct \(blue\)\. It never expands the endocrine/metabolic branch \(red, dashed\), so the gold codeE785\(hyperlipidemia\) is missed: the comorbidity is in the coded record but never documented in the note body, illustrating that recall is bounded by what the narrative states\.The model interacts with the environment through Qwen3’s function\-calling interface\. The rollout driver parses each<tool\_call\>, executes the requested tool, and returns the result as a tool\-response turn\. During training, model\-generated tokens are trainable and environment\-injected tokens are masked; during evaluation, active rollouts are batched by wave for tractability at 1,000\+ notes\.
### IV\-CShared Traversal Policy \(SFT\-1\)
A single traversal policy makes all selection decisions across hierarchy depths within one continuous context window, conditioning each level on the accumulated traversal history\. Prior LLM traversal of ICD hierarchies\[[22](https://arxiv.org/html/2607.11954#bib.bib22)\]relies on prompting alone; we train the policy directly on oracle demonstrations\. We fine\-tune on the 2,000\-note SFT split using the gold trajectories\. Each trajectory is replayed through the rollout driver to produce the trainable/masked split described in Section[IV\-B](https://arxiv.org/html/2607.11954#S4.SS2)\. All fine\-tuning uses QLoRA\[[7](https://arxiv.org/html/2607.11954#bib.bib7)\]with rankr=16r\{=\}16on a 4\-bit quantized base model; continuation arms warm\-start from the preceding adapter\.
Shared Traversal PolicyDischarge noteQwen3\-4B\-Instruct\-25074\-bit quantizedQLoRAr=16r\{=\}16\(single adapter\)Traversal env\.expand / selectICD\-10\-CM codesmulti\-turnCascade BaselineDischarge noteQwen3\-4B\-Instruct\-25074\-bit quantizedA2A\_\{2\}subsect\.A1A\_\{1\}chapterA3A\_\{3\}codeICD\-10\-CM codesFigure 2:Architecture comparison\.Left:shared traversal policy \(single adapter, multi\-turn\)\.Right:cascade baseline \(three specialist adapters, single\-turn each\)\. Both share the same 4\-bit quantized backbone\.
### IV\-DCascade Baseline \(SFT\-3\)
The cascade decomposes the same traversal into three independent single\-turn classification problems, one per hierarchy depth: a chapter specialist, a subsection specialist, and a code specialist\. Each specialist receives the discharge note and a flat list of candidates at its assigned level, and responds with a single line \(SELECTED: id1, id2, \.\.\.\) choosing every candidate the note supports\. There is no tool\-call protocol, no multi\-turn interaction, and no shared context across levels\. At inference time the three are executed sequentially, so the chapter specialist’s selections determine which candidates the subsection specialist sees, and in turn which billable codes are presented to the code specialist\. All three share a common system prompt \(Supplementary Materials, Appendix A\)\.
The candidate lists at the first two levels \(L1: chapters, L2: subsections\) are modest at both label\-space scales, but the third level \(L3: billable codes\) presents all billable\-leaf descendants of selected subsections in one batched prompt\. This makes the cascade efficient, requiring three model calls per note, but creates a scaling limitation\. Under the*all*label space, the L3 prompt often exceeds the 32,000\-token cap used for a single 46GB GPU\. During training, over\-cap L3 examples are dropped rather than truncated; during evaluation, the candidate tail is truncated, imposing a hard recall ceiling on 28% of SFT\-3 test notes and 32% of SFT\-3\+test notes\. The shared traversal policy has 0% truncation at the same scale\.
We train the three specialists on the same 2,000\-note SFT split used for SFT\-1, with each note re\-expressed as three single\-turn examples \(one per level\), yielding 6,000 training examples total\.
### IV\-EContinuation and RL Arms
The shared traversal policy and cascade baseline each use 2,000 notes for initial training\. To isolate the effects of data, architecture, and training algorithm, we define three continuation arms on the same 3,000\-note continuation split\.
SFT\-1\+warm\-starts from SFT\-1 and continues supervised fine\-tuning on 3,000 additional gold trajectories constructed as in Sections[III\-D](https://arxiv.org/html/2607.11954#S3.SS4)and[IV\-B](https://arxiv.org/html/2607.11954#S4.SS2)\.
SFT\-3\+warm\-starts each SFT\-3 specialist and continues supervised fine\-tuning on the same continuation split, re\-expressed as per\-level single\-turn examples\.
GRPO\-1warm\-starts from SFT\-1 and applies Group Relative Policy Optimization on the same 3,000 notes\. The model generates live rollouts against the environment, and each rollout is scored with the terminal reward from Section[IV\-A](https://arxiv.org/html/2607.11954#S4.SS1)\. We run GRPO\-1 on the*frequent*label space only; the*all*variant was omitted because training time scales to approximately 2–3 GPU\-days and the*frequent*result already shows no benefit over supervised continuation\.
### IV\-FEvaluation Protocol
All arms are evaluated on the same held\-out test set \(1,000 notes\) and validation set \(500 notes\), shared across both label spaces and all training arms \(Section[III\-C](https://arxiv.org/html/2607.11954#S3.SS3)\)\. Evaluation rollouts are executed through the batched rollout driver \(Section[IV\-B](https://arxiv.org/html/2607.11954#S4.SS2)\) under greedy decoding \(temperature 0\) for determinism and paired comparability\.
Zero\-shot floors\.To anchor the fine\-tuning results, we evaluate two zero\-shot baselines using the base model base model \(Qwen3\-4B\-Instruct\-2507\[[11](https://arxiv.org/html/2607.11954#bib.bib11)\]\) with no task\-specific training\. The traversal\-format base \(Base in Table[III](https://arxiv.org/html/2607.11954#S5.T3)\) receives the same system prompt, tool schema, and discharge note as SFT\-1, scoring near\-zero micro\-F1 on both label spaces \(0\.001\), reflecting protocol\-following failure rather than lack of clinical knowledge\. The cascade base \(Base\-3\) evaluates the same untrained model under the single\-turnSELECTED:format of SFT\-3, clearing the floor meaningfully \(0\.229 on*frequent*, 0\.058 on*all*\)\. The gap between the two floors quantifies the protocol\-following overhead that the gold trajectories \(Section[III\-D](https://arxiv.org/html/2607.11954#S3.SS4)\) are designed to close\.
Metrics\.We report micro\-averaged precision, recall, and F1 \(pooled across all notes and codes, weighting each prediction equally\) and macro\-averaged precision, recall, and F1 \(averaged across the codes present in the gold set, weighting rare and common codes equally\)\. Micro\-F1 is the headline metric; macro\-F1 surfaces rare\-code performance, which dominates the*all*label space\. We additionally report path\-F1, a hierarchical variant that awards partial credit for predictions landing in the correct region of the ICD\-10\-CM hierarchy rather than requiring an exact leaf match\. For a codecc, letA\(c\)=\{c\}∪anc\(c\)A\(c\)=\\\{c\\\}\\cup\\mathrm\{anc\}\(c\)denote its augmentation with all ancestors along the path toward the root \(the root itself excluded\), and extend this to a set of codes byA\(C\)=⋃c∈CA\(c\)A\(C\)=\\bigcup\_\{c\\in C\}A\(c\)\. WritingC^i\\hat\{C\}\_\{i\}andCiC\_\{i\}for the predicted and gold codes of noteii, path\-F1 replaces exact\-match set membership with membership over the augmented sets:
path\-P=∑i\|A\(C^i\)∩A\(Ci\)\|∑i\|A\(C^i\)\|,path\-R=∑i\|A\(C^i\)∩A\(Ci\)\|∑i\|A\(Ci\)\|,\\text\{path\-P\}=\\frac\{\\sum\_\{i\}\|A\(\\hat\{C\}\_\{i\}\)\\cap A\(C\_\{i\}\)\|\}\{\\sum\_\{i\}\|A\(\\hat\{C\}\_\{i\}\)\|\},\\;\\;\\text\{path\-R\}=\\frac\{\\sum\_\{i\}\|A\(\\hat\{C\}\_\{i\}\)\\cap A\(C\_\{i\}\)\|\}\{\\sum\_\{i\}\|A\(C\_\{i\}\)\|\},\(1\)with path\-F1 their harmonic mean\. Equation[1](https://arxiv.org/html/2607.11954#S4.E1)gives the micro variant \(sums pooled over notes\); the macro variant averages the per\-code F1 over the augmented gold codes\. Because a mispredicted leaf still shares ancestors with the true code, this rewards a prediction that lands on the correct branch over one that is entirely unrelated\.
Statistical testing\.We use a paired bootstrap to assess differences between arms\. Because micro\-F1 and macro\-F1 are pooled statistics that do not decompose into per\-note scores, a paired t\-test is not applicable\. For each of 5,000 resamples we draw 1,000 test notes with replacement and recompute both arms’ pooled metrics on the identical resampled set\. We report the mean difference, 95% confidence interval, and two\-sided p\-value \(values below1/5,0001/5\{,\}000reported asp<0\.001p<0\.001\)\. We report uncorrected p\-values for approximately fifteen planned comparisons\.111The large effects \(supervised data\) survive Bonferroni correction\. The small architecture deltas that reach nominal significance \(p=0\.002p=0\.002–0\.0260\.026\) would not all survive correction atα=0\.05/15≈0\.003\\alpha=0\.05/15\\approx 0\.003, and we interpret them accordingly\.
We additionally report per\-level selection accuracy \(chapter, subsection, billable code\) to diagnose where recall is lost, and for the cascade under the*all*label space, the fraction of test notes affected by L3 candidate\-list truncation \(Section[IV\-D](https://arxiv.org/html/2607.11954#S4.SS4)\)\.
## VResults
### V\-AOverall Performance
Table[III](https://arxiv.org/html/2607.11954#S5.T3)reports micro\- and macro\-averaged precision, recall, and F1 for all arms on the held\-out test set \(n=1,000n\{=\}1\{,\}000\) at both label\-space scales\. The results are summarized below\.
TABLE III:Micro\- and macro\-averaged performance on the held\-out test set \(n=1,000n\{=\}1\{,\}000\) by training arm and label space\. CAML, LAAT, and PLM\-ICD are flat multi\-label baselines retrained on the same splits\. Path\-F1 gives hierarchical partial credit \(Section[IV\-F](https://arxiv.org/html/2607.11954#S4.SS6)\)\.The traversal policy outperforms flat baselines, with the margin growing at scale\.Three flat multi\-label baselines anchor these results externally\. CAML and PLM\-ICD both reach competitive micro\-F1 on*frequent*\(0\.668 and 0\.691\) but collapse to near\-identical scores on*all*\(0\.194 and 0\.196\), with macro\-F1 falling to 0\.002 for both—indicating that neither convolutional attention nor pretrained transformer encoders can discriminate among 15,761 candidates without structural guidance\. Only LAAT survives at scale \(0\.483 on*all*\), and SFT\-1 trained on just 2,000 trajectories already matches it \(0\.478 vs\. 0\.483\)\. SFT\-1\+significantly outperforms all three baselines at both scales \(Table[IV](https://arxiv.org/html/2607.11954#S5.T4)\), with the margin widening sharply on*all*:\+0\.044\+0\.044micro\-F1 over LAAT and\+0\.157\+0\.157macro\-F1, confirming that the traversal formulation’s advantage is concentrated where flat scoring is weakest\. Notably, PLM\-ICD’s pretrained encoder offers no advantage over CAML’s CNN at the full scale, suggesting that the bottleneck is the flat output layer rather than the representation \- with only 2,000 training notes the vast majority of the 15,761 output neurons see fewer than five positive examples\.
More supervised data helps, and the effect grows with scale\.The\+3,000\+3\{,\}000\-note continuation \(SFT\-1\+over SFT\-1\) is the largest gain in the table:\+0\.022\+0\.022micro\-F1 on*frequent*\(0\.687→0\.7090\.687\\rightarrow 0\.709\) and\+0\.049\+0\.049on*all*\(0\.478→0\.5270\.478\\rightarrow 0\.527\)\. The lift roughly doubles from the 50\-code to the full label space\. SFT\-1\+is the best arm in both settings, and the same pattern holds for the cascade \(SFT\-3\+over SFT\-3 by comparable margins\)\. Macro\-F1 gains track micro, so the improvement is not confined to high\-frequency codes\.
GRPO does not improve over supervised continuation\.GRPO\-1 \(0\.688 micro\-F1\) ties its SFT\-1 warm\-start \(0\.687\) and falls below SFT\-1\+\(0\.709\) on both micro and macro\. Training directly on the coding metric via RL on 3,000 notes offers no benefit over simply continuing supervised fine\-tuning on the same data from the same initialization\.
Neither architecture wins consistently\.The cascade \(SFT\-3, SFT\-3\+\) lands within roughly one micro\-F1 point of the corresponding shared\-policy arm \(SFT\-1, SFT\-1\+\) at both scales\. On*all*, SFT\-3\+scores 0\.513 vs\. SFT\-1\+’s 0\.527 on micro but 0\.260 vs\. 0\.256 on macro; the deltas flip sign across metrics and variants\. Statistical tests follow in Section[V\-B](https://arxiv.org/html/2607.11954#S5.SS2)\. Where the cascade does lose is structurally: 28–32% of test notes suffer L3 candidate\-list truncation under the*all*label space \(Section[IV\-D](https://arxiv.org/html/2607.11954#S4.SS4)\), a constraint the shared policy never encounters\.
### V\-BStatistical Comparisons
Table[IV](https://arxiv.org/html/2607.11954#S5.T4)reports paired\-bootstrap estimates \(5,000 resamples\) of the difference in pooled micro\- and macro\-F1 between matched arms, with 95% confidence intervals and two\-sided p\-values\. All external comparisons are significant \(p<0\.001p<0\.001\)\. The margins are modest on*frequent*\(\+0\.014\+0\.014to\+0\.042\+0\.042micro\-F1\) but large on*all*\(\+0\.044\+0\.044to\+0\.333\+0\.333micro\-F1\), confirming that the traversal formulation’s advantage grows with label\-space scale\. On*frequent*, SFT\-1 without continuation data already matches LAAT \(p=0\.08p=0\.08, n\.s\.\)\.
Among the internal comparisons, the data effect is the only robust finding: SFT\-1\+over SFT\-1 is significant on all four tests \(p<0\.001p<0\.001\), with the lift doubling from*frequent*to*all*\. GRPO\-1 ties SFT\-1 \(microp=0\.892p=0\.892\) and falls significantly below SFT\-1\+\(p<0\.001p<0\.001\)\. Architecture deltas are small \(≤0\.014\\leq 0\.014\), inconsistent in sign, and would not survive Bonferroni correction\.
TABLE IV:Paired bootstrap comparisons \(5,000 resamples\):Δ\\Deltamicro\- and macro\-F1 withpp\-values\. External baselines \(CAML, LAAT, PLM\-ICD\) above the rule; internal architecture, algorithm, and data comparisons below\.VariantComparisonMetricNewBaseΔ\\DeltappVerdictFrequentSFT\-1\+vs CAMLMicro0\.70900\.6675\+0\.042\+0\.042<0\.001<0\.001↑\\uparrowhigherMacro0\.65140\.5819\+0\.070\+0\.070<0\.001<0\.001↑\\uparrowhigherSFT\-1\+vs LAATMicro0\.70900\.6951\+0\.014\+0\.014<0\.001<0\.001↑\\uparrowhigherMacro0\.65140\.6157\+0\.036\+0\.036<0\.001<0\.001↑\\uparrowhigherSFT\-1\+vs PLM\-ICDMicro0\.70900\.6912\+0\.018\+0\.018<0\.001<0\.001↑\\uparrowhigherMacro0\.65140\.6125\+0\.039\+0\.039<0\.001<0\.001↑\\uparrowhigherSFT\-1 vs LAATMicro0\.68740\.6951−0\.008\-0\.0080\.0840\.084n\.s\.Macro0\.62290\.6157\+0\.007\+0\.0070\.1800\.180n\.s\.SFT\-1\+vs SFT\-1Micro0\.70900\.6874\+0\.022\+0\.022<0\.001<0\.001↑\\uparrowhigherMacro0\.65140\.6229\+0\.029\+0\.029<0\.001<0\.001↑\\uparrowhigherGRPO\-1 vs SFT\-1Micro0\.68780\.6874\+0\.000\+0\.0000\.8920\.892n\.s\.Macro0\.61740\.6229−0\.006\-0\.0060\.1040\.104n\.s\.GRPO\-1 vs SFT\-1\+Micro0\.68780\.7090−0\.021\-0\.021<0\.001<0\.001↓\\downarrowlowerMacro0\.61740\.6514−0\.034\-0\.034<0\.001<0\.001↓\\downarrowlowerSFT\-3 vs SFT\-1Micro0\.68910\.6874\+0\.002\+0\.0020\.6380\.638n\.s\.Macro0\.62990\.6229\+0\.007\+0\.0070\.1420\.142n\.s\.SFT\-3\+vs SFT\-1\+Micro0\.70280\.7090−0\.006\-0\.0060\.0860\.086n\.s\.Macro0\.63940\.6514−0\.012\-0\.0120\.0170\.017↓\\downarrowlowerAllSFT\-1\+vs CAMLMicro0\.52660\.1938\+0\.333\+0\.333<0\.001<0\.001↑\\uparrowhigherMacro0\.25620\.0018\+0\.254\+0\.254<0\.001<0\.001↑\\uparrowhigherSFT\-1\+vs LAATMicro0\.52660\.4829\+0\.044\+0\.044<0\.001<0\.001↑\\uparrowhigherMacro0\.25620\.0988\+0\.157\+0\.157<0\.001<0\.001↑\\uparrowhigherSFT\-1\+vs PLM\-ICDMicro0\.52660\.1958\+0\.331\+0\.331<0\.001<0\.001↑\\uparrowhigherMacro0\.25620\.0018\+0\.254\+0\.254<0\.001<0\.001↑\\uparrowhigherSFT\-1\+vs SFT\-1Micro0\.52660\.4775\+0\.049\+0\.049<0\.001<0\.001↑\\uparrowhigherMacro0\.25620\.2172\+0\.039\+0\.039<0\.001<0\.001↑\\uparrowhigherSFT\-3 vs SFT\-1Micro0\.47500\.4775−0\.003\-0\.0030\.5980\.598n\.s\.Macro0\.22690\.2172\+0\.010\+0\.0100\.0260\.026↑\\uparrowhigherSFT\-3\+vs SFT\-1\+Micro0\.51260\.5266−0\.014\-0\.0140\.0020\.002↓\\downarrowlowerMacro0\.26020\.2562\+0\.004\+0\.0040\.5450\.545n\.s\.
### V\-CHierarchical Recall by Depth
To diagnose where in the hierarchy codes are lost, we compute micro\-pooled recall at each depth of the ICD\-10\-CM tree: at depthdd, each predicted and gold leaf code is collapsed to its depth\-ddancestor, and recall is computed over the resulting sets\.222Codes whose paths are shorter thandddo not contribute at that depth, so deeper columns reflect only codes that actually reach that level of the hierarchy\. The denominators therefore shrink at deeper levels\.Figure[3](https://arxiv.org/html/2607.11954#S5.F3)plots this for the best shared\-policy and cascade arms at both label\-space scales\.
ChapterSubsect\.CategorySubcat\.Exact code0\.40\.40\.60\.60\.80\.832% truncatedTraversal depthRecallSFT\-1\+freq\.SFT\-1\+allSFT\-3\+freq\.SFT\-3\+allFigure 3:Micro\-pooled recall by traversal depth for SFT\-1\+\(blue\) and SFT\-3\+\(red\) under*frequent*\(solid\) and*all*\(dashed\)\. At the*all*scale, 32% of cascade test notes suffer L3 truncation\. Depths d5\-d6 omitted \(ragged paths\)\.Recall degrades monotonically with depth, from 0\.80\-0\.83 at the chapter level to 0\.70 \(frequent\) and 0\.49 \(all\) at the billable leaf\. The*frequent*→\\rightarrow*all*gap widens at each depth, confirming that rare\-code discrimination at the bottom of the tree drives the macro\-F1 collapse\. The cascade tracks the shared policy at shallow depths but diverges where L3 truncation \(32% of test notes\) imposes a hard recall ceiling the shared policy never encounters\.
## VIDiscussion
### VI\-AThe Rare\-Code Bottleneck
The macro\-F1 collapse from*frequent*to*all*\(0\.651→0\.2560\.651\\rightarrow 0\.256for the best arm\) is the starkest result in the study\. Micro\-F1 also drops \(0\.709→0\.5270\.709\\rightarrow 0\.527\), but the far larger macro gap shows the degradation is concentrated in the long tail of rare codes rather than spread uniformly\. The flat baselines degrade far more sharply, confirming that the traversal formulation’s hierarchical decomposition partially insulates against the rare\-code problem without solving it\. The per\-depth analysis \(Figure[3](https://arxiv.org/html/2607.11954#S5.F3)\) corroborates this: chapter\-level recall is comparable across scales, but the gap widens at each depth, largest at the billable leaf where rare codes dominate\.
### VI\-BSupervision Is the Bottleneck
Of the three variables we isolate, only additional supervised data reliably improves performance \(Table[IV](https://arxiv.org/html/2607.11954#S5.T4)\)\. Neither RL nor architectural specialization produces a comparable or even consistent gain\. This is worth stating plainly because the natural design impulse points the other way: once a hierarchy is available, it is tempting to assign specialist models to each level and to replace hand\-built trajectories with learned rewards\. Our factorial study tests both assumptions and finds neither justified at this scale\. Our results suggest that, at least for extraction\-bound graph\-constrained coding at this model scale, the complexity does not pay off\. GRPO ties its supervised warm\-start and loses to supervised continuation on the same data; the cascade matches the shared policy but never beats it\. The bottleneck is not how the model is trained or how the traversal is decomposed, but how many gold demonstrations it sees\.
This conclusion is scoped to the regime we test \(a single 4B model, one reward design, 2,000–5,000 notes\), and RL or specialization may yet help at larger scales or with richer reward shaping\. What our study contributes is a controlled baseline: any added machinery should be shown to beat simply adding more supervised data at matched compute\.
### VI\-CConsolidation Scales Where the Cascade Cannot
On headline accuracy the shared traversal policy and the three\-specialist cascade are statistically indistinguishable at both label\-space scales\. The consolidation argument therefore rests not on accuracy but on deployment simplicity \(one adapter, no inter\-stage orchestration\) and a structural scaling advantage\.
The cascade hits a scaling wall at the full code space: its batched L3 prompt forces every billable\-leaf descendant of the selected subsections into one context window, overflowing it for 28\-32% of test notes and making those codes unselectable regardless of model quality\. A larger memory budget would raise this ceiling, but the pressure is structural \- the single\-shot design’s prompt cost grows with label\-space size, whereas the shared policy’s per\-turn fan\-out distributes the same load across turns \(0% truncation\)\.
### VI\-DLimitations
All results are from a single 4B\-parameter model \(Qwen3\-4B\-Instruct\-2507\); RL and architectural specialization may behave differently at larger scales or with other model families\. The GRPO result is based on a single reward design, one hyperparameter configuration, and one warm\-start on the*frequent*label space only; the full\-scale RL cell is argued but not directly measured\. All data comes from MIMIC\-IV, a single US academic medical center, and the gold labels are billed administrative codes known to be incomplete, so measured precision is a lower bound when the model assigns a clinically valid code that was not billed\. The 32,000\-token context cap driving the cascade’s 28\-32% truncation rate is a memory budget constraint \(single 46GB GPU\), not an architectural limit of the model\. Finally, our training budget of 5,000 notes is small relative to the 122,197 available; the data\-scaling claim rests on just two training budgets \(2,000 and 5,000 notes\) and we do not determine where the lift saturates\.
## VIIConclusion
We formulated ICD\-10\-CM code prediction as a single graph\-constrained traversal policy that significantly outperforms flat multi\-label baselines at both label\-space scales \(by\+0\.044\+0\.044micro\-F1 and\+0\.157\+0\.157macro\-F1 over the strongest \(LAAT\) on the full 15,761\-code space\) and conducted a controlled study isolating architecture, training algorithm, and data\. Across both scales, only additional supervised data reliably helped, RL offered no benefit over supervised continuation, and one shared policy matched the three\-specialist cascade while avoiding its context\-window truncation\. Practitioners should therefore deploy one model, not three, and invest in annotation rather than architectural or algorithmic complexity\.
Several directions remain open\. The data\-scaling curve has only two budgets \(2,000 and 5,000 notes\); larger annotation budgets and targeted sampling of rare\-code\-rich notes could determine where the supervised lift saturates\. Sampling\-based decoding \(self\-consistency, best\-of\-nn\) and per\-code threshold tuning may improve rare\-code recall, which the macro\-F1 collapse identifies as the primary bottleneck\. RL deserves re\-examination with richer reward shaping and at larger model scales before a general negative conclusion is warranted\. Finally, validation on a second institution and coding system would test whether these findings generalize beyond a single US academic center\.
## Acknowledgment
This material is based upon work supported by the National Science Foundation under Grant No\. 2431840\.
## References
- \[1\]J\. Mullenbach, S\. Wiegreffe, J\. Duke, J\. Sun, and J\. Eisenstein, “Explainable prediction of medical codes from clinical text,” in*Proc\. NAACL\-HLT*, 2018, pp\. 1101–1111\.
- \[2\]F\. Li and H\. Yu, “ICD coding from clinical text using multi\-filter residual convolutional neural network,” in*Proc\. AAAI*, vol\. 34, no\. 5, 2020, pp\. 8180–8187\.
- \[3\]T\. Vu, D\. Q\. Nguyen, and A\. Nguyen, “A label attention model for ICD coding from clinical text,” in*Proc\. IJCAI*, 2020, pp\. 3335–3341\.
- \[4\]P\. Cao, Y\. Chen, K\. Liu, J\. Zhao, S\. Liu, and W\. Chong, “HyperCore: Hyperbolic and co\-graph representation for automatic ICD coding,” in*Proc\. ACL*, 2020, pp\. 3105–3114\.
- \[5\]L\. Liu, O\. Perez\-Concha, A\. Nguyen, V\. Bennett, and L\. Jorm, “Automated ICD coding using extreme multi\-label long text transformer\-based models,”*Artif\. Intell\. Med\.*, vol\. 144, 2023, Art\. no\. 102662\.
- \[6\]E\. J\. Hu et al\., “LoRA: Low\-rank adaptation of large language models,” in*Proc\. ICLR*, 2022\.
- \[7\]T\. Dettmers, A\. Pagnoni, A\. Holtzman, and L\. Zettlemoyer, “QLoRA: Efficient finetuning of quantized LLMs,” in*Adv\. Neural Inf\. Process\. Syst\.*, 2023\.
- \[8\]A\. E\. W\. Johnson et al\., “MIMIC\-IV, a freely accessible electronic health record dataset,”*Sci\. Data*, vol\. 10, no\. 1, 2023, Art\. no\. 1\.
- \[9\]A\. E\. W\. Johnson, L\. Bulgarelli, T\. Pollard, S\. Horng, L\. A\. Celi, and R\. G\. Mark, “MIMIC\-IV\-Note,” PhysioNet, 2023, doi: 10\.13026/1n74\-ne17\.
- \[10\]J\. Edin et al\., “Automated medical coding on MIMIC\-III and MIMIC\-IV: A critical review and replicability study,” arXiv:2304\.10909, 2023\.
- \[11\]A\. Yang et al\., “Qwen3 technical report,” arXiv:2505\.09388, 2025\.
- \[12\]E\. Alsentzer et al\., “Publicly available clinical BERT embeddings,” in*Proc\. 2nd Clin\. NLP Workshop*, 2019, pp\. 72–78\.
- \[13\]J\. Lee et al\., “BioBERT: A pre\-trained biomedical language representation model for biomedical text mining,”*Bioinformatics*, vol\. 36, no\. 4, pp\. 1234–1240, 2020\.
- \[14\]C\.\-W\. Huang, S\.\-C\. Tsai, and Y\.\-N\. Chen, “PLM\-ICD: Automatic ICD coding with pretrained language models,” in*Proc\. 4th Clin\. NLP Workshop*, 2022, pp\. 10–20\.
- \[15\]Z\. Yuan, C\. Tan, and S\. Huang, “Code synonyms do matter: Multiple synonyms matching network for automatic ICD coding,” in*Proc\. ACL \(Short Papers\)*, 2022, pp\. 808–814\.
- \[16\]W\. Ren, R\. Zeng, T\. Wu, T\. Zhu, and R\. G\. Krishnan, “HiCu: Leveraging hierarchy for curriculum learning in automated ICD coding,” in*Proc\. MLHC*, 2022\.
- \[17\]X\. Xie, Y\. Xiong, P\. S\. Yu, and Y\. Zhu, “EHR coding with multi\-scale feature attention and structured knowledge graph propagation,” in*Proc\. CIKM*, 2019, pp\. 649–658\.
- \[18\]Z\. Yang, S\. Wang, B\. P\. S\. Rawat, A\. Mitra, and H\. Yu, “Knowledge injected prompt based fine\-tuning for multi\-label few\-shot ICD coding,” in*Findings EMNLP*, 2022, pp\. 1767–1781\.
- \[19\]B\. T\. Willard and R\. Louf, “Efficient guided generation for large language models,” arXiv:2307\.09702, 2023\.
- \[20\]T\.\-T\. Nguyen et al\., “MIMIC\-IV\-ICD: A new benchmark for eXtreme multilabel classification,” arXiv:2304\.13998, 2023\.
- \[21\]A\. Soroush et al\., “Large language models are poor medical coders — benchmarking of medical code querying,”*NEJM AI*, vol\. 1, no\. 5, 2024\.
- \[22\]J\. S\. Boyle, A\. Kascenas, P\. Lok, M\. Liakata, and A\. Q\. O’Neil, “Automated clinical coding using off\-the\-shelf large language models,” in*Proc\. NeurIPS Workshop Deep Generative Models for Health \(DGM4H\)*, 2023, arXiv:2310\.06552\.
## Supplementary Materials
## Appendix ASystem Prompts
### A\-AShared Agentic Policy \(SFT\-1 / SFT\-1\+/ GRPO\-1\)
The following system prompt is used for the shared agentic policy across all training and evaluation conditions\. The model additionally receives the tool schema forexpand\_level\(\)andselect\(children\)as function definitions in Qwen3’s native tool\-calling format\.
AssignICD\-10\-CMcodesbytraversingtheontology
level\-by\-levelusingtools\.
Rules:
\-ONLYselectIDsexplicitlylistedinthemost
recentexpand\_level\(\)output\.
\-DoNOTguesscodesoruseexternalknowledgeof
theICDhierarchy\.
\-Onlythemostspecific\(leaf\)codesarebillable
\-\-drillasdeepasthenotesupports\.
Process:
1\.Callexpand\_level\(\)\-\-showsallpendingnodes
andtheirchildren,groupedbyparent\.Leafcodes
\(nosub\-codesexist\)areauto\-recordedimmediately\.
2\.Callselect\(children=\[\.\.\.\]\)withtheIDsfromANY
groupthatthenotesupports\.Passanemptylist
ifnoneofthecurrentoptionsapply\.
3\.Repeatuntilexpand\_level\(\)returnsDONE\.
4\.Outputexactly:ANSWER:code1,code2,\.\.\.
Thinkin1SHORTsentence,thenimmediatelycall
thetool\.
#### A\-A1Tool Schema
The model receives the following tool definitions in Qwen3’s native tool\-calling format\.
\{
"type":"function",
"function":\{
"name":"expand\_level",
"description":"Showallpendingnodesandtheir
childrengroupedbyparent\.Leafcodes\(no
sub\-codes\)areauto\-recordedasfinalanswers\.
ReturnsDONEwhentheagendaisempty\.",
"parameters":\{
"type":"object",
"properties":\{\},
"required":\[\]
\}
\}
\}
\{
"type":"function",
"function":\{
"name":"select",
"description":"QueuechildnodeIDsforfurther
expansion\.Onlyleafcodescanbecomefinal
answers\(auto\-recordedbyexpand\_level\)\.
Non\-leafcodesmustbedrilledinto\-\-they
cannotberecordeddirectly\.Passchildren=\[\]
toskipallcurrentoptionsifnoneapplyto
thenote\.",
"parameters":\{
"type":"object",
"properties":\{
"children":\{
"type":"array",
"items":\{"type":"string"\},
"description":"ChildnodeIDstodrill
into\.Passanemptylisttoskipall
optionsatthislevel\."
\}
\},
"required":\["children"\]
\}
\}
\}
#### A\-A2Example Trajectory
The user message contains only the clinical note\. The following illustrates the multi\-turn tool\-call interaction for a note with gold codes E785, I10, and J449\. A complete traversal typically spans 5\-7 depth levels, each requiring an expand\-and\-select round trip \(approximately 12\-14 model turns total\)\.
\[assistant\]
<tool\_call\>\{"name":"expand\_level",
"arguments":\{\}\}</tool\_call\>
\[toolresponse\]
UnderROOT:Ch3\(Diseasesoftheblood\.\.\.\),
Ch4\(Endocrine\.\.\.\(E00\-E89\)\),
Ch9\(Circulatory\.\.\.\(I00\-I99\)\),
Ch10\(Respiratory\.\.\.\(J00\-J99\)\),\.\.\.Ch21\(\.\.\.\)
Callselect\(children=\[\.\.\.\]\)withIDsfromANY
groupabove\.
\[assistant\]
<tool\_call\>\{"name":"select",
"arguments":\{"children":
\["Ch4","Ch9","Ch10"\]\}\}</tool\_call\>
\[toolresponse\]
Selected\[’Ch4’,’Ch9’,’Ch10’\]\.
Callexpand\_level\(\)todrilldeeper\.
\.\.\.expand\_level/selectrepeatsateachdepth,
drillingeachbranchtoitsleaf\.\.\.
\(Ch4\-\>E70\-E88\-\>E785auto\-recorded;
Ch9\-\>I10\-I1A\-\>I10auto\-recorded;
Ch10\-\>J40\-J4A\-\>J44\-\>J449auto\-recorded\)
\[assistant\]
ANSWER:I10,E785,J449
### A\-BCascade Baseline \(SFT\-3 / SFT\-3\+\)
All three cascade specialists share the following system prompt\. Each specialist receives a single\-turn user message containing the clinical note and a candidate list specific to its ontology level\. There is no tool\-call protocol; the model responds with a singleSELECTED:line\.
YouareanexpertICD\-10\-CMcoder\.Youassigncodes
bynavigatingtheICD\-10\-CMontologyonelevelata
time\.Givenaclinicalnoteandalistofcandidate
nodesatthecurrentlevel\(groupedbytheirparent\),
selectEVERYcandidatethenotesupports\.
Rules:
\-ONLYchooseIDsfromthecandidatelistbelow\.
DoNOTinventcodesoruseoutsideknowledge
ofthehierarchy\.
\-UsetheIDsverbatim\.
\-ReplywithONElineexactly:
SELECTED:id1,id2,\.\.\.
Ifnoneofthecandidatesapplytothenote,reply:
SELECTED:NONE
#### A\-B1L1: Chapter Specialist
User message format:
Clinicalnote:
<medical\_note\>
Position:topoftheICD\-10\-CMhierarchy\(chapters\)\.
Candidates\(chooseallthatapply\):
\-Ch3\(Diseasesoftheblood\.\.\.\(D50\-D89\)\)
\-Ch4\(Endocrine,nutritionalandmetabolic
diseases\(E00\-E89\)\)
\-Ch9\(Diseasesofthecirculatorysystem
\(I00\-I99\)\)
\.\.\.
Expectedresponse:
SELECTED:Ch4,Ch9,Ch10
#### A\-B2L2: Subsection Specialist
User message format:
Clinicalnote:
<medical\_note\>
Selectedchapters:Ch4\(Endocrine\.\.\.\(E00\-E89\)\),
Ch9\(Circulatory\.\.\.\(I00\-I99\)\),
Ch10\(Respiratory\.\.\.\(J00\-J99\)\)
Candidates\(chooseallthatapply\):
UnderCh4\(Endocrine\.\.\.\):
\-E00\-E07\(Disordersofthyroidgland\)
\-E08\-E13\(Diabetesmellitus\)
\-E70\-E88\(Metabolicdisorders\)
UnderCh9\(Diseasesofthecirculatorysystem\):
\-I10\-I1A\(Hypertensivediseases\)
\-I20\-I25\(Ischemicheartdiseases\)
\.\.\.
Expectedresponse:
SELECTED:E70\-E88,I10\-I1A,J40\-J4A
#### A\-B3L3: Code Specialist
User message format:
Clinicalnote:
<medical\_note\>
Selectedsubsections:E70\-E88\(Metabolicdisorders\),
I10\-I1A\(Hypertensivediseases\),
J40\-J4A\(Chroniclowerrespiratorydiseases\)
Candidates\(chooseallthatapply\):
UnderE70\-E88\(Metabolicdisorders\):
\-E785\(Hyperlipidemia,unspecified\)
\-E871\(Hypo\-osmolalityandhyponatremia\)
\-E875\(Hyperkalemia\)
UnderI10\-I1A\(Hypertensivediseases\):
\-I10\(Essential\(primary\)hypertension\)
\-I110\(Hypertensiveheartdiseasew/HF\)
\.\.\.
Expectedresponse:
SELECTED:E785,I10,J449
Note: under thealllabel space, the L3 candidate list can reach tens of thousands of tokens \(average≈\\approx26k, maximum≈\\approx74k\)\. The sequence length is capped at 32,000 tokens to fit within the memory budget of a single 46GB GPU; the model itself supports up to 262k positions\. See Section IV\-D of the main text for truncation handling\.
## Appendix BRollout Driver Implementation
### B\-AOverview
The traversal environment requires a multi\-turn loop: the model generates a tool call, the environment executes it and returns a result, and the model generates again conditioned on the full history\. The training framework \(TRL 0\.22\.2 with Unsloth\) issues a single generation call per rollout and provides no native tool\-calling loop\. We implement the multi\-turn cycle in two drivers that share identical per\-rollout stepping logic\.
### B\-BSingle\-Note Driver
Used for tokenizing gold trajectories, GRPO training \(per\-prompt sequential rollouts\), and single\-note debugging\. Each turn, the driver:
1. 1\.Calls the model \(or replays a gold action during SFT data construction\)\.
2. 2\.Parses the output for a<tool\_call\>block using a regex pattern matching Qwen3’s format:<tool\_call\>\{\.\.\.\}</tool\_call\>\.
3. 3\.Executes the parsed tool call against the environment \(expand\_level\(\)orselect\(children=\[\.\.\.\]\)\)\.
4. 4\.Injects the tool result as the next turn, wrapped in Qwen3’s<tool\_response\>format\.
Rather than re\-rendering the full message list through the chat template at each turn \(which would not round\-trip correctly for Qwen3’s<think\>and<tool\_call\>spans\), the driver maintains a running concatenation of raw token IDs, appending each new segment directly\. This produces the trainable/masked split by construction:
- •mask = 1 \(trainable\):all tokens the model generated \(tool\-call content, reasoning,ANSWERline\)\.
- •mask = 0 \(masked\):all tokens the driver injected \(tool\-result content, role headers such as<\|im\_start\|\>userand<\|im\_start\|\>assistant,<tool\_response\>wrappers, and any<think\>block opening tokens auto\-inserted by the generation prompt\)\.
### B\-CBatched Driver
Used for evaluation only \(not for GRPO training, which uses the sequential single\-note driver\)\. The per\-rollout state machine is a line\-for\-line port of the single\-note driver, verified byte\-identical on a batch of one against a deterministic mock\. The difference is operational: all active rollouts are pooled into a singlellm\.generate\(\)call per wave\. At each turn the driver:
1. 1\.Collects the running token context and remaining budget from every rollout that has not yet finished\.
2. 2\.Issues one batched generation call with per\-requestSamplingParams\(each rollout carries its ownmax\_tokensreflecting its remaining budget\)\.
3. 3\.Steps each rollout individually: parsing its tool call, executing against its own environment instance, injecting its tool result\.
4. 4\.Retires rollouts that emitANSWERor exhaust their token budget\.
The total number of generation calls equals the maximum number of turns across all rollouts \(capped at 16\), not the sum of per\-rollout turns\. This is what makes multi\-turn evaluation tractable at 1,000 test notes\.
### B\-DAgentic Guardrails
Two mechanisms prevent degenerate rollout behavior:
Forced first expansion\.When enabled, the driver injects a maskedexpand\_level\(\)call as the first assistant turn before the model generates anything\. This seeds the loop with the chapter\-level candidate list so the model cannot skip the traversal and emit a memorizedANSWERfrom the system prompt alone — a failure mode observed in early GRPO runs where the model learned to terminate immediately to avoid negative reward\.
Early\-answer rejection\.When enabled, if the model emitsANSWERbefore the environment has reportedDONE\(i\.e\., the agenda is not yet empty\), the driver injects a masked user\-turn nudge instructing the model to continue traversing:
Youhavenotfinishedtraversingtheontology\.
DoNOToutputANSWERyet\.Continuebycalling
select\(children=\[\.\.\.\]\)orexpand\_level\(\)\.Youmay
onlyoutputANSWERafterexpand\_level\(\)reportsDONE\.
The rollout remains active for the next wave\. This prevents the model from terminating mid\-traversal to collect a partial reward during GRPO, or from short\-circuiting the traversal during evaluation\.
### B\-EQwen3 Framing Detection
The injected text between turns \(closing the assistant turn, wrapping the tool result, re\-opening the next assistant turn\) must match the chat template exactly\. Qwen3\-Instruct and Qwen3\-Thinking variants differ in whether the generation prompt auto\-opens a<think\>block\. At initialization, the driver probes the tokenizer with a minimal message list and inspects whether the rendered generation prompt ends with<think\>\. It then builds the inter\-turn framing template accordingly:
\#Templateforinjectingtoolresultsbetweenturns:
<\|im\_start\|\>user
<tool\_response\>
\{result\}
</tool\_response\><\|im\_end\|\>
<\|im\_start\|\>assistant
<think\>\#onlyifmodelauto\-opens<think\>
This auto\-detection ensures the driver works with both Instruct and Thinking variants without manual configuration\.
## Appendix CTraining Hyperparameters
All arms use the same base model \(Qwen3\-4B\-Instruct\-2507\) with 4\-bit quantization and QLoRA adapters\. The data\-splitting seed is 42; the training seed is 3407\. Table[V](https://arxiv.org/html/2607.11954#A3.T5)reports the SFT configuration shared across SFT\-1, SFT\-1\+, SFT\-3, and SFT\-3\+\. Table[VI](https://arxiv.org/html/2607.11954#A3.T6)reports the GRPO\-1 configuration\.
TABLE V:QLoRA and SFT training configuration \(shared across all SFT arms\)\.TABLE VI:GRPO\-1 training configuration\. LoRA and base model settings are identical to Table[V](https://arxiv.org/html/2607.11954#A3.T5); only the training recipe differs\.The SFT and GRPO recipes are deliberately not matched on learning rate or epoch count: an SFT epoch \(teacher\-forced pass over gold trajectories\) and a GRPO epoch \(rollout, score, and policy update withGGgenerations per prompt\) are different objectives with different optima\. Fairness is ensured by the data\-matched splits and the SFT\-1\+compute control, not by identical hyperparameters across algorithms\.
Theallvariant uses larger sequence and completion budgets because gold traversals over the full 15,761\-leaf ontology produce 3–4×\\timeslonger token sequences\. Thefrequentbudgets are non\-binding for that variant\. The cascade’s L3 sequence cap of 32,000 tokens is a memory budget constraint for the single 46GB GPU used in this study; the base model \(Qwen3\-4B\-Instruct\-2507\) natively supports up to 262,144 positions\. The completion length budget does not apply to the cascade, whose specialists produce only a short single\-turnSELECTED:response\.
## Appendix DFlat Baseline Configurations
CAML, LAAT, and PLM\-ICD are retrained on the same data splits as the agentic and cascade arms\. All three use the pooled 5,000\-note training set \(SFT 2k \+ continuation 3k\), the 500\-note validation set for threshold tuning and early stopping, and the 1,000\-note test set\. A global decision threshold is selected by grid search over\[0\.05,0\.925\]\[0\.05,0\.925\]in steps of 0\.025, maximizing validation micro\-F1; if no code exceeds the threshold for a given note, the argmax is taken\. All experiments use a single NVIDIA L40S \(46 GB\) GPU and random seed 3407\.
### D\-ACAML \(Mullenbach et al\. 2018\)
Reimplementation in PyTorch 2\.12\.
TABLE VII:CAML training configuration\.
### D\-BLAAT \(Vu et al\. 2020\)
Reimplementation in PyTorch 2\.12\.
TABLE VIII:LAAT training configuration\.
### D\-CPLM\-ICD \(Huang et al\. 2022\)
Reimplementation in PyTorch 2\.12 and Transformers 5\.x\.
TABLE IX:PLM\-ICD training configuration\.Similar Articles
Hierarchical Modeling of ICD Codes in EHR Foundation Models
This paper investigates explicit encoding of ICD-10-CM hierarchy in EHR foundation models, using hierarchical token augmentation and graph-based code representations. Experiments on MIMIC-IV and eICU show improvements over flat code representations for in-domain and cross-dataset prediction tasks.
Trajectory-Aware Clinical Risk Prediction via Severity-Grounded Knowledge Graphs and Retrieval-Augmented Generation
Proposes TRACER, a framework that integrates severity-grounded knowledge graphs and retrieval-augmented generation for trajectory-aware clinical risk prediction, achieving large gains in mortality and readmission prediction on MIMIC-III and MIMIC-IV datasets.
Primary ICD Category Prediction using LLM-based Probing
This paper presents a method that uses frozen medical large language model (LLM) representations as a shared embedding space to predict primary ICD diagnosis categories from both structured and unstructured electronic health record data, achieving improved accuracy over baseline methods on MIMIC-IV and showing transferability to MIMIC-III.
Multimodal Routing for Interpretable, Robust, and Auditable Clinical Prediction
This paper proposes an explicit multimodal routing framework for clinical prediction using EHR data, enabling interpretable, robust, and auditable reasoning across structured variables, clinical notes, and chest X-rays via discrete unimodal, bimodal, and trimodal routes with inference-time route masking for missing modality simulation.
Can Post-Training Turn LLMs into Good Medical Coders? An Empirical Study of Generative ICD Coding
This empirical study investigates whether post-training (supervised fine-tuning and reinforcement learning) can improve LLMs' performance on automated ICD coding, introducing a diagnostic curriculum called PHI that extends GRPO to refine missed-code cases. Results show that prompting-only evaluation underestimates LLM potential, with SFT providing the main capability jump and RL further improving performance.