HypEHR: Hyperbolic Modeling of Electronic Health Records for Efficient Question Answering

1. Problem Definition

Electronic Health Record (EHR) Question Answering (EHR-QA) involves answering natural language clinical questions based on a patient's longitudinal medical history.

• Current Limitations: Existing approaches generally fall into three categories:
  1. Representation Learning: Sequential or graph-based models that often ignore the intrinsic hierarchical structure of medical data.
  2. Text-to-SQL: Converting questions to SQL queries, which relies on large-scale LLMs and struggles with complex schema alignment.
  3. LLM-based Pipelines: Retrieval-augmented generation using massive models (e.g., GPT-4, Llama-3). While accurate, these are computationally expensive, difficult to deploy in privacy-sensitive clinical settings, and fail to explicitly leverage the tree-like hierarchy of medical ontologies (e.g., ICD codes).
• Core Insight: Medical codes and patient trajectories exhibit hyperbolic geometry properties (hierarchical, tree-like structures). Euclidean embeddings distort these structures, whereas hyperbolic spaces can embed hierarchies with minimal distortion.
• Research Question: Can a compact model, explicitly aligned with the intrinsic hyperbolic geometry of EHRs, compete with billion-parameter LLMs in complex QA tasks?

2. Methodology: HypEHR Framework

HypEHR is a compact, two-stage framework based on Lorentzian (hyperbolic) modeling. It consists of a Patient Encoder and a Question-Answering Head.

A. Hyperbolic Clinical Sequence Encoder (Pretraining Stage)

The goal is to learn a patient representation in hyperbolic space ($H^d_L$) that respects the hierarchy of medical codes.

1. Embedding: Medical concepts (diagnosis, procedures, drugs) are embedded into a Lorentzian manifold via a learnable function $e_c \in H^d_L$.
2. Visit Aggregation: Within a visit, code embeddings are aggregated using a hyperbolic attention mechanism to form a visit representation $h_t$.
3. Sequence Modeling: A Lorentz Transformer processes the sequence of visits $\{h_t\}$, adapting self-attention, residuals, and normalization to the Lorentz manifold. This yields contextualized visit states $\{z_t\}$ and a global summary $z_{[CLS]}$.
4. Multi-Task Pretraining Objective: The encoder is trained to minimize a joint loss:
   $$L = L_{diag} + \lambda L_{hier}$$
   • $L_{diag}$ (Next-visit Diagnosis Prediction): Predicts the next diagnosis codes based on the current trajectory (Binary Cross-Entropy).
   • $L_{hier}$ (Hierarchy-Aware Regularization): Enforces the ICD ontology structure (Chapters $\to$ Blocks $\to$ Categories $\to$ Subcategories). It is decomposed into:
     • Radial Term ($L_{rad}$): Ensures deeper nodes in the hierarchy are farther from the origin (radius increases with depth).
     • Relative Term ($L_{rel}$): A triplet loss ensuring that an ancestor-related code is closer to a target code than a non-ancestor code in hyperbolic distance.

B. Hyperbolic EHR-QA Model (Inference Stage)

1. Question Encoding: A natural language question $q$ is encoded by a biomedical text encoder (Bio_ClinicalBERT) into Euclidean space, then projected to the hyperbolic manifold via an affine map and an exponential map, yielding $z_q$.
2. Hyperbolic Cross-Attention: The model computes attention scores between the question $z_q$ and visit states $\{z_t\}$ using negative scaled hyperbolic distances:
   $$s_t = -\gamma d_H(z_q, z_t)$$
   This identifies the most relevant visits for the query.
3. Aggregation: A Hyperbolic Weighted Fréchet Mean (HypAgg) aggregates the attended visits into a question-conditioned patient summary ($z_{p|q}^{visit}$). A secondary attention mechanism constructs a code-level rationale vector ($z_{p|q}^{code}$).
4. Answer Heads: Specialized heads process the hyperbolic representations (mapped back to tangent space via logarithmic map) to predict answers for four types:
   • Boolean: Existence of a condition.
   • Concept: Specific diagnosis or finding.
   • Numerical: Lab values or measurements.
   • Integer Count: Number of occurrences.

3. Key Contributions

1. Novel Architecture: Introduction of HypEHR, the first compact model to explicitly model EHRs and clinical questions entirely within a hyperbolic (Lorentzian) geometry, leveraging the intrinsic tree-like structure of medical ontologies.
2. Efficiency: The model achieves performance comparable to massive LLMs (e.g., Llama-3-8B, GPT-based pipelines) but uses orders of magnitude fewer parameters (~22M parameters vs. billions).
3. Geometry-Aware Regularization: Development of a hierarchy-aware loss function that explicitly aligns embeddings with the ICD code tree, ensuring that semantic similarity correlates with hyperbolic distance.
4. Theoretical Justification: A formal proof (Proposition 1) demonstrating that the ICD-10-CM hierarchy admits a low-distortion embedding into the Lorentz model, validating the geometric hypothesis.

4. Experimental Results

The model was evaluated on two MIMIC-IV-based benchmarks: EHRXQA and MIMIC-IV-Ext-Instr.

• QA Performance:
  • HypEHR achieved 89.53% accuracy on EHRXQA and 76.02% on MIMIC-IV-Ext-Instr.
  • It outperformed traditional representation learning models (RETAIN) and lightweight LLMs (Llama-3-8B).
  • It approached the performance of heavy LLM-based baselines (NeuralSQL with GPT-5.2, EHRAgent) without relying on massive parameter counts.
• Clinical Prediction:
  • When applied to standard tasks (Mortality, Readmission, Length-of-Stay, Phenotype prediction), HypEHR achieved state-of-the-art or comparable results to LLM-based baselines (e.g., best performance in Readmission and Phenotype prediction).
• Ablation Studies:
  • Pretraining: Removing the pretraining stage caused a significant drop in performance, confirming the necessity of learning the patient trajectory representation first.
  • Geometry: The hyperbolic model (HypEHR) significantly outperformed its Euclidean counterpart (EucEHR), proving the benefit of hyperbolic geometry for this data.
  • Hierarchy Loss: While not the primary driver of accuracy, the hierarchy loss ($L_{hier}$) significantly improved the model's ability to capture structural relationships, as evidenced by the monotonic increase of embedding radii with ICD tree depth.

5. Significance and Limitations

Significance:

• Privacy & Deployment: HypEHR offers a viable solution for on-premise, privacy-conscious clinical settings where deploying billion-parameter LLMs is infeasible due to cost and data sovereignty concerns.
• Data Efficiency: It demonstrates that explicitly modeling the geometric structure of medical data can yield high performance with minimal parameters, challenging the "bigger is better" paradigm in medical AI.
• Interpretability: The use of hyperbolic distances provides a geometrically intuitive way to measure similarity between medical concepts based on their hierarchical distance.

Limitations:

• Answer Type Restriction: Currently limited to discriminative answer types (Boolean, Concept, Numeric, Count). It does not yet support open-ended, generative text responses.
• Preprocessing Overhead: Requires restructuring datasets into predefined categories and mapping answers to specific vocabularies.
• Computational Complexity: Hyperbolic operations (exponential/logarithmic maps, Riemannian optimization) are computationally more complex than Euclidean operations, and public libraries for hyperbolic deep learning are less mature.

Conclusion

HypEHR successfully bridges the gap between the structural priors of medical data and efficient machine learning. By leveraging hyperbolic geometry, it achieves LLM-level performance in EHR question answering with a fraction of the computational cost, offering a promising direction for scalable and privacy-preserving clinical AI.