BoxLM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Authors: Yanchao Tan
Hang Lv
Yunfei Zhan
Guofang Ma
Bo Xiong
Carl Yang
Venue: ICML 2025 poster (2025)
Keywords: Diagnosis prediction
Box embeddings
Language model
Hierarchy modeling
Primary area: applications->health_medicine

Abstract: Language Models (LMs) have advanced diagnosis prediction by leveraging the semantic understanding of medical concepts in Electronic Health Records (EHRs). Despite these advancements, existing LM-based methods often fail to capture the structures of medical concepts (e.g., hierarchy structure from domain knowledge). In this paper, we propose BoxLM, a novel framework that unifies the structures and semantics of medical concepts for diagnosis prediction. Specifically, we propose a structure-semantic fusion mechanism via box embeddings, which integrates both ontology-driven and EHR-driven hierarchical structures with LM-based semantic embeddings, enabling interpretable medical concept representations.
Furthermore, in the box-aware diagnosis prediction module, an evolve-and-memorize patient box learning mechanism is proposed to model the temporal dynamics of patient visits, and a volume-based similarity measurement is proposed to enable accurate diagnosis prediction. Extensive experiments demonstrate that BoxLM consistently outperforms state-of-the-art baselines, especially achieving strong performance in few-shot learning scenarios, showcasing its practical utility in real-world clinical settings.

