Bi-Axial Transformer for EHR Analysis

• Bi-Axial Transformer (BAT) is a transformer architecture that applies attention along both time and clinical variable axes to preserve EHR structure and informative missingness.
• BAT embeds sensor readings, missingness indicators, sensor identities, and absolute time into a 3D tensor, processing it with parallel axial-attention tracks.
• Empirical results demonstrate BAT’s state-of-the-art performance on sepsis prediction and competitive outcomes on mortality classification in EHR datasets.

Searching arXiv for the BAT paper and closely related axial-transformer context. The Bi-Axial Transformer (BAT) is a transformer architecture for classification on sparse, irregular, multivariate EHR time series that preserves the clinically salient time $\times$ variable structure of patient records and applies attention along both axes rather than collapsing the data into a single sequence. In the formulation used for EHR classification, BAT embeds each variable–time cell with its value, missingness indicator, sensor identity, and observation time; processes the resulting tensor in two parallel axial-attention tracks with opposite attention orderings; and is reported to achieve state-of-the-art performance on sepsis prediction while remaining competitive to top methods for mortality classification (DeVries et al., 17 Aug 2025).

1. Clinical setting and modeling objective

BAT is defined on EHR classification datasets of the form

$$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$

where patient $i$ has a multivariate time series $X_i$, observation times $t_i$, static demographic data $s_i$, and label $y_i$ (DeVries et al., 17 Aug 2025). For patient $i$, the dynamic observations have shape $T_i \times D$, with $T_i$ time points and $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$0 clinical variables, and a binary observation indicator matrix $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$1 records whether each variable was observed at each time (DeVries et al., 17 Aug 2025).

The architectural motivation is tied to several properties of EHRs emphasized in the BAT paper: they are irregularly sampled, highly sparse, multivariate, multimodal, and often exhibit informative missingness (DeVries et al., 17 Aug 2025). The last point is especially central. BAT is motivated by the observation that the absence of a measurement can itself be predictive, because measurement patterns often reflect clinician behavior and patient state rather than random omission (DeVries et al., 17 Aug 2025).

The paper positions BAT against three common representation strategies. A conventional transformer over time embeds each timepoint as one token and therefore primarily models temporal relationships among timepoints. An inverted transformer represents each variable trajectory as one token and therefore primarily models relationships among variables. Tuple or set-based methods such as SeFT operate only on observed tuples and may therefore fail to preserve the full missingness pattern. BAT is proposed because it aims to preserve informative missingness, model dependencies along both the time axis and the clinical variable axis, and avoid flattening all $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$2 pairs into one prohibitively long sequence (DeVries et al., 17 Aug 2025).

2. Input representation and explicit encoding of missingness

BAT does not embed an entire timepoint or an entire variable trajectory into a single vector. Instead, after embedding, each patient trajectory becomes

$$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$3

so that each element $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$4 corresponds to one variable $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$5 at one time $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$6 (DeVries et al., 17 Aug 2025).

Each observation embedding is described as a learned combination of four ingredients: the sensor reading $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$7, the missingness indicator $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$8, the sensor identity $$\mathcal{D}:= \{(X_1, t_1, s_1, y_1), \ldots, (X_N, t_N, s_N, y_N)\},$$9, and a temporal encoding derived from the absolute observation time $i$0 (DeVries et al., 17 Aug 2025). The paper states that the sensor reading is concatenated with the binary mask, embedded with a linear layer, joined with a learned sensor identity embedding, and then augmented by a sinusoidal time encoding (DeVries et al., 17 Aug 2025). A reconstructed compact form, explicitly identified as such in the paper’s technical synthesis, is

$i$1

with the caveat that exact dimensional bookkeeping is not fully specified in the text (DeVries et al., 17 Aug 2025).

Temporal encoding is continuous-time rather than index-based. BAT uses

$i$2

where $i$3 is the absolute observation time, $i$4 is a hyperparameter representing the maximum possible time point, and $i$5 is the encoding dimension (DeVries et al., 17 Aug 2025). The explicit use of absolute time is intended to better fit irregularly sampled EHRs than positional encodings tied only to token order.

Missingness is preserved in two ways. First, BAT explicitly concatenates the indicator matrix $i$6 to the observed values before the linear embedding stage. Second, it retains the full $i$7 tensor rather than removing unobserved entries. The paper argues that this allows BAT to exploit informative missingness more directly than observed-only tuple representations, because the model sees both the explicit mask feature and the structured pattern of absent entries across time and variables (DeVries et al., 17 Aug 2025).

3. Bi-axial attention architecture

