BoxLM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare

We thank the reviewer for the overall positive evaluations and many detailed suggestions. In the following, we focus on several of the main issues to provide our feedback:

Reviewer yPUD.W1: The difference between BoxLM and BoxCare.

As described in Section 2, BoxCare only uses box embeddings to represent ontology hierarchies without fully integrating the rich semantics and hierarchical structures inherent in EHRs. In contrast, BoxLM innovatively integrates PLM with box embeddings, unifying ontology-driven and EHR-driven hierarchies with LM-based semantics for more accurate diagnosis prediction. We will clarify the difference in the introduction section.

Reviewer yPUD.W2: Explicitly stating the specific task and research objectives upfront after the introduction for better reader comprehension.

We have provided the problem definition and an overview of BoxLM in Section 3.1. To further improve readability, we will reorganize the paper structure by moving the methodology section to Section 2, immediately following the introduction.

Reviewer yPUD.W3: Box embedding concepts and operations are relatively complex.

Box embeddings only introduce one additional offset embedding compared with traditional vector representation, and the time complexity for diagnosis prediction remains comparable to that of traditional vector-based approaches. Specifically, the time complexity is $\mathcal{O}(N \times C \times d)$, scaling linearly with the embedding dimension $d$, where $N$ and $C$ denote the number of patients and CCS codes, respectively.

Reviewer yPUD.W4: Adding hyperparameter studies.

We have evaluated the impacts of the box embedding dimension $d$ (Table 1), the Gumbel distribution scale $\beta$ (Table 2), and the box volume calculation (Table 3), which will be included in the final version.

For box embedding dimension $d$, the results show that increasing $d$ slowly enhances model performance for both datasets. For a fair comparison with baselines, we set $d=16$.

Table 1: Experimental results (%) for diagnosis prediction with varying box dimension $d$.

For the Gumbel distribution scale $\beta$, a larger $\beta$ makes the distribution closer to uniform, while a smaller $\beta$ causes its probability density function to approach a hinge function, leading the random variable to degenerate into a constant [1]. In this paper, we set $\beta=0.2$, balancing model accuracy with the distinctiveness among boxes.

Table 2: Experimental results (%) for diagnosis prediction with varying Gumbel distribution scale $\beta$.

For the box volume calculation, we compare the soft volume calculation [2] with our used Bessel volume calculation (i.e., Eq. (9)). They both effectively mitigate the training difficulties that arise when disjoint boxes should overlap.

Table 3: Experimental results (%) for diagnosis prediction with varying box volume calculation.

Reference
[1] When Box Meets Graph Neural Network in Tag-aware Recommendation, KDD, 2024.
[2] Smoothing the geometry of probabilistic box embeddings, ICLR, 2018.

Reviewer yPUD.W5: Providing examples for the interpretability of the proposed method.

We have provided a real case study in Section 4.3, visualizing different levels of medical concepts (e.g., diagnosis, CCS, patient). As shown in Lines 403-420, BoxLM accurately predicts the development of coronary atherosclerosis and other heart diseases (CCS code 101) for the patient diagnosed with essential hypertension by calculating the overlap between patient and CCS boxes, which aligns with clinical knowledge that hypertension increases the risk of atherosclerosis.

We thank the reviewer for the overall summary. As for the several weaknesses and questions you mentioned, our responses are listed as follows:

Reviewer UQNn.W1: The evaluation metric suggests including recall.

Following [2-6], we adopt visit-level Precision@k (P@k) and code-level Accuracy@k (Acc@k) for the comprehensive evaluation of diagnosis prediction from coarse-grained and fine-grained perspectives, respectively (detailed in Appendix D). Additionally, we have included Recall@k (k=10, 20) in our evaluation (Table 1), further demonstrating the effectiveness of BoxLM.

Table 1: Recall@k results (%) for diagnosis prediction on both datasets with 5% training data.

Reference
[2] Unveiling Discrete Clues: Superior Healthcare Predictions for Rare Diseases, WWW, 2025.
[3] Stage-Aware Hierarchical Attentive Relational Network for Diagnosis Prediction, TKDE, 2024.
[4] Predictive Modeling with Temporal Graphical Representation on Electronic Health Records, IJCAI, 2024.
[5] MEGACare: Knowledge-guided Multi-view Hypergraph Predictive Framework for Healthcare, Information Fusion, 2023.
[6] INPREM: An Interpretable and Trustworthy Predictive Model for Healthcare, KDD, 2020.

Reviewer UQNn.W2: Missing baseline Doctor AI [1].

We have included Doctor AI [1] in our benchmarks for comparison. As shown in Table 2, BoxLM outperforms Doctor AI by up to 40.30% gains in Recall@10 on MIMIC-IV. Notably, both Doctor AI and RETAIN [7], developed by the same research group, use RNNs to capture the temporal dynamics of patient visits. Doctor AI slightly surpasses RETAIN with short visit sequences on MIMIC-III, while RETAIN performs better on MIMI-IV via its reverse time attention mechanism. In contrast, BoxLM integrates temporal modeling with richer representations of medical concepts from EHR data and ontological sources, leading to superior generalization. We will add this discussion to the final version.

Table 2: Experimental results (%) for diagnosis prediction on both datasets with 5% training data.

