Next Visit Diagnosis Prediction via Medical Code-Centric
Multimodal Contrastive EHR Modelling with Hierarchical Regularisation

Before introducing our novel fusion strategies, we first explain modality-specific encoders that extract features from each modality. We design them as simple as possible to highlight the efficacy of our proposed fusion strategies. In other words, our framework is modular, with the potential for performance enhancement if the encoders are switched to more representative ones.

We employ a simple embedding layer for both medical codes and demographics, and a combination of BioWord2Vec Zhang et al. (2019) and 1D CNN Kim (2014) to process clinical notes. Subsequently, the feature vector is passed to a fully connected layer (Linear) connected with ReLU activation function (Nair and Hinton, 2010).

where $m_{t}$ is a data of modality $m\in(C,H,W)$ at $t$-th visit and $\text{Encoder}_{m}$ is a modality-specialised encoder, passing the feature vector $M_{t}$ to MLP. Finally, a modality-specific feature $\bar{M}_{t}$ is yielded. Appendix A provides a detailed information on how each modality-specific encoder operates.

Cross-Modal Transformer After acquiring representations from all modalities, we entangle them using two cross-modal transformers (CMTs), introduced by MulT Tsai et al. (2019). It has verified its effectiveness in integrating meaningful information across different modalities. Initially, we put the each distinct representation to a temporal non-linear projector, 1D CNN:

where $\bar{M}_{t}$ is a representation from any modality $m$ and $\hat{H}_{t}^{m}$ is a resultant representation. Conv1D is equivalent to 1D CNN. Next, we introduce cross-modal attention, which facilitates the information transfer from the source modality to the target modality, e.g. medical codes $\to$ clinical notes.

Let two modalities as $m_{1}$ and $m_{2}$. Then, using trainable weights $W^{(\cdot)}$ with a dimension of $d_{k}$, we define the query, key and values as $Q^{m_{1}}=H^{m_{1}}W^{Q^{m_{1}}}$, $K^{m_{2}}=H^{m_{2}}W^{K^{m_{2}}}$, and $V^{m_{2}}=H^{m_{2}}W^{V^{m_{2}}}$, respectively. The cross-modal attention, denoted as CA, from $m_{1}$ to $m_{2}$ is then:

We omit $t$ for brevity. CMT is an extension of the CA. It is composed of a multi-head cross-modal attention block (MHA) and a Layer Normalisation layer (LM) Ba et al. (2016). It is computed feed-forwardly for $i=1,\ldots,D$ layers as follows:

During the process at MHA, the representations from the source modality are correlated with the target modality, enhancing the representational power across different modalities. As presented in Fig. 2, the fusion is performed in a medical code-centric fashion, thus we set $m_{1}$ as medical code $C$ and $m_{2}$ as either demographics $H$ or clinical notes $W$. Thus, we acquire two representations of $Z^{C\to H}_{t}$ and $Z^{C\to W}_{t}$ from the two CMTs.

Self-Attention Transformer To extract sequential feature representations effectively and boost dependencies from the above two cross-modal and medical code representations, a self-attention transformer (SA) is employed. It processes across the patient’s single visits:

Additionally, we perform a residual connection He et al. (2016) between the code representation before and after $\text{SA}^{C}$ to enhance the influence of the medical code modality representation.

Multimodal Adaptation Gate Rather than performing a simple concatenation of the three distinct representations, we modify and adopt previous multimodal adaptation gate (MAG) Rahman et al. (2020); Yang and Wu (2021) in the medical code-centric manner. First, we calculate the trimodal gating value $g\in\mathbb{R}$ and the displacement vector H by concatenating meaningful representations in the previous stage as:

This modification maximises the influence of medical code representation during the multimodal fusion process. Then, a weighted summation is performed between the medical code representation $\hat{y}_{t}^{C}$ and the displacement vector H to derive the multimodal representation M:

