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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Imputation-free transformer learning enables robust Alzheimer’s disease prediction and calibrated uncertainty quantification across heterogeneous clinical cohorts Christelle Schneuwly Diaz1,2∗, Narmina Baghirova1,2 , Duy-Thanh Vũ1,2 , Duy-Cat Can1,2 , Gilles Allali3 , Philippe Ryvlin4 , and Oliver Y. Chén1,2†

arXiv:2607.11656v1 [q-bio.NC] 13 Jul 2026

for the Alzheimer’s Disease Neuroimaging Initiative (ADNI) ‡, for the Australian Imaging, Biomarkers and Lifestyle (AIBL) study, and §. for the Open Access Series of Imaging Studies (OASIS) ¶ 1

Platform of Bioinformatics, Lausanne University Hospital. Faculty of Biology and Medicine, University of Lausanne. 3 Leenaards Memory Centre, Lausanne University Hospital. 4 Department of Clinical Neurosciences, Lausanne University Hospital. 2

∗ [email protected] † [email protected]

‡ Part of the data used in this article was from the Alzheimer’s Disease Neuroimaging Initiative

(ADNI) database (adni.loni.usc.edu). As such, the investigators within the ADNI contributed to the design and implementation of ADNI and/or provided data but did not participate in the analysis or writing of this paper. A complete list of ADNI investigators is available at: http: //adni.loni.usc.edu/wp-content/uploads/how to apply/ADNI Acknowledgement List.pdf. § Part of the data used in the preparation of this article was obtained from the Australian Imaging Biomarkers and Lifestyle flagship study of ageing (AIBL) funded by the Commonwealth Scientific and Industrial Research Organisation (CSIRO) which was made available at the ADNI database (www.loni.usc.edu/ADNI). The AIBL researchers contributed data but did not participate in analysis or writing of this report. AIBL researchers are listed at https://data.aibl.org.au/adni/ ¶ Part of the data used in this article was from the Open Access Series of Imaging Studies (OASIS) (https://sites.wustl.edu/oasisbrains/).

Abstract Accurate diagnostic classification and disease-severity prediction for Alzheimer’s disease are hampered by the pervasive 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 that are 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 the Alzheimer’s Disease Neuroimaging Initiative (ADNI; N=7858 scans), and evaluated it, without further retraining, on two independent cohorts: the Open Access Series of Imaging Studies (OASIS-3; N=2675 scans) and the Australian Imaging, Biomarkers and Lifestyle (AIBL; N=1286 scans) study. Across all cohorts and various tasks including binary and multi-class diagnostic classification as well as continuous cognitive score prediction, NITROGEN showed robust probability calibration and uncertainty quantification advantages over tree-based ensemble methods, while maintaining competitive discriminative and continuous cognitive score prediction performance. Moreover, 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 robust AD status classification. To address prediction reliability under incomplete data, we further introduced a modality-aware uncertainty adjustment that augments predictive uncertainty in proportion to the estimated importance of absent modalities, enabling models to express calibrated confidence when critical diagnostic information is unavailable. Taken together, our results show that imputation-free attention learning preserved meaningful discrimination under external cohort shift, while revealing the expected performance degradation on more distributionally different cohorts. These findings further demonstrate that evaluating machine learning models for neurological disease along dimensions of calibration, interpretability, and cross-cohort reliability—not predictive accuracy alone—is essential for better clinical deployment. Code to reproduce our experiments is publicly available at https://github.com/cschneuw/nitrogen. Keywords: Alzheimer’s disease, Machine learning, Transformers, Missing Data, Uncertainty Quantification, Interpretability, Explainable AI, Calibration

1 Introduction Alzheimer’s disease (AD) is a progressive neurodegenerative disorder characterised by a long preclinical phase, heterogeneous clinical presentation, and complex interactions between molecular pathology, brain structure, and cognitive decline [1]. Globally, ‖ NITROGEN: No-imputation Inter-sample-attention TRansformer Oriented for GENeral datasets.

c Absence-permissive attention via the NITROGEN

The reality of incomplete heterogeneous patient records Patient 1

Patient 2

Patient 4

Patient 3

MRI CSF

APOE

Demographics

Clinical reality handled by NITROGEN Heterogeneous patient data

Patient 6

MRI

CSF

Patient 5

Complete and incomplete data

MRI

Not perfectly-aligned information

CSF

Demographics

Various data can enter NITROGEN

Intersample attention

APOE Demographics

Demographics

The “synthetic evidence” problem

Conventional impute-then-predict

Absence-permissive learning

Imputed synthetic data

Predictor

CSF Synthetic data Cognition Demographics

CSF Masked tokens

Phantom evidence

Demographics

Overconfident

d Cross-cohort and -method validation

Masked feature attention

No imputation MRI

Probability of having AD: 75%

NAIM

MRI

NITROGEN

Probability of having AD: 75%

Confidence interval

Observed within-patient & between-patient context

Better calibrated

ADNI (test)

Missingness-aware confidence

Temporal-pole convergence

Case 1: Informative modalities available

ADNI

Case 2: Key modalities missing

Missing modality-weighted uncertainty

OASIS 3 AIBL NITROGEN

ADNI (training)

Consensus matrix

OASIS 3

All features Temporal pole and non-MRI features Temporal pole only

AIBL

Diagnostic classification & disease severity prediction

/ CN vs AD

CN / MCI / AD

MMSE / CDR

Bilateral temporal pole

MRI

CSF

APOE

Demo

MRI

CSF

APOE

Probability of having AD

0.75

Uncertainty Low

Uncertainty High

Low

High

Sufficiency

Reproducible and independent predictive signal of the temporal pole

Demo

Probability of having AD

Confidence: High

Selective prediction based on uncertainty filtering

Confidence: Low

Fig. 1: Overview of the NITROGEN framework for imputation-free multimodal learning from incomplete Alzheimer’s disease cohorts. (a) Alzheimer’s disease cohorts integrate complementary sources of information, including demographic, genetic, neuroimaging, cognitive, and molecular measurements, but these modalities are unevenly available across individuals and studies. This patient-level heterogeneity creates a major challenge for multimodal learning and cross-cohort generalisation. (b) Conventional impute-then-predict strategies replace missing measurements with synthetic values, potentially introducing artificial evidence, altering data structure, and affecting the reliability of model confidence estimates. (c) NITROGEN learns directly from partially observed patient records by combining masked feature-level attention within individuals with intersample attention across heterogeneous patients, allowing integration of available evidence without explicit imputation. (d) The NITROGEN framework is trained on ADNI and evaluated across held-out ADNI participants and independent OASIS-3 and AIBL cohorts, covering diagnostic classification, disease-state discrimination, and continuous cognitive score prediction. (e) Cross-cohort and cross-method attribution analyses identify temporal pole cortical thickness (right hemisphere: TempPole 1; left hemisphere: TempPole 3) as important MRI-derived features associated with Alzheimer’s disease classification, showing consistent importance rankings across datasets and models. (f ) Modality-aware uncertainty adjustment incorporates the importance of missing information into individual predictions, increasing uncertainty when highly informative modalities are unavailable and providing confidence estimates that reflect patient-specific evidence availability.

more than 55 million people live with dementia, with AD accounting for the majority of cases, a number projected to nearly triple by 2050 as populations age [2]. The personal cost is severe: patients lose independence, memory, and identity over years, while caregivers shoulder an enormous unpaid burden, estimated at approximately

82 billion hours of informal care annually provided to people with dementia living at home worldwide [3]. Health systems face mounting pressure, with dementia-related expenditure already exceeding $1313.4 billion per year worldwide and rising [4]. Despite this burden, therapeutic progress in AD has long been limited. However, the recent approval of several disease-modifying therapies marks an important transition from decades of largely unsuccessful drug development toward clinically available interventions [5]. These advances are expected to stimulate the development of additional therapeutic targets and treatment strategies. Nevertheless, the clinical benefits observed to date remain modest, a limitation that has been attributed in part to the late stage at which treatment is initiated, after substantial neurodegeneration has already occurred [6]. Consequently, improving disease detection, staging, and the prediction of disease severity may facilitate the identification of therapeutic windows during which interventions are most effective. In parallel to methods development, central to earlier and more reliable diagnosis are multimodal data. Accurate characterisation of AD progression increasingly relies on the integration of diverse biological measurements that capture the multifaceted and heterogeneous nature of the disease. Key variables include demographic information such as age, sex, and years of education, genetic risk factors such as APOE ε4 genotype, structural neuroimaging features derived from MRI, and molecular biomarkers of amyloid and tau pathology measurable via cerebrospinal fluid assays or positron emission tomography. Demographic and genetic factors shape baseline AD risk. Age remains an important non-genetic risk factor for late-onset AD, whereas APOE ε4 genotype substantially increases lifetime risk. AD is also more prevalent in women, although the basis of this difference remains debated and may involve both longevity and interactions between sex and APOE -associated risk [7]. Brain imaging data such as magnetic resonance imaging (MRI) data capture structural neurodegeneration; they provide a key observable marker of disease progression [6]. In particular, medial temporal lobe atrophy is among the earliest detectable changes, followed by involvement of parietal, frontal and cingulate cortices; the hippocampus, in particular, exhibits accelerated and nonlinear morphological decline [6, 8]. These structural alterations correlate with symptom severity and are widely used to track disease progression and treatment response. Although amyloid and tau pathology represent defining molecular hallmarks of AD, their direct quantification typically requires invasive or high-cost procedures, such as cerebrospinal fluid (CSF) sampling or positron emission tomography (PET) imaging. Consequently, disease characterisation often relies on heterogeneous and partially observed combinations of molecular, neuroimaging, genetic, demographic, and cognitive measurements. Because no single modality fully captures disease progression, integrating complementary sources of information has become increasingly important for studying AD heterogeneity and progression, particularly at early and prodromal stages [6]. In practice, however, multimodal datasets remain inherently incomplete due to differences in acquisition protocols, study design, patient compliance, and follow-up schedules, resulting in substantial variability in modality availability across individuals and cohorts [9].

Cognitive groups N = 1,286

b Available measurements

N = 2,675

N = 7,858

AD MCI

2.5

12.5

2.0

10.0

CDR

CDR-SB

MMSE

3.0

15.0

7.5

1.5

5.0

1.0

2.5

0.5

0.0

MCI

APOE4 genotype 1750

APOE4 allele count

1500 1250 1000 750 500 250 0

MCI

Longitudinal cortical thickness differences

e Mean Cortical Thickness [mm]

Cognitive scores 17.5

MRI

AIBL

Age and APOE

OASIS-3

MCI

CSF

ADNI MCI

Education

Control

Default 2.6

2.3

2.5

2.2

2.4 2.1 2.3 2.0

2.3

Dorsal Attention

F M

2.6

2.7

2.2 2.1

2.3

2.5

Somatomotor

Sal. Vent. Attention

2.2

2.6 2.4

2.1

2.3

2.0

Visual 2.1

2.0

1.9

2.4

2.0

2.3

2.2

Limbic 2.8

1.8

Female

*** p< 0.001

* p< 0.05

Male

** p< 0.01

n.s.

Fig. 2: Dataset characteristics and cortical thickness analyses. Panels (a–d) are colour-coded by cohort: Alzheimer’s Disease Neuroimaging Initiative (ADNI, red), Open Access Series of Imaging Studies-3 (OASIS-3, blue), and Australian Imaging, Biomarkers and Lifestyle Flagship Study of Ageing (AIBL, purple). (a) Distribution of cognitively normal (CN), mild cognitive impairment (MCI), and Alzheimer’s disease (AD). (b) Availability of multimodal data across cohorts (60 participants per cell; gray: missing). (c) Box-plots of cognitive scores: Mini-Mental State Examination (MMSE), Clinical Dementia Rating (CDR), and CDR Sum of Boxes (CDR-SB). (d) Distribution of APOE ε4 allele counts across cohorts. (e) Mean cortical thickness trajectories as a function of age for CN (blue) and AD (red) participants, by sex and functional brain network. Adjacent age bins were compared using Welch’s t -tests with Benjamini–Hochberg false discovery rate correction. Shaded bands indicate 95% confidence intervals; horizontal strips denote significant comparisons.

This is not a peripheral inconvenience; rather, it is the crucial reality of AD research. The use of multimodal data for AD prediction, therefore, necessitates addressing the challenge of missing data. Traditionally, missing data in clinical machine learning has been predominately handled through explicit data imputation prior to model training, a strategy commonly referred to as impute-then-classify for classification tasks and, broadly, impute-then-predict. While this approach is compatible with standard algorithms, imputation inevitably introduces synthetic measurements that can distort feature distributions, disrupt inter-feature dependencies, and compromise uncertainty estimation [10]. In diagnostic applications, these distortions are not merely technical inconveniences: overconfident or biased predictions can directly influence clinical decision-making, with potential downstream consequences for patients, including misdiagnosis or inappropriate treatment recommendations. To mitigate these shortcomings, a growing body of work has begun to explore missing-aware models that operate directly on partially observed data, including modified decision trees and ensemble methods with missingness-aware splitting criteria [11], linear models augmented with explicit missingness indicator variables [12], and neural network architectures that either incorporate masking mechanisms or explicitly modify the forward pass to exclude missing inputs contributions during computation [13, 14]. Broadly, these approaches fall into two categories: those that augment the input space with missingness indicators, and those that directly adapt model computation to exclude or down-weight missing contributions. Recently, attention-based deep learning models, particularly transformer architectures, have demonstrated considerable flexibility in handling tabular data, extending well beyond natural language processing [15, 16]. Leveraging self-attention and masked attention mechanisms, these models have been adapted to generate predictions directly from tabular datasets with missing values [17], removing the need for explicit imputation at the input stage. However, eliminating imputation addresses only how missing data are handled during representation learning; it does not by itself guarantee that the resulting predictions are reliable. Two further challenges therefore remain unresolved. The first concerns the gap between discrimination and calibration. Most existing approaches are optimised and evaluated primarily on discriminative accuracy, while the reliability of individual-level probability estimates, a critical requirement for clinical decision support, is rarely assessed. Although the effect of imputation on uncertainty has been examined relative to complete-case analyses [10], this question remains essentially unexplored in the harder setting when entire modalities, rather than scattered values, are missing. The second concerns the lack of modality-aware uncertainty. Existing uncertainty estimation methods typically treat all missing information as interchangeable, without accounting for how much a given absent modality would have contributed to the prediction. Here, we address these challenges within a unified framework called NITROGEN (No-imputation Inter-sample-attention TRansformer Oriented for GENeral datasets), which integrates Not ANother Imputation Method (NAIM) [17], a imputation-free method introducing masked self-attention, and SAINT, an inter-sample-attention transformer [18]. NITROGEN is designed to handle incomplete multimodal clinical data by jointly capturing within-patient feature dependencies and inter-patient

