G4H Half-Day

Embeddings : from NLP to medical pathways

Meilame Tayebjee

15 October 2025

Table of contents

• 1 From NLP…
  • Representing words
  • Introducing embeddings
  • How to train ?
  • Application of word embeddings: similarity
  • Application of word embeddings: Language Models
  • Introducing Transformers
• 2 … to patient embeddings
  • Our data
  • Embeddings
  • Tokenization
  • Pre-training task
  • First trained models
  • Plot of the embedding space
• 3 Next steps
  • Causal inference
  • Improving pre-training task
  • Scaling
• References

1 From NLP…

Representing words

• Historically (up to the 2000s), words in NLP were treated as symbols — not as entities with meaning:
  • A one-hot vector (e.g. “cat” = [0, 0, 1, 0, 0, …])
  • Or a string in a lookup table.
• Problems:
  • No notion of similarity:
    • “hotel” = [ 1 0 0 0] et “motel” = [0 1 0 0] are orthogonal !
  • Dimensional explosion

Introducing embeddings

• Each word w should be mapped to a dense vector v_w \in \mathbf{R}^d where d <<< |V|
• Taking context into account (distributional semantics):
  • “You shall know a word by the company it keeps” (J. R. Firth, 1957)

From Stanford NLP Course

• Count-based: co-occurrence matrix, PPMI…
• Deep-learning based: Word2Vec

From Lena Voita

From Lena Voita

Notice how the context length is important and can quickly fill the memory.

How to train ?

From Lena Voita

Application of word embeddings: similarity

From GloVe project page

Application of word embeddings: Language Models

• Language Models = models that predict the next word
  • Fixed-window LM: does not scale when context size increases
  • RNNs: model size independent from context size, but sequential processing + vanishing gradient

From Stanford NLP Course

From Stanford NLP Course

Introducing Transformers

• Attention matrix \alpha_{i, j}(X) to produce contextualized embeddings
  • Each embedding is transformed, via a weighted average
  • No dimension change, but the output is now context-aware
• Parallelizable, scalable
• Two types:
  • BERT
    • pre-trained with Masked Language Modeling
  • GPT
    • causal attention matrix
    • pre-trained with next token prediction task

2 … to patient embeddings

Our data

• Our data: patient (medical) pathways
• Very similar to NLP, with differences
  • unstructured (just as in NLP)
  • temporally irregular (not like NLP!)

Embeddings

• Represent tokens and patient as a dense vector, to:
  • make visualizations
  • use the vector for downstream task
    • prediction, classification
    • clustering
    • causal inference

Tokenization

Dataset of medical pathways: For i = 1, \dots, n, X_i \coloneqq (e_i, t_i, f_i), where:

• e_i \in \mathbb{R}^{cs} are the tokenized events, containing integers between 0 and the vocabulary size \lvert \mathcal{V} \rvert
• t_i \in \mathbb{R}^{cs} is the temporal vector, containing the dates
• f_i \in \mathbb{R}^{2} corresponds to the patient general features, for now age and gender

Pre-training task

Objectives

• Having a time-aware Transformers (\neq NLP)
• Keeping in mind causal inference
  • We want to be able to efficiently compute the embedding of the pathway at any time (the trajectory of the patient)
• Problems with the next token task:
  • a lot of simultaneous events, especially for medications
  • different types of events: does it make sense to have a competition among them ?

Description

We predict for each token, if it will appear within the next t days.

We consider, simulatenously, several t \in \mathcal{T}, \mathcal{T} being the sorted set of short term horizons, typically 14, 30, 90, 180 days.

First trained models

We trained a “large” model (40M params) and a “small” one (5M).

ROC curve of the pre-training task for a given specialty on a test set of 10,000 patients.

Plot of the embedding space

3 Next steps

Causal inference

• Goal: Estimate the impact of losing one’s Médecin Traitant (Primary Care Provider, PCP), and looping back with Titouan’s work

• Up to a given time t, we have:

  • The covariate X_{<t} \coloneqq (e_{<t}, t_{<t}, f_{<t}), the medical history of the patient
  • W_t is the (binary) treatment at time t, say: the patient has lost its PCP (death or retirement) before t
  • Y_t (W_t), the binary outcome, indicating whether death has occured within a given time window [t, t+ \delta_d]

• We want to estimate the Conditional Average Treatment Effect (CATE): \tau(x, t) \coloneqq \mathbf{E}[Y_t(1) - Y_t(0) | X_{<t} = x]

Improving pre-training task

• Exploring Posson loss - ground truth being number of events for each token, for each time window
• Exponential loss for time-to-event task (as in MOTOR (Steinberg et al. (2024)))

We want ideally to have a generative model.

Scaling

• Hoffmann et al. (2022) propose a scaling law for LLMs

  • Trade off between number of training tokens (train dataset size) and number of parameters in the model (model size)

  • Estimated optimal ratio: \frac{\text{\# training tokens}}{\text{\# parameters}} = 20

    • is it the same for medical pathways ?

From Shmatko et al. (2025)

From Hoffmann et al. (2022)

References

Hoffmann, Jordan, Sebastian Borgeaud, Arthur Mensch, et al. 2022. Training Compute-Optimal Large Language Models. https://arxiv.org/abs/2203.15556.

Shmatko, Artem, Alexander Wolfgang Jung, Kumar Gaurav, et al. 2025. ‘Learning the Natural History of Human Disease with Generative Transformers’. Nature, ahead of print, September 17. https://doi.org/10.1038/s41586-025-09529-3.

Steinberg, Ethan, Jason Alan Fries, Yizhe Xu, and Nigam Shah. 2024. ‘MOTOR: A Time-to-Event Foundation Model for Structured Medical Records’. The Twelfth International Conference on Learning Representations. https://openreview.net/forum?id=NialiwI2V6.