Multi-label Few-shot ICD Coding as Autoregressive Generation with Prompt

TL;DR

This paper introduces a novel autoregressive generation approach with prompts for multi-label ICD coding, significantly improving few-shot and full data performance by transforming the task into text generation and leveraging a new ensemble reranker.

Contribution

The study proposes a new generation-based method with prompt design and a novel pretraining objective for improved ICD coding, especially in few-shot scenarios.

Findings

Outperforms previous models in few-shot ICD coding with a macro F1 of 30.2.

Significantly improves full-data ICD coding macro F1 to 14.6 with ensemble reranker.

Demonstrates effectiveness of text generation and prompt-based methods in complex medical coding tasks.

Abstract

Automatic International Classification of Diseases (ICD) coding aims to assign multiple ICD codes to a medical note with an average of 3,000+ tokens. This task is challenging due to the high-dimensional space of multi-label assignment (155,000+ ICD code candidates) and the long-tail challenge - Many ICD codes are infrequently assigned yet infrequent ICD codes are important clinically. This study addresses the long-tail challenge by transforming this multi-label classification task into an autoregressive generation task. Specifically, we first introduce a novel pretraining objective to generate free text diagnoses and procedure using the SOAP structure, the medical logic physicians use for note documentation. Second, instead of directly predicting the high dimensional space of ICD codes, our model generates the lower dimension of text descriptions, which then infer ICD codes. Third, we designed a novel prompt template for multi-label classification. We evaluate our Generation with Prompt model with the benchmark of all code assignment (MIMIC-III-full) and few shot ICD code assignment evaluation benchmark (MIMIC-III-few). Experiments on MIMIC-III-few show that our model performs with a marco F1 30.2, which substantially outperforms the previous MIMIC-III-full SOTA model (marco F1 4.3) and the model specifically designed for few/zero shot setting (marco F1 18.7). Finally, we design a novel ensemble learner, a cross attention reranker with prompts, to integrate previous SOTA and our best few-shot coding predictions. Experiments on MIMIC-III-full show that our ensemble learner substantially improves both macro and micro F1, from 10.4 to 14.6 and from 58.2 to 59.1, respectively.