A Foundational Brain Dynamics Model via Stochastic Optimal Control

TL;DR

This paper presents a scalable, robust brain dynamics model using stochastic optimal control and amortized inference, achieving state-of-the-art results in various neuroimaging tasks and demonstrating strong generalization across datasets.

Contribution

It introduces a novel SOC-based brain dynamics model with efficient inference, integrating self-supervised learning for improved representation and scalability.

Findings

Achieves state-of-the-art performance on multiple neuroimaging tasks.

Demonstrates robustness and transferability across diverse datasets.

Provides a scalable framework for understanding brain dynamics.

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.