A Foundational Brain Dynamics Model via Stochastic Optimal Control

2025-02-07

AI Summary

This work addresses key challenges in fMRI signal understanding—namely, high noise levels, substantial inter-subject heterogeneity, and poor generalization across downstream tasks—by introducing the first foundational model for brain dynamics grounded in functional MRI. We propose a continuous-discrete state-space model rooted in stochastic optimal control (SOC), integrating simulation-free latent dynamics with local linear approximations to enable efficient and scalable inference. Crucially, we pioneer the incorporation of SOC principles into an amortized variational inference framework, deriving an SOC-driven evidence lower bound (ELBO) objective that supports self-supervised pretraining and cross-population generalization. Pretrained on the UK Biobank dataset, our model achieves state-of-the-art performance on diverse downstream tasks—including demographic prediction, trait analysis, disease diagnosis, and prognosis. Robustness and transferability are further validated across independent datasets: HCP-A, ABIDE, and ADHD200.

Technology Category

• Machine Learning: Probabilistic Circuits and Graphical Models
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

We introduce a foundational model for brain dynamics that utilizes stochastic optimal control (SOC) and amortized inference. Our method features a continuous-discrete state space model (SSM) that can robustly handle the intricate and noisy nature of fMRI signals. To address computational limitations, we implement an approximation strategy grounded in the SOC framework. Additionally, we present a simulation-free latent dynamics approach that employs locally linear approximations, facilitating efficient and scalable inference. For effective representation learning, we derive an Evidence Lower Bound (ELBO) from the SOC formulation, which integrates smoothly with recent advancements in self-supervised learning (SSL), thereby promoting robust and transferable representations. Pre-trained on extensive datasets such as the UKB, our model attains state-of-the-art results across a variety of downstream tasks, including demographic prediction, trait analysis, disease diagnosis, and prognosis. Moreover, evaluating on external datasets such as HCP-A, ABIDE, and ADHD200 further validates its superior abilities and resilience across different demographic and clinical distributions. Our foundational model provides a scalable and efficient approach for deciphering brain dynamics, opening up numerous applications in neuroscience.

Problem

• Model brain dynamics using stochastic optimal control.
• Handle noisy fMRI signals with continuous-discrete SSM.
• Promote robust, transferable representations via ELBO.

Innovation

• Stochastic Optimal Control
• Continuous-Discrete State Space
• Simulation-Free Latent Dynamics