PaperlayerView exact revision on arXiv

Source-linked AI summary

M$^2$PFN: End-to-End Disentangled Alignment for Generalizable Multimodal In-Context Learning in Alzheimer's Disease

Lujia Zhong, Shuo Huang, Jianwei Zhang, Xinyu Nie, Yonggang Shi

arXiv:2609.28836v1cs.CVcs.AIcs.LG

TL;DR

Multimodal AD models can struggle to generalize across cohorts, and TabPFN’s tabular prior does not naturally match image-derived features. M2PFN trains multimodal representations through a frozen TabPFN engine with disentangled alignment and a gated tabular shortcut. It reports state-of-the-art ADNI diagnosis and MMSE regression, plus strong transfer to external cohorts without retraining.

  • Problem

    Multimodal AD models face limited cross-cohort generalization, while TabPFN’s synthetic tabular prior does not naturally match image-derived features.

    Source: §1 · pp. 1–2

  • Method

    M2PFN keeps TabPFN frozen while differentiable inference trains MRI and tabular encoders, aligns their representations, and incorporates a gated tabular shortcut.

    Source: §1 · pp. 2–3

  • Results

    M2PFN reports state-of-the-art ADNI performance for three-class diagnosis and MMSE regression, and the best gradient-free adaptation performance on OASIS-3 and SCAN without retraining.

    Source: §1 · p. 3, §5 · p. 15

  • Takeaways & Limitations

    Replacing the support set enables transfer to external cohorts while leaving the frozen ICL engine unchanged.

    Source: §1 · pp. 2–3

Abstract

from arXiv · show

While various multimodal methods combining imaging and tabular data for Alzheimer's disease (AD) diagnosis were proposed, they are often limited in generalization across cohorts. In-context learning (ICL) has demonstrated excellent generalization performances and high flexibility in foundational tabular models such as TabPFN. To extend TabPFN's ICL to multimodal AD analysis, the main obstacle is that TabPFN is meta-trained on synthetic tabular priors that do not naturally match the statistical structure of image-derived features. We propose M$^2$PFN, an end-to-end framework that turns this tabular foundation model into a multimodal AD predictor. M$^2$PFN (i) performs differentiable inference through TabPFN's transformer, back-propagating task gradients into 3D-MRI and tabular encoders; (ii) aligns the two modalities into a shared subspace, via disentanglement and a contrastive objective, matched to the ICL engine's prior; and (iii) folds in a frozen tabular-only prediction through a learnable gated shortcut. Because the ICL engine stays frozen, its in-context mechanism is preserved for test-time generalization, while end-to-end training shapes the encoders into features it can exploit. On ADNI ($n=2240$, three-class CN/MCI/AD), M$^2$PFN attains $65.55\%$ macro-F1 and $82.21\%$ macro-AUC, surpassing a comprehensive set of unimodal and multimodal baselines. By swapping only the head for a TabPFN regressor, the same architecture regresses baseline MMSE on a $1250$-subject sub-cohort to test MAE $1.743$, outperforming every multimodal baseline. On two external cohorts (OASIS-3 and SCAN) with no retraining, M$^2$PFN achieves the best AUC and the lowest MMSE MAE across all baselines, and transfers even when the cognitive instrument changes.

1 Introduction

Multimodal AD models can struggle to transfer across cohorts with different imaging protocols and available tabular features, while tabular ICL models cannot directly use imaging. M2PFN adapts a frozen TabPFN through differentiable multimodal learning and reports strong ADNI performance and external-cohort transfer.

Motivation

  • AD multimodal models often require retraining or extensive fine-tuning across cohorts with different imaging protocols, scanners, populations, and tabular feature availability.Tabular inputs range from more than 100 variables in ADNI to fewer than 10 in some external cohorts.Source: §1 · p. 2
  • Task-specific AD models provide limited support for flexible multimodal integration and robust cross-cohort deployment because their parameters are optimized for training distributions.The passage identifies CNNs, transformers, and mixture-of-experts models as examples.Source: §1 · p. 2
  • TabPFN’s synthetic tabular priors do not naturally match image-derived features, and its non-differentiable preprocessing prevents joint encoder training.The passage argues that end-to-end training is needed to shape image embeddings for the engine’s prior.Source: §1 · p. 2

