NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Authors: Changshuo Liu
Lingze Zeng
Kaiping Zheng
Shaofeng Cai
Beng Chin Ooi
James Wei Luen Yip
Venue: ICML 2025 poster (2025)
Keywords: Healthcare
Cohort
Representation Learning
EHR
Primary area: applications->health_medicine
TL;DR: A cohort-aware neural representation learning method that precisely segments patients into finer-grained cohorts and captures both intra- and inter-cohort information.

Abstract: Electronic health records (EHR) aggregate extensive data critical for advancing patient care and refining intervention strategies. EHR data is essential for epidemiological study, more commonly referred to as cohort study, where patients with shared characteristics or similar diseases are analyzed over time. Unfortunately, existing studies on cohort modeling are limited, struggling to derive fine-grained cohorts or effectively utilize cohort information, which hinders their ability to uncover intrinsic relationships between cohorts. To this end, we propose NeuralCohort, a cohort-aware neural representation learning method that precisely segments patients into finer-grained cohorts via an innovative cohort contextualization mechanism and captures both intra- and inter-cohort information using a Biscale Cohort Learning Module. Designed as a plug-in, NeuralCohort integrates seamlessly with existing backbone models, enhancing their cohort analysis capabilities by infusing deep cohort insights into the representation learning processes. The effectiveness and generalizability of NeuralCohort are validated across extensive real-world EHR datasets. Experimental results demonstrate that NeuralCohort consistently improves the performance of various backbone models, achieving up to an 8.1% increase in AUROC.

NohralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Changshuo Liu
Lingze Zeng'
Kaiping Zheng'
Shaofeng Cai
Beng Chin Ooi 12
James Wei Luen Yip
Abstract
Division
Outcome
Electronic health records (EHR) aggregate exten- sive data critical for advancing patient care and refining intervention strategies. EHR data is es- sential for epidemiological study, more commonly referred as cohorl study, where patients with shared characteristics similar diseases alyzed over time. Unfortunately, existing stud- ies on cohort modeling are limited, struggling to derive fine-grained cohorts or effectively utilize cohort information, which hinders their ability to uncover intrinsic relationships between cohorts. To this end, we propose NeuralCohort, a cohort- aware neural representation learning method that precisely segments patients into finer-grained co- horts via an innovative cohort contextualization mechanism and captures both intra- and inter- cohort information using a Biscale Cohort Learn- ing Module. Designed as a plug-in, NeuralCohort integrates seamlessly with existing backbone mod- els, enhancing their cohort analysis capabilities by infusing deep cohort insights into the repre- sentation learning processes. The effectiveness and generalizability of NeuralCohort are validated across extensive real-world EHR datasets. Ex- perimental results demonstrate that NeuralCohort consistently improves the performance of vari- ous backbone models, achieving up 8.1 % increase AUROC.
(a) Traditional Cohort Studies
Temporal Modeling
Representation Learning
Pesudo Similrity
Godee
Pre-context Cohort Synthests Module
Cohort Derivation
Patlent
Cohort Leariing
Cohort Fuston
Prodiction
2 Biscale Cohort Learning Module
(b) Paradigm of NeuralCohort
Figure Traditional cohort studies and the paradigm of Neural- Cohort to achieve interpretable cohort intelligence.
2016). EHR data is instrumental in advancing data-driven approaches within the healthcare sector and facilitating crit- ical clinical decision-making (Yan et al., 2020; Wynants et al., 2020; Ooi et al., 2015) for more optimized patient management, including assisting doctors in assessing pa- tients' health conditions, developing treatment plans, and proactively averting adverse outcomes with heightened effi- cacy and intelligence (Hou al., 2020; Li et al., 2025; Lin et al., 2025), among others.
Contemporary approaches (Huang et al., 2019; Choi et al., 2017; Shang et al., 2019) primarily focus on extracting representations from the intrinsic information encapsulated within EHR data, which are subsequently used to enhance different downstream healthcare analytic tasks. Unfortu- nately, one indispensable consideration to achieve effective EHR representation learning has been largely overlooked by most existing studies Patient Cohorts. In healthcare research, as depicted in Figure (a), patient cohorts pertain to groups of individuals identified by shared characteristics. ranging from demographics to specific health conditions. Patient cohorts are integral to researching and developing effective medical interventions. For instance, when COVID-
1. Introduction
Electronic health records (EHR) are systematized collec- tions of patients' medical histories, stored electronically within healthcare systems, which typically consist of patient demographics and temporal medical features (Johnson et al.,
'National University of Singapore, Singapore 'Zhejiang Uni- versity, China Nationaa University Hospital. Singapore. Corre- spondence to: Changshuo Liu <shuo@ nuw.nuu..u.. ed..
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
19 ravages the world, doctors may opt construct two distinct patient cohorts: one comprising COVID-19 patients aged between 20 and 40, and the other encompassing those aged between 50 and 70. Notably, elderly patients are more likely to exhibit severe symptoms such as breathing distress, cognitive decline, and elevated fever, while younger pa- tients typically present with milder symptoms, such as fever, cough, and sore throat, or even remain asymptomatic (Al- imohamadi et al., 2020). In addition, doctors may observe consistent pattern among patients with fever symptoms, pre- dominantly diagnosing them with COVID-19 (Viner et al., 2021). The former analysis indicates that the elderly pop- ulation is prone to more severe symptoms, whereas the latter reveals a discernible association between fever and COVID-19, which serves as a valuable diagnostic indicator. Drawing from these insights, significant optimization goal of EHR representation learning is to develop an effective cohort construction approach that advances from individ- ual intelligence to interpretable cohort intelligence, thereby unleashing the profound potential of EHR data to facili- tate EHR analysis and deepen the understanding of patient insights.
an innovative cohort synthesis task to model the temporal EHR data with pseudo patient similarity, thereby paving the way for constructing fine-grained cohorts. Biscale Co- hort Learning Module: synergistically distill the cohort information both from local intra-cohort and global inter- cohort perspectives and finally derives the overall cohort information to enhance EHR data representation learning. Our contributions are summarized as follows.
We propose NeuralCohort, cohort-aware neural rep- resentation learning method to support fine- ane-grained co- hort generation and simultaneously exploit local intra- cohort and global inter-cohort information, which have been overlooked by prior studies in EHR analysis.
We design two-module paradigm for NeuralCohort. In Pre-context Cohort Synthesis Module, we model patients' temporal EHR data with pseudo patient simi- larity to construct fine-grained cohorts. Subsequently, in Biscale Cohort Learning Module, we capitalize on local intra-cohort and global inter-cohort information and encode them into augmented representations for prediction.
A straightforward approach cohort division relies exclu- sively on explicit features and uses either a single feature or a combination of several features to derive cohorts across different healthcare analytic tasks. Despite being intuitive, this approach may not satisfactorily tackle EHR data in practice, as it fails to achieve the following two crucial desiderata: (1) Fine-grained Cohort Division. EHR data records patients' temporal visits to the hospital for different reasons, thereby engendering different diagnoses, medica- tions, etc. Consequently, coarse-grained cohort division may place similar patients into different cohorts at varying time points, potentially introducing noise into the repre- sentation learning process. This underscores the need for a fine- ane-ggrined cohort division that takes into account the patient similarity in the time dimension, rather than rely- ing solely on explicit features. (2) Local Intra-cohort and Global Inter-cohort Information Exploitation. As previously discussed, the similarity of patients within a co- hort, i.e., intra-cohort information, is valuable to medical decisions. Additionally, contrasting the characteristics of a patient against those from different cohorts, i.e., inter-cohort information, can provide valuable insights and enhance di- agnostic accuracy. Therefore, effective EHR representation learning should comprehensively and jointly leverage both intra- and inter-cohort information to utilize EHR data to the utmost. However, existing methods fail to simultaneously fulfill both desiderata.
NeuralCohort can be seamlessly integrated with vari- ous backbone models, serving as a versatile plug-in to incorporate cohort information into healtheare analyt- ics, thereby enhancing overall performance.
We evaluate the effectiveness of NeuralCohort on three real-world EHR datasets across two tasks. The exper- imental results demonstrate that NeuralCohort, when integrated with backbone models, outperforms its coun- terparts by a large margin of up 8. %% in AUROC, confirming the efficacy of the comprehensive cohort information learned by NeuralCohort.
2. Related Work
EHR data encompasses a wide array of information, includ- ing structured data such as lab results and medications, as well unstructured data such as clinical notes (Sauer et al., 2022). Employing representation learning methods is cru- cial for comprehending the intricate relationships inherent within heterogeneous EHR data (Landi et al., 2020; Cai et al., 2021; Steinberg et al., 2021; Zheng et al., 2021; 2022; Cai et al., 2022; Cheng et al., 2016; Su et al., 2025). For instance, CompNet (Wang et al., 2019) uses a graph convolu- tional model, enhanced by reinforcement learning method, to learn the interactions between different medications and predict medication combinations. Beyond deriving represen- tations from the EHR data directly, several studies integrate patient similarity to enhance the efficacy of representation learning methods. A cosine similarity patient graph is constructed to aggregate the patient similarity information
To this end, we propose NeuralCohort, a cohort-aware neu- ral representation learning method for healthcare analytics. As illustrated in Figure (b), NeuralCohort is structured into two key modules. Pre-context Cohort Synthesis Module:
2

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healtheare Analytics
Pre-context Cohort Synthesis Module
EHR Data
Hierarchical Visit Engine
PseudoSim Training
Visit
VILA DOSTE PERCE
Bottom-up
Diagnoses DD,
Reverse Time Attention
poer Visit-level Medical Codes
Medications M,
Laboratory 4
22
Cohort Derivation
Intra-cohort
GCN
Fusion
Inter-cohort
Reconstructinn
Representation Initialization
Local Global Cohort Modeling
Prediction
2 Biscale Cohort Learning Module
Figure 2. Overview of NeuralCohort. NeuralCohort first models patients' visit-level representations to derive patient-level representations under PseudoSim Training task. The derived representations are used generate cohorts. NeuralCohort fuses the representations of the intra- and inter-cohort graphs with those derived from the backbones for the final prediction.
for enhancing length of stay prediction (Zang al., 2023). GRASP (Zhang et al., 2021) introduces framework that characterizes patient similarity through learned inter-cohort representations for various downstream clinical tasks. Un- fortunately, these methods predominantly focus on EHR data without exploiting the potential impact the similarity relationships among patients on EHR representation learn- ing. Further, some restrict their analytics to cohort-level similarity, and hence fall short of fully capturing the com- plex similarity relationships among patients. NeuralCohort mitigates the problem by generating fine-grained cohorts and leveraging cohort information from both intra- and inter- cohort perspectives. This dual-module method allows for more nuanced understanding and utilization of cohort dy- namies, significantly enhancing the effectiveness of EHR data representation learning.
ally aware EHR representation learning.
For more comprehensive discussion on traditional cohort studies and deep clustering, please refer to Appendix B.
3. Methodology
3.1. Problem Formulation
Given an EHR dataset comprising NP patients P (21.P2...- PNr). each patient p has a temporal se- quence of visits Vp {uI, U2. uN }. where N The features is the number of visits for patient of each visit v encompass three categories of medical codes: diagnoses D, = dd, ... dya), medications M, = (m), m2, ... mnx"), J. and laboratory tests C, (11, 12. ... ,1x:), where N'd Nm and N denote the num- bers of diagnosis codes, medication codes and laboratory test codes associated with visit respectively.
There have been patient similarity prediction proposals (Sun et al., 2012; Suo et al., 2018; Navaz et al., 2022). For in- stance, a Mahalanobis distance learning method (Wang & Sun, 2015) proposed integrate human experts' knowl- edge and patients' historical data to measure patient similar- ity. NeuralCohort differs from these works, which typically leverage deep learning models to discern intricate patterns among patients for deriving similarity, in that we employ pseudo patient similarity to drive the generation of cohorts. This difference enhances NeuralCohort's s ability to integrate and leverage cohort dynamics for more robust and contextu-
Given backbone model, our objective is to enhance the representation learning EHR data by incorporating pa- tient cohorts. To attain this goal, we propose NeuralCohort, cohort-aware neural representation learning method for healthcare analytics. As depicted in Figure 2, we present the training pipeline M(P; Op, Oe) as follows:
M(P: Br. Oc)= Mr(S(P)PPep))
Ro:
(I)
NeurstCabort
Pre-osetare: Cabart Syntheeie
Riwesie Colver Learning
where Mp and Mc correspond the modules for the Pre-
3

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
context Cohort Synthesis and Biscale Cohort Learning, re- spectively, with associated parameters Op and Oc. The Pre-context Cohort Synthesis Module aims to derive cohorts C via modeling pair-wise pseudo patient similarity S(P). The Biscale Cohort Learning Module leverages these co- horts C' alongside the representations from the backbone Ro. resulting in enhanced representations.
Visit level
V--1
Vt+1
Bottum-up up Modeling
Diagnosis level
d1
d
ds
de
d
Medication & Laboratory Test
3.2. Pre-context Cohort Synthesis Module
m
m2
Myn
h
The Pre-context Cohort Synthesis Module is geared towards learning patient-level representations and consequently de- riving cohorts. In essence, NeuralCohort first constructs a hierarchical network tailored for processing medical codes to extract visit-level features R2 from the EHR data. Sub- sequently, NeuralCohort proceeds to learn the patient-level features Ry based on both the visit-level features R2 and the pair-wise pseudo patient-level similarity S(P), serving as the basis for Biscale Cohort Learning.
Figure The hierarchy for visit-level features.
Finally, the representation for diagnosis codes is derived using the hierarchical architecture described above.
To), Rv. = 77. To).
(2)
where D, is the set of diagnosis codes for visit v, and To denotes the medical ontology tree. The functions 9p and Is respectively model the hierarchical path and semantics. The function fo is responsible for transforming these rep- resentations into diagnosis- representations. Further details regarding the implementations of these functions are provided in Appendix G.
Hierarchical Visit Engine. Each visit contains a set of di- agnosis codes D.,, medication codes Mr and laboratory test codes c. Notably, diagnosis codes adhere to specific on- tologies that standardize the classification and interpretation of medical data. For example, the ICD-9 (International Clas- sification of Diseases, Ninth Revision) employs a tree-like structure where each leaf corresponds to a standardized di- agnosis code and a unique medical term (Quan et al., 2005). The leaves and their respective ancestors delineate hierar- chical paths within the ontology. A specific instance of path Root + 240-279 + 249-259 > 250 - 250.4 represents the hierarchical path of the leaf 250.4, with 250 designating Diabetes mellitus and 250.4 denoting Diabetes with renal manifestations. To distinguish leaves with analogous struc- tures and derive informative diagnosis code representations, we partition the ontology into two constituents:
In addition to the hierarchical ontology structure for diag- nosis codes, a further code-level hierarchy is discernible within the patients' visits. During a hospital visit, diagnosis codes represent diseases, medications are prescribed to treat these diseases, and laboratory tests assist in confirming the diagnoses. Correspondingly, akin to the approach adopted in MiME (Choi et al., 2018), we establish a code-level hier- archy for visit-level representation, as illustrated in Figure 3. The visit-level representation is computed via:
R2 = (v JMC SMR Rc.))
(3)
Path. This component traces the path in the ontology tree from the Roos to a leaf. We construct a graph based on the tree structure of the ontology and employ a graph-based method obtain the representation of the path. However, certain leaves, such 250.4 and 250.5, may share similar paths, which necessitates additional differentiation strategies.
where RD. is the representation of diagnosis codes calcu- lated by Equation 2. R.M, and Re. denote the representa- tions of medication codes, and laboratory test codes for visit v, where RM, = (MM), Re, fc(L.). The function fv is employed to generate the visit-level representation R,.
To derive patient-level representations, we take the current visit as the anchor visit and adopt an adapted reverse time attention mechanism.
Semantics. Each leaf within the ontology embodies a unique medical term. Ancestors convey board seman- tics, whereas leaves provide specific semantic details. Further, the semantic information inherent in the same ancestor vary for different children. To tackle this, we adopt a semantic similarity metric between the text representations of an ancestor and a leaf as the weight for aggregating their representations. This combined consideration of paths and semantics enables the dif- ferentiation of individual leaves.
ON=, ,0: sofl (R
,Ro))) ))
(4)
3, = siim (Re,, Reyx), ), for j = 1,2,.... No'
(5)
N" R₂ = ΣM' 3, 2 G R.,
(6)
where the attention a for each visit is calculated in reverse order. We introduce a cosine similarity score 3 to explicitly
4

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
model the similarity between past visits and the most recent visit. The combination of a and B enhances the model's ability to capture and utilize temporal relationships within the visit sequence.
derivation; however, it is beyond the primary scope of this paper. Here, Jensen-Shannon divergence (Lin, 1991) and student's t-distribution (Student, 1908) are employed. Please refer to Appendix G for detailed implementation.
PseudoSim Training. The central consideration for cohort construction is patient similarity. Traditionally, patient simi- larity is labeled manually by experts, which is a resource- intensive process demanding substantial human efforts. To address this challenge, we propose a pseudo patient similar- ity prediction task, leveraging pseudo labels generated from patients' diagnosis codes to facilitate the learning of infor- mative patient representations. grounded on the fact that diagnosis codes play a pivotal role as the primary criteria in medical treatment.
LJs = Ep Ee &pe log
2tp
+ 9pc log
20po Gyc t
(9)
tpe
(1+11,,-6e11)
Ppcl 22, 9cc
where qpe Σ: E. (1+111 Epo Σ; (o2,// 2, 9.y)
Opc is the soft cohort assignment of patient p. Epc is the self-training target cohort distribution, |C] is the number of cohorts and je is the centroid of cohort c. Finally, the cohort e for patient p is arg max Ope- c
3.3. Biscale Cohort Learning Module
Given the temporal context and diagnosis code frequency, we utilize weighted Jaccard Similarity function enhanced with time decay factor to determine patient similarity;
Cohorts play a crucial role in EHR analysis. The Biscale Cohort Learning Module is employed to encode cohort in- formation into the representations derived from backbones and harness local and global cohort information to its fullest extent for diverse healthcare analytic tasks.
(Dlcta1)) Camm (d. Pp.)
uCPp,
Ppt.
y, )py))
JPp
Individual Representation Initialization. We first derive individual representations from EHR data using established backbones as the initial representations. For a given pa- tient p. the initial representation RV is generated by the backbones, which will serve as fundamental inputs for Neu- ralCohort, facilitating its subsequent operations.
where Dp, and Dpy are the sets of diagnosis codes across all visits for patient P. and Pj. respectively. At; and Aty represent the time intervals between the last occurrence of the diagnosis code d and the most recent visit for patients P, and Pj. respectively. C'Fiinn (CFmax) computes the min- imum (maximum) frequency of the diagnosis code d in Dp. and Dp,
Local Intra-cohort Modeling, To effectively capture inter- actions among patients within a cohort, Local Intra-cohort Modeling depicts the cohort as a graph in which each patient is represented as node. In this graph, each pair of patients in the same cohort is connected. The weight between two nodes is calculated based on the similarity between patient representations R, computed in the Pre-context Cohort Syn- thesis Module.
Finally, we optimize the representation learning by learn- ing patient relationships through maximization of Mutual Information Neural Estimation (Belghazi et al., 2018). For each patient p. we select the top Kp patients with the high- est similarity scores as positive samples Stt . and randomly choose Kp patients from the remaining pool as negative samples Sy During the training process, we utilize the learned patient-level representations Ry. positive samples Stt and negative samples Sp for PseudoSim Training.
A(i.j) = I|
R R
(10)
+ b*)
(11)
where A is the adjacency matrix of the graph. H² is the node representations at layer k, H. = Ro. W* and 6² are the trainable weights and bias. A = D 3 (A+ D is the normalized adjacency matrix and D is the degree matrix. The output of the final layer is the local intra-cohort representation R2,
Lipes = EEE R.). R.) HogE (X.(
+ + (R R3) exp(f
(8)
where the first line represents the mutual information of positive samples composed of expectation of similarity and nog--oomalization, and the second line serves as a penalty for the negative samples. fo is the similarity prediction network.
Global Inter-cohort Modeling. The global inter-cohort modeling aims to elucidate the contrasting knowledge across dissimilar cohorts. The inter-cohort representations should maintain distinet separation while preserving the semantic integrity of each cluster. In light of this, we employ an encoder-decoder architecture and cohort loss function to reconstruct the representations for cohorts.
The strategy above not only yields valuable information on patient similarity but also facilitates the construction of cohorts C from the patient-level representations Rp. which is subsequently harnessed in downstream tasks.
Cohort Derivation. The leteent-level representations could potentially be utilized with various methods for cohort
Rs = fdee Σ R() pEc
(12)
5

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Patients within the same cohort generally exhibit similar behaviors, in contrast those from different cohorts, who may manifest distinct patterns. To supervise the represen- tations of both intra- and inter-cohort graphs, we introduce a cohort loss function composed of reconstruction loss and contrastive loss:
Table I. Overall performance of NeuralCohort against baselines for readmission prediction on the MIMIC- IIII dataset.
Readmission Task on MIMIC-III
Model
AUPRC
AUROC
Accuracy
ClinicaBBRRT
0.630=0.005
0.651=00.066
58.7%=0055%
Σ Σ (R)) Je Σ s, RRO) 0,
KNN
0.628 0.628+0.002
0.651+0.0002
58.6 58.6%
Lco
E IRG CEC
E
(13)
K-Means
0.62910.001
0.650.10.001
18.3%-1..3%
+ DEC
0.632+0.005
0.654=+0.002
58.4%=0.11%
PEc
DEKM
0.638+0.003
0.659++0.005
58.9%
GRASP
0.61810.002
0.617:10.001
56.2%
where sim is the cosine similarity function. The first term is the reconstruction loss and the other term is the cohort contrastive loss.
.DGLOS
0.635+0.000
0.533++0.002
58. 58.0%
IDC
0.638 0.008
0.657::0.004
59 59.0%
NeuralCohorr
0.662.10.033
0.600
61.2% 0.4%
Med2Vec
0.55410.005
0.614:10004
54.1%%=0.7%
Cohort Fusion. To consolidate the information from the representations generated by the backbones, the intra-cohort graph, and the inter-cohort graph, we devise a cross-domain attention mechanism to obtain the augmented EHR data representation for subsequent healthcare analytic tasks:
KNN
0.541=0.004
0.598==0.003
54.9%=0.5%
K-Means
0.544+0.005
0.600=00.004
54.5%=00.4%
DEC
0,550=0.003
0.611=#0.002
54.3%=00.5%
+ -DEKM
0.547=00004
0.608==0.006
54.6%%=044%
-GRASP
0.542+0.003
0.601
53.8%
+GGLOS
0.559=0.004
0.542=00.002
54.5%
IDC
0.562=000044
0.622:+0.003
54.5%
NeuralCohort
0.574+0.003
0.634 0.005
56,9 % 0.2%
RE, RG v (d))(
RRG
aL
GG
(14)
MiME
0.543.10.006
0.602.10.005
56.8%
v /diet(ng)
KNN
0.543110.004
0.610:10.003
56 0.55%
K-Means
0.546110000
50.5% 0.605
Re final = Sattn ( RR or,' RT og R&)
(15)
DEC
0.54910.007
0.608110.004
57.3%
+ DEKM
0,548110002
0.611110.003
57. 0.55%
where R², and R& pertain to the representations extracted from the intra- and inter-cohort graphs, respectively. The final loss function for NeuralCohort is as follows:
+ GRASP
0.530=1.0009
0.589=100110
57 00,95
DGLoS
0.5011000006
0.543.10.004
57. 10.6%
IDC
0.542:10.006
0.605-10.003
57.2 -0.4%
+ NeuralCohort
0.508:20004
0.629:10.003
58.6 56 0.33
L = Apes Tepes + AIsTJs + reo + Las (16)
where Apes Às and Aco denote the weights assigned to the loss of the PseudoSim Training Lpss cohort derivation LJs and the Cohort Loss Leo. Lds signifies the loss associated with the downstream task.
et al., 2018), and ClinicalBERT (Huang et al., 2019) to de- rive initial representation Ro. Subsequently, we employ KNN, K-Means, DEC (Xie et al., 2016). DEKM (Guo et al., 2021), GRASP (Zhang et al., 2021), DGLoS (Zang et al., 2023) and IDC (Svirsky& & Lindenbaum, 2024) as baselines for comparison with our proposed NeuralCohort, which are directly applied to the representations of the backbones to integrate the cohort information into R2 for enhanced pre- dictive accuracy. The efficacy of the models on a given task primarily hinges on the backbones. The objective of NeuralCohort is to infuse patient similarity information into these backbones. Therefore, our analysis focuses on assess- ing the performance enhancements achieved by integrating NeuralCohort and other baselines into the backbones.
4. Experiments
4.1. Experimental Setup
We conduct extensive experiments across three widely- recognized real-world EHR datasets: MIMIC-III (John- son et al., 2016), MIMIC-IV (Johnson et al., 2020) and Diabetes 30 (Clore John & Beata, 2014), addressing two pivotal medical tasks, cross-visit Hospital Readmis- sion Prediction (Huang et al., 2019) and within-visit Long Length-of-Stay (LOS) Prediction (Alsinglawi et al., 2022). Through these experiments, we aim to assess the generaliz- ability and efficacy of NeuralCohort in predicting critical healthcare outcomes, thereby demonstrating its applicability across diverse patient populations and clinical scenarios.
In our experimental setup, we repeat each model across the datasets five times and compute the average result. We measure the performance of the models using three metrics: Area Under the Precision- Curve (AUPRC), Area Un- der the Receiver Operating Characteristic (AUROC), and accuracy. Additionally, we conduct an ablation study of NeuralCohort and a sensitivity study on its key hyperparam- eters including Kp and the number of cohorts |C'.
To assess the versatility of NeuralCohort as a plug-in for enhancing the representations derived from various back- bones, we compare the performance of basclines with their enhanced counterparts using cohort information provided by NeuralCohort. We employ well-established backbone models, namely Ved2 (Choi et al., 2016), MiME (Choi
Detailed descriptions of the datasets, tasks, backbones, base- lines, and experimental settings are provided in Appendix from D to G.
6

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Table 2. Overall performance of NeuralCohort against baselines for long LOS prediction on the MIMIC-III dataset.
Table 3. The ablation study results of NeuralCohort on long LOS prediction utilizing the ClinicalBERT backbone on the MIMIC-IIZ III dataset. My, Mintrus and Menter represent Pre- context Cohort Synthesis Module, Local Intra-cohort Modeling, and Global Inter-cohort Modeling of Biscale Cohort Learning Modulc.
Long LOS Task on MIMIC- III
Model
AUPRC
AUROC
Accuracy
ClinicalBERT
0.658=0.002
0.590+=0.002
59.7%=++.33%
+ KNN
0.659+0.002
0.594+0.001
59.9%+00.4%
+ K-Means
0.605-10.001
0,590.10.00
59.8%%-1.33%
Model
Modules
MIKIC III AUROC
+ DEC
0.652=0.008
0.603++00009
60.2%-+0.3%%
Mp Mosten
Manter
AUPRC
Accuracy
+ DERM
0.661+0.000
0.592++0.000
60.1%++0.3%
+ GRASP
0.665.10.002
0.584.10.00
60.0%%110.4%
-Mi -Maters Minter
0,682
0.613
60.4%
+ DGLLOS
0.685=0.005
0,548+0.000
61 6..9%%+
0.675
0.603
59.2%
+ IDC
0.681 +..004
0.612+00.003
60.2%%++0.44% 0.4%%
0.714
0.649
61.8%
+ NeuralCohort
0.738 0.003
0.671 0.004
63.7% 0.7%%
NeuralCohort
0.738
0.671
63.7%
Med2Vec
0.908:10.001
0.894110000
64.4%
+ KNN
0.876=0.002
0.80
75.1% 0..%%
K-Means
0.877=0.002
0.875+00.003
74.8%+++033%
+ DEC
0.886=10.004
0.889
05.4%
analysis. Specifically, KNN and K-Means, when applied directly to backbones, do not operate within a similarity- aware feature space. Additionally, DGLoS constructs a coarse-grained global graph, and GRASP focuses exclu- sively on inter-cohort modeling. DEC, DEKM and IDC cannot model the medical semantics. As a result, instead of accurately modeling patient similarity, these baselines may introduce noise regarding patient similarity into the back- bones, consequently degrading the overall performance.
+ DEKM
0.903=00.022
0.897++0.000
73.8%++0.44%
GRASP
0,903=0.003
0.887+0.002
0..5%
+ DGLLS
0.907=00.004
0.837=000006
0..5%
+ IDC
0.906=#0.033
0.899+00003
7..2%%
+ NeuralCohort
0.919=0.002
0,906
80.7 70.7%%-01 + 0.3%%
MiME
0.913-10.003
0. 1.904.10.00
78.4%%-10..5
+ KNN
0.910=10.002
0.895110.000
80.2% 4%
K-Means
0.904=0.003
0.00
0..0%
+ DEC
0.912.10.005
0. 0.907-10.008
79.6%
+ DEKM
0.911=10.005
0. 0.908
79.6%--10.4%
+ GRASP
0.898:20.004
0.896:10.000
81.1%:10.44
DGLLOS
0.917-10.002
0.854110.000
80.3%.10.3%
+ IDC
0.919-10.005
0.900
80.3%-10..2%
Among the backbones, ClinicalBERT demonstrates superior performance compared to Med2Vee and MiME in the hospi- tal readmission task, whereas the opposite trend is observed in the long LOS task. This divergence can be attributed to the distinct characteristics of the tasks and the underlying model architectures. Med2Vec and MiME are adept at mod- eling relationships among large number of medical codes, which provide comprehensive information about a patient's s current visit, thus making them more suitable for within- visit long LOS prediction. Conversely, clinical notes contain extensive records of patients' past medical histories. Clin- icalBERT, which excels at learning representations from such data, is better equipped to capture the relationships across multiple visits and enhance its efficacy for cross-visit readmission prediction.
+ -NeuralCohort
0.936
0.923+00.02
82.8% 1.22%
We provide additional experimental evaluations, including sensitivity study in Appendix H and an efficiency study in Appendix J.
4.2. Overall Results
We present the overall results of NeuralCohort and baselines for the two prediction tasks on MIMIC- III dataset in Ta- bles 1 and 2, respectively. The conclusions drawn from dif- ferent datasets are consistent, confirming the robustness of our findings. Results for MIMIC-IV IV and Diabetes130 datasets are provided in Appendix I.
As demonstrated by the results, NeuralCohort consistently attains the best performance, outperforming the baselines by significant margin with improvements up to 8.0% in AUPRC, 8.1% in AUROC, and 16.3% in accuracy. This considerable performance advantage underscores the ef- ficacy of NeuralCohort over diverse backbones. Further, NeuralCohort enhances the performance of each backbone on both tasks on all evaluated metrics, which substantiates its capability to derive fine-graines cohorts and effectively infuse cohort information into the representations, leading to improved predictive accuracy.
4.3. Ablation Study
We conduct an ablation study on the long LOS prediction task with the ClinicalBERT backbone on the DIMIC-IIC dataset. in order to evaluate the impact of each module and component of NeuralCohort. The experimental results are presented in Table 3. Comparing NeuralCohort against NeuralCohort (-Mp) (Row 1 vs. 4), the Pre-context Co- hort Synthesis Module demonstrates notable contributions, i.e., 5.6%, 5.8%, and 3.3% improvements in AUPRC, AU- ROC, and accuracy. This underscores the efficacy of the Pre-context Cohort Synthesis Module in capturing nuanced patient similarity, thus emphasizing the importance of ac- counting for fine- fine-graing grained cohorts for EHR Representation Learning. The results in Rows and of Table 3 elucidate
We observe the baselines fail to yield consistent performance improvements primarily due to their inadequate modeling of cohort information at fine-grained level for accurate
7

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Table Comparison between traditional cohorts and NeuralCo- hort on the backbone ClinicalBERT and Ved2Vec for readmission prediction on the MIMIC- III dataset.
Table 5. Comparison of C-H scores between NeuralCohort and baselines with the backbones ClinicalBERT and Med2Vec on long LOS prediction.
MIMIC-IN III
Model
Sc-u
Model
Sc-u
Model
ClinicaBBRTT
57.6
Med2Vec
18.4
AUPRC
AUROC
Accuracy
KNN
56.8
+ KNN
18.5
K-Means
59.7
K-MMans
20.1
ChiicalBERT
0.630
0.651
58.7%
DEC
62.6
+ DEC
20.9
MCG
0.629
0.651
58.5%
+ DEKM
60.2
+ +DEKM
19.6
MMCA
0.631
0.651
58.3%
GRASP
65.5
GGRASP
21.4
+ .MCD
0.629
0.652
58.7%
DGLLS
61.7
+ +DGLoS
21.5
*MCH
0.621
0.643
58.1%
+ IDC
69.4
+DCC
23.2
NeuralCohort
80.3
NeuralCohort
25.7
NouralCohort
0.662
0.681
61.2% %
Med2Vec
0.554
0.614
54.1%
MCGG
0.546
0.611
54.3%
+MCA
0.548
0.612
53.4%
MMD
0.557
0.619
54.5%
*MCH
0.551
0.616
53.7%
NeuralCohort
0.574
0.634
56.9%
the respective contributions of Local Intra-cohort Modeling and Global Inter-cohort Modeling. In comparison, the re- moval of the Local Intra-cohort Modeling results in a more pronounced decline in performance compared to the re- moval of Global Inter-cohort Modeling. This highlights the necessity of unveiling implicit information within similar patients. In nutshell, these results affirm the indispens- able roles played by distinct modules and components in NeuralCohort in ensuring accurate EHR analysis.
(a) Med2Vec
(b) GRASP
4.4. Traditional Cohorts vs. NeuralCohort
(c) IDC (d) NeuralCohort
To evaluate NeuralCohort against traditional medical co- hort construction approaches effectiveness, we employ a widely recognized cohort division approach using two dis- tinct criteria: gender (G), age (A) and two disease-specific criteria: diabetes diagnosis (D) and hypertension diagno- sis (H). The results displayed in Table 4. As shown, traditional medical cohorts tend to perform comparably to, and occasionally less accurately than the backbones. This primarily arises from the reliance of traditional medical cohort construction on a limited set of features, resulting in coarse- grainedd cohorts that are not effective enough for cohort pattern mining. Consequently, such traditional ap- proaches will inevitably group highly dissimilar patients into the same cohort, introducing noise to the model that adversely affects the overall performance. In stark contrast, NeuralCohort leverages a patient's sequential visit-level rep- resentations from both intra- and inter-cohort perspectives at a fine-grained level, thereby boosting the prediction accu- racy.
Figure 4. t-SNE visualization the final representations before prediction for selected eight-cohort points, where backhone is Med2Vec.
which further results in the absence of ground truth for evaluating the performance of clustering- based methods on patient similarity. To address this, we use the Calinski- Harabasz score (C-H score) (Caliński Harabasz, 1974) to quantitatively assess the performance of NeuralCohort and other baselines. The C-H score Sc-# is calculated
Sc (m- W [[C[-(c--1
(17)
where m is the sample size, |C| is the number of clusters, B is the cross-cluster variance, W is the within-cluster vari- ance, and tr is the trace of matrix.
The comparison of C-H scores between NeuralCohort and baselines with ClinicalBERT and Med2Vec for the long LOS prediction task is detailed in Table We note that KNN, DGLoS, along with the bare backbones do not involve the concept of "clusters" and as a result, we employ K-
4.5. Interpretability Analysis
Quantitive Analysis of Cohorts. It non-trivial to define patient similarity for comparison in healthcare analytics,
8

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Means their representations prior to prediction layer to evaluate C-H scores. The results clearly demonstrate that integrating a backbone with NeuralCohort yields higher C- H scores, which indicates that the cohorts generated through NeuralCohort are more interpretable and effective.
Table 6. Distinctive feature of four cohorts
Cohort
Distinctive features
Visualization. The visualization of selected eight cohorts (clusters) points using t-SNE (Van der Maaten & Hinton, 2008), based on the final representation R Rina of Med2Vec for the long LOOS prediction task, is illustrated in Figure 4. Extended visualization of all baselines are shown in Ap- pendix K. Notably, significant overlaps among the points from Figure 4 (a) to (c) suggest that the backbone and base- lines lack robust discriminative capabilities when utilizing R Pigiaa However, the integration of NeuralCohort substan- tially reduces this overlap, enhancing the distinctness of the cohorts. This demonstrates that NeuralCohort is capable of encoding the data into representation similarity space and the fine-grained cohort information integrated by Neu- ralCohort effectively facilitates downstream EHR analysis tasks.
#1
HF, Arrhythmia, CHD, Angina Diabetes, DLM, CD, Obesity OUA, PSL, UF, F&E Disorders OPI, CKD, Thyroiditis, Flu
#2
#3
#4
Table 7. Detailed Descriptions of the Abbreviations
Abbreviation
Detailed Description
HF
Heart Failure
CHD
Chronic Heart Disease
DLM
Disorders of Lipoid Metabolism
CD
Coagulation Defects
OUA
Other and Unspecified Anemias
PSL
Pneumonitis due to Solids and Liquids Urethral Fistula Fluid and Flectrolyte Disorders
4.6. Clinical Signiticance
UF
NeuralCohort identifies cohort-specing features that are directly linked to clinical outcomes, therehy enhancing decision-making for patient management. For example, the distinctive features of the four cohorts visualized in Figure 4 are provided in Table 6, where the features are identified through t-test. The detailed descriptions of the used abbreviations are listed in Table 7. Specifically, Co- hort #1 includes patients with cardiovascular conditions and is at high risk for acute cardiac events and extended hospital stays (Tigabe Tekle et al., 2022). Early identifica- tion enables the prioritization of telemetry beds, cardiology consults, and monitoring. Clinically, this enables timely diuretics, echocardiography, and discharge planning with heart failure follow-up. Cohort #2 features patients with chronic metabolic and hematologic conditions that often co- occur and need interdisciplinary care. Identifying this cohort allows hospitals to mobilize diabetes educators, lipid clinics, and anticoagulation monitoring services. Clinically, it en- ables insulin titration, lipid therapy, and bleeding risk man- agement to reduce complications and improve outcomes. Cohort #3 is characterized by renal and urological issues re- quiring frequent labs, fluid monitoring, and nephrology sup- port. Early identification allows hospitals to allocate renal panels, schedule imaging for uropathy, and prepare dialysis resources-helping manage AKI risk and avoid care escala- tion. Cohort #4 presents complex chronic and acute condi- tions, requiring coordination across pulmonology, nephrol- ogy, endocrinology, and infectious discase. Identifying this cohort supports planning for respiratory support, hormone therapy, infection control, and opioid withdrawal. Opera- tionally, it enables resource allocation for respiratory teams,
& E Disorders
OPI
Other Pulmonary Insufficiency
CKD
Chronic Kidney Disease
isolation rooms, and endocrine/renal labs. Such insights allow clinicians to optimize resource allocation and tailor in- terventions, thus significantly improving hospital efficiency and patient care.
5. Conclusions
In EHR representation learning, existing studies primar- ily focus on deriving informative representations to facili- tate downstream task predictions. However, two essential desiderata, fine-grained cohort division and effective ex- ploitation of both intra- and inter-cohort information, remain unresolved. To address the problem, we propose NeuralCo- hort, a cohort-aware neural representation learning method for EHR representation learning. The core thrust of Neural- Cohort centers around two pivotal modules: the Pre-context Cohort Synthesis Module, responsible for cohort deriva- tion, and the Biscale Cohort Learning Module, dedicated to capturing and encoding local intra-cohort and global inter- cohort insights. Extensive experimental results confirm the superiority of NeuralCohort over baselines across various backbones, achieving improvements of up to 8. 1% in AU- ROC. More importantly, NeuralCohort contributes valuable medical insights to EHR analysis.
9

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Acknowledgments
network for structured data. In SIGMOD Conference, pp. 207-220. ACM, 2021.
The work of NUS researchers is supported by the Lee Foun- dation in terms of NUS Lee Kong Chian Centennial Profes- sorship grant.
Caliñski, T. and Harabasz, A dendrite method for clus- ter analysis. Communications in Statistics-thent and Methods, 3(1):1-27, 1974.
Impact Statement
Cheng, Y., Wang, F., Zhang, P., and Hu, J. Risk prediction with electronic health records: A deep learning approach. Proceedings of the 2016 SIAM international confer- ence on data mining, pp. 432 440. SIAM, 2016.
This paper advances deep learning for healthcare by address- ing critical challenges in cohort discovery. The proposed NeuralCohort extends beyond traditional individual intelli- gence paradigms, evolving into interpretable cohort intelli- gence to enhance patient population analysis. Being highly adaptable and generalizable, NeuralCohort can be seam- lessly integrated with diverse backbone models for wide variety of healthcare applications. More importantly, Neural- Cohort holds significant implications for large-scale cohort research, fostering new insights into population health dy- namies, supporting advancements of precision medicine, and ultimately contributing the development of more effective and targeted healthcare interventions.
Choi, E., Bahadori, M. T., Searles, E., Coffey, C., Thomp- son, M., Bost, J., Tejedor-Sojo, J., and Sun, J. Multi-layer representation learning for medical concepts. In proceed- ings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 1495-1504, 2016.
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 knowledge discovery and data mining, pp. 787-795, 2017.
References
Alimohamadi, Y., Sepandi, M., Taghdir, M., and Hosamirud- sari, H. Determine the most common clinical symp- toms in covid-19 199 patients: systematic review and meta- analysis. Journal of preventive medicine and hygiene, 61 (3):E304, 2020.
Choi, E., Xiao, C., Stewart, W., and Sun, Mime: Multi- level medical embedding electronic health records for predictive healthcare. Advances in neural information processing systems, 31, 2018.
Clore John, Cios Krzysztof, D. and Beata, Diabetes 130-US Hospitals for Years 1999-2008. UCI Machine Learning Repository, 2014. DOI:
Alsinglawi, B., Alshari, O., Alorjani, M., Mubin, O., Al- najjar, F., Novoa, M., and Darwish, O. An explainable machine learning framework for lung cancer hospital length of stay prediction. Scientific reports, 12(1):607, 2022.
Eldredge, J. Cohort studies in health sciences librarianship. Journal of the Medical Library Association, 90(4):380, 2002.
Belghazi, M. I., Baratin, A., Rajeswar, S., Ozair, S., Bengio, Y., Courville, A., and Hjelm, R. D. Mine: mutual informa- tion neural estimation. arXiv preprint 1801 0000,00 2018.
Fisher, J. icd-codex/icd-codexx v0.4.5, December 2020. URL https:://oi.... www.dddo.ooogg/o// /eoo/do 4300935.
Cai, Q., Xie, Z., Zhang, M., Chen, G., Jagadish, H., and Ooi, B. C. Effective temporal dependence discovery in time series data. Proceedings of the VLDB Endowment, 11(8):893-905, 2018.
Gabrielaite, M., Bennedbæe,, M., Rasmussen, M. S., Kan, V., Furrer, H., Flisiak, R., Losso, M., Lundgren, J. D., Group, I. S. S., and Marvig. R. L. Deep-sequencing of viral genomes from a large and diverse cohort of treatment- naive hiv-infected persons shows associations between intrahost genetic diversity and viral load. PLoS computa- tional biology, 19(1);e 1010756 2023.
Cai, Q., Zheng, K., Ooi, B. C., Wang, W., and Yao, C. ELDA: learning explicit dual-interactions for healthcare analytics. In ICDE, pp. 393-406, IEEE, 2022.
Cai, Q., Zheng, K., Jagadish, H. V., Ooi, B. C., and Yip, W. L. Cohortnet: Empowering cohort discovery for interpretable healthcare analytics. Proc. VLDB Endow., 17(10)2487-2500 2024.
Guo, W., Lin, K., and Ye, W. Deep embedded k-means clustering. In 2021 International Conference on Data Mining Workshops (ICDMW), pp. 686-694, IEEE, 2021.
Hou, N., Li, M., He, L., Xie, B., Wang, L., Zhang, R., Yu, Y.. Sun, X., Pan, Z., and Wang, K. Predicting 30-days mortality for mimic-iii patients with sepsis-3: a machine
Cai, S., Zheng, K., Chen, G., Jagadish, H. V., Ooi, B. C., and Zhang, M. Arm-net: Adaptive relation modeling
10

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
learning approach using xgboost. Journal of translational medicine, 18(1):1-14, 2020.
random walks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 14193-14201, 2024.
Huang. K., Altosaar, J., and Ranganath, R. Clinicalbert: Modeling clinical notes and predicting hospital readmis- sion. arXiv preprint arXiv: 05334, 2019.
Navaz, A. N., El-Kassabi, H., Mohamed, A. S., Oulhaj, A., and Khalil, K. A novel patient similarity network (psn) framework based on multi-moder deep learning for precision medicine. Journal of Personalized Medicine, 12(5):768, 2022.
Jiang, D., Cai, Q., Chen, G., Jagadish, H., Ooi, B. C., Tan, K.,L.,, and Tung, A. K. Cohort query processing. Pro- ceedings of the VLDB Endowment, 10(1), 2016.
Ooi, B. C., Tan, K.-L., Wang. S., Wang, W., Cai, Q.. Chen, G., Gao, J., Luo, Z.,, Tung, A. K., Wang, Y., et al. Singa: A distributed deep learning platform. In Proceedings of the 23rd ACM international conference Multimedia, pp. 685-688, 2015.
Johnson, A., Bulgarelli, L.,, Pollard, T., Horng, S., Celi, L. A., and Mark, R. Mimic-iv. Phy- sioNer. Available online at: https://physionet. org/comten/mimiciv/ll http:///miiiiiiie...feeeeeee August 23, 2021), 2020.
Paneth, N. and Monk, C. The importance of cohort re- search starting early in life to understanding child health. Current opinion in pediatrics, 30(2):292, 2018.
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, a freely ac- cessible critical care database. Scientific data, 3(1):1-9, 2016.
Pennington, J., Socher, R., and Manning, C. D. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), pp. 1532-1543, 2014.
Landi, I., Glicksberg, B. S., Lee, H.-C., Cherng, S., Landi, G., Danieletto, M., Dudley, J. T., Furlanello, C., and Miotto, R. Deep representation learning of electronic health records to unlock patient stratification at scale. NPJ digital medicine, 3(1):96, 2020.
Quan, H., Sundararajan, V., Halfon, P., Fong, A., Burnand, B., Luthi, J.-C., Saunders, L. D., Beck, A., Feasby, T. E., and Ghali, W. A. Coding algorithms for defining comorbidities in icd-9-cm and ied-10 administrative data. Medical care, pp. 1130-11 1 139, 2005.
Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
Ros, F. and Riad, R. Feature and Dimensionality Reduction for Clustering with Deep Learning. Springer, 2024.
Li, S., Lin, T., Lin, L.,, Zhang, W., Liu, J., Yang, X., Li, J., He, Y., Song, X., Xiao, J., et al. Eyecaregpt: Boost- ing comprehensive ophthalmology understanding with tailored dataset, benchmark and model. arXiv preprint ar0iv:2504. 13650, 2025.
Sauer, C. M., Chen, L.-C., Hyland, S. L., Girbes, A., Elbers, P., and Celi, L. A. Leveraging electronic health records for data science: common pitfalls and how to avoid them. The Lancet Digital Health, 4(12):=833<888, 2022.
Setia, M. S. Methodology series module Cohort studies. Indian journal of dermatology, 61(1):21-22 2016.
Li, Wu, Y., Ester, M., Kao, B., Wang, X., and Zheng, Y. Semi- supervupe clustering attributed heterogeneous information networks. In Proceedings of the 26th inter- national conference on world wide web, pp. 1621-1629, 2017.
Shang. J., Ma, T., Xiao, C., and Sun, J. Pre-training of graph
augmented transformers for medication recommendation. arXiv preprint arXiv: /906.00366 2019.
Lin, Divergence measures based on the shannon entropy. IEEE Transactions on Information theory. 37(1):145-151 1991.
Song, W. and Chung, K. C. Observational studies: co- hort and case-control studies. Plastic and reconstructive surgery, 126(6):2234-2243) 2010.
Lin, T., Zhang, W., Li, S., Yuan, Y., Yu, B., H., He, W., Jiang, H., Li, M., Song, X., et al. Healthgpt: medical large vision- langugge model for unifying comprehension and generation via heterogeneous knowledge adaptation. arXiv preprint arXiv:2502.0965
Steinberg, E., Jung, K., Fries, J. A., Corbin, C. K., Pfohl, S. R., and Shah, N. H. Language models are an effective representation learning technique for electronic health record data. Journal of biomedical informatics, 113: 103637, 2021.
2025.
Lu, Y., Lin, Y., Yang, M., Peng, D., Hu, P., and Peng. X. De- coupled contrastive multi-view clustering with high-order
Student. The probable error of a mean. Biometrika, pp. 1-25, 1908.

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Su, G., Zheng, K, Zhao, T., and Yin, J. CLBAR: addressing representation contamination in multimodal healthcare analytics. In KDD (1), pp. 1289-1300. ACM, 2025.
Xiao, F,, Cai, S., Chen, G., Jagadish, H., Ooi, B. C., and Zhang, M. Vecaug: Unveiling camouflaged frauds with cohort augmentation for enhanced detection. In Proceed- ings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 6025-6036, 2024.
Sun, J., Wang, F., Hu, J., and Edabollahi, S. Super- vised patient similarity measure of heterogeneous pa- tient records. SIGKDD Explor. Newsl., 14(1):16-24, dee 2012. ISSN 1931-0145. doi: 10.1145/2408736. 2408740. URL attps://dok www.dii...g.....o. 2408736. 2408740.
Xie, J., Girshick, R., and Farhadi, A. Unsupervised deep embedding for clustering analysis. In International con- ference on machine learning, pp. 4784448888 118-487. PMLR, 2016.
Xie, Z., Cai, Q., He, F., Ooi, G. Y., Huang, W., and Ooi, B. C. Cohort analysis with ease. In Proceedings of the 2018 International Conference on Management of Data, pp. 1737-1740, 2018.
Suo, Q., Ma, F., Yuan, Y., Huai, M., Zhong, W., Gao, J., and Zhang, A. Deep patient similarity learning for personal- ized healthcare. IEEE Transactions on NanoBioscience, 17(3):219-277, 2018.
doi: N.11099TNB.2011 2883272622...
Xie, Z., Ying, H., Yue, C.,, Zhang, M., Chen, G., and Ooi, B. C. Cool, cohort online analytical processing system. In 2020 IEEE 36th International Conference on Data Engineering (ICDE), pp. 577-588. IEEE, 2020.
Svirsky, J. and Lindenbaum, O. Interpretable deep clustering for tabular data. In ICML. OpenReview.net, 2024.
Tigabe Tekle, M., Bekalu, A. F., and Tefera, Y. G. Length of hospital stay and associated factors among heart failure patients admitted to the university hospital in northwest ethiopia. Plos one, 17(7):c2278009 2022.
Yan, L., Zhang, H.-T., Goncalves, J., Xiao, Y., Wang, M., Guo, Y.. Sun, C., Tang. X., Jing, L., Zhang, M.. al. An interpretable mortality prediction model for covid- 19 patients. Nature machine intelligence, 2(5):283-288 2020.
Van der Maaten, L. and Hinton, G. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Zang, T., Zhu, Y., Huang, X., Yang, X., Chen, Q. Yu, J., and Tang, F. Enhancing length of stay prediction by learning similarity aware representations for hospitalized patients. Artificial Intelligence in Medicine, 144: 1026600, 2023.
Viner, R. M., Ward, J. L., Hudson, L. D., Ashe, M.. Patel, S. V., Hargreaves, D., and Whittaker, E. Systematic re- view of reviews of symptoms and signs of covid-19 in children and adolescents. Archives of disease in child- hood, 106(8):802-80) 807. 2021.
Zhang, C., Gao, X., Ma, L., Wang, Y., Wang, J., and Tang, W. Grasp: Generic framework for health status represen- tation learning based on incorporating knowledge from similar patients. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 715-723, 2021.
Wang, F. and Sun, J. Psf: a unified patient similarity evalua- tion framework through metric learning with weak super- vision. IEEE journal of biomedical and health informat- ics, 19(3):1053-1060, 2015.
Wang. S., Ren, P., Chen, Z., Ren, Z., Ma, J., and de Rijke, M. Order-free medicine combination prediction with graph convolutional reinforcement learning. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management, pp. 1623-1632. 2019.
Zheng, K., Chen, G., Herschel, M., Ngiam, K. Y., Ooi, B. C., and Gao, J. PACE: learning effective task decomposition for human-in-the-loowe healthcare delivery. In SIGMOD Conference, pp. 2156-2168. 68. ACM, 2021.
Zheng. K., Cai, S., Chua. H. R., Herschel, M., Zhang, M.. and Ooi, B. C. Dyhealth: Making neural networks dy- namic for effective healthcare analytics. Proc. VLDB Endow, 15(12):34453444 2022.
Watanabe, C. and Kameoka, H. X-de: Explainable deep clustering based on learnable spectrogram templates. Neu- ral Computation,
2021.
Wu, J., Chen, X., Xu, K., and Ii, S. Structural entropy guided graph hierarchical pooling. In International con- ference on machine learning, pp. 24017-24030. PMLR, 2022.
Zheng, K., Chua, H.-R., Herschel, M., Jagadish, H., Ooi, B. C., and Yip, J. W. L. Exploiting negative samples: a catalyst for cohort discovery in healthcare analytics. In Proceedings of the s1 International Conference on Machine Learning, pp. 61287-6330, 2024.
Wynants, L., Van Calster, B., Collins, G. S., Riley, R. D., Heinze, G., Schuit, E., Bonten, M. M., Dahly, D. L., Damen, J. A., Debray, T. P., et al. Prediction models for diagnosis and prognosis of sovid-19; systematic review and critical appraisal. hmj. 369, 2020.
Zhu, J. Deng, F. Zhao, J., and Chen, J. Adaptive eagregation-distillattn autoencoder for unsupervised anomaly detection. Pattern Recognition, 131:108897, 2022.
12

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healtheare Analytics
Zhu, J., Deng, F., Zhao, J., Liu, D., and Chen, J. Uaed: Un- supervised abnormal emotion detection network based wearable mobile device. IEEE Transactions Network Science and Engineering, 10(6))3682-360 -399,, 2023.
A. Notations
ensure clarity and consistency throughout the paper, a notation table has been provided. This table serves quick reference for readers to easily understand and recall the meaning of each symbol.
Table 8. Frequently used notations
Notation
Mcaning
The all patients
P4
i-th patient in set of patient P Number of patients
N
Set visits for patient
2
i-th visit V,
Number visits for patient p i-th diagnosis code within visit similar to m, and I Number diagnosis codes for visit similar to and N Diagnosis codes visit similar to M, and Cs
N
De
De
the sets diagnosis codes across all visits for patient Pt network to generate the visit-level representation
similar to other functions fMs: Jcc fo and fMC
fo
the similarity prediction network in PseudoSim Training
fan
the encoder used in Global Inter-cohort Modeling
face
the decoder used in Global Inter-cohort Modeling
Representation of diagnosis codes of visit similar to RM, and R... Representation of visit
Rd.
R.
Rp
Representation of patient based on R.. pscudo similarity between patient and p)
S(p..p)
St
Positive patients for patient based on their similarity Negative patients for patient based their similarity Medical ontology
Generated cohorts
RE
Initial representation of patient derived from backbone representation of patient p in Local Intra-cohort Modeling Global Inter-cohort Modeling
R
R
representation of cohort Number of patients for Sttt Sp
Kp
al
attention for Local Intra-cohort Modeling the atlention for Global Inter-cuhort Modeling
ac
B. Extended Related Work
Traditional Cohort Study. Cohorts have been broadly applied across various practical domains analyze group- based patterns and outcomes (Xiao et al., 2024), with the help of cohort processing techniques (Jiang et al., 2016; Xie et al., 2018; Cai et al., 2018; Xie et al., 2020). Among these domains, healthcare stands out as an especially criti- cal area where cohort analysis plays central role. Cohort studies in healtheare are a type of observational research design that longitudinally tracks individuals sharing spe- cific characteristics over time (Eldredge, 2002; Setia, 2016). The primary objective of these studies is to investigate the emergence of particular health-related outcomes or events. Cohort studies hold substantial significance in healthcare research, they facilitate the examination of relationships
13

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
entific discovery. As deep clustering continues evolve, addressing interpretability concerns will be crucial for their practical adoption in real-world applications.
Features
Diagnosis
Medication
C. EHR data
Age
The introduction EHR data structure is illustrated Fig- ure 5. EHR data encapsulates patients' medical histories, providing holistic view of their healthcare journey. As centralized and structured repository, EHR data integrates various aspects of patient care, including detailed records of hospital admissions, outpatient visits, and emergency depart- ment encounters. EHR data comprises a diverse set of in- formation, such as clinical diagnoses, preseribed treatments, medication history, laboratory test results, imaging reports, and clinical notes. Furthermore, EHR systems facilitate longitudinal tracking of patient health, enabling clinicians to monitor disease progression, assess treatment efficacy, and support clinical decision-making EHR data serve as a critical resource for personalized medicine, healthcare analytics, and large-scale population health studies.
Time
Patients
Figure 5. Complex EHR data structure from three dimensions. EHR data consists of time-series medical information that varies across patients.
between potential risk factors or interventions and various health outcomes (Paneth & Monk, 2018; Song & Chung, 2010; Cai et al., 2024; Zheng et al., 2024). For example, substantial and demographically diverse cohort of HIV- infected treatment-naive participants is leveraged explore correlations between viral load and amino acid entropy, considering variables such as sex, age. race, duration of infection, and HIV population structure (Gabrielaite et al., 2023). However, such traditional cohort studies in health- care generally employ a coarse parse-grained approach, leading to an insufficient exploration of nuanced patient insights.
D. Dataset Descriptions and Statistics
MIMIC-III (Johnson et al., 2016) is publicly available medical database, readily available to researchers world- wide. MIMIC-III encompasses 53,423 distinct hospital admissions, catering to adult patients aged 16 years or above, who were admitted to critical care units in the Beth Israel Deaconess Medical Center between 2001 and 2012. Each admission comprises varying number of vital signs, diag- nosis codes, medication codes, and laboratory test codes. Further, there are 2,083,180 de-identified notes linked to these admissions.
Deep Clustering. Clustering is fundamental technique to group similar data points together based on their intrin- sic structure. The clusters provide a comprehensive global characterization of instances, which could be used for in- terpretable feature learning (Svirsky& Lindenbaum, 2024), dimensionality reduction (Ros & Riad, 2024), anomaly de- tection (Zhu al., 2022; 2023) and community detection (Li al., 2017; Wu al., 2022). With the rapid development and remarkable success of deep learning, deep clustering has emerged as powerful paradigm to enhance clustering performance through deep neural networks. DEC (Xie et al., 2016) refines cluster assignments using an autoencoder with soft assignment. DIVIDE (Lu et al., 2024) employs high- order random walks progressively identify reliable data pairs. Despite these advancements, interpretability remains significant challenge in deep clustering. Unlike traditional clustering approaches such as K-Meam and hierarchical clustering, where cluster assignment rules are explicitly defined, deep clustering methods typically operate within high-dimensional latent spaces, making difficult to inter- pret how individual samples are grouped. To address this issue, X-DC (Watanabe & Kameoka, 2021) interprets deep clustering as fitting learnable spectrogram templates an in- put spectrogram followed by Wiener filtering. IDC (Svirsky & Lindenbaum, 2024) introduce a self-supervised approach to enable sample-specift feature selection, thereby refining cluster assignments. Interpretability is of vital importance in critical domains such as healthcare, finance, and sci-
MIMIC-IV (Johnson et al., 2020) is an openly accessible medical database comprising patient admission data col- lected from 2008 to 2022. It includes contemporary data and is structured with modular data organization, with an emphasis on data provenance. Its modular design facilitates the independent use of disparate data sources as well as their combined utilization.
Diabetes130 Dataset (Clore John & Beata, 2014) consists of ten years (1999-2008) of clinical care data from 130 U.S. hospitals and integrated delivery networks. The dataset is provided for the pattern analysis in historical diabetes care, with the potential to achieve safe and personalized healthcare for patients.
The detailed statistics for MIMTC-TTT, MIMTC-TV IV and Diabet estee datasets, including the numbers of patients, visits, and medical codes are presented in Table 9.
14

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Each dataset contains diverse types of clinical data, and tailor data usage each backbone. Med2Vec and MiME primarily leverage medical codes while ClinicalBERT uti- lizes clinical notes to model the representations of EHR data.
Table Data statistics of MMIMI MMIC-111 1
and MIMIC-IH IV.
Dataset
MIMIC-III
MIMIC-IV
Diabetes130
38,597
299.713
71,518
V
53,423
431,231
101,766
notes
2,000,000
su 330,000
Avg. # of V per patient
1.38
1.44
1.42
Baselines are used to compare with NeuralCohort for fair evaluation. KNN, K-Means, DEC, DEKM, GRASP, DGLoS and IDC are adopted baseline methods in our evaluation.
unique
6,984
25,809
915
unique M
4,686
6,911
24
# unique C
726
1,623
Avg. # of D per visit
12.2
11
2.98
Avg. # of M per visit
78
36
1.95
KNN selects the nearest KN nodes for each node its candidate neighbors where KN is the neighbor size in the Local Intra-cohort Modeling. After obtaining the initial representation Ro. KNN is directly employed to identify the KN neighbors, which are then aggregated using the mean function update the representations of the nodes.
Avg. # of C per visit
521
274
0.22
E. Downstream Tasks
Hospital Readmission Prediction Task involves predicting whether a patient will be readmitted to the hospital within 30-day timeframe per visit (Huang et al., 2019). This cross- visit prediction task could provide healthcare professionals with early insights into patients who are at an elevated risk of being readmitted, allowing for timely interventions and improved patient outcomes.
K-Means groups all the nodes into |C] clusters, where |C] represents the number of cohorts. The nodes within the same group share identical neighbors. Follow- ing the acquisition of the initial representation Ro, K- Means employed to form [C] clusters. Subsequently, a random selection of K N neighbors is performed for mean aggregation, leading the update of the node's representation.
Long Length-of-Stay (LOS) Prediction Task pertains classification task that predicts whether the patient will re- main hospitalized for a duration exceeding seven days (Als- inglawi et al., 2022). Within-visit long LOS prediction could aid in optimizing resource allocation, improving pa-
DEC (Xie et al., 2016) is an unsupervised clustering method that jointly learns feature representations and cluster assignments using deep neural networks. DEC consists of two primary phases: pretraining and clus- tering. In the pretraining phase, DEC uses deep au- toencoder to initialize a low-dimensional latent space and during the clustering phase, DEC refines both the latent space and cluster assignments iteratively.
tient flow, and enhancing patient care planning.
F. Backbones and Baselines
Backbones are employed to derive the initial representa- tions R.. In the experiments, Ved2Vec, MiME, and Clini- calBERT are employed as baselines for performance evalu- ation.
DEKM (Guo al., 2021) a deep clustering method that trains an autoencoder to generate embedding space, transforms the embedding space to a new space that reveals the cluster-structure information and optimizes the representation to increase the cluster- structure information in the new space.
Med2Vec (Choi et al., 2016) not only models the co- occurrence information of medical codes but also cre- ates a hierarchical structure based on the sequential order of visits to learn the representations.
MiME (Choi et al., 2018) constructs hierarchical architecture that decomposes the representations into treatment level, diagnosis level, visit level, and patient level, order to render the learned representations interpretable and effective.
GRASP (Zhang et al., 2021) is designed harness the information from similar patients through the rep- resentations of backbones via the Gumbel- -ax tech- nique and employs GCN with an inter-cohort graph to enhance representation learning. However, GRASP merely considers the relationship between cohorts, which limits its performance. In our evaluation, we implement GRASP using its open-source codes with the original settings adopted.
ClinicalBERT (Huang et al., 2019) adapts the BERT model for healthcare, focusing solely on medical notes. To handle long sequences of medical notes, Clinical- BERT splits them into fixed-length slices, meaning that the slices within a visit share the same label. The final predictions are computed by aggregating the predic- tions from each slice.
DGLoS (Zang et al., 2023) implements a module grounded in graph representation learning to produce similarity-awarr representations of patients, thereby
15

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healtheare Analytics
1.00
BERT
Med2Vec
MIME
CinicalBERT
Med2Vec
MIME
0.675
0.650
0.90
0.625
AUPRC
AUPRC
0.600
0.85
0.80
0.550
0.75
0.525
0.500
0.70
11
33
15
1.7
19
5
11
1.3
1.5
17
19
(a) Readmission Task
Long LOS Task
Figure 6. Effects of hyperparameters Kr for both tasks the MIMIC-II dataset.
0.700
1.00
ClinicalBERT
Med2Vec
MIME
ClinicalBERT
Med2Vec
MIME
0.675
0.95
0.650
0.90
0.625
AUPRC
AUPRC
0.600
0.85
0.575
0.80
0.550
0.75
0.525
0.500
0.70
900
1000
1100
1200
1300
900
1000
1100
1200
1300
(a) Readmission Task
(b) Long LOS Task
Figure 7. Effects of the cohort numbers for both tasks on MIMIC-III dataset.
enhancing predictive accuracy. In DGLoS, patients are represented as nodes in the graph, with patient similarity as the weights of the edges. We adopt the deep modeling module, graph modeling module and prediction module with the backbone achieve fair comparison.
G. Experimental Settings
All the experiments are conducted on a server with Intel(R) Xeon(R) W-2133 CPU @ 3.60GHz, 64G memory, and NVIDIA GeForce RTX 2080 Ti.
All datasets are pre-processed with similar logic to the ex- perimental settings (Huang et al., 2019). We randomly split each dataset into training, validation, and test sets with a ra- tio of 8: 1: For each model, we train 20 epochs (the mod- els all converge within 20 epochs). The Pre-context Cohort Synthesis Module and Biscale Cohort Learning Module trained in serial. In the Pre-context Cohort Synthesis Mod- ule, we employ icdcodex (Fisher, 2020) and GloVe (Pen- nington et al., 2014) as 9p and gs in the Path and Semantics components of Hierarchical Visit Engine to distinguish sim- ilar paths the ontology tree and capture the semantics, respectively. JD. JM- Sc. JME, fv. Saut and S are all standard MIP layers. For cohort derivation, K-Means is adopted derive the initial centroid of each cohort. The dimensions of embedding for diagnosis, medication and
IDC (Svirsky Lindenbaum, 2024) is deep cluster- ing model that accurately assigns clusters on tabular data. IDC employs self-supervised approach train autoencoder and gating network, enabling sample- specific feature selection by learning sparse, local fea- ture representations that reconstruct data efficiently and then IDC refines cluster assignments using a clustering head optimized through coding rate reduction loss to achieve compact and separaed clusters.
16

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Table 10. Overall performance of NeuralCohort against baselines for readmission prediction on the MIMIC-IV IV dataset.
Table 11. Overall performance of NeuralCohort against baselines for long LOS prediction on the MIMIC-IV dataset.
Readmission Task on MIMIC
Long LOS Task MIMIC -IV
Model
Model
AUPRC
AUROC
Accuracy
AUPRC
AUROC
Accuracy
ClinicaBBRRT
0.612+0.006
0.629+00.004
56.2%
CliniaaBBRR
0.650+0.004
0.578+0.000
38.1%++0.355
+ KNN
0.614=0.004
0.003
56.4%::0.2%%
KNN
0.654=0.003
0.581=0.005
58, 1%.
+ K-Means
0.614=0.004
0.629:10.003
56.4%::00.2%
K-Means
0.651=0.005
0.580=200004
58. 08.0%
+ DEC
0.609:1..00
0.635:10.002
56.5%
DEC
0.647::0004
0.582110.03
58. 1..2%%
+ DEKM
0.617=0.004
0.631
55.9%:=0.2%%
+ DDEKM
0.662=000003
0.589=0.004
59. 09.2%%=+0.44%
+ GRASP
0.626=0.005
0.641=000004
0..3%
GRASP
0.659=0.004
0.590=0.003
58. 5888%%+-0.5%
DGLoS
0.62310.002
0.619110000
57.1%
DGLoS
0.671100008
0.55810.0006
58.5% 4%%
+ IDC
0.627
0.643110003
57.5%
IDC
0.669100005
0.596110.004
59.2 59.2%-10.225
+ NeuralCohort
0.643
0.657+00003
5..0%
NeuralCohort
0.724=0.003
0.635=0.002
61.8%++00.6%
Med2Vec
0.561=0.007
0.6171:0004
05.4%
Mcd2Vec
0.834=0.005
0.815=0.007
62.5%;:10.66%
+ KNN
0.560+0.0044
0.618+00.044
0..3%
KNNN
0.839+00009
0.820+0.004
6..6%
+ K-Means
0.564++0.004
0.619++0.033
5.22%
K-MMeans
0.831+0.004
0.814+0.004
6..5%
+ DEC
0.557=00005
0.612=#0.01
55. 0..3%
DEC
0.827=00003
0.810=0.002
62. 0.77%
+ DEKM
0.562.10.007
0.61811.000
5..6%
DEKM
0.84110.002
0.822110.006
64.5%.10.4%
+ GRASP
0.55810.000
0.624.10.002
55.6%.10.2%%
GRASP
0.839110.003
0.82410.005
63.0%.10.33%
+ DGLoS
0.563=0.005
0.573+00.006
56.5%
DGLoS
0.837=00003
0.784=00003
64.2%:::0.2%%
+ IDC
0.566=00004
0.621+0.004
56.4%%++.05%
IDC
0.842=000006
0.823+0.008
65.8%:++0.2%
+ NeuralCohort
0.579
0.63910.005
57.6%
NeuralCohort
0.855_10.003
0.833_0.004
68.1%.10.3%
MiME
0.572+00.05
0.624+0.004
55.6%
MIME
0.867+0.005
0.849+00.007
76.9%
+ KNN
0.571=0.003
0.619:10.004
05.3%
KNNN
0.867=0.002
0.848=0.004
76. 0..2%
+ K-Mcans
0.501=0.002
0.621
55.5%
K-Mcans
0.869=0.004
0.850=0.003
76. 0,6
+ DEC
0.504
0.603-10.002
55.4%
DEC
0,867=10.02
0.856:10.004
77. 77.2%%:10.44%
DEKM
0.500
0.619110.004
55.2%
+EKKM
0.874:10.002
0.842:100066
76. 76.4%.110.44
+ GRASP
0.565=0.004
0.617=00004
55.2%==0.3%
GRASP
0.873=0.003
0.859=0.004
77.5% 0.5%%
+ DGLoS
0.582110.066
0.5791100033
50.2%
+GLLoS
0.872.10.004
0.83710.002
78. 2..40.%%%
+ IDC
0.57710.006
0.62411..05
10.3%
+ IDC
0.873110.002
0.85310.004
77. 07.5%
+ NeuralCohort
0.604=00.005
0.641+000066
57.9% 0.55
NeuualCCoort
0.889=0.0003
0.868=0.004
laboratory test, visit and patient 128, 64, 64, 128 and 128. The dimensions of Rg, R2, Rg, and Rinal are the same dimension as the setting of the backbone. Apes Ars and Aco are set to 0.1 in our experiments. We use the Adam optimizer with an initial learning rate of Ir = le 1e-3. 3. Hyper- parameters K2- = 5, |C] = 1100 remain consistent across all datasets and downstream tasks.
is excessively large, dissimilar patient pairs may be incor- rectly identified as similar. Conversely, when K₂ is small, NeuralCohort fails effectively learn representations due to insufficient maximization of Mutual Information Neural Estimation. Finally, these figures suggest that the optimal choice of Kp is approximately 5, either increasing or de- creasing these values will result in degraded performance.
Effects of Cohort Number (C]. We evaluate the impac of cohort numbers on both tasks using the MIMIC-III dataset. The corresponding experimental results are de- picted in Figure 7. It is clearly shown that the performance of all evaluated models first increases and then decreases as the number of cohorts increases. The reason for such phenomena related to the target of NeuralCohort, i.e., ex- tracting cohort information for integration with backbones to achieve boosted performance. If the number of cohorts is smaller, is harder for the Global Inter-cohort Model- ing to capture the difference between the coarse- grainedd cohorts and if the size of cohorts tends to be larger, it is more likely to aggregate dissimilar patients, which leads to noisy information when utilizing the Local Intra-cohort Modeling. Conversely, with the number of cohorts increas- ing, fewer patients will be assigned to a cohort, which may
H. Sensitivity Study
We conduct a sensitivity study to examine the impact of the key hyperparameters |C| and Kp on the performance of NeuralCohort for both tasks using the MIMIC-IN III dataset, with all the backbones.
Effects of Hyperparameter Ky. In our investigation of the crucial hyperparameters' influence on NeuralCohort's per- formance, we first focus on the number of positive samples or negative samples Kp in the Pre-context Cohort Synthesis Module. The AUPRC results, obtained by varying these hyperparameters using three datasets for both tasks, are il- lustrated in Figure 6. From this figure, we observe that the performance of NeuralCohort fluctuates at the beginning and subsequently decreases with the increase of Kp- When Kp
17

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics
Table 12. Overall performance of NeuralCohort against baselines for readmission prediction on the D Diabetes13s 30 dataset.
Table 14, Parameter Size and Running time of NeuralCohort.
Model
Parmmeter Size (M)
Training Time (6)
Inference Time (ins)
Readmission Task on labbetss111 30
Model
CinicalBERT
109,4
4299.6
4.32
AUPRC
AUROC
Accuracy
KNN
109,4
6962.8
6.53
K-MMans
109,4
5525.2
Med2Vec
0.623
0.628.10.003
59.4%
+ DDCC
114.1
6225.2
5.95
+ KNN
0.625-0.003
0.625++0.044
59.3%
-EKM
115.4
6851.7
6.37
7.29
+ K-Means
0.622+0.004
0.623++0.003
59.0%%+002%%
+ IIGL o8
113.6
8418.8
7.93
+ DEC
0.602
0.605
09.4%
IDC
118.0
6546.2
7.04
DEKM
0.624+00.004
0.63 +0.0003
59.7%++0.5%%
NeuralCohort
121.3
6432.4
7.16
+ DGLoS
0.629+00.005
0.614++0.007
59. +9.5%
+ GRASP
0.628=0.003
0.632:+0.003
59 5%.4%%
I. Experimental Results on MIMIC-IV and Diabetes133 30 Datasets
+ IDC
0.627=0.006
0.634=00003
60.1% 6..3%%
+ NeuralCohort
0.641 1 0.003
0.004
61.5% 1 0.3%
MiME
0.662+40.004 0.657=0.002
62.7%:+0.3%%
+ KNN
0.657.10.004 0.649-1.000
0..4%
Experimental results of MIMIC-TV and Di Diabee 1 130 datasets are displayed in Table 10 to Table 13. It should be noted that there is no clinical text data in Diabetes 1 30 dataset, therefore, backbone ClinicalBERT cannot be em- ployed in experiments with Di Diabetesl
+ K-Means
0.608
0, 0.664400006
63 03.2%%
+ DEC
0.662+00005
0.669++0.002
63.4%%+0..%%
DEKM
0.658=0.002
0. 0.000
0..7%:10.2%
+ GRASP
0.653-10.002 0.656=1..05
0..5%
dataset. The results demonstrate that NeuralCohort consistently outper- forms the baselines, confirming its effectiveness and gener- alizability.
+ DGLLS
0.662:+0.006
0. 0.649:+0.007
6..2%
+ IDC
0.664-10.006
0.675:+0.003
6..7%
+ NeuralCohort
0.676=0.0066
0.689=+0.004
65.4%==0.%%
J. Efficiency Study
Table 13. Overall performance of NeuralCohort against baselines for long LOS prediction the Diabetes 30 dataset.
Table 14 compares the efficiency of NeuralCohort with back- bone ClinicalBERT and baselines in terms of parameter size, training time, and inference time. NeuralCohort ex- hibits a comparable parameter size but achieves a reasonable training time, significantly outperforming computationally intensive methods such as DGLoS, In terms of inference time, NeuralCohort is competitive with DGLoS, GRASP and IDC but shows longer inference time than other meth- ods. However, the trade-off in inference speed is balanced by the potential gains in modeling accuracy and represen- tational power offered by NeuralCohort. Overall, the ta- ble demonstrates that NeuralCohort provides effective balance between parameter complexity and computational efficiency, making it a promising choice for applications requiring cohort-aware learning.
Long LOS Task on Diabet es essl3 30
Model
AUPRC
AUROC
Accuracy
Med2Vec
0.696 0.006
0.005
65.4%
+ KNN
0.700=0.004
0. 0.706=0.002
65.4%
+ K-Means
0.701
0.709
65.8%
+ DEC
0.700
0. 0.702-00.04
65.2%
+ DEKM
0.692=0.006
0.716=0.005
66.2%%+0.4%
+ GRASP
0.6931..003
0.700-0.003
65.4%
+ DGLoS
0.705+0.003
0..006
66.0%
+ IDC
0.71710.007
0.714=0.004
66.5%--10.4%
+ NeuralCohort
0.732: 0.004
0.739 0.003
68.2%
MiME
0.719=0.003
0.729 +0.005
78.3%
+ KNN
0.71710.002
0.72510.003
78.0%
+ K-Mcans
0.720=0.004
0.732=0.005
08.5%
+ DEC
0.724=0.008
0.739=0.004
79.4%%
+ DEKM
0.718100002
0.724=0.005
78.7%
K. Extened Visualizaion
+ DGL.S
0.727+00007
0.704+0.002
78.9%
+ GRASP
0.724=00005
0.731=0.005
7..5%
Extended visualization of all the methods is illustrated in Figure 8. The significant overlap among the points in Fig- ure 8 (a) through (h) suggests that both the backbone and the baselines exhibit limited discriminative power. However, the incorporation of NeuralCohort substantially reduces this overlap, thereby enhancing the distinctiveness of the co- horts. This indicates that NeuralCohort effectively encodes the data into representation similarity space.
+ IDC
0.732+0.005
0.740=0.002
79.6%
+ NeuralCohort
0.745-00004
0.762-0.007
82.1 0033%
result in inadequate intra-cohort information. Therefore, the number of cohorts exerts a vitally important influence on both tasks using the MIMIC-III dataset, which means a suitable value should be selected for the number of cohorts. In our setting, the best [C] should be about 1100.
18