BoxLM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Yanchao Tan 123 Hang Ly 123
Yunfei Zhan 123
Guofang Ma
Bo Xiong
Carl Yang?
Abstract
Language Models (LMs) have advanced diagno- sis prediction by leveraging the semantic under- standing of medical concepts Electronic Health Records (EHRs). Despite these advancements, existing LM- methods often fail to capture the structures of medical concepts (e.g., hierar- chy structure from domain knowledge). In this paper, we propose BoxLM, a novel framework that unifies the structures and semantics of med- ical concepts for diagnosis prediction. Specif- ically, we propose a strmetuue-semantic fusion mechanism via box embeddings, which integrates both ontology-driven and EHR-driven hierarchi- cal structures with LM- AM--aaeed semantic embed- dings, enabling interpretable medical concept rep- resentations. Furthermore, in the box-aware diag- nosis prediction module, an sedlve-and-memorrie patient box learning mechanism is proposed model the temporal dynamics of patient visits, and volume-based similarity measurement is proposed to enable accurate diagnosis predic- tion. Extensive experiments demonstrate that BoxLM consistently outperforms state-of-the- art baselines, especially achieving strong perfor- mance in few-shot learning scenarios, showcasing its practical utility in real-world clinical settings.
Model A
Model
Semanties
Text Common Cold:
Unified Model
Text
Hierarchy
Allergic Rhinitis:
Hierarchy
Common Cold: Allergic Rhinitis;
C Ontology
100-06
130-39
Visit Sequence
EHR
E
Semantics
Structures
Struchures
(n)
(b)
(c)
Figure I. Exploring semantics and structures in healthcare. (a) modeling semantical similarity with vector embeddings. (b) en- coding structural relationships with vector embeddings. (c) inte- grating semantic and structural dimensions with box embeddings.
1. Introduction
Accurate diagnosis prediction based on Electronic Health Records (EHRs) has emerged pivotal task in modern healthcare, enabling personalized treatment planning and improved patient outcomes. However, the widespread appli- cation of EHR- based systems is hindered by concerns over data privacy and limited accessibility to patient-specific datasets, resulting in the few-shot learning problem (Xu et al., 2024). This problem necessitates innovative ap- proaches with limited patient EHR data while maintain- ing predictive performance. Recent advancements in Lan- guage Models (LMs) have demonstrated their potential to address the few-shot learning problem by incorporating the semantic understanding of medical concepts. I.Ms, such as BioBERT (Lee et al., 2020) and PubmedBERT (Gu et al., 2021), introduce pre-trained biomedical embeddings. Moreover, cutting-edge models like MedBench (Cai et al., 2024) and BioMistral (Labrak al., 2024) further refine this approach through large-scale training medical corpora, achieving notable improvements in semantic representation and contextual understanding of limited medical concepts.
'Engineering Research Center of Big Data Intelligence, istry Education, Fuzhou, China FFujian Key Laboratory of Net- work Computing and Intelligent Information Processing, Fuzhou University, Fuzhou, China Coollege of Computer and Data Sci- ence, Fuzhou University, Fuzhou, China *School of Computer Sci- ence, Zhejiang Gongshang University, Zhejiang, China ZZeejiang Key Laboratory of Big Data and Future E-Commerce Technology, Zhejiang, China CCetter for Biomedical Informatics Research, Stanford University, Palo Alto, USA 'Department of Computer Science, Emory University, Atlanta, USA. Correspondence to: Guofang Ma <maguofang zigsu.edu.cn..
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Despite these advancements, existing LM- based methods often fail capture the structures of medical concepts, such

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
as the hierarchy of medical ontologies and the hierarchy of EHR-driven visit patterns. These structures are crucial for accurate diagnosis prediction (Luo et al., 2020; Lv et al., 2024). For instance, consider a patient with history of common cold and chromic simusitis. Based solely on se- mantics, LM-based models might predict the next diagnosis as common cold due to symptom similarity and its higher frequency in the training data. By unifying semanties with structures, the model can leverage hierarchical relationships to improve accuracy. In the ontology hierarchy, allergic rhinitis and chronic sinusitis are both classified under other diseases of the upper respiratory tract. Moreover, chronic sinusitis is often a downstream complication of poorly man- aged allergic rhinitis in the EHR hierarchy. Recognizing these relationships, the model can more accurately predict allergic rhinitis as the next diagnosis, illustrating the impor- tance of incorporating structural knowledge.
prediction accuracy, outperforming state-of-the-art models.
In summary, we make the following contributions:
Unified BoxIM Framework. We propose a novel frame- work that unifies structures and semantics of medical concepts via box embeddings, capturing the complex relationships among diagnoses, CCS codes, and patient visits at the conceptual level.
Effective Model Designs. We introduce a structure- semantic fusion mechanism to jointly integrate ontology-driven and EHR-driven hierarchical struc- tures with PLM-based semantic embeddings for medi- cal concept representation. Furthermore, propose an evolve-and-memorize patient box learning mech- anism and a volume- based diagnosis prediction, to quantify complex relationships between the patient and CCS codes.
However, effectively modeling these complex structures within the framework of LM-based diagnosis prediction remains a significant challenge. Standard embedding tech- niques, as shown in Figure 1 (a) and Figure 1(b), represent concepts as single points in a vector space. These vector embeddings, while effective at capturing similarity relation- ships, fail to encode complex relationships such inclusion relations inherent in the hierarchy. For instance, vector em- beddings can model the similarity between hypertension and cardiovascular diseases by encoding semantic mean- ing. Whereas, they do not explicitly capture the hierarchical relationship where cardiovascular diseases encompasses hypertension. Inspired by the success of box embeddings across various domains (Vilnis et al., 2018; Huang et al., 2023; Lv et al., 2024; Lin et al., 2024), we find that they offer a promising solution by representing entities as high- dimensional hyperrectangles, where two boxes can clearly "enclose" or "intersect" each other. As shown in Figure (c)), box embeddings are well-suited to capturing complex relationships, including hierarchical inclusion and set-based intersections, making them ideal choice for unifying structures and semantics in medical concepts.
Extensive Experiments on Real EHR Datasets. We con- duct extensive experiments on real-world EHR datasets, demonstrating that BoxIM outperforms state-of-the- art models in both visit-level and code-level diagnosis prediction. Particularly, BoxLM achieves strong perfor- mance in few-shot learning scenarios, highlighting its practical applicability in real-world clinical settings.
2. The BoxLM Framework
2.1. Problem Definition and BoxLM Overview
Given a patient's EHR data, which consists of a sequence of visits, each associated with set of medical concepts (e.g.. diagnoses and CCS codes), the goal is predict the relevant CCS codes for the patient's next visit. Notably, CCS is classification system that groups diagnoses 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.
Let a patient p's EHR data be represented as a sequence of visits (U1, u2..... w where each visit ee is associated with set of medical concepts, including diagnoses d E Ma and CCS codes E MMe. The hierarchical relationships among these concepts are derived from two sources: (1) Ontology-driver hierarchy graph 9m = (Vi.: Em,) where Vm denotes the set of medical concepts (i.e., diagnoses and CCS codes) and e, represents the parent-child relationships. With Gm. we obtain diagnosis box b and CCS code box b (2) EHR-driven hierarchy graph G. = (V.,, E.). where V, includes all visits and their associated medical concepts, and &, captures the relationships between visits and their constituent concepts. With Go, we obtain visit box b
In this paper, we propose BoxIM, a novel Box-aware Language Model for diagnosis prediction. BoxLM lever- ages the expressive power of box embeddings to seamlessly unify structural and semantic representations of medical con- cepts, enabling precise diagnosis prediction. Specifically, our framework consists of two key modules: (1) structure- semantic fusion module that integrates both ontology-driven and EdR-driven hierarchical structures with semantic em- beddings from Pre-trained Language Models (PLMs), and (2) a box-aware diagnosis prediction module that quantifics the similarity between patient visits and Clinical Classifica- tions Software (CCS) codes using the intersection volume of their boxes. By modeling medical concepts as high- dimensional hyperrectangles, BoxLM significantly enhances
The patient's overall representation is then modeled as a box embedding b, by aggregating the sequence of visits.
2

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Ontology-driven Hierarchy Modeling
(c) Evolve-and-Memorize Patient Modeling
ed
Diagnosis names
Center
Visit sequence
BioBERT
MLP
of ba
Diagnosis box
ATg
N(C)
Offset
ba
Cen
b bur
Hp
Jog
Center
Patient box
Ontology hierarchy
AGG
of b
of
by Cen
fm
(d)
Offset of ba
CCS box
Max
Offset
of bp
P
AGG
b Cen
b bof
(d) Volume-based Diagnosis Prediction
ec
CCS code
Center
Intersection
Volume
Similarity score
names
BioBERT
MLP
of b
Juen
Vol(b pnc)
8p.c
(b) EHR-driven Hierarchy Modeling
M
ai
Center of bp:
Visit box
Ppni
7
EHR hierarchy Gv
by byy
Gumbel
Diagnosis Prediction
Offsct
of ofbb,
ppnc
Max
Structure-Semantic Fusion via Box Embeddings
Box-aware Diagnosis Prediction
Figure The overall framework our proposed BoxLM.
The prediction is based on the intersection volume- similarity score Up.c between the patient box b, and the boxes of candidate CCS codes b..
during fusion and used to encode hierarchical structures.
box is defined as: b, Cen bD*)), where b,en E Rdd and b idim represent the center and offset of the box, respectively.
Definition 1 (Box Embedding).
We summarize the main modules of the BoxIM framework in Figure 2 to provide overview. BoxIM consists of two components. The first component, Peructure-Semantie Fu- sion via Box Embeddings, integrates semantic embeddings from PLMs with structural relationships derived from both ontology and EHR data. The second component, Box-aware Diagnosis Prediction, models the temporal progression of patient visits and predicts future CCS codes based the intersection volume between patient and CCS code boxes.
To preserve the semantic information and adapt it for clinical tasks, we compute the box center using PLM- seman- tic embeddings (shown in Figure 2(a)). Specifically, we re- trieve the e, and e, embeddings for medical concepts from BioBERT (Lee et al., 2020) by using their clinical names (e.g., disease names and CCS names) indices. These em- beddings then passed through a learnable Multi-Layer Perceptron (MLP) to reduce dimensionality and obtain the box center b Cen as follows: (d.e)
2.2. Structure-Semantic Fusion via Box Embeddings
To effectively capture both semantic and structural rela- tionships in medical concepts, we introduce box-aware language model that models entities as high-dimensional hy- perrectangles. This method integrates semantic embeddings from PLMs with structural relationships, directly producing interpretable box representations the conceptual level. We now proceed describe how BoxLM designed capture hierarchical structures from both ontology and EHR data.
Cen MLP(Ce), h MIP((e.) b
(1)
While PLMs effectively capture the semantic meaning of medical concepts through box centers, they do not inher- ently encode the parent-child relationships found in medical ontologies. To address this, we propose to leverage clinical knowledge to calculate offset, so as to represent the relative position of concepts within hierarchy. For instance, the offset for parent concept, such as cardiovascular disease, should be larger than that of its child concept, such as hyper- tension, reflecting the broader nature of the parent concept. This approach ensures that hierarchical relationships are explicitly captured within the box embedding.
2.2. 1. ONTOLOGY-DRIVEN HIERARCHY MODELING
Box embeddings model medical concepts as high- dimensional hyperrectangles, where the center and offset of each box explicitly represent the semantic meaning and hierarchical relationships of these concepts, respectively. this part, we detail how these box embeddings are computed
To encode hierarchical relationships into offset embeddings, we propose a directed graph- based convolution mechanism
3

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
operating over a hierarchical graph. Specifically, in the ontology-driven hierarchy graph 9... (Vm; ), each medical concept E Vm (e.g., cardiovascular disease or CCS code) is initialized with their center embedding b Con (d.a)' and we adopt relation-aware graph convolutional networks to compute their offset embedding b bfff (d.c)' which jointly embeds both nodes and relations directed graph. The offset embedding is computed as:
by the diversity of its associated medical concepts. A visit with a narrow range of diagnoses and CCS codes will have smaller offset, while visit with diverse or conflicting medical concepts will have a larger offset. The offset is computed as the maximum of the offsets of its associated medical concepts:
by max off i M,,
(5)
bfff= AGG ((o((, Cen W.) (n,r) E (dd))) be AGG (t6(b) Cen, W.) (d,r) E ((c)))
(2)
where 6 is the offiet for each medical concept associ- ated with the visit i/ and max(-) is an element-wise maxi- mum operation.
where AGG is an aggregation function (e.g.. summation or mean) that combines the information from neighboring concepts. (d) represents the set of neighboring concepts of diagnosis d, and (c) denotes the set of diagnosis con- cepts associated with the CCS code c. ob;, Wr) is compo- sition function that combines neighboring concept with the relation W, E Rdimxdimm is the learnable weight, and r encompasses two types of directed edges (i.e., parent-to- child and child-to-pared -areeettt relationships).
2.3. Box-aware Diagnosis Prediction
In clinical practice, a patient's medical history is repre- sented sequence of visits over time, where each visit contains multiple medical concepts such as diagnoses and CCS codes. As a patient's condition evolves, visits often exhibit overlaps in medical concepts, reflecting recurring conditions or related diagnoses. Capturing these overlaps while maintaining an awareness of the temporal progression is essential for accurate diagnosis prediction. Traditional point- embeddings struggle to model such complex re- lationships effectively, as they fail to represent the inherent hierarchical structure and semantic relationships of medi- cal concepts. This motivates the use of box embeddings, which naturally model overlaps, hierarchies, and seman- tic relationships through high-dimensional hyperrectangles, providing more expressive and interpretable representation for diagnosis prediction.
2.2.2. -RIVED HIERARCHY MODELING
In EHR data, each patient visit is represented as a set of med- ical concepts, including diagnoses and CCS codes. The ag- gregation of diagnoses and CCS codes into a visit naturally forms a hierarchical structure G.. where the visit serves parent node encompassing its associated diagnoses and CCS codes as child nodes. Unlike ontology-driven hierarchies, which are predefined by medical ontologies, EHR-driven hierarchies emerge from the clinical context of patient visits.
To model this structure, we represent each visit 12 as box embedding b., = (bCen, boff), ), where the center b bee cap- tures the semantic meaning of the visit, and the offset bf reflects the variability introduced by the diversity of its asso- ciated medical concepts (shown in Figure 2(b)). The center of the visit's box embedding is computed as a weighted aggregation of the box centers of its associated diagnoses and CCS codes:
1.
PATIENT MODELING
To effectively model the dynamic nature of patient vis- its while preserving structural and semantic relationships among medical concepts, we propose an evolve-and- memorize patient box learning mechanism (shown in Fig- ure 2(c)). This mechanism leverages the geometric proper- ties of boxes, where the center evolves over time to capture the temporal progression of patient visits Hp. and the offset memorizes the structural and semantic boundaries of the pa- tient's 's medical history. Specifically, the evolve component dynamically updates the patient representation by weighting historical visits based on their time intervals, while the mem- orize component aggregates information from all historical visits to preserve hierarchical and semantic relationships. By integrating these two perspectives, our BoxLM achieves robust and interpretable understanding of patient data:
b Cen
ag ag-b; CCn
M
where been represents the center of the box for each diag- nosis or CCS code E Me,, and a; is the attention score that captures the relevance of each medical concept to the visit. The attention score is computed as:
exp (MLPP:(CCe Cen) 0gi E Ejem,, exp (M)P)( )
(4)
bCeu
Σ oeH,,
T Cen Doy
The offset of the visit's 's box embedding, Off represents the range of the visit's representation, which influenced
(6)
b Off = mmax(((b)) Off | 10g E Hpl)),
4

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
where the influence of historical visits Te, is calculated as:
Table Statistics of the datasets used in our experiments.
exp (MLP)(AT,))
Dataset
MIMIC-III
MIMIC-IV
Tog H exp (MLP) (AT,))
(7)
# of patients
5,449
79,393
# of visits
14,141
329,597
where ATg is the time interval between visit vy and vt.
Avg. # visits per patient
2.60
4.15
Max # visits per patient
29
169
2.3.2. VOLUME-BASH BASED DIAGNOSIS PREDICTION
# of unique diagnoses
3,874
37,917
of CCS codes
285
842
Given the high-dimensional and structural nature of box embeddings, we propose to leverage the intersection vol- ume between the patient box and CCS code boxes as measure of their similarity (shown in Figure 2(d)). By quan- tifying the overlap between boxes, our method provides intuitive insights into the relationships between medical con- cepts, thereby improving the accuracy and interpretability of model predictions.
high-dimensional box is the product of its side lengths along each dimension. This allows us to approximate the volume as follows:
dóm
Vop(e)
= II 2 22 exp (Mamee - ma 2/28))
Asm
Definition 2 (Box Minimun/Maximum Corner). The minimum and maximum corners of the box b, are given by byui biam Cen - bitt and b = =b
divm
II B Blog
1 + exp (A (Hane Mpnoo )
min
/8
+ b.#t.
Definition (Box Volume). The volume of box b, is calcu- lated by Vol(b.) = I max (b.me bmin 0), where k the indicator of dimension.
where Kg is the modified Bessel function of the second kind of order zero and y is the Muser-Maacheroni constant. This formulation leverages noise ensembles over a large collection of boxes, enabling the model to escape plateaus, alleviate gradient vanishing, and stabilize the optimization process (Huang et al., 2023). The detailed proof of this formulation is provided in Appendix B.
Definition 4 (Box Intersection). If there is overlap be- tween boxes b, and bj. their intersection box is denoted as banj = b, n by. where the minimum and maximum corners of the intersection box are bigi = max(b b and bing = min(b."nxx max "bpaa bynax).
The box volume- similarity score for the patient box bp and CCS code box b is computed as:
Despite the advantages of using box volume-based sim- ilarity, a common challenge is the potential for gradient vanishing when the boxes do not overlap (Lin et al., 2024; Liang et al., 2023; Jiang et al., 2023). To address this is- sue, we draw inspiration from GumbelBox (Dasgupta et al., 2020) and introduce an effective approach that integrates the Gumbel distribution g(z; μ. 5) @ exp __-e -5 to calculate the intersection box volume, ensuring non-zero gradients and stable optimization.
gp.c Log(Vol(bpne)).
(10)
where the logarithm function log further prevents the gradi- ent vanishing issue.
To model the multi-label nature of diagnosis prediction, we adopt Binary Cross- Cross-Entropy Loss, ensuring robust training by providing a probabilistic interpretation of the similarity scores. The objective function is defined as follows:
Specifically, the minimum and maximum corners of the intersection box between the patient box and CCS code boxes are derived using Gumbel distributions as follows:
[Mall
C C- Σ yn. log +(1 (1- Vp.c) 1og(1-Un..)) (11)
b miin = max maxxbmiin miin bmin) ~ MaxCumbel (p min $),
where pp.c is the ground truth label, and ÿp.c the predicted probability for patient p and CCS code c, which are normal- ized using the softmax function to ensure they represent valid probabilities for multi-label classification.
maxx min( (iinbbbie biax bmax) MinGumbel min ane B In (exp
max , (),
buin
buin /8)
+ exp
max "nn -311 exx -b b exp
b
(8)
3. Experiments
where MaxGumbel/Minnumbae are the max/min stable Gumbel distribution with location la and scale 3.
3.1. Experimental Setup
The intersection box expected volume is computed by con- sidering each dimension independently, as the volume of
Datasets and Evaluation Protocols. We use two real-world EHR datasets to verify the effectiveness of compared meth-
5

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Table Experimental results for diagnosis prediction (%) on the MIMIC-III and MIMIC-IV datasets with 5% training data. The hest performances are highlighted in boldface and the second runners are underlined. Improv. denotes the relative improvements of our proposed BoxIM over the second runners.
Dataset
MIMIC-I1I
MIMIC-IV
Vesit-Level
Code-Leve
Visit-Level
Code-Level
Metric
P@11 10
P@20
Acc@
Acc@20
P@10
P@20
Acc
Acc@20 20
DoctorAl
35. 36.72_0.16
43.70
26.31
42.6420.21
30.65+0.05
37 37.29+0.07
22. 22.6940.06
34.8320.00
RETAIN
34.
43.10 43.1040.18
25.73+0.17
42.00+0.19
33.4720.04
40.5640.05
24.52+0.03
37.5420.06
StageNet
36.02-0.19
43.91
26.47
12.7420.13
31.4620.07
38.02
22. 22.9440.06
34.8420.077
TRANS
36.64_0.22
44.85
26.96
43.35
28. 28.05
35.3610.05
21.64 21.6440.06
33.31
KAME
35.01=0.17
43.17
25.95
42.16
30.544..08
37.1640.11
22. 22.2940.09
34.10=z0.08
CGL
35.54_0.22
43.78+0.25
26.1540.19
42.49-0.20
31.7240.13
38.5140.00
23. 23.14
35.76+0.12
HiTANet
36.23_0.15
10.19
26.5640.17
43.0220.16
28.83+0.09
35.40,4011
21. 21.75+0.12
33.39
BoxCare
20.11
40.11
28.1240.19
43.84
35.1340.00
41.944..110
25, 25.58+0.09
38.80.0.09
BERT
20.77
32.1940.23
16.14
31.88+0.19
12.92++012
20.2440013
10.1640.12
19 19.27
BERT
23.81=0.19
32.26+0024
17.8440:26
32.15
13. 13.1620.10
20.85++012
10.95+0.11
19.59 14
BioBERT
31.17=0.21
43.21
25.35+0.19
42.06-0.20
18.81
27.37+0012
15.32+0.10
25.35=20112
BioBERT
34.42=0.225
43.44
25.49
42.27
18. 18.95
28.08:40.14
15.4440.13
25.61 20.
VecoCare
35.27=0.20
43.54
25.98
42.30
29.17
35.8440.12
22.06g0.09
33.89=0.12
BoxIM
43.88_0.23
51.62_0.25
31.7440.21
48.7410.19
42.04=0..4
49.65440.08
29.9440.06
44.5220.08
Improv. 14.84%
13.93%
12.87%
11.18%
19.67%
18.38%
17.04%
14.74%
ods, i.e., MIMIC-III (Johnson et al., 2016) and MIMIC- IV (Johnson al., 2018). Following (Ma et al., 2018; Chen et al., 2024), we chose patients who made at least two visits for both datasets and predicted the CCS codes for the next visit of patients. The statistics are summa- rized in Table 1. For evaluation metrics, we use visit-level Precision @K & (P@k) and code-level Accuracy @k (Aco@ from coarse-grained and fine-grained perspectives, which consistent with (Ma et al., 2018; Zhang et al., 2020; Wang al., 2023; Chen et al., 2024). As suggested in (Choi et al., 2016a; Lu et al., 2021; Lv et al., 2024), we also report Recall@ @ (see Appendix G) for more comprehensive evalu- ation of diagnosis prediction. Details on evaluation metrics are provided in Appendix
tients as the unit of segmentation, consistent with (Lv et al., 2024). We optimize the compared baselines with the stan- dard Adam optimizer and carefully tune their hyperparame- ters as suggested in the original papers. In particular, we set the embedding dimension dim as 16 and the scale of Gum- bel distribution 33 as 0.2. We evaluate each model with the f-fold cross- validaaio strategy. Both the mean and standard deviation of test performance are reported. All experiments are conducted using two NVIDIA GTX 3090 Ti GPUs. The full code for this work is available!.
3.2. Overall Diagnosis Prediction Results
In this section, we assess the effectiveness of our BoxLM for diagnosis prediction under the few-shot and varying ratios of training data scenarios, respectively.
Baselines. To comprehensively evaluate proposed BoxLM, we adopt 13 representative state-of-the-art meth- ods as baselines for comparison from three main perspec- tives: (1) temporal-aware methods: DoctorAI (Choi et al., 2016a), RETAIN (Choi et al., 2016b), StageNet (Gao et al., 2020), and TRANS (Chen et al., 2024); (2) hierarchy- aware methods: KAME (Ma et al., 2018), CGL (Lu al., 2021), HiTANet (Luo al., 2020), and BoxCare (Lv al.,
3.2.1. FEW-SHOT DIAGNOSIS PREDICTION
To comprehensively evaluate the performance of our pro- posed BoxLM, we compare it against 13 state-of-thhe-r baselines from three main perspectives: temporal-aware, hierarchy-aware, and semantic-aware methods. The results, as shown in Table 2, demonstrate that BoxIM consistently outperforms all baselines across both the MIMIC-III and MIMIC-IV datasets.
2024); (3) semantic-aware methods: BERT (Kenton & Toutanova, 2019), BERT", BioBERT (Lee et al., 2020), BioBERT and VecoCare (Xu et al., 2023b). Here, the superscript indicates methods that incorporate patients' historical visits during training. Additional details about these baselines are provided in Appendix
Comparison with hierarchy-aware models. Among the baselines, BoxCare stands out as a strong hierarchy- aware method that leverages box embeddings model structural relationships among diagnoses. BoxCare primarily focuses
Implementation Details. Both datasets are split into train- ing, validation, and test sets at ratio of 7:1:2, with pa-
6

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
RETAIN TRANS
REFFAIN TRANS
RETAIN
(88) BurClred
35
BeaCCare [Real.M
1500
BoxCart BOLLM
Borl
1000
30
500
30
1%%
5%
10%
1.5%
50%
100%
1%
5%
10%
1.5%
50%
100%
1%
5%
10%
1.5%
30%
100%
(a) Lesie-LevelF P@@ (%)
(b) Code-Level Acc@ (%)
(c) Training Runtime (S)
Figure 3. Experimental results for diagnosis prediction on the MIMIC-III dataset with varying ratios of training data.
on ontology- driven hierarchies (e.g., ICD codes), which cap- ture standardized relationships between medical concepts. However, it does not fully incorporate EHR-driven hier- archies, which reflect real-world patient visit patterns and medical concept co-occurrences. In contrast, BoxLM explic- itly integrates knowledge from pre-trained language models with both ontology-driven and EHR-driven hierarchies. As result, BoxLM outperforms BoxCare by average of 16.70% at the visit level and 13.96% at the code level across both datasets, demonstrating its superior ability to unify structural and semantic information of medical concepts for accurate diagnosis prediction.
3.2.2. VARYING RATIOS OF TRAINING DATA
To investigate the impact of training data size model performance, we conduct experiments on the MIMIC-III dataset with varying ratios of training data, e.g., 1%, 5%, 10%, 15%, 50%, and 100% (shown in Figure 3).
In general, BoxIM significantly outperforms all baselines across all data ratios, demonstrating its ability to general- ize even with limited training data. At both the visit and code levels (Figure 3(a) and Figure 3(b)), basclines show improvements as the training data ratio increases, but their performance gains are modest compared to BoxLM. For in- stance, BoxCare, the second- best method, improves a P@ 10 at the visit level from 34.78% 74% to 38.21 %, as the training data increases from 1% to 5%. However, even with 100% train- ing data, BoxCare's performance in P@ 10 remains slightly below BoxLM's performance with only 15% training data.
Comparison with temporal-aware models. One repre- sentative baseline is RETAIN, which employs RNNs with reverse time attention mechanisms to model the temporal dynamics of patient visits in EHRs. Another notable base- line is TRANS, a temporal graphic method, that leverages transformer- bassed architectures to model the sequential na- ture of patient visits. Compared with RETAIN and TRANS, BoxLM combines temporal modeling with a rich represen- tation of medical concepts derived from both EHR data and ontological sources through our structure-semantic fu- sion mechanism. This integration is particularly evident in BoxLM's superior performance on the MIMIC-IV dataset with more medical entities (shown in Table 1), surpassing TRANS by up to 46.48% in P@ 10 on MIMIC-IV.
Additionally, the runtime analysis (Figure 3(c)) further highlights the efficiency of BoxIM. Notably, RETAIN and TRANS exhibit significantly longer training times, where RETAIN is based on the reverse time attention mechanisms and TRANS relies on temporal graphical representations. In contrast, BoxIM achieves superior efficiency by leveraging box embeddings to unify structural and semantic represen- tations in a compact and interpretable manner. This design not only reduces computational overhead but also enables faster convergence during training. Compared to BoxCare, which also utilizes box embeddings, BoxLM introduces the strieture-semantic fusion mechanism and the evolve-and- memorize patient box learning mechanism. These innova- tions employ external semantic knowledge from PLM to further enhance efficiency. Moreover, they allow BoxIM to model temporal dynamics and hierarchical relationships effectively with an average 36.57% lower runtime. More ex- perimental results and analyses on MIMIC-IV can be found in Appendix H.
Comparison with semantic-aware models. Semantic- aware methods, such VecoCare and BioBERT, primar- ily rely on vector-based representations to encode medical concepts. These methods focus on capturing the semantic meaning of individual concepts. For example, BioBERT enhances semantic understanding through domain-specifti pre-training, while VecoCare uses PLM embeddings to rep- resent medical concepts. However, their rector-based design inherently limits their ability to model complex hierarchical relationships. In contrast, BoxLM models medical concepts as high-dineensional geometric structures (hyperrectangles) to unify structural and semantic information. This enables BoxLM to achieve an average improvement of 28.76% over semantic-aware methods on both datasets.
3.3. Analysis of Our Framework
Ablation Studies. To better understand our proposed tech- niques, i.e., Ontology- based hierarchy modeling (Onto),
7

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Visit
98,
98.
Visit 2
98. 101....
Visit (Ground-Truii
BoxCare's Hierarchy
401-405
401
-
401.1
RJAST
Predictive CCS
Visit 98
410-4114
414.01
CCS code
98
Essential hypertension
101
Coronary atherosclerosis and other heart diseases
Predictive CCS of Visit
Diagnosis
BoxLM's Hierarchy
98
98 101
401
Malignant essential hypertension
4
Benign essential hypertension
414.01
Coronury utherosclerosis of native
401.1
coronary artery
414.01
Figure 4. An illustrative example of diagnosis prediction for patient y on MIMIC-III generated by BoxCare and BoxLM, respectively.
complexity of overlapping and related medical concepts.
Table Ablation results (%) of our BoxIM on two EHR datasets,
Level-Level
Code-Level
Rather than representing visits single points, modeling visits as box embeddings by capturing EHR ddiven hier- archy achieves further improvement, with up to 25.33% gains in P@ 10 on MIMIC-IV. By capturing the relation- ships between visits and their associated medical concepts (e.g., diagnoses), BoxLM provides more robust and in- terpretable framework, especially in datasets with larger number of visits and medical entities.
Metric
P@
2@20
Acc@ 10
Acc @
MIMIC-II
Base
34.4220.25
43.44
25.49_
40.27
Onto
37.07 10.10
44.69 10.15
26.83_o.. _013
43.1810.n
FHR
41.40=0.24
48.7440.22
30.05_a.18
46.37
BoxLM
13.28
11.6210.25
31.74_0.21
48.7410.19
MIMIC-IV
Base
18.95
28.0810..1
15.44_0.13
25.61
Onto
32. 14..00
08.03
23.29_o.08
30.71
EHR
20.28
47.94
28.73_0.04
43.09
By extending box embeddings to patient-level representa- tions, we are able to effectively measure the overlap and complex relationships (e.g., structure, semantic, and tem- poral relations) among medical entities, leading to more precise diagnosis predictions. These results also affirm that our evolve-and-memorte patient box learning mechanism captures the evolving nature of patient visits.
BoxIM
22.04.0.04
10.65
29.94_0.06
44.5220.08
EHR-based hierarchy modeling (EHR), and evolve-and- memorize patient modeling, we study BoxIM on two EHR datasets as follows:
(1) Base model leverages BioBERT embeddings for medical concepts without any structure modeling, where visits are represented as single points; (2) "+ Onto" introduces box embeddings for CCS codes and diagnoses with our ontology- based hierarchy modeling (cf., Section 2.2.1) I). Visit represen- tations are constructed by summing the centers and offsets of associated CCS code and diagnosis boxes, followed by concatenation rather than using box representations; (3) "+ EHR" extends to visit-level box embeddings with our EHR-base based hierarchy modeling (cf., Section 2.2.2). Patient are vector embeddings similar to the visit embeddings in the "+ Onto" setup: (4) BoxLM models diagnoses, CCS codes, visits, and patients as box embeddings, and employs volume-based similarity for diagnosis prediction.
Hyperparameter Studies. We evaluate how three hyperpa- rameters (i.e., the box embedding dimension dim, the Gum- bel distribution scale 3, and the box volume calculation) impact the performance and clarify how to set them. The de- tailed results are shown in Appendix For box embedding dimension dim, the results show that increasing dim slowly enhances model performance for both datasets. For fair comparison with baselines, we set dim 16. For the Gum- bel distribution scale 8, a larger 3 makes the distribution closer to uniform, while a smaller / causes its probability density function to approach a hinge function, leading the random variable to degenerate into a constant (Lin et al., 2024). In this paper, we set 8 0.2, balancing model ac- curacy with the distinctiveness among boxes. For the box volume calculation, we compare the soft volume calcula- tion (Li et al., 2019) with our used Bessel volume calculation (cf., Eq. (9)). They both effectively mitigate the training difficulties that arise when disjoint boxes should overlap.
From Table 3, we have the following observations:
Incorporating box embeddings for hierarchical modeling achieves 7.70% gains with P@ 10 on MIMIC-III, with even more significant gains on MIMIC-IV (i.c., 69.60 99.60%). Such results highlight that datasets with more medical entities (e.g., MIMIC-IV) benefit greatly from modeling ontology- driven hierarchy structures, as they effectively address the
Case Studies. To highlight the advantages of BoxLM in diagnosis prediction, we provide real case study from the MIMIC-III dataset. Figure 4 contrasts the predictive results
8

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
of visit 3 for patient p from BoxCare and BoxIM.
concepts (Kwon al., 2024; Kim al., 2024). However, existing L.M-hased models face challenges in effectively encoding and interpreting hierarchical structures latent in language, which can be helpful for diagnosis prediction.
For patient p's diagnosed with essential hypertension dur- ing the first two visits (i.e., ii and 02), BoxCare relies on ontology-driven hierarchy for diagnosis boxes and models patients using vector embeddings. Consequently, it can- not capture the co-occurrence patterns present in EHR data, leading to inaccurate diagnosis predictions. In contrast, our proposed BoxIM make full use of ontology-driven and EHR-driven hierarchy by designing the structure-semantic fusion mechanism and poolv--and-memoriz patient box learning mechanism, empowering the model to uncover fine- grained, interpretable associations among medical entities. Therefore, BoxLM accurately forecasts the possibility of de- veloping coronary atherosclerosis and other heart diseases (with CCS code 101) by calculating the overlap between the accurate patient and CCS mexes-consistent with clinical knowledge that hypertension increases the heart's workload on arteries (Dzau, 1990; Libby et al., 2009), thereby ele- vating the risk of atherosclerosis. Additionally, we provide a quantitative metric Consistency @ metric to verify the interpretability of BOxIM (see Appendix J).
Box Embedding. Box embeddings have emerged as a pow- erful geometric representation approach for capturing com- plex hierarchical and relational information. Unlike tradi- tional vector embeddings, box embeddings model entities as high-dimensional hyperrectangles, enabling the repre- sentation of inclusion, exclusion, and intersection relation- ships (Mei et al., 2022). Motivated by the box embed- dings, which can clearly capture the complex structural rela- tions, real-world applications such as logical query answer- ing (Ren et al., 2020), knowledge base completion (Huang et al., 2023), taxonomy expansion (Jiang et al., 2023), recommendation (Lin et al., 2024), and diagnosis predic- tion (Lvv et al., 2024) have been proposed. As closest to us, BoxCare (Lv et al., 2024) applied box embeddings to repre- sent disease hierarchies for clinical predictions. However, BoxCare primarily focused on ontology hierarchies without fully integrating the rich semantics and temporal structures inherent in EHRs. For a more comprehensive discussion of related work, please refer Appendix A.
4. Related work
Diagnosis Prediction. Diagnosis prediction has been exten- sively studied due to its pivotal role in healthcare (Hendrik- sen et al., 2013; Yang et al., 2023). A significant category of deep learning-based based methods focuses on modeling con- textual dependencies within patient visit sequences using architectures like Recurrent Neural Networks (RNNs) and transformers, such DoctorAl (Choi et al., 201 RE- TAIN (Choi et al., 2016b), Dipole (Ma et al., 2017), and TRANS (Chen et al., 2024). Another prominent direction in- volves integrating external knowledge to enhance represen- tation learning (Choi et al., 2017; Ye et al., 3021; Xu et al., 2023a). For example, hierarchy-awar
5. Conclusion
In this paper, we propose BoxLM, a novel framework that unifies structural and semantic representations of medical concepts for diagnosis prediction. BoxLM in- cludes structure-semantic fusion mechanism that inte- grates antology-driver and EHR-driven hierarchies with LM-based semantic embeddings, and box-aware diagnosis prediction module that models temporal dynamics and quan- tifies relationships between patient and medical concepts. Extensive experiments demonstrate that BoxLM consistently outperforms state-of-the-an baselines, particularly in few- shot learning scenarios, showcasing its practical utility in real-workd clinical settings.
methods (Lu et al., 2021; Wang et al., 2023; Lv et al., 2024) have demonstrated strong performance by explicitly encoding disease domain knowledge. However, these methods often neglect the rich semantic information among medical concepts, which is valuable for diagnosis prediction modeling the contextual relationships and similarities between medical terms.
For future work, we aim to explore the application of BoxLM to other medical tasks (c.g., medication recommen- dation, medical report generation, and medical image seg- mentation) and investigate its integration with multimodal data sourees (e.g., clinical notes, X-rays, and sensor records) to further enhance its capabilities and application value.
Clinical Language Model. LMs have significantly ad- vanced clinical tasks by incorporating rich semantic infor- mation (Kenton & Toutanova, 2019; Yan & Pei, 2022). For example, BioBERT (Lee et al., 2020) and PubMed- BERT (Gu et al., 2021) utilized biomedical text corpora for pre-training, enabling improved representation of domain- specific medical terminology. VecoCare (Xu et al., 2023b) aggregated semantic information from patients through a dual-channel retrieval mechanism. In recent years, medical- specific L.LMs have leveraged large-scale medical datasets to further enhance the contextual understanding of clinical
Disclosure of Generative AI
We acknowledge the use of ChatGPT (OpenAl)) to assist in proofreading and improving the structure of this essay. The prompts we used include: "Please check the grammar and suggest improvements for clarity."
httpss//chat.openai.coo

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Acknowledgements
health records. In Proceedings of the Thirty-Third In- ternational Joint Conference on Artificial Intelligence, IJCAI-24, pp. 5763-5771, 2024.
This work was supported in part by the National Natural Sci- ence Foundation of China under Grants 6236223,,, Fu- jian Provincial Natural Science Foundation of China under Grants (2025101540), Zhejiang Provincial Natural Science Foundation of China under Grants (L-Q23F020007), and Zhejiang Provincial Department of Agriculture and Rural Affairs Project under Grants (2024SNJF044). Carl Yang was not supported by any fund from China.
Chen, T., Yin, H., Long. J., Nguyen, Q. V. H., Wang. Y., and Wang, M. Thinking inside the box: learning hypercube representations for group recommendation. In Proceed- ings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, PP. 1664-1673, 2022.
Choi, E., Bahadori, M. T., Schuetz, A., Stewart, W. F., and Sun, J. Doctor ai: Predicting clinical events via recurrent neural networks. In Machine learning for healthcare conference, pp. 301-318. PMLR, 2016a.
Impact Statement
This paper presents work whose goal is to advance the field of machine learning in healthcare by proposing BoxIM, novel framework for diagnosis prediction. BoxIM unifies structural and semantic representations of medical concepts, enabling accurate and interpretable predictions while mini- mizing the reliance on large-scale aatient-specifie data. By leveraging anonymized Electronic Health Records (EHRs) and incorporating hierarchical ontologies, BoxLM reduces the risk of privacy infringement and addresses the challenges of data scareity in clinical settings.
Choi, E., Bahadori, M. T., Sun, J., Kulas, J., Schuetz, A., and Stewart, W. Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. Advances in Neural Information Processing Systems, 29: 2016b.
Choi, E., Bahadori, M. T., Song, L., Stewart, W. F., and Sun, J. Gram: graph-based attention model for health- care representation learning. In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 787-795, 2017.
As a decision-support tool, BoxIM is designed to assist healthcare professionals rather than replace them, aiming to reduce their workload and improve diagnostic efficiency. Its strong performance in few-shot learning scenarios makes it particularly valuable in resource-constrained environments, such as rural or underdeveloped areas, where access to ex- tensive patient data is limited. By advancing the integra- tion of structural and semantic knowledge in medical AL, BoxLM contributes to the development of intelligent, ac- cessible, and ethical diagnostic tools, with the potential to improve patient outcomes and healthcare equity.
Dasgupta, S., Boratko, M., Zhang, D., Vilnis, L., Li, X., and
McCallum, A. Improving local identifiability in proba- bilistic box embeddings. Advances in Neural Information Processing Systems, 13:182-99, 2020.
Dzau, J. Atherosclerosis and hypertension: mechanisms and interrelationshiis. Journal of cardiovascular phar- macology, S64, 1990.
Gao, J., Xiao, C.,, Wang, Y., Tang, W., Glass, L. M., and Sun, J. Stagenet: Stage-aware neural networks for health risk prediction. In Proceedings of The Web Conference, pp. 530-540, 2020.
References
Abboud, R., Ceylan, I., Lukasiewicz, T., and Salvatori, T. Boxe: A box embedding model for knowledge base com- pletion. Advances in Neural Information Processing Sys- tems, 33:9649-9661. 2020.
Gu, Y., Tinn, R., Cheng, H., Lucas, M., Usuyama, N., Liu, X., Naumann, T., Gao, J., and Poon, H. Domain-specific language model pretraining for biomedical natural lan- guage processing. ACM Transactions on Computing for Healthcare (HEALTH), 3(1):1-23, 2021.
Bai, T., Zhang, S., Egleston, B. I.., and Vucetic, S. Inter- pretable representation learning for healthcare via cap- turing disease progression through time. In Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 43-51, 2018.
Hendriksen, J. M., Geersing, G.-J,, Moons, K. G., and de Groot, J. A. Diagnostic and prognostic prediction models. Journal of Thrombosis and Haemostasis, 11: 129-141, 2013.
Cai, Y., Wang, L., Wang, Y., de Melo, G., Zhang, Y., Wang, Y., and He, L. Medbench: A large-scale chinese bench- mark for evaluating medical large language models. In Proceedings of the AAAI Conference on Artificial Intelli- gence, pp. 17709-17717, 2024.
Huang. Z., Wang, D.,, Huang, B., Zhang, C., Shang, J., Liang, Y., Wang, Z., Li, X., Faloutsos, C., Sun, Y., et al. Concept2box: Joint geometric embeddings for learning two-view knowledge graphs. Findings of the Association for Computational Linguistics, pp. 10105-10118, 2023.
Chen, J.. Yin, C., Wang, Y., and Zhang, P. Predictive model- ing with temporal graphical representation on electronic
10

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Jiang, S., Yao, Q.. Wang, Q., and Sun, Y. A single vector is not enough: Taxonomy expansion via box embeddings. In Proceedings of the ACM Web Conference 2023, pp. 2467-2476 2476, 2023.
Li, Y., Li, Z.,, Zhang, K., Dan, R., Jiang, S., and Zhang, Y. Chatdoctor: A medical chat model fine-tuned large language model meta-ai (Iamaa) using medical domain knowledge. Cureus, 15(6), 2023.
Johnson, A. E., Pollard, T. J., Shen, L., Lehman, L.-w. H., Feng, M., Ghassemi, M., Moody, B., Szolovits, P., An- thony Celi, L.,, and Mark, R. G. Mimic-iii, freely ac- cessible critical care database. Scientific data, 3(1):1-9, 2016.
Liang, T., Zhang, Y., Di, Q. Xia, C., Li, Y., and Yin, Y. Contrastive box embedding for collaborative reasoning. In Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 38-47, 2023.
Johnson, A. E., Stone, D. J., Celi, L. A., and Pollard, T. J. The mimic code repository: enabling reproducibility in critical care research. JAMIA, 25(1):32-39, 2018.
Libby, P., Ridker, P. M., Hansson, G. K., and on Atherothrombosis, L. T. N. Inflammation in atheroscle- rosis: from pathophysiology practice. Journal of the American college of cardiology, 54(23)2112-2133 2009,
Kenton, J.
D. M.-W. C. and Toutanova, K. Bert: Pre- training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pp. 4171- 4186, 2019.
Lin, F., Zhao, Z., Zhu, X., Zhang, D., Shen, S., Li, X., Xu, T., Zhang, S., and Chen, When box meets graph neural network in tag-aware recommendation. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 1770-1780, 2024.
Kim, Y.. Xu, X., McDuff, D., Breazeal, C., and Park, H. W. Health-Ilm: Large language models for health prediction via wearable sensor data. CoRR, abs/2401 2tls22001.06666, 2024.
Liu, L.,, Yang, X., Lei, J., Liu, X., Shen, Y., Zhang, Z.,, Wei, P., Gu, J., Chu, Z., Qin, Z., et al. A survey on medical large language models: Technology, application, trustworthiness, and future directions. arXiv preprint arXiv:2406.03712, 2024.
Kwon, T., Ong, K. T.-i., Kang. D., Moon, S., Lee, R., Hwang, D., Sohn, B., Sim, Y., Lee, D., and Yeo, Large language models are clinical reasoners: Reasoning-awart diagnosis framework with prompt-generann In Proceedings of the AAAI Conference on Artificial In- telligence, pp. 18417-18425, 2024.
rationales.
Lu, C., Reddy, C. K., Chakraborty, P., Kleinberg, S., and Ning, Y. Collaborative graph learning with auxiliary text for temporal event prediction in healthcare. In Proceed- ings of the Thirtieth International Joint Conference on Artificial Intelligence, pp. 3529-3535, 2021.
Labrak, Y., Bazoge, A.. Morin, E.,, Gourraud, P., Rouvier, M., and Dufour, R. Biomistral: A collection of open- source pretrained large language models for medical do- mains. In Ku, L., Martins, A., and Srikumar, V. (eds.), Findings of the Association for Computational Linguis- tics, pp. 5848-5864 2024.
Lu, C., Han, T., and Ning, Y. Context-aware health event prediction via transition functions on dynamic disease graphs. In Thirty-Sixth AAAI Conference on Artificial Intelligence, pp. 4567-4574, 2022.
Langley, P. Crafting papers machine learning. In Proceed- ings of the 17th International Conference on Machine Learning, pp. 1207-1216, 2000.
Luo, J., Ye, M., Xiao, C., and Ma, F. Hitanet: Hierarchi- cal time-aware attention networks for risk prediction on electronic health records. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery data mining, pp. 647-656, 2020.
Lee, J., Yoon, W., Kim, S., Kim, D., Kim, S., So, C. H., and Kang, J. Biobert: a pre-trained biomedical language representation model for biomedical text mining. Bioin- formatics, 36(4): (6(4)::234413) 1234- 2020.
Lv, H., Chen, Z., Yang, Y., Ma, G., Yanchao, T., and Yang, C. Boxcare: A box embedding model for disease repre- sentation and diagnosis prediction in healthcare data. In Companion Proceedings of the ACM on Web Conference 2024, pp. 1130-1133, 2024.
Li, X., Vilnis, L., Zhang. D., Boratko, M., and McCallum, A. Smoothing the geometry of probabilistic box embeddings. In International conference on learning representations, 2019.
Ma, F., Chitta, R., Zhou, J., You, Q., Sun, T., and Gao, J. Dipole: Diagnosis prediction in healthcare via attention- based bidirectional recurrent neural networks. In Proceed- ings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1903-1911, 2017.
Li, Y., Qian, B., Zhang, X., and Liu, H. Graph neural network-based diagnosis prediction. Big data, 8(5):379- 390, 2020.

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Ma, F., You, Q., Xiao, H., Chitta, R., Zhou, J., and Gao, J. Kame: Knowledge-based attention model for diagnosis prediction in healthcare. In Proceedings of the 27th ACM international conference on information and knowledge management, pp. 743-752, 2018.
Vilnis, L., Li, X., Murty, S., and McCallum, A. Proba- bilistic embedding of knowledge graphs with box lattice measures. In ACL, pp. 263-272, 2018.
Vladika, J., Schneider, P., and Matthes, F. Medreqal: Exam- ining medical knowledge recall of large language models via question answering. In Ku, L.., Martins, A., and Srikumar, V. (eds.), Findings of the Association for Com- putational Linguistics, pp. 14459-14466 2024.
Mci. L., Mao, J., Guo, G., and Wen, J.-R. Learning proba- bilistic box embeddings for effective and efficient rank- ing. In Proceedings of the ACM Web Conference 2022,
pp. 473-482, 2022.
Wang, L., Liu, Q. Zhang, M., Hu, Y., Wu, S., and Wang, Stage-aware hierarchical attentive relational network for diagnosis prediction. IEEE Transactions Knowledge and Data Engineering, 36(4))1773-11)) 2023.
Onoe, Y., Boratko, M., McCallum, A., and Durrett, G. Mod- eling fine-grained entity types with box embeddings. In IJCNLP, pp. 2051-2064 2021.
Patel, D., Dangati, P., Lee, J., Boratko, M., and McCal- lum, A. Modeling label space interactions in multi-label classification using box embeddings. In ICLR 2022-10th International Conference on Learning Representations, 2022.
Wu, C., Shi, S., Wang. C., Liu, Z., Peng, W., Wu, W., Kong. D., Li, H., and Gai, K. Enhancing recommendation ac- curacy and diversity with box embedding: A universal framework. In Proceedings of the ACM on Web Confer- ence 2024, pp. 3756-3766, 2024.
Pavlović, A. and Sallinger, E. Expressive: A spatio- functional embedding for knowledge graph completion. In The Eleventh International Conference on Learning Representations, 2023.
Xiong, B., Cochez, M., Nayyeri, M.. and Staab, S. Hy- perbolic embedding inference for structured multi-label prediction. Advances in Neural Information Processing Systems, 35:33016-33028 2022.
Ren, H., Hu, W., and Leskovec, J. Query2box: Reasoning over knowledge graphs in vector space using box em- beddings. In 8th International Conference on Learning Representations, 2020.
Xu, X., Yao. B., Dong, Y., Gabriel, S., Yu. H., Hendler, J., Ghassemi, M., Dey, A. K., and Wang, D. Mental-lIm: Leveraging large language models for mental health pre- diction via online text data. Proceedings of the ACM on Interactive, Mobile, Wearable and Ubiquitous Technolo- gies, 8(1):1-32, 2024.
Saab, K., Tu, T., Weng, W.-H., Tanno, R., Stutz, D., Wul- czyn, E., Zhang, F. Strother, T., Park, C., Vedadi, E., et al. Capabilities of gemini models in medicine. arXiv preprint arXiv:2404.15 18416, 2024.
Xu, Y., Chu, X., Yang, K., Wang. Z., Zou, P., Ding, H., Zhao, J., Wang, Y., and Xie, B. Seqcare: Sequential training with external medical knowledge graph for diagnosis prediction in healthcare data. In Proceedings of the ACM Web Conference 2023, pp. 2819-2830, 2023a.
Shin, H.-C., Zhang, Y., Bakhturina, E., Puri, R., Patwary, M., Shoeybi, M., and Mani, R. Biomegatron: larger biomedical domain language model. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 4700-4706, 2020.
Xu, Y., Yang. K., Zhang, C., Zou, P., Wang, Z., Ding, H., Zhao, J., Wang, Y., and Xie, B. Vecocare: Visit sequences- clinical notes joint learning for diagnosis prediction in healthcare data. In Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence, IJCAI-23, pp. 4921-4929, 2023b.
Subramanian, S. and Chakrabarti, S. New embedded rep- resentations and evaluation protocols for inferring tran- sitive relations. In The 4sst International ACM SIGIR Conference on Research & Development in Information Retrieval, pp. 1037-1040, 2018
Yan, B. and Pei, M. Clinical-bert: Vision-language pre- training for radiograph diagnosis and reports generation. In Proceedings of the AAAI Conference on Artificial In- lelligence, pp. 2982-299229999 29902 2022.
Tang, X., Zou, A.. Zhang, Z., Li, Z., Zhao, Y., Zhang, X., Cohan, A., and Gerstein, M. Medagents: Large language models as collaborators for zero-shot medical reasoning. In Ku, L., Martins, A., and Srikumar, V. (eds.), Findings of the Association for Computational Linguistics, 599- 621, 2024.
Yang, K., Xu, Y., Zou, P., Ding, H., Zhao, J., Wang, Y., and Xic, B. Kerprint: local-global knowledge graph enhanced diagnosis prediction for retrospective and prospective interpretations. In Proceedings of the AAAI Conference Artificial Intelligence, pp. 5357-5365, 2023.
Vilnis, I.. and McCallum, A. Word representations via gaussian embedding. In ICLR, 2015.
12

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
Ye, M., Cui, S., Wang. Y., Luo, J., Xiao, C., and Ma, F. Medpath: Augmenting health risk prediction via medical knowledge paths. In Proceedings of The Web Conference, pp. 1397-1409, 2021.
Zhang, X., Qian, B., Cao, S., Li, Y., Chen, H., Zheng, Y., and Davidson, Inprem: An interpretable and trustwor- thy predictive model for healthcare. In Proceedings of the 26th ACM SIGKDD International Conference on Knowl- edge Discovery & Data Mining, pp. 450-460, 2020.
13

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
A. Detailed Related Work
Traditional Diagnosis Prediction. Early diagnosis prediction methods primarily relied on traditional machine learning approaches, such as logistic regression and random forest, which struggled to effectively handle the sighddmenssioal and sparse nature of EHR data (Hendriksen et al., 2013). The advent of deep learning has significantly advanced predictive modeling in diagnosis (Bai al., 2018; Li et al., 2020). For example, RETAIN (Choi et al., 2016b) employed a two-level attention-based RNN. Dipole (Ma et al., 2017) proposed a bidirectional RNN architecture with three attention mechanisms. TRANS (Chen et al., 2024) modeled the patient's ' EHR as a temporal heterogencous graph and utilized time-aware visit nodes to capture patient health status changes. Timeline (Bai et al., 2018) utilized time-aware attention mechanisms in RNNs for health event predictions. Chet (Lu et al., 2022) designed a context-aware dynamic graph learning method to learn disease combinations and disease development schemes. Another prominent direction involves integrating external knowledge graphs to enhance representation learning, such GRAM (Choi et al., 2017), MedPath (Ye et al., 2021), and SeqCare (Xu et al., 2023a).
Structure-aware Diagnosis Prediction. Existing structure-aware models mainly use the hierarchy of medical ontolo- gies(e.g., ICD-9) for diagnosis prediction. For example, HiTANet (Luo el al., 2020) adopted the time-aware Transformer and attention mechanism to capture correlations between visits from local and global views. KAME (Ma et al., 2018) proposed a graph-based attention model to obtain good performance with insufficient data. CGL (Lu et al., 2021) captured structural features of both patients and diseases, integrating attentive text features into sequential learning process.
Medical Large Language Models. With the advent of LLMs, medical artificial intelligence has undergone significant technological advancements and paradigm shifts, demonstrating the potential of LLMs to enhance healthcare delivery and improve patient outcomes (Liu et al., 2024). Starting with BioBERT(Lee et al., 2020) in 2019, the field has seen the development of specialized models for various medical tasks. Key milestones include BioMegatron(Shin et al., 2020) and PubMedBERTT(Gn et al., 2021), followed by the introduction of models like ChatDoocorr(Li et al., 2023) and ClinicalBERT(Yan & Pei, 2022). The field continues to evolve with advanced models such as MedAgents (Tang et al., 2024), MedREQAL (Vladika et al., 2024), PubMedBERT(Saab et al., 2024), and Health- Health-LLMM (Kim al., 2024), underscoring the ongoing innovation in leveraging LLMS for diverse healthcare applications.
Geometric Embedding, Geometric embedding models have attracted significant attention for their ability to preserve the intrinsic geometric structure of data (Lin et al., 2024; Chen et al., 2022). Among them, box embeddings represent entities as high-dimensional hyperrectangles. This makes them particularly effective for tasks involving hicrarchical (Onoe et al., 2021), transitive relations (Subramanian & Chakrabarti, 2018), entailments (Li et al., 2019), and uncertainty (Vilnis & McCallum, 2015), thereby inferring missing or incomplete hierarchical links in data. Motivated by the box embeddings. various real-world applications have been proposed, such as laxonomy expansion (Jiang et al., 2023), knowledge base completion (Huang al., 2023; Abboud et al., 2020), taxonomy expansion (Jiang et al., 2023), recommendation (Lin et al., 2024; Wu et al., 2024), and diagnosis prediction (Lv et al., 2024). Additionally, recent advances extend this idea to more expressive geometric forms such as polygons and hypercubes (Pavlović & Sallinger, 2023; Chen et al., 2022), further enhancing performance in domains like knowledge graphs and recommender systems.
B. Proof of Intersection Box Expected Volume Formulation
The detailed derivation of porie and pone is provided in (Dasgupta et al., 2020). According to Definition 3 in Section 2.3.2, we obtain the intersection box expected volume by calculating the expected length for each dimension as follows:
dim.
Vol(byne) = E [maaa((b (ba max
-b min pne preet 0)] = II 2,3Ko (2exp - (4 (umme - Myyric,,) )
minn
(12)
1=1
where dim is the dimension of embeddings and Ko is the modified Bessel function of the second kind of order zero. The proof of this formulation is detailed in (Dasgupta et al., 2020).
Next, we define f(x) = 2,3K, e2 Obviously, the function fx) is essentially exponential as x increases. The volume function approaches a hinge function as 3 + 0, which leads to numerical stability concerns (Dasgupta et al., 2020). Consequently, we employ the softplus function to approximate f(x)::
((x) R Blog (1 + exp=
(13)
14

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
where 7 is the Euler-Mascheroni constant. Based this approximation, the formulation of the intersection box expected volume can be derived follows:
dém
tiin (1 (upre - pri, (0.6,)) 1/822)))
Vol(bpne)
3 log
+ exp
(14)
R=1
C. Complexity Analysis of BoxLM
We analyze the time complexity of the BoxIM framework by four major components: Ontology-drives Hierarchy Modeling, EHR-driven Hierarchy Modeling, Evolve-and Memoorie Patient Modeling, and Volume-based Diagnosis Prediction.
(1) Ontology-drivee Hierarchy Modeling: the ontology hierarchy graph 9m (Vm E,), the offset embeddings of medical concepts (i.e., diagnoses and CCS codes) computed through relation-aware graph convolution networks. The complexity per layer is O (1 (Emml dim [Rml dim²), where |eml is the number of edges in Gm. [mm] is the number of relation types (e.g., parent-chil child relationship), and dim is the embedding dimension.
(2) EHR-driven Hierarchy Modeling: In the EHR-driven hierarchy graph, each visit vy is associated with a set of medical concepts. We represent the visit's box embedding via an attention attention-weighteed aggregation and maximum operation. The complexity is O(M,,,, dim), where [M,,, is the number of medical concepts associated with visit u..
(3) Evolve-and-Memorriz Patient Modeling: The patient box aggregates information from all historical visits Hp via temporal weighting and the complexity is O(Hpp dim). where |Hp] is the number of the patient's historical visits.
(4) Volume-basee basuee Diagnosis Prediction: In the diagnosis prediction phase, each patient box is compared to [Mael candidate CCS boxes via the -umbellapproxximated intersection volume. For N,, patients, the time complexity O(N, x [Mal x dim), which grows linearly with the embedding dim and similar to traditional vector-bass approaches.
D. Notations and Corresponding Description
As shown in Table 4, we summarize the key notations used in BoxIM and their corresponding descriptions.
Table 4. Notations
Notation
Description
d, v ed, ec
The medical concepts, including diagnosis d, CCS code c, and visit v. The embedding of medical concepts d, derived from PLM. The box embedding of concept i, where b Cen and bo are the center and offset of the box. The minimum corner of the box b,. The maximum corner of the box b;.
b, b==0 (bCo bfff
biiin
bynax
Vol(b:)
The volume of the box b,.
bin
The intersection box between boxes b, and by.
MinGumbel(p..) (3,
The min stable Gumbel distribution.
MaxGumbel (4,
The max stable Gumbel distribution.
The location vector of the Gumbel distribution.
3
The scale parameter of the Gumbel distribution. The set of neighboring concepts of diagnosis and CCS code The set of medical concepts d, c. The set of associated medical concepts of visit
((d).N((c)
Ma. Me
Mo Hp
The temporal progression of patient visits.
E. Details of Evaluation Metrics
Following previous studies (Ma et al., 2018; Chen et al., 2024; Choi et al., 2016a; Lu et al., 2021; Lv et al., 2024), three evaluation metrics are adopted for diagnosis prediction: Visit-level Precision @+ k (P@k) measures prediction effectiveness by counting accurate medical code predictions within the top & ranked results, normalized by the minimum of & and the
15

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
actual number of categories present in a patient visit. Code-level Accuracy @ k (Acc@k) quantifies prediction reliability as the proportion of correct diagnoses among all predicted medical codes at the individual code level. Recall @ measures the ratio of true medical codes in the top & predictions by the total number of ground-truth medical codes, reflecting the model's is ability to retrieve relevant diagnoses. The value of k ranges from 10 to 20 across both metrics, with higher values indicating better model effectiveness, where visit-level precision evaluates broader performance and code-level accuracy reflects more detailed correctness. Recall mimics the clinical process of differential diagnosis, where physicians identify the most probable conditions for evaluation. The calculation formulas are detailed in (Chen et al., 2024; Choi et al., 2016a).
F. Details of Compared Baselines
We compare 13 representative state-of-the-arr methods as baselines for comparison from three main perspectives:
(I) Temporal-aware methods concentrate on modeling temporal dynamics and sequences in patient data. DoctorAl (Choi et al., 2016a) encodes visits into vector representations and then processes them with GRUs to predict diagnoses in future visits. RETAIN (Choi et al., 2016b) employs RNNs that integrate reverse time attention mechanism to predict patient diagnoses. StageNet (Gao al., 2020) proposes stage-aware LSTM module and stage-adaptive convolutional module for diagnosis prediction. TRANS (Chen al., 2024) constructs temporal heterogeneous graph to jointly capture temporal dynamics and structural relationships in EHR data.
(2) Hierarchy methods emphasize capturing hierarchical structures within medical data. KAME (Ma et al., 2018) focuses on predicting patients' future health information by incorporating medical ontology knowledge of ICD codes into the sequence model. HiTANet (Luo et al., 2020) proposes a hicrarchical time-aware transformer for risk prediction based on EHRs. CGL (Lu et al., 2021) designs a collaborative graph learning model to explore patient-disease interactions and external medical knowledge. BoxCare (Lv al., 2024) employs box embeddings model both inclusive and exclusive relations among diseases and ICD-9 codes in EHR data.
(3) Semantic-aware methods primarily focus on leveraging semantic understanding for medical concepts through PLMs. BERT (Kenton & Toutanova, 2019) and BERT* pre-train the bidirectional encoder representations from transformers on large-scale unlabeled text data to understand the context of words by training on data bidirectionally. BioBERT (Lee et al., 2020) and BioBERT utilize biomedical text corpora for pre-training BERT, enabling improved representation of -meicca--peccfii terminology. For each visit, we first use these PLMs to encode a textual sequence formed by the names of all associated medical concepts. Then, we separately train Multi-Layer Perceptron (MIP) based on the obtained visit embeddings for diagnosis prediction. Here, the superscript indicates methods that incorporate patients' historical visits during training. VecoCare (Xu et al., 2023b) combines structured EHR data and clinical notes through the dual-channel retrieval mechanism to reduce heterogeneous semantic biases.
For a fair comparison, all baselines use the same EHR data without clinical notes, medications, and procedures, relying solely on the diagnostic information of patients.
G. Experimental Results and Analyses with Top-k Recall Metric
To comprehensively evaluate the performance of the proposed BoxLM, we compare it with four representative state-of-the-art baselines in Recall @ k (k=10, 20). As shown in Table 5, BoxLM consistently outperforms all baselines across both the MIMIC-III and MIMIC-IV datasets. These results are consistent with the observations reported in Section 3.2, further demonstrating the effectiveness of BoxIM.
H. Experimental Results and Analyses on MIMIC-IV with Varying Ratio of Training Data
To investigate the impact of training data size on model performance, conduct experiments on the MIMIC-IV dataset with varying ratios of training data, e.g., 1 1% 5%, 10%, and 15% (shown in Figure 5).
In general, BoxLM significantly outperforms all baselines across all data ratios, demonstrating its ability to generalize even with limited training data. At both the visit and code levels (Figure 5(a) and Figure 5(b)), baselines show improvements as the training data ratio increases, but their performance gains are modest compared to BoxLM. For instance, BoxCare, the second- best method, improves a P@ 10 the visit level from 30.60% 35.13%, as the training data increases from 1% to 5%. However, even with 15% training data, BoxCare's performance in P@ 10 remains slightly below BoxLM'S performance
16

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
with only 1% training data.
Additionally, the runtime analysis (Figure 5(c)) further highlights the efficiency of BoxIM. Notably, RETAIN and TRANS exhibit significantly longer training times, where RETAIN is based on the reverse time attention mechanism and TRANS relies on temporal graphical representations, In contrast, BoxIM achieves superior efficiency by leveraging box embeddings to unify structural and semantic representations in a compact and interpretable manner. This design not only reduces computational overhead but also enables faster convergence during training. Compared to BoxCare, which also utilizes box embeddings, BoxLM introduces the structure-semantic fusion mechanism and the evolve-and and-memorize patient box learning mechanism. These innovations employ external semantic knowledge from PLM to further enhance efficiency. Moreover, they allow BoxLM model temporal dynamics and hierarchical relationships with average 33.48% lower runtime.
Table 5. Recall @k results for diagnosis prediction (%) on the MIMIC-III and MIMIC-IV datasets with 5% training data. The best performances are highlighted boldface and the second runners are underlined. Improv. denotes the relative improvements our proposed BoxLM over the second runners.
Dataset
MIMIC-III
MIMIC-IV
Metric
Recall@ 10
Recall @ 20
Recall 10
Recall
DoctorAI
27.36
43.06
24.84
37.09
RETAIN
26.68
42.52
27.22
39.95
TRANS
27.97
43.82
23.38
34.84
BoxCare
28.61
44.63
29.76
42.33
BoxLM
34.49
50.71
34.85
48.65
Improw.
20.55%
13.62%
17.10%
14.93%
0SA RETAIN
35
0950 REETAIN
25000
RCTAIN
TRANS
20000
Buct M
Desl.M
Dacl M
15000
25
5000
1%
5%
10%
15%
1%
5%
10%
15%
1%
556
10%
15%
(a) Levit-Level P@ 10 (%)
(b) Code-Level Acc@ 10 (%)
(c) Training Runtime (8)
Figure 5. Experimental results for diagnosis prediction on the MIMIC-IV dataset with 1%, 5%, 10%. and 15% training data.
Experimental Results and Analyses of Hyperparameter Studies
As shown in Table 6, we evaluate the impacts of the box embedding dimension dim, the Gumbel distribution scale 3, and the box volume calculation. For box embedding dimension dim, the results show that increasing dim slowly enhances model performance for both datasets. For a fair comparison with baselines, we set dim = 16. For the Gumbel distribution scale B, larger makes the distribution closer uniform, while a smaller 3 causes its probability density function to approach hinge function, leading the random variable to degenerate into a constant (Lin et al., 2024). In this paper, we set = 0.2, balancing model accuracy with the distinctiveness among boxes. For the box volume calculation, compare the soft volume calculation (Li et al., 2019) with our used Bessel volume calculation (cf., Eq. (9)). They both effectively mitigate the training difficulties that arise when disjoint boxes should overlap.
J. Experimental Results and Analyses with Top-k Consistency Metric
In our study, CCS (Clinical Classification Software) codes from the MIMIC-IV dataset can be grouped into 22 distinct body systems', serving hierarchical classification labels for the predicted CCS codes. To better assess the interpretability of
http:://laapuu..aaeeeoooooooosse refined..jpp jefinedd
17

---

BoxIM: Unifying Structures and Semantics of Medical Concepts for Diagnosis Prediction in Healthcare
BoxLM, introduce the Consistency@@r @ metric (k=10, 20) (Patel al., 2022; Xiong et al., 2022), which quantifies the alignment between model predictions and the underlying hierarchical structure. The metric be computed as:
Σ l(h h4), Consistency@kk min(k: HDD)
1
(15)
where HHl is the number of ground-truth hierarchical labels associated with the CCS codes in visit e, and the numerator counts the number correct hierarchical predictions in the top-k. We compare proposed BoxLM with three representative ofate-of-the-ar baselines. Table demonstrates that BoxIM effectively preserves consistency on hierarchy, highlighting its strong interpretability and predictive capabilities.
Table 6. Hyperparameter study results (%) on the MIMIC-III and MIMIC-IV datasets with 5% training data.
Datasct
MIMIC-III
MIMIC-IV
Visit-Lewe
Code-Level
Visit-Level
Code-Level
Metric
Recall@ 10
Recall
Recall @
Recall
Pw 10
P@20
Acc 10
Acc
P@ 10
P@
Acc 10
Acc
dim
40.33
47.48
29.53
46.05
31.08
46.80
36.14
43.53
26.24
40.17
29.37
42.86
dim
41.88
49.35
30.57
32.39
48.68
41.00
48.55
29.36
43.83
33.62
47.82
dim
43.88
51.62
31.74
50.71
29.94
44.52
34.85
48.8
dim 32
44.84
51.68
32.52
48.99
34.88
50.99
50.72
57.48
35.74
51.01
42.31
56.71
42.91
50.12
31.05
47.47
33.50
49.48
41.10
48.92
29.18
43.43
34.27
47.98
8 0.2
43.88
51.62
31.74
48.74
34.49
50.71
42.04
49.65
29.94
44.52
34.85
48.65
8 0.3
45.39
52.37
32.81
49.55
35.56
51.69
41.71
49.35
29.78
44,30
34,79
48.60 60
45.55
52.73
33.00
49.80
$2.05
41.96
49.64
29.89
44.41
34.50
48.34
8
45.81
52.84
33.25
50.05
35.66
52.14
40.70
48.06
29.14
43.23
33.63
47.08
Bessel Volume
43.88
51.62
31.74
48.74
34.49
50.71
42.04
49.65
29.94
44,52
34.85
48.65
Soft Volume
43.91
51.05
31.79
48.22
34.31
50.39
42.23
49.76
30.10
44.62
34,93
48.74
Table Consistency @ results for diagnosis prediction (%) on the MIMIC-IV datasets with 5% training data. The best performances are highlighted boldface and the second runners underlined. Improv. denotes the relative improvements of our proposed BoxIM the second runners.
Dataset
MIMIC-IV
Metric
Consistency 10
Consistency@220
RETAIN
41.81
46.55
TRANS
35.13
38.43
BoxCare
47.36
52,06
BoxLM
56.53
62.44
Improw.
19.36%
19.94%
18


Reviews:
[{"review_id": "Dp84T9jyNJ", "reviewer": "Reviewer_skUY", "summary": "The paper proposes BoxLM, a framework unifying structural (ontology and EHR hierarchies) and semantic (language model embeddings) representations of medical concepts using box embeddings. Key contributions include a structure-semantic fusion mechanism, an evolve-and-memorize patient box learning module for temporal dynamics, and volume-based similarity for diagnosis prediction. Experiments on MIMIC-III/IV datasets demonstrate state-of-the-art performance, particularly in few-shot scenarios. The authors emphasize interpretability through geometric overlap analysis and case studies.\n\n## update after rebuttal\n\nAuthors addressed most of my concerns, I have raised my Overall Recommendation score from 2 to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> **Reviewer skUY.W1:** Adding quantitative metrics for interpretability.\n\nIn 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.\n\nThe 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.\n\nThe results (Table 1) demonstrate that BoxLM preserves consistency on hierarchy.\n\n**Table 1: Consistency@k results (%) for diagnosis prediction.**\n\n| Dataset |                    | MIMIC-IV           |       |        |\n| ------- | ------------------ | ------------------ | ----- | ------ |\n| Metric  | **Consistency@10** | **Consistency@20** | P@10  | Acc@10 |\n| RETAIN  | 41.81              | 46.55              | 33.47 | 24.52  |\n| TRANS   | 35.13              | 38.43              | 28.70 | 21.64  |\n| BoxCare | 47.36              | 52.06              | 35.13 | 25.58  |\n| BoxLM   | 56.53              | 62.44              | 42.04 | 29.94  |\n\n**Reference**  \n[1] Hyperbolic Embedding Inference for Structured Multi-Label Prediction, NeurIPS, 2022.  \n[2] Modeling Label Space Interactions in Multi-Label Classification Using Box Embeddings, ICLR, 2021.\n\n> **Reviewer skUY.W2:** Lacking hardware details and code availability.\n\nWe 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.\n\n> **Reviewer skUY.W3:** Adding related works on recent geometric approaches.\n\nWe 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.\n\n**Reference**  \n[3] ExpressivE: A Spatio-Functional Embedding for Knowledge Graph Completion, ICLR, 2023.  \n[4] Thinking inside The Box: Learning Hypercube Representations for Group Recommendation, SIGIR, 2022.\n\n> **Reviewer skUY.W4:** The saying for introducing box embedding might not be proper, as it's not straightforward to encode highdimensional embeddings to hyperrectangle.\n\nWe 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.\n\n> **Reviewer skUY.W5:** Some typos exist.\n\nThanks for pointing this out. We will fix it in the final version.\n\n> **Reviewer skUY.W6:** Section 3.3.1 needs a diagram for patient box evolution.\n\nWe 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.\n\n> **Reviewer skUY.Q1:** How does BoxLM handle missing or incomplete hierarchical links in EHRs?\n\nBoxLM 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.\n\n**Reference**  \n[5] A Single Vector Is Not Enough: Taxonomy Expansion via Box Embeddings, WWW, 2023.  \n[6] BoxE: A Box Embedding Model for Knowledge Base Completion, NeurIPS, 2020.\n\n> **Reviewer skUY.Q2:** What is the impact of box dimensionality on performance and computational cost?\n\nWe 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$.", "claims_and_evidence": "Mostly yes. Interpretability claims rely on a single case study (Figure 4); quantitative metrics for interpretability (e.g., user studies or consistency scores) are absent.", "methods_and_evaluation_criteria": "Generally make sense.\n- Box embeddings effectively model hierarchical inclusion and EHR-driven visit patterns. The fusion of ontology/EHR structures and evolve-and-memorize mechanism are novel and appropriate.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "- Validity: 5-fold cross-validation and Adam optimizer are standard. Few-shot experiments (1%–15% training data) robustly demonstrate generalizability.\n- Reproducibility: Runtime comparisons (Figure 3c) lack hardware details. Code availability is promised but not provided.", "relation_to_broader_scientific_literature": "Properly credits prior work (e.g., BoxCare, BioBERT) but might overlook recent geometric approaches in other fields (not limited to EHR / healthcare, I'm not sure).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "My concerns are mainly the interpretability, and your motivation for designing the method. It might not be proper to say`Luckily, we find that box embeddings offer a promising solution by representing entities as highdimensional hyperrectangle.` as it's not straightforward to encode highdimensional embeddings to hyperrectangle from my side.", "other_comments_or_suggestions": "- Typos: \"Exampled\" → \"Example\" (Figure 4 caption).\n- Section 3.3.1 needs a diagram for patient box evolution.\n- Include ablation on box dimensionality (see Questions).", "questions_for_authors": "- How does BoxLM handle missing or incomplete hierarchical links in EHRs (e.g., unrecorded parent-child relationships)?\n- What is the impact of box dimensionality (dim) on performance and computational cost?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "l6AgFNRAtb", "reviewer": "Reviewer_UQNn", "summary": "In this paper, the authors focus on a critical task: diagnosis prediction. They introduce a unique approach, the BoxLM representation, to represent Electronic Health Records (EHR) and diseases. While the paper is well-written and easy to follow, it has two significant shortcomings: 1. The evaluation metric used is not comprehensive. Typically, recall is considered an essential metric in AI-assisted diagnosis, as missing a diagnosis can have serious implications. 2. There is a lack of discussion and comparison with existing approaches in diagnosis prediction.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the overall summary. As for the several weaknesses and questions you mentioned, our responses are listed as follows:\n\n> **Reviewer UQNn.W1:** The evaluation metric suggests including recall. \n\nFollowing [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.\n\n**Table 1: Recall@k results (%) for diagnosis prediction on both datasets with 5% training data.** \n\n| Dataset |               | MIMIC-III     |       |        |\n| ------- | ------------- | ------------- | ----- | ------ |\n| Metric  | **Recall@10** | **Recall@20** | P@10  | Acc@10 |\n| RETAIN  | 26.68         | 42.52         | 34.73 | 25.73  |\n| TRANS   | 27.97         | 43.82         | 36.64 | 26.96  |\n| BoxCare | 28.61         | 44.63         | 38.21 | 28.12  |\n| BoxLM   | 34.49         | 50.71         | 43.88 | 31.74  |\n| **Dataset** |               | **MIMIC-IV**      |       |        |\n| Metric  | **Recall@10** | **Recall@20** | P@10  | Acc@10 |\n| RETAIN  | 27.22         | 39.95         | 33.47 | 24.52  |\n| TRANS   | 23.38         | 34.84         | 28.70 | 21.64  |\n| BoxCare | 29.76         | 42.33         | 35.13 | 25.58  |\n| BoxLM   | 34.85         | 48.65         | 42.04 | 29.94  |\n\n**Reference**  \n[2] Unveiling Discrete Clues: Superior Healthcare Predictions for Rare Diseases, WWW, 2025.  \n[3] Stage-Aware Hierarchical Attentive Relational Network for Diagnosis Prediction, TKDE, 2024.  \n[4] Predictive Modeling with Temporal Graphical Representation on Electronic Health Records, IJCAI, 2024.  \n[5] MEGACare: Knowledge-guided Multi-view Hypergraph Predictive Framework for Healthcare, Information Fusion, 2023.  \n[6] INPREM: An Interpretable and Trustworthy Predictive Model for Healthcare, KDD, 2020.\n\n> **Reviewer UQNn.W2:** Missing baseline Doctor AI [1].\n\nWe 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.\n\n**Table 2: Experimental results (%) for diagnosis prediction on both datasets with 5% training data.**\n\n| Dataset       |           | MIMIC-III |       |        |\n| ------------- | --------- | --------- | ----- | ------ |\n| Metric        | Recall@10 | Recall@20 | P@10  | Acc@10 |\n| **Doctor AI** | 27.36     | 43.06     | 35.72 | 26.31  |\n| RETAIN        | 26.68     | 42.52     | 34.73 | 25.73  |\n| BoxLM         | 34.49     | 50.71     | 43.88 | 31.74  |\n| **Dataset**       |           | **MIMIC-IV**  |       |        |\n| Metric        | Recall@10 | Recall@20 | P@10  | Acc@10 |\n| **Doctor AI** | 24.84     | 37.09     | 30.65 | 22.69  |\n| RETAIN        | 27.22     | 39.95     | 33.47 | 24.52  |\n| BoxLM         | 34.85     | 48.65     | 42.04 | 29.94  |\n\n**Reference**  \n[7] RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism, NeurIPS, 2016.\n\n> **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.\n\nOur 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.\n\n> **Reviewer UQNn.Q2:** What's the difference between CCS and diagnosis?\n\nSorry 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.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experimental design also has some flaws. In line with previous works, the authors should include recall as an evaluation metric.", "relation_to_broader_scientific_literature": "no", "essential_references_not_discussed": "Moreover, the paper lacks comparison with existing works on diagnosis prediction, such as [1] and subsequent works referencing [1].\n\n[1] Choi, Edward, et al. \"Doctor AI: Predicting Clinical Events via Recurrent Neural Networks.\" Machine Learning for Healthcare Conference. PMLR, 2016.", "other_strengths_and_weaknesses": "no", "other_comments_or_suggestions": "no", "questions_for_authors": "1. In Section 4.3's ablation study, the contribution of BoxLM should be evaluated. For instance, it could be compared with the classic vector representation.\n\n2. What's the difference between CCS and diagnosis", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QHUvBj6ozC", "reviewer": "Reviewer_yPUD", "summary": "This paper introduces BoxLM, a novel framework for diagnosis prediction in healthcare that aims to unify the semantic understanding of medical concepts with their underlying structural relationships. This approach integrates ontology-driven hierarchies and EHR-driven visit patterns with semantic embeddings from pre-trained LMs, and it further proposes a box-aware diagnosis prediction module and volume-based similarity measurement to model the temporal dynamics of patient visits and enable accurate diagnosis prediction. Overall, this paper has a clear motivation, a relatively complete theoretical section, and provides sufficiently comprehensive experimental results.\n\n## update after rebuttal\n\nThe authors have addressed my comments and I will maintain the original score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> **Reviewer yPUD.W1:** The difference between BoxLM and BoxCare.\n\nAs 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.\n\n> **Reviewer yPUD.W2:** Explicitly stating the specific task and research objectives upfront after the introduction for better reader comprehension.\n\nWe 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. \n\n> **Reviewer yPUD.W3:**  Box embedding concepts and operations are relatively complex.\n\nBox 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. \n\n> **Reviewer yPUD.W4:** Adding hyperparameter studies.\n\nWe 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. \n\nFor 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$.\n\n**Table 1: Experimental results (%) for diagnosis prediction with varying box dimension $d$.**\n\n| Dataset  | MIMIC-III |        | MIMIC-IV |        |\n| -------- | --------- | ------ | -------- | ------ |\n| Metric   | P@10      | Acc@10 | P@10     | Acc@10 |\n| $d=4$  | 40.33     | 29.53  | 36.14    | 26.24  |\n| $d=8$  | 41.88     | 30.57  | 41.00    | 29.36  |\n| $d=16$ | 43.88     | 31.74  | 42.04    | 29.94  |\n| $d=32$ | 44.84     | 32.52  | 50.72    | 35.74  |\n\nFor 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.\n\n**Table 2: Experimental results (%) for diagnosis prediction with varying Gumbel distribution scale $\\beta$.**\n\n| Dataset     | MIMIC-III |        | MIMIC-IV |        |\n| ----------- | --------- | ------ | -------- | ------ |\n| Metric      | P@10      | Acc@10 | P@10     | Acc@10 |\n| $\\beta=0.1$ | 42.91     | 31.05  | 41.10    | 29.18  |\n| $\\beta=0.2$ | 43.88     | 31.74  | 42.04    | 29.94  |\n| $\\beta=0.4$ | 45.55     | 33.00  | 41.71    | 29.78  |\n| $\\beta=0.6$ | 45.81     | 33.25  | 40.70    | 29.14  |\n\nFor 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.\n\n**Table 3: Experimental results (%) for diagnosis prediction with varying box volume calculation.**\n\n| Dataset       | MIMIC-III |        | MIMIC-IV |        |\n| ------------- | --------- | ------ | -------- | ------ |\n| Metric        | P@10      | Acc@10 | P@10     | Acc@10 |\n| Soft volume   | 43.91     | 31.79  | 42.23    | 30.10  |\n| Bessel volume | 43.88     | 31.74  | 42.04    | 29.94  |\n\n**Reference**  \n[1] When Box Meets Graph Neural Network in Tag-aware Recommendation, KDD, 2024.  \n[2] Smoothing the geometry of probabilistic box embeddings, ICLR, 2018.\n\n> **Reviewer yPUD.W5:** Providing examples for the interpretability of the proposed method.\n\nWe 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.", "claims_and_evidence": "BoxLM's main claim of outperforming baselines for diagnosis prediction is well-supported by MIMIC-III/IV experiments. Evidence is strong across metrics and data settings, including few-shot. Ablations and case studies validate components. However, the introduction overstates the novelty of box embeddings for hierarchical concepts, which has been proposed by BoxCare. The introduction should acknowledge BoxCare and clarify BoxLM's specific innovations. Despite this minor intro overclaim, core claims are valid. BoxLM builds on and improves prior work, showing strong results. Evidence supports the paper's conclusions.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are appropriate and make sense. The use of MIMIC-III and MIMIC-IV is well-justified and provides a realistic benchmark. The evaluation metrics are standard metrics in the field of EHR-based diagnosis prediction and are suitable for assessing the performance of the models.", "theoretical_claims": "The proofs in Appendix C are correct.", "experimental_designs_or_analyses": "The experimental designs are well-controlled, and the analyses are appropriate for supporting the claims of the paper. No major issues are identified in the experimental design or analyses.", "relation_to_broader_scientific_literature": "This work significantly builds upon the recent advancements in representing medical concepts and their relationships using box embeddings, most notably following the approach introduced by BoxCare. It extends BoxCare by more comprehensively leveraging hierarchy from both ontologies and EHR data. Crucially, it integrates pre-trained language models by leveraging their inherent prior knowledge.", "essential_references_not_discussed": "No essential related works appear to be missing.", "other_strengths_and_weaknesses": "## Strength\n\n- The paper is well-motivated, clearly articulating the research problem. The explanation of the proposed methods and formulas is presented with sufficient clarity, facilitating understanding.\n- The experimental evaluation yields compelling results, demonstrating consistent and statistically significant outperformance against a comprehensive suite of baselines across two large-scale, real-world EHR datasets. This robust empirical validation strongly supports the proposed approach.\n- The runtime analysis indicates that BoxLM achieves competitive computational efficiency, and even outperforms some baselines in terms of speed. This efficiency is a significant advantage for practical deployment and real-world application within clinical settings.\n\n## Weakness\n\n- While the paper effectively leverages box embeddings to model hierarchical medical concepts, building upon the foundational work of BoxCare, the primary innovation appears to reside in the enhanced integration of EHR information. The framing in the introduction may inadvertently overemphasize the novelty of the core box embedding concept itself, potentially leading to a perception of overclaim regarding the overall contribution relative to BoxCare.\n- To improve initial reader comprehension, the paper would benefit from a more explicit and upfront articulation of the specific task being addressed and the primary research objectives immediately following the introduction. This would provide essential context and guide the reader more effectively.\n- While box embeddings offer advantages, they are inherently more complex than traditional vector embeddings. The paper could benefit from further simplification or more intuitive explanations of box embedding concepts and operations, especially for readers less familiar with this technique.\n- While the paper mentions setting the embedding dimension and Gumbel scale, there appears to be limited discussion on the sensitivity analysis of these hyperparameters, particularly those specific to box embeddings (e.g., dimensionality, different distance/volume metrics). Exploring the impact of these choices would strengthen the robustness of the findings.\n- While paper claims interpretability, the authors should provide an example in the experimental section to visualize how different levels of concepts are represented in box embeddings, thereby enhancing the interpretability of BoxLM.", "other_comments_or_suggestions": "NA", "questions_for_authors": "See weakness.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The three reviewers are uniformly positive about the paper, rating it as 'weak accept' as a final recomendation.
The AC agrees with the reviewers and recommend 'Accept'.
Decision: Accept (poster)