Here, $\alpha$ is a scaling factor, modulating the influence of the displacement vector H and $\beta$ is a trainable parameter that is randomly initialised. Both $\left\|\hat{y}_{t}^{C}\right\|_{2}$ and $\left\|\text{H}\right\|_{2}$ are the $L_{2}$ norm of their respective entities. Finally, we apply a layer normalisation and dropout to M.

Prediction To predict next visit diagnosis, we feed the representation M in the previous stage into a single linear layer with a Sigmoid activation function to calculate the predicted probability $\hat{y}_{t+1}$.

where cross-entropy loss $\mathcal{L}_{\text{{ce}}}$ is applied as the loss function. $y_{t+1}$ is a ground truth with elements $|\mathbb{C}|$, which takes a value of 1 if the $i$-th code exists in $V_{t+1}$, otherwise 0.

Contrastive learning has been leveraged in multimodal pre-training literature (Radford et al., 2021; Zhang et al., 2022) to align diverse modalities effectively. Inspired by prior works, we apply two bimodal contrastive losses to further intricately entangle the different modalities by anchoring on the medical code representations.

Again, let two distinct modalities of $m_{1}$ and $m_{2}$, where representation vectors derived from each modality be $\hat{H}_{i}^{m_{1}}$ and $\hat{H}_{i}^{m_{2}}$. Given a $i$-th pair of $(\hat{H}_{i}^{m_{1}},\hat{H}_{i}^{m_{2}})$, our bimodal contrastive loss scheme incorporates two asymmetric losses, $m_{1}$-to-$m_{2}$ contrastive loss for the $i$-th pair and its inverse.

where $\langle,\rangle$ is cosine similarity and temperature $\tau\in\mathbb{R}^{+}$ is a parameter modulating distribution’s concentration and Softmax function’s gradient. Subsequently, a bimodal contrastive loss is determined by a weighted combination of $l_{i}^{(m_{1}\to m_{2})}$ and $l_{i}^{(m_{2}\to m_{1})}$ using a weighting parameter $\alpha\in[0,1]$ and averaging over the mini-batch $N$ as:

We apply this to two pairs, one between medical codes and demographics, and the other between medical codes and clinical notes.

Note that, our multimodal contrastive loss is applied inter-modally, in line with the CLIP Radford et al. (2021), rather than intra-modally. Moreover, we consider at the patient level rather than at the visit level. This is because patient level representations share similar patterns between their visits.

We conduct experiments on a publicly available large-scale, deidentified real-world EHR data, MIMIC-III (Johnson et al., 2016). It is acquired from intensive care units (ICU) patients at Beth Israel Deaconess Medical Center between 2001 and 2012. It contains multifaceted data, including ICD-9 medical codes, demographics, clinical notes, and so on. We provide descriptions on data pre-processing and the corresponding statistics to Appendix B.

We describe the details for implementation. First, we set 256 and 0.1 as a hidden dimension and a dropout rate across the entirety of the model (e.g. medical code and demographics feature extraction modules, Transformers including CMT and SA, and MAG), respectively. In the clinical note extraction module, filter sizes are set to [2, 3, 4], and the hidden dimension is 512. For the CMTs and SAs, we set the number of heads and encoder layers to be 4 and 3, respectively.

Also, following the previous work (Radford et al., 2021), the temperature $\tau$ and alpha $\alpha$ are 0.1 and 0.25 for the contrastive loss. The coefficients of loss terms, $\lambda_{\text{ce}}$, $\lambda_{\text{bi-con}}$, and $\lambda_{\text{hrchy}}$ are set to 1, 1, and 0.1, respectively. Especially, the $\lambda_{\text{hrchy}}$ is set relatively small to weakly regularise each modality-specific encoder to learn the parental levels of ICD-9 codes, without overly constraining them. We provide the experimental results on the different $\lambda_{\text{hrchy}}$ to Appendix C.