relationships, while explicitly accommodating partially missing observations through masked attention (Fig. 1 and 6). Unlike imputation-based pipelines, NITROGEN does not require missing values to be filled before training or inference. Instead, masked attention explicitly limits the contribution of missing observations during representation learning. At the same time, intersample attention enables the model to incorporate relational information across patients, which helps stabilise predictions in heterogeneous clinical cohorts with modality-level data loss. To evaluate the efficacy of the model, we first applied it to three predictive tasks: a binary disease classification, multi-class disease classification and continuous cognitive score prediction. We then assessed generalisation and the reproducibility of the framework via cross-cohort validation as well as cross-cohort and cross-method analyses across three independent datasets: the ADNI dataset (Nsubjects = 1, 776), the OASIS-3 data (Nsubjects = 1, 078) and AIBL (Nsubjects = 1, 286), first applying the model trained on a part of the ADNI data on the held-out ADNI data, and subsequently to the latter two cohorts without further training or fine-tuning. Furthermore, we identified brain regions consistently selected across all three datasets, as well as those robustly implicated across all analytical methods and three cohorts. To address the reliability of individual-level probability estimates, we further examine the robustness of individual-level confidence estimates under incomplete multimodal inputs. Specifically, we proposed an uncertainty adjustment, through which global estimates of modality importance are combined with subject-specific missingness patterns to increase uncertainty when informative modalities are absent. We designed the uncertainty adjustment because, intuitively, missing MRI, CSF, demographic, genetic, or cognitive information should not affect confidence equally: the absence of a modality that contributes strongly to prediction should lead to greater uncertainty regarding the final prediction than the absence of a less informative modality. Practically, this adjustment accounts not only for how much information is missing, but also for which type of information is missing. Collectively, by learning directly from incomplete, heterogeneous records rather than imputed ones, NITROGEN performs classification tasks and disease severity estimation across independently acquired, diverse cohorts while remaining calibrated and transparent about which information a given prediction is missing. This demonstrates that AD prediction models can remain trustworthy under missing values, missing modalities, and cohort heterogeneity—conditions under which conventional approaches typically degrade.

2 Results In this section, we first present the results regarding three tasks: binary disease status classification, multi-class disease classification, and continuous cognitive score prediction. We then demonstrate the reliability and reproducibility of NITROGEN and derived neuroimaging features in the binary classification task as well as the impact of uncertainty adjustment techniques. To do that, we perform cross-cohort analysis using three independent datasets from the Alzheimer’s Disease Neuroimaging Initiative (ADNI), the Open Access Series

of Imaging Studies (OASIS-3) and the Australian Imaging, Biomarkers and Lifestyle (AIBL) study. We report predictive performance across binary classification, multiclass diagnostic classification, and continuous cognitive score regression tasks. To understand the potential drivers of model predictions, we perform permutationbased feature attribution analysis, identifying the top 30 most influential MRI regions on each dataset independently. We then examine the consistency of these rankings across all seven models and all three cohorts, retaining only regions that appear among the top attributed features across all methods and datasets. The cross-model and cross-cohort analysis identifies temporal pole cortical thickness as the most consistently attributed MRI feature, suggesting robustness of this finding to modelling choices and cohort characteristics. Furthermore, to assess whether this feature retains predictive value in isolation, we conduct two complementary sufficiency analyses in which feature attributions are recomputed on a held-out data partition, distinct from that used for the original attribution analysis, to prevent data leakage between attribution computation and sufficiency evaluation. Non-selected features are then progressively replaced with missing values, and model performance is evaluated on the remaining held-out subjects across all three cohorts. This allows us to quantify the extent to which the identified temporal pole feature, alone or in combination with a minimal set of non-imaging variables, supports diagnostic classification, independently of the remaining feature set. Finally, we evaluate the proposed uncertainty adjustment by measuring the improvement in predictive accuracy obtained after excluding predictions with high estimated uncertainty, together with the corresponding reduction in the number of evaluated samples on the validation set.

2.1 Dataset Characteristics and Cortical Thickness Analyses We first characterised the three cohorts included in this study. Figure 2 summarises cohort composition (Fig. 2a), modality-specific missingness patterns (Fig. 2b), cognitive score distributions (Fig. 2c), and APOE ε4 allele frequencies (Fig. 2d), revealing substantial cohort heterogeneity across data modalities and datasets. Because cortical thickness varies with age and sex across brain regions [19], we examined how these effects differ between diagnostic groups (cognitively normal (CN) vs. AD) and between sexes across age strata using Welch t-tests (Fig. 2e). We found that AD participants exhibit significantly reduced cortical thickness relative to CN individuals across all functional brain networks and age groups, with this difference already present in the youngest stratum examined (55–64 years), suggesting that structural changes are present at the time of clinical diagnosis across the full age range studied. These findings support the use of MRI-derived cortical thickness as a discriminative feature for classification and regression tasks.

2.2 Binary Disease Group Classification We evaluated binary diagnostic classification in two settings. First, we considered a classification task distinguishing AD from CN subjects. Second, we considered a

broader impairment-detection task distinguishing CN subjects from subjects with either mild cognitive impairment (MCI) or AD. For each task, the NITROGEN model was trained and evaluated on the ADNI cohort using a 60%/20%/20% train-validation-test split. To assess generalisation under cohort shift, the ADNI-trained model was then applied, without further training or fine-tuning, to the independent OASIS-3 and AIBL cohorts. The same evaluation procedure was repeated for baseline models, including the original NAIM model, XGBoost, LightGBM, Random Forest, MA-Lasso, and MA-GBT. Classification metrics and model comparison results are reported in Table A2. Our results reveal several insights. For the CN vs. AD classification task, all models achieved strong discriminative performance on the ADNI test set, with AUC values ranging from 0.913 to 0.935 (Table A2, Fig. 3a). Performance transferred well to AIBL data (AUC: 0.901–0.918), likely reflecting the similar recruitment strategies and clinical protocols shared between ADNI and AIBL. Performance degraded on OASIS-3 data (AUC: 0.689–0.719), consistent with the marked distributional shift between these cohorts (Fig. 2), as OASIS-3 specifically targets preclinical populations resulting in a higher proportion of CN participants and fewer AD cases. Across cohorts, attention-based models (NITROGEN, NAIM) and MA-Lasso consistently achieved higher sensitivity than tree-based ensemble methods, despite the latter achieving the highest AUC. On ADNI test data and AIBL data, MA-Lasso achieved the highest sensitivity among all models (0.833 and 0.761 respectively), followed closely by NITROGEN (0.832 and 0.697) and NAIM (0.808 and 0.695), while tree-based models showed comparatively reduced sensitivity, particularly on AIBL data (XGBoost: 0.610; LightGBM: 0.645). On OASIS-3 data, this pattern partially reversed, with XGBoost (0.770) and NAIM (0.763) achieving the highest sensitivity, followed closely by NITROGEN (0.751), while MA-Lasso (0.650) and Random Forest (0.636) performed comparatively worse. This pattern suggests that, with the exception of OASIS-3 data, tree-based ensemble methods are more conservative in predicting AD status, favouring specificity over sensitivity, whereas attention-based architectures maintain a more balanced operating point without explicit threshold tuning. Notably, Random Forest exhibited markedly elevated variance on OASIS-3 data across all metrics (accuracy standard deviation: 0.168; specificity standard deviation: 0.222), suggesting that its point estimates on this cohort are highly seed-dependent and should be interpreted with caution. Second, we evaluated all models on a more challenging binary classification task distinguishing cognitively normal participants from those with any level of cognitive impairment (MCI or AD). This setting reflects a relevant screening scenario but is inherently more difficult than CN vs. AD classification due to the heterogeneous and transitional nature of MCI, which exhibits overlapping characteristics with both CN and AD along the disease continuum [20, 21]. Our results show that it is possible to generally distinguish cognitively normal participants from those with any level of cognitive impairment. As expected, performance decreased across all models and cohorts relative to the CN vs. AD task, with ADNI test data AUC ranging from 0.748 to 0.802 and OASIS-3 data AUC ranging from 0.639 to 0.700 (Table A3). This decline likely reflects both cross-cohort differences in acquisition protocols and the increased

Binary classification

a 1.0

1.0

False positive rate

True positive rate

1.0

False positive rate

1.0

False positive rate

1.0

Correct and incorrect prediction probabilities Correct predictions

Incorrect predictions 1.0

1.0

Discrimination vs. calibration

NITROGEN

0.4

Random Forest

MA-LASSO

0.1

Expected calibration error (ECE)

True Positive Rate

1.0

Mean predicted probability

0.00

Calibration

Difference between correct and incorrect predictions

Difference

Prediction probability

0.5

0.12

Prediction Probability

1.0

0.70

0.90

ADNI (val)

NAIM

ADNI (test)

XGBoost

AIBL

LightGBM

OASIS-3

MA-GBT

ROC AUC (discrimination)

Fig. 3: Classification performance, calibration, and discriminationcalibration trade-off across models and cohorts. (a) ROC curves for binary CN vs. AD classification on ADNI test set, AIBL, and OASIS-3. All models trained on ADNI. Curves represent the mean curve and shaded area the standard deviation across 10 random seeds (seeds = [123, 42, 0, 7, 13, 21, 37, 55, 77, 99]). (b) Prediction efficacy analysis on ADNI validation set (seed=123). Box-plots of predicted probabilities and 100 randomly drawn samples for correctly classified, incorrectly classified, and mean prediction probability gap bar plot. Prediction probability gap is the difference between mean probabilities for correct versus incorrect predictions. (c) Reliability diagrams showing model calibration on ADNI validation set (seed=123). Mean predicted probability per bin (x-axis) versus observed fraction of AD cases (y axis). Dashed diagonal indicates perfect calibration. (d) Discrimination-calibration trade-off. x-axis: ROC-AUC; y -axis: Expected Calibration Error (ECE; lower is better). Evaluated on ADNI validation set, ADNI test set (seed=123), AIBL data, and OASIS-3 data . Lower-right region: strong discrimination and good calibration. 10

diagnostic ambiguity introduced by MCI, which represents an intermediate stage of cognitive decline with variable clinical presentation. Third, beyond discrimination, we assessed model calibration, meaning the alignment between predicted class probabilities and observed outcome frequencies, as calibration reveals differences in model reliability not captured by AUC alone (Fig. 3b, c, d). Among tree-based models, LightGBM, MA-GBT and XGBoost achieved the best competitive discriminative performance but exhibited some overconfidence as shown by the poor alignment with identity line on the calibration plot (Fig. 3c) and on OASIS-3 and ADNI test sets with slightly higher Expected Calibration Error (ECE) values (Fig. 3d). In contrast, NAIM and NITROGEN produced better-calibrated probabilities across datasets, with NITROGEN achieving the largest prediction probability gap between correct and incorrect predictions on the validation set, followed by Random Forest, indicating superior uncertainty awareness (Fig. 3d). Most models maintained reasonable calibration on the ADNI and AIBL sets; however, calibration degraded on OASIS-3 for all models. These results suggest that whilst NITROGEN’s sensitivity limitations at the default decision threshold merit consideration, its superior calibration, argue for its utility as a candidate for reliable probabilistic prediction in heterogeneous clinical settings. To further improve uncertainty estimates in settings with incomplete data, in Section 2.8, we explored entropy-based uncertainty quantification adjusted for modality-level missingness.

2.3 Multi-class Diagnostic Classification Following the binary classification analyses, we evaluated model performance in a three-class diagnostic task distinguishing CN, MCI, and AD participants. This setting provides a relevant yet more challenging problem, as models must simultaneously differentiate normal cognition, mild cognitive impairment, and AD. We note that these categories correspond to clinician-assigned diagnoses rather than biologically defined disease stages. Global and per-class performance metrics are reported in Tables A4 and A5. Per-class analysis revealed interpretable heterogeneity in class-wise performance (Table A5). Across all models and cohorts, AD discrimination was consistently the strongest task: per-class AUC on ADNI test ranged from 0.830–0.863, and NITROGEN achieved the highest AD AUC of any model on both ADNI test (0.863) and the fully held-out AIBL cohort (0.907), indicating that the captured signal associated with established disease generalises robustly even under domain shift. CN discrimination was likewise stable (0.744–0.799 on ADNI test), reflecting a consistent separation between cognitively normal individuals and clinical syndromes. MCI was noticeably harder to discriminate (AUC 0.547–0.654 on ADNI test, narrowing further on OASIS-3 to 0.361–0.441 across models), a pattern consistent across every method we evaluated, including the tree-based ensembles. This is in line with MCI’s established clinical status as an intermediate, heterogeneous category rather than a discrete diagnostic entity, and the reduced MCI sample size in OASIS-3 likely compounds this further. The convergence of all five models on this same difficulty pattern suggests it reflects a property of the diagnostic task itself rather than a limitation specific to any one architecture. 11

At the global level (Table A4), tree-based ensembles achieved the highest macro AUC on ADNI test (0.762–0.771), with NITROGEN at 0.752 and NAIM at 0.707. Under genuine domain shift to AIBL, NAIM improved slightly (0.707 to 0.729), while NITROGEN’s macro AUC dropped modestly (0.752 to 0.719), smaller than XGBoost’s decline (0.046). On OASIS-3, LightGBM and NAIM performed strongest (0.594 and 0.596), while NITROGEN and XGBoost ranked lower (0.583 and 0.574). These results demonstrate that in-distribution accuracy does not predict cross-cohort robustness. Taken together, these results indicate that all models capture AD-related signal, and that NITROGEN in particular combines competitive in-distribution discrimination with stable transfer to external cohorts, while MCI detection remains the principal opportunity for improvement across the field, consistent with its recognised status as the most clinically ambiguous diagnostic category in this three-class setting.

2.4 Cognitive Score Prediction Next, we sought to predict continuous cognitive scores obtained during neuropsychological assessments that measure a more fine-grained characterisation of cognitive decline. Here, we consider the prediction of the Mini Mental State Examination (MMSE) scores and Clinical Dementia Rating (CDR) score (or their Sum of Boxes (CDR-SB) version). We aim to predict these two scores simultaneously in an end-toend fashion using a single model. Full results are reported in Supplementary Table A6 and NITROGEN’s regression results are plotted in Figure 4b. On the ADNI test set, all models showed moderate-to-strong Pearson correlation with clinician-rated MMSE scores, with NAIM achieving the strongest correlation (r = 0.627, p = 1.58 × 10−174 ), closely followed by NITROGEN (r = 0.621, p = 1.58 × 10−170 ) and Random Forest (r = 0.610, p = 1.58 × 10−163 ). On the fully external AIBL cohort, correlations for MMSE prediction improved for every model relative to ADNI test, with Random Forest and NITROGEN performing best (r = 0.666 and 0.664, respectively), indicating that the ordinal relationship between predicted and true MMSE scores generalise to this cohort. Performance on OASIS-3 was weaker for all models (r between 0.245 and 0.303), consistent with the markedly different case mix in this cohort (only 49 MCI and 398 AD cases, versus 184 and 155 in AIBL) rather than a cohort-specific failure of any one method. For CDR-SB prediction, NITROGEN and NAIM were again the best-performing models on the ADNI test set (r = 0.641 and 0.635, respectively), followed by Random Forest (r = 0.633) and XGBoost (r = 0.604). Due to the unavailability of CDR-SB in AIBL and OASIS-3, performance was assessed against the global CDR score as a surrogate target. R2 values for this surrogate comparison were negative for all models in both external cohorts, reflecting the low variance of the global CDR scale in samples that are predominantly cognitively normal; we therefore focus on Pearson correlation, which remained informative, as the primary cross-cohort metric. On AIBL, Random Forest achieved the highest correlation (r = 0.656), followed closely by NITROGEN (r = 0.603). On OASIS-3, correlations were more attenuated but remained significant for all models (r = 0.247–0.284), likely reflecting the coarse, low-variance nature of the global CDR scale in a cohort with very few non-CN cases.

Notably, NITROGEN’s correlations on the ADNI test set and the two fully external cohorts were consistently close to one another for both MMSE and CDR-SB (e.g., r = 0.621 on ADNI test versus r = 0.664 on AIBL for MMSE), in contrast to the tree-based ensembles, whose strong in-distribution fit did not always translate into a comparable advantage once evaluated externally. This pattern echoes the classification results above and suggests that NITROGEN’s representation of disease severity, while not always the single best performer on any one held-out set, transfers consistently across cohorts. Taken together, these results indicate that agreement between predicted and observed clinical severity scores, as measured by Pearson correlation, was largely preserved across cohorts despite the pronounced scale and case-mix differences in OASIS-3 and AIBL. NITROGEN achieved correlations that were competitive with those of the other methods and, on AIBL, ranked among the best-performing models. Moreover, its performance changed comparatively little between in-distribution and external evaluation, suggesting greater robustness to cohort differences that may be beneficial for deployment in populations that differ from the training data.

