Self-Supervised Transformer for Sparse and Irregularly Sampled Multivariate Clinical Time-Series

In plain words

Hospital monitoring data is messy. A patient has many measurements, such as heart rate and lab results, and each one is taken at its own irregular times. For most patients most measurements are missing. The usual fix is to squeeze everything onto a regular time grid and fill the gaps with estimated values. That blurs exactly the timing details that matter.

This paper keeps each measurement as it was recorded: when it was taken, which measurement it was, and its value. The model reads the whole collection of these records with a Transformer, which lets every record look at every other one. It turns the numbers into vectors with small learned networks instead of sorting them into bins. It also learns first from the unlabeled records, by predicting what each patient's measurements look like shortly afterwards. That helps most when few patients have labels. A companion version of the model reports how much each record pushed the prediction up or down.

TL;DR

• Represents a clinical time series as a set of observation triplets (time, variable, value), with no aggregation and no imputation §3.1 §3.2.
• Embeds time and value with learned one-hidden-layer networks (the paper's continuous value embedding), the variable with a lookup table, then applies Transformer blocks, attention pooling and a sigmoid head §3.2.1-§3.2.5.
• Pretrains on forecasting the variables observed in a later window, then fine-tunes on in-hospital mortality §3.2.6.
• Author-reported: best ROC-AUC, PR-AUC and min(Re,Pr) on MIMIC-III (24-hour mortality) and PhysioNet-2012 (48-hour mortality), with PR-AUC 3.2% and 3.5% above the best baseline §4.5 tab. 4.
• The released code departs from the paper in many places (hidden size, heads, learning rates, normalisation, residual form, loss scaling). Read ## Paper vs code before implementing.

Problem

Clinical time series are multivariate, very sparse and sampled at irregular times. The average variable missing rate is 89.7% for MIMIC-III and 79.7% for PhysioNet-2012 tab. 2. Aggregating to a regular grid discards fine timing, and interpolation or imputation adds noise and overhead that grows as the data gets sparser §1. Labelled outcomes are also scarce in many healthcare settings §1.

The target task is binary: predict in-hospital mortality from demographics and the series of one ICU stay §3.1. The evaluated windows are the first 24 hours for MIMIC-III and 48 hours for PhysioNet-2012 §4.1. The paper says the framework could also serve other supervised or unsupervised settings, but only binary classification is evaluated §1.

Method

Overview: each stay becomes a set of triplets, which are embedded, contextualised by a Transformer, pooled by attention and concatenated with a demographics embedding fig. 3 §3.2. The pieces, in order:

1. Input. A stay is a set of n triplets (t, f, v), with t ≥ 0, f a variable from a fixed set F and v a real value, plus demographics d and a label y §3.1. The unit of t is not stated.
2. Initial triplet embedding. e_i is the sum of a feature embedding (lookup table), a value embedding and a time embedding §3.2.1. Value and time use the continuous value embedding (CVE): a network with one input, a hidden layer of floor(sqrt(d)) tanh units and d outputs, of the form U tanh(Wx + b) §3.2.1. The output has no bias term in the stated form.
3. Contextual embedding. M Transformer blocks (M = 2 in the final tables) tab. 3. Each block has multi-head self-attention with h heads eq. 1 eq. 2, then a feed-forward layer with hidden width 2d eq. 3. Dropout, residual connections and layer normalisation are added to each sub-layer. During training, attention dropout masks some key positions before the softmax §3.2.2.
4. Fusion attention. For each contextual embedding c_i, a scalar score a_i and a softmax weight α_i are computed, and the series embedding e^T is the weighted sum of the c_i eq. 4 eq. 5 eq. 6 §3.2.3. The hidden width of this attention network is named d_a but its value is not stated in the compiled paper.
5. Demographics. A two-layer tanh network with hidden width 2d maps d to e^d eq. 7 §3.2.4. The authors report that a separate network worked better than encoding demographics as triplets with a default time §3.2.4.
6. Prediction. The logit is a linear layer over the concatenation of e^d and e^T, followed by a sigmoid eq. 8 §3.2.5. The target loss is cross-entropy §3.2.5.
7. Self-supervision. The forecasting head predicts all |F| variables from the same concatenation eq. 9 §3.2.6. The loss is a squared error over the variables that the forecast window observed, averaged by the number of samples N' eq. 10 §3.2.6. Forecasting was preferred to masking as a pretext task §3.2.6. The model is pretrained on forecasting, then the whole model is fine-tuned on the target task §3.2.6.
8. Forecast windows. For MIMIC-III, the observation window ends at x hours for x = 20, 24, ..., 124, starts at min(0, x - 24) hours, and is followed by a 2-hour prediction window. For PhysioNet-2012, the window is [0, x) for x = 12, 16, ..., 44 hours, with the same 2-hour prediction window. Several windows per series are used §4.1 fig. 2.
9. Interpretable variant (I-STraTS). The fusion attention is applied to the initial embeddings e_i instead of the contextual ones, and the raw demographics vector replaces the demographics embedding eq. 11 §3.3. Output is then a sum of per-demographic terms and per-observation terms, called contribution scores eq. 12 §3.3.

Pseudocode for the forward pass, following the paper's description §3.2:

# Inputs: times (n,), variables (n,), values (n,), demographics (D,)
# Output: probability of in-hospital mortality (scalar)  [§3.2.5]

def strats_predict(times, variables, values, demographics):
    # Continuous value embedding (CVE) of time and value  [§3.2.1]
    time_emb = cve(times, time_params)                  # (n, d)
    value_emb = cve(values, value_params)               # (n, d)
    # Learned lookup table for the variable name  [§3.2.1]
    variable_emb = variable_table[variables]            # (n, d)
    # The initial triplet embedding is the sum of the three parts  [§3.2.1]
    triplet_emb = time_emb + value_emb + variable_emb   # (n, d)

    # Transformer blocks let every triplet attend to every other  [§3.2.2]
    contextual_emb = triplet_emb
    for block in transformer_blocks:
        contextual_emb = transformer_block(contextual_emb)   # (n, d)

    # Attention pooling of the triplets into one series vector  [§3.2.3]
    scores = fusion_mlp(contextual_emb)                 # (n,)
    weights = softmax(scores)                           # (n,), sums to 1
    series_emb = weights @ contextual_emb               # (d,)

    # Demographics through their own tanh network  [§3.2.4]
    demo_emb = demographics_mlp(demographics)           # (d,)

    # Concatenate both embeddings and map to one logit  [§3.2.5]
    joint = concatenate(demo_emb, series_emb)           # (2d,)
    return sigmoid(output_weights @ joint + output_bias)

# Continuous value embedding: one scalar to d numbers  [§3.2.1]
def cve(x, weights, bias, out_weights):
    # One hidden layer of floor(sqrt(d)) tanh units, no output bias
    hidden = tanh(outer(x, weights) + bias)     # (n, floor(sqrt(d)))
    return hidden @ out_weights                 # (n, d)

# One Transformer block over the set of triplets  [§3.2.2]
def transformer_block(emb):
    # Each head works in d/h dimensions  [eq. 1]
    heads = []
    for head in range(num_heads):
        queries = emb @ query_proj[head]        # (n, d/h)
        keys = emb @ key_proj[head]             # (n, d/h)
        values = emb @ value_proj[head]         # (n, d/h)
        # Scaled dot-product scores, then softmax over the keys  [eq. 1]
        scores = queries @ keys.T / sqrt(d / num_heads)   # (n, n)
        # Training only: randomly mask some key positions  [§3.2.2]
        scores = attention_dropout(scores)
        heads.append(softmax(scores) @ values)  # (n, d/h)
    # Concatenate heads and project back to width d  [eq. 2]
    attended = concatenate(heads) @ output_proj   # (n, d)
    # Residual connection, layer normalisation and dropout  [§3.2.2]
    emb = layer_norm(emb + dropout(attended))
    # Feed-forward layer with hidden width 2d  [eq. 3]
    hidden = relu(emb @ w1 + b1)                  # (n, 2d)
    emb = layer_norm(emb + dropout(hidden @ w2 + b2))   # (n, d)
    return emb

# Target loss: cross-entropy on the mortality label  [§3.2.5]
def target_loss(probability, label):
    return -(label * log(probability)
             + (1 - label) * log(1 - probability))

# Self-supervised forecasting loss  [§3.2.6]
def forecast_loss(series_emb, demo_emb, forecast_values, forecast_mask,
                  num_samples):
    # Linear head over all variables, shared encoder  [eq. 9]
    joint = concatenate(demo_emb, series_emb)   # (B, 2d)
    predicted = forecast_head(joint)            # (B, V)
    # Squared error only where the window observed the variable  [eq. 10]
    squared = forecast_mask * (predicted - forecast_values) ** 2
    # Sum over the batch, divided by N' (number of samples)  [eq. 10]
    return sum(squared) / num_samples

Key equations, transcribed from the paper. Here ∘ denotes concatenation, as in the paper, not an element-wise product:

\mathbf{H}_j = \mathrm{softmax}\!\left(\frac{(\mathbf{E}\mathbf{W}^q_j)(\mathbf{E}\mathbf{W}^k_j)^T}{\sqrt{d/h}}\right)(\mathbf{E}\mathbf{W}^v_j), \quad j=1,\dots,h eq. 1

MHA(\mathbf{E}) = (\mathbf{H}_1 \circ \cdots \circ \mathbf{H}_h)\,\mathbf{W}_c eq. 2

\mathbf{F}(\mathbf{X}) = \mathrm{ReLU}(\mathbf{X}\mathbf{W}^f_1 + \mathbf{b}^f_1)\,\mathbf{W}^f_2 + \mathbf{b}^f_2 eq. 3

a_i = \mathbf{u}_a^T \tanh(\mathbf{W}_a\mathbf{c}_i + \mathbf{b}_a) eq. 4

\alpha_i = \frac{\exp(a_i)}{\sum_{j=1}^{n}\exp(a_j)} eq. 5

\mathbf{e}^T = \sum_{i=1}^{n}\alpha_i\mathbf{c}_i eq. 6

\mathbf{e}^d = \tanh\!\big(\mathbf{W}^d_2\,\tanh(\mathbf{W}^d_1\mathbf{d} + \mathbf{b}^d_1) + \mathbf{b}^d_2\big) eq. 7

\tilde{y} = \mathrm{sigmoid}\big(\mathbf{w}_o^T[\mathbf{e}^d \circ \mathbf{e}^T] + b_o\big) eq. 8

\tilde{\mathbf{z}} = \mathbf{W}_s[\mathbf{e}^d \circ \mathbf{e}^T] + \mathbf{b}_s eq. 9

\mathcal{L}_{ss} = \frac{1}{|N'|}\sum_{k=1}^{N'}\sum_{j=1}^{|\mathcal{F}|} m^k_j\,(\tilde{z}^k_j - z^k_j)^2 eq. 10

\tilde{y} = \mathrm{sigmoid}\Big(\mathbf{w}_o^T\big[\mathbf{d} \circ \sum_{i=1}^{n}\alpha_i\mathbf{e}_i\big] + b_o\Big) eq. 11

\tilde{y} = \mathrm{sigmoid}\Big(\sum_{j=1}^{D}\mathbf{w}_o[j]\,\mathbf{d}[j] + \sum_{i=1}^{n}\sum_{j=1}^{d}\alpha_i\,\mathbf{w}_o[j+D]\,\mathbf{e}_i[j] + b_o\Big) eq. 12

Equation numbers are counted from the LaTeX source in source order, because the compiled PDF was not available to check them.

Hyperparameters

Per-model values from the paper's hyperparameter table tab. 3. The paper gives one learning rate per model and dataset, with no separate values for pretraining and fine-tuning.

Training settings stated in the text:

Results

Author-reported. Setup: two EHR datasets, mortality prediction, STraTS compared against six baselines §4.1 §4.2. MIMIC-III is split 64:16:20 by patient. PhysioNet-2012 uses set-a as the test set and sets b and c split 80:20 into train and validation §4.1. Metrics are ROC-AUC, PR-AUC and min(Re, Pr), the best value of min(recall, precision) over thresholds §4.3. Each model is trained 10 times on a random 50% of the labelled data, and the table shows mean ± standard deviation tab. 4.

Dataset statistics tab. 2:

Mortality prediction tab. 4:

Headline claim: STraTS improves PR-AUC by 3.2% on MIMIC-III and 3.5% on PhysioNet-2012 over the best baseline §4.5. From tab. 4, the best baseline is GRU on MIMIC-III (0.559) and GRU-D on PhysioNet-2012 (0.481). The paper also says that STraTS is best on all metrics in both datasets §4.5.

Ablation, with and without self-supervision (ss+ and ss-) tab. 5:

The paper reads the ablation as: the interpretable variant scores slightly lower than STraTS, self-supervision improves both, and I-STraTS (ss+) beats STraTS (ss-) on all metrics on MIMIC-III and on PR-AUC on PhysioNet-2012 §4.6.

Label efficiency: fig. 4 (MIMIC-III) and fig. 5 (PhysioNet-2012) show performance for different fractions of labelled data, averaged over 10 runs. All models degrade as labels shrink. The paper says STraTS keeps a clear advantage when labels are scarce, which it attributes to self-supervision §4.5. The figures carry the values; the text gives no numbers for them.

Interpretability case study (one MIMIC-III stay, a patient who died on day 6). The I-STraTS probability of in-hospital mortality from day 1 is 0.94 §4.7. The top five variables by average contribution score tab. 6:

The paper presents this as one case, not as a validated attribution method §4.7.

Paper vs code

The official code is at commit e936cda2cff365f15febd8a71d1ca71242868091 (MIT). The paper's stated values were checked against src/ and run_main.sh. The repository's README says that its changes make some results deviate from the published numbers README.md.

• Attention scaling. The paper divides the query-key product by sqrt(d/h) eq. 1 §3.2.2. The code takes the raw dot product, without scaling modeling_strats.py.
• Residual connections and layer norm. The paper adds residual connections and layer normalisation after each sub-layer §3.2.2. The code averages the residual branch, (sub-layer + x) / 2, and has the layer-norm calls commented out modeling_strats.py modeling_strats.py. The README says LayerNorm was removed README.md.
• Feed-forward activation. The paper uses ReLU eq. 3. The code uses GELU modeling_strats.py.
• Demographics network. The paper's demographics embedding ends with a tanh eq. 7. The code is Linear, Tanh, Linear, with no output activation models.py.
• Time input. The paper uses raw t §3.1. The code converts minutes to [-1, 1] by dividing by a maximum (1440 for MIMIC-III and 2880 for PhysioNet-2012) dataset_pretrain.py dataset_pretrain.py dataset.py. Fine-tuning reuses the pretraining maximum dataset.py.
• Fine-tuning head. The paper's target head is one linear layer over the concatenation eq. 8. When a checkpoint is loaded, the code inserts the forecasting head (linear, d to V) before the binary head (linear, V to 1), with no nonlinearity between them models.py modeling_strats.py.
• Target loss. The paper uses cross-entropy §3.2.5. The code uses binary cross-entropy with a positive-class weight equal to negatives divided by positives in the training subset models.py dataset.py.
• Forecast loss normalisation. The paper divides by the number of samples N' eq. 10. The code divides by the number of observed forecast entries in the batch models.py.
• Forecast windows. The paper uses a fixed grid of observation windows §4.1. The code draws one random prediction time per stay from the observed times at or after 12 hours, with a 2-hour prediction window dataset_pretrain.py dataset_pretrain.py dataset_pretrain.py. The observation window reaches back at most 24 hours (MIMIC-III) or 48 hours (PhysioNet-2012), and at most max_obs observations dataset_pretrain.py. Pretraining on MIMIC-III uses only the first five days of data dataset_pretrain.py.
• Forecast target value. For each variable, the code keeps the earliest observation in the prediction window dataset_pretrain.py. The paper does not say which observation to use.
• Pretraining population. The code pretrains on every stay except the test split dataset_pretrain.py. The paper says unlabelled and labelled data are used §3.2.6 §4.5 but does not mention the test split.
• Pretraining patience and epoch. The paper uses patience 5 and an epoch of 256,000 samples §4.4. The code's default patience is 10 and it does not override it for pretraining main.py run_main.sh run_main.sh. An epoch is one pass over the training IDs main.py. Pretraining runs for 30 epochs on MIMIC-III and 100 on PhysioNet-2012 run_main.sh run_main.sh.
• Batch size. The paper uses 32 §4.4. The code's default is 16 and run_main.sh does not override it main.py. Evaluation uses 32 main.py.
• Hidden size, heads and learning rate. The paper's STraTS uses d = 50 with h = 4 heads and lr 5e-4 on both datasets tab. 3. The released scripts use hidden size 64 and 16 heads (head width 4), with lr 5e-4 for pretraining and for the ss- model, and lr 5e-5 for fine-tuning run_main.sh run_main.sh. The README says all models use hidden size 64 README.md.
• Gradient clipping. Not in the paper. The code clips the gradient norm to 0.3 main.py.
• Variable dropout. Not in the paper. During training, with probability equal to the dropout rate (0.2), the code removes all observations of a variable from a stay, in both pretraining and fine-tuning modeling_strats.py. The README lists this as an addition README.md.
• Observation cap. The paper caps observations at the 99th percentile for MIMIC-III §4.4. The code's default is 880 for both datasets, and run_main.sh does not change it. In fine-tuning, the code keeps a random subset of the observations of each stay dataset.py main.py.
• Training subsets. The paper uses 10 random 50% samples §4.5. The code's run index picks a contiguous block of the stored ID list, with blocks spaced evenly across the list dataset.py. The same loop also runs fractions 0.4 down to 0.1 run_main.sh.
• Framework. The paper says Keras with TensorFlow §4.4. The code is PyTorch modeling_strats.py.
• Consistent with the paper. The CVE hidden size, the I-STraTS contribution form eq. 11 modeling_strats.py, the demographic preprocessing for PhysioNet-2012 §4.1 dataset.py, and the metrics §4.3 evaluator.py match the paper.

Pitfalls

• Dimensions of the output layer. The paper writes w_o^T as R^d §3.2.5 but concatenates two d-vectors eq. 8. The input to the output layer is 2d (or d + D for I-STraTS). The code uses 2d modeling_strats.py models.py.
• Notation. ∘ means concatenation in equations 2, 8, 9 and 11, not an element-wise product.
• Unstated values. The fusion hidden width d_a §3.2.3, the learning-rate schedule, the Adam betas, the time unit §3.1, the MIMIC-III observation cap as a number §4.4 and the separate pretraining and fine-tuning learning rates are not stated in the compiled paper. The code's values are listed under Paper vs code. For d_a, the code uses d modeling_strats.py.
• Head dimension. The paper does not define the per-head width. The scaling and the code both imply d/h eq. 1 modeling_strats.py.
• Attention scaling. Without sqrt(d/h) scaling, attention logits grow with width. A reimplementation that follows the paper's equation will differ from the released code modeling_strats.py.
• Training loop skips NaN losses. Batches with a NaN loss are skipped without a warning main.py. Check the log for skipped steps.
• PR-AUC definition. The code computes the area under the precision-recall curve with the trapezoidal rule, not average precision evaluator.py.
• Set size and compute. Attention cost grows with the number of observations, which is why the paper caps observations for MIMIC-III. The conclusion names longer series as future work §5.
• Data access. MIMIC-III requires credentialed access on PhysioNet. The preprocessing scripts are the only route from raw data to the processed split README.md. The processed data is not in the repository.
• Age cleaning. The code replaces ages above 200 with 91.4 in MIMIC-III dataset.py. The paper does not mention this.
• Variable set. When training from scratch, the code keeps only variables seen in the labelled training subset dataset.py, so the feature count changes with the training fraction. Fine-tuning fixes it to the pretraining set.
• Normalisation statistics. Fine-tuning reuses the pretraining statistics for values, and the code loads them from the pretraining run dataset.py.
• Results are author-reported and the code differs. Expect different numbers from the released scripts than from Table 4 README.md.
• Interpretability. The case study is one stay §4.7. The contribution scores are linear in the output logit, so they explain the model's logit, not the probability.

Lineage

• Builds on: the Transformer encoder with multi-head attention §3.2.2, which the paper cites as Vaswani et al. Continuous-time models it contrasts with: GRU-D §2, and interpolation networks §2.
• Compared against: GRU, TCN, SAnD, GRU-D, InterpNet and SeFT §4.2. SeFT is the closest in input form, since it also takes a set of observation triplets, but it uses sinusoidal time encodings §4.2.
• Self-supervision for time series: the paper says it is the first to explore self-supervised learning for sparse, irregular multivariate series §2. It contrasts this with a denoising Transformer pretraining method for dense multivariate series §2.