NeuroRAD-FM: A Foundation Model for Neuro-Oncology with Distributionally Robust Training

2025-09-18

AI Summary

In neuro-oncology, data heterogeneity across institutions and the rarity of key molecular biomarkers (e.g., CDKN2A/2B, ATRX) severely impair the generalizability and clinical utility of machine learning models. To address this, we introduce NeuroRAD-FM—the first foundation model specifically designed for neuro-oncological applications—integrating self-supervised learning (BYOL, DINO, MAE, MoCo) with distributionally robust optimization (DRO) to mitigate multi-center data shifts and severe class imbalance. Experiments demonstrate substantial improvements: CDKN2A/2B prediction AUC increases from 0.73 to 0.92; ATRX prediction AUC reaches 0.82. Survival prediction c-index improves significantly across all sites, and inter-site representation divergence is markedly reduced. NeuroRAD-FM establishes a generalizable foundation model paradigm for precise molecular subtyping and prognostic assessment in neuro-oncology.

Technology Category

• Machine Learning: Other Foundations of Machine Learning
• Computer Vision: Other Foundations of Computer Vision
• Cognitive Modeling & Cognitive Systems: Other Foundations of Cognitive Modeling & Systems

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Neuro-oncology poses unique challenges for machine learning due to heterogeneous data and tumor complexity, limiting the ability of foundation models (FMs) to generalize across cohorts. Existing FMs also perform poorly in predicting uncommon molecular markers, which are essential for treatment response and risk stratification. To address these gaps, we developed a neuro-oncology specific FM with a distributionally robust loss function, enabling accurate estimation of tumor phenotypes while maintaining cross-institution generalization. We pretrained self-supervised backbones (BYOL, DINO, MAE, MoCo) on multi-institutional brain tumor MRI and applied distributionally robust optimization (DRO) to mitigate site and class imbalance. Downstream tasks included molecular classification of common markers (MGMT, IDH1, 1p/19q, EGFR), uncommon alterations (ATRX, TP53, CDKN2A/2B, TERT), continuous markers (Ki-67, TP53), and overall survival prediction in IDH1 wild-type glioblastoma at UCSF, UPenn, and CUIMC. Our method improved molecular prediction and reduced site-specific embedding differences. At CUIMC, mean balanced accuracy rose from 0.744 to 0.785 and AUC from 0.656 to 0.676, with the largest gains for underrepresented endpoints (CDKN2A/2B accuracy 0.86 to 0.92, AUC 0.73 to 0.92; ATRX AUC 0.69 to 0.82; Ki-67 accuracy 0.60 to 0.69). For survival, c-index improved at all sites: CUIMC 0.592 to 0.597, UPenn 0.647 to 0.672, UCSF 0.600 to 0.627. Grad-CAM highlighted tumor and peri-tumoral regions, confirming interpretability. Overall, coupling FMs with DRO yields more site-invariant representations, improves prediction of common and uncommon markers, and enhances survival discrimination, underscoring the need for prospective validation and integration of longitudinal and interventional signals to advance precision neuro-oncology.

Problem

• Improving generalization of neuro-oncology foundation models across institutions
• Enhancing prediction accuracy for uncommon molecular markers in brain tumors
• Addressing data heterogeneity and class imbalance with robust training

Innovation

• Distributionally robust optimization for cross-institution generalization
• Self-supervised pretraining of multiple backbone architectures
• Improved prediction of both common and uncommon molecular markers