Contributions

  • M2PFN keeps TabPFN frozen while optimizing multimodal representations through its in-context engine, allowing support-set replacement for transfer without retraining.Its contributions include differentiable prompt tuning and disentangled alignment of MRI and tabular features.Source: §1 · pp. 2–3
  • M2PFN achieves state-of-the-art performance on ADNI for both three-class CN/MCI/AD diagnosis and MMSE regression.The framework also reports robust zero-retraining transfer to OASIS-3 and SCAN, with the best performance among gradient-free adaptation baselines.Source: §1 · p. 3

2 Related Work

Prior AD multimodal models combine MRI and tabular data but remain parametric, limiting adaptation to new cohorts without retraining. TabPFN offers in-context prediction for tabular data, while M2PFN adapts its frozen engine by directing gradients into the modality encoders.

2.1 Multimodal Alzheimer’s Disease Analysis

  • Multimodal AD models combine MRI and clinical/tabular data, but their fixed parameters cannot reference labeled exemplars at inference or effectively adapt to new cohorts without retraining or fine-tuning.Prior work spans cross-modal transformers and mixture-of-experts architectures.Source: §2.1 · p. 3

2.2 In-Context Learning

  • TabPFN performs well on small-to-medium tabular datasets without test-time optimization, but its structured-input design and non-differentiable preprocessing exclude MRI and block gradients to imaging encoders.The section places TabPFN among prior-data fitted networks that meta-learn Bayesian inference over synthetic task distributions.Source: §2.2 · p. 3

2.3 Prompt Tuning for ICL Engines

  • M2PFN keeps the ICL engine frozen and sends task gradients through TabPFN’s transformer into imaging and tabular encoders, unlike MMPFN’s frozen precomputed image embeddings.The cited comparison notes that MMPFN lightly fine-tunes a mixer but cannot adapt the imaging encoder end to end.Source: §2.3 · p. 4

3 Method

M2PFN uses frozen TabPFN as an in-context prediction head while learning multimodal representations for its inputs. The method combines shared and modality-specific representations, differentiable task optimization, and a gated tabular shortcut.

3.1 Problem Formulation

  • Replacing the support set at inference allows the frozen TabPFN to predict new queries without updating its parameters.Gradients through the query path optimize the modality encoders, disentangler, projection, and fusion gate.Source: §3.1 · p. 5, §3.3 · p. 6

3.2 Alignment-Based Multimodal Representation Learning

  • The visual and tabular encoders produce modality representations that two parallel heads transform into shared and modality-specific representations, which are concatenated following the ordering in Fig. 1(a).The visual branch uses a MONAI [33] 3D ResNet-18 [10] initialized from MedicalNet [34], while the tabular branch uses an FT-Transformer [35].Source: §3.2 · pp. 5–6
  • Class-prototype alignment brings shared MRI and tabular representations together, while modality-specific representations remain optimized by the downstream task loss.The alignment targets are bounded by the number of diagnostic categories, rather than the number of subjects.Source: §3.2 · pp. 5–6

3.3 Differentiable ICL Classification and Regression

  • The projected multimodal representation enters the frozen TabPFN, whose differentiable transformer path passes task gradients back to the trainable front end.For regression, the TabPFN classifier is replaced with a regressor; its posterior mean is the multimodal prediction.Source: §3.3 · p. 6

3.4 Gated TabPFN Shortcut

  • A learnable scalar gate blends the multimodal prediction with a frozen tabular-only TabPFN prediction.The gate is initialized to an equal blend, and the task loss updates the gate while gradients pass through only the multimodal branch.Source: §3.4 · p. 7

3.5 Training Objective

  • Training optimizes the encoders, disentangler heads, projection, and gate while keeping all TabPFN parameters fixed.The stated alignment weight is 0.1 and β is 1 unless stated otherwise.Source: §3.5 · p. 7

4 Experiments

Experiments use ADNI for training and evaluation, with OASIS-3 and SCAN held out as external test cohorts. Cohorts vary in feature availability, and the evaluation covers diagnosis and cognitive-score regression.

