GRAIL: Geometry-Aware Retrieval-Augmented Inference with LLMs over Hyperbolic Representations of Patient Trajectories

2026-02-13

Technology Category

• Reasoning under Uncertainty: Graphical Models
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning
• Machine Learning: Graph-based Machine Learning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Predicting future clinical events from longitudinal electronic health records (EHRs) is challenging due to sparse multi-type clinical events, hierarchical medical vocabularies, and the tendency of large language models (LLMs) to hallucinate when reasoning over long structured histories. We study next-visit event prediction, which aims to forecast a patient's upcoming clinical events based on prior visits. We propose GRAIL, a framework that models longitudinal EHRs using structured geometric representations and structure-aware retrieval. GRAIL constructs a unified clinical graph by combining deterministic coding-system hierarchies with data-driven temporal associations across event types, embeds this graph in hyperbolic space, and summarizes each visit as a probabilistic Central Event that denoises sparse observations. At inference time, GRAIL retrieves a structured set of clinically plausible future events aligned with hierarchical and temporal progression, and optionally refines their ranking using an LLM as a constrained inference-time reranker. Experiments on MIMIC-IV show that GRAIL consistently improves multi-type next-visit prediction and yields more hierarchy-consistent forecasts.

Problem

• next-visit prediction
• electronic health records
• clinical event forecasting
• hierarchical medical vocabularies
• LLM hallucination

Innovation

• hyperbolic representation
• structured retrieval
• clinical event prediction
• hierarchical medical vocabularies
• LLM-augmented inference