NeuralCohort: Cohort-aware Neural Representation Learning for Healthcare Analytics

Changshuo Liu, Lingze Zeng, Kaiping Zheng, Shaofeng Cai, Beng Chin Ooi, James Wei Luen Yip

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.

Lay Summary:

Electronic health records (EHRs) store large amounts of patient data that can help doctors improve patient care and identify patterns in diseases. One way researchers use this data is through cohort studies, where they analyze groups of patients who share similar traits or conditions over time. However, existing methods often struggle to group patients precisely and make full use of this cohort information, which limits their ability to reveal meaningful insights.

To address this, we developed **NeuralCohort**, a new AI-based method that can better group patients into more detailed categories and learn from both the differences and similarities across these groups. NeuralCohort works alongside existing models and helps them understand patient data more deeply.

Our experiments can help researchers and healthcare providers make more accurate predictions and better decisions based on EHR data.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Applications->Health / Medicine

Keywords: Healthcare, Cohort, Representation Learning, EHR

Submission Number: 3642