EHR Future Condition Forecasting Transformer
Custom PyTorch causal Transformer trained to forecast newly diagnosed conditions from longitudinal synthetic EHR timelines.
Model
The model represents each patient as a chronological sequence of diagnosis, medication start/stop, procedure, observation, and encounter events ending with an [ANCHOR] token.
Architecture:
- Causal Transformer decoder
- Hidden dimension: 256
- Transformer layers: 4
- Attention heads: 8
- Maximum sequence length: 512
- Temporal rotary positional embeddings using patient age
- Nonlinear 256 -> 256 -> 40 classification head
- 40 multilabel future-condition targets
The Transformer backbone was first pretrained using causal next-event prediction and subsequently fine-tuned using Asymmetric Loss.
Final validation performance
- Macro AUROC: 0.731870
- Mean Average Precision (mAP): 0.203563
- Best fine-tuning epoch: 19
- Fine-tuning validation loss: 0.017285
The final pretraining checkpoint was selected at epoch 29/30 with validation next-event loss 1.470598.
Preprocessing
The fitted preprocessing state is included in
preprocessing_state.joblib.
Numeric observations are represented using training-fitted quantile bins. Categorical observations use training-fitted categorical mappings. Medication starts and stops are represented separately.
Vocabulary size: 1383
Maximum sequence length: 512
Loading
This is a custom PyTorch model and is not directly loadable with
transformers.AutoModel.
The model implementation and evaluation scripts are available at:
https://github.com/nateRaw3/ehr-forecasting-transformer.git
The model checkpoint is stored as model.pt.
To use this checkpoint with the accompanying GitHub repository, download
model.pt and preprocessing_state.joblib and provide them to the repository's
inference pipeline.
Intended use
This model was developed for the provided synthetic EHR forecasting assessment. It predicts the probability of 40 future first-ever condition diagnoses during the defined post-anchor prediction horizon.
Limitations
- Trained and evaluated on synthetic Synthea-derived EHR data.
- Not validated for clinical use.
- Performance varies substantially between target conditions.
- Rare outcomes have high uncertainty in estimated performance.
- The model should not be interpreted as providing causal clinical effects.