🧬 biology

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

The paper introduces NITROGEN, an imputation-free transformer model that achieves robust Alzheimer's disease prediction and calibrated uncertainty quantification across heterogeneous clinical cohorts by directly learning from partially observed data without relying on bias-inducing imputation strategies.

Original authors: Christelle Schneuwly Diaz, Narmina Baghirova, Duy-Thanh Vu, Duy-Cat Can, Gilles Allali, Philippe Ryvlin, Oliver Y. Chén

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Christelle Schneuwly Diaz, Narmina Baghirova, Duy-Thanh Vu, Duy-Cat Can, Gilles Allali, Philippe Ryvlin, Oliver Y. Chén

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to solve a massive, complex puzzle: predicting whether a person has Alzheimer's disease. Usually, doctors and computers try to solve this by looking at a huge collection of clues: brain scans (MRI), genetic tests (like the APOE gene), spinal fluid samples, memory test scores, and basic info like age and education.

But here's the messy reality of real life: not everyone has all the puzzle pieces. Some patients have great brain scans but no spinal fluid test. Others have genetic data but no memory scores. Some are missing almost everything except their age.

The Old Way: "Fake It Till You Make It" (And Why It Fails)

For years, the standard way to handle these missing pieces was to guess what they might be. Computers would fill in the blanks with "synthetic" or made-up numbers based on what they knew about other people. The paper calls this the "impute-then-predict" strategy.

The authors argue this is a bad idea. It's like trying to solve a mystery by inventing fake evidence. If you guess a patient's spinal fluid levels, you might accidentally create a "phantom" clue that tricks the computer. The computer might then say, "I am 99% sure this person has Alzheimer's!" with total confidence, even though that confidence is built on a lie. This leads to overconfident predictions, which is dangerous in medicine.

The New Hero: NITROGEN

Enter NITROGEN (No-imputation Inter-sample-attention TRansformer Oriented for GENeral datasets). Think of NITROGEN as a super-smart detective who refuses to guess.

Instead of filling in the blanks, NITROGEN looks at the clues that are there and says, "Okay, I'll work with what I have." It uses a special trick called masked attention. Imagine you are reading a story where some words are crossed out. Instead of trying to guess the missing words, NITROGEN simply ignores the crossed-out parts and focuses on the words that are visible. It learns to understand the relationships between the clues it does see, without ever inventing fake ones.

But NITROGEN has a second superpower: Intersample Attention. While it looks at one patient's missing pieces, it also looks at other patients in the room. If Patient A is missing a brain scan, but Patient B (who looks very similar) has one, NITROGEN can borrow that context to make a better guess. It learns from the whole crowd, not just the individual.

The Big Test: Can It Handle Real Chaos?

The researchers trained NITROGEN on a massive dataset called ADNI (which had 7,858 scans). Then, they threw it into the deep end without any extra training to see if it could handle two completely different groups of people:

