A Foundational Brain Dynamics Model via Stochastic Optimal Control

📄 arXiv: 2502.04892v1 📥 PDF

作者: Joonhyeong Park, Byoungwoo Park, Chang-Bae Bang, Jungwon Choi, Hyungjin Chung, Byung-Hoon Kim, Juho Lee

分类: cs.LG, q-bio.NC, stat.ML

发布日期: 2025-02-07

备注: The first two authors contributed equally

💡 一句话要点

提出基于随机最优控制的脑动力学基础模型，用于脑疾病诊断与预测。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 脑动力学建模 随机最优控制 fMRI信号处理 深度学习 脑疾病诊断

📋 核心要点

  1. 现有方法难以有效处理fMRI信号的复杂性和噪声，且计算成本高昂，限制了脑动力学模型的应用。
  2. 论文提出基于随机最优控制的脑动力学模型，通过连续-离散状态空间模型和局部线性近似实现高效推断。
  3. 该模型在多个数据集上预训练，并在人口统计学预测、疾病诊断等任务上取得了领先结果，验证了其有效性。

📝 摘要（中文）

本文提出了一种基于随机最优控制（SOC）和分摊推断的脑动力学基础模型。该方法采用连续-离散状态空间模型（SSM），能够稳健地处理fMRI信号的复杂性和噪声。为了解决计算限制，论文在SOC框架下实施了一种近似策略。此外，提出了一种无模拟的潜在动力学方法，该方法采用局部线性近似，从而实现高效且可扩展的推断。为了有效的表征学习，从SOC公式中推导出一个证据下界（ELBO），该下界与自监督学习（SSL）的最新进展顺利集成，从而促进了稳健且可迁移的表征。该模型在UKB等大型数据集上进行预训练，并在包括人口统计学预测、特征分析、疾病诊断和预后等各种下游任务中取得了最先进的结果。在HCP-A、ABIDE和ADHD200等外部数据集上的评估进一步验证了其在不同人口统计学和临床分布中的卓越能力和鲁棒性。该基础模型为解码脑动力学提供了一种可扩展且高效的方法，为神经科学开辟了众多应用。

🔬 方法详解

问题定义：现有脑动力学模型在处理fMRI信号时，难以兼顾信号的复杂性、噪声以及计算效率。传统方法通常需要大量的计算资源，并且泛化能力有限，难以适应不同数据集和任务的需求。因此，如何构建一个既能准确捕捉脑动力学特征，又能高效进行推断和泛化的模型是一个关键问题。

核心思路：本文的核心思路是利用随机最优控制（SOC）框架来建模脑动力学，并结合分摊推断和局部线性近似来提高计算效率。SOC框架能够自然地处理系统中的随机性和不确定性，而分摊推断则可以通过学习一个推断网络来加速后验分布的估计。局部线性近似则进一步简化了潜在动力学的建模，使其更易于处理。

技术框架：该模型采用连续-离散状态空间模型（SSM）来描述脑动力学。整体架构包含以下几个主要模块：1）编码器：将fMRI信号编码为潜在状态的表示；2）状态转移模型：基于SOC框架，描述潜在状态随时间的演化；3）解码器：将潜在状态解码为可观测的fMRI信号；4）推断网络：用于近似后验分布，加速推断过程。整个模型通过最大化证据下界（ELBO）进行训练。

关键创新：该论文的关键创新在于将随机最优控制（SOC）框架引入到脑动力学建模中，并提出了一种无模拟的潜在动力学方法，该方法采用局部线性近似，从而实现高效且可扩展的推断。此外，通过从SOC公式中推导出一个证据下界（ELBO），该下界与自监督学习（SSL）的最新进展顺利集成，从而促进了稳健且可迁移的表征。

关键设计：在状态转移模型中，使用了局部线性近似来简化计算。损失函数采用证据下界（ELBO），包含重构损失、KL散度和控制成本。网络结构方面，编码器和解码器可以使用卷积神经网络或循环神经网络等结构，具体选择取决于fMRI信号的特性和任务需求。关键参数包括状态空间维度、控制成本系数等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

该模型在UKB数据集上预训练后，在人口统计学预测、特征分析、疾病诊断和预后等多个下游任务中取得了最先进的结果。在HCP-A、ABIDE和ADHD200等外部数据集上的评估进一步验证了其在不同人口统计学和临床分布中的卓越能力和鲁棒性。具体性能提升数据未知。

🎯 应用场景

该研究成果可广泛应用于神经科学领域，例如脑疾病的早期诊断、个性化治疗方案的制定、认知功能评估等。通过对脑动力学的深入理解，可以开发更有效的干预手段，改善患者的生活质量。此外，该模型还可以用于研究不同人群的脑功能差异，揭示脑发育和衰老的机制。

📄 摘要（原文）

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.