Horses to Zebras: Ontology-Guided Data Augmentation and Synthesis for ICD -9 Coding (2022)

Abstract: Medical document coding is the process of assigning labels from a structured label space (ontology – e.g., ICD-9) to medical documents. This process is laborious, costly, and error-prone. In recent years, efforts have been made to automate this process with neural models. The label spaces are large (in the order of thousands of labels) and follow a big-head long-tail label distribution, giving rise to few-shot and zero-shot scenarios. Previous efforts tried to address these scenarios within the model, leading to improvements on rare labels, but worse results on frequent ones. We propose data augmentation and synthesis techniques in order to address these scenarios. We further introduce an analysis technique for this setting inspired by confusion matrices. This analysis technique points to the positive impact of data augmentation and synthesis, but also highlights more general issues of confusion within families of codes, and underprediction.

## introduction large-scale multi-labelled text classification (lmtc) tasks, such as automated icd-9 coding of mimic-iii discharge summaries, suffer from a big-head long-tail distribution of classes. this phenomenon naturally arises due to some labels being more frequent than others. this can further be affected by the source of the data -in the case of clinical natural language processing (nlp), this is often a single institution. for instance, hospitals in switzerland are unlikely to have cases of injuries caused by shark bites (code w56.41xd in icd-10). hence, depending on the data source, some labels will have a very small population -few-shot (fs), or even no population at all -zero-shot (zs) scenario. furthermore, adding new labels into a standard by splitting/fusing/altering existing concepts, or introducing new concepts also creates a zs scenario. medical coding methods need to be able to adapt to these scenarios. medical coding methods can be broadly divided on the task (document-level/entity-level) and the approaches (rule-based/machine learning) they use. methods like apache ctakes @xcite , medcat @xcite , or se-mehr @xcite perform named entity recognition and linking (ner+l), identifying specific spans of text within the document and associating them with concepts in a knowledge base, such as umls @xcite . they primarily use rule-based methods, with some inclusion of machine learning -e.g., medcat uses contextual word embeddings for disambiguating homographic strings, such as context-sensitive abbreviations (e.g., the string "hr" can mean "hour" or "heart rate"). neural lmtc models, such as caml @xcite , predict labels on the document level. labels are not associated with any particular string in the text, but rather appear as document-level sets. rule-based ner+l methods, assuming machine learning is not used, are not affected by the fs/zs scenarios. there either is a suitable rule designed for a given situation, or not. if a new code is introduced into the label space, the rules need to be adjusted to reflect this. neural learning approaches are data-driven. the populations of labels available during training and the variety in the inputs to which they are associated affect the model's generalisability, especially if the model is not designed with the few-shot/zeroshot scenario in mind. previous work has tried to address this with setting non-trainable parameters within the network as representations of icd-9 codes enriched with knowledge from the ontologies @xcite . while the fewshot/zero-shot performance improved, the overall performance deteriorated. an alternative to model adjustments is to avoid the fs/zs scenarios by supplying more data, i.e., through data augmentation or synthesis. aug-mentation through synonym replacement has been previously done using wordnet (ollagnier and williams, 2020), with improvements coming from the use of a medical knowledge base (umls) @xcite . simple natural language generation techniques were also employed (ollagnier and williams, 2020). these techniques, while expanding the vocabulary, are only capable of producing synthetic documents with labels present in the original training data. synthesising new documents with alternative labels has been done based on document templates in the scope of radiology reports -however, human experts were involved in the process @xcite . we propose a novel type of data synthesis for icd-9 coding, and medical lmtc tasks in general with the aim to replace concepts of underspecified codes with more specific, and often less frequent, alternatives. similar to @xcite we recognise the value of augmenting concepts of interest. rather than using templating in order to determine concept location, we use pre-existing ner+l techniques (medcat and semehr) to identify spans relevant to the gold standard labelling. furthermore, we introduce an error analysis technique for this setting inspired by confusion matrices. this technique associates codes within the prediction set with codes in the gold standard set according to the ontological structure allowing us to track mispredictions co-occurring with unmatched gold-standard codes indicating confusion -which codes tend to be mispredicted as others. our work provides the following contributions: ## background in this section we will introduce medical ontologies, both as a label set, and source of external knowledge. we will describe named entity recognition and linking used for determining relevant spans of text, introduce lmtc as our task, discuss previous data augmentation techniques in clinical nlp, and finally comment on the current approaches to evaluation and analysis of lmtc models. ## data we employ the discharge summaries of mimic-iii @xcite due to their common use in medical lmtc tasks. mimic-iii is a multimodal medical dataset acquired from the intensive care units of the beth israel deaconess medical center in boston, massachusetts between years 2001 and 2012. access of the data was granted through physionet foot_8 after completing the ethical training by the collaborative institutional training initiative program. the dataset is coded with icd-9 codes on the document level. these labels do not perfectly represent the content of the text -mimic-iii is significantly under-coded for specific conditions @xcite , and sometimes incorrect codes are assigned -e.g., in the case of smoking status @xcite . the data has been pre-processed and split following @xcite 's procedures. the label distribution within the dataset follows a bighead long-tail distribution. we divide the labels, similar to @xcite , into three subsets according to their population size: of the total 8,929 unique labels 4,351 appear in more than 5 documents within the training set; 4,341 at least once but at most 5 times (few-shot); and 237 labels do not appear in the training set, while existing in the development or test set (zero-shot). the training set consists of 47,719 documents. ## methods our methods include data augmentation and synthesis strategies based on the synonyms and adjacent concepts respectively, and an analytic tool for lmtc based on a set of assumptions to adapt confusion matrices with ontological structure. ## experiments we have applied medcat and semehr to the training set, producing candidate spans associated with cuis. post cui-to-icd-9 conversion, we have removed all candidates not matching the gold standard of their source discharge summary. we have produced augmented and synthetic data according to our description in sections 4.1.1 and 4.1.2 and combined them with the original training set (dropping any duplicates) to produce enriched training sets as presented in section 4.1.3. we have further created a baseline-like dataset of a similar size to our largest datasets -semehr-both and medcat-both -as a controlled experiment. this was done by concatenating two baseline datasets (2xbaseline). assuming a constant number of epochs, training on 2xbaseline corresponds to training on the baseline for double the number of epochs. we train caml models based on the implementation of @xcite 11 for 15 epochs on the training sets ( ## results for each experiment we report results averaged across 5 runs ( to the official results of @xcite (micro-f1 of 0.53), due to our limited number of training epochs (while @xcite ceases training after the precision@8 does not improve for 10 epochs). all the proposed methods improve on the baseline with regard to standard and hierarchical micro-f1. augmentation (da) sets, while comparatively worse than synthetic (ds) on r@10-zero and standard and hierarchical micro-f1, perform better on r@10-few. this was to be expected as the da methods provide little for the zero codeset, while producing more of the labels in the freq and few codesets. interestingly, semehr-da performs on par with medcat-da despite having a smaller training set. the combination of da and ds methods (both) report the best f1 results, with medcat-both performing best in 5 our of the 8 reported metrics (including mac-cor, mac-oof and match). both of these methods' results are at least as good as those of 2xbaseline, which is of comparable size. the best r@10-few and r@10-zero performance was achieved by 2xbaseline, which corresponds to training the baseline for twice as many epochs. while the different improvement of da and ds in r@10-few and r@10-zero implies our methods enhance these subsets, 2xbaseline dominating these metrics suggests that a better performance fs/zs can be achieved with more training epochs. the difference between the standard and hierarchical (cophe) f1 scores remained largely the same, which implies partial errors were not addressed by these methods -this is further supported by the changes in mac-oof dominating compared to those of mac-conf. the lowest mac-conf was achieved by the original baseline, but was coupled with a high oof implying that this low confusion is mostly due to a higher proportion of codes not being predicted at all. ## conclusion and discussion the data enrichment methods have improved on the baseline showing potential in approaching the fewshot/zero-shot scenario through data, rather than the model. however, our approach relied on the use of external ner+l tools, whose predictions are imperfect, and may not be available for all domains of interest. other avenues of finding relevant entities, e.g., the attention outputs of lmtc models, should be explored in future work. while the data enrichment results are encouraging, further analysis on fully trained lmtc models is desirable. whcm results point to a major issue with most false negatives coming from underprediction of a family, rather than within-family confusion. further analysis should be conducted on false positives. the analysis from the whcm tool can provide possible explanation of the errors of a model and may shed light on the design of more accurate models for lmtc.