💬 NLP

From Documents to Spans: Code-Centric Learning for LLM-based ICD Coding

This paper proposes "Code-Centric Learning," a training framework that improves the accuracy, interpretability, and cost-efficiency of LLM-based ICD coding by shifting supervision from full clinical documents to scalable evidence spans, enabling smaller models to outperform larger baselines and generalize to unseen codes.

Original authors: Xu Zhang, Wenxin Ma, Chenxu Wu, Rongsheng Wang, Kun Zhang, S. Kevin Zhou

Published 2026-03-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xu Zhang, Wenxin Ma, Chenxu Wu, Rongsheng Wang, Kun Zhang, S. Kevin Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a medical librarian whose job is to read thousands of patient stories (clinical notes) and assign them specific "library codes" (ICD codes) so insurance companies know how to pay and researchers can study diseases.

There are about 70,000 different codes in the library, but the stories you read are long, messy, and full of medical jargon.

The Old Way: The "Overwhelmed Intern"

Traditionally, computers tried to learn this job by reading the entire patient story every single time.

The Problem: Imagine trying to teach a student to recognize every specific type of tree in a forest by only showing them photos of the entire forest. It's slow, expensive, and the student gets confused.
The Result: The computer was slow to train, it only learned the few trees it saw often (missing the rare ones), and when it guessed a code, it couldn't explain why it picked it. It was like a black box.

