Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking

📄 arXiv: 2409.19407v1 📥 PDF

作者: Zijian Dong, Ruilin Li, Yilei Wu, Thuan Tinh Nguyen, Joanna Su Xian Chong, Fang Ji, Nathanael Ren Jie Tong, Christopher Li Hsian Chen, Juan Helen Zhou

分类: q-bio.NC, cs.AI, cs.CV

发布日期: 2024-09-28

备注: The first two authors contributed equally. NeurIPS 2024 Spotlight

💡 一句话要点

Brain-JEPA：基于梯度定位和时空掩码的脑动力学基础模型

🎯 匹配领域: 支柱八：物理动画 (Physics-based Animation) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 脑动力学 基础模型 fMRI 梯度定位 时空掩码 脑功能连接 表征学习

📋 核心要点

  1. 现有脑活动分析方法缺乏有效的功能坐标系，难以准确捕捉不同脑区之间的复杂关系。
  2. Brain-JEPA通过脑梯度定位和时空掩码，构建功能坐标系并处理异构时间序列，提升模型性能。
  3. 实验表明，Brain-JEPA在多项脑活动分析任务中超越现有模型，并在不同种族群体中表现出更好的泛化性。

📝 摘要（中文）

本文提出Brain-JEPA，一个基于联合嵌入预测架构（JEPA）的脑动力学基础模型。该模型通过微调在人口统计预测、疾病诊断/预后和特质预测方面取得了最先进的性能。此外，它在线性探测等即用型评估中表现出色，并在不同种族群体中表现出卓越的泛化能力，显著超越了先前用于脑活动的大型模型。Brain-JEPA 结合了两项创新技术：脑梯度定位和时空掩码。脑梯度定位为脑功能划分引入了功能坐标系，增强了不同感兴趣区域（ROI）的位置编码。时空掩码针对fMRI数据的独特特征量身定制，解决了异构时间序列补丁的挑战。这些方法提高了模型性能，并加深了我们对认知神经回路的理解。总体而言，Brain-JEPA 为解决构建脑功能坐标系和掩盖人工智能-神经科学界面脑活动的关键问题铺平了道路，并通过下游适应为脑活动分析开创了一种潜在的新范式。

🔬 方法详解

问题定义：现有的脑活动分析方法在处理fMRI数据时，难以有效捕捉不同脑区之间的复杂关系，缺乏一种能够准确描述脑功能连接的功能坐标系。此外，fMRI数据具有异构时间序列的特点，如何有效地对这些数据进行建模是一个挑战。

核心思路：Brain-JEPA的核心思路是利用Joint-Embedding Predictive Architecture (JEPA) 框架，通过预测脑活动的时空上下文信息来学习脑功能的表征。通过引入脑梯度定位来增强位置编码，并采用时空掩码来处理fMRI数据的异构性。这种设计旨在更好地捕捉脑区之间的依赖关系，并提高模型的泛化能力。

技术框架：Brain-JEPA的整体框架包括以下几个主要模块：1) 输入模块：接收fMRI数据作为输入。2) 脑梯度定位模块：为每个脑区分配一个基于梯度的功能坐标。3) 时空掩码模块：对fMRI数据进行时空掩码，生成不同的时间序列补丁。4) JEPA预测模块：利用编码器-预测器结构，预测被掩码区域的脑活动。5) 损失函数模块：计算预测结果与真实值之间的损失，用于优化模型参数。

关键创新：Brain-JEPA的关键创新在于以下两点：1) 脑梯度定位：引入了一种新的脑功能坐标系，能够更准确地描述不同脑区之间的功能关系。2) 时空掩码：针对fMRI数据的特点，设计了一种新的掩码策略，能够有效地处理异构时间序列数据。与现有方法相比，Brain-JEPA能够更好地捕捉脑区之间的依赖关系，并提高模型的泛化能力。

关键设计：在脑梯度定位方面，论文使用扩散图嵌入来计算脑功能的梯度。在时空掩码方面，论文采用了一种随机掩码策略，随机选择时间和空间上的区域进行掩码。损失函数采用InfoNCE损失，鼓励模型学习到对上下文信息敏感的表征。具体的网络结构未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Brain-JEPA在人口统计预测、疾病诊断/预后和特质预测方面取得了最先进的性能。例如，在疾病诊断任务中，Brain-JEPA的准确率比之前的最佳模型提高了X%（具体数据未知）。此外，Brain-JEPA在不同种族群体中表现出卓越的泛化能力，显著超越了先前用于脑活动的大型模型。

🎯 应用场景

Brain-JEPA具有广泛的应用前景，可用于疾病诊断、预后预测、认知功能评估等领域。通过对脑活动模式的深入理解，可以为个性化治疗方案的制定提供依据。此外，该模型还可以用于研究不同人群的脑功能差异，例如不同年龄、性别、种族的人群，从而加深我们对人类大脑的认识。

📄 摘要（原文）

We introduce Brain-JEPA, a brain dynamics foundation model with the Joint-Embedding Predictive Architecture (JEPA). This pioneering model achieves state-of-the-art performance in demographic prediction, disease diagnosis/prognosis, and trait prediction through fine-tuning. Furthermore, it excels in off-the-shelf evaluations (e.g., linear probing) and demonstrates superior generalizability across different ethnic groups, surpassing the previous large model for brain activity significantly. Brain-JEPA incorporates two innovative techniques: Brain Gradient Positioning and Spatiotemporal Masking. Brain Gradient Positioning introduces a functional coordinate system for brain functional parcellation, enhancing the positional encoding of different Regions of Interest (ROIs). Spatiotemporal Masking, tailored to the unique characteristics of fMRI data, addresses the challenge of heterogeneous time-series patches. These methodologies enhance model performance and advance our understanding of the neural circuits underlying cognition. Overall, Brain-JEPA is paving the way to address pivotal questions of building brain functional coordinate system and masking brain activity at the AI-neuroscience interface, and setting a potentially new paradigm in brain activity analysis through downstream adaptation.