CARE: Outcome-Contrastive Case Retrieval for Clinical Trajectory Prediction

This repository contains the implementation for CARE, a retrieval-augmented clinical trajectory model for longitudinal structured EHR data. CARE is designed for sparse, irregular clinical event streams where similar observed patient histories may lead to different future outcomes.

CARE trains a time-aware patient encoder, builds a train-only memory bank of outcome-labeled case prefixes, retrieves similar positive and negative cases for each query, and fuses those outcome-contrastive neighbors to produce clinical risk predictions with case-based evidence.

The main benchmark target is EHRSHOT, a longitudinal EHR benchmark from Stanford Medicine with 6,739 patients and 15 clinical prediction tasks. CLMBR and TrajGPT scripts are included only to reproduce comparison baselines for CARE.

CARE currently supports:

• CARE encoder feature extraction for all EHRSHOT tasks in the shared few-shot protocol.
• Retrieval-augmented CARE risk training/evaluation for boolean EHRSHOT tasks.
• Masked multi-horizon retrieval-risk prediction over configured horizons.
• Self-supervised masked-code CARE pretraining on train-split EHR trajectories.
• Supervised CARE training on EHRSHOT labels, optionally initialized from the self-supervised checkpoint.
• Train-only outcome-contrastive retrieval memory.
• Full-train direct evaluation of the retrieval-augmented CARE risk model.
• Retrieval coverage, evidence similarity, calibration, conformal uncertainty, and patient-clustered confidence intervals.
• Retrieval-augmented few-shot evaluation on boolean EHRSHOT tasks, labeled by checkpoint training scope.
• Protocol-matched supervised feature-head baselines for CARE, CLMBR, and TrajGPT embeddings.

CARE supports masked multi-horizon prediction with configured horizons. Current EHRSHOT release files do not expose outcome timestamps, so each benchmark task supervises its configured horizon (task_horizon_days) unless a future label file provides horizon_days directly. Keep supervised-transfer encoder, supervised-transfer retrieval few-shot, and full-train retrieval-risk CARE results separate.

Create a Python environment and install dependencies:

git clone <repo-url>
cd ehrshot-care
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Optional CLMBR/FEMR regeneration dependencies are listed separately:

pip install -r requirements-clmbr-optional.txt

Optional FAISS HNSW retrieval dependencies are only needed for retrieval_backend: faiss experiments:

pip install -r requirements-retrieval-optional.txt

Most experiments do not need CLMBR regeneration because EHRSHOT assets include precomputed CLMBR features. Exact CARE retrieval is the default and does not require FAISS.

Request EHRSHOT access through the official data release process and place the extracted files under data/:

data/
  EHRSHOT_ASSETS/
    benchmark/     # task labels and few-shot splits
    features/      # precomputed CLMBR embeddings
    splits/        # train/val/test patient split files
  EHRSHOT_MEDS/
    data/          # MEDS-format structured EHR events
    labels/        # task label files
    metadata/      # code metadata

The default configs expect:

data/EHRSHOT_ASSETS
data/EHRSHOT_MEDS/data

All commands assume the virtual environment is active:

source .venv/bin/activate

CARE can first learn trajectory representations from train-split MEDS event streams without task labels:

python scripts/care/pretrain.py --config configs/care_ehrshot.yaml

Output:

results/care/checkpoints/pretrained_model.pt
results/care/checkpoints/pretrain_tokenizer.json
results/care/checkpoints/pretraining_manifest.json

This masked-code modeling stage is the cleanest way to compare CARE against foundation-style CLMBR and TrajGPT representations.

CARE supervised training has two stages:

1. Supervised encoder/risk warm-up without retrieval using masked task/horizon labels.
2. Reload the best warm-up checkpoint, build train-only multi-horizon memory, then run retrieval-enabled fine-tuning with the encoder frozen.

Checkpoints are selected by macro task-horizon validation loss, while the checkpoint payload also stores the pooled micro validation loss for audit. This keeps rare tasks from being dominated by high-volume labels during model selection.

python scripts/care/train.py \
  --config configs/care_ehrshot.yaml \
  --init_checkpoint results/care/checkpoints/pretrained_model.pt

Outputs:

results/care/checkpoints/best_model.pt
results/care/checkpoints/best_warmup_model.pt
results/care/checkpoints/tokenizer.json
results/care/checkpoints/task_to_index.json
results/care/memory/care_memory.pt