4.1 Dataset and Preprocessing

  • ADNI supplies 172 tabular features, while OASIS-3 and SCAN supply 7 and 27, respectively; all cohorts use processed T1-weighted MRI.Available features are mapped into ADNI slots, with continuous values mean-imputed and categorical missingness assigned a reserved index.Source: §4.1 · p. 8
  • All models are trained exclusively on ADNI, with OASIS-3 and SCAN held out as external test sets.The ADNI classification cohort contains 2,240 subjects; regression includes 1,250 ADNI subjects with valid baseline cognitive scores.Source: §4.1 · p. 8

4.2 Implementation Details

  • Classification is evaluated with accuracy, macro-F1, macro-AUC, sensitivity, and specificity; regression uses MAE, Pearson r, and R2.Results are reported on held-out test sets; experiments use five random seeds and report mean ± standard deviation.Source: §4.2 · pp. 8–9

4.3 Baselines

The comparison includes unimodal MRI and tabular models alongside multimodal fusion methods. Baselines use the same data splits, with a shared 3D image encoder supplied to applicable multimodal models for comparison.

4.3 Baselines

  • Multimodal baselines lacking a released 3D image encoder receive the same MedicalNet-pretrained 3D ResNet-18 used by M2PFN.Their own tabular encoders are retained, and the image encoder is fine-tuned end-to-end.Source: §4.3 · p. 9
  • MRI baselines include ResNet-3D, DenseNet-3D, SE-ResNet-3D, and SwinUNETR, while tabular baselines include XGBoost, AutoGluon, and TabPFN-v2.5.All baselines are trained and evaluated on the same splits.Source: §4.3 · pp. 9–10
  • Multimodal baselines span cross-modal and mixture-of-experts fusion methods, including InterpretCC, MoE++, SwitchGate, mmFormer, MulT, Flex-MoE, I2MoE, and MMPFN.Methods retain their standard public implementations and published-code default hyperparameters.Source: §4.3 · pp. 9–10

4.4 Classification on ADNI

On ADNI three-way diagnosis, M2PFN outperforms the evaluated unimodal and multimodal baselines across all reported metrics. Its improvement over MMPFN is statistically significant.

4.4 Classification on ADNI

  • M2PFN achieves 65.55% macro-F1 and 82.21% macro-AUC, the best performance across all reported metrics and above both unimodal and multimodal baselines.MMPFN is the strongest multimodal competitor; its improvement over MMPFN is statistically significant.Source: §4.4 · p. 10

4.5 Ablation Study

The ablation study finds that preserving TabPFN’s pretrained in-context learning capabilities contributes most to M2PFN’s performance. End-to-end optimization and representation alignment are complementary, while disentangled alignment and the frozen shortcut also improve results.

4.5 Ablation Study

  • Replacing the frozen TabPFN heads with trainable MLP classifiers produces the largest performance drop.The study identifies preserving TabPFN’s pretrained in-context learning capabilities as the largest source of M2PFN’s performance gain.Source: §4.5 · p. 10
  • Jointly removing end-to-end optimization and representation alignment degrades performance more than removing either component alone, indicating complementarity.End-to-end optimization lets gradients from the frozen ICL engine guide multimodal representation learning.Source: §4.5 · p. 10
  • Removing disentangled alignment causes a modest but significant reduction in AUC (p < 0.05) and reduces training stability.The alignment module consistently improves predictive performance and training stability.Source: §4.5 · p. 10
  • The frozen TabPFN shortcut encourages the multimodal branch to learn complementary information beyond the tabular features, yielding richer representations and better performance.Source: §4.5 · p. 10

4.6 Subject-level Interpretability

M2PFN exposes the support subjects most associated with each query through branch-weighted attention attribution. Case examples and support-removal tests examine how these neighborhoods relate to predictions.

4.6 Subject-level Interpretability

  • Figure 2 compares the top-five support subjects retrieved by the multimodal and shortcut branches for five test queries, with bar length showing normalized attention and color showing diagnosis.In Cases B–D, the branches retrieve different diagnostic evidence that their gated combination uses in the predictions.Source: Figure 2 · p. 11, §4.6 · p. 11
  • Unified attribution combines attention from both branches to represent each support subject’s relative attribution to a query.Both branches use the same indexed support set; their attention distributions are extracted across TabPFN layers and heads, averaged, and renormalized.Source: §4.6 · pp. 10–11
  • Case E is a failure case: both branches mainly retrieve AD supports for an MCI query, producing an AD prediction.The authors describe the query and its highly attended supports as having AD-like clinical profiles.Source: §4.6 · p. 12
  • Removing the top-800 attended supports reduces macro-AUC from 0.822 to 0.344, while removing the same number at random leaves macro-AUC at 0.780.Across all 351 ADNI test queries, the deletion curve shows faster macro-AUC decline for top-attended than random supports.Source: §4.6 · pp. 12–13, Figure 3 · p. 12
  • On a class-stratified 30-query subset, top-50 supports have 4.52±1.29 times the mean leave-one-out influence of random controls, and every query is above unity.Per-query attention and leave-one-out influence have Spearman correlation 0.639 (median 0.647).Source: §4.6 · p. 12

