Cohort-Anchored Foundation Models for Electronic Health Records: From Risk Scores to Auditable Peer Cohorts

Technical Summary: Cohort-Anchored Foundation Models for Electronic Health Records

Problem Statement
Despite the emergence of billion-parameter foundation models for Electronic Health Records (EHRs) that demonstrate impressive zero- and few-shot capabilities, a structural gap remains between benchmark accuracy and reliable clinical deployment. Current models typically output a scalar risk score or a latent representation, treating patient comparison as an emergent side-effect rather than the primary unit of evidence. This creates three critical failure modes:

1. Opacity and Hallucination: Models often produce fluent but substantively incorrect medical content, and standard explanations (e.g., attention weights) fail to address the clinician's actual reasoning process, which relies on comparing a patient to a cohort of similar peers with documented outcomes.
2. Brittleness under Distribution Shift: Predictive performance degrades significantly when models are applied to new institutions, demographics, or protocols, as the underlying representation is not anchored to clinically meaningful population structures.
3. Lack of Auditable Evidence: Clinicians reason through cohorts (groups of similar patients), yet current systems do not return a "cohort-with-outcome" as the inference output, making the decision process difficult to audit or verify against established clinical infrastructure.

Methodology: The CAFM Framework
The paper proposes cafm (Cohort-Anchored Foundation Model), a framework designed to elevate the patient cohort to a "first-class object" in the machine learning pipeline. Rather than rebuilding existing backbones, cafm acts as a compositional retrofit that integrates cohort-aware capabilities into existing EHR foundation models through four interlocking stages:

1. Deviation-Aware Cohort Construction (Stage 1):
   Before pretraining, the corpus undergoes curation to detect and handle two distinct issues: deviations (statistically anomalous records due to technical errors or documentation drift) and bias (systematic inconsistencies arising from clinical workflows, such as informative missingness). This stage produces a curated dataset where cohorts are assembled using rule-based phenotypes, anchor-and-learn methods, or learned prototypes, ensuring the training data reflects clinically valid populations.

2. Cohort-Conditioned Pretraining (Stage 2):
   The framework introduces a composite contrastive objective to the standard masked or autoregressive pretraining loss.

   • Attraction: Pulls together representations of patients with similar soft cohort distributions.
   • Repulsion (Near-Miss Negatives): Pushes apart patients who are phenotypically similar but have divergent outcomes. This "near-miss" sampling is critical for sharpening cohort boundaries.
   • Regularization: Includes entropy regularization to prevent the model from collapsing all patients into a single cohort.
   • Curriculum Learning: The framework employs a warm-start strategy and a curriculum that gradually increases the weight of contrastive losses to stabilize joint training between the backbone encoder and the cohort module.

3. Multimodal Cohort Alignment (Stage 3):
   To address "modality contamination" (where one modality, such as free-text notes, dominates the joint representation), the framework enforces intra-cohort similarity within each modality (text, structured codes, imaging, physiological signals) separately. This ensures that weaker modalities retain their discriminative capacity and do not collapse onto the dominant modality, preserving the integrity of the cohort structure across data types.

4. Clinician-in-the-Loop Refinement (Stage 4):
   At inference, the model does not output a scalar score but retrieves a peer cohort and its associated outcome distribution. This output serves as the explanation, audit unit, and drift signal. The framework supports:

   • Auditable Evidence: Clinicians can inspect the retrieved cohort's statistics and outcome rates.
   • Task Decomposition: Clinicians can modify cohort definitions (e.g., adding constraints) in real-time, triggering a re-computation of the outcome distribution without re-running the full backbone.
   • Lifecycle Monitoring: The system monitors for distribution shifts by tracking changes in the retrieved cohort distributions, offering a more sensitive drift signal than aggregate performance metrics.

Key Contributions

• Structural Re-framing: The paper argues that cohort anchoring is a structural prerequisite for EHR foundation models, distinct from in-context learning, retrieval-augmented generation (RAG), or post-hoc explanation methods. It posits that the cohort, not the individual embedding, should be the unit of clinical evidence.
• The CAFM Framework: A concrete, four-stage methodology that can be retrofitted onto existing backbones (e.g., BEHRT, Med-BERT) without requiring a complete architectural redesign.
• Mathematical Formulation: The authors provide explicit mathematical formulations for the contrastive objectives, including the definition of "near-miss" negatives and the composite loss function (Eq. 6 and Eq. 9).
• Empirical Hypotheses: The paper outlines five falsifiable hypotheses (H1–H5) that the framework predicts, covering cohort coverage vs. calibration, the efficacy of near-miss negatives, early drift detection, mitigation of modality collapse, and clinician preference for cohort evidence in high-uncertainty scenarios.
• Case Studies: The framework is illustrated through four operational case studies: Acute Kidney Injury (AKI) prediction, cardiovascular risk stratification from ECGs, optic neuropathy triage from orbital imaging, and ERG-grounded report generation.

Results and Illustrations
The paper does not present a single unified experimental benchmark with new state-of-the-art scores. Instead, it illustrates the framework's potential through four case studies drawn from operational settings:

• AKI Prediction: Demonstrates how cohort anchoring transforms a scalar risk into a structured evidence object, allowing clinicians to query specific factors (e.g., nephrotoxic medication exposure) within the retrieved cohort.
• Cardiovascular Risk (ECG): Shows how multimodal alignment prevents the ECG signal from overwhelming structured EHR data, grounding risk estimates in institution-specific historical trajectories.
• Optic Neuropathy Triage: Highlights how cohort retrieval supports triage decisions by providing a curated set of prior referral cases, addressing the trust gap in primary care settings.
• ERG Report Generation: Proposes a generative approach where reports are grounded in retrieved patient cohorts rather than just text snippets, reducing hallucination by constraining generation to clinically comparable populations.

Significance and Claims
The paper positions itself as a perspective backed by a concrete framework, aiming to bridge the gap between impressive benchmarks and trustworthy bedside use. Its primary claims are:

1. Methodological Necessity: Cohort anchoring is the most direct route to making foundation models interpretable, robust to distribution shifts, and auditable. It shifts the paradigm from "representation quality" to "population-grounded evidence."
2. Practical Feasibility: The framework is designed to be a retrofit, requiring engineering effort to integrate cohort modules and retrieval indices rather than the development of entirely new backbone architectures.
3. Testability: The framework is defined by specific, falsifiable empirical hypotheses. The authors argue that the validity of cohort anchoring should be judged not by rhetorical appeal but by whether these hypotheses hold under prospective evaluation.
4. Operational Reality: The paper acknowledges that the barriers to adoption are not conceptual but organizational. Realizing this vision requires sustained effort in data curation, governance, and the maintenance of cohort definitions as clinical practices evolve.

The authors conclude that while the necessary modeling components exist, the next phase of progress in clinical AI depends on equipping existing models with the population-grounded substrate that clinical reasoning has traditionally relied upon.