Reference
[7] RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism, NeurIPS, 2016.

Reviewer UQNn.Q1: In Section 4.3's ablation studies, the contribution of BoxLM should be evaluated, including the comparison with the classic vector representation.

Our ablation studies in Section 4.3 have included the comparison with classic vector representations (i.e., the "Base" model in Table 3). Specifically, the "Base" model uses BioBERT embeddings without structure modeling, which represents visits as single points. By comparison, "+ Onto" and "+ EHR" introduce box embeddings for CCS codes, diagnoses, and visits with our ontology-based and EHR-based hierarchy modeling, achieving significant improvement by up to 25.33% gains in P@10 on MIMIC-IV. These experimental results provide a comprehensive understanding of the importance and reliability of our designed techniques.

Reviewer UQNn.Q2: What's the difference between CCS and diagnosis?

Sorry for the confusion. CCS is a classification system that groups diagnosis codes into clinically meaningful categories for analytical purposes, whereas a diagnosis refers to the identification of a specific disease or condition in a patient's visit. We will add this explanation in the final version.

We thank the reviewer for the constructive comments on our paper. In response to the shortcomings mentioned, we provide the following answers to several questions:

Reviewer skUY.W1: Adding quantitative metrics for interpretability.

In our study, CCS (Clinical Classification Software) codes from the MIMIC-IV dataset can be grouped into 22 distinct body systems, serving as hierarchical classification labels for the predicted CCS codes. To assess the interpretability of BoxLM, we have introduced the Consistency@k metric (k=10, 20) [1, 2], which quantifies the alignment between model predictions and the underlying hierarchical structure.

The metric can be computed as: $Consistency@k=\frac{1}{\min(k,|H_v|)}\sum_{i=1}^k\mathbb{I}(\hat{h}_i=h_i)$, where $|H_v|$ is the number of ground-truth hierarchical labels associated with the CCS codes in visit $v$, and the numerator counts the number of correct hierarchical predictions in the top-k.

The results (Table 1) demonstrate that BoxLM preserves consistency on hierarchy.

Table 1: Consistency@k results (%) for diagnosis prediction.

Reference
[1] Hyperbolic Embedding Inference for Structured Multi-Label Prediction, NeurIPS, 2022.
[2] Modeling Label Space Interactions in Multi-Label Classification Using Box Embeddings, ICLR, 2021.

Reviewer skUY.W2: Lacking hardware details and code availability.

We have provided the code link in the bottom left of Page 6 for reproducibility. Additionally, we have added the following hardware details: All experiments were conducted on two NVIDIA GTX 3090 Ti GPUs.

Reviewer skUY.W3: Adding related works on recent geometric approaches.

We have included relevant work on box embeddings from other fields (see Appendix A). Additionally, we will add the discussion and proper citations for [3, 4] in our related work, covering polygon and hypercube embeddings in areas such as knowledge graphs and recommendation systems.

Reference
[3] ExpressivE: A Spatio-Functional Embedding for Knowledge Graph Completion, ICLR, 2023.
[4] Thinking inside The Box: Learning Hypercube Representations for Group Recommendation, SIGIR, 2022.

Reviewer skUY.W4: The saying for introducing box embedding might not be proper, as it's not straightforward to encode highdimensional embeddings to hyperrectangle.

We understand that the sentence "... representing entities as high-dimensional hyperrectangles" (Lines 67-68) may not have clearly conveyed our design. Specifically, we propose using box embeddings to capture complex relationships (e.g., inclusion and intersection) in medical data. This representation aligns well with the hierarchical structure of medical ontologies and EHR-driven visit patterns. We will revise this sentence in the final version.

Reviewer skUY.W5: Some typos exist.

Thanks for pointing this out. We will fix it in the final version.

Reviewer skUY.W6: Section 3.3.1 needs a diagram for patient box evolution.

We have included a diagram for the evolve-and-memorize patient box learning mechanism in the upper right of Figure 2. For clarity, we will explicitly reference it in Section 3.3.1.

Reviewer skUY.Q1: How does BoxLM handle missing or incomplete hierarchical links in EHRs?

BoxLM uses box embeddings to model medical concepts from ontology and EHR data. The geometric properties of these embeddings allow the model to infer implicit relational structures, such as transitivity and entailment. This capability, demonstrated in tasks like taxonomy expansion [5] and knowledge base completion [6], enables BoxLM to recover missing or incomplete hierarchical links.

Reference
[5] A Single Vector Is Not Enough: Taxonomy Expansion via Box Embeddings, WWW, 2023.
[6] BoxE: A Box Embedding Model for Knowledge Base Completion, NeurIPS, 2020.

Reviewer skUY.Q2: What is the impact of box dimensionality on performance and computational cost?

We have added experiments to analyze the impact of box embedding dimension $d$ (please refer to Table 1 in Reviewer yPUD.W4), which will be included in the final version. The results show that increasing $d$ slowly enhances model performance for both datasets. Notably, the time complexity of diagnosis prediction scales as $\mathcal{O}(N \times C \times d)$ ($N$ and $C$ denote the number of patients and CCS codes), which grows linearly with $d$ and is similar to traditional vector-based approaches. For a fair comparison with baselines, we set $d=16$.