BEHRT: Transformer for Electronic Health Records

In plain words

Doctors' records are a long history of diagnoses, each dated to a visit. BEHRT treats one patient's history like a document. Each diagnosis is a word, each visit is a sentence, and the whole history is read by a model that was built for language, a Transformer in the style of BERT.

Each word gets a mix of signals: which diagnosis it is, how old the patient was at that point, which visit it belongs to, and where it sits in the sequence. Before it is used for prediction, the model plays a fill-in-the-blank game: some diagnoses are hidden and it must guess them from the rest of the history. This teaches it how diagnoses tend to go together and in what order they tend to appear.

After that pre-training, a small output layer is trained to predict which conditions a patient will be diagnosed with at the next visit or in the next 6 or 12 months. The authors report that this beats earlier deep-learning models for medical records on this task, and that the learned representations of diagnoses look clinically sensible.

TL;DR

• BEHRT is BERT adapted to diagnosis histories from CPRD, a UK primary-care database, with 301 diagnosis codes and about 1.6 million patients §2.2.
• Each token is the sum of four learned or fixed embeddings: diagnosis, age, segment (visit parity) and visit position §2.3.
• Pre-training uses masked disease prediction; fine-tuning is multi-label prediction at the next visit and within 6 and 12 months §2.4 §2.5.
• Author-reported average precision (APS) is higher than DeepR and RETAIN on all three tasks tab. 1. The code is public (MIT) but the data is not, so the results cannot be reproduced from this repository alone.

Problem

Deep models on EHR data usually take a hand-built feature vector per patient, and earlier sequence models (RNN, CNN) struggle with long histories §1. The authors want a model that reads the full ordered history, handles irregular gaps between visits and several kinds of concept, and yields representations that can be reused across many diseases §1 §2.3.

The data are assumed to be coded diagnoses from linked primary and hospital care, with visit order known and visit dates available §2.2. The paper does not discuss other data sources or small datasets.

Method

Data §2.1 §2.2 fig. 1 fig. 2. The source is CPRD primary care, linked to Hospital Episode Statistics (HES) for hospital care, using only GP practices that consented to HES linkage §2.1. The authors start from about 8 million patients, keep those eligible for linkage that meet CPRD's quality flags and have at least five visits, leaving 1.6 million §2.2 fig. 1. Primary-care Med Codes are mapped to Read Codes (procedure codes dropped); ICD-10 level-4 codes and Read Codes are mapped to Caliber disease codes. This gives G = 301 diagnosis codes §2.2.

Sequence format §2.2 fig. 2. For patient p with n_p visits, the history becomes one token sequence: a CLS token, then each visit's diagnoses in time order, with a SEP token closing every visit. Visits with no diagnosis are dropped fig. 2.

Embedding §2.3 fig. 3 NextXVisit.py. Each token has four vectors that are added together, then layer-normalised and passed to dropout:

• disease embedding, learned, 288 dimensions §3 tab. 2;
• age embedding, the patient's age at that token; the code uses months, see Pitfalls;
• segment embedding, two learned vectors, A and B, alternating from one visit to the next, marking visit boundaries §2.3;
• position embedding, fixed sinusoidal values taken from Vaswani et al. §2.3. In the code the position index counts visits, not tokens. See Paper vs code.

Encoder §3 tab. 2 MLM.py. A 6-layer Transformer encoder with 12 attention heads, hidden size 288 and intermediate size 512 [tab. 2, row 21]. The attention mask hides only padding, so attention is bidirectional over the whole history §2.3.

Pre-training with masked disease prediction §2.4 utils.py. Following BERT, 15% of the tokens are selected. Of those, 80% become MASK, 10% become a random vocabulary token, and 10% stay unchanged. The loss is computed only at selected positions. The output head predicts a token over the full vocabulary and shares weights with the input disease embedding MLM.py. The reported metric is precision at threshold 0.5 §2.4.

Fine-tuning and labels §2.5. Patients are split at random 80/10/10 into train, validation and test §2.5. Three tasks share one model:

• T1, next visit: a random cut point j with 3 < j < n_p; the input is visits 1 to j and the target is the multi-hot vector of diagnoses in visit j+1 §2.5.
• T2 and T3, next 6 or 12 months: the cut point is drawn from a range where at least 6 or 12 months of record follow; the target is the multi-hot vector of diagnoses in that window §2.5. Patient counts are 699K (T1), 391K (T2) and 342K (T3) §2.5.

The predictor is one feed-forward classifier layer on the BEHRT output §3. The code reads the first-token output through an extra pooler layer, see Paper vs code.

Metrics §2.5. APS and AUROC are computed per patient and then averaged over patients. APS is the main metric. The paper cites the scikit-learn-style definitions it used §2.5.

# BEHRT input for one patient  [§2.2] [code: dataLoader/NextXVisit.py#L21-L44]
# Input:  visits, a list of visits each a list of disease ids, with ages
# Output: token ids, ages, segments, positions, attention mask (length L)

