[Paper Review] PLM-ICD: Automatic ICD Coding with Pretrained Language Models

Chao-Wei Huang, Shang-Chi Tsai+1 more|arXiv (Cornell University)|Jul 12, 2022

Topic Modeling4 citations

TL;DR

This paper proposes PLM-ICD, a framework that leverages pretrained language models (PLMs) for automatic ICD coding by addressing three key challenges: large label space, long input sequences, and domain mismatch. By integrating domain-specific pretraining, segment pooling, and label attention mechanisms, PLM-ICD achieves state-of-the-art performance on the MIMIC-3 dataset, outperforming prior PLM-based methods with a micro-F1 of 59.8% and micro-AUC of 88.65%.

ABSTRACT

Automatically classifying electronic health records (EHRs) into diagnostic codes has been challenging to the NLP community. State-of-the-art methods treated this problem as a multilabel classification problem and proposed various architectures to model this problem. However, these systems did not leverage the superb performance of pretrained language models, which achieved superb performance on natural language understanding tasks. Prior work has shown that pretrained language models underperformed on this task with the regular finetuning scheme. Therefore, this paper aims at analyzing the causes of the underperformance and developing a framework for automatic ICD coding with pretrained language models. We spotted three main issues through the experiments: 1) large label space, 2) long input sequences, and 3) domain mismatch between pretraining and fine-tuning. We propose PLMICD, a framework that tackles the challenges with various strategies. The experimental results show that our proposed framework can overcome the challenges and achieves state-of-the-art performance in terms of multiple metrics on the benchmark MIMIC data. The source code is available at https://github.com/MiuLab/PLM-ICD

Motivation & Objective

• Address the underperformance of pretrained language models (PLMs) in automatic ICD coding despite their success in other NLP tasks.
• Identify and resolve three core challenges: large label space, long input sequences exceeding PLM context length, and domain mismatch between general-domain pretraining and clinical text.
• Develop a unified framework that integrates tailored solutions to enhance PLM effectiveness for multi-label, long-sequence, domain-specific ICD coding.
• Establish best practices for fine-tuning PLMs on medical NLP tasks with similar characteristics.
• Demonstrate that PLM-based models can achieve competitive or superior performance on ICD coding benchmarks when properly adapted.

Proposed method

• Apply domain-specific pretraining using medical corpora to reduce domain mismatch between pretraining and fine-tuning data.
• Implement segment pooling to handle long clinical notes by dividing input into overlapping segments, encoding each with the PLM, and aggregating representations.
• Integrate a label attention mechanism that computes attention scores between the PLM’s hidden representations and each ICD label independently, enabling effective multi-label classification.
• Use subword tokenization with custom vocabularies (e.g., PubMedBERT, RoBERTa-PM) to reduce over-fragmentation and improve coverage of medical terminology.
• Optimize the fine-tuning process with learning rate scheduling and warmup, showing sensitivity to hyperparameters.
• Evaluate and compare different PLM architectures and attention mechanisms to identify the most effective configuration.

Experimental results

Research questions

• RQ1 Why do standard fine-tuning schemes fail to leverage the full potential of PLMs in automatic ICD coding?
• RQ2 How do long input sequences and large label sets impact PLM performance in clinical text classification?
• RQ3 To what extent does domain mismatch between general-domain pretraining and clinical text affect PLM performance?
• RQ4 Which combination of architectural and training strategies yields optimal performance for ICD coding with PLMs?
• RQ5 What are the key hyperparameters and design choices that significantly influence the final performance of PLM-based ICD coding models?

Key findings

• PLM-ICD achieves a micro-F1 of 59.8% and micro-AUC of 88.65% on the MIMIC-3 full test set, outperforming prior PLM-based models.
• The label attention mechanism from LAAT yields the best performance (10.4% macro-F1), significantly outperforming CAML (8.7%) and BERT-XML (8.2%).
• Segment pooling outperforms both HIER-BERT (2.8% macro-F1) and Longformer (5.1% macro-F1), demonstrating its effectiveness for long clinical notes.
• Using a maximum document length of 3,072 tokens and segment length of 128 achieves optimal trade-off between performance and computational cost.
• The learning rate and warmup schedule are critical: reducing the learning rate to 2e-5 degrades performance by ~3%, and a constant learning rate reduces performance by ~4%.
• Models with domain-specific vocabularies (e.g., PubMedBERT, RoBERTa-PM) show better performance and reduced over-fragmentation compared to BERT-based models with standard vocabularies.