4.7 Generalization to Cognitive Score Regression

For baseline MMSE regression on ADNI, M2PFN outperforms the evaluated multimodal methods. The regression setup retains the classification architecture while replacing the prediction head and task loss.

4.7 Generalization to Cognitive Score Regression

  • The regression setup replaces the TabPFN classifier with a regressor and cross-entropy with smooth-L1 loss, while retaining the disentangler, alignment loss, shortcut, and end-to-end training.Source: §4.7 · p. 13
  • The target is baseline MMSE on its original 0–30 clinical scale, with lower scores indicating greater cognitive impairment.The ADNI test-set regression uses a 1,250-subject MMSE sub-cohort.Source: §4.7 · p. 13
  • M2PFN achieves the lowest MAE, highest Pearson correlation, and highest coefficient of determination, outperforming conventional multimodal methods and MMPFN.It also exhibits lower variance across random seeds.Source: §4.7 · p. 13

4.8 External Validation on OASIS-3 and SCAN

M2PFN transfers diagnosis and cognitive-score prediction to OASIS-3 and SCAN without retraining, with support-set replacement improving cross-cohort and cross-instrument performance. The results distinguish MCI diagnosis and SCAN's MoCA transfer as important tests of generalization.

4.8 External Validation on OASIS-3 and SCAN

  • Evaluation uses frozen ADNI checkpoints without gradient updates, in zero-shot and support-swap modes, with the same labeled support samples and held-out test subjects for all methods.Parametric baselines receive class-prior correction for diagnosis or affine rescaling for regression; ICL methods use the support samples as input.Source: §4.8 · p. 13
  • With support-set replacement, M2PFN achieves the highest MCI AUC and best macro-AUC among parametric learning-based methods on both external cohorts.The diagnosis table reports one-vs-rest AUC by class and macro-average for OASIS-3 and SCAN.Source: Table 5 · p. 14, §4.8 · p. 14
  • M2PFN achieves the lowest MAE on OASIS-3 and SCAN, significantly outperforming all baselines; on SCAN, support replacement with MoCA-labeled examples also substantially improves correlation.OASIS-3 uses MMSE, while SCAN requires transfer to MoCA; rescaling cannot improve Pearson correlation because prediction rankings remain unchanged.Source: Table 6 · p. 14, §4.8 · p. 14
  • Effective cross-instrument transfer relies on transferable multimodal representations that remain compatible with the ICL engine, alongside support-set adaptation.The regression table covers OASIS-3 MMSE and SCAN MoCA.Source: §4.8 · pp. 14–15, Table 6 · p. 14

5 Conclusion

M2PFN repurposes frozen TabPFN for multimodal AD prediction by optimizing representations through its ICL transformer. It reports strong ADNI diagnosis and MMSE regression results, and transfers to independent cohorts without retraining.

5 Conclusion

  • M2PFN achieves state-of-the-art ADNI performance on three-class diagnosis and MMSE regression, consistently outperforming strong multimodal baselines.The frozen ICL engine enables gradient-free support-set replacement for transfer to independent clinical cohorts without retraining.Source: §5 · p. 15
  • The framework keeps TabPFN frozen while learning multimodal representations through its ICL transformer, using disentangled alignment and a gated tabular shortcut.The representation learner bridges heterogeneous MRI and tabular biomarkers while encouraging complementary information learning.Source: §5 · p. 15
  • The authors conclude that combining multimodal representation learning with ICL is an effective paradigm for robust cross-cohort clinical prediction.They present foundation-model-based ICL as a potential alternative to task-specific fine-tuning for multimodal medical AI.Source: §5 · p. 15

Related papers

Loading 2609.28836v1…