def build_sequence(visits, ages, max_len):
    # Keep the most recent max_len - 1 tokens
    # [code: dataLoader/NextXVisit.py#L32-L33]
    tokens = flatten_with_sep(visits)[-(max_len - 1):]
    # Prepend CLS; it takes the age of the first token
    # [code: dataLoader/NextXVisit.py#L36-L40]
    tokens = ["CLS"] + tokens
    # Segment flips after every SEP, giving alternating visit parity
    # [§2.3] [code: dataLoader/utils.py#L43-L56]
    segments = alternate_segment_ids_per_visit(tokens)
    # Position counts visits: it increments after every SEP
    # [§2.3] [code: dataLoader/utils.py#L59-L69]
    positions = count_visits(tokens)
    # Pad to max_len and mark real tokens with 1
    # [code: dataLoader/NextXVisit.py#L43-L58]
    mask = pad_and_mask(tokens, max_len)
    return tokens, ages, segments, positions, mask

# Input embedding: sum of four embeddings, then LayerNorm and dropout  [§2.3]
# [code: model/NextXVisit.py#L40-L57]
# Input:  token ids, ages, segments, positions, each (B, L)
# Output: hidden states (B, L, 288)

def embed(tokens, ages, segments, positions):
    # Sum disease, age, segment and visit-position vectors  [§2.3]
    summed = (
        disease_table[tokens]
        + age_table[ages]
        + segment_table[segments]
        + position_table[positions]
    )
    # Normalise, then apply dropout  [code: model/NextXVisit.py#L55-L56]
    return dropout(layer_norm(summed))

# Masked disease prediction, one step
# [§2.4] [code: dataLoader/utils.py#L14-L40]
# Input:  tokens (L,), vocabulary
# Output: corrupted tokens (L,), labels (L,), IGNORE where no loss is taken

def mask_tokens(tokens, vocabulary):
    # Start from copies so the originals stay available for labels
    corrupted = list(tokens)
    labels = []
    for position, token in enumerate(tokens):
        # Select 15% of positions for prediction
        # [§2.4] [code: dataLoader/utils.py#L19-L20]
        if random() >= 0.15:
            labels.append(IGNORE)
            continue
        # Predict the original token at this position  [§2.4]
        labels.append(vocabulary[token])
        # Draw a fraction in [0, 1) and split it 80/10/10
        # [§2.4] [code: dataLoader/utils.py#L21-L29]
        fraction = random()
        if fraction < 0.8:
            corrupted[position] = MASK
        elif fraction < 0.9:
            corrupted[position] = random_vocabulary_token(vocabulary)
        # Otherwise leave the token unchanged
    return corrupted, labels

# Next-visit / multi-horizon prediction head  [§2.5] [§3]
# [code: model/NextXVisit.py#L119-L137]
# Input:  hidden (B, L, 288), multi-hot targets (B, G)
# Output: loss (scalar), logits (B, G)

def predict_diagnoses(hidden, targets):
    # Take the first-token (CLS) output  [code: model/NextXVisit.py#L113]
    first = hidden[:, 0]                       # (B, 288)
    # Pooler: dense layer and tanh. The paper does not describe it.
    # [code: model/NextXVisit.py#L113]
    pooled = tanh(dense(first))                # (B, 288)
    # Dropout, then one linear layer to G logits
    # [code: model/NextXVisit.py#L131-L132]
    logits = linear(dropout(pooled))           # (B, G)
    # Multi-label soft-margin loss over the G labels
    # [code: model/NextXVisit.py#L134-L137]
    loss = multilabel_soft_margin(logits, targets)
    return loss, logits

Hyperparameters

Values marked "code" come from the official notebooks, not the paper. The paper's text states only the architecture sizes and the 15% masking rule.

Results

Author-reported. Setup: CPRD, 1.6 million patients, 301 diagnosis codes, one linear head per task on the pre-trained encoder §3 §2.5. Each cell is APS | AUROC, averaged over patients tab. 1.

MLM pre-training reaches precision 0.6597 with the tuned architecture §3 tab. 2 row 21. Disease embeddings projected to two dimensions with t-SNE group co-occurring and same-chapter diseases, and place sex-specific diseases apart §3 fig. 4. Attention maps show links between a diagnosis and distant past diagnoses, illustrated for two patients §3 fig. 5. Per-disease results for the next-6-month task are plotted as APS against AUROC, with circle size showing prevalence fig. 6. The appendix table app. B tab. 5 lists 57 diseases with an occurrence ratio of at least 0.01; it is not reproduced here.

Best tuned baseline configurations: DeepR row 16 (average precision 0.3588) and RETAIN row 8 (average precision 0.3740) app. A tab. 3 tab. 4. Their full search tables are in the paper and are not reproduced here.

Paper vs code

