Technical Summary: PRA-PoE for Robust Multimodal Alzheimer's Disease Classification

1. Problem Statement

Accurate classification of Alzheimer's Disease (AD) relies on integrating multimodal evidence, including structural MRI (sMRI), PET scans (FDG and Amyloid), and clinical tabular data. However, real-world clinical cohorts suffer from arbitrary modality missingness due to cost, contraindications, and protocol heterogeneity. This results in a "missingness pattern mismatch" where the distribution of observed modality subsets during training differs significantly from those encountered during deployment.

This mismatch induces two primary challenges:

1. Conditional Representation Shift: Models tend to co-adapt to frequent modality subsets, causing the representation of a specific modality to drift depending on which other modalities are present (context-dependent drift). This degrades generalization on rare or unseen modality combinations.
2. Overconfident Reliance on Synthesized Features: Existing imputation or cross-modal translation methods often treat synthesized (imputed) features as if they were fully observed. This lack of explicit uncertainty modeling leads to "blind fusion," where noisy imputations corrupt the aggregation process, causing negative transfer even when strong observed modalities are available.

2. Methodology: PRA-PoE

The authors propose PRA-PoE, a unified framework designed to handle incomplete multimodal inputs through two core components: Prototype-anchored Representation Alignment (PRA) and an Uncertainty-aware Product of Experts (UA-PoE) fusion mechanism.

2.1 Unified Feature Encoding

Heterogeneous inputs (3D neuroimaging and tabular data) are mapped into a shared latent space using modality-specific encoders ($E_m$). Missing modalities are represented as zero placeholders, which are excluded from cross-modal attention mechanisms via an availability mask.

2.2 Prototype-Anchored Representation Alignment (PRA)

To mitigate representation shift, PRA enforces subset-invariant semantics before fusion. It operates via:

• Learnable Global Prototypes: Each modality is anchored to a learnable global prototype ($P_m$) to establish a consistent semantic space.
• Availability-Conditioned Tokens: The framework uses two learnable tokens ($t_{obs}$ and $t_{mis}$) to explicitly encode whether a modality is observed or missing. These tokens are concatenated with the global prototype to form a query ($q_{i,m}$).
• Masked Cross-Attention: A multi-head cross-attention mechanism refines the representation of a target modality using features from only the observed modalities. An additive bias derived from the observation mask ensures missing modalities do not contribute to the attention calculation.
• Adaptive Calibration: For observed modalities, the original features are adaptively fused with the aligned features to reduce subset-induced drift. For missing modalities, features are synthesized solely from the observed context.

2.3 Uncertainty-Aware Product of Experts (UA-PoE)

Following PRA, the framework treats each modality as a Gaussian expert within a Product of Experts (PoE) fusion scheme.

• Gaussian Modeling: Each modality's aligned representation is mapped to a diagonal Gaussian distribution ($N(\mu, \sigma^2)$) via a lightweight head.
• Implicit Uncertainty Down-weighting: Crucially, all $M$ experts participate in the fusion, including those based on PRA-synthesized features. Because synthesized features carry higher epistemic uncertainty, their corresponding experts naturally learn larger variances ($\sigma^2$). In the closed-form PoE fusion, higher variance results in lower precision, automatically down-weighting the contribution of unreliable (imputed) modalities without requiring explicit uncertainty thresholds.
• Prediction: The fused posterior distribution is sampled via Monte Carlo (MC) sampling to generate the final classification probabilities.

2.4 Loss Function

The model is trained using a combination of Cross-Entropy loss and a Kullback-Leibler (KL) divergence penalty against a standard normal prior. This regularization prevents over-confident experts and encourages the fused distribution to remain well-calibrated.

3. Key Contributions

• Explicit Modeling of Availability: Unlike prior methods that rely on implicit imputation, PRA-PoE explicitly encodes modality availability using learnable tokens, distinguishing observed from missing data at the feature level.
• Representation Alignment: The PRA module reduces context-dependent representation shift by anchoring modality embeddings to global prototypes and refining them via availability-conditioned attention.
• Uncertainty-Aware Fusion: The UA-PoE mechanism integrates synthesized features into the fusion process but automatically suppresses their influence through variance-based precision weighting, preventing negative transfer from noisy imputations.
• Clinically Realistic Evaluation Protocol: The framework is evaluated under a protocol that trains on naturally missing data and tests on all non-empty modality combinations, including rare, long-tailed patterns, rather than just complete cases or random dropouts.

4. Experimental Results

The authors evaluated PRA-PoE on two large-scale benchmarks: ADNI (CN vs. MCI vs. AD) and OASIS-3 (CN vs. AD).

• Performance: PRA-PoE consistently outperformed state-of-the-art baselines (including mmFormer, ShareSpec, AnyMod, and various MoE frameworks) across all non-empty modality subsets.
  • On ADNI, it achieved a 5.4% relative improvement in average accuracy and a 4.0% gain in AUC over the strongest baseline (MoE-retriever).
  • On OASIS-3, it achieved a 10.9% relative improvement in average F1 score over the strongest baseline.
• Efficiency: The model is parameter-efficient (13.8M parameters) and computationally lightweight (18.4G FLOPs), with an inference speed of 8.1 ms/sample, making it suitable for low-latency deployment.
• Ablation Studies: Removing PRA caused the largest performance drop (Average ACC dropped from 64.1% to 61.0%), confirming its critical role in mitigating distribution shift. Removing UA-PoE also degraded performance, validating the necessity of uncertainty-aware fusion.

5. Significance and Claims

The paper claims that PRA-PoE addresses a critical gap in multimodal learning for clinical settings: the inability of current models to generalize under arbitrary and heterogeneous missingness patterns. By explicitly modeling modality availability and integrating uncertainty directly into the fusion mechanism, the framework delivers robust predictions even when rare or incomplete data subsets are encountered.

The authors position PRA-PoE as a step toward practical clinical application, noting its ability to handle the "long-tailed" distribution of modality combinations found in real-world cohorts. The work suggests that future extensions could include longitudinal modeling and external multi-center validation to further assess generalizability.