The term bi-axial refers to attention along two axes of the EHR tensor: the time axis and the clinical variable or sensor axis (DeVries et al., 17 Aug 2025). Instead of flattening the entire $i$8 grid into one sequence, BAT uses axial attention so that attention is computed along one axis while holding the other fixed.

For a fixed time $i$9, sensor-axis attention allows one variable’s embedding to attend to other variables measured at that same time. For a fixed sensor $X_i$0, time-axis attention allows one timestamp to attend to other timestamps for that same variable (DeVries et al., 17 Aug 2025). The model processes the embedded tensor in two parallel tracks that apply these operations in opposite orders:

• Track A: sensor attention $X_i$1 time attention
• Track B: time attention $X_i$2 sensor attention

This ordering is treated as architecturally meaningful. After the first axis-specific attention pass, each embedding is already contextually enriched along one axis; the second pass then propagates that contextual information along the other axis. A reconstructed summary provided in the paper’s technical synthesis is

$X_i$3

where $X_i$4 denotes a sensor-axis encoder pass and $X_i$5 denotes a time-axis encoder pass (DeVries et al., 17 Aug 2025). The paper identifies this as a reconstruction from the architecture description rather than a directly typeset equation.

Implementation is described as straightforward: BAT transposes either the time or sensor dimension onto the batch dimension and then applies standard self-attention machinery to the resulting axis-specific slices (DeVries et al., 17 Aug 2025). Each axis attention pass uses a bi-directional transformer encoder with self-attention heads, normalization layers, and FFNs, and the paper notes that encoder parameters are shared across axis attention passes within the same track (DeVries et al., 17 Aug 2025).

After both rounds of attention in each track, the weighted embeddings are pooled with either mean or max pooling, the two pooled track outputs are concatenated and merged through a linear layer plus ReLU, and a separate linear embedding of static demographics is fused before a two-layer binary classification head produces the final prediction (DeVries et al., 17 Aug 2025).

4. Mathematical formulation and computational profile

BAT adopts standard scaled dot-product attention,

$X_i$6

with $X_i$7 the key dimension (DeVries et al., 17 Aug 2025). The novelty is not a new attention kernel but the decision to apply attention axis-wise to a $X_i$8 tensor that preserves variable identity and missingness structure.

For a fixed time $X_i$9, one may write the sensor slice as

$t_i$0

and for a fixed sensor $t_i$1, the time slice as

$t_i$2

which is the notation used in the BAT synthesis to describe sensor-axis and time-axis attention, respectively (DeVries et al., 17 Aug 2025). The paper does not present a single official block equation for the entire bi-axial encoder, but it explicitly contrasts BAT with full flattened attention over all $t_i$3 cells.

That contrast is computationally consequential. If all cells were flattened into one sequence, the sequence length would be $t_i$4, giving attention cost

$t_i$5

BAT instead applies axial attention, which the paper interprets in EHR terms as requiring $t_i$6 work for sensor attention and $t_i$7 work for time attention, for total

$t_i$8

The paper therefore argues that BAT is materially more scalable than full attention over all variable–time cells (DeVries et al., 17 Aug 2025).

A concrete example is given for MIMIC-III. With maximum $t_i$9 and $s_i$0, flattening would produce

$s_i$1

tokens, which the paper treats as prohibitive for standard full self-attention (DeVries et al., 17 Aug 2025). This computational argument is integral to BAT’s design: the model seeks to preserve the native EHR grid while avoiding the quadratic cost of unrestricted attention over all grid cells.

5. Training protocol and empirical evidence

BAT is evaluated on three public EHR datasets and one auxiliary robustness benchmark. The EHR tasks are mortality classification on PhysioNet Challenge 2012 (P12) and MIMIC-III, and patient-level binary sepsis classification on PhysioNet Challenge 2019 (P19) (DeVries et al., 17 Aug 2025). A separate HAR dataset with artificially induced missingness is used for robustness analysis (DeVries et al., 17 Aug 2025).

The paper reports the following dataset characteristics: P12 has 11,988 samples, maximum 215 time points, 37 sensors, 14.23% positive class, and 84.32% sparsity; P19 has 40,333 samples, maximum 336 time points, 34 sensors, 7.27% positive class, and 80.13% sparsity; MIMIC-III has 21,107 samples, maximum 2881 time points, 16 sensors, 13.22% positive class, and 65.50% sparsity (DeVries et al., 17 Aug 2025).

Training and evaluation use an 8:1:1 random split repeated over 5 different random splits, with metrics averaged across splits (DeVries et al., 17 Aug 2025). Class imbalance is handled by using all positive samples, resampled 3 times, together with an equal number of randomly selected negative samples in each epoch (DeVries et al., 17 Aug 2025). BAT and the transformer baselines are trained with AdamW, and hyperparameters are selected by random sweep with 20 sampled configurations per dataset/model, choosing the best by AUROC (DeVries et al., 17 Aug 2025).