Official code: https://github.com/deepmedicine/BEHRT, commit ca0163faf5ec09e5b31b064b20085f6608c2b6d1, MIT. The repository belongs to the Deep Medicine group listed in the author affiliations. It holds the model, data loaders and notebooks. It does not hold the CPRD preprocessing (the module it imports is missing), the diagnosis vocabulary, the trained weights, or the DeepR and RETAIN baselines.

• Position index. The paper gives the position encoding as the sinusoidal rule of Vaswani et al. §2.3. The code uses that rule, but the index is the visit number, which increments after each SEP, not the token index utils.py NextXVisit.py.
• Masked tokens. The paper masks 15% of disease words §2.4. The code applies the mask to every token in the sequence, including CLS and SEP MLM.py utils.py.
• Prediction head. The paper describes one feed-forward classifier layer on the BEHRT output §3. The code inserts a dense+tanh pooler on the first token before dropout and the linear layer NextXVisit.py NextXVisit.py.
• Loss functions. The paper gives no loss. The code uses cross-entropy over the vocabulary, ignoring unmasked positions, for MLM MLM.py. For prediction it uses the multi-label soft-margin loss NextXVisit.py.
• Validation split. The paper uses train, validation and test §2.5. The notebooks load train and test only, and the checkpoint is kept when the test score is highest NextXVisit.ipynb NextXVisit.ipynb NextVIsit-6month.ipynb.
• Sequence length. Not stated in the paper. The code uses 64 for MLM MLM.ipynb and 100 for the prediction tasks NextXVisit.ipynb.
• Warmup. The paper states no optimiser. The code calls BertAdam with a warmup fraction of 0.1 and no total step count optimiser.py.
• Attention direction. The paper calls the model bidirectional §2.3. The code masks padding only, so attention is bidirectional MLM.py. No discrepancy.
• Segment and embedding sum. The paper's segment alternates A and B per visit §2.3. The code does the same utils.py. The code adds the four vectors and then normalises NextXVisit.py. No discrepancy.

Pitfalls

• Data and vocabulary are not public. CPRD access is licensed. The preprocessing script imports a module that is not in the repository and reads paths left blank behrtFormat.py. The MLM notebook loads its vocabulary from a blank path MLM.ipynb. You must build your own 301-code (or other) vocabulary from scratch.
• Position index is the visit number. Tokens within one visit share the same position id, and positions are fixed (from_pretrained freezes them by default). If you use token indices, you deviate from the code NextXVisit.py.
• Age unit. Ages are stored in months and capped at 110 years, giving 1,320 month values NextXVisit.ipynb utils.py behrtFormat.py. The paper does not state the unit.
• Random replacement pool. The 10% random-token case draws from all vocabulary values, including special tokens utils.py.
• Label space. Labels are built from the vocabulary with PAD, SEP, CLS and MASK removed, but UNK kept NextVIsit-6month.ipynb. The label size is therefore the vocabulary size minus 4, which may differ from G = 301. The vocabulary file is not in the repository, so this cannot be checked.
• Checkpoint selection on the test set. The reported best checkpoint is selected on the test loader NextXVisit.ipynb NextXVisit.ipynb. This biases the reported numbers upward. Hold out a validation split and select on it.
• Learning-rate schedule. BertAdam receives no total step count optimiser.py. In pytorch_pretrained_bert (not vendored here, and not checked in this session), a missing total step count usually means no warmup or decay. Verify this before relying on warmup.
• Baseline name conflict. The introduction attributes Deepr to Nguyen et al. (a CNN) §1. The results cite DeepR as Miotto et al. §3. The baseline search tables use filters and kernel sizes, which fits the CNN app. A tab. 3. Which model was compared is therefore unclear. The baseline code is not in the repository.
• Inconsistent baseline numbers. DeepR's 6-month and 12-month APS and AUROC are identical in tab. 1 (0.393 / 0.943). RETAIN's best tuning run (average precision 0.3740 tab. 4 row 8) does not match its next-visit value of 0.382 tab. 1. The paper does not explain either.
• Headline gain. The abstract gives an 8.0–10.8% absolute gain. Against RETAIN, tab. 1 gives 8.0 (next visit), 10.8 (6 months) and 9.3 (12 months) points. Against DeepR the gains reach 13.2 points (6 months). The abstract's range does not name a baseline. Quote the table, not the range.
• Missing hyperparameters. The paper gives no learning rate, batch size, epochs, dropout, seed, or sequence length §2.4 §3. The values above come from notebooks, not the paper. No random seed is set in the MLM or prediction notebooks.
• Pre-training precision is not a proxy. The MLM precision (0.6597) is at threshold 0.5, and it is not comparable to the APS in tab. 1.

Lineage

• Builds on: Transformer encoders §2.3; BERT and its masked pre-training §1 §2.4; the attention visualisation of Vig §3.
• Compared against: RETAIN (Choi et al., reverse-time attention) §1 tab. 1; DeepR, under the name conflict above §1 §3 tab. 1.
• Superseded by: not stated in this card.