PaperlayerView exact revision on arXiv

Source-linked AI summary

CoLa-ICD: A Knowledge-Enhanced Framework for Long-Tail Automated Medical Coding

Yihang Cheng, Veronica Liesaputra, Andrew Trotman

arXiv:2608.30234v1cs.AI

TL;DR

Automated medical coding must assign ICD codes despite long notes, imbalanced code frequencies, and varied clinical terminology, with rare codes especially difficult to learn and distinguish. CoLa-ICD enriches label representations, models code relations, and aligns clinical evidence with label semantics; experiments report improved overall and long-tail coding performance on ICD-9 and ICD-10 datasets.

  • Problem

    Rare ICD codes have limited training instances and can be confused with semantically similar labels, while long notes and terminology mismatch complicate automated coding.

  • Method

    CoLa-ICD combines LLM-generated label terms, a co-occurrence graph neural network, multi-head term attention, and label-aware contrastive learning.

  • Results

    CoLa-ICD improves overall coding performance and achieves state-of-the-art results with robust long-tail gains on ICD-9 and ICD-10 datasets.

  • Takeaways & Limitations

    Knowledge-enhanced evidence–label alignment is a promising direction for rare-code prediction in automated medical coding.

  • Takeaways & Limitations

    Evaluation is limited to MIMIC datasets from a single medical center, primarily covering ICU discharge summaries, so performance may vary across settings.

Abstract

from arXiv · show

Automatic medical coding assigns ICD codes to clinical notes, but it remains challenging due to long documents, imbalanced label distributions, and diverse terms. These challenges are especially severe for rare codes, which have limited training instances and are easily confused with semantically similar labels. We introduce CoLa-ICD, a knowledge-enhanced framework for long-tail prediction. CoLa-ICD enriches ICD labels with external terms, models dependencies among related codes, and learns stronger alignment between label semantics and clinical evidence for long-tail prediction. Experiments show that CoLa-ICD improves long-tail prediction with larger gains in larger and sparser label spaces and achieves state-of-the-art performance in AUC, F1, and P@k. Our code is available at https://github.com/youwillbethebest/Cola-ICD.

1 Introduction

Automated medical coding is a difficult multi-label task because ICD codes are long-tailed, clinical notes are lengthy, and clinical wording may differ from official code descriptions. CoLa-ICD addresses these challenges by enriching labels, modeling code relations, and aligning evidence with label semantics, improving long-tail performance.

  • More than half of ICD codes appear fewer than 10 times, making rare codes difficult to learn and important for capturing precise clinical details.
  • Clinical notes contain 1,500–2,000 words, while each ICD code depends on dense, short evidence such as diagnoses, procedures, medications, or abnormal findings.
  • Clinical concepts may use wording such as “former smoker” or “quit tobacco” instead of the official description “History of tobacco use,” complicating rare-code prediction.
  • CoLa-ICD combines LLM-generated terms, co-occurrence relations, and label-aware contrastive learning to improve evidence–label alignment for long-tail codes.
  • The framework enriches label representations with LLM-generated terms, integrates a co-occurrence GNN, and uses multi-head term attention to retrieve supporting evidence.
  • CoLa-ICD achieves state-of-the-art performance on ICD-9 and ICD-10 benchmarks, with robust improvements on long-tail codes.

2 Related Work

Prior automated coding methods use neural encoders, pretrained language models, long-document strategies, external knowledge, and contrastive learning. CoLa-ICD targets gaps in jointly modeling semantic and structural sparsity while grounding predictions in short evidence spans.

  • Early automated coding methods use convolutional or recurrent neural networks, establishing strong baselines for assigning ICD codes to discharge summaries.
  • Label-wise and cross-attention methods improve interactions between individual codes and clinical text by measuring label-specific token relevance.
  • Pretrained biomedical language models improve clinical representations, while segment pooling and chunking strategies address the token limits of long notes.
  • Knowledge-injection methods enrich labels through code structure, synonyms, textual definitions, generated terms, or internal label correlations.
  • Existing knowledge-injection methods often model either semantic sparsity or structural sparsity, but seldom both within the same evidence–label matching process.
  • Most contrastive methods treat text and labels globally, overlooking short, specific evidence needed for accurate ICD coding.

3 Problem Formulation

Automated medical coding is formulated as multi-label classification: given a clinical document, the model predicts all applicable ICD codes from a large candidate set. The task is difficult because the label space is large and code frequencies are highly imbalanced.

  • Given a clinical document X, automated coding predicts all relevant ICD codes from a candidate set C containing N possible codes.
  • Each document is represented by a binary label vector Y, where each entry indicates whether the corresponding ICD code applies.
  • The task combines a large label space with a highly imbalanced, long-tailed distribution of code frequencies.

