[Paper Review] Towards Semi-Structured Automatic ICD Coding via Tree-based Contrastive Learning

Chang Lü, Chandan K. Reddy+2 more|arXiv (Cornell University)|Oct 14, 2023

Medical Coding and Health Information4 citations

TL;DR

This paper proposes a novel framework for semi-structured automatic ICD coding by first automatically segmenting clinical notes into sections, then applying tree-based contrastive learning with a soft multi-label similarity metric based on tree edit distance to reduce data variability. The method enhances existing ICD coding models through contrastive pre-training and masked section training, achieving significant performance gains on MIMIC-III benchmarks, especially in low-data regimes.

ABSTRACT

Automatic coding of International Classification of Diseases (ICD) is a multi-label text categorization task that involves extracting disease or procedure codes from clinical notes. Despite the application of state-of-the-art natural language processing (NLP) techniques, there are still challenges including limited availability of data due to privacy constraints and the high variability of clinical notes caused by different writing habits of medical professionals and various pathological features of patients. In this work, we investigate the semi-structured nature of clinical notes and propose an automatic algorithm to segment them into sections. To address the variability issues in existing ICD coding models with limited data, we introduce a contrastive pre-training approach on sections using a soft multi-label similarity metric based on tree edit distance. Additionally, we design a masked section training strategy to enable ICD coding models to locate sections related to ICD codes. Extensive experimental results demonstrate that our proposed training strategies effectively enhance the performance of existing ICD coding methods.

Motivation & Objective

• To address the challenge of limited and variable clinical note data in automatic ICD coding.
• To exploit the semi-structured nature of clinical notes by automatically segmenting them into meaningful sections.
• To reduce variability in clinical notes caused by differing writing styles and section ordering through contrastive pre-training and masked section training.
• To improve the performance of existing ICD coding models under data scarcity using structure-aware representation learning.
• To develop a versatile framework applicable to multi-label classification tasks with semi-structured data.

Proposed method

• Proposes a content-based algorithm to automatically extract section titles and segment clinical notes into sections, enabling order-agnostic modeling.
• Introduces a contrastive pre-training framework using a soft multi-label similarity metric based on tree edit distance to measure similarity between section pairs.
• Designs a masked section training strategy to help models focus on sections relevant to ICD code predictions.
• Employs tree edit distance in the loss function to better capture semantic and structural relationships between sections.
• Applies the proposed training strategies (contrastive pre-training and masked section training) as data augmentation to enhance existing ICD coding models.
• Uses backbones like JointLAAT, EffectiveCAN, and MSMN to evaluate the effectiveness of the proposed training strategies.

Experimental results

Research questions

• RQ1 Can automatic section segmentation improve the robustness of ICD coding models under data scarcity?
• RQ2 How effective is tree-based contrastive learning with soft multi-label similarity in reducing variability in clinical notes?
• RQ3 Does masked section training enhance model performance by improving attention to relevant clinical sections?
• RQ4 How does the proposed method compare to baseline ICD coding models in low-data settings?
• RQ5 Can the proposed framework be generalized to other multi-label classification tasks with semi-structured data?

Key findings

• The proposed section segmentation algorithm achieves high F1 scores on section title extraction, demonstrating its effectiveness in identifying clinical note structure.
• Contrastive pre-training with tree edit distance significantly improves ICD coding performance, while replacing it with Jaccard similarity degrades performance, indicating the importance of semantic-aware similarity.
• Masked section training yields a slightly better performance gain than contrastive pre-training alone, suggesting direct optimization on the prediction head is highly effective.
• The combination of both contrastive pre-training and masked section training leads to the best performance across all evaluated backbones (JointLAAT, EffectiveCAN, MSMN) on the MIMIC-50 task.
• Case studies show that the model with the proposed strategies successfully predicts all four ICD codes by attending to relevant sections like 'physical exam' and 'social history', whereas the baseline fails to detect codes not in the main diagnosis sections.
• The framework demonstrates strong generalization, with consistent improvements across multiple ICD coding models, confirming its versatility and effectiveness in low-data scenarios.