Bi-Axial Transformers: Addressing the Increasing Complexity of EHR Classification

TL;DR

The paper introduces Bi-Axial Transformer (BAT), a novel model for EHR classification that captures complex data relationships, improves robustness to missing data, and achieves state-of-the-art results in sepsis prediction.

Contribution

We propose the Bi-Axial Transformer, a new architecture that attends to both variable and time axes in EHR data, enhancing modeling of data relationships and handling missingness.

Findings

BAT achieves state-of-the-art sepsis prediction performance.

BAT is more robust to data missingness than other transformers.

Re-implemented baseline models are publicly available for benchmarking.

Abstract

Electronic Health Records (EHRs), the digital representation of a patient's medical history, are a valuable resource for epidemiological and clinical research. They are also becoming increasingly complex, with recent trends indicating larger datasets, longer time series, and multi-modal integrations. Transformers, which have rapidly gained popularity due to their success in natural language processing and other domains, are well-suited to address these challenges due to their ability to model long-range dependencies and process data in parallel. But their application to EHR classification remains limited by data representations, which can reduce performance or fail to capture informative missingness. In this paper, we present the Bi-Axial Transformer (BAT), which attends to both the clinical variable and time point axes of EHR data to learn richer data relationships and address the difficulties of data sparsity. BAT achieves state-of-the-art performance on sepsis prediction and is competitive to top methods for mortality classification. In comparison to other transformers, BAT demonstrates increased robustness to data missingness, and learns unique sensor embeddings which can be used in transfer learning. Baseline models, which were previously located across multiple repositories or utilized deprecated libraries, were re-implemented with PyTorch and made available for reproduction and future benchmarking.