An Investigation of Memorization Risk in Healthcare Foundation Models

TL;DR

The paper develops a contextual, risk-based framework for auditing memorization in healthcare foundation models trained on structured EHR data. It introduces six black-box tests (T1–T6) that separately assess generative and embedding memorization (T1–T4) and patient-level privacy risk via perturbation and sub-population analyses (T5–T6), using a publicly available EHR-FM (EHRMamba2) trained on MIMIC-IV and a MedBERT-based embedding space to quantify sequence similarity with $d_{EMD}$. The results show that increasing prompt information enhances memorization signals in some cases, while embedding leakage remains weak on average; perturbation and sub-population tests reveal high-risk, patient-specific memorization in certain scenarios, underscoring the need for safeguards. The authors provide an open-source toolkit to enable reproducible privacy assessments and discuss implications for clinical deployment, emphasizing that memorization does not always equate to privacy risk but can in targeted contexts. Overall, the work offers practical methods to quantify and mitigate privacy risks in EHR foundation models prior to release and real-world use.

Abstract

Foundation models trained on large-scale de-identified electronic health records (EHRs) hold promise for clinical applications. However, their capacity to memorize patient information raises important privacy concerns. In this work, we introduce a suite of black-box evaluation tests to assess privacy-related memorization risks in foundation models trained on structured EHR data. Our framework includes methods for probing memorization at both the embedding and generative levels, and aims to distinguish between model generalization and harmful memorization in clinically relevant settings. We contextualize memorization in terms of its potential to compromise patient privacy, particularly for vulnerable subgroups. We validate our approach on a publicly available EHR foundation model and release an open-source toolkit to facilitate reproducible and collaborative privacy assessments in healthcare AI.

Figures (8)

• Figure 1: Our proposed tests (T1–T6) to evaluate memorization in EHR-FM, grouped by two core objectives. T1–T4 measure the extent to which models reveal training data, targeting both the embedding and the generated data domain. T5–T6 perform patient-level evaluations to quantify individual privacy risks from memorization or data leakage.
• Figure 2: Model predictions in trajectory memorization. (a) Average distance between generated and ground truth EHR sequences of 100 codes for different prompting setups. (b) Predicted tokens and their frequencies in the best-performing trajectory prediction. All tokens are related to common lab measures, and most often normal values.
• Figure 3: Code generation and membership inference. (a) Frequency of EHR code randomly generated by EHRMamba2 versus the original distribution of codes in the training dataset. (b) Distribution of detection scores from a membership inference attack based on model confidence. The x-axis shows scores, the y-axis shows counts, with colors indicating member (True) vs non-member (False) samples.
• Figure 4: Predictive results for the perturbation test on 2 identified samples from infectious disease and substance abuse. Each bar represents the probability of generating the sensitive token, when the model is prompted with the 10 prompts and with different ages. The horizontal line is the threshold of 0.3 used in T2, and the red vertical line represents the original patient age.
• Figure 5: Likelihood of predicting the codes related to sensitive attributes (Infectious disease, mental health, and substance abuse) in EHR trajectories generated by prompting the benchmark model with rare diagnoses.

Theorems & Definitions (2)

• Definition 3.1: Generative memorization
• Definition 3.2: Embedding memorization