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

2026-07-13

AI Summary

This study addresses model bias, distorted relationships, and overconfidence arising from missingness and heterogeneity in clinical data for Alzheimer’s disease prediction. To this end, the authors propose NITROGEN, a Transformer-based model that operates without data imputation. By integrating masked attention and inter-sample attention mechanisms, NITROGEN jointly models intra-patient feature dependencies and inter-patient relationships directly from partially observed multimodal clinical data. Furthermore, it incorporates a modality-aware uncertainty calibration mechanism to enable robust prediction and reliable uncertainty quantification across heterogeneous cohorts. Evaluated with ADNI for training and OASIS-3 and AIBL for validation, NITROGEN achieves competitive discriminative performance while significantly outperforming tree-based ensemble methods in calibration and uncertainty estimation, and identifies diagnostically relevant yet insufficient biomarkers.

Abstract

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.

Problem

• Alzheimer's disease
• data incompleteness
• data heterogeneity
• clinical cohorts
• diagnostic prediction

Innovation

• imputation-free learning
• transformer architecture
• uncertainty quantification
• multimodal clinical data
• cross-cohort generalization