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

arXiv: https://arxiv.org/abs/2508.12418 · HTML: https://mlsanity.com/paper/2508.12418/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 55 references that are in the map.
- First posted on arXiv 17 Aug 2025.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Are Transformers Effective for Time Series Forecasting?](https://mlsanity.com/paper/2205.13504/) (2022)
- [iTransformer: Inverted Transformers Are Effective for Time Series Forecasting](https://mlsanity.com/paper/2310.06625/) (2023)
- [Recurrent Neural Networks for Multivariate Time Series with Missing Values](https://mlsanity.com/paper/1606.01865/) (2016)
- [ClinicalBERT: Modeling Clinical Notes and Predicting Hospital Readmission](https://mlsanity.com/paper/1904.05342/) (2019)
- [Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction](https://mlsanity.com/paper/2005.12833/) (2020)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [Long-term Forecasting with TiDE: Time-series Dense Encoder](https://mlsanity.com/paper/2304.08424/) (2023)
- [BEHRT: Transformer for Electronic Health Records](https://mlsanity.com/paper/1907.09538/) (2019)
- [TSMixer: Lightweight MLP-Mixer Model for Multivariate Time Series Forecasting](https://mlsanity.com/paper/2306.09364/) (2023)
- [Set Functions for Time Series](https://mlsanity.com/paper/1909.12064/) (2019)
- [Graph-Guided Network for Irregularly Sampled Multivariate Time Series](https://mlsanity.com/paper/2110.05357/) (2021)
- [Phased LSTM: Accelerating Recurrent Network Training for Long or Event-based Sequences](https://mlsanity.com/paper/1610.09513/) (2016)
- [Self-Supervised Transformer for Sparse and Irregularly Sampled Multivariate Clinical Time-Series](https://mlsanity.com/paper/2107.14293/) (2021)
- [Interpolation-Prediction Networks for Irregularly Sampled Time Series](https://mlsanity.com/paper/1909.07782/) (2019)
- [Multi-layer Representation Learning for Medical Concepts](https://mlsanity.com/paper/1602.05568/) (2016)
- [EHRXQA: A Multi-Modal Question Answering Dataset for Electronic Health Records with Chest X-ray Images](https://mlsanity.com/paper/2310.18652/) (2023)
- [Deep learning for temporal data representation in electronic health records: A systematic review of challenges and methodologies](https://mlsanity.com/paper/2107.09951/) (2021)

## Built on by (most cited first, 2 of 2)

- [Towards Self-Supervised Foundation Models for Critical Care Time Series](https://mlsanity.com/paper/2509.19885/) (2025)
- [MATA-Former & SIICU: Semantic Aware Temporal Alignment for High-Fidelity ICU Risk Prediction](https://mlsanity.com/paper/2604.01727/) (2026)

## Abstract (arXiv)

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.