The New Way: "Code-Centric Learning" (The Paper's Solution)

The authors of this paper propose a smarter way to train the computer, which they call Code-Centric Learning.

Think of it like training a detective instead of a general reader. Instead of reading the whole book to find the answer, the detective learns to spot tiny, specific clues (evidence spans) that point directly to a specific code.

Here is how they did it, using three simple steps:

1. The "Flashcard" Strategy (Mixed Training)

Instead of forcing the computer to read 100-page patient files, they broke the learning down into two types of flashcards:

The Full Story Cards: A few real patient stories where the computer learns to find the clues and pick the code. This teaches it how to handle the whole picture.
The Clue Cards: Millions of tiny snippets. One side says "Broken Leg," the other side says "Patient fell off a bike and can't walk." The computer just learns: "Oh, 'can't walk' + 'bike' = Broken Leg code."
Why it works: It's like learning vocabulary by studying individual words and sentences rather than trying to memorize whole novels. It's 80% faster to train because the computer doesn't have to process huge files every time.

2. The "Imagination Engine" (Data Expansion)

The biggest problem with medical coding is that there are 70,000 codes, but most datasets only have examples for about 10% of them. What happens when the computer sees a code it has never seen before?

The Old Way: The computer would just guess or fail.
The New Way: The authors used a powerful AI to invent training examples for the missing codes.
- Analogy: Imagine you want to teach a child what a "Pangolin" looks like, but you've never seen one. You show them a picture of a "Sloth" (a similar animal) and say, "A Pangolin is like a Sloth, but with armor scales." The AI does this for medical codes. It looks at a known code, finds a similar unknown one, and synthesizes a realistic patient story for it.
- This allows the computer to learn 100% of the codes, even the ones it has never seen in real life.

3. The "Highlighter" (Interpretability)

In the old days, the computer would just spit out a code: "Diagnosis: Diabetes." You had no idea if it was right or why.

The New Way: The computer is now forced to highlight the evidence first.
- Output: "Diagnosis: Diabetes. Evidence: 'Patient has high blood sugar and frequent thirst.'"
Why it matters: This is like a student showing their work on a math test. If the highlight is wrong, a human doctor can fix it immediately. This makes the AI a helpful partner, not a mysterious oracle.

The Big Win

By switching from "reading the whole book" to "studying the clues," the researchers achieved three amazing things:

Speed: They trained the model in 40 hours instead of 210 hours.
Smarts: A small, cheap computer model trained this way performed just as well as massive, expensive, proprietary models.
Trust: The model can now point to the exact sentence in the patient's note that justified its decision, making it safe for real hospitals to use.

In a nutshell: They stopped trying to teach the AI to read the whole library and started teaching it to be a master detective who knows exactly which clues to look for, even for cases it has never seen before.

1. Problem Statement

Automatic International Classification of Diseases (ICD) coding is a critical task in healthcare for insurance reimbursement and data analysis, but it faces significant challenges when applying Large Language Models (LLMs):

Limited Code Coverage: Existing public datasets (e.g., MIMIC-IV) cover only a small fraction (approx. 10%) of the ~70,000 ICD-10-CM codes. This restricts models from generalizing to "unseen" codes not present in the training data.
Lack of Interpretability: Standard fine-tuning on document-level data predicts codes directly without explicit supporting evidence. Most public datasets lack human-annotated evidence spans, making it difficult for clinicians to verify or correct the model's reasoning.
Training Inefficiency: Fine-tuning LLMs on long clinical documents (often thousands of tokens) is computationally expensive due to the quadratic complexity of attention mechanisms, making it impractical for many institutions.

2. Methodology: Code-Centric Learning (CCL)

The authors propose Code-Centric Learning (CCL), a training framework that shifts supervision from full clinical documents to scalable, short evidence spans. The core intuition is that if an LLM learns to assign the correct ICD code to a specific evidence span, it implicitly learns to recognize that evidence within a full document.

The framework consists of two main strategies:

A. Mixed Training Strategy

Instead of training solely on full documents, CCL utilizes a hybrid dataset:

Document-Level Data: A small set of documents annotated with both ICD codes and their corresponding evidence spans (e.g., from the MDACE dataset). This teaches the model to aggregate evidence from a full context and assign multiple codes.
Span-Level Data: A large collection of (evidence span, single ICD code) pairs. These are short, focused inputs that allow for efficient training and inject specific code knowledge.

Training Objective: The model is fine-tuned to predict codes given either a full document (with instructions to extract evidence first) or a specific evidence span.

B. Code-Centric Data Expansion

To address the "unseen code" problem and ensure full coverage of the ICD space, the authors construct a multi-tier knowledge base ( $D = D_{gold} \cup D_{silver} \cup D_{syn}$ ):

Gold Pairs: Extracted from authoritative official resources (ICD Alphabetic Index and Tabular List), pairing index terms with their default codes.
Silver Pairs: Mined from noisy public datasets (like MIMIC-IV) using a two-stage LLM pipeline. First, an LLM extracts evidence spans for codes in the documents; second, it consolidates these into representative evidence expressions.
Synthetic Pairs: For codes not found in official or public data, the authors use an LLM (GPT-5.1) to synthesize plausible evidence spans. This is done by retrieving the nearest neighbor code in the ICD hierarchy and inferring evidence for the target code based on the neighbor's knowledge.

3. Key Contributions

Novel Training Paradigm: Proposes shifting from document-level to span-level supervision, demonstrating that span-level classification yields transferable gains for document-level evidence extraction and coding.
Comprehensive Data Expansion: Introduces a method to achieve near 100% ICD code coverage by synthesizing evidence for unseen codes using official guidelines and LLM inference, overcoming the scarcity of annotated data.
Efficiency and Interpretability: The approach drastically reduces training costs by focusing on short spans and explicitly generates evidence spans, enabling human-in-the-loop verification.
Small Model Performance: Demonstrates that small-scale open-source LLMs (e.g., Llama-3.1-8B) trained with CCL can outperform much larger proprietary models (e.g., GPT-4, Qwen-235B) on standard benchmarks.

4. Experimental Results

The method was evaluated on MDACE (expert-annotated) and ACI-Bench (out-of-distribution) datasets.

Performance:
- On MDACE, CCL with Llama-3.1-8B achieved a Micro-F1 of 59.3% and Macro-F1 of 35.2%, outperforming the standard SFT baseline (57.4% Micro-F1) and significantly beating Chain-of-Thought (CoT) prompting (26.5% Micro-F1).
- On ACI-Bench (unseen distribution), CCL achieved a Micro-F1 of 65.8%, surpassing the SFT baseline (63.4%) and the large-scale CLH agent (47.5%).
- Unseen Codes: The method showed a massive improvement on codes not seen in the document-level training data, increasing accuracy from 8.8% to 50.0%.
Efficiency:
- Training time was reduced by ~80% (from 210 hours to 40 hours for Llama-3.1-8B) by using only 200 annotated documents and 150k spans instead of 100k full documents.
Interpretability:
- The model achieved an Evidence Recall of 74.0% on MDACE, compared to 0% for code-only methods.
- When human-annotated evidence was provided as input (simulating human-AI collaboration), performance jumped to a Micro-F1 of 78.0%, proving the system's controllability.

5. Significance

This paper fundamentally rethinks the ICD coding pipeline for LLMs. By decoupling the "evidence extraction" and "code assignment" tasks and training them on scalable span-level data, the authors solve the three major bottlenecks of current LLM-based coding:

Scalability: It makes full ICD code coverage feasible without requiring massive, manually annotated datasets for every code.
Accessibility: It allows smaller, locally deployable models to achieve state-of-the-art performance, addressing privacy and compute constraints in clinical settings.
Trust: By generating explicit evidence spans, it bridges the gap between AI automation and clinical workflow, allowing human coders to review and correct specific evidence rather than just the final code.

The work suggests a new direction for medical NLP where "learning from spans" is more effective than "learning from documents" for complex, evidence-driven classification tasks.