---

NeuralCohort: Cohort-aware Neural Representation Learning for Healtheare Analytics
(a) Med2Vec
(b) KNN
(c) K-Means
(d) DEC
(e) DEKM
(f) GRASP
(g) DGLoS
(h) IDC
(i) NeuralCohort
Figure The extended t-SNE visualization of the final representations hefore prediction selected eight-cohort points, where backbone is Med2Vec.
19


Reviews:
[{"review_id": "YUOiFNeLzZ", "reviewer": "Reviewer_uWw6", "summary": "The paper proposes NeuralCohort -- a cohort-aware neural representation learning method designed to improve electronic health record (EHR) analysis. It addresses the challenges of fine-grained cohort segmentation and effectively utilizes both intra- and inter-cohort information. By incorporating the Pre-context Cohort Synthesis for pseudo patient similarity prediction task, and Biscale Cohort Learning modules to learn the intra- and inter- cohort interactions, NeuralCohort enhances existing models, boosting their performance. Its approach not only improves predictive accuracy but also provides valuable medical insights for EHR studies.", "questions": "", "limitations": "", "rating": 4, "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": "Thank you for your thoughtful and encouraging review. We sincerely appreciate your recognition of NeuralCohort as a cohort-aware neural representation learning framework that effectively addresses the challenges of fine-grained cohort segmentation and the modeling of both intra- and inter-cohort relationships. We are especially pleased that you found the two-stage design — modeling at the patient visit level followed by biscale cohort learning — to be both novel and interpretable. Our goal is to not only improve predictive performance across benchmark EHR datasets such as MIMIC-III, MIMIC-IV, and Diabetes130, but also to provide clinically meaningful insights through interpretable cohort structures. Your support affirms the value of our approach for advancing cohort-based analysis in real-world healthcare applications.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes, the methods had been extensively tested on popular MIMIC-III, MIMIC-IV and Diabates130 datasets which are benchmarks for the EHR domain.", "theoretical_claims": "Yes mostly, did not go into details.", "experimental_designs_or_analyses": "The experiments are well designed to showcase the effectiveness of the approach. The suuplementary material has more addendum experiments.", "relation_to_broader_scientific_literature": "Prior works have not sufficiently explored the cohort formulation for EHR tasks. Patient similarity has been tackled based on expert labellings and more recently cosine similarity between embeddings from EHR data. This work models patient's visit level representation to then model cohort interactions. It shows the worth of cohort studies. This could be useful for most downstream healthcare analysis tasks.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "Strengths\n- The ideas of modeling in 2 stage manner - first visit level then modeling the intra- and inter- cohort signals are novel.\n- The method is clearly effective and also interpretible which is a key consideration for healthcare domain.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "n/a", "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": "gmztO3lAOH", "reviewer": "Reviewer_6GpW", "summary": "This paper proposes NeuralCohort, a cohort-aware neural representation learning method for healthcare analytics. The approach segments patients into fine-grained cohorts and captures both intra- and inter-cohort information through a two-module paradigm.", "questions": "", "limitations": "", "rating": 2, "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": "Thanks for your constructive feedback. This is the link to our supplemented results: https://anonymous.4open.science/r/ICML2025-93F9.\n\nQ1: Lack clear evidence on how cohort insights translate into clinical interventions. \n\nPlease refer to Appx. K for detailed analysis. To be more specific, Cohort #1 includes patients with cardiovascular conditions and is at high risk for acute cardiac events and extended hospital stays[1]. Early identification enables the prioritization of telemetry beds, cardiology consults, and monitoring. Clinically, this enables timely diuretics, echocardiography, and discharge planning with heart failure follow-up. Cohort #2 features patients with chronic metabolic and hematologic conditions that often co-occur and need interdisciplinary care. Identifying this cohort allows hospitals to mobilize diabetes educators, lipid clinics, and anticoagulation monitoring services. Clinically, it enables insulin titration, lipid therapy, and bleeding risk management to reduce complications and improve outcomes. Cohort #3 is characterized by renal and urological issues requiring frequent labs, fluid monitoring, and nephrology support. Early identification allows hospitals to allocate renal panels, schedule imaging for uropathy, and prepare dialysis resources—helping manage AKI risk and avoid care escalation. Cohort #4 presents complex chronic and acute conditions, requiring coordination across pulmonology, nephrology, endocrinology, and infectious disease. Identifying this cohort supports planning for respiratory support, hormone therapy, infection control, and opioid withdrawal. Operationally, it enables resource allocation for respiratory teams, isolation rooms, and endocrine/renal labs.\n\nQ2: Practical applications beyond the MIMIC dataset\n\nWe evaluated NeuralCohort on a real-world Diabetes130 dataset, which is shown in Appx. I. We also analyze the practical application, similar to Appx. K but on the readmission task. The distinctive features are shown in Table A.\n\nCohort #1 is characterized by longer hospital stays, insulin use, and frequent prior outpatient visits, indicating complex or unstable diabetes requiring intensive inpatient care. These patients often need time to stabilize their glucose levels, adjust medications, and manage comorbid conditions. They benefit from diabetes management plans, glucose monitoring, and endocrinology support. Hospitals should allocate extended-stay beds, involve diabetes educators, and initiate early discharge coordination. Close follow-up via outpatient care or telemedicine helps reduce readmission risk. Cohort #2 comprises low-acuity patients. Likely admitted for scheduled care or monitoring, they benefit from preventive care, education, and routine screening. Hospitals can streamline care using standardized protocols, fast-track discharge workflows, and minimal specialty consults, optimizing resource use for this stable cohort. Defined by a high rate of emergency admissions, younger age, and few prior inpatient visits, cohort #3 likely represents underserved or poorly engaged patients who rely on the ER for primary care needs. They may face social or behavioral health challenges. Interventions should focus on screening for social determinants, community referrals, and behavioral health support. Hospitals can prioritize care navigation and ER diversion efforts to reduce unnecessary admissions and alleviate emergency departments burden. Cohort #4 consists of high-risk, high-utilization patients. Tailored interventions should focus on intensive transitional care management, including medication reconciliation, post-discharge follow-ups within 48–72 hours, and multidisciplinary care planning. Hospitals can assign care managers, enroll patients in chronic care programs, and coordinate outpatient services to reduce readmissions and enhance long-term outcomes.\n\nQ3: Novelty of the approach.\n\nIt should be noted that NeuralCohort is an application-driven ML submission targeting practical healthcare applications, where originality could mean \"a novel combination of existing methods to solve the task at hand so as to match the needs of the user\" rather than wholly novel methods[2]. Our key innovation lies in the cohort-aware representation learning framework, which explicitly disentangles and models both local intra-cohort and global inter-cohort structures — a perspective not addressed in prior work. Moreover, NeuralCohort not only enhances predictive performance but also yields clinically interpretable embeddings that align with real-world needs, such as intervention planning and resource optimization.\n\n[1] Tigabe Tekle, Masho, Abaynesh Fentahun Bekalu, and Yonas Getaye Tefera. \"Length of hospital stay and associated factors among heart failure patients admitted to the University Hospital in Northwest Ethiopia.\" Plos one 17.7 (2022): e0270809.\n\n[2] https://icml.cc/Conferences/2025/ReviewerInstructions", "claims_and_evidence": "The paper presents a well-supported evaluation of its method using multiple EHR datasets, demonstrating performance improvements. However, it lacks clear evidence on how cohort insights translate into clinical interventions.", "methods_and_evaluation_criteria": "The methodology integrates existing techniques (hierarchical modeling, GNNs, contrastive learning) without introducing fundamentally new innovations.", "theoretical_claims": "NA", "experimental_designs_or_analyses": "The experimental setup has no significant methodological flaws identified. However, the study would benefit from demonstrating practical applications beyond the MIMIC dataset.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "The paper's strength lies in its comprehensive evaluation across multiple EHR datasets, demonstrating consistent performance improvements when integrated with existing backbone models.\n\nI do not find significant methodological or experimental flaws in the study. However, my primary concern is the novelty of the approach. The method appears to be an integration of existing techniques—hierarchical modeling, graph neural networks, and contrastive learning—rather than introducing fundamentally new methodological innovations. Both the core idea and the task itself are not particularly novel compared to existing work, especially given the extensive exploration of similar tasks on the MIMIC dataset.\n\nAdditionally, it remains unclear how the proposed work and the insights derived from cohort segmentation translate into tangible clinical interventions or improved patient management strategies. I suggest that the authors clearly articulate what is conceptually novel beyond the integration of existing techniques. Moreover, providing concrete examples of how the identified cohorts could inform specific clinical decisions or extending the application beyond standard prediction tasks—already well-studied on MIMIC datasets—would strengthen the contribution of this work.", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "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": "Pp2iyW4OoH", "reviewer": "Reviewer_8wHJ", "summary": "This paper introduces the NeuralCohort framework, a novel method for neural cohort generation and selection based on two distinct strategies that model and generate (local) intra- and (global) inter-cohort information using a pre-context synthesis module employing pseudo similarity and a biscale cohort learning module to fuse and predict downstream classification tasks.", "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 appreciate your insightful comments. This is the link to our supplemented results: https://anonymous.4open.science/r/ICML2025-93F9.\n\nQ1: The possibilities of techniques. How would SimCLR help?\n\nKey components like reverse-time attention and mutual information have been validated in prior work (e.g., RETAIN[1], FASTDeepHit[2]). While other architectures are possible, our choice achieves SOTA performance, generalization, and interpretability. Many of the components are orthogonal to our primary contribution so NeuralCohort can readily incorporate more advanced alternatives as they become available.\n\nSimCLR, where we randomly mask EHR to construct positive pairs, degrades performance as shown in Tables D and E. The reason lies in SimCLR is designed for image data, where augmentations like cropping retain semantic meaning. In EHR, creating label-preserving augmentations is more challenging considering the complex medical semantics of EHR. \n\nQ2: Method only works well for the two datasets? How to use NeuralCohorts for other datasets?\n\nNeuralCohort is tailored for EHR data and suitable for domains where subpopulation patterns are important, such as diabetes or Multiple Sclerosis. We evaluate NeuralCohort on Diabetes130 dataset (Appx. I), with clinical insights detailed in Q2 of Reviewer 6GpW.\n\nWe have done the clinical application on Multiple Sclerosis dataset. The distinctive features are listed in Table B.\n\nCohort #1 shows early neurological abnormalities. They should be prioritized for early initiation of DMTs and placed on close clinical surveillance schedules. Hospitals and neurology departments should ensure these patients have expedited access to high-resolution neuroimaging, neurologist consultations, and structured follow-up. Cohort #2 shows slower progression and may follow atypical paths. A conservative management approach with lifestyle interventions is suitable. Operationally, they can be monitored in outpatient settings with less frequent MRI, reserving resources for higher-risk patients. Patients in Cohort #3 display minimal abnormalities, indicating a stable, low-risk, or non-converting group. The optimal strategy for managing this cohort involves conservative care focused on patient reassurance, symptom education, and minimal intervention unless clinical deterioration is observed. Cohort #4 includes younger patients with positive LLSSEP and VEP, suggesting early subclinical neurological involvement despite low EDSS scores. This profile signals elevated risk for MS conversion. Clinical strategies should focus on prevention through lifestyle counseling, low-risk DMTs, and regular monitoring. Besides, this group warrants access to early neurology resources and is well-suited for clinical trials aimed at delaying progression before disability develops.\n\nQ3: Which dimensions feature span, how to derive the number of clusters, the number of pairs, and how the baselines were tuned? \n\nThe dimensions of embedding for diagnosis, medication and laboratory test, visit and patient are 128, 64, 64, 128 and 128. The dimensions of $\\mathcal{R} _ 0$, $\\mathcal{R} _ L$, $\\mathcal{R} _ G$, $\\mathcal{R} _ {final}$ are the same dimension as the setting of the backbone.\n\nWe use Student's t-distribution as a soft assignment function to map patient to cohort, computing similarity via a t-kernel to produce normalized soft assignments. The number of clusters determines the number of centroids and is initialized using KMeans. A sensitivity study on cohort number and pair numbers is provided in Appx. H. All experiments are repeated five times, and the average results are reported.\n\nThe following protocol is used to ensure fair comparisons. We followed the original papers' recommended hyperparameters or adopted defaults from the official implementation when tuning details were unavailable. To ensure fairness, shared hyperparameters—such as the number of cohorts or clusters—were kept consistent across all methods.\n\nQ4: It is a bit suboptimal to support the claims.\n\nWe explicitly address four key aspects: improved task performance (Sec. 4.2, Appx. I), CH-score for similarity evaluation (Sec. 4.5), t-SNE visualizations of cohort (Sec. 4.5), and clinical validation on MIMIC-III (Appx. K) and Diabetes130 (Q2, Reviewer 6GpW), which could validate the generated cohorts.\n\nQ5: How does Figure 8 look for a different number of clusters? \n\nThe number of clusters is chosen to support clear visualization. Below, we present results for 6 and 10 clusters in Figure A and Figure B. While some baselines form clusters, they lack clear boundaries, whereas NeuralCohort produces more distinct groupings.\n\n[1] Choi, Edward, et al. \"Retain: An interpretable predictive model for healthcare using reverse time attention mechanism.\" Advances in neural information processing systems 29 (2016).\n\n[2] Do, Hyungrok, et al. \"Fair survival time prediction via mutual information minimization.\" Machine Learning for Healthcare Conference. PMLR, 2023.", "claims_and_evidence": "The authors provide interesting experiments and ablation studies for NeuralCohort. Firstly, they claim and demonstrate that NeuralCohort can consistently perform with different backbone neural architectures in their experiments and can outperform current state-of-the-art embedding/clustering methods for two classification tasks, namely \"cross-visit Hospital Readmission Prediction\" and \"within-visit Long\nLength-of-Stay\". This empirically supports their claim that NeuralCohort is a strong encoding/feature aggregation method.", "methods_and_evaluation_criteria": "All architectural design considerations are properly introduced with mathematical formulation. While the formulations fundamentally make sense, some of the decisions \"Hierarchical Visit Engine\" or \"PseudoSim Training\" are not adequately motivated. It is hard to understand what other possibilities have been tested in related literature and why the authors believe that this constitutes the best and/or most robust method to model the discussed concept.  For instance, how would a contrastive learning approach such as SimCLR [1] help or impede the modeling compared to the PseudoSim Training? For instance, it is unclear to me if this architectural design choice is something that works well for the two datasets used in this study or if it is clinically meaningful in other studies, such as Multiple Sclerosis research. I also find it very hard to understand - based on the manuscript and not the codebase - the intermediate steps of the framework and the dimensionality of different feature spaces. This makes it hard to assess if the evaluation of experiments is fair and how baseline algorithms have been tuned or if the experimental design is posed in a way that favors NeuralCohort.\n\n[1] Chen T, Kornblith S, Norouzi M, Hinton G. A simple framework for contrastive learning of visual representations. InInternational conference on machine learning 2020 Nov 21 (pp. 1597-1607). PmLR.", "theoretical_claims": "The manuscript includes no novel proofs but mathematical formulations of architectural design choices.", "experimental_designs_or_analyses": "The authors provide two experiments on three different datasets with three different backbones, which makes the case that their method works reliably and is adaptable to other backbones. \nAs mentioned previously, it is hard to understand from the main paper/appendix which dimensions feature spaces span, how the authors derive the number of clusters for their experiments, and how architectural design choices (e.g. cluster number) affect the baseline methods. I have read the interesting ablation study in the Appendix. Still, to me, it is not ultimately clear how the baseline methods were tuned (if they were tuned) and how robust the baselines are across a parameter set.\n\nWhile the paper claims to provide a method for cohort generation, it is validated on two downstream segmentation tasks that serve as a proxy for this problem. While I am unaware of a better experiment/task to demonstrate it, I believe this is a bit suboptimal to support the claims.", "relation_to_broader_scientific_literature": "The authors claim that NeuralCohort is (one of the) first methods to look into cohort selection apart from traditional methods used by clinicians and in clinical studies. They correctly set their work into the context of neural clustering algorithms, which they also use as baselines. I feel the methods section could explain more about why the authors have decided to model certain relationships the way they have decided to do and how they ground this in prior work.", "essential_references_not_discussed": "I have not found any relevant papers that the authors have missed.", "other_strengths_and_weaknesses": "Strengths:\n- I would like to highlight that the authors have already shared their code in the reviewing phase. \n\nWeaknesses:\n- To me this paper lacks details regarding how to use NeuralCohorts for other datasets and the intended use case - the formation and selection of cohorts in studies.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "- How do you typically select hyperparameters such as a positive number of pairs for PseudoSIM for Neural Cohorts?\n- How were they selected for this paper's experiments? Given the variability of the ablation study in the appendix, are the results in the reported tables the best-case scenario?\n- Were the number of clusters empirically selected for Figure 8, Supplementary Material? How do the plots look for a different number of clusters ? Given that the performance of NeuralCohort in your experiments is better than baselines but not always by a high margin, I am surprised to see such distinct cluster in comparison. \n- Did you assess features in the neural cohort clusters? Are they medically meaningful, which would be significant considering that it should be used in medical cohorts?", "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": "XbrFkF4h0K", "reviewer": "Reviewer_UU4H", "summary": "The paper proposes NeuralCohort, a cohort-aware neural representation learning method for healthcare analytics. It introduces two modules: (1) a Pre-context Cohort Synthesis Module to derive fine-grained cohorts via pseudo patient similarity, and (2) a Biscale Cohort Learning Module to integrate intra- and inter-cohort information. The method is evaluated on EHR datasets (MIMIC-III, MIMIC-IV, Diabetes130) for tasks like hospital readmission and length-of-stay prediction, demonstrating performance gains of up to 8.1% in AUROC when integrated with backbone models like ClinicalBERT and Med2Vec.", "questions": "", "limitations": "", "rating": 2, "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": "Thanks for the detailed reviews. This is the link to our supplemented results: https://anonymous.4open.science/r/ICML2025-93F9.\n\nQ1: Lack validation of fine-grained cohorts. No qualitative analysis of cohort granularity or interpretability; visualizations lack clinical context.\n\nIn practice, clinicians typically use coarse-grained cohorts (Sec. 4.4), making fine-grained grouping like that in Appx. K difficult to define manually. NeuralCohort enables the automatic discovery of such cohorts and identifies key features for clinical insight and decision support.\n\nPlease refer to Appx. K, Q1, Q2 of Reviewer 6GpW and Q2 of Reviewer 8wHJ for the distinctive features of cohorts on three datasets.\n\nQ2: Hyperparameter and cohort derivation choices lack justification; no sensitivity analysis or tuning details.\n\nWe conducted sensitivity analysis (Tables F and C), varying one hyperparameter at a time while keeping others fixed. JS divergence outperformed KL. The optimal values are: $\\lambda_{pcs}$=0.1, $\\lambda_{JS}$=0.1, $\\lambda_{co}$=0.1, dropout=0.1, and learning rate=1e-3. When training, dropout is set to 0.1, and the setting of other parameters is shown in Appx. G.\n\nQ3: Cohort derivation lacks detail, and reliance on pseudo-labels without expert validation is a limitation.\n\nPlease refer to Q3 of Reviewer 8wHJ for cohort derivation details. Pseudo-labeling is widely used in representation learning[1], especially in EHR, where manual labeling is costly and impractical. Expert-defined similarity requires extensive clinical review, which is challenging at scale. Our method remains flexible and can incorporate expert input in future semi-supervised or clinician-in-the-loop settings.\n\nQ4: Baselines are outdated; recent methods and references are missing. No comparison to temporal graph models, contrastive learning, or dynamic cohorting approaches.\n\nWe briefly discuss related methods and will expand in revision. TRANS[2] models  EHR as a temporal heterogeneous graph to capture both temporal and structural information. SimSig[3] uses a contrastive learning approach to learn similar embeddings of patients with physiological signal data. Longitudinal K-Means[4] presents K-Means approaches for subtyping opioid use trajectories from EHR data and interprets the resulting subtypes using decision trees.\n\nWe clarify that temporal graph networks for EHR modeling are complementary to our method and could serve as backbones rather than baselines. Comparisons with SimSig and Longitudinal K-Means in Table G show they underperform our method.\n\nQ5: Limited novelty; the method appears to be an incremental extension of GRASP, which may already capture both local and global cohort-level information.\n\nUnlike GRASP's flat KNN graph, NeuralCohort introduces a distinct, bi-level architecture that explicitly models both intra-cohort and inter-cohort structures through dedicated modules. This enables fine-grained cohort discovery and more effective cohort-aware representation learning. Unlike GRASP, NeuralCohort integrates medical semantics into the learning process, which our experiments show leads to more stable performance—whereas GRASP may degrade backbone performance due to less clinically grounded similarity signals.\n\nQ6: Were ROC-AUC improvements statistically significant?\n\nWe evaluate AUROC difference with MIMIC-III dataset on long LOS task in Table H, showing statistically significant improvements.\n\nQ7: How do confounding factors explain inverse lab test correlations?\n\nCounterintuitive correlations, such as low LDL linked to worse outcomes, may result from confounding factors like medication use. Patients with low LDL often have prior cardiovascular risk and are treated with statins, making low LDL a marker of treated high-risk status rather than low risk. NeuralCohort captures such latent clinical structures by modeling intra-cohort similarities and inter-cohort differences. For instance, it can distinguish between statin-treated high-risk patients and untreated low-risk individuals, offering more clinically meaningful stratification than models treating patients independently. Please refer to Appx. K, Q1, Q2 of Reviewer 6GpW and Q2 of Reviewer 8wHJ for practical applications.\n\n[1] Xi, Liang, et al. \"Semi-supervised time series classification model with self-supervised learning.\" Engineering Applications of Artificial Intelligence 116 (2022): 105331.\n\n[2] Chen, Jiayuan, et al. \"Predictive modeling with temporal graphical representation on electronic health records.\" IJCAI: proceedings of the conference. Vol. 2024. 2024.\n\n[3] Shanto, Subangkar Karmaker, et al. \"Contrastive Self-Supervised Learning Based Approach for Patient Similarity: A Case Study on Atrial Fibrillation Detection from PPG Signal.\" arXiv preprint arXiv:2308.02433 (2023).\n\n[4] Mullin, Sarah, et al. \"Longitudinal K-means approaches to clustering and analyzing EHR opioid use trajectories for clinical subtypes.\" Journal of biomedical informatics 122 (2021): 103889.", "claims_and_evidence": "The claims are partially supported by experiments showing consistent improvements over baselines. However:\n\nProblematic Claim: The assertion that NeuralCohort \"simultaneously fulfills fine-grained cohort division and information exploitation\" lacks direct validation. While ablation studies show performance drops when modules are removed, there is no qualitative analysis of cohort granularity or clinical interpretability.\n\nUnjustified Choices: Hyperparameters (e.g., $\\lambda_{p e s}$ =0.1) and cohort derivation methods (e.g., Jensen-Shannon divergence) are not rigorously justified. Sensitivity analyses for these parameters are missing.", "methods_and_evaluation_criteria": "Methods: The hierarchical modeling of EHR data (diagnoses, medications, lab tests) and pseudo-similarity training are appropriate for temporal healthcare data. However, the cohort derivation process is under-detailed, and the reliance on pseudo-labels (without validation against expert annotations) is a limitation.\n\nEvaluation: Standard tasks (readmission, LOS) and metrics (AUROC, AUPRC) are well-chosen. However, baselines like KNN and K-Means are outdated; comparisons with recent methods (e.g., contrastive learning or graph-based models) are missing.", "theoretical_claims": "Lack of theoretical proof.", "experimental_designs_or_analyses": "Verified: The ablation study (Table 3) confirms the contribution of each module. The consistent gains across datasets (Tables 1–2, 8–11) support generalizability.\n\nIssues:\n\nThe cohort visualization (Figure 8) lacks clinical context (e.g., what distinguishes the cohorts in practice?).", "relation_to_broader_scientific_literature": "The work connects to EHR representation learning (e.g., Med2Vec, ClinicalBERT) and patient similarity research. However, it overlooks recent advances:\n\nMissing References: Temporal graph networks (e.g., TGN) for dynamic EHR modeling, contrastive learning frameworks (e.g., SimCLR) for patient similarity.\n\nGaps: No discussion of how NeuralCohort compares to dynamic cohorting methods (e.g., longitudinal clustering).", "essential_references_not_discussed": "See Relation To Broader Scientific Literature Above", "other_strengths_and_weaknesses": "pros:\n\nS1: Novel integration of cohort dynamics into neural representation learning.\n\nS2: Modular design compatible with diverse backbones.\n\nS3: Rigorous evaluation across multiple datasets and tasks.\n\ncons:\n\nW1: Limited novelty. This work only makes incremental modifications based on GRASP. I believe that GRASP's network architecture is also capable of utilizing both local intra-queue information and global inter-queue information, ultimately extracting comprehensive queue information.\n\nW2: Lack of interpretability analysis for generated cohorts (e.g., clinical relevance).\n\nW3: Incomplete baseline comparisons (e.g., state-of-the-art graph models).\n\nW4: Insufficient details on cohort derivation and hyperparameter sensitivity.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Q1: Were ROC-AUC improvements statistically significant (e.g., via bootstrapping or permutation tests)?\n\nQ2: How were hyperparameters (e.g., learning rate, dropout) selected during fine-tuning?\n\nQ3: How do confounding factors (e.g., medication use) explain inverse lab test correlations (e.g., LDL)?", "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: This paper introduces NeuralCohort, a cohort-aware neural representation learning for electronic health record (EHR) analysis. The method integrates cohort segmentation through two modules: the Pre-context Cohort Synthesis and the Biscale Cohort Learning Module. Evaluations on diverse EHR datasets (MIMIC-III, MIMIC-IV, Diabetes130) demonstrate notable improvements (up to 8.1% AUROC gains) when integrated with baseline models such as ClinicalBERT and Med2Vec.

Reviewers agree that this paper provide a novel integration of cohort dynamics into neural representation learning and has solid evaluation across multiple EHR datasets and prediction tasks.
Decision: Accept (poster)