4 Methods

CoLa-ICD combines long-document encoding, enriched ICD label knowledge, code co-occurrence modeling, label-wise evidence extraction, and contrastive alignment for automated coding. The resulting framework is designed to improve semantic matching, particularly for sparse and long-tail codes.

  • Long-document encoding: CoLa-ICD preserves all tokens in long discharge summaries by segmenting them before encoding and concatenating segment-level hidden states.This enables attention to evidence scattered across long documents without truncation.
  • Label knowledge injection: Each ICD code is enriched with UMLS descriptions and synonyms plus LLM-generated abbreviations, extended definitions, and common clinical phrases.The terms form a fixed-length list Ti for each code.
  • Co-occurrence graph modeling: The co-occurrence GNN transforms code co-occurrence statistics into PPMI-weighted dependencies and updates flattened code-term embeddings.PPMI down-weights frequency-driven associations, while the GNN produces knowledge-enhanced label representations for downstream attention.
  • Multi-head term attention: Multi-head term attention uses each code’s terms as distinct queries over contextual token representations, then concatenates the resulting evidence across heads.The final evidence representation fi is passed through an MLP and sigmoid to produce a prediction score for code ci.
  • Label-aware contrastive learning: Label-Aware Contrastive Learning aligns attention-extracted evidence fi with its knowledge-enhanced label embedding pi while separating it from semantically similar negatives.Hard-negative mining selects highly similar incorrect codes, encouraging discrimination among semantically similar labels.
  • Training objective: The model combines binary cross-entropy with a weighted contrastive loss for multi-label ICD prediction.The contrastive-loss weight is set to λ = 0.05 using validation performance to balance the two loss terms.

5 Experiments

CoLa-ICD is evaluated on MIMIC-III ICD-9 and MIMIC-IV ICD-10 using long-tail, overall, and standard classification and ranking metrics. It achieves the strongest reported performance across frequency buckets and overall evaluation, including in the larger, sparser ICD-10 label space.

  • Performance on Long-tailed Codes: CoLa-ICD achieves the highest Micro-F1 across all five training-frequency buckets on both MIMIC-III ICD-9 and MIMIC-IV ICD-10.The evaluation covers both the standard ICD-9 benchmark and the larger ICD-10 label space.
  • Performance on Long-tailed Codes: 0.077 improvement over GKI-ICD in the MIMIC-IV ICD-10 1–10 frequency bucket demonstrates the largest reported long-tail gain.The corresponding improvement in the 11–50 bucket is 0.056.
  • Performance on Long-tailed Codes: 0.194 Micro-F1 in the MIMIC-III ICD-9 1–10 bucket improves on GKI-ICD’s 0.132.This provides a parallel rare-code result on ICD-9.
  • Overall Performance: CoLa-ICD obtains the highest Macro-F1 (0.135), Micro-F1 (0.621), P@8 (0.781), and P@15 (0.633) on ICD-9.These results outperform the strongest knowledge-enhanced baseline.
  • Overall Performance: On MIMIC-IV ICD-10, CoLa-ICD achieves the best Macro-F1 (0.108), Micro-F1 (0.600), and P@8 (0.723).Its improvements over the strongest baseline are 0.039, 0.021, and 0.021 absolute points, respectively.

6 Analysis & Discussion

Ablations, span-removal tests, and a case study examine which CoLa-ICD components support rare-code prediction and evidence selection. The results associate knowledge injection, label-aware contrastive learning, and code relations with stronger rare-code performance and more prediction-relevant evidence.

  • Ablation Study: Removing any CoLa-ICD component reduces Macro-F1, Micro-F1, P@8, and Rare Micro-F1 on both datasets.The full model achieves the best F1 and rare-code performance across the ICD-9 and ICD-10 ablations.
  • Ablation Study: 0.049 and 0.049 Rare Micro-F1 drops follow removal of knowledge injection on ICD-9 and ICD-10, respectively.Rare Micro-F1 decreases from 0.194 to 0.153 on ICD-9 and from 0.204 to 0.155 on ICD-10.
  • Ablation Study: 0.027 and 0.031 Rare Micro-F1 drops follow removal of LACL on ICD-9 and ICD-10, respectively.The result supports LACL’s role in aligning evidence with labels and distinguishing clinically similar codes.
  • Ablation Study: 0.020 and 0.023 Rare Micro-F1 drops follow removal of the co-occurrence GNN on ICD-9 and ICD-10, respectively.The ablation attributes the component’s contribution to structural information among correlated ICD codes.
  • Evidence Utility of Attention-Selected Spans: 0.113 target-code probability drop occurs when CoLa-ICD’s top-attended rare-code span is removed, versus 0.037 for a random rare-code span.Overall, the corresponding Top-1 and Random drops are 0.097 and 0.037.
  • Evidence Utility of Attention-Selected Spans: The full model’s Top-1–Random gap is 0.074, exceeding the gaps after removing KI, LACL, or the co-occurrence GNN.The component removals reduce the overall gap to 0.048, 0.029, and 0.027, respectively.