We train models using Adam optimiser (Kingma and Ba, 2014) with a constant learning rate of 1e-4 and mini-batch size of 4, for a maximum of 50 epochs. The training is stopped if there is no gain for consecutive 5 epochs on validation data. Also, following the previous work (Choi et al., 2017), our proposed framework is evaluated using top-$k$ accuracy, ranging $k$ from 5, 10, 20 to 30. This is consistent with how physicians consider a comprehensive set of potential diagnoses, and is suitable for multi-label classification scenarios where multiple diseases often co-occur. Details on other baselines are provided to Appendix D.

Our proposed framework is implemented using PyTorch (Paszke et al., 2019) and accelerated via a single NVIDIA GeForce RTX 3090 GPU.

Table 1 provides quantitative results of the proposed NECHO in comparison to the baselines on the MIMIC-III data for the diagnosis prediction task. NECHO notably excels over all existing baselines in EHR modelling and multimodal fusion strategies. Its effectiveness is attributed to its ability to leverage unique and complementary information from other modalities, which especially improves top-30 accuracy ranging from 0.5% to 10.7% over modality-specific encoders that constitute NECHO.

As shown in Table 1, the multimodal fusion is imperative. It’s noteworthy that whilst MAIN An et al. (2021) employs a trimodal representation learning, its performance falls short compared to the bimodal MNN Qiao et al. (2019). This discrepancy might arise from the harmful effects of improperly fusing demographic data lately. Especially, bimodal MNN shows comparable performance to trimodal fusion strategies baselines. This confirms the limitations of the tertiary symmetric multimodal fusion methodologies and raises the need for a medical code-centric approach, taking into account the modality imbalance.

To validate the efficacy of our fusion strategy, we compare NECHO that excludes the hierarchical regularisation (NECHO ${}_{\text{w/o}\ \mathcal{L}_{\text{hrchy}}}$) amongst multimodal EHR modelling and fusion baselines. Our method demonstrates superior performance over them, including NECHO ${}_{\text{w/o code centring}}$. These findings highlight the significance of designing multimodal fusion framework by centring medical codes representation that ensures a seamless aggregation of diverse data modalities. Furthermore, we also provide a comparative study on our novel code-centric MAG with others Rahman et al. (2020); Yang and Wu (2021) to Appendix E.

Next, we delve into the significance of regularising modality-specific encoders using parental level of medical codes. We juxtapose NECHO with NECHO ${}_{\text{w/o}\ \mathcal{L}_{\text{hrchy}}}$ and ULGM Yu et al. (2021), at which modality-specific encoders learn the same level of medical ontology as the final prediction. They two show inferior performance, emphasising the importance of our novel strategy. It is discussed further in Ablation Studies (Section 4.2.2).

Furthermore, whilst NECHO does not completely surpass MIPO, replacing its simple medical code encoder with MIPO (NECHO ${}_{\text{w/ MIPO}}$) outperforms MIPO. It especially achieves a 1.01% increase in top-30 accuracy, indicating that 1) our framework is modular, and 2) NECHO can predict additional accurate diseases than MIPO by leveraging complementary information from various modalities, emphasising its significance in real clinical settings. We provide a regarding case study (Section 4.2.3).

Another noteworthy point outside the multimodal strategies is that, amongst the clinical note baselines, Clinical BERT (Alsentzer et al., 2019) that is trained with a maximum of 512 tokens surpasses the combination model of BioWord2Vec (Zhang et al., 2019) and 1D CNN (Kim, 2014) with equivalent number of tokens but is inferior to that model trained with 10k tokens. This suggests that enhancing performance is more about processing a large number of tokens than increasing model complexity in EHR learning. This also justifies our preference for BioWord2Vec over Clinical BERT within the realm of Pretrained Language Models.

We conduct ablation studies to discern influence of each module on the overall performance as: 1) individual modalities, 2) the multimodal fusion strategies (including Transformers, MAG, and bimodal contrastive losses), and 3) the hierarchical regularisation. The results are reported in Table 2.

