Exploiting hierarchy in medical concept embedding

Objective

To construct and publicly release a set of medical concept embeddings for codes following the ICD-10 coding standard which explicitly incorporate hierarchical information from medical codes into the embedding formulation.

Materials and Methods

We trained concept embeddings using several new extensions to the Word2Vec algorithm using a dataset of approximately 600,000 patients from a major integrated healthcare organization in the Mid-Atlantic US. Our concept embeddings included additional entities to account for the medical categories assigned to codes by the Clinical Classification Software Revised (CCSR) dataset. We compare these results to sets of publicly-released pretrained embeddings and alternative training methodologies.

Results

We found that Word2Vec models which included hierarchical data outperformed ordinary Word2Vec alternatives on tasks which compared naïve clusters to canonical ones provided by CCSR. Our Skip-Gram model with both codes and categories achieved 61.4% Normalized Mutual Information with canonical labels in comparison to 57.5% with traditional Skip-Gram. In models operating on two different outcomes we found that including hierarchical embedding data improved classification performance 96.2% of the time. When controlling for all other variables, we found that co-training embeddings improved classification performance 66.7% of the time. We found that all models outperformed our competitive benchmarks.

Discussion

We found significant evidence that our proposed algorithms can express the hierarchical structure of medical codes more fully than ordinary Word2Vec models, and that this improvement carries forward into classification tasks. As part of this publication, we have released several sets of pretrained medical concept embeddings using the ICD-10 standard which significantly outperform other well-known pretrained vectors on our tested outcomes.

Methods
This dataset includes trained medical concept embeddings for 5428 ICD-10 codes and 394 Clinical Classification Software (Revised) (CCSR) categories.  We include several different sets of concept embeddings, each trained using a slightly different set of hyperparameters and algorithms.

To train our models, we employed data from the Kaiser Permanente Mid-Atlantic States (KPMAS) medical system.  KPMAS is an integrated medical system serving approximately 780,000 members in Maryland, Virginia, and the District of Columbia.  KPMAS has a comprehensive Electronic Medical Record system which includes data from all patient interactions with primary or specialty caregivers, from which all data is derived. Our embeddings training set included diagnoses allocated to all adult patients in calendar year 2019.

For each code, we also recovered an associated category, as assigned by the Clinical Classification Software (Revised).

We trained 12 sets of embeddings using classical Word2Vec models with settings differing across three parameters.  Our first parameter was the selection of training algorithm, where we trained both CBOW and SG models.  Each model was trained using dimension k of 10, 50, and 100.  Furthermore, each model-dimension combination was trained with categories and codes trained separately and together (referred to hereafter as ‘co-trained embeddings’ or ‘co-embeddings’).  Each model was trained for 10 iterations.  We employed an arbitrarily large context window (100), since all codes necessarily occurred within a short period (1 year).

We also trained a set of validation embeddings only on ICD-10 codes using the Med2Vec architecture as a comparison.  We trained the Med2Vec model on our data using its default settings, including the default vector size (200) and a training regime of 10 epochs.  We grouped all codes occurring on the same calendar date as Med2Vec ‘visits.’  Our Med2Vec model benchmark did not include categorical entities or other novel innovations.

Word2Vec embeddings were generated using the GenSim package in Python.  Med2Vec embeddings were generated using the Med2Vec code published by Choi.  The JSON files used in this repository were generated using the JSON package in Python.