7 Conclusion

CoLa-ICD combines enriched ICD labels, code-relation modeling, and evidence–label alignment for long-tail automated medical coding. Experiments on ICD-9 and ICD-10 show improved overall coding performance and support knowledge-enhanced evidence–label alignment as a direction for rare-code prediction.

  • Conclusion: CoLa-ICD enriches ICD label representations with synonyms and LLM-generated terms, models code relations with a co-occurrence GNN, and aligns evidence with label semantics.The framework uses multi-head term attention and label-aware contrastive learning for alignment.

8 Limitations

The evaluation is limited to MIMIC datasets, chiefly ICU discharge summaries from one medical center, so performance may vary across settings and splits.

  • The experiments cover ICD-9 and ICD-10 but use only MIMIC datasets from a single medical center, primarily ICU discharge summaries.Performance may vary across institutions, specialties, and note types.
  • Rare-code results are sensitive to split composition because many codes appear only a few times in training.
  • Administrative and billing requirements can make target labels imperfectly aligned with explicitly stated clinical evidence.

A Additional Hyperparameter Details

Hyperparameters are selected using validation performance, with settings adjusted for the size of the label space in the main and supplementary experiments.

  • Hyperparameter selection: ρ is selected from {0.1, 0.3, 0.5}, and the main full-label experiments use ρ = 0.3.
  • Hyperparameter selection: K = 128 is used for main full-label experiments, while K = 50 is used for the MIMIC-III Top50 experiment.The main setting is chosen to cover the large label space, whereas the supplementary setting uses all available negative codes.
  • Hyperparameter selection: The temperature parameter τ is set to 0.1 following previous contrastive-learning work.
  • Hyperparameter selection: λ = 0.05 is selected from validation performance, and δ is chosen by maximizing validation Micro-F1.
  • Dataset statistics: Dataset statistics for MIMIC-III(1.4) and MIMIC-IV(3.1) are provided in Table 5.

B.2 Knowledge-Source Ablation

The ablation compares official descriptions, UMLS synonyms, LLM-generated terms, and their combination, finding complementary terminology benefits that are strongest for rare codes.

  • Knowledge-source comparison: Full Knowledge Injection achieves the highest value for every reported metric across both datasets.The comparison includes official descriptions, UMLS synonyms, LLM-generated terms, and their combination; results are averaged over ten random seeds.
  • Rare-code results: Rare Micro-F1 rises from 0.153 to 0.194 on MIMIC-III and from 0.155 to 0.204 on MIMIC-IV with Full Knowledge Injection.
  • Knowledge-source comparison: Both UMLS synonyms and LLM-generated terms improve over official ICD descriptions alone across both datasets.Their combination provides complementary lexical coverage for ICD labels.
  • Supplementary evaluations: MIMIC-III Top50 results are supplementary because the subset focuses on frequent labels rather than long-tail coding.
  • Supplementary evaluations: MIMIC-IV ICD-9 results are supplementary evidence for comparability with prior work, not the paper’s main ICD-10 generalization evidence.
  • Term-generation procedure: The appendix describes Gemini 2.5 Flash term generation and evaluates five candidate term sets, selecting a strategy on MIMIC-III for unchanged use on MIMIC-IV.

D.1 Classifier architecture and inputs

The classifier encodes discharge summaries and candidate ICD terms with BioClinicalBERT, using trainable term attention and classification components while selecting the term source on validation data.

  • Inputs and encoder: BioClinicalBERT encodes discharge-summary text and candidate ICD terms, producing representations used by the classifier.
  • Trainable components: BioClinicalBERT remains frozen, while the term-attention module and linear classifier are trainable.
  • Term-source selection: The candidate set maximizing validation Macro-F1 is selected, and the optimal source trains the final CoLa-ICD model.
Loading 2608.30234v1…