A General Knowledge Injection Framework for ICD Coding

TL;DR

This work tackles the challenge of assigning accurate ICD codes to clinical texts under long-tail code distributions and limited code-specific evidence. It introduces GKI-ICD, a general, module-free knowledge injection framework that unifies ICD Description, Synonym, and Hierarchy through guideline synthesis and multi-task learning, allowing joint utilization of diverse knowledge without bespoke architectures. Empirical results on MIMIC-III benchmarks demonstrate state-of-the-art performance across multiple metrics, including improvements on rare codes, and ablation confirms the value of integrating multiple knowledge sources. This approach offers a scalable, effective means to enhance ICD coding and other multi-label classification tasks by leveraging structured domain knowledge during training.

Abstract

ICD Coding aims to assign a wide range of medical codes to a medical text document, which is a popular and challenging task in the healthcare domain. To alleviate the problems of long-tail distribution and the lack of annotations of code-specific evidence, many previous works have proposed incorporating code knowledge to improve coding performance. However, existing methods often focus on a single type of knowledge and design specialized modules that are complex and incompatible with each other, thereby limiting their scalability and effectiveness. To address this issue, we propose GKI-ICD, a novel, general knowledge injection framework that integrates three key types of knowledge, namely ICD Description, ICD Synonym, and ICD Hierarchy, without specialized design of additional modules. The comprehensive utilization of the above knowledge, which exhibits both differences and complementarity, can effectively enhance the ICD coding performance. Extensive experiments on existing popular ICD coding benchmarks demonstrate the effectiveness of GKI-ICD, which achieves the state-of-the-art performance on most evaluation metrics. Code is available at https://github.com/xuzhang0112/GKI-ICD.

Figures (4)

• Figure 1: An example of ICD coding: Occurrence of multiple codes and noisy content in a long medical text document makes it hard to link each ICD code to its corresponding evidence (marked in same color), explaining the necessity of incorporating code-specific knowledge.
• Figure 2: Our proposed general knowledge injection training framework for ICD coding, GKI-ICD. For each training sample, we first retrieve code-specific knowledge to synthesize a guideline, and then use this guideline and multi-task learning to inject knowledge into the model. Note that our method only incorporates knowledge in the training stage, which has no effect on the computation cost of the model during the inference stage.
• Figure 3: The model architecture adopted in our work.
• Figure 4: Case Study on MIMIC-III-Top-50 Dataset. We visualize the predicted ICD codes and the retrieved evidence of PLM-CA and our method. The red means the token which gains the greatest attention weight.