A Foundational Brain Dynamics Model via Stochastic Optimal Control

TL;DR

BDO tackles the challenge of learning scalable, transferable brain-dynamics representations from noisy fMRI by integrating stochastic optimal control with amortized inference and self-supervised learning in a continuous-discrete state-space model. It adopts a locally linear, simulation-free inference scheme to produce a universal feature for downstream tasks, pre-trained on the UK Biobank and validated across diverse datasets, achieving state-of-the-art results in demographics, trait prediction, and psychiatric diagnosis. The approach offers a principled, efficient foundation model for brain dynamics with robust cross-dataset generalization and interpretable latent structure, enabling broad neuroscience and clinical applications.

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.

Figures (12)

• Figure 1: Conceptual illustration of our proposed Brain Dynamics with Optimal control (BDO). The ROI signals observed at discrete time points are encoded into an optimal control policy, which steers the continuous latent state dynamics. The pre-trained optimal control policy is then utilized for various downstream tasks.
• Figure 2: Our BDO surpasses other foundation models, demonstrating outstanding efficiency. Even the smallest BDO (5M), achieves comparable performance while being significantly efficient in both parameters and resource usage.
• Figure 3: Scalability results of HCP-A age regression in LP.
• Figure 4: BDO captures a latent space that encodes clinically relevant information from fMRI recordings. For each fMRI scan, a universal feature $\mathbb{A}$ is extracted as a summary representation. The $\mathbb{A}$ is then projected into a 2D space using PCA and UMAP. The resulting embedding reveals a structured organization across both internal and external datasets.
• Figure 5: (Left) Training curve (Right) Pearson correlation $\rho$ as the mask ratio $\gamma$ and balancing factor $\tau$ are varied.

Theorems & Definitions (5)

• Proposition 4.1: Evidence lower bound
• Theorem 4.2: Simulation-free inference
• Theorem 1.1: Girsanov Theorem
• proof
• proof