2.5 Model Explanation via Feature Attribution Analysis To better understand the potential biological underpinnings of the model predictions, we analysed feature attributions and modality contributions using permutation importance. We focus on feature attributions in the binary CN versus AD classification setting, as this contrast provided the clearest discrimination across cohorts. We quantify feature importance by measuring the drop in model performance induced by randomly permuting input feature values. The top 30 features with the largest impact on NITROGEN predictions on the ADNI test set are shown in Figure 4b, with APOEε4 identified as the most important feature, followed by MRI-derived cortical thickness features from bilateral temporal poles, other temporal and limbic regions (parahippocampal cortex, default mode network areas), and demographic variables including education and age. When the same permutation-based attribution analysis was applied independently to held-out test sets from OASIS-3 and AIBL, the top-ranked features showed consistencies despite variations. Features that appeared in the top 30 across all three cohorts for NITROGEN are displayed in Figure 4c and include: bilateral temporal poles (Left (L)/Right (R) Limbic Temporal Pole 1/3 and L Temporal Pole 2), parahippocampal cortex (L Default ParaHipp 1), additional temporal and default mode regions (L/R Default Temporal 1/2), somatomotor cortex (L SomMot 4), visual cortex (L Vis 12), a control network temporal region (L Cont Temp 1), APOEε4, and age. The prominence of temporal cortical regions is consistent with the well-established spatial progression of AD pathology. Amyloid-β deposition begins years to decades before symptom onset and preferentially accumulates within regions of the default mode network, including lateral temporal, posterior cingulate, and medial prefrontal cortices [21]. In contrast, tau pathology is thought to originate within medial temporal structures and subsequently propagate to temporal and parietal association cortices before involving the broader neocortex, a process that more closely parallels the emergence of cognitive symptoms [8, 21]. 13

Predicting cognitive scores using NITROGEN

MCI

Observed MMSE

Identity line

Least-squares fit

Observed CDR-SB

d Shared MRI features across datasets and models

r = 0.26 p<0.001

r = 0.60 0 p<0.001 0.0 0.5 1.0

2.0

Observed CDR

3.0

0.0

0.5

1.0

2.0

Shared ROIs

Observed CDR

Top 30 features with highest permutation importances

DNA

MRI

Binary classification using all features 1.0

True Positive Rate

0.10

Demographics

0.05

MRI (temporal pole)

Union of ROIs

Binary classification using only features from the temporal pole 1.0

False Positive Rate

0.00

Mean Absolute Permutation Importance

Union of ROIs

Shared ROIs

OASIS-3

AIBL 6

r = 0.64 p<0.001 0

Observed MMSE

ADNI (Test )

1.0

OASIS-3 AIBL

Observed MMSE

r = 0.25 p<0.001

ADNI (Test)

OASIS-3

r = 0.66 p<0.001

r = 0.62 p<0.001 0

Predicted CDR-SB

AIBL

True Positive Rate

OASIS-3 AIBL

ADNI (Test )

Predicted MMSE

c Shared MRI features across datasets using NITROGEN

ADNI (Test)

False Positive Rate

1.0

all features temporal pole cortical thickness and non-MRI data only temporal pole cortical thickness

Fig. 4: NITROGEN performance and feature attribution analysis. (a) Predicted vs. clinician-rated cognitive scores on ADNI, AIBL, and OASIS-3. Top: MMSE. Bottom: CDR-SB (ADNI test set) or CDR (OASIS-3 and AIBL data). Pearson r (R) and p-values (p) shown. (b) Top 30 features by mean absolute permutation importance on ADNI test set. (c) Top 30 MRI ROI features across datasets (ADNI test, AIBL and OASIS-3). (d) Consensus MRI ROIs identified in bilateral temporal pole across cohorts and models. (e) Classification accuracy on ADNI test set, AIBL and OASIS-3 data. (f ) Test of sufficiency using data from temporal pole only. ROC curves compare all features, consensus features from temporal pole plus non-imaging features, and features from temporal pole only.

Our results extend these observations by showing that temporal and limbic features were consistently prioritised across independent cohorts and diverse machine learning models, suggesting that their predictive value is not specific to a single dataset or modelling framework. Although their relative rankings varied across cohorts, likely reflecting differences in participant characteristics and acquisition protocols, the crosscohort and cross-method analyses (Section 2.6) identified a subset of features that remained robustly selected despite these sources of variability.

2.6 Consistently Attributed MRI Regions Across Cohorts and Methods To assess whether feature attributions yield consistent patterns despite cohort heterogeneity and methodological differences, we computed the intersection of top 30 features across all seven models and all three cohorts. This analysis identified a minimal core set of four features consistently selected across all 21 model–cohort combinations: bilateral temporal pole cortical thickness (L Limbic Temporal Pole 1, R Limbic Temporal Pole 3), APOEε4, and age (Fig. 4e). The persistence of such a small set across all combinations suggests these features consistently carry important signal for AD status, with limited sensitivity to methodological choice or to distributional differences between cohorts. The association between temporal pole volume loss and cognitive decline is relatively well established [8], as is its correlation with plasma tau in AD/MCI participants in ADNI [22]. That the temporal poles carry consistent signal despite differences across datasets and methodological frameworks suggests that they may serve as consistently informative structural features in the context of AD.

2.7 Sufficiency Analysis of Consistently Attributed Regions To evaluate whether the consistently attributed temporal pole region retained predictive value in isolation, we performed a sufficiency analysis using a two-stage evaluation protocol. Note that the feature selection performed here differs from the previous section to further guard against data leakage between attribution computation and sufficiency evaluation, attributions were recomputed on a restricted, non-overlapping data partition and extended to the full multimodal feature space rather than cortical-thickness features alone, yielding a distinct, stricter feature set. First, feature attributions were re-calculated with 70% of the held-out test sets (ADNI test, OASIS-3, and AIBL). Second, performance was evaluated on the remaining 30% of each test set using previously trained models, thereby preventing data leakage. Applying this stricter, multimodal feature selection identified two robustly attributed features: APOEε4 and left limbic temporal pole cortical thickness (L Limbic Temporal Pole 3). We evaluated performance under two feature-restricted scenarios: (i) MRI, where the identified temporal pole MRI feature was retained but all other MRI features were replaced with missing values, while demographics, CSF biomarkers, and genotyping data remained available, and (ii) feature-masked (ALL), where only the temporal pole measure was retained and all other features were masked. Results are summarised in Table A7. Retaining the temporal pole measure together with non-imaging variables (MRI setting) resulted in respectable discriminative performance. Across models, AUC values ranged from 0.76–0.80 on the ADNI test set (versus 0.90–0.92 under the full feature configuration), 0.59–0.69 on OASIS-3, and 0.57–0.81 on AIBL, with several models (e.g., NAIM, MA-Lasso) approaching or exceeding 0.80 on AIBL. That a single cortical region, combined only with routinely collected demographic and genetic variables, recovers the large majority of full-model discriminative performance underscores the value of the temporal pole as a parsimonious structural biomarker for AD.

Further, we investigated whether it is possible to classify patients using only the temporal pole feature (Table A7). Although AUC values suggested ranking-level discrimination (0.76—0.77 on ADNI test, 0.83—0.87 on AIBL), this did not translate into clinically usable classification . When required to make binary decisions, most models defaulted to constant predictions. NAIM Default, NITROGEN, and Random Forest converged to constant predictions of class AD, while LightGBM and MA-GBT converged to constant CN predictions. XGBoost produced random balanced accuracy (0.53 on ADNI test, 0.40 on AIBL, 0.28 on OASIS-3) without exhibiting constant predictions. Random Forest was the clear exception to the AUC pattern above, with AUC collapsing to 0.41-–0.57, consistent with its tree-splitting behaviour degenerating under a near-constant single-feature input. MA-Lasso was unique: it produced probability distributions tightly clustered around 0.50 (reflecting model uncertainty), yet still achieved non-random accuracy, most notably a balanced accuracy of 0.81 on AIBL. However, this performance must be contextualised: AIBL exhibits severe class imbalance (87.6% CN, 12.4% AD), and inspection of MA-Lasso’s predictions reveals that its 0.81 balanced accuracy reflects assignment of different labels despite genuine probability uncertainty (probs 0.50), rather than confidence-driven discrimination. Constant-prediction models on AIBL appear deceptively high in AUC (0.80+) due to class imbalance, but achieve balanced accuracy of 0.50 and are unusable. In other words, these results suggest that most classifiers, lacking complementary features to anchor a calibrated decision boundary, default to a single constant label at the standard 0.5 threshold rather than exploiting the genuine ranking ability the temporal pole feature provides on its own—underscoring that AUC and balanced accuracy must be interpreted jointly when characterising single-feature sufficiency. Taken together, these analyses indicate a feature-richness gradient with two distinct findings. First, temporal pole cortical thickness alone seems to carry a reproducible ranking signal for AD status. Second, this signal is not, by itself, sufficient to support stable decision-level classification, which requires the additional context provided by demographics, genotype, and complementary MRI regions. Comprehensive multimodal assessment remains preferable when all modalities are available, providing the highest diagnostic confidence. Nevertheless, the robustness of the temporal pole signal and its contribution alongside minimal non-imaging information indicate that relevant predictions can be maintained under reduced modality availability.

2.8 Uncertainty adjustment Although softmax-derived probabilities provide a default measure of confidence, they conflate two distinct sources of uncertainty: uncertainty arising from the inherent difficulty of a case, and uncertainty arising from missing diagnostic information. A model evaluated on a patient with all modalities present and a model evaluated on a patient missing a highly informative modality (e.g., CSF biomarkers) may report similar softmax confidence despite having access to substantially different amounts of evidence. This is undesirable in a clinical setting, where the reason for low confidence— case difficulty versus missing data—has direct implications for what action should follow (e.g., obtaining the missing test versus seeking a specialist opinion). Motivated by this gap, we sought to adjust predictive uncertainty explicitly as a function of which 16

Selective prediction by uncertainty threshold

Confusion matrices across selection regimes

Full n = 882

Base

Selection threshold (τ = 0.900)

Base selection 88% retained

Adjusted selection 86% retained

Adjusted

2.0

1.0

471

289

435

256

429

251

Observed

Density

1.5

0.5

0.0 0.6

0.8

1.0

Full

Base selection (88%) Base selection (88%)

Predicted

Classification performance by selection regime

100

Uncertainty

Predicted

Permutation test: improvement not due to chance

Adjusted selection (86%) Adjusted selection (86%)

Score (%)

0 70

Accuracy Accuracy

Sensitivity Sensitivity

Specificity

Precision Precision

−0.02

Predicted

0.00 Mean of random improvement = −0.0002

0.02

p~0

0.4

Observed improvement (adj.) = 0.0378

0.2

Observed improvement (base) = 0.0311

0.0

0.04

Accuracy improvement

Fig. 5: Missingness-aware uncertainty adjustment improves selective prediction reliability. (a) Distribution of base (blue) and missingness-adjusted (red) uncertainty scores on the ADNI validation set (N = 882). The selection threshold (τ = 0.900) is shown as a solid vertical line; samples with uncertainty above this threshold are excluded from selective prediction. (b) Confusion matrices for CN vs. AD classification under the full validation set (Full, N = 882), the subset retained under base-uncertainty thresholding (Base selection, 88% retained), and the subset retained under missingness-adjusted uncertainty thresholding (Adjusted selection, 86% retained). (c) Classification performance (accuracy, sensitivity, specificity, precision, and F1) for the full validation set and for the base- and adjusted-selection subsets. (d) Permutation test assessing whether the accuracy improvement achieved through adjusted-uncertainty selection exceeds chance. Null distribution (gray bars, 1,000 permutations) of accuracy improvement under random sample exclusion at matched coverage; dashed line indicates the mean of the null distribution, dash-dotted line indicates the observed improvement under base-uncertainty selection (+0.0311), and solid line indicates the observed improvement under adjusted-uncertainty selection (+0.0378, p < 0.001).

modalities are absent for a given subject, rather than relying on raw model confidence alone. In clinical decision-support systems, reliable estimates of prediction trustworthiness are as important as predictive accuracy, particularly in multimodal settings where data completeness varies across subjects. Shannon entropy computed from predicted class probabilities provides a simple sample-level measure of uncertainty. When critical diagnostic modalities are unavailable, a model should express higher uncertainty even if its raw output probabilities appear confident. As such, we inflated these sample entropy values based on the amount of important missing features and an inflation factor β . Inflated uncertainty values can then be thresholded; the logic is that if uncertainty is too high, it is better to withhold a prediction until more assessments

are made. We fix the uncertainty threshold at τ =0.900 for selective prediction (reject option). The scaling factor β is selected by maximising the Spearman rank correlation between adjusted uncertainty and prediction errors, weighted by the retention gap (the difference in accuracy between retained and rejected samples). Retained samples are those whose adjusted uncertainty falls below τ and are therefore used for prediction, while rejected samples exceed the threshold and are abstained from (i.e., prediction is not trustworthy). We present the results in Fig. 5. On the validation set (N = 882), base Shannon entropy averaged 0.428 ± 0.309. After adjustment (β = 0.400), entropy increased to 0.446 ± 0.323. Without selective abstention, predictions on the full dataset achieved accuracy 0.862 and F1 0.826. Applying uncertainty thresholding at τ = 0.900, the adjusted entropy outperformed raw entropy-based thresholding: adjusted uncertainty retained 85.7% of samples (756/882) with accuracy 0.900 and F1 0.869, compared to raw entropy which retained 87.8% (774/882) with accuracy 0.893 and F1 0.861. Beyond this comparison, the adjusted uncertainty scores correlated significantly with prediction errors (Spearman ρ = 0.348, p < 10−26 ), and rejected samples showed markedly lower accuracy (0.635) than retained samples, indicating that the adjustment has identified a genuinely less reliable subset of predictions rather than discarding samples arbitrarily. The missingness-aware adjustment offers measurable, if modest, improvement over raw entropy filtering on the validation set (+0.007 in accuracy) at the cost of a 2.0% reduction in coverage. This increment is small in absolute terms but consistent in direction across the metrics examined, and it is achieved on top of a selective-prediction mechanism that is already effective at isolating unreliable predictions, as shown by the retained-versus-rejected accuracy gap above. The marginal benefit is expected given the limited missingness patterns in ADNI, where only CSF biomarkers are systematically absent (70.1% missing) while other modalities remain near-complete. The adjustment may become more impactful in cohorts with heterogeneous or multimodality missingness patterns, where importance-weighted scaling can more effectively capture the relative contribution of unavailable features to diagnostic uncertainty.

3 Discussion In this paper, we presented NITROGEN, an attention-based architecture, for the joint tasks of performing AD diagnosis and cognitive score prediction from incomplete multimodal clinical data. To evaluate its generalizability and reproducibility, we applied NITROGEN across data from three independent cohorts with heterogeneous characteristics. Computationally, our results demonstrated that NITROGEN conferred meaningful advantages in reliability and cross-cohort robustness over standard tree ensemble methods while achieving competitive classification and prediction performance. Additionally, our results highlighted important differences in model reliability that are not captured by AUC alone. Tree-based ensemble methods, particularly LightGBM, achieved the highest discriminative performance on ADNI but exhibited systematic overconfidence and poor probability calibration, most pronounced under domain shift,

