SlowGuess/NeurIPS_2023_OCR huggingface dataset:
is  $1 - \frac{2\times 2}{6 - 1} = 0.2$ , since  $|\mathcal{T}_i\cup \mathcal{T}_j| = |\{\hat{1},2,3,5,6,7\}| = 6$ . Although  $\mathcal{L}_1$  and  $\mathcal{L}_2$  do not have the same label nodes, they still share some similarities given their topological categories in the ICD hierarchical structure: the label node 2 in  $\mathcal{L}_2$  is the parent of 5 in  $\mathcal{L}_1$ , while the label node 6 in  $\mathcal{L}_2$  is a sibling of 7 in  $\mathcal{L}_1$ .
Figure 2: Examples of two spanning super-trees. Nodes with orange and green colors denote label nodes and their ancestor nodes.
Contrastive pre-training In the pre-training of the clinical note encoder Encnote, given a quadruple  $(s_k^i,s_{k'}^i,s_k^j,s_{k'}^j)$ , we first calculate the note representation  $\mathbf{h}_{\mathrm{sec}}$  and apply a max pooling layer to obtain section representations for all these sections:
$$
\mathbf {s} _ {\{k, k ^ {\prime} \}} ^ {\{i, j \}} = \operatorname {M a x P o o l i n g} \left(\operatorname {E n c} _ {\text {n o t e}} \left(s _ {\{k, k ^ {\prime} \}} ^ {\{i, j \}}\right)\right) \in \mathbb {R} ^ {d}. \tag {5}
$$
Then, we aim to utilize the soft similarity between the labels of two clinical notes to guide the similarity of these four section representations. Note that the similarity between sections in a positive pair is
Figure 3:Contrastive loss.
1 while the similarity of negative pairs is  $\alpha$ . Thus, we design a contrastive loss  $\mathcal{I}$  using the mean absolute error (MAE) as follows:
$$
\mathcal {J} = \mathcal {J} _ {m} (1, \beta (\mathbf {s} _ {k} ^ {i}, \mathbf {s} _ {k ^ {\prime}} ^ {i})) + \mathcal {J} _ {m} (\alpha_ {i j}, \beta (\mathbf {s} _ {k} ^ {i}, \mathbf {s} _ {k} ^ {j})) + \mathcal {J} _ {m} (1, \beta (\mathbf {s} _ {k} ^ {j}, \mathbf {s} _ {k ^ {\prime}} ^ {j})) + \mathcal {J} _ {m} (\alpha_ {i j}, \beta (\mathbf {s} _ {k ^ {\prime}} ^ {i}, \mathbf {s} _ {k ^ {\prime}} ^ {j})). \tag {6}
$$
Here,  $\mathcal{I}_m$  denotes the MAE loss,  $\alpha_{ij}$  denotes the aforementioned tree edit distance similarity between two label sets, and  $\beta (\cdot ,\cdot)$  is the cosine similarity function between two vectors. Figure 3 intuitively shows the proposed contrastive loss for the quadruple with soft similarity.
In summary, there are two main benefits of using contrastive learning on sections:
(1) The model gains a better understanding of the relationships between sections within a clinical note. This is particularly beneficial when a clinical note lacks explicit indicators of ICD codes, such as the "discharge diagnoses" section. With contrastive pre-training, the model can infer ICD codes by the content of related sections.
(2) It helps the model adapt to different writing styles of medical professionals and various demographic features of patients so that the model can focus on the text related to ICD codes.
4.3 Masked section training
As discussed earlier, existing methods mostly consider clinical notes as long sequences. Inspired by the denoising techniques for text [12], we develop a simple yet effective training strategy with permutation and masking on sections. It further mitigates the variability caused by the section order.
Given a clinical note that has been segmented into sections, denoted by  $S = \{t_k : s_k\}_{k=1}^T$ , we first shuffle the order of the sections to create an order-agnostic structure. Then, similar to the dropout technique [24] used to avoid overfitting in training deep learning models, we randomly mask a subset of the sections, subject to a threshold  $\gamma$  where  $0 \leq \gamma < 1$ . The remaining sections are concatenated back into a long sequence  $S'$  that is suitable for the input of existing ICD coding models. It is important to note that we do not aim to modify the original model architecture, but rather to generate samples that can help to reduce the variability of clinical notes in training. In summary, this section masking process can be described as follows:
$$
S ^ {\prime} = \bigoplus_ {k \in \operatorname {p e r m} (T)} s _ {k} ^ {\prime}, \quad \text {w h e r e} s _ {k} ^ {\prime} = \left\{ \begin{array}{l l} s _ {k} & \text {i f} | s _ {k} | > 0 \text {a n d} \theta \sim U [ 0, 1 ] \geq \gamma , \\ \text {e m p t y} & \text {o t h e r w i s e .} \end{array} \right. \tag {7}
$$
Table 1: Data statistics for the MIMIC-50, MIMIC-rare-50, and MIMIC-full tasks.
Task	Item	Train	Dev	Test
MIMIC-50	#Docs.	8,066	1,573	1,729
Avg.# words per Doc.	1,478	1,739	1,763
Avg.# codes per Doc.	5.7	5.9	6.0
Total # codes	50	50	50
MIMIC-rare-50	#Docs.	249	20	142
Avg.# words per Doc.	1,770	1,930	2,071
Avg.# codes per Doc.	1.0	1.0	1.0
Total # codes	50	50	50
MIMIC-full	#Docs.	47,723	1,631	3,372
Avg.# words per Doc.	1,434	1,724	1,731
Avg.# codes per Doc.	15.7	18.0	17.4
Total # codes	8,692	3,012	4,085
Here,  $\oplus$  denotes the concatenation operation, and  $U[0,1]$  means the uniform distribution from 0 to 1. By using perm, we can generate a random permutation of  $(1,2,\dots,T)$  of section indices, which is a random shuffle of all sections in a clinical note.
With shuffling and masking in training, the ICD coding model is no longer limited by the order of the clinical notes. Additionally, certain sections, such as "discharge diagnoses", may not always play a deterministic role in the prediction. This allows the model to focus more on other sections that are also relevant to the predicted ICD codes. Note that, in the inference step, we do not perform shuffling and masking, but use the original sequence as input for an ICD coding model.
5.1 Dataset, tasks, and evaluation metrics
The MIMIC-III [11] dataset is a popular publicly available EHR dataset that contains the discharge summaries and corresponding ground-truth ICD codes. We follow the ICD coding tasks in prior work [29, 30] and conduct three prediction tasks:
- MIMIC-50 prediction: Predicting the top 50 frequent ICD codes in the MIMIC-III dataset.
- MIMIC-rare-50 prediction: Predicting the rare 50 ICD codes that occur less than 10 times.
- MIMIC-full prediction: Predicting the entire (8,692) ICD codes in the MIMIC-III dataset.
The detailed training/dev/test dataset statistics for each task are listed in Table 1. Our experiments are conducted with cross-validation on the dev set to adjust hyper-parameters.
We use the following evaluation metrics which have been used in prior ICD coding studies [29, 30]. The metrics for MIMIC-full prediction are Macro/Micro  $F_{1}$  and precision at 8/15 (P@8, P@15). For MIMIC-50 prediction, we use Macro/Micro  $F_{1}$  and precision at 5 (P@5). For MIMIC-rare-50 prediction, we use Macro/Micro  $F_{1}$ .
5.2 Backbone models
To verify the effectiveness of our proposed section-based contrastive pre-training and masked section training (CM), we choose the following state-of-the-art ICD coding models as backbones<sup>1</sup>:
- MultiResCNN [13]: It encodes clinical notes with multi-filter residual CNN and label attention.
- HyperCore [4]: It also uses a convolutional encoder for text. Moreover, it applies hyperbolic embedding for ICD codes and uses GCN to model code co-occurrence.
<sup>1</sup>We do not include KEPT [29] here because our devices do not support the training of KEPT due to its high complexity. We list the result of KEPT in Appendix C.1 for reference.
Table 2: Results (\%) of MIMIC-50 when trained with and without the proposed contrastive pretraining and masked training (CM) strategies. Cells with the green color denote an improvement of w/ CM compared to w/o CM.
Model	w/o CM	w/ CM
Macro F1	Micro F1	P@5	Macro F1	Micro F1	P@5
MultiResCNN	60.8 (0.3)	67.1 (0.1)	64.3 (0.3)	62.2 (0.3)	68.1 (0.1)	65.1 (0.2)
HyperCore	61.1 (0.2)	66.2 (0.2)	63.5 (0.3)	62.0 (0.2)	67.4 (0.2)	64.5 (0.3)
JointLAAT	66.4 (0.1)	71.6 (0.2)	67.3 (0.4)	67.2 (0.2)	72.0 (0.3)	67.9 (0.1)
EffectiveCAN	66.7 (0.1)	71.5 (0.2)	66.4 (0.2)	67.5 (0.2)	71.8 (0.1)	67.8 (0.1)
PLM-ICD	64.5 (0.3)	69.3 (0.2)	64.5 (0.4)	65.2 (0.1)	70.3 (0.2)	65.6 (0.2)
Hierarchical	65.3 (0.1)	70.6 (0.3)	66.5 (0.1)	66.1 (0.2)	71.8 (0.4)	67.2 (0.3)
MSMN	68.1 (0.2)	72.0 (0.1)	67.5 (0.1)	69.1 (0.1)	72.5 (0.1)	68.3 (0.2)
- JointLAAT [25]: It uses a bidirectional LSTM to encode clinical notes and proposes a joint learning method to predict ICD codes and their parent codes in the ICD hierarchical structure.
- EffectiveCAN [14]: Similar to MultiResCNN, it also applies a convolutional encoder with multiple residual squeeze-and-excitation networks.
- PLM-ICD [9]: It is a transformer model (Roberta-base) that splits clinical notes into chunks to satisfy the maximum length of pre-trained large language models.
- Hierarchical [5]: It is hierarchical a transformer model (Roberta-large) by splitting clinical notes into paragraphs.
- MSMN [30]: It is an LSTM text encoder and incorporates the synonym of code descriptions to make the model better understand the variety of code names.
For DF-IAPF, we set the maximum word number  $(\mathcal{N})$  in n-gram to 5. We set  $K$  in top- $K$  candidates to 50 for the review of medical professionals. For contrastive pre-training, the batch size is 16, the learning rate is  $5 \times 10^{-4}$ , the optimizer is AdamW, and the epoch number is 20. The contrastive pre-training only uses the training dataset of each task to avoid data leakage. For the masked section training, we set  $\gamma$  to 0.2 for MIMIC-full prediction and 0.3 for MIMIC-50/MIMIC-rare-50 prediction.
The backbone models except HyperCore [4] and EffectiveCAN [14] are implemented using their publicly released code and the optimal parameters reported in their papers. For HyperCore and EffectiveCAN, the authors do not release the code. Therefore, we implemented a version that has a close performance to the original paper. For the MIMIC-50 and MIMIC-rare-50 tasks, we run every baseline 5 times and report their average and standard deviations (std).
All programs are executed using a machine with Python 3.9.3, CUDA 11.7, an Intel i9-11900K CPU, 64GB memory, and an NVIDIA RTX 3090 GPU. The code of the proposed DF-IAPF method and training strategies can be found at: https://github.com/LuChang-CS/semi-structured-icd-coding.
5.4 Experimental results
Extracted section titles To demonstrate the effectiveness of our proposed DF-IAPF algorithm to extract section titles, we compare it with a rule-based extraction algorithm [26]. It designs special rules for every observed section title based on colons and occurrence frequencies to segment clinical notes into sections. We list the extracted section titles and analyze the effectiveness and advantages of the proposed DF-IAPF algorithm in Appendix C.2.
MIMIC-50-prediction We report the results of MIMIC-50 in Table 2. Here, we run each backbone model 5 times and report their mean value and std. Among all backbone models, MSMN achieves the best result on all metrics without the proposed CM. Additionally, with the proposed CM strategies, the performance of all backbone models is improved, and Macro  $F_{1}$  is improved by  $1.5\%$  on average. Additionally, we also run a paired t-test on the Macro  $F_{1}$  score between the backbone models w/ CM and w/o CM. The  $p$ -values for all backbone models are less than  $5 \times 10^{-2}$ , indicating that the improvement brought by the CM strategies is statistically significant over the original models.
Figure 4: Abalation studies of the MIMIC-50 task when trained without or with the proposed contrastive pre-training (C) and masked section training (M). We also report a variant by replacing the tree edit distance with the simple Jaccard similarity (J).
MIMIC-rare-50-prediction
We report the results of MIMIC-rare-50 in Table 3. We observe that, with limited training samples in this task, the performance of backbone models is not as good as the MIMIC-50 prediction. However, the proposed CM strategies can significantly improve the prediction results. On average, the Macro  $F_{1}$  score is improved by  $55.8\%$ . Note that, the contrastive pre-training is conducted on every task instead of the entire clinical notes. Therefore, we can conclude that
Table 3: Results (\%) of MIMIC-rare-50 when trained with and without the proposed CM strategies. Cells with the green color denote an improvement of w/ CM compared to w/o CM.
Model	w/o CM	w/ CM
Macro F1	Micro F1	Macro F1	Micro F1
MultiResCNN	11.2 (2.1)	13.1 (1.8)	22.8 (1.3)	23.3 (1.9)
HyperCore	12.5 (1.3)	15.6 (2.2)	23.4 (1.9)	25.2 (1.2)
JointLAAT	20.2 (1.9)	21.6 (1.5)	28.6 (1.1)	27.8 (1.4)
EffectiveCAN	19.8 (1.4)	22.5 (2.1)	27.1 (2.4)	28.0 (1.5)
PLM-ICD	22.6 (2.5)	24.3 (1.9)	30.3 (1.5)	29.5 (1.3)
Hierarchical	23.1 (1.7)	24.6 (1.4)	32.0 (1.2)	31.3 (2.2)
MSMN	23.7 (1.0)	24.1 (1.9)	31.2 (1.3)	30.6 (1.7)
the proposed CM strategies can learn a good initialization of the ICD coding models in pre-training and serve as an effective data augmentation method in training with limited data.
Due to space constraints, we show the results of MIMIC-full in Appendix C.3. In summary, these experiments prove the capability of the DF-IAPF algorithm in extracting section titles and demonstrate the effectiveness of the proposed CM strategies in enhancing existing ICD coding models.
To validate the effectiveness of the contrastive pre-training and masked section training, we conduct an ablation study by only applying one strategy in the MIMIC-50 prediction task and replacing the tree edit distance with Jaccard similarity. Here, we choose MIMIC-50 prediction and use JointLAAT, EffectiveCAN, and MSMN as the backbones. Figure 4 demonstrates the prediction results when training without CM (w/o CM), with only contrastive pre-training using Jaccard similarity (w/ J), with only contrastive pre-training (w/ C) using tree edit distance, with only masked section training (w/ M), and with both two strategies (w/ CM). In this figure, pre-training with simple Jaccard similarity even decreases the performance. It shows that this similarity cannot appropriately guide contrastive learning because many clinical notes have disjoint labels. We notice that both contrastive pre-training and masked section training contribute to improving the performance of the original ICD coding models. Specifically, the masked section training is slightly better than the contrastive pre-training. We think it is because the masked section training is directly applied to training ICD coding models, while the contrastive pre-training learns good initialization before training.
5.6 Case studies
To intuitively demonstrate the effectiveness of the proposed contrastive pre-training and masked training, in Figure 5, we give an example of a clinical note snippet and the predictions w/o CM and w/ CM using MSMN on the MIMIC-50 prediction task. Here, the ground truth contains 4 ICD
Figure 5: An example of prediction without and with the proposed CM strategies using MSMN.
codes, 45.13 (procedure), 530.81 (disease), 96.04 (procedure), and V15.82 (disease). The MSMN model w/o CM predicts two codes correctly while failing to predict 96.04 and V15.82 because they do not occur in the discharge diagnosis and procedure sections. However, with CM, the MSMN model successfully predicts all four ICD codes by locating them in related sections, including "CXR Endotracheal tube" in physical exam and "former smoker" in social history.
In this work, we aim to minimize the variability of clinical notes in the ICD coding task by studying the semi-structured format of clinical notes. To reduce human effort, we propose an automatic algorithm to extract section titles and segment clinical notes into sections. We also design contrastive pre-training and masked section training to let the ICD coding model better locate sections related to predictions. Additionally, a tree-edit distance is designed in the loss function to measure the similarity of positive/negative pairs. Extensive experiments demonstrate the effectiveness of the proposed section title extraction algorithm and training strategies. It is worth emphasizing that our proposed methodology is versatile, as it can not only be applied to clinical notes but also employed in general multi-label classification tasks that involve semi-structures such as sections. In the future, we are committed to exploring the broader applicability of our approach across various domains.
Limitations: Although the proposed training strategies are able to enhance existing ICD coding models, they are dependent on the design of these models. If the model is well-designed and has many parameters, it is generally overfitting with limited training data. In this case, our proposed training strategies are a good enhancement. Additionally, we only focus on the variability caused by the order of sections in this work, but there are other formats of variability such as typos and synonyms. In the future, we plan to design new ICD coding models based on sections and consider more types of variability to further improve the robustness of the training process.
This work is supported in part by the US National Science Foundation under grants 1948432, 2047843, and 2245907. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation.
[1] Tian Bai, Shanshan Zhang, Brian L Egleston, and Slobodan Vucetic. Interpretable representation learning for healthcare via capturing disease progression through time. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pages 43-51, 2018.
[2] Tal Baumel, Jumana Nassour-Kassis, Raphael Cohen, Michael Elhadad, and Noemie Elhadad. Multi-label classification of patient notes: Case study on ICD code assignment. In 2018 AAAI Joint Workshop on Health Intelligence (W3PHIAI 2018), pages 409-416. AAAI press, 2018.
[3] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150, 2020.
Ethical considerations While EHR data contains private information of patients, the MIMIC-III dataset used in this work as well as all backbone models is a publicly available dataset. It de-identified the sensitive information of patients and doctors with masks, including admission/discharge date, name, and hospital name (e.g.,  $[{}^{**}\text{first name}3^{**}]$ ) to protect privacy. Therefore, the data we use will not leak such information even if we publish our code and model parameters.
Societal Impacts Incorrect ICD coding can lead to medical billing errors which can affect patients and healthcare costs. However, as an enhancement of existing ICD coding models, our work aims to improve the prediction accuracy of ICD coding. We believe our method does not bring additional negative societal impacts to ICD coding.
B Pseudo code of the DF-IAPF algorithm
We present the proposed DF-IAPF in Algorithm 1. In lines 7-9, this algorithm uses a span in clinical notes to obtain an n-gram as a phrase  $t$  and updates its occurrence in the local clinical note. Lines 10-12 update the global document frequency and phrase frequency. Finally, lines 15-16 calculate the DF-IAPF score for every phrase. In line 17 we sort all phrases descendingly by the DF-IAPF score and select the top phrases with the highest scores. Finally, we filter out shorter titles that are subsequences of longer titles with high scores in lines 18-20.
Note that this algorithm is an offline extraction of phrases before training. The computation procedures of DF-IAPF are similar to the TF-IDF, except that we add a for-loop of n-gram in line 3. Therefore, the time complexity of the DF-IAPF algorithm is  $\mathcal{N} \times \mathcal{O}(\mathrm{TF - IDF})$ . In our experiments, the running time of the DF-IAPF algorithm is about 2 minutes.
Algorithm 1: Section Title Extraction
Input :A set  $\mathcal{S}$  of clinical notes  $S = \{S\}$  An integer  $\mathcal{N}$  to control the maximum word count in n-grams; An integer  $K$  to select top-  $K$  phrases Output:A candidate set  $\mathcal{C}$  of section titles   
1 NT  $\leftarrow$  an empty mapping from phrases to counts with a default value of 0   
2 APF  $\leftarrow$  an empty mapping from phrases to a frequency list with a default value of an empty list   
3 for  $N\gets 1$  to  $\mathcal{N}$  do for  $S\in S$  do n  $\leftarrow$  the number of words in  $S$  PF  $\leftarrow$  an empty mapping from phrases to frequencies with a default value of 0 for  $i\gets 1$  to  $n - N + 1$  do t  $\leftarrow (w_{i},w_{i + 1},\dots ,w_{i + N - 1}) / /$  N-gram PF(t)  $\leftarrow$  PF(t)  $+1$  // Update the frequency of t in this document S for  $t\in PF$  do NT(t)  $\leftarrow$  NT(t)  $+1$  // Update the frequency of documents containing t Append PF(t) to APF(t) // Update the frequency list of t   
13 nd  $\leftarrow |S|$    
14 C  $\leftarrow$  an empty mapping from phrases to scores   
15 for  $t\in NT$  do  $\mathcal{C}(t)\gets \frac{\mathrm{NT}^2(t)}{n_d\times\sum_{i = 1}^{\mathrm{NT}(t)}\mathrm{APF}(t)_i}$  // DF-IAPF, Equation (2)   
17 C  $\leftarrow$  Sort C descendingly by the score and select  $K$  phrases with the highest scores   
18 for  $(t_1,t_2)\in \mathcal{C}\times \mathcal{C}$  do if  $t_1\subsetneq t_2$  then  $\mathcal{C}\gets \mathcal{C}\setminus \{t_1\} /$  Remove shorter titles that are subsequences of longer titles with high scores.   
21 return C
Table 4: Top 20 section titles extracted by our proposed DF-IAPF algorithm and a rule-based method using colons and occurrence frequencies.
Rank	DF-IAPF	Frequency	Rank	Rule-based	Frequency
1	history of present illness	0.95	1	admission date	1.00
2	date of birth	0.87	2	- service	0.95
3	+ sex	0.87	3	date of birth	0.87
4	+ discharge date	1.00	4	history of present illness	0.95
5	admission date	1.00	5	- allergies	0.87
6	social history	0.82	6	past medical history	0.90
7	past medical history	0.90	7	social history	0.82
8	discharge medications	0.83	8	- discharge disposition	0.75
9	medications on admission	0.77	9	discharge medications	0.83
10	discharge diagnosis	0.94	10	discharge diagnosis	0.94
11	discharge condition	0.85	11	medications on admission	0.77
12	discharge instructions	0.71	12	attending	0.71
13	major surgical or invasive procedure	0.78	13	family history	0.74
14	brief hospital course	0.98	14	discharge condition	0.85
15	pertinent results	0.68	15	discharge instructions	0.71
16	followup instructions	0.89	16	major surgical or invasive procedure	0.78
17	family history	0.74	17	physical exam	0.94
18	+ chief complaint	0.77	18	brief hospital course	0.98
19	attending	0.71	19	pertinent results	0.68
20	physical exam	0.94	20	followup instructions	0.89
23	service	0.95	38	chief complaint	0.77
28	discharge disposition	0.75	664	discharge date	1.00
29	allergies	0.87	1726	sex	1.00
C.1 Results of KEPT
We do not include KEPT [29] in the backbone models because our devices do not support the training of KEPT due to its high complexity. We list the result of KEPT (w/o CM) here for reference. It is worth noting our proposed contrastive pre-training and masked section training are also applicable to KEPT.
- MIMIC-full prediction:
- Macro  $F_{1}$  : 11.8
- Micro  $F_{1}$ : 59.9
- P@8:77.1
- P@15:61.5
- MIMIC-50 prediction:
- Macro  $F_{1}$  : 68.9
- Micro  $F_{1}$  : 72.9
- P@5: 67.3
- MIMIC-rare-50 prediction:
- Macro  $F_{1}$  : 30.4
- Micro  $F_{1}$ : 32.6
C.2 Extracted section titles
To demonstrate the effectiveness of our proposed DF-IAPF algorithm to extract section titles, we compare it with a rule-based extraction algorithm [26]. It designs special rules for every observed section title based on colons and occurrence frequencies to segment clinical notes into sections. We list the top 20 extracted section titles in Table 4.
Qualitative analysis Here, the rank is obtained using DF-IAPF scores (left) or occurrence frequencies (right). The symbol “+” indicates the title extracted by our DF-IAPF algorithm but not by the rule-based algorithm, while the symbol “-” means the title extracted by the rule-based algorithm but not the DF-IAPF algorithm in the top 20 section titles. In this Table, we observe that 17 titles are commonly extracted by both algorithms, indicating that our automatic section title algorithm is comparable to the hand-crafted rule-based method in terms of effectiveness. We further analyze the rank of missing section titles from both algorithms in the top 20 titles. All the titles that are not