To evaluate fixed CARE encoder embeddings with the same logistic-regression few-shot evaluator used for CLMBR and TrajGPT, extract CARE features and label the checkpoint training scope explicitly. With the default supervised CARE checkpoint, report these as supervised-transfer features, not pure few-shot/self-supervised representation features:

python scripts/care/extract_embeddings.py \
  --config configs/care_ehrshot.yaml \
  --checkpoint results/care/checkpoints/best_model.pt \
  --checkpoint_training_scope supervised_ehrshot_train

python scripts/shared/run_evaluation.py \
  --assets_dir data/EHRSHOT_ASSETS \
  --features results/care/embeddings/care_features.pkl \
  --model_name care_pooled_encoder \
  --output_dir results/care/few_shot \
  --standardize

This produces a summary.csv using the same logistic-regression few-shot evaluator used for CLMBR and TrajGPT embeddings. The checkpoint scope in the feature metadata determines how the result should be reported. The extraction script resolves data paths and model dimensions from the checkpoint-merged CARE config. This encoder-feature path can cover all EHRSHOT tasks because it does not use the boolean-only retrieval-risk head.

To evaluate outcome-contrastive retrieval under the EHRSHOT few-shot splits for boolean tasks:

python scripts/care/run_fewshot_retrieval.py \
  --config configs/care_ehrshot.yaml \
  --checkpoint results/care/checkpoints/best_model.pt \
  --checkpoint_training_scope supervised_ehrshot_train \
  --output_dir results/care/few_shot_retrieval

This script builds each replicate memory bank only from that replicate's k-shot training examples and reports retrieval, encoder, positive_only, and negative_only ablations. It also reports positive/negative/contrastive retrieval coverage and top-neighbor similarity. Default CARE retrieval is time-valid; --disable_time_valid_retrieval exists only for an explicit protocol-parity ablation and should be labeled that way. It requires a modern retrieval_finetune checkpoint. With the default supervised CARE checkpoint, report this as supervised-transfer few-shot retrieval, not as pure few-shot representation learning.

To evaluate the retrieval-augmented CARE risk model directly:

python scripts/care/evaluate.py \
  --config configs/care_ehrshot.yaml \
  --checkpoint results/care/checkpoints/best_model.pt \
  --memory results/care/memory/care_memory.pt \
  --output_dir results/care/direct

This mode is labeled as full_train_direct_retrieval and should not be reported as the same protocol as EHRSHOT few-shot evaluation. It resolves paths/tasks from the checkpoint-merged CARE config, requires a retrieval_finetune checkpoint for retrieval evaluation, fits a validation-set temperature scaler by default, and reports per-task/per-horizon AUROC, AUPRC, Brier score, ECE, prevalence, class counts, retrieval coverage, top-neighbor similarity, conformal interval width/coverage, and patient-clustered bootstrap confidence intervals. Pass --no_calibration or --no_conformal only for explicit ablations.

The main evaluation script requires a modern checkpoint payload and strict state loading. Retrieval evaluation requires both a retrieval_finetune checkpoint and a matching memory bank.

Training builds the memory automatically. To rebuild memory for a saved checkpoint:

python scripts/care/build_memory.py \
  --config configs/care_ehrshot.yaml \
  --checkpoint results/care/checkpoints/best_model.pt

The memory builder requires a modern trained checkpoint and the checkpoint tokenizer. Evaluation rejects memory files whose checkpoint path, checkpoint identity hash/run ID, task map, or horizon metadata do not match the requested run.

Use the released CLMBR-T-base EHRSHOT embeddings as a comparison baseline for CARE:

python scripts/shared/run_evaluation.py \
  --assets_dir data/EHRSHOT_ASSETS \
  --features data/EHRSHOT_ASSETS/features/clmbr_features.pkl \
  --model_name clmbr \
  --output_dir results/clmbr

Optional CLMBR embedding regeneration:

python scripts/clmbr/extract_embeddings.py --config configs/clmbr_ehrshot.yaml

Use TrajGPT as a trajectory-modeling comparison baseline for CARE:

python scripts/trajgpt/pretrain.py --config configs/trajgpt_ehrshot.yaml

python scripts/trajgpt/extract_embeddings.py --config configs/trajgpt_ehrshot.yaml

python scripts/shared/run_evaluation.py \
  --assets_dir data/EHRSHOT_ASSETS \
  --features results/trajgpt/embeddings/trajgpt_features.pkl \
  --model_name trajgpt \
  --output_dir results/trajgpt

For a protocol-matched CARE comparison table, use the same supervised logistic-regression feature head on CARE, CLMBR, and TrajGPT embeddings:

python scripts/shared/run_matched_supervised_head.py \
  --assets_dir data/EHRSHOT_ASSETS \
  --features data/EHRSHOT_ASSETS/features/clmbr_features.pkl \
  --model_name clmbr_matched_supervised \
  --output_dir results/clmbr_matched_supervised

Run the same command for CARE and TrajGPT feature files. These outputs use protocol=matched_supervised_feature_head and can be compared directly with each other.

After CARE and baseline runs write summary.csv, compare result folders with an explicit protocol:

python scripts/shared/compare_models.py \
  --results-dir results/ \
  --protocol few_shot_feature

compare_models.py refuses to mix protocols by default. Use one protocol per table, such as few_shot_feature, matched_supervised_feature_head, few_shot_retrieval_time_valid_supervised_ehrshot_train, or full_train_direct_retrieval. When reporting results, keep CARE supervised-transfer encoder few-shot results, CARE supervised-transfer retrieval few-shot results, matched supervised feature-head results, and CARE full-train retrieval-risk results separate.

ehrshot-care/
  configs/
    care_ehrshot.yaml           # CARE model/training/evaluation config
    clmbr_ehrshot.yaml          # Optional CLMBR extraction config
    trajgpt_ehrshot.yaml        # TrajGPT config
  ehrshot/
    data_loading.py             # MEDS parquet loading utilities
    evaluation.py               # shared evaluation helpers
    tasks.py                    # EHRSHOT task definitions
  models/
    CARE/                       # CARE model package
      config.py
      encoders.py
      memory.py
      retriever.py
      fusion.py
      heads.py
      losses.py
      model.py
      evidence.py
      calibration.py
      utils.py
    trajgpt/                    # TrajGPT baseline model code
    trajgpt_embedder.py
    embedder.py
  scripts/
    care/
      pretrain.py               # CARE self-supervised masked-code pretraining
      train.py                  # CARE supervised training + retrieval fine-tuning
      extract_embeddings.py     # CARE encoder features for few-shot evaluation
      run_fewshot_retrieval.py  # retrieval-augmented few-shot evaluation
      build_memory.py           # train-only retrieval memory construction
      evaluate.py               # full-train retrieval-risk evaluation
    clmbr/
      extract_embeddings.py     # optional CLMBR extraction
    trajgpt/
      pretrain.py               # TrajGPT pretraining
      extract_embeddings.py     # TrajGPT feature extraction
    shared/
      run_evaluation.py         # EHRSHOT few-shot feature evaluator
      run_matched_supervised_head.py # matched supervised feature-head baselines
      compare_models.py         # protocol-safe result summary comparison
  data/                         # local EHRSHOT data, not committed
  results/                      # local outputs, not committed

CARE uses masked multi-horizon outcome-contrastive case retrieval:

• Encode each patient history prefix with a time-aware transformer encoder.
• Expand each labeled prefix into a masked horizon target using horizon_days when available, otherwise task_horizon_days.
• Build a training-only memory bank containing prefix embeddings, task IDs, prediction times, horizons, patient IDs, and binary outcomes.
• For a query prefix, retrieve similar train cases with the target outcome and similar train cases without the target outcome.
• Fuse the query with positive and negative retrieved cases using a divergence-aware fusion module.
• Predict clinical risk from the fused representation.
• Select checkpoints using macro task-horizon validation loss and calibrate direct risk predictions with validation-set temperature scaling.
• Report retrieval coverage/evidence similarity and conformal uncertainty intervals for retrieval-risk predictions.

The retrieval bank is filtered by task, horizon, patient ID, and prefix time to reduce leakage risk. Retrieval fine-tuning freezes the encoder so query embeddings remain aligned with the stored memory embeddings. Direct risk evaluation reports patient-clustered bootstrap confidence intervals. Exact retrieval is the default; retrieval_backend: faiss enables optional FAISS HNSW partition indexes when faiss-cpu or faiss-gpu is installed.

On macOS, if OpenMP libraries conflict, run commands with:

KMP_DUPLICATE_LIB_OK=TRUE OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python scripts/...

The CARE and TrajGPT scripts set KMP_DUPLICATE_LIB_OK automatically.

• EHRSHOT: An EHR Benchmark for Few-Shot Evaluation of Foundation Models, Wornow et al., NeurIPS 2023.
• CLMBR-T-base, Stanford Shah Lab.
• TrajGPT: Irregular Time-Series Representation Learning for Health Prediction, Song et al., IEEE JBHI 2025.