# Imputation-free transformer learning enables robust Alzheimer's disease prediction and calibrated uncertainty quantification across heterogeneous clinical cohorts

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 42 references that are in the map.
- First posted on arXiv 13 Jul 2026.

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

- [Captum: A unified and generic model interpretability library for PyTorch](https://mlsanity.com/paper/2009.07896/) (2020)
- [TabTransformer: Tabular Data Modeling Using Contextual Embeddings](https://mlsanity.com/paper/2012.06678/) (2020)
- [SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training](https://mlsanity.com/paper/2106.01342/) (2021)
- [A Survey on Uncertainty Quantification Methods for Deep Learning](https://mlsanity.com/paper/2302.13425/) (2023)
- [Understanding Softmax Confidence and Uncertainty](https://mlsanity.com/paper/2106.04972/) (2021)
- [Not Another Imputation Method: A Transformer-based Model for Missing Values in Tabular Datasets](https://mlsanity.com/paper/2407.11540/) (2024)
- [Confidence estimation of classification based on the distribution of the neural network output layer](https://mlsanity.com/paper/2210.07745/) (2022)
- [Entropy methods for the confidence assessment of probabilistic classification models](https://mlsanity.com/paper/2103.15157/) (2021)
- [Imputation-free Learning of Tabular Data with Missing Values using Incremental Feature Partitions in Transformer](https://mlsanity.com/paper/2504.14610/) (2025)

## Abstract (arXiv)

Accurate diagnostic classification and disease-severity prediction for Alzheimer's disease are hampered by the incompleteness and heterogeneity of real-world clinical data. Left unaddressed, these barriers prevent reliable disease modelling and hinder effective clinical evaluation. Conventional imputation strategies introduce systematic bias, distort inter-feature relationships, and yield overconfident predictions, limitations especially consequential in diagnostic settings. Here, we propose NITROGEN, an imputation-free transformer that jointly models within-patient feature dependencies and between-patient relational structure through masked and intersample attention, enabling robust multimodal learning directly from partially observed records. We trained NITROGEN on ADNI (N=7858 scans), and evaluated it on two independent cohorts: OASIS-3 (N=2675 scans) and AIBL (N=1286 scans). Across cohorts and diagnostic and cognitive score prediction tasks, NITROGEN showed robust calibration and uncertainty quantification advantages over tree-based ensemble methods, while maintaining competitive discriminative performance. Cross-cohort and cross-method analyses identified cortical thickness in the temporal pole, age, and APOE genotype as important, though not individually sufficient, features for AD classification. We further introduced a modality-aware uncertainty adjustment that augments predictive uncertainty proportionally to the importance of absent modalities, enabling calibrated confidence when diagnostic information is unavailable. Together, our results show that imputation-free attention learning preserved meaningful discrimination under cohort shift, revealing expected degradation on more distributionally different cohorts, and demonstrate that evaluating models along calibration, interpretability, and cross-cohort reliability, not accuracy alone, is essential for clinical deployment.