This is the official repository for the ICML 2026 paper “How Should Transformers Encode Numeric Values in Electronic Health Records?”
The paper is available here: arXiv
This repository contains code for:
- implementing the numeric encoding schemes evaluated in the paper
- generating the synthetic tasks used to evaluate numeric reasoning
The implementation builds on the repository described in:
Montgomery et al. (2025). BONSAI: A framework for processing and analysing Electronic Health Records (EHR) data using transformer-based models. Journal of Open Source Software, 10(114), 8869. https://doi.org/10.21105/joss.08869
The repository supports four methods for handling numeric values (e.g., lab results, vital signs) in EHR sequences:
-
discrete: Discretises numeric values into categorical tokens (requires binning). Values are converted to categorical tokens and embedded using the concept embedding layer. -
combined: Uses separate embeddings for categorical concepts and continuous values, placing them in different sequence positions. Concept embeddings are used for categorical positions, while value embeddings are used for positions with numeric values. They occupy separate positions in the sequence and are not combined. -
concat: Uses the same sequence positions for both concepts and values (joint embeddings). Replaces concept embeddings with a projected concatenation of concept and value embeddings. When a value exists, the final embedding isconcat_proj([concept_embeds, value_embed]), which projects the concatenated embeddings to the hidden size. When no value exists, the original concept embeddings are used. -
FiLM: Uses the same sequence positions for both concepts and values (joint embeddings). Replaces concept embeddings with Feature-wise Linear Modulation of value embeddings. When a value exists, the final embedding isgamma * value_embed + beta, where gamma and beta are learned linear transformations of the concept embeddings. When no value exists, the original concept embeddings are used.
Binning: Numeric values can optionally be binned before embedding. Binning is required for discrete mode and optional for combined, concat, and FiLM modes.
Key Implementation Files:
corebehrt.modules.model.embeddings- Embedding layer implementations for all value representation methodscorebehrt.modules.model.model- Model architecture with value prediction heads and training logiccorebehrt.modules.features.values- Value processing and integration into the feature pipelinecorebehrt.functional.features.values- Utility functions for binning and discretising numeric values
For more details on the overall pipeline, see the sections below.
See the synthetic data workflow guide for how to create base cohorts and variant synthetic datasets.
- End-to-end EHR Pipeline: Tools for data ingestion, cleaning, and feature extraction.
- BERT-based Modeling: Pretraining on massive EHR corpora followed by task-specific finetuning.
- Cohort Management: Flexible inclusion/exclusion logic, temporal alignment, outcome definition.
- Scalable: Designed to run both locally or on cloud infrastructure (Azure).
- Built-in Validation: Cross-validation and out-of-time evaluation strategies.
Below is a high-level overview of the most important directories:
- main: Primary pipeline scripts (create_data, pretrain, finetune, etc.)
- modules: Core implementation of model architecture and data processing (detailed overview)
- configs: YAML configuration files for each pipeline stage
- functional: Pure utility functions supporting module operations (detailed overview)
- azure: Cloud deployment and execution utilities (azure instructions)
- synthetic_data: Synthetic EHR/lab data generators, post-processing scripts, and analysis helpers (workflow guide)
For running tests and pipelines, create and activate a virtual environment, then install the required dependencies:
python -m venv .venv
source .venv/bin/activate
(.venv) pip install -r requirements.txt
Below is a high-level description of the steps in the BONSAI pipeline. For detailed configuration options, see the main README. The pipeline can be run from the root directory by executing the following commands:
(.venv) python -m corebehrt.main.create_data
(.venv) python -m corebehrt.main.prepare_training_data --config_path corebehrt/configs/prepare_pretrain.yaml
(.venv) python -m corebehrt.main.pretrain
(.venv) python -m corebehrt.main.create_outcomes
(.venv) python -m corebehrt.main.select_cohort --config_path corebehrt/configs/select_cohort.yaml
(.venv) python -m corebehrt.main.prepare_training_data --config_path corebehrt/configs/prepare_finetune.yaml
(.venv) python -m corebehrt.main.finetune_cv
(.venv) python -m corebehrt.main.select_cohort --config_path corebehrt/configs/select_cohort_held_out.yaml
(.venv) python -m corebehrt.main.prepare_training_data --config_path corebehrt/configs/prepare_held_out.yaml
(.venv) python -m corebehrt.main.evaluate_finetune --config_path corebehrt/configs/evaluate_finetune.yaml(.venv) python -m corebehrt.main.evaluate_pretrain --config_path corebehrt/configs/evaluate_pretrain.yaml
Before using BONSAI, you need to convert your raw healthcare data into the MEDS (Medical-Event-Data-Standard) format format. We provide a companion tool ehr2meds to help with this conversion:
- Converts source data (e.g., hospital EHR dumps, registry data) into MEDS
- Performs code normalization and standardization
- Provides configuration options for handling different data sources
- Includes validation to ensure data quality
- Goal: Convert MEDS into tokenized features suitable for model training.
- Key Tasks:
- Vocabulary Mapping: Translates raw medical concepts (e.g., diagnoses, procedures) into numerical tokens.
- Temporal Alignment: Converts timestamps into relative positions (e.g., hours or days from an index date).
- Incorporate Background Variables: Incorporates static features such as age, gender, or other demographics.
- Efficient Output: Produces a structured parquet format that can be rapidly loaded in subsequent steps.
- Goal: Train a ModernBERT model via masked language modeling.
- Key Tasks:
- Large scale self-supervised training on EHR sequences
- Embedding temporal relationships between medical events
- Saves checkpoints for downstream finetuning
- Goal: Generate outcomes from the formatted data for supervised learning.
- Key Tasks:
- Search for specific concepts (medications, diagnoses, procedures) in the data
- Optionally create exposure definitions for more complex study designs
-
Goal: Define the study population
-
Key Tasks:
- Apply inclusion/exclusion criteria (e.g. age, prior outcomes)
- Generate index dates for each patient
- Produce folds and test set for cross-validation
- Goal: Adapt the pretrained model for specific binary outcomes
- Key Tasks:
- K-fold cross-validation
- Includes early stopping and evaluation on test set
For a detailed overview of the pipeline, see the main README.
For running BONSAI on Azure cloud infrastructure using SDK v2, refer to the Azure guide. This includes:
- Configuration setup for Azure
- Data store management
- Job execution in the cloud
- Environment preparation
BONSAI supports multiple methods for handling numerical values (e.g., lab results, risk scores) in EHR sequences. See the Numeric Value Representation section above for details on the four supported methods (discrete, combined, concat, and FiLM). For implementation details and configuration options, refer to the main README.
We welcome contributions! Please see our Contributing Guidelines for details on:
- Code style and formatting
- Testing requirements
- Pull request process
- Issue reporting
This project is licensed under the MIT License - see the LICENSE file for details.
If you use BONSAI in your research, please cite the following paper:
@article{Montgomery2025,
author = {Montgomery, A. and others},
title = {BONSAI: A framework for processing and analysing {E}lectronic {H}ealth {R}ecords ({EHR}) data using transformer-based models},
journal = {Journal of Open Source Software},
volume = {10},
number = {114},
pages = {8869},
year = {2025},
doi = {10.21105/joss.08869}
}