OASIS-3 (2,675 scans): A group with very different characteristics (mostly healthy people, very few Alzheimer's cases).
AIBL (1,286 scans): An Australian study with its own unique quirks.

The Results:

It didn't break. Even when the data was messy and different from what it learned, NITROGEN stayed calm.
It was honest. Unlike other computer models (like the popular tree-based ones) that got overconfident and made mistakes while sounding sure, NITROGEN gave calibrated answers. This means if it said there was a 75% chance of disease, it was actually right about 75% of the time.
It knew when to say "I don't know." The team added a special feature: if a crucial clue (like a brain scan) was missing, NITROGEN would automatically raise its "uncertainty" alarm. It's like a detective saying, "I can't solve this yet because the most important evidence is missing," rather than guessing and being wrong.

The "Magic" Clue: The Temporal Pole

The researchers asked: "What is the single most important clue?" They found that across all the different groups and computer models, one specific part of the brain kept popping up as the star player: the temporal pole (a tiny region at the very front of the brain's side).

The Finding: The thickness of this specific brain area was a huge indicator of Alzheimer's.
The Catch: The paper tested if this one clue was enough to solve the whole puzzle. The answer? No. While the temporal pole was a strong signal, it wasn't enough on its own to make a perfect diagnosis. You still need the other clues (age, genetics, etc.) to get a reliable answer. The paper suggests that while this brain region is a consistent "hero," it needs a team to win the game.

The Bottom Line

This paper suggests that for medical AI to be truly useful, it needs to stop pretending it knows everything. It needs to handle missing data without faking it, admit when it's unsure, and be honest about its confidence. NITROGEN shows that by doing this, we can build models that are not just smart, but also trustworthy enough to help doctors in the real world, where data is rarely perfect.

The authors measured this using real data from thousands of people, not just simulations. They found that while NITROGEN wasn't always the absolute fastest or highest-scoring model in every single test, it was the most reliable and honest one, especially when the data was messy or different from what it was trained on.

Technical Summary: NITROGEN for Imputation-Free Alzheimer's Disease Prediction

Problem Statement
Accurate diagnostic classification and disease-severity prediction for Alzheimer's disease (AD) are hindered by the pervasive incompleteness and heterogeneity of real-world clinical data. Multimodal datasets typically integrate demographic, genetic, neuroimaging, and molecular measurements, yet these modalities are unevenly available across individuals and studies due to varying acquisition protocols and patient compliance. Conventional "impute-then-predict" strategies address this by replacing missing values with synthetic data. However, the authors argue that such imputation introduces systematic bias, distorts inter-feature relationships, and yields overconfident predictions. These limitations are particularly consequential in diagnostic settings where reliable uncertainty quantification is essential for clinical decision-making. Furthermore, existing missing-aware models often fail to distinguish between uncertainty arising from the inherent difficulty of a case and uncertainty arising from the absence of critical diagnostic information.

Methodology: The NITROGEN Framework
The authors propose NITROGEN (No-imputation Inter-sample-attention TRansformer Oriented for GENeral datasets), an imputation-free transformer architecture designed to learn directly from partially observed records. The framework integrates two key mechanisms:

Masked Feature Attention (NAIM): Based on the NAIM model, NITROGEN employs a masked self-attention mechanism that explicitly ignores missing values during representation learning. Instead of imputing data, the model uses a masking matrix to suppress the contribution of missing features within the attention computation, ensuring that synthetic evidence is never introduced.
Intersample Attention (SAINT-inspired): To stabilize predictions in heterogeneous cohorts, the architecture incorporates an intersample attention module. This allows the model to learn relational structures across different patients within a mini-batch, leveraging shared patterns to compensate for missing or noisy inputs at the individual level.

The model was trained on the Alzheimer's Disease Neuroimaging Initiative (ADNI) dataset ( $N=7,858$ scans) and evaluated without retraining on two independent external cohorts: OASIS-3 ( $N=2,675$ ) and AIBL ( $N=1,286$ ). The evaluation covered three tasks: binary classification (CN vs. AD), multi-class classification (CN vs. MCI vs. AD), and continuous cognitive score prediction (MMSE and CDR-SB).

To address the reliability of probability estimates under incomplete data, the authors introduced a modality-aware uncertainty adjustment. This method scales the predictive uncertainty (Shannon entropy) of a sample based on the estimated importance of the missing modalities. If a highly informative modality (e.g., CSF biomarkers) is absent, the model's confidence is explicitly reduced, providing a calibrated measure of trustworthiness that reflects data availability.

Key Results

Robustness and Calibration: Across all cohorts and tasks, NITROGEN demonstrated superior probability calibration compared to tree-based ensemble methods (XGBoost, LightGBM, Random Forest). While tree-based models often achieved high discriminative accuracy (AUC) on in-distribution data, they exhibited systematic overconfidence, particularly under domain shift (e.g., on the OASIS-3 cohort). NITROGEN maintained well-calibrated predictions, showing a larger gap in prediction probabilities between correct and incorrect cases, indicating better uncertainty awareness.
Cross-Cohort Generalization: The model showed robust transfer to external cohorts. Performance on AIBL was strong, likely due to shared recruitment strategies with ADNI. On OASIS-3, which targets preclinical populations with a different distribution, performance degraded as expected for all models, but NITROGEN maintained competitive sensitivity without the extreme variance observed in some tree-based models.
Feature Attribution and Sufficiency: Permutation-based feature attribution identified cortical thickness in the bilateral temporal pole, APOE $\epsilon$ 4 genotype, and age as the most consistently important features across all models and cohorts. A sufficiency analysis revealed that while the temporal pole feature alone carries a reproducible ranking signal for AD status, it is insufficient for stable decision-level classification without complementary demographic and genetic context.
Uncertainty Adjustment: The proposed modality-aware uncertainty adjustment improved selective prediction reliability. By excluding predictions with high adjusted uncertainty, the model achieved higher accuracy on the retained subset compared to using raw entropy, demonstrating that the adjustment effectively identifies less reliable predictions when critical modalities are missing.

Significance and Claims
The paper claims that NITROGEN demonstrates that imputation-free attention learning can preserve meaningful discrimination under external cohort shifts while revealing expected performance degradation on distributionally different data. The authors emphasize that evaluating machine learning models for neurological disease requires dimensions beyond predictive accuracy, specifically calibration, interpretability, and cross-cohort reliability.

The study concludes that:

Imputation is not necessary: Meaningful multimodal learning can occur directly from incomplete records without introducing synthetic bias.
Uncertainty must be context-aware: In clinical settings, a model's confidence should reflect not just the difficulty of the case but also the completeness of the diagnostic evidence.
Temporal pole as a robust biomarker: The consistent selection of temporal pole cortical thickness across diverse datasets and methods suggests it is a robust structural feature for AD, though it functions best in conjunction with other modalities.

The authors position their work as a step toward trustworthy clinical deployment, arguing that models must be transparent about missing information to be useful in real-world diagnostic workflows where data completeness varies.