SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chang Lu*, Chandan K. Reddy†, Ping Wang*, Yue Ning*  
* Department of Computer Science, Stevens Institute of Technology † Department of Computer Science, Virginia Tech  
* {clu13, ping.wang, yue.ning}@stevens.edu † reddy@cs.vt.edu
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.
The adoption of electronic health records (EHR) data has become widespread in modern healthcare facilities as they provide a centralized platform to maintain comprehensive medical information of patients, including diagnoses, procedures, laboratory tests, and clinical notes [1]. To efficiently manage and categorize diseases and procedures, EHR data utilizes the International Classification of Diseases (ICD) system developed by the World Health Organization. The ICD system provides a hierarchical structure that maps diseases/procedures to digital codes. Clinical notes in EHR data are generally stored as free text, while diagnosis and procedure codes are extracted from these notes and saved as structured data. The process of extracting ICD codes from clinical notes is referred to as ICD coding and is a crucial task in medical services such as medical records management, medical billing [23], and insurance reimbursement [19]. It also supports healthcare research endeavors such as diagnosis prediction [1, 15] and medication recommendation [22].
The traditional ICD coding task relies on human effort, which is both time-consuming and prone to errors [27]. Incorrect code assignments can be costly. For instance, the error payout rate due to wrong code assignment reached  $6.8\%$  in 2000, as stated by the US Centers for Medicare and Medicaid's statistics [17]. Consequently, researchers are exploring automated ICD coding methods to assign ICD codes to medical documents with algorithms. Recent methods generally treat the ICD coding task as a multi-label classification problem [13, 32, 30], as one clinical note can contain multiple diagnosis/procedure codes. To capture the relationship between text and codes, code representations have been studied by incorporating the semantic information of code names [18] with hierarchical structures, synonyms, and co-occurrence of codes to provide fine-grained code representations and
relationships [4, 30, 29]. However, due to the nature of clinical notes, automatic ICD coding tasks still present certain challenges when it comes to learning the representation of clinical notes:
1. Limited availability of data. Due to privacy constraints, EHR data can be challenging to access. For example, a publicly available EHR dataset, MIMIC-III [11], only contains around 50,000 clinical notes that can be used for ICD coding. Furthermore, the appearance of codes in an EHR dataset follows a long-tail distribution. Around  $51.6\%$  of codes occur less than 6 times, and  $60.2\%$  of codes occur less than 10 times, making it even more difficult to train data-driven deep learning models for these codes due to data paucity.
2. Ignoring structural information. Based on our observation, most clinical notes have common sections such as "physical exam", "history of present illness", "discharge followups", and "brief hospital course", as depicted in Figure 1(a). These sections reflect correlated health information in long clinical notes. However, most of the existing ICD coding models treat clinical notes as a single sequence, disregarding these semi-structures that represent essential elements of diagnoses.
3. Variability of clinical notes. Clinical notes are written by various medical professionals, who may have different writing habits. Different clinical notes may present different orders and combinations of sections, adding to the variability of the data. Additionally, different patients may also lead to a diverse content of clinical notes due to their unique diseases or physical exams. This variability becomes a more negative factor in training an ICD coding model with limited data.
In this paper, our goal is to utilize the semi-structured format and reduce the variability of clinical notes with limited data. As shown in Figure 1(b), existing methods typically treat clinical notes as long sequences of words [16, 30, 29] without considering the semi-structured format of clinical notes. As a result, these models can easily be affected by the variability of clinical notes with limited data. To overcome this challenge, we propose to automatically segment a clinical note into multiple sections to build an order-agnostic structure. Based on the extracted sections, we introduce a contrastive learning framework to initially reduce the variability in clinical notes in pre-training and allow the text encoder to better understand the relationship of sections with limited training data. This proposed contrastive
(a) Sections in clinical note
Figure 1: An example of a clinical note in the format of multiple sections and one sequence of words.
(b) A clinical note as a sequence
learning method defines a soft multi-label similarity between section pairs from the same and different clinical notes. Finally, we design a masked section training method to further minimize the variability of clinical notes in the training of ICD coding models.
Main contributions In summary, the main contributions of this paper are listed as follows:
- We propose a content-based algorithm that automatically segments clinical notes into sections. To the best of our knowledge, our work is one of the first to investigate automatic semi-structured segmentation for clinical notes in ICD coding.
- We present a contrastive learning framework based on a soft multi-label similarity with tree edit distance and a masked section training strategy to alleviate the variability of clinical notes with limited EHR data.
- We conduct extensive experiments on real-world EHR datasets and demonstrate that our proposed section-based learning can enhance the performance of existing ICD coding methods.
The task of ICD coding involves predicting ICD codes from discharge summaries and can be approached as a multi-label text categorization problem. In the past, natural language processing (NLP) techniques have been commonly used to learn the representation of clinical notes. For instance, Perotte et al. [21] used the term frequency-inverse document frequency (TF-IDF) features of clinical notes and employed support vector machine (SVM) classifiers for ICD coding. Mullenbach et al. [18] proposed CAML based on convolutional neural networks (CNN), while Baumel et al. [2] employed a two-layer recurrent neural network (RNN) to encode clinical notes. Liu et al. [14] applied squeeze-and-excitation networks in CNN and used the focal loss to deal with rare codes.
In addition to the conventional RNN/CNN-based models, attention and graph neural networks have also been explored in the context of ICD coding. Li et al. [13] proposed a multi-filter residual CNN model incorporating label attention between codes and text. Xie et al. [28] utilized the hierarchical structure of ICD codes and developed a graph neural network to capture the relation between codes. Cao et al. [4] considered both the hierarchical structure and the co-occurrence of ICD codes, embedding the codes into a hyperbolic space. Yuan et al. [30] proposed to improve the matching of the ICD code names that occur in clinical notes using attention with code synonyms.
To alleviate the lack of training data and labels, weak supervision has also been studied for ICD coding. It aims to automatically generate weakly labeled training data using rules, heuristics, or medical domain knowledge. Dong et al. [7] adopted an existing named entity linking tool called SemEHR to identify rare diseases from clinical notes. Gao et al. [8] proposed a labeling function called KeyClass by extracting n-grams as keywords and computing the cosine similarity of word embeddings between keywords and labels.
Although large language models (LLMs) are popular and effective in many NLP tasks such as machine translation and question-answering systems, it has been shown by Pascual et al. [20] and Ji et al. [10] that pre-trained LLMs such as BERT [6] do not help to improve the performance of ICD coding due to the long text of clinical notes and the difference of training data between ICD coding and pre-training tasks of LLMs. To overcome these problems, Liu et al. [9] split the clinical notes into chunks to fit the pre-defined maximum input length of transformer-based models. However, splitting a document into chunks can break coherent information in clinical notes. Yang et al. [29] introduced KEPT, a transformer-based model that uses Longformer [3] to encode the long text, pre-trained with a contrastive learning method for code synonyms. They also designed a prompt learning framework for the prediction. However, the Longformer and prompt learning used in KEPT require a huge number of model parameters and extremely long input, which is barely applicable in training.
As previously discussed, most current methods only treat clinical notes as long sequences without considering their semi-structured format, making it challenging to handle the variability of clinical notes. Additionally, while some models, like KEPT, incorporate pre-training, it is designed only for labels but not for clinical notes. Thus, these models are not effective in comprehending the relationship among different sections of clinical notes. In light of these limitations, our paper aims to investigate the semi-structured format of clinical notes and improve the model's ability to learn the representations of long-text clinical notes with limited data.
Problem formulation Consider the ICD codes as a set  $\mathcal{L} = \{l_i\}_{i=1}^L$ , where  $L = |\mathcal{L}|$  is the number of codes. Specifically,  $l_i = \{w_j\}_{j=1}^m$  with  $m$  tokens is the description of the  $i$ -th label, where  $w_j \in \mathcal{V}$  and  $\mathcal{V}$  is the vocabulary of all tokens in the ICD code descriptions and clinical notes. Given a clinical note  $S = \{w_j\}_{j=1}^n$  with  $n$  tokens, the ICD Coding task is to train a model  $\mathcal{M}$  to predict a binary vector  $\hat{\mathbf{y}} \in \{0,1\}^L$ , where  $\hat{\mathbf{y}}_i = 1$  means the code  $l_i$  exists in the clinical note  $S$ .
General ICD coding framework To better demonstrate key parts of the ICD coding, we simplify it as flat multi-label classification. A general ICD Coding framework  $\mathcal{M}$  contains three modules:
- Clinical note encoder  $(\mathsf{Enc}_{\mathsf{note}})$ : Given a clinical note  $S$ , the clinical encoder is a text encoder  $\mathsf{Enc}_{\mathsf{text}}$  that first encodes the words into embeddings and uses RNN, CNN, or Transformer encoder to compute hidden representations  $\mathbf{h}_{\mathsf{note}}$  of words:  $\mathbf{h}_{\mathsf{note}} = \mathsf{Enc}_{\mathsf{text}}(S) \in \mathbb{R}^{n \times d}$ .
- ICD code encoder  $(\mathsf{Enc}_{\mathsf{code}})$ : This module can be regarded as a domain knowledge encoder that incorporates the text description of all codes (i.e., code names) in the ICD system, which are agnostic to the training-data. It is also a text encoder that first calculates the hidden word representations of a code name and then uses a pooling layer (e.g., mean/max pooling) on words to get the code representation for one ICD code:  $\mathbf{h}_{\mathrm{code}}^i = \mathrm{Pooling}(\mathsf{Enc}_{\mathrm{text}}(l_i)) \in \mathbb{R}^d$ . Eventually, we have the hidden representations of all the ICD codes:  $\mathbf{h}_{\mathrm{code}} \in \mathbb{R}^{L \times d}$ .
- Fusion between note and code (Fusion): This module aggregates the representations of clinical notes and ICD codes to generate predictions, denoted as  $\hat{\mathbf{y}} = \mathrm{Fusion}(\mathbf{h}_{\mathrm{note}},\mathbf{h}_{\mathrm{code}})$ . To achieve this, it first applies an attention mechanism between the codes and notes by calculating  $\mathbf{q}_{\mathrm{code}} = \mathrm{Attn}(\mathbf{h}_{\mathrm{code}},\mathbf{h}_{\mathrm{note}},\mathbf{h}_{\mathrm{note}}) \in \mathbb{R}^{L\times d}$ , where the query is the code representation  $\mathbf{h}_{\mathrm{code}}$  and the key and
value are note representation  $\mathbf{h}_{\mathrm{note}}$ . It then takes the dot product between the attention output and the code representation to obtain the final output  $\mathbf{o} = \mathbf{q}_{\mathrm{code}} \odot \mathbf{h}_{\mathrm{code}} \in \mathbb{R}^L$ . Finally, a sigmoid function is applied to get the final prediction  $\hat{\mathbf{y}}$ .
Both Encnote and Enccode contain a text encoder Enctext. It is a common practice to share the parameters of these two text encoders including word embeddings and model weights.
We first present an algorithm to automatically extract section titles and segment clinical notes into sections. Then, we introduce the proposed training strategies for existing ICD coding models: contrastive pre-training and masked section training based on the extracted sections to reduce the variability of clinical notes with limited training data.
4.1 Automatic section-based segmentation
As mentioned in Section 1, clinical notes typically contain sections with standard titles, but the order of these sections may vary depending on the writing style of medical professionals. To reduce the variability in clinical notes, it is important to extract sections related to ICD codes. The initial step is to identify all possible section titles for further segmentation. However, since clinical notes are written in plain text, there are no universal rules to extract these titles. Consequently, an automatic segmentation algorithm based on the content of clinical notes is needed to extract the section titles.
Inspired by TF-IDF which can retrieve keywords in a document, we propose an n-gram document frequency-inverse average phrase frequency (DF-IAPF) algorithm to extract section titles. TF-IDF captures the unique importance of a word for a document. In TF-IDF, a word becomes a keyword of a document when it has a high term frequency in this document while few documents contain this word. However, extracting section titles is different from extracting keywords for the following reasons:
(1) Section titles are usually phrases instead of single words (e.g., "history of present illness");
(2) Unlike keywords that are common in a document but less frequent in a corpus, most clinical notes have similar section titles but they often appear only once within a clinical note.
Based on these two properties of section titles, we introduce DF-IAPF to automatically extract section titles based on the corpus-level frequency and uniqueness of phrases in the document. We first define the DF-IAPF score for a phrase  $t = (w_{1}, w_{2}, \dots, w_{N})$  that contains  $N$  words (n-gram).
Document frequency-inverse average phrase frequency We first let  $\mathrm{DF}(t)$  be the relative frequency of documents containing  $t$ , and  $\mathrm{IAPF}(t)$  be the inverse average phrase frequency of  $t$  in all documents containing  $t$ :
$$
\mathrm {D F} (t) = \frac {n _ {t}}{n _ {d}}, \quad \mathrm {I A P F} (t) = \frac {1}{\frac {1}{n _ {t}} \sum_ {i = 1} ^ {n _ {d}} f _ {t , i}} = \frac {n _ {t}}{\sum_ {i = 1} ^ {n _ {d}} f _ {t , i}}, \tag {1}
$$
where  $n_d$  is the total number of documents,  $n_t$  is the number of documents containing  $t$ , and  $f_{t,i}$  is the occurrence number of  $t$  in the document  $i$ . The document frequency-inverse average phrase frequency (DF-IAPF) is defined as follows:
$$
\operatorname {D F - I A P F} (t) = \operatorname {D F} (t) \times \operatorname {I A P F} (t) = \frac {n _ {t}}{n _ {d}} \times \frac {n _ {t}}{\sum_ {i = 1} ^ {n _ {d}} f _ {t , i}} = \frac {n _ {t} ^ {2}}{n _ {d} \sum_ {i = 1} ^ {n _ {d}} f _ {t , i}}. \tag {2}
$$
The DF-IAPF algorithm assigns a higher score to phrases that appear frequently across all documents but occur less frequently within each document on average. For example, the formal section title, "brief hospital course", should have a higher score than a random phrase "this patient has". This is because most clinical notes contain a unique section titled "brief hospital course", while the phrase "this patient has" is more commonly used and appears multiple times in a clinical note, which lowers its score in the DF-IAPF algorithm. Then, we iterate through all n-grams with a maximum word count of  $\mathcal{N}$  in clinical notes to select candidates for section titles. Since we use n-gram to extract phrases, finally, we filter out shorter titles that are subsequences of longer titles with high scores. The specific algorithm to extract candidates and complexity analysis are presented in Appendix B.
Once the section title candidates with the highest DF-IAPF scores have been retrieved, we manually select phrases from this small candidate set to form a title subset  $\{t_1, t_2, \ldots, t_T\}$  with  $T$  titles. This selection process is completed by medical experts to ensure the correctness of selected titles. Since section titles are mostly unique within a clinical note, we use the first occurrence position of the extracted section titles as anchors to segment each clinical note into multiple sections  $\{s_k\}_{k=1}^T$  and build an order-agnostic structure. Given a clinical note  $S$ , the segmentation process from the plain text  $S$  to sections  $s_k$  with  $n_k$  words can be summarized as follows:
$$
S \xrightarrow {\text {D F - I A P F s e g m e n t a t i o n}} \left\{t _ {k}: s _ {k} \right\} _ {k = 1} ^ {T}, \tag {3}
$$
where  $t_k$  denotes a section title and  $s_k = (w_1, w_2, \ldots, w_{n_k})$  is the content under the section  $t_k$ .
4.2 Supervised tree-based contrastive learning on sections
In a general ICD coding framework, Fusion is an attention mechanism that enables the selection of significant words in clinical notes related to code descriptions. Ideally, a clinical note should contain sections called "discharge diagnoses" and "major procedures", which encompass all code descriptions corresponding to the labels. In this case, the model can accurately extract codes from these sections. However, many clinical notes lack these two sections. Even if they exist in some clinical notes, the descriptions may be incomplete. Typically, these sections contain only primary diagnosis or procedure codes, while the labels include all secondary codes. Under these circumstances, the model must locate related records from other sections such as "physical exam" or "discharge medications", given that these sections may imply the key expressions for the diagnoses or procedures. Thus, it is necessary to improve the model's ability to comprehend the content of each section.
To accomplish this, we design a contrastive learning framework based on sections. It makes the clinical note encoder distinguish sections from the same clinical note or different clinical notes so that the model can be aware of similar clinical notes and recognize related sections.
Construction of contrastive samples Since we formulate ICD coding as a multi-label classification task, it is hard to find two clinical notes with the same labels as a positive pair, and it is too trivial to obtain negative pairs using two clinical notes with different labels. Therefore, we construct positive/neighbor section pairs for further training. Positive pairs: For each clinical note  $S_{i} = \{t_{k}: s_{k}^{i}\}_{k=1}^{T}$ , we randomly select an anchor section  $s_{k}^{i}$ . To increase the connectivity between sections in the same note, we then sample a different section  $s_{k'}^{i}$  from the same clinical note  $S_{i}$  to build a positive pair  $(s_{k}^{i}, s_{k'}^{i})$ , where  $t_{k'} \neq t_{k}$ . Furthermore, from a different clinical note  $S_{j}$ , we sample two sections  $s_{k}^{j}$  and  $s_{k'}^{j}$  that correspond to  $t_{k}$  and  $t_{k'}$  in  $S_{i}$ , to build two neighbor pairs:  $(s_{k}^{i}, s_{k}^{j})$  and  $(s_{k'}^{i}, s_{k'}^{j})$ . Note that,  $s_{k}^{j}$  and  $s_{k'}^{j}$  are also a positive pair. Finally, a new sample for contrastive learning is a quadruple including the anchor, positive section, and two neighbor sections, i.e.,  $(s_{k}^{i}, s_{k'}^{i}, s_{k}^{j}, s_{k'}^{j})$ .
Soft multi-label similarity To achieve the goal of contrasting section pairs in the same or different clinical notes, an intuitive idea is to calculate the Jaccard similarity between two label sets or the cosine similarity between label vectors. Unfortunately, these metrics cannot capture the underlying disease relationships in the ICD system. For example, the cosine or Jaccard similarity for two sets of disease labels:  $\{\text{diabetes type} I\}$  and  $\{\text{diabetes type} II\}$ , will be zero because they have no overlap. However, these two diseases both belong to diabetes in the ICD system. Instead of assigning hard contrastive label 0 or 1 to two label sets, we design a soft similarity of two label sets that considers disease relationships in the ICD hierarchical structure by utilizing the tree edit distance [31] that measures the minimum number of node edit operations (add, delete, and replace) required to transform one tree into another.
Definition 1 (Spanning super-tree). Given the ICD hierarchical structure  $\mathcal{H}$  as a tree, the label set  $\mathcal{L}_i$  of clinical notes  $S_i$ , the spanning super-tree  $\mathcal{T}_i$  of  $\mathcal{L}_i$  is defined as a minimum tree that has the same root as  $\mathcal{H}$  and contains all label node of  $\mathcal{L}_i$  and all ancestors of  $\mathcal{L}_i$ :  $\mathcal{T}_i = \bigcup_{l \in \mathcal{L}_i} \rho(l) \cup \mathcal{L}_i \subset \mathcal{H}$ . Here,  $\rho(l)$  denotes all the ancestors of one label node  $l$ .
Based on the spanning super-tree, the similarity  $\alpha_{ij}$  between two label sets  $\mathcal{L}_i$  and  $\mathcal{L}_j$  is defined as:
$$
\alpha_ {i j} = 1 - \frac {2 \times \operatorname {d i s t} \left(\mathcal {T} _ {i} , \mathcal {T} _ {j}\right)}{\left| \mathcal {T} _ {i} \cup \mathcal {T} _ {j} \right| - 1} \in [ - 1, 1 ], \tag {4}
$$
where  $\text{dist}$  denotes the tree edit distance [31] between  $\mathcal{T}_i$  and  $\mathcal{T}_j$ .
In this similarity metric, we consider both the tree edit distance and the cardinality of trees. In the denominator, we use  $|\mathcal{T}_i \cup \mathcal{T}_j| - 1$  because every pair of  $\mathcal{T}_i$  and  $\mathcal{T}_j$  share the same root. Figure 2 shows two super-tree examples. In the first tree, nodes 5 and 7 are the labels of the clinical note  $S_1$ , while nodes 1, 2, and 3 are their ancestors. In the second tree, nodes 2 and 6 are labels of another clinical note, and nodes 1 and 3 are their ancestors. The tree with all colored nodes forms the spanning super-tree. The distance between these two spanning super-trees is 2 because we can delete node 5 and replace node 7 with 6 in the first tree to transform it into the second tree. Thus, the similarity of these two trees