a setting where inflated confidence on incorrect predictions is particularly consequential for clinical decision-making. In contrast, NITROGEN maintained well-calibrated predictions across cohorts, as reflected in the lowest ECE on OASIS-3 and a larger confidence gap between correct and incorrect predictions, indicating meaningful uncertainty awareness rather than uniform overconfidence. These findings reinforce the argument that clinical machine learning benchmarks should systematically evaluate models along multiple dimensions, discrimination, calibration, and cross-cohort generalisation, particularly for decision support applications where reliability is no less critical than predictive accuracy [9]. Recent advances in artificial intelligence have substantially improved the performance of automated AD classification using MRI data. Specialised three-dimensional convolutional architectures trained on large imaging cohorts have reported near-ceiling performance for binary CN versus AD discrimination under controlled experimental settings, with accuracies approaching 99% in some studies [23]. More recently, MRI foundation models pre-trained on large collections of brain scans have demonstrated improved data efficiency and transferability, achieving strong performance even with limited downstream training data [24]. These advances highlight the substantial diagnostic information contained within structural MRI alone and demonstrate the capabilities of modern imaging-based deep learning models. However, such approaches address a different problem setting from the one considered here. MRI foundation models and specialised three-dimensional convolutional architectures typically operate on a single, highly informative modality and assume the availability of complete imaging data. By contrast, our objective was to model heterogeneous clinical datasets composed of demographic, genetic, biomarker, and neuroimaging variables that are frequently incomplete and unevenly sampled across individuals. Although this tabular multimodal setting may not achieve the near-perfect classification performance reported for dedicated MRI models, it addresses a broader class of prediction problems encountered in clinical research, where multiple data sources must be integrated despite substantial missingness. Rather than competing directly with imaging foundation models, our results suggest that these approaches may ultimately be complementary. Future work may leverage powerful modality-specific representations, such as pre-trained MRI encoders, while preserving downstream missingness-aware integration and uncertainty modelling. Such a modular strategy may combine the representational strength of foundation models with the flexibility required for heterogeneous clinical data. Moreover, predicting continuous cognitive outcomes proved more challenging than binary classification across all models. This reduced external generalisation likely reflects the inherent difficulty of cross-cohort score prediction in the absence of harmonised assessment protocols, compounded by a scale mismatch between CDR-SB, the training target, and the global CDR used as observed scores in OASIS-3. Furthermore, predicted values exhibited a compressed dynamic range relative to the clinician rated scores, with a systematic tendency toward regression to the population mean, a well-characterised phenomenon in neural network regression. Taken together, these observations suggest that targeted methodological refinements, such as post-hoc output calibration, for example, that of temperature scaling, distributional loss functions,

or architectural modifications, may be useful to better capture the full variability of cognitive outcomes. Scientifically, cross-cohort and cross-method validation revealed key brain regions consistently identified across datasets and machine learning methods. Notably, cortical thickness in the bilateral temporal pole emerged as the most consistently implicated MRI region across models and cohorts. While hippocampal atrophy is classically considered the hallmark structural signature of AD, the prominence of temporal pole thinning is consistent with evidence that cortical thickness loss in this area correlates with cognitive decline and plasma tau levels in both MCI and AD populations [8, 22]. As a core node of the limbic network, the temporal pole subserves semantic memory and social cognition and undergoes early atrophy in both AD and frontotemporal dementia; its predictive relevance may therefore extend across the full diagnostic spectrum. However, sufficiency analysis indicated that robust AD discrimination arises not from temporal pole atrophy in isolation, but from its integration with complementary clinical, genetic, biomarker, and imaging information. Beyond convergent attribution evidence, our results revealed a clear dissociation in generalisation performance between AIBL and OASIS-3. Generalisation to AIBL was consistent with the shared recruitment strategies and clinical protocols between ADNI and AIBL. By contrast, the marked performance degradation on OASIS-3 when applying ADNI-trained models is consistent with known differences between the two cohorts. Specifically, unlike ADNI and AIBL, OASIS-3 enrols preclinical populations, comprising a higher proportion of cognitively normal participants, fewer AD patients, and no MCI subjects [25]. Nevertheless, the sensitivity improvements observed for NITROGEN on both external cohorts suggest that cross-sample attention may confer a form of implicit regularisation that enhances robustness to distributional shift, even in the absence of explicit domain adaptation. Finally, we introduced a missingness-aware adjustment of Shannon entropy by scaling prediction uncertainty according to modality-level importance, in order to better reflect confidence under incomplete multimodal inputs and while this approach yielded only marginal improvements in accuracy over standard uncertainty filtering, it provides a simple mechanism for more reliable uncertainty estimation in decisionsupport settings where robustness to missing modalities is essential. We acknowledge several limitations of our study. First, all models were trained exclusively on ADNI. We do so because ADNI is a well-characterised research cohort with relatively uniform acquisition protocols; training the model only on one cohort, despite evaluating externally on two other datasets, however, may limit generalizability of the trained model to clinical populations with greater demographic and technical heterogeneity. Consistent with this, we observe substantial performance degradation on OASIS-3, likely attributable to distributional shift between cohorts. Future work may consider transfer learning or domain adaptation strategies to improve robustness across datasets, sites and scanners. Second, prediction of cognitive scores remained moderate across all architectures, potentially reflecting the inherent variability of psychometric instruments and the complex, non-linear relationship between neuroimaging features and cognitive outcomes [26]. In parallel, the training instability of NITROGEN on MMSE and CDR-SB prediction suggests that task-specific modifications, such

as auxiliary losses or score-stratified sampling, may be useful to improve continuous outcome prediction. Third, the sufficiency analysis revealed limitations of conventional performance metrics for diagnosing feature importance under extreme feature ablation. Although the identified temporal pole region was consistently selected across datasets and methods, restricting the models to this feature frequently produced degenerate behaviours, including constant predictions or uniformly uncertain outputs. In several cases, standard metrics alone failed to adequately characterise these failure modes, requiring direct inspection of prediction distributions and probability outputs. More informative evaluation strategies may therefore be required to distinguish genuine feature sufficiency from prediction collapse under severe information restriction. Fourth, the proposed missingness-aware uncertainty adjustment may partially capture information already encoded in the predictive distribution itself. Missing features can reduce confidence both directly, through the absence of informative inputs, and indirectly, because incomplete feature combinations deviate from patterns observed during training. Consequently, the additional uncertainty induced by missing modalities may overlap with uncertainty already reflected in the predicted probabilities. More generally, disentangling epistemic uncertainty from uncertainty attributable to missing or incomplete observations remains an open challenge. Future work may therefore investigate alternative uncertainty estimation approaches, including deep ensembles, Monte Carlo dropout, or explicit aleatoric-epistemic decompositions. Finally, while our study included individuals spanning the CN, MCI, and AD spectrum, the analyses remained cross-sectional and therefore cannot directly characterise within-subject disease trajectories, the proposed NITROGEN architecture, however is naturally suitable to modelling relationships across samples and disease stages. Future work may extend the framework to longitudinal settings using mixed-effects formulations, repeatedmeasures designs, neural dynamical systems, or physics-informed approaches to better capture the temporal progression of neurodegeneration.

4 Methods 4.1 Data Three cohorts were included in this study: ADNI [27] as the primary dataset, and OASIS-3 [25] and AIBL [28] for external validation. Structural MRI scans were processed using our previously established preprocessing pipeline [29]. The feature set used in this study comprises MRI-derived regional cortical thickness measures, demographic variables (age, sex, and years of education), APOE ε4 allele counts, and cerebrospinal fluid (CSF) biomarkers, including amyloid-β (Aβ ), total tau (t-tau), and phosphorylated tau (p-tau).

Alzheimer’s Disease Neuroimaging Initiative (ADNI) Data used in the preparation of this article were obtained from the Alzheimer’s Disease Neuroimaging Initiative (ADNI) database (adni.loni.usc.edu). The ADNI was launched in 2003 as a public-private partnership, led by Principal Investigator Michael W. Weiner, MD. The primary goal of ADNI has been to test whether serial

magnetic resonance imaging (MRI), positron emission tomography (PET), other biological markers, and clinical and neuropsychological assessment can be combined to measure the progression of mild cognitive impairment (MCI) and Alzheimer’s disease (AD). The ADNI dataset comprised 7,858 cross-sectional samples derived from repeated visits from 1,776 unique individuals, of whom 56.4% were male. Diagnostic groups comprise 2,497 cognitively normal (CN), 3,583 mild cognitive impairment (MCI), and 1,778 Alzheimer’s disease (AD) subjects. APOEε4 genotyping was available for 1,769 unique subjects, of whom 53.3% carried no risk allele, 36.8% carried one allele, and 9.9% carried two alleles. Participants had a mean age of 73.7 years (range: 54.4–91.4) and a mean educational attainment of 16.1 years (range: 4– 20). CSF biomarker measurements were available for 2,353 samples or observations, with mean concentrations of 956.5 pg/mL for Aβ (range: 200.0 -1,700.0 pg/mL), 293.2 pg/mL for total tau (range: 80.0–1,300.0 pg/mL), and 28.1 pg/mL for phosphorylated tau (range: 8.0–120.0 pg/mL), consistent with previously reported ADNI biomarker distributions [27].

Open Access Series of Imaging Studies (OASIS)-3 The OASIS-3 dataset comprised 2,675 cross-sectional samples from 1,078 unique individuals, of whom 54.9% were female. Diagnostic groups comprise 2,228 CN, 49 MCI, and 398 AD samples. Participants had a mean age of 71.4 years (range: 45–97) and a mean educational attainment of 15.8 years (range: 6–29). APOEε4 genotyping was available for all 1,078 unique subjects, of whom 59.7% carried no risk allele, 34.8% carried one allele, and 5.5% carried two alleles. CSF biomarker variables were included as fully missing columns to harmonise the feature space with ADNI. Australian Imaging, Biomarkers and Lifestyle (AIBL) Data was collected by the AIBL study group. AIBL study methodology has been reported previously (Ellis et al. 2009). The AIBL dataset comprised 1,286 crosssectional samples from 684 unique individuals, of whom 56.0% were female. Diagnostic groups included 947 CN, 184 MCI, and 155 AD samples. Participants had a mean age of 73.7 years (range: 55–96). MRI-derived cortical thickness features were available for 1,282 of 1,286 samples. APOEε4 genotyping was available for all 684 subjects, of whom 70.0% carried no risk allele, 24.9% carried one allele, and 5.1% carried two alleles. Education and CSF biomarker variables were included as fully missing columns to harmonise the feature space with ADNI. Missing Variables Not all modalities were available across cohorts: CSF biomarkers were missing in OASIS-3, while years of education were unavailable in AIBL. To ensure a consistent feature space and avoid unnecessary sample exclusion, these variables were retained as missing features. Key cohort characteristics are summarised in Fig. 2.

4.2 Models All modalities in this study were represented as tabular data. Clinical tabular data in healthcare often suffer from missing values, high variability, and noise, which pose significant challenges for predictive modelling [9]. In the context of AD, where integrating diverse data sources—including clinical assessments, imaging, and fluid biomarkers, is crucial, improving current architectures becomes a pressing need. As such, here we introduce an attention based architecture that operate directly on incomplete tabular representations, enabling robust integration of heterogeneous modalities without requiring explicit imputation. This approach supports predictive tasks of clinical relevance under realistic conditions of missingness and noise.

4.2.1 Architecture The main modules in our architecture combine NAIM [17], which introduces a reformulation of masked self-attention to handle missing values by ignoring their contributions during training, and SAINT [18], which leverages intersample attention to stabilise predictions in the presence of noisy or incomplete inputs. The backbone of the proposed architecture is illustrated in Figure 6. The NAIM model introduces a masked self-attention mechanism specifically designed to handle missing values in tabular data [17]. In NAIM, input features are first transformed into embeddings, with categorical features being transformed into learnable token embeddings that encode both feature identity and value, while numerical variables are projected into the same embedding space through learned linear fully connected layer, yielding a unified token representation suitable for attention-based modelling. Then, to explicitly exclude contributions from missing data, NAIM introduces feature-wise masking directly within the attention computation. The modified attention is defined in [17] as:

QK T F eatureM askedAttention(Q, K, V ) = ReLU softmax √ + M + M ′ V, dh (1) where Q, K, V ∈ Rn×dh are query, key, and value projections, respectively, derived from input embeddings, and M ∈ Rn×n is a masking matrix encoding missing features. The additional mask M ′ = M T further suppresses contributions from missing values after the attention weights are computed. This “double masking“ ensures that missing features do not influence the learned dependencies, either directly or indirectly, allowing the model to operate robustly on incomplete inputs without requiring explicit imputation. A graphical representation is in Figure 6. While the masked attention mechanism operates at the feature level within each sample, it does not explicitly leverage relationships between different samples. To address this, in our architecture, we extend NAIM with an intersample attention module inspired by SAINT [18]. The key idea is to enable the model to learn similarities 23

Classification

Transformer

Validation AD status

Test

ADNI (test) OASIS AIBL

Regression

MMSE CDR

Feature analysis

Masked Attention Encoder Block

…

Intersample Attention Encoder Block

Categorical tokenizer

Classifier

Sigmoid

Linear Embedding

Intersample Attention

Norm

GELU

Norm

Masked Multihead Attention

Norm

Categorical tokenizer

… Norm

Linear Embedding

Feature analysis

Fig. 6: Top: Overview of the experimental design across datasets (ADNI, OASIS3, AIBL), illustrating training and cross-cohort evaluation settings used to assess generalisation across independently acquired clinical populations. Bottom : Overview of the NITROGEN architecture. The model is composed of stacked Masked Attention Encoder Blocks, which handle feature-wise masked attention as described in Eq. (4.2.1), followed by Intersample Attention Blocks, which allow the model to incorporate information across samples according to the mechanism detailed in Eq. (4.2.1). This modular design enables modelling of both feature-level dependencies and crosssample interactions, enhancing performance in noisy and heterogeneous tabular data settings.

across samples (e.g., across subject), allowing it to exploit shared patterns within a

mini-batch. This may be beneficial in clinical settings, where subjects may exhibit similar profiles despite missing or noisy measurements. Given feature embeddings e ∈ Rb×n×de , where b is the batch size, n the number of features, and de the embedding dimension, each sample is reshaped into a single token of dimension dh = n · de , yielding e′ ∈ Rb×dh . Query, key, and value projections are computed as: ′ Q′ = e′ WQ′ , K ′ = e′ WK , V ′ = e′ WV′ ∈ Rb×dh ,

(2)

′ respectively, with dh = n · de , and WQ′ , WK , WV′ ∈ Rd×d . Self-attention is then applied across the batch dimension in the following way:

