OmniMouse: Scaling properties of multi-modal, multi-task Brain Models on 150B Neural Tokens

这篇论文介绍了 OmniMouse，这是一个用于建模小鼠视觉皮层神经元活动的多模态、多任务基础模型。该研究基于迄今为止最大的单神经元数据集之一（包含 73 只小鼠、323 次记录、超过 300 万个神经元，总计超过 1500 亿个神经 Token），探讨了数据与模型规模扩展（Scaling）在神经科学建模中的规律。

以下是该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 背景： 在语言（LLM）和计算机视觉领域，通过扩展数据量和模型参数量（Scaling Laws）已经取得了突破性进展，形成了统一的基础模型。然而，在神经科学领域，是否也存在类似的扩展规律尚不清楚。
• 挑战：
  • 现有的神经数据集通常规模较小、碎片化且多样性不足。
  • 之前的模型大多局限于单一模态（仅预测神经活动或仅解码行为）、单一任务，或者缺乏跨数据集的扩展性。
  • 缺乏能够同时处理神经预测、行为解码和神经预测（Forecasting）的统一架构。
• 核心问题： 构建一个统一的多模态模型，能否在大规模神经数据上实现性能提升？神经建模的瓶颈是数据量还是模型容量？

2. 方法论 (Methodology)

2.1 数据集 (Dataset)

• 规模： 323 次记录，73 只清醒、头固定的小鼠，超过 300 万个单神经元记录。
• 模态：
  • 神经活动： 通过双光子钙成像记录（6-14 Hz 采样，后上采样至 30 Hz），覆盖初级视觉皮层及周围高级视觉区域。
  • 视觉刺激： 自然图像（ImageNet）、电影视频、Sports-1M 数据集、参数化刺激（如 Gabor 波、随机点运动等）。
  • 行为数据： 跑步速度、瞳孔位置（x, y）、瞳孔直径及其导数（20 Hz 采样）。
• Token 化： 总 Token 数超过 1500 亿。

2.2 模型架构 (OmniMouse Architecture)

OmniMouse 采用基于 Transformer 的统一架构，支持灵活的掩码（Masking）策略，能够处理任意模态组合的输入和输出。

• Token 化 (Tokenization)：
  • 神经响应： 采用类似 POYO+ 的单神经元 Token 化方法。使用步长卷积将钙信号转换为 Token，并添加神经元、会话和动物的身份嵌入（Identity Embeddings）。
  • 视觉刺激： 使用 Hiera 视觉 Transformer（前 10 层）提取时空特征。
  • 行为轨迹： 将行为变量投影为共享的线性层嵌入，并添加通道和身份嵌入。
• 掩码策略 (Masking)：
  • 模型训练时采用 119 种结构化掩码配置。
  • 核心任务： 神经预测（基于过去预测未来）、子群预测（基于部分神经元预测目标神经元）、刺激编码（基于视频预测神经活动）、行为解码（基于神经活动预测行为）。
  • 灵活性： 支持任意组合，例如“基于部分神经元 + 视频 + 行为”预测未来神经活动。
• 网络结构：
  • 编码器： 使用交叉注意力（Cross-Attention）将可变长度的输入压缩为固定长度的潜在序列（Latent Sequence），包含全局注册 Token（Global Registers）以保留全局信息。
  • 多模态融合： 堆叠多层 Transformer（局部滑动窗口注意力 + 全局注意力），融合神经、行为和视觉特征。
  • 解码器： 使用交叉注意力机制，根据查询 Token（Query Tokens）重建目标神经活动或行为轨迹。
• 训练目标：
  • 神经预测：泊松损失（Poisson Loss）。
  • 行为解码：均方误差（MSE Loss）。
  • 采用 Warmup-Stable-Decay 训练策略，确保模型收敛并生成密集的中间检查点用于扩展性分析。

3. 主要贡献 (Key Contributions)

1. 系统性的扩展性分析 (Systematic Scaling Analysis)：
   • 发现数据量是提升预测精度的主要驱动力，而模型规模在中等规模（约 80M 参数）后会出现性能饱和。
   • 这与语言模型（LLM）的扩展规律相反：在 LLM 中，大数据集使得参数扩展成为进步的主要驱动力；而在神经建模中，即使数据量巨大，模型仍受限于数据多样性而非计算能力。
2. 统一的多模态多任务架构：
   • 提出了第一个能够灵活处理神经预测、行为解码、刺激编码及其任意组合的单一模型。
   • 通过灵活的掩码机制，实现了跨神经元、跨时间、跨刺激和跨动物的统一建模。
3. 最先进性能 (State-of-the-Art Performance)：
   • 在几乎所有评估任务中，OmniMouse 均超越了专门的基线模型（如 MtM, Neuroformer, CEBRA, POYO+ 等）。
   • 即使在数据量匹配（Data-matched）的条件下，OmniMouse 的架构优势依然明显。
   • 在 Sensorium 2022 和 2023 竞赛中刷新了记录。

4. 实验结果 (Results)

• 模型规模扩展：
  • 随着参数量从 1M 增加到 80M，所有神经预测任务的性能显著提升。
  • 超过 80M 参数后，性能提升微乎其微，甚至出现过拟合，表明当前模型受限于数据而非计算力。
  • 行为解码表现出最积极的扩展动态，大模型在行为预测上仍有提升空间，未完全饱和。
• 数据规模扩展：
  • 在 8 到 323 次会话的嵌套数据集中，性能随数据量增加而稳定提升。
  • 更大的模型（80M）从额外数据中获益更多，特别是在包含视频输入的任务中。
  • 数据扩展对行为解码的增益最大，且未见饱和迹象。
• 基线对比：
  • 在 8 只小鼠的小规模数据集上，OmniMouse-5M 已超越所有基线模型。
  • 在 323 次会话的全量数据上，OmniMouse-80M 在神经预测和行为解码的所有指标上均达到 SOTA。
• 泛化能力：
  • 模型能够泛化到训练期间未见过的掩码配置（如不同的上下文神经元数量），证明了其学习到了可泛化的神经表征。

5. 意义与展望 (Significance)

• 重新定义神经建模的扩展规律： 论文指出，对于目前的神经数据（即使是小鼠视觉皮层这种相对简单的系统），数据多样性和数据量仍是瓶颈，而非模型容量。这提示未来神经科学的基础模型发展需要更丰富、更多样化的数据集。
• 统一框架的潜力： OmniMouse 证明了单一模型可以统一处理多种神经科学任务，为构建神经科学领域的“基础模型”铺平了道路。
• 相变的可能性： 作者推测，随着数据集的进一步丰富（更多样化的刺激、行为、脑区），可能会像大语言模型一样触发“相变”，解锁定性新的能力（Emergent Abilities）。
• 局限性： 模型参数随神经元数量线性增长（每个神经元有独立嵌入），限制了向更大规模数据集扩展的计算效率；目前主要基于钙成像数据，尚未整合电生理数据；行为数据主要限于自发活动，复杂行为的泛化性待验证。

总结： OmniMouse 通过构建一个大规模、多模态的统一模型，揭示了神经科学建模中“数据驱动”而非“参数驱动”的扩展特性，为未来构建更强大的神经基础模型提供了重要的实证依据和架构蓝图。代码已开源。