The best BAT hyperparameters selected in the supplementary table are shallow. For P12, BAT uses dropout 0.1, attention dropout 0.4, 2 heads, 1 layer, max pooling, learning rate 0.0001, embedding size 128, and batch size 16. For MIMIC-III, it uses dropout 0.4, attention dropout 0.2, 1 head, 1 layer, mean pooling, learning rate 0.0001, embedding size 16, and batch size 4. For P19, it uses dropout 0.4, attention dropout 0.4, 1 head, 1 layer, max pooling, learning rate 0.0001, embedding size 32, and batch size 16 (DeVries et al., 17 Aug 2025).

Empirically, BAT is strongest on sepsis prediction. On P19 it achieves 77.05 $s_i$2 1.92 AUPRC and 95.41 $s_i$3 0.58 AUROC, which the paper describes as a 2.32-point AUPRC and 0.58-point AUROC improvement over the second-best model (DeVries et al., 17 Aug 2025). On P12 mortality it achieves 55.24 $s_i$4 1.45 AUPRC and 86.87 $s_i$5 0.86 AUROC; on MIMIC-III mortality it achieves 52.09 $s_i$6 2.38 AUPRC and 84.50 $s_i$7 0.74 AUROC (DeVries et al., 17 Aug 2025). These mortality results are competitive, but GRU-D remains best overall on both mortality datasets in the paper’s comparison (DeVries et al., 17 Aug 2025).

Relative to transformer baselines, the paper emphasizes that BAT outperforms both the standard Transformer and iTransformer on average, reporting an average 8.95-point AUPRC gain over iTransformer and 1.46-point AUPRC gain over the standard Transformer (DeVries et al., 17 Aug 2025). The robustness analysis on HAR further reports that BAT outperforms Transformer and iTransformer at all induced missingness levels from 0% to 99%, with lower variance across splits, although all models approach random performance near 99% sparsity (DeVries et al., 17 Aug 2025).

6. Interpretation, transferability, and nomenclature

A distinctive claim of BAT is that it learns explicit sensor embeddings rather than relying only on fixed input positions to identify variables (DeVries et al., 17 Aug 2025). The paper argues that these embeddings may support transfer learning, cross-dataset harmonization, dataset integration, and some degree of interpretability. The strongest direct evidence is a joint P12–MIMIC-III experiment comparing separate and shared sensor embeddings. With shared embeddings for overlapping sensor identities, P12 improves from 50.71 $s_i$8 3.11 to 52.07 $s_i$9 1.51 AUPRC and from 84.22 $y_i$0 2.14 to 84.98 $y_i$1 1.33 AUROC, while MIMIC-III improves from 80.15 $y_i$2 1.92 to 81.95 $y_i$3 1.62 AUROC but decreases from 45.48 $y_i$4 5.54 to 43.87 $y_i$5 2.18 AUPRC (DeVries et al., 17 Aug 2025). The paper interprets the reduced variance under shared embeddings as evidence of more stable learning.

The same section of the paper also underscores BAT’s limits. It is not uniformly best against all model classes, because GRU-D still leads on the mortality benchmarks. Several implementation details remain implicit, including exact embedding concatenation shapes, pooling dimensions, and some parameter-sharing specifics. The best-performing BAT configurations are often only 1 layer deep, so the depth-scaling behavior remains unclear. The P19 evaluation is simplified from timepoint-level sepsis detection to patient-level binary sepsis classification, which leaves early-event prediction as an open question (DeVries et al., 17 Aug 2025).

The BAT acronym is also overloaded on arXiv. In (Cao et al., 2023), BAT denotes Bi-directional Adapter for Multi-modal Tracking rather than a bi-axial transformer; in (Ghaffari et al., 18 Feb 2026), it denotes Better Audio Transformer; and in (Ji et al., 2024), it denotes a Binarized Transformer Accelerator. A separate paper, P2AT, uses axial attention over vertical and horizontal image axes for semantic segmentation (Elhassan et al., 2023). In contrast, the Bi-Axial Transformer for EHRs attends over the time axis and the clinical variable axis, making its notion of “bi-axial” specific to structured medical time series rather than spatial vision grids (DeVries et al., 17 Aug 2025).

In that sense, BAT reframes EHR classification as structured sparse tensor modeling. Its central proposition is that EHRs should not be treated only as sequences of visits or only as collections of variable trajectories, but as sparse time-by-variable matrices whose missingness patterns, variable identities, and temporal structure all carry predictive signal (DeVries et al., 17 Aug 2025).