Q′ K ′T ′ )V . IntersampleAttention(Q′ , K ′ , V ′ ) = sof tmax( √ dh

(3)

Finally, the resulting representations are reshaped back to the original feature-wise structure Rb×n×de . This allows each sample to attend to other samples in the batch, capturing interpatient relationships that can help stabilise predictions in the presence of missing or noisy data.

4.2.2 Attention-Based Models We evaluated two attention-based models across all tasks: the proposed NITROGEN model, which extends NAIM by incorporating intersample attention inspired by SAINT [18] and a baseline NAIM model having the same parameters as in the original paper [17]. Both models share the same input processing pipeline: categorical variables are embedded using dataset-specific indices and cardinalities, continuous features are projected into a shared token space via learnable linear transformations, and missing values are handled natively through a dedicated −∞ token that is excluded by the masked self-attention mechanism. Both NAIM and NITROGEN were trained using the AdamW optimiser with a batch size of 128 and early stopping based on validation loss, with a patience of 50 epochs for both models. Output layers and loss functions were adapted per task: a single sigmoid unit with binary cross-entropy loss for binary classification tasks, and two linear output units with mean squared error loss for simultaneous MMSE and CDR-SB prediction. NAIM comprised six transformer encoder layers with token embedding dimension dtoken = 6, three self-attention heads, a feed-forward dimension of 1,000, GELU activation, and no dropout. NAIM was trained with a learning rate of 1 × 10−4 , final model was saved after 31 epochs as it reached minimum validation loss. NITROGEN extended this architecture with four transformer encoder layers, dtoken = 32, four self-attention heads, a feed-forward dimension of 463, and a single intersample attention layer following the encoder layers, enabling the model to attend jointly across feature tokens and across samples within a batch. NITROGEN was

trained with a learning rate of 1 × 10−6 . The final model was saved after 7 epochs, due to the very fast convergence of NITROGEN despite reduced learning rate.

4.3 Explainable AI To characterise feature relevance across models, we computed post-hoc perturbationbased feature attributions using the Captum interpretability framework [30]. Specifically, we employed feature permutation attribution, computed independently for each trained model on the held-out test sets from ADNI, OASIS-3, and AIBL. Attributions were computed with respect to the per-sample log-loss rather than raw predicted probability, providing a loss-sensitive measure of feature relevance that accounts for both the direction and confidence of predictions. For each feature, attribution was obtained by independently permuting that feature across all test samples while keeping all other features fixed, and measuring the resulting change in model log-loss. To obtain global importance profiles, sample-level attributions were aggregated using the mean absolute attribution score across the test set. This perturbation-based approach is modelagnostic and applicable across both transformer-based and tree-based architectures, enabling a consistent attribution procedure across all evaluated models [31]. We note that feature permutation shuffles all feature values uniformly across samples, including missing value tokens (−∞), redistributing missingness patterns across samples. Attribution scores for partially observed features therefore reflect the joint contribution of feature magnitude and missingness structure to model predictions.

4.4 Uncertainty quantification For clinical deployment, predictive performance alone is insufficient; models must also provide reliable estimates of predictive uncertainty. In healthcare applications, uncertainty estimation is essential for assessing model reliability, identifying potentially incorrect predictions, and supporting safe decision-making. Uncertainty can arise from two primary sources: data uncertainty (aleatoric) and model uncertainty (epistemic) [32, 33]. In this work, we focus on predictive uncertainty derived from model outputs, with particular attention to how this uncertainty is affected by missing modalities in incomplete multimodal clinical data. In standard neural network classifiers, predictive confidence is commonly approximated using the softmax probabilities produced by the final output layer. Although these probabilities are not guaranteed to represent well-calibrated posterior probabilities, prior work has shown that maximum softmax probability can serve as a simple yet effective baseline for detecting misclassified or out-of-distribution samples [34]. Beyond mis-classification detection, predictive entropy and related measures derived from the softmax distribution have been proposed as indicators of model reliability and generalisation performance [35]. However, softmax-based confidence estimates are often overconfident, particularly under dataset shift, motivating the need for calibration-aware uncertainty assessment. Alternative approaches have therefore explored transformations of the logits—such as kurtosis or normalised logit differences—as improved proxies for prediction correctness [36].

More generally, predictive uncertainty can be modelled explicitly. For instance, discriminative models may be extended to output parameters of predictive distributions (e.g., mean and variance), while generative approaches estimate uncertainty by modelling the underlying data distribution [33]. Although these methods provide principled uncertainty estimates, they typically require architectural modifications and additional training constraints. We start with the assumption the predicted class probabilities can serve as a preliminary reflection of predictive confidence [34]. Specifically, for a model output h(x), the softmax probability for class i is defined as follows: exp(hi (x)) πi = PC , (4) c=1 exp(hc (x)) where πi denotes the predicted probability of class i, for i = 0, 1, . . . , C , and the uncertainty is subsequently summarised using the entropy of the predictive distribution. Further, to quantify uncertainty, we compute the Shannon entropy of the predictive distribution [37]:

H (π (x)) = −

C X

πi (x) log πi (x).

(5)

i=1

In the binary classification setting (C = 2), this reduces to:

H (p) = −p log p − (1 − p) log(1 − p),

(6)

where p = π1 (x) denotes the predicted probability of the positive class. The entropy is maximal at p = 0.5 (maximum uncertainty) and vanishes as p → 0 or p → 1, reflecting increasingly confident predictions.

4.5 Adjusting Prediction Confidence for Missing Modalities In multimodal clinical settings, entire diagnostic modalities are frequently unavailable due to cost, acquisition constraints, or clinical workflows. While many machine learning models are designed to handle missing data, they typically produce predictions without explicitly quantifying how the absence of specific modalities impacts predictive confidence. To the best of our knowledge, existing approaches either focus on handling missing inputs during training or on estimating predictive uncertainty, but do not connect these two aspects in a principled manner. To address this gap, we propose a missingness-aware uncertainty framework that provides a simple framework for incorporating modality-level missingness into prediction confidence estimates.

Modality importance estimation. To quantify the contribution of each modality, we used the feature attribution scores obtained with permutation importance in the explainable AI analysis. Let pi,f denote

the permutation importance associated with feature f for sample i, and let Mk be the set of features belonging to modality k . For each sample, the importance of modality k is computed by summing the absolute permutation importance values of all features belonging to that modality and averaging by the number of features:

mi,k =

1 X f ∈ M∥ |pi,f |. |M∥ |

(7)

The overall importance of modality k is then obtained by averaging across all samples: N 1 X Ik = mi,k , (8) N i=1 where N is the number of samples. Finally, to enable comparison across modalities, the modality importance scores are normalised: Ik I˜k = P , (9) j Ij where the sum runs over all modalities j . The normalised importance I˜k ∈ [0, 1] represents the relative contribution of modality k to the model predictions.

Missingness-aware uncertainty adjustment. For a subject x with a set of missing modalities Mmiss , we define a missingness-aware uncertainty inflation factor IF (x) = 1 + β

I˜k ,

(10)

k∈Mmiss

where β ≥ 0 controls the sensitivity of the uncertainty adjustment. The adjusted uncertainty interval width is then given by

P Iadj (x) = IF (x) · P Iwidth .

(11)

This formulation increases predictive uncertainty in proportion to the estimated importance of the missing modalities. As a result, the absence of highly informative modalities leads to wider uncertainty intervals, providing a method for incorporating missing modality information into confidence estimates.

Code availability Code to reproduce the experiments described in this study is publicly available at https://github.com/cschneuw/nitrogen.

5 Experiments 5.1 Baseline Models In addition to NAIM and NITROGEN, five baseline models were evaluated for benchmarking: XGBoost, LightGBM, Random Forest, MA-LASSO, and MA-GBT. These baselines were selected to include both widely used tree-based methods and recently proposed missing-aware tabular learning approaches. XGBoost [38] was configured with 400 boosting rounds, unrestricted maximum depth, a learning rate of 0.15, subsampling and column sampling ratios of 0.6, and histogram-based tree construction. LightGBM [39] was used with 400 boosting rounds, a maximum depth of 8, a learning rate of 0.1, and full row and column sampling. Both XGBoost and LightGBM handle missing values natively and were restricted to classification tasks. Random Forest [40] comprised 200 trees with Gini impurity criterion, unrestricted depth, and a minimum of 4 samples per leaf, with missing values handled through mean imputation. Two missing-aware baselines—MA-LASSO and MA-GBT [41]—were evaluated on classification tasks only, as their formulation does not extend to regression. MA-LASSO was trained with regularisation parameters α = 100 and β = 100; MA-GBT used 200 estimators, a maximum depth of 5, a learning rate of 0.1, and a subsampling ratio of 0.6. Both missing-aware models explicitly penalise reliance on missing features through modified objective functions, with input features zero-imputed at training and inference time. Owing to its substantially longer training time, MA-GBT was not evaluated in all experimental settings.

5.2 Binary Classification Tasks Two binary classification tasks were defined and evaluated on stratified 60/20/20 train/validation/test splits derived from ADNI only, with OASIS-3 and AIBL reserved exclusively as held-out external cohorts. Splits were performed using group-stratified shuffling to ensure that repeated measurements from the same individual were always assigned to the same partition, preventing subject-level data leakage. Continuous features were standardised using a scaler fitted on the training set only and applied to all subsequent splits and external cohorts. Classification performance was assessed over ten independent training runs initialised with different random seeds (123, 42, 0, 7, 13, 21, 37, 55, 77, and 99). Metrics reported in Table A2 and the ROC curves in Figure 3 are presented as the mean and standard deviation across runs. All other analyses were performed using a fixed random seed (123).

Task 1: CN vs AD Retaining only CN and AD samples from ADNI, the training set comprised 2,560 samples (1,452 CN, 1,108 AD), the validation set 882 samples (524 CN, 358 AD; 59.4% vs 40.6%), and the test set 833 samples (521 CN, 312 AD; 62.5% vs 37.5%). External cohorts comprised 2,625 samples from OASIS-3 (2,228 CN, 398 AD) and 1,102 samples from AIBL (947 CN, 155 AD).

Task 2: CN vs MCI/AD Using all available ADNI samples, the training set comprised 4,704 samples (1,438 CN, 3,266 MCI/AD), the validation set 1,561 samples (538 CN, 1,023 MCI/AD), and the test set 1,593 samples (521 CN, 1,072 MCI/AD). The class imbalance reflects the natural diagnostic distribution in ADNI and was preserved across splits without resampling. External cohorts comprised 2,675 samples from OASIS-3 (2,228 CN, 447 demented) and 1,286 samples from AIBL (947 CN, 339 MCI/AD).

5.3 Multi-class Diagnostic Classification All architectures were adapted for three-class classification (CN / MCI / AD) by changing the output to 3 classe while retaining the hyper-parameters established in the binary setting. For tree-based models, objectives were updated to their multi-class equivalents. For NAIM and NITROGEN, the output layer was adapted to produce three logits passed through a softmax activation, with the predicted class determined by argmax over the probability vector and optimised with cross-entropy loss. Probabilities were formatted as an N × C matrix to enable per-class and macro-averaged AUC computation under the one-vs-rest (OvR) strategy. Following the same group-stratified 60/20/20 splitting strategy as in Section 5.2, the ADNI training set comprised 4,704 samples (CN: 1,438, MCI: 2,196, AD: 1,070), the validation set 1,561 samples (CN: 538, MCI: 680, AD: 343), and the test set 1,593 samples (CN: 521, MCI: 707, AD: 365). OASIS-3 (2,675 samples; CN: 2,228, MCI: 49, AD: 398) and AIBL (1,286 samples; CN: 947, MCI: 184, AD: 155) were used as held-out external validation cohorts. The near-absence of MCI samples in OASIS-3 reflects differences in recruitment strategy relative to ADNI and AIBL, with OASIS3 specifically targeting preclinical populations, resulting in a higher proportion of cognitively normal participants and fewer dementia cases.

5.4 Cognitive Score Prediction To extend the evaluation beyond diagnostic classification, models were trained on ADNI to jointly predict Mini-Mental State Examination (MMSE) and Clinical Dementia Rating Sum of Boxes (CDR-SB) scores from multimodal input features. For NAIM and NITROGEN, the output layer was replaced with a two-dimensional linear unit optimised with mean squared error loss, while all other architectural components and training settings remained identical to the classification setting; tree-based models were similarly adapted to predict both targets simultaneously. The ADNI cohort (7,858 samples; CN: 2,497, MCI: 3,583, AD: 1,778) was partitioned into training, validation, and test sets using a 60/20/20 group-stratified split (seed = 123), yielding 4,704 training samples (CN: 1,256, MCI: 2,236, AD: 1,212), 1,561 validation samples (CN: 446, MCI: 723, AD: 392), and 1,593 test samples (CN: 521, MCI: 707, AD: 365). OASIS-3 (CN: 2,228, MCI: 49, AD: 398) and AIBL (CN: 947, MCI: 184, AD: 155) were reserved exclusively as external validation cohorts. As OASIS-3 and AIBL provide only the global Clinical Dementia Rating (CDR) rather than CDR-SB, CDR was used as a surrogate target for out-of-distribution evaluation; while CDR and CDR-SB differ in scale and granularity, they are expected to be

strongly correlated, allowing cross-cohort regression performance to be approximately assessed.

5.5 Test of Sufficiency To verify that the identified MRI features are sufficient to sustain classification performance, we performed a sufficiency test on held-out subsets not used during feature attribution computation. The training and validation splits were identical to those described in Section 5.2. To prevent data leakage between attribution computation and sufficiency evaluation, each test cohort was further partitioned into a 70% attribution subset and a 30% held-out evaluation subset using stratified splits with subject identifiers as grouping variable. The attribution subsets comprised 580 ADNI samples (CN: 353, MCI/AD: 227), 1,833 OASIS-3 samples (CN: 1,554, MCI/AD: 279), and 762 AIBL samples (CN: 649, MCI/AD: 113). The corresponding held-out evaluation subsets comprised 253 ADNI samples (CN: 168, MCI/AD: 85), 793 OASIS-3 samples (CN: 674, MCI/AD: 119), and 340 AIBL samples (CN: 298, MCI/AD: 42). Permutation-based feature attributions were computed exclusively on the 70% attribution subsets. Features were ranked by mean absolute attribution score for each model-dataset combination, and the most consistent MRI features were identified by retaining those appearing within the top-30 across all combinations. This procedure yielded one cortical thickness features from the temporal pole: L Limbic TempPole 3. Sufficiency was evaluated on the remaining 30% of the ADNI test set using models trained on the original ADNI training data described in Section 5.2. The trained models were applied without retraining under three input configurations: (i) the full multimodal feature set, used as the baseline; (ii) the two selected temporal pole features together with all non-MRI features, with all other MRI cortical thickness features masked; and (iii) the two selected temporal pole features only, with all remaining features masked. Performance metrics and ROC curves are reported in Figure 4f and Table A7.

5.6 Hyper-parameter sweeps To ensure fair comparison across architectures, hyperparameter selection was performed separately for each model family using the Task 1 (CN vs AD) training and validation sets (splits were generated with fixed seed = 123). The resulting hyperparameter configurations were then kept fixed across all tasks and cohorts to maintain a consistent experimental framework and limit the influence of repeated task-specific optimisation. All experiments were tracked and managed using Weights & Biases (W&B) sweeps [42]. For NITROGEN, Bayesian search with hyperband-based early termination (minimum 15 iterations, η = 2, s = 5) was performed to maximise balanced accuracy on the validation set, over the following hyperparameter space: number of encoder layers (2, 3, 4, 6), token embedding dimension dtoken (32, 48, 64), number of self-attention heads (4, 8), feed-forward hidden dimension (integer-uniform between

128 and 1,024), dropout rate (0, 0.1, 0.2), activation function (ReLU, GELU), presence of bias terms (learnable intercepts), mask type (0: missing-feature masking; 1: missing-feature masking with masking of interactions involving missing features; 2: both forward and reverse missing-feature masks), number of intersample attention layers (1, 2), and intersample layer placement (applied after the encoder, or interleaved with encoder layers); intersample attention was enabled in all sweep configurations. The optimiser (AdamW), embedder initialisation (uniform), and maximum number of training epochs (200, with early stopping) were fixed across runs, and the learning rate was sampled log-uniformly between 1 × 10−5 and 1 × 10−3 . During both hyperparameter search, the learning rate was adaptively reduced using a scheduler (factor = 0.5, patience = 5 epochs), and training was terminated early if the monitored validation metric failed to improve by at least 1 × 10−4 over a fixed patience of 20 epochs, with the same configuration retained for all downstream training runs. For tree-based models, grid search was conducted over commonly used ranges. XGBoost sweeps varied number of estimators (100, 200, 400), maximum tree depth (4, 6, 8 or None), learning rate (0.01–0.2), and row and column subsampling ratios (0.6, 0.8, 1.0). LightGBM sweeps additionally explored number of leaves and subsampling ratios. Random Forest sweeps varied number of estimators, maximum depth, maximum features per split, and minimum samples per leaf or split. For missing-aware models, MA-LASSO sweeps explored combinations of regularisation parameters α and β ∈ {0.1, 1, 10, 100}. MA-GBT sweeps varied number of estimators, maximum depth, learning rate, sub-sampling ratio, and missingness penalty α, with optional node-wise regularisation enabled or disabled. In all cases, the final hyperparameter configuration was selected based on validation balanced accuracy and held fixed for all subsequent tasks and cross-dataset evaluations, with only output layer and loss function adapted per task.

5.7 Selection of the Uncertainty Sensitivity Parameter The sensitivity parameter β controls the influence of missing modalities on predictive uncertainty. To select an appropriate value, we performed a grid search over β ∈ [0, 2] on the ADNI validation set (N = 882) for the CN versus AD classification task, with a fixed uncertainty threshold τ = 0.900, chosen to retain a meaningful proportion of samples while allowing selective abstention on uncertain predictions. For each value of β , we computed the adjusted uncertainty Hadj and retained samples satisfying: Hadj (x) ≤ τ. (12) To identify the optimal β , we maximised the product of two complementary criteria: the Spearman rank correlation between adjusted uncertainty and prediction errors, ρ(Hadj , 1[ŷ ̸= y ]), which measures how well uncertainty ranks incorrect predictions; and the retention gap, defined as the difference in accuracy between retained and rejected samples. Let Accretained denote the classification accuracy on samples with Hadj (x) ≤ τ (retained for prediction) and Accrejected denote the accuracy on samples

with Hadj (x) > τ (abstained from):

S (β ) = max(0, ρ(Hadj , 1[ŷ ̸= y ])) × max(0, Accretained − Accrejected ) .

(13)

This joint criterion rewards configurations where uncertainty is both well-calibrated with respect to errors and effectively separates reliable from unreliable predictions, while remaining insensitive to trivial solutions at β = 0 where no modality adjustment is applied. The optimal parameter was identified at β ∗ = 0.400 with τ = 0.900, achieving a Spearman correlation of ρ = 0.348 (p < 10−26 ) between adjusted uncertainty and prediction errors (Fig. A6). On the validation set, uncertainty thresholding retained 85.7% of samples (756/882), improving accuracy from 0.862 to 0.900 (∆ = +0.038), with rejected samples exhibiting substantially lower accuracy (0.635), confirming that the adjusted uncertainty effectively identifies unreliable predictions.

5.8 Uncertainty adjustment evaluation To assess whether the improvement in selective accuracy obtained through uncertainty thresholding is statistically significant, we conducted permutation tests on both the validation (see Fig. 5). For each permutation, true diagnostic labels were randomly shuffled and the same fraction of samples as retained by the uncertainty threshold was selected at random, without any uncertainty-based filtering. This procedure was repeated 1,000 times to construct a null distribution of accuracy improvements expected by chance. The observed accuracy improvement, defined as the difference between selective accuracy on uncertainty-retained samples and accuracy on the full dataset, was then compared against this null distribution. A one-sided empirical pvalue was computed as the fraction of permuted improvements exceeding the observed value. In addition to selective accuracy, we assessed the quality of the adjusted uncertainty estimates through two complementary analyses. First, we computed the Spearman rank correlation ρ between adjusted uncertainty scores and binary prediction errors, measuring whether higher uncertainty reliably identifies misclassified samples. Second, we evaluated calibration by partitioning samples into decile bins of uncertainty and computing the Spearman correlation between mean bin uncertainty and mean bin error rate, providing a coarser but more stable measure of monotonic calibration across the uncertainty range. Both analyses were performed for raw and missingness-adjusted uncertainty to quantify the improvement attributable to the modality adjustment.

Acknowledgements. C.S.D., N.B. and O.Y.C. conceptualised the study. C.S.D. developed the methods, wrote the codes, performed data processing and analysis. N.B. contributed to uncertainty quantification adjustment. C.S.D. and O.Y.C. wrote the manuscript, with comments from all other authors. This work was supported by Swiss National Science Foundation under Grants CR00I5-235987 and 3200-0-239967. Data collection and sharing for the Alzheimer’s Disease Neuroimaging Initiative (ADNI) is funded by the National Institute on Aging (National Institutes of

Health Grant U19 AG024904). The grantee organisation is the Northern California Institute for Research and Education. In the past, ADNI has also received funding from the National Institute of Biomedical Imaging and Bioengineering, the Canadian Institutes of Health Research, and private sector contributions through the Foundation for the National Institutes of Health (FNIH) including generous contributions from the following: AbbVie, Alzheimer’s Association; Alzheimer’s Drug Discovery Foundation; Araclon Biotech; BioClinica, Inc.; Biogen; Bristol-Myers Squibb Company; CereSpir, Inc.; Cogstate; Eisai Inc.; Elan Pharmaceuticals, Inc.; Eli Lilly and Company; EuroImmun; F. Hoffmann-La Roche Ltd and its affiliated company Genentech, Inc.; Fujirebio; GE Healthcare; IXICO Ltd.; Janssen Alzheimer Immunotherapy Research & Development, LLC.; Johnson & Johnson Pharmaceutical Research & Development LLC.; Lumosity; Lundbeck; Merck & Co., Inc.; Meso Scale Diagnostics, LLC.; NeuroRx Research; Neurotrack Technologies; Novartis Pharmaceuticals Corporation; Pfizer Inc.; Piramal Imaging; Servier; Takeda Pharmaceutical Company; and Transition Therapeutics. Data used in the preparation of this article was obtained from the Australian Imaging Biomarkers and Lifestyle flagship study of ageing (AIBL) funded by the Commonwealth Scientific and Industrial Research Organisation (CSIRO) which was made available at the ADNI database (www.loni.usc.edu/ADNI). Data were provided in part by OASIS-3: Longitudinal Multimodal Neuroimaging: Principal Investigators: T. Benzinger, D. Marcus, J. Morris; NIH P30 AG066444, P50 AG00561, P30 NS09857781, P01 AG026276, P01 AG003991, R01 AG043434, UL1 TR000448, R01 EB009352. AV-45 doses were provided by Avid Radiopharmaceuticals, a wholly owned subsidiary of Eli Lilly.

Declarations Data availability Three cohorts were used in this study: ADNI (primary dataset), and OASIS3 and AIBL (external validation). ADNI data are available at https://adni.loni. usc.edu/. Access can be requested for scientific investigation, teaching, or clinical research planning via the Image and Data Archive (IDA) platform (https://ida. loni.usc.edu). AIBL data are available to authorised users through a request process integrated with the ADNI access procedure via the LONI infrastructure, at https://ida.loni.usc.edu/collaboration/access/appApply.jsp?project=AIBL, or jointly with ADNI at https://ida.loni.usc.edu/collaboration/access/appApply.jsp?project= AIBL&project=ADNI, subject to agreement with the relevant project terms of use. OASIS-3 data are available at https://sites.wustl.edu/oasisbrains/home/oasis-3/ and can also be accessed via the LONI/IDA platform subject to a Data Use Agreement (https://ida.loni.usc.edu/collaboration/access/appLicense.jsp)

References [1] Safiri, S., Ghaffari Jolfayi, A., Fazlollahi, A., Morsali, S., Sarkesh, A., Daei Sorkhabi, A., Golabi, B., Aletaha, R., Motlagh Asghari, K., Hamidi, S., 34

et al.: Alzheimer’s Disease: A Comprehensive Review of Epidemiology, Risk Factors, Symptoms Diagnosis, Management, Caregiving, Advanced Treatments and Associated Challenges. Frontiers in Medicine 11, 1474043 (2024) https: //doi.org/10.3389/fmed.2024.1474043 [2] Nichols, E., Steinmetz, J.D., Vollset, S.E., et al.: Estimation of the Global Prevalence of Dementia in 2019 and Forecasted Prevalence in 2050: An Analysis for the Global Burden of Disease Study 2019. Lancet Public Health 7, 105–125 (2022) https://doi.org/10.1016/S2468-2667(21)00249-8 [3] Wimo, A., Gauthier, S., Prince, M.: Global Estimates of Informal Care. Technical report, Alzheimer’s Disease International (2018). https://www.alzint.org/u/ global-estimates-of-informal-care.pdf [4] Wimo, A., Seeher, K., Cataldi, R., et al.: The Worldwide Costs of Dementia in 2019. Alzheimer’s & Dementia 19(7), 2865–2873 (2023) https://doi.org/10.1002/ alz.12901 [5] Cheng, H., Kang, L., Xu, Y., Tanzi, R.E., Zhang, C., Wang, C.: Innovation breakthrough in the Alzheimer’s disease pharmaceutical industry. npj Drug Discovery 3(1), 9 (2026) https://doi.org/10.1038/s44386-026-00044-7 [6] Hunter, T.R., Santos, L.E., Tovar-Moll, F., De Felice, F.G.: Alzheimer’s Disease Biomarkers and Their Current Use in Clinical Research and Practice. Mol Psychiatry 30(1), 272–284 (2025) https://doi.org/10.1038/s41380-024-02709-z [7] Dessy, T., Barhdadi, A., Cyr, M.-C., Sandoval, J., Bherer, L., Rouleau, J., Provost, S., Lemieux Perreault, L.-P., Sylvestre, M.-P., Gagliano Taliun, S.A., Dubé, M.-P.: Disentangling the Effects of Sex and Gender on APOE ε4–Related Neurocognitive Impairment. Alzheimer’s & Dementia: Diagnosis, Assessment & Disease Monitoring 17(2), 70111 (2025) https://doi.org/10.1002/dad2.70111 [8] Frisoni, G.B., Fox, N.C., Jack, C.R., Scheltens, P., Thompson, P.M.: The Clinical Use of Structural MRI in Alzheimer Disease. Nature Reviews Neurology 6(2), 67–77 (2010) https://doi.org/10.1038/nrneurol.2009.215 [9] Vavekanand, R., Kumar, T., Kumar, S., Kumar, G., Laghari, A.A.: Multimodal Machine Learning Approaches in Predictive Healthcare Analytics: A Comprehensive Survey. Archives of Computational Methods in Engineering (2026) https: //doi.org/10.1007/s11831-026-10560-4 [10] Larose, C., Dey, D.K., Harel, O.: The Impact of Missing Values on Different Measures of Uncertainty. Statistica Sinica 29(2), 551–566 (2019) https://doi.org/ 10.5705/ss.202017.0113 26705478 [11] Twala, B.E.T.H., Jones, M.C., Hand, D.J.: Good Methods for Coping with Missing Data in Decision Trees. Pattern Recognition Letters 29(7), 950–956 (2008)

https://doi.org/10.1016/j.patrec.2008.01.010 [12] Van Ness, M., Bosschieter, T.M., Halpin-Gregorio, R., Udell, M.: The Missing Indicator Method: From Low to High Dimensions. In: Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. KDD ’23, pp. 5004–5015. Association for Computing Machinery, New York, NY, USA (2023). https://doi.org/10.1145/3580305.3599911 [13] Śmieja, M., Struski, L., Tabor, J., Zieliński, B., Spurek, P.: Processing of Missing Data by Neural Networks. In: Advances in Neural Information Processing Systems, vol. 31. Curran Associates, Inc., Canada (2018) [14] Zeng, Z., Liu, Y., Yao, S., Liu, J., Xiao, B., Liu, C., Gong, X.: Neural Networks Based on Attention Architecture Are Robust to Data Missingness for Early Predicting Hospital Mortality in Intensive Care Unit Patients. DIGITAL HEALTH 9, 20552076231171482 (2023) https://doi.org/10.1177/20552076231171482 [15] Gorishniy, Y., Rubachev, I., Khrulkov, V., Babenko, A.: Revisiting Deep Learning Models for Tabular Data. In: Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P.S., Vaughan, J.W. (eds.) Advances in Neural Information Processing Systems, vol. 34, pp. 18932–18943. Curran Associates, Inc., Online (2021) [16] Huang, X., Khetan, A., Cvitkovic, M., Karnin, Z.: TabTransformer: Tabular Data Modeling Using Contextual Embeddings. https://arxiv.org/abs/2012.06678v1 (2020). https://doi.org/10.48550/arXiv.2012.06678 [17] Caruso, C.M., Soda, P., Guarrasi, V.: Not Another Imputation Method: A Transformer-based Model for Missing Values in Tabular Datasets. arXiv (2025). https://doi.org/10.48550/arXiv.2407.11540 [18] Somepalli, G., Goldblum, M., Schwarzschild, A., Bruss, C.B., Goldstein, T.: SAINT: Improved Neural Networks for Tabular Data via Row Attention and Contrastive Pre-Training. arXiv (2021). https://doi.org/10.48550/arXiv.2106. 01342 [19] Vu, D.-T., Can, D.-C., Diaz, C.S., Bodelet, J.S., Blanc, G.E., Phan, H., Allali, G., Nguyen, V.-D., Cao, H., He, X., et al.: Residual Partial Least Squares Learning: Brain Cortical Thickness Simultaneously Predicts Eight Non-pairwise-correlated Behavioural and Disease Outcomes in Alzheimer’s Disease. bioRxiv, 2024–03 (2024) https://doi.org/10.1101/2024.03.11.584383 [20] Wei, X., Zhang, T., Xiong, R., Zhang, Q., Zhang, J., Jin, Z., Li, L.: Mapping Heterogeneous Brain Structural Subtypes in Alzheimer’s Disease and Mild Cognitive Impairment Using Normative Models. Translational Psychiatry 16(1), 168 (2026) https://doi.org/10.1038/s41398-026-03902-0 [21] Knopman, D.S., Amieva, H., Petersen, R.C., Chételat, G., Holtzman, D.M.,

Hyman, B.T., Nixon, R.A., Jones, D.T.: Alzheimer Disease. Nature Reviews Disease Primers 7(1), 33 (2021) https://doi.org/10.1038/s41572-021-00269-y [22] Advancing Research and Treatment for Frontotemporal Lobar Degeneration (ARTFL) investigators, Thijssen, E.H., La Joie, R., Wolf, A., Strom, A., Wang, P., Iaccarino, L., Bourakova, V., Cobigo, Y., Heuer, H., Spina, S., VandeVrede, L., Chai, X., Proctor, N.K., Airey, D.C., Shcherbinin, S., Duggan Evans, C., Sims, J.R., Zetterberg, H., Blennow, K., Karydas, A.M., Teunissen, C.E., Kramer, J.H., Grinberg, L.T., Seeley, W.W., Rosen, H., Boeve, B.F., Miller, B.L., Rabinovici, G.D., Dage, J.L., Rojas, J.C., Boxer, A.L.: Diagnostic Value of Plasma Phosphorylated Tau181 in Alzheimer’s Disease and Frontotemporal Lobar Degeneration. Nature Medicine 26(3), 387–397 (2020) https://doi.org/10. 1038/s41591-020-0762-2 [23] Shukla, A., Tiwari, R., Tiwari, S.: Alz-ConvNets for Classification of Alzheimer Disease Using Transfer Learning Approach. SN Computer Science 4(4), 404 (2023) https://doi.org/10.1007/s42979-023-01853-7 [24] Tak, D., Garomsa, B.A., Zapaishchykova, A., Chaunzwa, T.L., Climent Pardo, J.C., Ye, Z., Zielke, J., Ravipati, Y., Pai, S., Vajapeyam, S., Mahootiha, M., Parker, M., Pike, L.R.G., Smith, C., Familiar, A.M., Liu, K.X., Prabhu, S., Arnaout, O., Bandopadhayay, P., Nabavizadeh, A., Mueller, S., Aerts, H.J., Huang, R.Y., Poussaint, T.Y., Kann, B.H.: A generalizable foundation model for analysis of human brain MRI. Nature Neuroscience 29(4), 945–956 (2026) https://doi.org/10.1038/s41593-026-02202-6 [25] LaMontagne, P.J., Benzinger, T.L., Morris, J.C., Keefe, S., Hornbeck, R., Xiong, C., Grant, E., Hassenstab, J., Moulder, K., Vlassenko, A.G., Raichle, M.E., Cruchaga, C., Marcus, D.: OASIS-3: Longitudinal Neuroimaging, Clinical, and Cognitive Dataset for Normal Aging and Alzheimer Disease. medRxiv (2019). https://doi.org/10.1101/2019.12.13.19014902 [26] McDonough, I.M., Popp, T.E.: Linear and nonlinear relationships between cognitive subdomains of ability discrepancy and Alzheimer’s disease biomarkers. Neuropsychology 34(2), 211–226 (2020) https://doi.org/10.1037/neu0000606 [27] Shaw, L.M., Vanderstichele, H., Knapik-Czajka, M., Clark, C.M., Aisen, P.S., Petersen, R.C., Blennow, K., Soares, H., Simon, A., Lewczuk, P., Dean, R., Siemers, E., Potter, W., Lee, V.M.-Y., Trojanowski, J.Q.: Cerebrospinal Fluid Biomarker Signature in Alzheimer’s Disease Neuroimaging Initiative Subjects. Annals of Neurology 65(4), 403–413 (2009) https://doi.org/10.1002/ana.21610 [28] Ellis, K.A., Bush, A.I., Darby, D., De Fazio, D., Foster, J., Hudson, P., Lautenschlager, N.T., Lenzo, N., Martins, R.N., Maruff, P., Masters, C., Milner, A., Pike, K., Rowe, C., Savage, G., Szoeke, C., Taddei, K., Villemagne, V., Woodward, M., Ames, D.: The Australian Imaging, Biomarkers and Lifestyle (AIBL) Study of Aging: Methodology and Baseline Characteristics of 1112 Individuals 37

Recruited for a Longitudinal Study of Alzheimer’s Disease. International Psychogeriatrics 21(4), 672–687 (2009) https://doi.org/10.1017/S1041610209009405 . Accessed 2026-02-16 [29] Diaz, C.S., Vũ, D.-T., Bodelet, J., Can, D.-C., Blanc, G., Jiang, H., Yao, L., Pantaleo, G., Chén, O.Y.: OPTIMUS: Predicting Multivariate Outcomes in Alzheimer’s Disease Using Multi-modal Data Amidst Missing Values. IEEE Transactions on Biomedical Engineering, 1–45 (2025) https://doi.org/10.1109/ TBME.2025.3608207 [30] Kokhlikyan, N., Miglani, V., Martin, M., Wang, E., Alsallakh, B., Reynolds, J., Melnikov, A., Kliushkina, N., Araya, C., Yan, S., Reblitz-Richardson, O.: Captum: A Unified and Generic Model Interpretability Library for PyTorch. arXiv (2020). https://doi.org/10.48550/arXiv.2009.07896 [31] Fisher, A., Rudin, C., Dominici, F.: All Models are Wrong, but Many are Useful: Learning a Variable’s Importance by Studying an Entire Class of Prediction Models Simultaneously. Journal of Machine Learning Research : JMLR 20, 177 (2019) [32] Gawlikowski, J., Tassi, C.R.N., Ali, M., Lee, J., Humt, M., Feng, J., Kruspe, A., Triebel, R., Jung, P., Roscher, R., Shahzad, M., Yang, W., Bamler, R., Zhu, X.X.: A Survey of Uncertainty in Deep Neural Networks. Artificial Intelligence Review 56(1), 1513–1589 (2023) https://doi.org/10.1007/s10462-023-10562-9 [33] He, W., Jiang, Z.: A Survey on Uncertainty Quantification for Deep Learning: An Uncertainty Source Perspective. arXiv (2024). https://doi.org/10.48550/arXiv. 2302.13425 [34] Hendrycks, D., Gimpel, K.: A Baseline for Detecting Misclassified and Outof-Distribution Examples in Neural Networks. Proceedings of International Conference on Learning Representations (2017) [35] Tornetta, G.N.: Entropy Methods for the Confidence Assessment of Probabilistic Classification Models. arXiv (2021). https://doi.org/10.48550/arXiv.2103.15157 [36] Taha, A.A., Hennig, L., Knoth, P.: Confidence Estimation of Classification Based on the Distribution of the Neural Network Output Layer. arXiv (2022). https: //doi.org/10.48550/arXiv.2210.07745 [37] Pearce, T., Brintrup, A., Zhu, J.: Understanding Softmax Confidence and Uncertainty. arXiv (2021). https://doi.org/10.48550/arXiv.2106.04972 [38] Chen, T., Guestrin, C.: XGBoost: A Scalable Tree Boosting System. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 785–794. ACM, San Francisco, CA, USA (2016). https://doi.org/10.1145/2939672.2939785

[39] Ke, G., Meng, Q., Finley, T., Wang, T., Chen, W., Ma, W., Ye, Q., Liu, T.-Y.: LightGBM: A Highly Efficient Gradient Boosting Decision Tree. In: Advances in Neural Information Processing Systems 30 (NeurIPS 2017), vol. 30. Curran Associates, Inc., USA (2017) [40] Breiman, L.: Random Forests. Machine Learning 45(1), 5–32 (2001) https://doi. org/10.1023/A:1010933404324 [41] Samad, M.D., Akhter, K.F.B., Rabbani, S.B., Kowsar, I.: Imputation-Free Learning of Tabular Data with Missing Values Using Incremental Feature Partitions in Transformer. arXiv (2025). https://doi.org/10.48550/arXiv.2504.14610 [42] Biewald, L.: Experiment Tracking with Weights and Biases. Software available from wandb.com (2020). https://www.wandb.com/

Appendix A

Supplementary information

Table A1: Performance of models trained on ADNI for CN vs. AD classification and evaluated on the ADNI train, validation, and test splits, as well as on the external OASIS-3 and AIBL cohorts. Models are sorted by AUC within each split. Seed = 123. Data/Split

Model

Acc.

Bal. Acc.

Spec.

Sens.

AUC

ADNI (Train)

XGBoost LightGBM MA-GBT Random Forest NITROGEN NAIM MA-Lasso

1.000 1.000 1.000 0.986 0.914 0.917 0.910

1.000 1.000 1.000 0.985 0.908 0.912 0.908

1.000 1.000 1.000 0.994 0.952 0.951 0.924

1.000 1.000 1.000 0.975 0.863 0.873 0.893

1.000 1.000 1.000 0.984 0.896 0.901 0.896

1.000 1.000 1.000 0.999 0.972 0.969 0.964

ADNI (Val)

NITROGEN XGBoost LightGBM NAIM MA-GBT MA-Lasso Random Forest

0.862 0.856 0.848 0.865 0.844 0.854 0.814

0.853 0.853 0.845 0.859 0.841 0.850 0.808

0.899 0.870 0.863 0.889 0.853 0.868 0.840

0.807 0.835 0.827 0.830 0.830 0.832 0.777

0.826 0.825 0.815 0.833 0.811 0.822 0.772

0.929 0.925 0.924 0.919 0.916 0.910 0.894

ADNI (Test)

LightGBM Random Forest MA-GBT XGBoost NITROGEN NAIM MA-Lasso

0.854 0.863 0.856 0.861 0.846 0.826 0.822

0.846 0.852 0.848 0.854 0.841 0.824 0.826

0.875 0.896 0.879 0.881 0.862 0.831 0.812

0.817 0.808 0.817 0.827 0.821 0.817 0.840

0.807 0.816 0.810 0.816 0.800 0.779 0.780

0.925 0.925 0.923 0.922 0.915 0.912 0.909

AIBL

NITROGEN MA-Lasso LightGBM XGBoost NAIM MA-GBT Random Forest

0.894 0.902 0.911 0.917 0.905 0.913 0.897

0.811 0.840 0.786 0.782 0.812 0.806 0.765

0.926 0.926 0.960 0.970 0.941 0.955 0.949

0.697 0.755 0.613 0.594 0.684 0.658 0.581

0.649 0.684 0.660 0.669 0.669 0.680 0.614

0.936 0.931 0.929 0.923 0.919 0.917 0.913

OASIS-3

MA-GBT NITROGEN NAIM XGBoost LightGBM MA-Lasso Random Forest

0.716 0.560 0.639 0.542 0.561 0.572 0.456

0.673 0.638 0.642 0.630 0.621 0.622 0.573

0.735 0.526 0.638 0.504 0.535 0.550 0.405

0.611 0.751 0.646 0.756 0.706 0.693 0.741

0.395 0.341 0.352 0.334 0.328 0.329 0.292

0.738 0.718 0.714 0.714 0.701 0.697 0.667

Table A2: Performance of models trained on ADNI for CN vs. AD classification on ADNI test set, AIBL, and OASIS-3 cohorts, reported as mean ± standard deviation across 10 random seeds (seeds = [123, 42, 0, 7, 13, 21, 37, 55, 77, 99]). Models are sorted by AUC within each split. Data/Split

Model

Bal. Acc.

Spec.

Sens.

AUC

ADNI (Test)

LightGBM XGBoost NITROGEN Random Forest NAIM MA-Lasso

0.869 ± 0.018 0.864 ± 0.019 0.848 ± 0.023 0.851 ± 0.025 0.846 ± 0.024 0.839 ± 0.024

0.861 ± 0.020 0.857 ± 0.019 0.845 ± 0.022 0.841 ± 0.025 0.839 ± 0.024 0.837 ± 0.023

0.895 ± 0.031 0.887 ± 0.032 0.858 ± 0.038 0.888 ± 0.037 0.870 ± 0.036 0.841 ± 0.038

0.827 ± 0.041 0.827 ± 0.034 0.832 ± 0.035 0.794 ± 0.039 0.808 ± 0.037 0.833 ± 0.037

0.829 ± 0.030 0.824 ± 0.030 0.809 ± 0.032 0.806 ± 0.036 0.802 ± 0.037 0.800 ± 0.035

0.935 ± 0.014 0.933 ± 0.012 0.924 ± 0.016 0.916 ± 0.020 0.915 ± 0.019 0.913 ± 0.018

AIBL

XGBoost MA-Lasso LightGBM Random Forest NAIM NITROGEN

0.917 ± 0.004 0.894 ± 0.011 0.911 ± 0.006 0.907 ± 0.002 0.890 ± 0.010 0.886 ± 0.020

0.789 ± 0.012 0.838 ± 0.013 0.800 ± 0.014 0.771 ± 0.009 0.809 ± 0.009 0.807 ± 0.023

0.967 ± 0.004 0.916 ± 0.014 0.954 ± 0.008 0.960 ± 0.005 0.922 ± 0.017 0.916 ± 0.033

0.610 ± 0.025 0.761 ± 0.030 0.645 ± 0.031 0.582 ± 0.022 0.695 ± 0.032 0.697 ± 0.072

0.674 ± 0.018 0.669 ± 0.024 0.670 ± 0.020 0.637 ± 0.009 0.641 ± 0.014 0.633 ± 0.026

0.918 ± 0.005 0.917 ± 0.008 0.916 ± 0.008 0.913 ± 0.003 0.905 ± 0.012 0.901 ± 0.011

OASIS-3

LightGBM XGBoost NITROGEN NAIM Random Forest MA-Lasso

0.613 ± 0.050 0.551 ± 0.074 0.556 ± 0.060 0.490 ± 0.056 0.610 ± 0.168 0.593 ± 0.021

0.660 ± 0.032 0.641 ± 0.039 0.636 ± 0.025 0.602 ± 0.026 0.621 ± 0.048 0.616 ± 0.019

0.593 ± 0.060 0.512 ± 0.090 0.521 ± 0.078 0.442 ± 0.071 0.606 ± 0.222 0.583 ± 0.029

0.727 ± 0.036 0.770 ± 0.028 0.751 ± 0.054 0.763 ± 0.033 0.636 ± 0.140 0.650 ± 0.047

0.365 ± 0.034 0.345 ± 0.036 0.341 ± 0.023 0.314 ± 0.021 0.349 ± 0.058 0.326 ± 0.016

0.719 ± 0.020 0.713 ± 0.031 0.710 ± 0.018 0.703 ± 0.012 0.700 ± 0.031 0.689 ± 0.017

Acc.

Table A3: Performance of models trained on ADNI for cognitively normal (CN) vs. cognitively impaired (MCI/AD) classification and evaluated across datasets. Models are sorted by AUC within each split. Seed = 123. Data/Split

Model

Acc.

Bal. Acc.

Spec.

Sens.

AUC

ADNI (Train)

XGBoost LightGBM MA-GBT Random Forest NAIM MA-Lasso NITROGEN

1.000 1.000 0.998 0.992 0.825 0.780 0.781

1.000 1.000 0.998 0.987 0.773 0.770 0.706

1.000 1.000 0.996 0.974 0.640 0.746 0.514

1.000 1.000 0.999 1.000 0.906 0.795 0.899

1.000 1.000 0.999 0.994 0.878 0.834 0.851

1.000 1.000 1.000 1.000 0.880 0.862 0.846

ADNI (Val)

LightGBM MA-GBT XGBoost Random Forest NITROGEN NAIM MA-Lasso

0.744 0.721 0.718 0.700 0.703 0.702 0.658

0.677 0.655 0.647 0.584 0.633 0.646 0.645

0.459 0.442 0.416 0.214 0.409 0.467 0.602

0.894 0.868 0.877 0.955 0.857 0.826 0.687

0.821 0.803 0.803 0.806 0.791 0.784 0.725

0.810 0.795 0.781 0.769 0.761 0.749 0.725

ADNI (Test)

MA-GBT LightGBM NITROGEN XGBoost Random Forest NAIM MA-Lasso

0.741 0.744 0.727 0.733 0.706 0.731 0.681

0.668 0.672 0.648 0.662 0.581 0.668 0.660

0.457 0.463 0.418 0.459 0.221 0.488 0.599

0.880 0.881 0.877 0.866 0.942 0.849 0.721

0.821 0.822 0.812 0.813 0.812 0.809 0.753

0.802 0.799 0.794 0.790 0.785 0.781 0.748

AIBL

Random Forest NITROGEN LightGBM MA-GBT XGBoost NAIM MA-Lasso

0.400 0.453 0.529 0.519 0.570 0.510 0.588

0.575 0.591 0.632 0.627 0.636 0.622 0.635

0.206 0.300 0.414 0.399 0.496 0.385 0.535

0.944 0.882 0.850 0.855 0.776 0.858 0.735

0.454 0.460 0.487 0.484 0.487 0.480 0.484

0.766 0.757 0.744 0.738 0.733 0.732 0.709

OASIS-3

LightGBM MA-Lasso NITROGEN Random Forest XGBoost MA-GBT NAIM

0.288 0.383 0.250 0.215 0.276 0.269 0.304

0.545 0.577 0.536 0.516 0.548 0.543 0.561

0.159 0.286 0.107 0.064 0.140 0.132 0.175

0.931 0.868 0.964 0.969 0.955 0.955 0.946

0.304 0.320 0.301 0.292 0.306 0.304 0.312

0.700 0.697 0.690 0.668 0.661 0.660 0.639

Table A4: Global classification performance of models trained on ADNI for three-class diagnostic classification (CN / MCI / AD), evaluated across datasets. ADNI comprised 7,858 samples divided into stratified 60/20/20 train/validation/test splits (Train: 4,704; Val: 1,561; Test: 1,593). OASIS-3 (n=2,675) and AIBL (n=1,286) were used as held-out external validation cohorts. Metrics reported are balanced accuracy (Bal. Acc.), macro-averaged F1, macro-averaged AUC under the one-vs-rest (OvR) scheme, Matthews Correlation Coefficient (MCC), and Cohen’s κ. Macro-averaged metrics weight each class equally regardless of class frequency, which is appropriate given the class imbalance present across cohorts. Models are sorted by Macro AUC within each split. Seed = 123. Full per-class metrics are reported in Table A5. Split

Model

Accuracy

Balanced Acc.

Macro F1

Macro AUC (OvR)

MCC

Cohen

ADNI (Train)

XGBoost LightGBM Random Forest NITROGEN NAIM

1.000 1.000 0.982 0.692 0.626

1.000 1.000 0.976 0.681 0.601

1.000 1.000 0.981 0.691 0.615

1.000 1.000 0.999 0.861 0.801

1.000 1.000 0.972 0.511 0.399

1.000 1.000 0.972 0.509 0.393

ADNI (Val)

XGBoost LightGBM Random Forest NITROGEN NAIM

0.582 0.596 0.570 0.541 0.523

0.580 0.594 0.550 0.543 0.519

0.589 0.602 0.557 0.552 0.526

0.782 0.778 0.765 0.737 0.719

0.342 0.366 0.326 0.277 0.247

0.339 0.362 0.304 0.276 0.242

ADNI (Test)

XGBoost LightGBM Random Forest NITROGEN NAIM

0.591 0.588 0.560 0.556 0.504

0.573 0.579 0.524 0.543 0.478

0.588 0.590 0.534 0.556 0.492

0.771 0.762 0.762 0.752 0.707

0.351 0.350 0.302 0.295 0.203

0.345 0.347 0.278 0.292 0.196

OASIS-3

LightGBM Random Forest NAIM NITROGEN XGBoost

0.265 0.116 0.209 0.233 0.206

0.383 0.340 0.331 0.369 0.339

0.247 0.150 0.177 0.247 0.209

0.596 0.595 0.594 0.583 0.574

0.117 0.081 0.052 0.123 0.091

0.062 0.024 0.026 0.056 0.044

AIBL

NAIM Random Forest LightGBM XGBoost NITROGEN

0.394 0.374 0.467 0.473 0.363

0.479 0.468 0.492 0.470 0.463

0.410 0.395 0.454 0.442 0.395

0.729 0.728 0.726 0.725 0.719

0.195 0.189 0.209 0.192 0.169

0.127 0.115 0.158 0.149 0.104

Table A5: Per-class classification performance (Recall or sensitivity, specificity, and AUC) for three-class diagnostic classification (CN / MCI / AD), evaluated across datasets. Sensitivity (Sens.) corresponds to the true positive rate for each class; specificity (Spec.) corresponds to the true negative rate. AUC is reported per class under the one-vs-rest (OvR) scheme. Models are sorted by Macro AUC within each split (same order as Table A4). Seed = 123. Global metrics are reported in Table A4. CN

MCI

Split

Model

Rec.

Spec.

AUC

Rec.

Spec.

AUC

Rec.

Spec.

AUC

ADNI (Train)

XGBoost LightGBM Random Forest NITROGEN NAIM

1.000 1.000 0.989 0.648 0.510

1.000 1.000 0.996 0.882 0.870

1.000 1.000 1.000 0.876 0.803

1.000 1.000 0.998 0.740 0.736

1.000 1.000 0.973 0.674 0.569

1.000 1.000 1.000 0.799 0.713

1.000 1.000 0.941 0.654 0.555

1.000 1.000 0.999 0.933 0.931

1.000 1.000 0.998 0.907 0.886

ADNI (Val)

XGBoost LightGBM Random Forest NITROGEN NAIM

0.496 0.498 0.342 0.513 0.390

0.828 0.848 0.910 0.778 0.814

0.798 0.802 0.781 0.748 0.735

0.635 0.660 0.765 0.550 0.615

0.572 0.578 0.439 0.557 0.496

0.670 0.667 0.646 0.589 0.565

0.609 0.624 0.542 0.566 0.551

0.918 0.915 0.930 0.918 0.910

0.877 0.866 0.869 0.874 0.856

ADNI (Test)

XGBoost LightGBM Random Forest NITROGEN NAIM

0.516 0.532 0.332 0.478 0.382

0.854 0.833 0.918 0.809 0.825

0.799 0.781 0.788 0.775 0.744

0.680 0.644 0.779 0.631 0.642

0.541 0.573 0.404 0.535 0.423

0.654 0.647 0.643 0.619 0.547

0.523 0.562 0.460 0.521 0.411

0.928 0.919 0.931 0.926 0.926

0.859 0.859 0.856 0.863 0.830

OASIS-3

LightGBM Random Forest NAIM NITROGEN XGBoost

0.223 0.092 0.136 0.201 0.154

0.895 0.951 0.937 0.899 0.911

0.677 0.684 0.686 0.680 0.653

0.449 0.755 0.245 0.531 0.388

0.454 0.139 0.648 0.323 0.406

0.429 0.361 0.441 0.378 0.380

0.477 0.173 0.613 0.374 0.475

0.787 0.965 0.489 0.899 0.771

0.684 0.740 0.656 0.691 0.690

AIBL

NAIM Random Forest LightGBM XGBoost NITROGEN

0.332 0.304 0.447 0.470 0.291

0.873 0.888 0.805 0.785 0.873

0.750 0.762 0.724 0.728 0.745

0.745 0.777 0.630 0.598 0.723

0.351 0.311 0.460 0.467 0.312

0.555 0.537 0.557 0.543 0.505

0.361 0.323 0.400 0.342 0.374

0.981 0.993 0.979 0.984 0.984

0.881 0.886 0.896 0.905 0.907

(a) ADNI test subset

(b) OASIS-3 test set

Fig. A1: Classification performance in multi-class cognitive group classification. ROC curves (One vs. Rest) for multi-class classification on (a) ADNI test set, (b) OASIS-3 and (c) AIBL. For each test set, the plots represent the ROC curves for the CN, MCI, AD vs. the rest and macro average from left to right. AUC values are reported in the legends. All models trained on ADNI.

Table A6: Performance of models trained on ADNI for regression of MMSE and CDR-SB scores, evaluated across datasets. ADNI comprised 7,858 samples (CN: 2,497, MCI: 3,583, AD: 1,778) divided into 60/20/20 train/validation/test splits. OASIS-3 (CN: 2,228, MCI: 49, AD: 398) and AIBL (CN: 947, MCI: 184, AD: 155) were used as held-out external validation cohorts. Within each split, models are sorted by MMSE R2 . Seed = 123. MMSE 2

CDR-SB

Data/Split

Model

MSE

MAE

MSE

MAE

ADNI (Train)

XGBoost Random Forest NITROGEN NAIM

0.000 3.075 6.721 7.151

0.000 1.159 1.843 1.876

1.000 0.788 0.536 0.507

1.000 0.918 0.733 0.714

0 0 0 0

0.000 1.626 3.079 3.121

0.000 0.828 1.185 1.223

1.000 0.736 0.500 0.493

1.000 0.892 0.713 0.703

0 0 0 0

ADNI (Val)

NITROGEN Random Forest NAIM XGBoost

7.745 8.015 8.389 9.049

2.027 2.076 2.081 2.239

0.446 0.427 0.400 0.353

0.670 0.661 0.636 0.617

1.18e−203 6.46e−197 5.69e−178 2.68e−164

3.586 3.633 3.815 3.811

1.309 1.363 1.384 1.384

0.415 0.407 0.377 0.378

0.648 0.650 0.616 0.632

1.10e−186 4.09e−188 1.41e−163 3.58e−175

ADNI (Test)

NAIM NITROGEN Random Forest XGBoost

8.206 8.216 8.417 9.450

2.007 1.999 2.069 2.205

0.383 0.382 0.367 0.289

0.627 0.621 0.610 0.569

1.58e−174 1.67e−170 8.87e−163 2.41e−137

3.486 3.517 3.578 3.853

1.337 1.302 1.352 1.409

0.402 0.397 0.386 0.339

0.635 0.641 0.633 0.604

9.47e−181 3.19e−185 6.63e−179 1.10e−158

AIBL

NITROGEN NAIM Random Forest XGBoost

8.255 8.571 8.744 9.196

1.727 1.775 1.767 1.871

0.384 0.360 0.347 0.313

0.664 0.635 0.666 0.579

2.67e−164 4.06e−146 2.45e−165 3.39e−116

1.551 2.402 1.642 2.127

0.901 1.246 1.095 1.076

-8.771 -14.136 -9.344 -12.405

0.603 0.554 0.656 0.568

2.87e−128 1.87e−104 2.91e−159 8.88e−111

OASIS-3

Random Forest NAIM XGBoost NITROGEN

9.738 11.928 13.555 16.326

2.426 2.770 2.270 3.004

-0.637 -1.005 -1.278 -1.744

0.285 0.303 0.278 0.245

3.02e−51 5.02e−58 8.68e−49 5.36e−38

4.797 7.896 7.253 5.725

1.879 2.456 1.692 2.079

-62.886 -104.163 -95.592 -75.252

0.284 0.247 0.263 0.255

6.31e−51 2.24e−38 1.09e−43 4.77e−41

ADNI (Test)

AIBL

OASIS-3

Fig. A2: Regression performance across models and datasets. Predicted versus observed values for MMSE and CDR-SB scores across test datasets (columns: ADNI, AIBL, OASIS) and models (rows). Each point represents an individual, coloured by diagnostic group (CN: green, MCI: purple, AD: red). The dashed line indicates perfect regression, and the solid line shows the least-squares fit. Pearson correlation coefficients (r) are reported in each panel. Models were trained on ADNI and evaluated on internal (ADNI) and external (AIBL, OASIS) cohorts.

(a) ADNI validation subset

(b) ADNI test subset

(d) AIBL test set

Fig. A3: Predicted class probabilities across internal and external datasets. Bar plots show mean predicted class probabilities across models, with error bars representing one standard deviation across samples (not shown for the correct–incorrect confidence gap, which reflects a derived difference between two conditions). Within each dataset, results are displayed for all samples, stratified by predicted class (CN and AD), by correctness of prediction (correct vs. incorrect), and by the correct–incorrect confidence gap. The confidence gap is defined as the difference between the mean predicted probability of correctly classified samples and that of incorrectly classified samples.

Fig. A4: Prediction probabilities gap across models and test datasets. Scatter plot summarising the prediction confidence gap, defined as the difference between the mean predicted probability of correctly classified samples and that of incorrectly classified samples (correct − incorrect). Each point represents a model–dataset pair. Colours indicate the classification model, and marker shapes denote the evaluation dataset (ADNI, OASIS-3, AIBL). Higher values indicate stronger separation between correct and incorrect predictions, reflecting improved confidence discrimination.

(a) ADNI validation set

(b) ADNI test set

(d) External AIBL cohort

Fig. A5: Calibration performance across internal and external datasets. Reliability diagrams showing model calibration for the positive class (AD). The mean predicted probability within each bin (x-axis) is plotted against the observed fraction of positives (y-axis). Colours correspond to different classification models. The dotted black diagonal represents perfect calibration. Models were trained on ADNI and evaluated on the validation set (a), ADNI test set (b), OASIS-3 (c), and AIBL (d).

Table A7: Sufficiency analysis using the identified features (APOEε4 and left temporal pole cortical thickness). Performance is reported for three input configurations: (i) full multimodal input (Base), (ii) only the identified MRI feature (LH Limbic TempPole 3) with all other MRI modalities removed (set to NA) (MRI), and (iii) only the MRI feature with all others masked (ALL). Metrics include accuracy (Acc.), balanced accuracy (Bal. Acc.), specificity (Spec.), sensitivity (Sens.), and area under the ROC curve (AUC). Seed = 123. Dataset

ADNI (Test)

AIBL

OASIS-3

Modality Model

Acc. Bal. Acc. Spec. Sens. AUC

Base

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.87 0.85 0.87 0.85 0.85 0.82 0.83

0.86 0.84 0.86 0.84 0.84 0.83 0.82

0.88 0.86 0.88 0.86 0.88 0.80 0.85

0.85 0.82 0.84 0.82 0.81 0.86 0.79

0.91 0.91 0.91 0.92 0.92 0.91 0.90

MRI

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.53 0.58 0.65 0.70 0.51 0.57 0.75

0.63 0.66 0.71 0.72 0.61 0.63 0.71

0.30 0.41 0.54 0.65 0.30 0.45 0.84

0.96 0.91 0.87 0.80 0.92 0.81 0.59

0.76 0.77 0.79 0.78 0.77 0.77 0.80

ALL

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.34 0.34 0.53 0.66 0.34 0.72 0.66

0.50 0.50 0.62 0.50 0.50 0.71 0.50

0.00 0.00 0.33 1.00 0.00 0.76 1.00

1.00 1.00 0.92 0.00 1.00 0.66 0.00

0.77 0.77 0.76 0.76 0.57 0.77 0.76

Base

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.90 0.90 0.93 0.93 0.90 0.92 0.93

0.78 0.81 0.79 0.80 0.75 0.88 0.81

0.94 0.93 0.98 0.98 0.95 0.94 0.97

0.62 0.69 0.60 0.62 0.55 0.83 0.64

0.92 0.96 0.94 0.94 0.92 0.97 0.93

MRI

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.38 0.57 0.81 0.71 0.15 0.74 0.87

0.63 0.71 0.60 0.66 0.47 0.75 0.64

0.30 0.53 0.87 0.72 0.05 0.74 0.95

0.95 0.88 0.33 0.60 0.88 0.76 0.33

0.81 0.79 0.73 0.70 0.57 0.81 0.79

ALL

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.12 0.12 0.40 0.88 0.12 0.82 0.88

0.50 0.50 0.65 0.50 0.50 0.81 0.50

0.00 0.00 0.32 1.00 0.00 0.83 1.00

1.00 1.00 0.98 0.00 1.00 0.79 0.00

0.83 0.86 0.86 0.86 0.41 0.86 0.87

Base

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT 51

0.49 0.54 0.51 0.52 0.43 0.55 0.68

0.60 0.65 0.60 0.57 0.56 0.61 0.65

0.44 0.49 0.47 0.49 0.38 0.52 0.69

0.76 0.81 0.73 0.66 0.74 0.70 0.60

0.68 0.72 0.69 0.67 0.64 0.70 0.72

MRI

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.28 0.39 0.33 0.38 0.19 0.46 0.69

0.53 0.59 0.58 0.54 0.51 0.57 0.64

0.18 0.30 0.23 0.31 0.05 0.42 0.71

0.88 0.87 0.93 0.77 0.97 0.72 0.56

0.65 0.65 0.65 0.62 0.59 0.60 0.69

ALL

NAIM NITROGEN XGBoost LightGBM Random Forest MA-Lasso MA-GBT

0.15 0.15 0.28 0.85 0.15 0.76 0.85

0.50 0.50 0.56 0.50 0.50 0.57 0.50

0.00 0.00 0.16 1.00 0.00 0.84 1.00

1.00 1.00 0.96 0.00 1.00 0.30 0.00

0.62 0.55 0.60 0.56 0.51 0.57 0.54

Fig. A6: Grid search results for the sensitivity parameter β at fixed uncertainty threshold τ = 0.900. (Left) Spearman rank correlation ρ between adjusted uncertainty and prediction errors as a function of β , used as the primary optimisation criterion. The optimal value β ∗ = 0.400 is indicated by the vertical red dashed line. (Centre) Accuracy on retained and rejected samples as a function of β , with the baseline accuracy on the full dataset shown as a horizontal dotted line. A well-calibrated β should increase the gap between retained and rejected accuracy. (Right) Coverage (fraction of samples retained) as a function of β , showing how the missingness adjustment progressively shifts samples above the fixed threshold τ as β increases. The optimal configuration (β ∗ = 0.400, τ = 0.900) retains 85.7% of validation samples (756/882), improving accuracy from 0.8617 to 0.8995 (∆ = +0.0378, p < 0.001 by permutation test), with rejected samples exhibiting substantially lower accuracy (0.6349).

Record · ID 363256 · SHA-256 797569916033bd46
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.