---
title: Bi-Axial Transformer for EHR Analysis
url: https://www.emergentmind.com/topics/bi-axial-transformer-bat
type: topic
---

# Bi-Axial Transformer for EHR Analysis

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** [2508.12418].

## 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\) [2508.12418]. For patient \(i\), the dynamic observations have shape \(T_i \times D\), with \(T_i\) time points and \(D\) clinical variables, and a binary observation indicator matrix \(M_i \in \mathbb{R}^{T_i \times D}\) records whether each variable was observed at each time [2508.12418].

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** [2508.12418]. 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 [2508.12418].

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 \((t,d)\) pairs into one prohibitively long sequence [2508.12418].

## 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
\[
E_i \in \mathbb{R}^{T_i \times D \times E},
\]
so that each element \(e_{i,d}^{t} \in \mathbb{R}^E\) corresponds to one variable \(d\) at one time \(t\) [2508.12418].

Each observation embedding is described as a learned combination of four ingredients: the sensor reading \(x_{i,d}^{t}\), the missingness indicator \(m_{i,d}^{t}\), the sensor identity \(d\), and a temporal encoding derived from the absolute observation time \(t\) [2508.12418]. 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 [2508.12418]. A reconstructed compact form, explicitly identified as such in the paper’s technical synthesis, is
\[
e_{i,d}^{t} = \psi\!\left([x_{i,d}^{t}, m_{i,d}^{t}, \mathrm{onehot}(d)]\right)+PE(t),
\]
with the caveat that exact dimensional bookkeeping is not fully specified in the text [2508.12418].

Temporal encoding is continuous-time rather than index-based. BAT uses
\[
PE_k(t) =
\begin{cases}
\sin\left(\frac{t}{T^{k/\tau}}\right), & \text{if } k \equiv 0 \pmod 2 \\
\cos\left(\frac{t}{T^{(k-1)/\tau}}\right), & \text{if } k \equiv 1 \pmod 2
\end{cases}
\]
where \(t\) is the absolute observation time, \(T\) is a hyperparameter representing the maximum possible time point, and \(\tau\) is the encoding dimension [2508.12418]. 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 \(M_i\) to the observed values before the linear embedding stage. Second, it retains the full \(T \times D\) 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 [2508.12418].

## 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** [2508.12418]. Instead of flattening the entire \(T \times D\) 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 \(t\), sensor-axis attention allows one variable’s embedding to attend to other variables measured at that same time. For a fixed sensor \(d\), time-axis attention allows one timestamp to attend to other timestamps for that same variable [2508.12418]. The model processes the embedded tensor in **two parallel tracks** that apply these operations in opposite orders:

- **Track A**: sensor attention \(\rightarrow\) time attention  
- **Track B**: time attention \(\rightarrow\) 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
\[
H^{(1)} = F_t(F_s(E)), \qquad H^{(2)} = F_s(F_t(E)),
\]
where \(F_s\) denotes a sensor-axis encoder pass and \(F_t\) denotes a time-axis encoder pass [2508.12418]. 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 [2508.12418]. 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 [2508.12418].

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 [2508.12418].

## 4. Mathematical formulation and computational profile

BAT adopts standard scaled dot-product attention,
\[
\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^T}{\sqrt{z_k}}\right)V,
\]
with \(z_k\) the key dimension [2508.12418]. The novelty is not a new attention kernel but the decision to apply attention axis-wise to a \(T \times D \times E\) tensor that preserves variable identity and missingness structure.

For a fixed time \(t\), one may write the sensor slice as
\[
E^{(t)} = E[t,:,:] \in \mathbb{R}^{D \times E},
\]
and for a fixed sensor \(d\), the time slice as
\[
E^{(d)} = E[:,d,:] \in \mathbb{R}^{T \times E},
\]
which is the notation used in the BAT synthesis to describe sensor-axis and time-axis attention, respectively [2508.12418]. 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 \(TD\) cells.

That contrast is computationally consequential. If all cells were flattened into one sequence, the sequence length would be \(L = TD\), giving attention cost
\[
O((TD)^2)=O(T^2D^2).
\]
BAT instead applies axial attention, which the paper interprets in EHR terms as requiring \(O(TD^2)\) work for sensor attention and \(O(DT^2)\) work for time attention, for total
\[
O(TD(T+D)).
\]
The paper therefore argues that BAT is materially more scalable than full attention over all variable–time cells [2508.12418].

A concrete example is given for MIMIC-III. With maximum \(T=2881\) and \(D=16\), flattening would produce
\[
2881 \times 16 = 46{,}096
\]
tokens, which the paper treats as prohibitive for standard full self-attention [2508.12418]. 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)** [2508.12418]. A separate **HAR** dataset with artificially induced missingness is used for robustness analysis [2508.12418].

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 [2508.12418].

Training and evaluation use an **8:1:1** random split repeated over **5 different random splits**, with metrics averaged across splits [2508.12418]. 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 [2508.12418]. 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 [2508.12418].

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** [2508.12418].

Empirically, BAT is strongest on sepsis prediction. On P19 it achieves **77.05 \(\pm\) 1.92 AUPRC** and **95.41 \(\pm\) 0.58 AUROC**, which the paper describes as a **2.32-point AUPRC** and **0.58-point AUROC** improvement over the second-best model [2508.12418]. On P12 mortality it achieves **55.24 \(\pm\) 1.45 AUPRC** and **86.87 \(\pm\) 0.86 AUROC**; on MIMIC-III mortality it achieves **52.09 \(\pm\) 2.38 AUPRC** and **84.50 \(\pm\) 0.74 AUROC** [2508.12418]. These mortality results are competitive, but **GRU-D** remains best overall on both mortality datasets in the paper’s comparison [2508.12418].

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 [2508.12418]. 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 [2508.12418].

## 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 [2508.12418]. 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 \(\pm\) 3.11** to **52.07 \(\pm\) 1.51** AUPRC and from **84.22 \(\pm\) 2.14** to **84.98 \(\pm\) 1.33** AUROC, while MIMIC-III improves from **80.15 \(\pm\) 1.92** to **81.95 \(\pm\) 1.62** AUROC but decreases from **45.48 \(\pm\) 5.54** to **43.87 \(\pm\) 2.18** AUPRC [2508.12418]. 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 [2508.12418].

The BAT acronym is also overloaded on arXiv. In [2312.10611], **BAT** denotes **Bi-directional Adapter for Multi-modal Tracking** rather than a bi-axial transformer; in [2602.16305], it denotes **Better Audio Transformer**; and in [2407.12070], it denotes a **Binarized Transformer Accelerator**. A separate paper, **P2AT**, uses axial attention over **vertical and horizontal** image axes for semantic segmentation [2310.15025]. 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 [2508.12418].

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 [2508.12418].

Source: https://www.emergentmind.com/topics/bi-axial-transformer-bat