Firstly, we assess the contribution of each modality within our proposed framework. The results demonstrate a clear superiority of the trimodal approach over its unimodal and bimodal ones. This underscores the unique representations from each modality are complementary to one another. Also, the significant performance degradation is observed upon the exclusion of medical code representation (w/o code), highlighting its pivotal role and rationalising our medical code-centred strategy. Additionally, whilst the exclusion of either notes or demographics similarly harms the performance, the note contains more meaningful information necessary than demographics, as shown in Table 1.

Secondly, we evaluate the impact of our medical code-centred strategies by removing each component. The resultant performance decline highlights their importance. Intriguingly, the performance disparities between models lacking transformers (w/o Transformers), lacking MAG (w/o MAG), and the full model (NECHO) widen as the value of $k$ increases, suggesting an amplified effect in scenarios involving a broader range of disease sampling. Conversely, the influence of contrastive losses (w/o $\mathcal{L}_{\text{bi-con}}$) remains relatively stable across different top-$k$ accuracies, indicating that they effectively align the distinct modalities in a semantically consistent fashion. These observations show that the adaptation of the proposed modules simultaneously is essential for effective inter-modality interaction and integration, thereby yielding significant performance enhancements.

Finally, the effectiveness of our novel parental level hierarchical regularisation is investigated. Its omission (w/o $\mathcal{L}_{\text{hrchy}}$) affects adversely model’s accuracy across various top-$k$ accuracies. This suggests that enforcing the encoders for three distinct modalities, guided by the parental levels of medical codes using an ICD-9 hierarchy, is essential for enhancing performance as it injects the general information and thus prevents the possible transmission of erroneous information when combining representations from distinct data modalities, thereby encouraging effective and accurate training.

To qualitatively evaluate the predictive performance between MIPO (Peng et al., 2021) and our NECHO, we present a case study (Table 3) using a patient whose medical history shows a progression from a mitral valve issue to complications after surgery and cardiac rhythm disturbances. In the study, codes are formatted according to the Clinical Classifications Software (CCS) and are sequenced based on their priority, significantly influencing the reimbursement for treatment. We prefix them with "D" to make them appear akin to diagnosis codes.

Notably, our NECHO model accurately predicts 6 out of the top-10 diagnosis, outperforming MIPO, which predicts only 3. Firstly, both successfully identify D53 (Disorders of lipid metabolism), D106 (Cardiac dysrhythmias) and D101 (Coronary atherosclerosis and other heart disease), likely due to these diagnoses being part of the patient’s prior medical codes. However, NECHO uniquely predicts D238 (Complications of surgical procedures or medical care), D49 (Diabetes mellitus without complication), D2616 (E Codes: Adverse effects of medical care) and D96 (heart valve disorder) which MIPO fails to identify.

Additionally, our model predicts D238 and D2616 using multifaceted information of both demographics and notes. D238 should be predicted for two points: 1) the patient was initially hospitalised due to emergency health problem according to demographics, and 2) his notes states visual hallucinations, monitoring for pericardial and pleural effusions. The prediction of D2616 aligns with potential risks associated with mitral valve replacement. On the contrary, MIPO’s prediction of D259 (Residual codes; unclassified) and D131 (Respiratory failure; insufficiency; arrest (adult)), which is considered less informative and a simple repetition from previous patient visits. D2 (Septicemia) and D3 (Bacterial infection) are not explicitly mentioned in the patient’s history, thus extremely challenging to predict. Hence, this demonstrates the necessity of the effective multimodal fusion strategy for its capability of capturing complementary and unique information in other modalities, verifying the effectiveness of the NECHO.

Apart from multimodal EHR learning, the content following the "Impression" in the preceding notes is only explicitly found in radiology reports. This indicates the importance of considering all available clinical note types to acquire a thorough understanding of a patient’s information. This contrasts with previous findings (Hsu et al., 2020; Husmann et al., 2022) suggesting that certain specific note types are representative in EHR learning.