跳转至

Leveraging unlabelled data for generalizable neural population decoding

会议: NeurIPS 2026
arXiv: 2607.14086
OpenReview: https://openreview.net/forum?id=GgTOMpnG1o
代码: https://github.com/braingeneers/mojo(⚠️ 论文未在正文给出可点击开源仓库,以官方后续公布为准)
领域: 其他
关键词: 神经群体解码、掩码自编码器、自监督联合训练、脉冲分词、脑机接口

一句话总结

本文提出基于掩码自编码器的自监督与监督联合预训练框架 MOJO,首次使单脉冲级别神经分词模型能够利用海量无标注神经电生理数据,在跨动物运动皮层解码、跨模态小鼠视觉/决策和人类语音脑电解码中均超越纯监督基线。

研究背景与动机

脑机接口(BCI)与闭环神经技术的核心是神经解码器,即从高维神经群体放电活动中还原出肢体运动、视觉感知或语言意图等行为变量。传统神经解码主要依赖卡尔曼滤波、MLP 或循环神经网络,受限于单 session 训练,难以适应电极漂移与跨会话神经变异性。近年来,以 POYO 和 POSSM 为代表的单脉冲分词(Spike Tokenization)基础模型,通过将每个动作电位直接编码为包含神经元 ID 与精确发生时间戳的 token,摆脱了传统固定时间 bin 计数拼接矩阵的僵硬结构,在跨会话甚至跨数据集预训练中取得了前沿突破。

然而,现有脉冲分词架构存在一个根本性短板:它们完全依赖配对的监督学习(Supervised Learning, SL),即预训练数据必须严格包含同步采集的行为标签(例如猴子手臂的 2D 抓取速度)。在实际神经科学实验中,高质量带标注的行为数据采集极其耗时且昂贵,而神经电生理实验室中存储着海量未经行为标注的自发或背景脑电记录。由于神经放电具有极高稀疏性、不规则时间间隔,直接在输入 token 序列做遮蔽会导致模型无法区分“被遮蔽的真实放电”与“原本就没有放电的静息区间”,导致传统掩码自编码(MAE)难以直接应用于脉冲序列。

这一矛盾的本质在于:脉冲分词模型渴望无监督数据的通用先验,但离散稀疏的时间点过程无法直接套用基于网格或连续序列的经典自监督方案。核心 idea:将自监督掩码操作从离散脉冲输入端后移至局部时间块的压缩隐变量层,构建掩码隐空间重构与多任务监督解码的联合训练目标,实现无标注神经动力学与任务行为表征的端到端协同。

方法详解

整体框架

MOJO 维护自监督学习(SSL)与监督学习(SL)两条协同路径。输入数据为多通道离散动作电位序列,每个脉冲被映射为神经元可学习嵌入与时间旋转位置编码(RoPE)之和。模型先通过局部输入交叉注意力(Input Cross-Attention)将每个时间块内的可变长脉冲序列压缩为定长隐空间 token;在 SSL 路径上对隐空间应用时间块随机掩码并用骨干网络重构被掩盖时段的泊松放电率;在 SL 路径上保留完整隐变量并结合骨干网络,通过任务专用查询与输出交叉注意力预测多任务行为目标。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入离散脉冲序列<br/>(Unit ID + 时间 RoPE)"] --> B["时间块局部交叉注意力<br/>(POSSM 局部编码压缩)"]
    B --> C["分路与隐空间时间掩码<br/>(50% Bernoulli 隐变量掩码)"]
    B --> D["未掩码隐空间序列<br/>(保留完整上下文)"]
    C --> E["共享骨干网络<br/>(Transformer / SSM)"]
    D --> E
    E --> F["输出交叉注意力与泊松重构<br/>(单神经元 Query 预测放电率)"]
    E --> G["多任务输出交叉注意力<br/>(任务 Query 解码行为变量)"]
    F --> H["联合损失反向传播<br/>L_MOJO = α_SSL * L_SSL + α_SL * L_SL"]
    G --> H

关键设计

1. 时间块局部输入交叉注意力:统一离散脉冲向定长隐空间的无损压缩

面对多通道、时间不均匀分布的离散脉冲输入,直接计算全序列注意力不仅引发二次方复杂度,还会破坏实时因果解码结构。MOJO 采用 POSSM 风格的局部交叉注意力设计,将连续神经记录切分为连续的时间块(默认 20 ms)。在每个时间块 \(t\) 内,所有脉冲 token \(\mathbf{X}_t\) 作为键(Key)和值(Value),与 \(N_c\) 个可学习的隐空间查询向量 \(\mathbf{c}_i^{(t)}\) 进行缩放点积注意力交互:

\[\mathbf{z}_i^{(t)} = \mathrm{softmax}\left(\frac{\mathbf{q}^{(t)} \mathbf{K}_t^\top}{\sqrt{D}}\right)\mathbf{V}_t\]

其中 \(\mathbf{q}^{(t)} = \mathbf{c}_i^{(t)}\mathbf{W}_q\)。此外,每个时间块中心还注入了特殊的无放电指示 token,确保未放电单元的沉默信息也能被显式编码。该局部设计将可变长度的高频脉冲彻底压缩为规整的隐变量序列 \(\mathbf{Z}\),同时保留精确的时间对齐。

2. 隐变量时间掩码自编码器:规避点过程离散空缺二义性

如果在输入脉冲端直接实施随机掩码,模型无法从数学上辨别是“神经元在当前毫秒未产生放电”,还是“脉冲存在但被算法故意掩蔽”。MOJO 的破局点在于将 MAE 机制后移至输入交叉注意力生成的隐变量 \(\mathbf{z}_i^{(t)}\) 上。算法在时间维度引入伯努利掩码变量 \(m^t \sim \mathrm{Bernoulli}(p)\)(默认掩码率 \(p = 0.5\)):

\[\tilde{\mathbf{z}}_i^{(t)} = (1 - m^t)\mathbf{z}_i^{(t)} + m^t \mathbf{e}_{\mathrm{mask}}\]

其中 \(\mathbf{e}_{\mathrm{mask}}\) 为可学习的掩码替换向量,并叠加指示空间查询来源的位置嵌入。由于局部交叉注意力严格限制在每个时间块内,一旦 \(m^t = 1\),该时间窗内的所有原始脉冲信息在 SSL 路径中被物理隔绝,迫使骨干网络仅依赖未掩盖上下文的群体动力学推断被掩蔽窗口的状态。

3. 神经元查询驱动的泊松放电率重构:无偏表征神经群体发放统计

被掩盖时段的重构目标需要符合神经放电的内在随机点过程特征。在经过共享骨干网络(Transformer 或 SSM)演化后,SSL 解码器使用神经元查询向量 \(\mathbf{u}\)(由神经元嵌入映射获得)通过受限输出交叉注意力重构放电率。为降低计算复杂度并避免过拟合,模型在每个被掩盖的时间块内随机无放回抽取至多 10 个神经元作为预测目标,且交叉注意力仅允许查询关注同一时间块内的隐空间 token。自监督损失函数采用神经脉冲预测标准的泊松负对数似然:

\[L_{\mathrm{SSL}} = \frac{1}{|M \times \mathcal{U}_{\mathrm{sub}}|} \sum_{t \in M} \sum_{u \in \mathcal{U}_{\mathrm{sub}}} \left( \hat{\lambda}_{u,t} - y_{u,t} \log \hat{\lambda}_{u,t} \right)\]

其中 \(\hat{\lambda}_{u,t}\) 为模型预测的瞬时放电率条件期望,\(y_{u,t}\) 为该时窗内的实际放电计数。该机制驱动网络学习群体神经元之间的功能连接与协同发放先验。

4. 路径深度集成与双目标联合优化:参数高复用与端到端对齐

MOJO 拒绝将 SSL 预训练与 SL 微调机械割裂为串行两阶段,而是构建了全参数共享的双分支联合架构。SL 路径使用无掩码的完整隐变量序列 \(\mathbf{Z}\) 输入骨干网络,通过任务专用查询与读出头最小化行为预测误差(均方误差或多类别交叉熵 \(L_{\mathrm{SL}}\));SSL 路径则输入掩码后的隐变量 \(\tilde{\mathbf{Z}}\) 并计算泊松重构损失。整个框架共享输入编码层与核心骨干网络(SSM 或 Transformer),在 POSSM 和 POYO 架构下新增的自监督参数开销分别仅约 10% 和 2%。总损失为双任务加权和:

\[L_{\mathrm{MOJO}} = \alpha_{\mathrm{SSL}} L_{\mathrm{SSL}} + \alpha_{\mathrm{SL}} L_{\mathrm{SL}}\]

实验表明,无需复杂的超参搜索,直接设置 \(\alpha_{\mathrm{SSL}} = \alpha_{\mathrm{SL}} = 1.0\) 即可在不同有标签/无标签数据比例下取得稳定增益。

损失函数 / 训练策略

在多数据集预训练阶段,模型针对不同会议与动物数据集启用多任务线性头(如猴子运动速度 MSE、小鼠刺激分类交叉熵)。在将预训练 MOJO 迁移至新会话或新个体时,论文支持两种下游微调机制: - 单元识别(Unit Identification, UI):冻结除神经元嵌入层(Unit Embedding)与会话嵌入(Session Embedding)之外的所有模型权重,仅优化数千至上万个参数(< 18K 参数),实现即插即用的极速少样本校准。 - 全微调(Full Finetuning, FT):先运行 100 epoch 的 UI 微调对齐新记录单元的特征空间,随后解冻全部骨干与读出头进行全局端到端反向传播。

实验关键数据

主实验

在跨会话与跨个体的猴子伸手运动手部 2D 速度解码任务中,MOJO 与纯监督基线(POYO、POMAMBA、POGRU)及传统基于时间 bin 计数的自监督大模型(NDT-2、NDT-3、NEDS)进行了系统对比。各方法在测试集上的平均决定系数 \(R^2\) 如下表所示:

方法 微调模式 同一动物不同日期 (C-CO 2010) \(R^2\) 未见新动物目标 (T-CO) \(R^2\) 未见新动物随机目标 (T-RT) \(R^2\)
NDT-2 FT 0.7846 ± 0.1167 0.7173 ± 0.0443 0.6323 ± 0.1339
NDT-3 FT 0.7524 ± 0.1322 0.8576 ± 0.0313 0.7066 ± 0.0980
NEDS FT 0.5968 ± 0.0760 0.7635 ± 0.0758 0.6121 ± 0.0918
POYO UI 0.7759 ± 0.1003 0.8123 ± 0.0419 0.7011 ± 0.0975
POMAMBA UI 0.7283 ± 0.1138 0.8574 ± 0.0225 0.7283 ± 0.0846
POGRU UI 0.7632 ± 0.1013 0.8587 ± 0.0216 0.7331 ± 0.0775
MOJO-POYO UI 0.7846 ± 0.0897 0.8387 ± 0.0251 0.7470 ± 0.0617
MOJO-POMAMBA UI 0.7937 ± 0.0735* 0.8753 ± 0.0160* 0.7591 ± 0.0588*
MOJO-POGRU UI 0.7949 ± 0.0892* 0.8772 ± 0.0189* 0.7570 ± 0.0693*
POYO FT 0.8244 ± 0.0753 0.8817 ± 0.0352 0.7624 ± 0.0815
POMAMBA FT 0.8142 ± 0.0763 0.8949 ± 0.0152 0.7580 ± 0.0745
POGRU FT 0.8126 ± 0.0892 0.8936 ± 0.0212 0.7575 ± 0.0875
MOJO-POYO FT 0.8438 ± 0.0888* 0.9131 ± 0.0177* 0.7964 ± 0.0663*
MOJO-POMAMBA FT 0.8153 ± 0.0925 0.9043 ± 0.0182* 0.7675 ± 0.0806
MOJO-POGRU FT 0.8222 ± 0.0950 0.9103 ± 0.0168* 0.7776 ± 0.0710*

注:* 表示与非 MOJO 最佳基线相比配对 t 检验 \(p < 0.05\)。

在第二项主实验中,模型在小鼠 Allen Visual Coding Neuropixels 数据集上进行 119 类自然场景(Natural Scenes, NS)与漂移光栅方向(Drifting Grating, DG)的视觉解码分类评估:

方法 微调模式 自然场景 (NS, 4 sessions) 准确率 (%) 漂移光栅 (DG, 4 sessions) 准确率 (%)
MLP 从头训练 83.17 ± 4.92 90.18 ± 2.03
POYO UI 88.17 ± 3.17 93.54 ± 2.06
POGRU UI 88.87 ± 2.55 99.17 ± 0.76
NEDS FT 93.96 ± 2.38 99.20 ± 0.57
MOJO-POYO UI 91.53 ± 1.99 99.74 ± 0.31
MOJO-POGRU UI 92.54 ± 3.41 99.53 ± 0.20
MOJO-POYO (J) UI 93.15 ± 2.39 99.43 ± 1.15
MOJO-POYO-L(J) UI 94.18 ± 2.00 99.84 ± 0.20
POYO FT 91.03 ± 4.07 93.39 ± 4.50
POGRU FT 91.09 ± 5.64 99.48 ± 0.62
MOJO-POYO FT 94.12 ± 2.79 100.00 ± 0.00
MOJO-POGRU FT 94.47 ± 2.83 99.22 ± 0.91
MOJO-POYO (J) FT 94.62 ± 2.74 99.69 ± 0.64
MOJO-POYO-L(J) FT 95.48 ± 2.05 99.74 ± 0.52

注:(J) 表示经过猴运动与小鼠视觉跨物种联合预训练;L 表示大模型版本。

消融实验

论文针对不同的自监督设计选择,在猴子运动伸手数据集上进行了严格的消融对比(表 13、14、15):

消融配置 变体细节 预训练 \(R^2\) 监督损失 (MSE) 微调 \(R^2\) (新会话) 说明
MOJO-POGRU (完整) 隐空间时间掩码 (50%) 0.9178 0.01299 0.8449 默认联合训练方案
替代自监督: 输入随机掩码 在脉冲 token 端随机丢弃 0.9080 0.01491 0.8318 离散脉冲输入端掩码表现较差
替代自监督: 输入神经元掩码 掩盖整条神经元脉冲流 0.9068 0.01502 0.8322 破坏特定通道神经动力学
替代自监督: 输入脑区掩码 随机掩盖整个脑区输入 0.3419 0.12182 崩塌 (未收敛) 脑区信息丢失导致严重欠拟合
替代自监督: 对比学习 (CPC) 同序列负样本对比预测 0.9114 0.01441 0.8274 对比损失出现严重过拟合
路径架构消融 (FT) 共享输入编码器 (默认) 0.8438 - 0.9131 (T-CO) 编码计算无冗余,FT泛化更强
路径架构消融 (FT) 独立双路输入编码器 0.8324 - 0.9030 (T-CO) 时间分辨率不一致导致微调掉点

关键发现

  • 低资源少样本微调优势极其显著:在猴子新会话极端少样本微调场景下(图 2a),仅使用 2 个和 4 个有标注校准试次,MOJO 结合 32 个无标注试次即可分别恢复出完全监督基线 60% 与 75% 的解码表现;在未见动物新任务中,微调参数量少于 18K 的 MOJO-POSSM (UI) 即可匹敌超过 7.6M 参数量的纯监督全微调模型。
  • 无标注数据大幅提升表征可解释性与单细胞生理特征预测:尽管模型在预训练中从未接触脑区标签与电极空间坐标,经 MOJO 学习的神经元嵌入在线性探测下能够以 82.72% 的准确率区分小鼠 3 大脑区、以 54.78% 准确率区分 7 个精细脑区(比纯监督 POYO 的 27.99% 翻倍);在线性回归 18 项单神经元放电统计特征时,MOJO 对对数中位数放电间隔(ISI)的预测 \(R^2\) 达 0.61(POYO 仅 0.33),对平均放电率预测达 0.88(POYO 为 0.54),并在探针相距 4 mm 的远端视觉皮层与丘脑/中脑功能耦合神经元对上呈现出紧密聚类。
  • 跨物种正向知识迁移(Cross-species Positive Transfer):利用灵长类猴运动皮层与小鼠视觉皮层进行联合预训练(MOJO-POYO-L(J)),小鼠视觉自然场景分类在 199 轮早期验证集收敛速度提升明显(准确率从 73.01% 提升至 79.50%),且完全保留了猴运动解码的原始精度(\(R^2\) 保持在 0.855),证明了隐空间自监督在异构神经动力学中提取通用先验的可行性。

亮点与洞察

  • 精巧规避脉冲点过程掩码困境:以往自监督掩码图像或连续时间序列模型在遇到神经离散放电时均告失效,本文避开直接在输入端强行定义脉冲缺失,而是将掩码操作安排在时间块局部的隐变量层面,既与主流 POYO 骨干完美兼容,又彻底解决了“无放电”与“被掩盖”的语义混淆。
  • 参数开销极低的双路径统一建模:通过共享局部时间块输入交叉注意力与深层骨干序列模型,仅通过外挂一个轻量级神经元查询解码头(参数增量仅 2%~10%),便在统一的梯度计算流中达成了自监督与监督任务的表征共赢,摆脱了繁琐的二阶段预训练-微调调参壁垒。
  • 打破脉冲与连续脑电模态隔阂:虽然 MOJO 是为离散神经脉冲量身打造,但通过引入线性值嵌入 patch 化设计,该框架无缝迁移至人类开颅高密度 ECoG 语音发音解码任务,音节分类准确率较纯监督提升高达 47%,性能比肩专用连续信号神经大模型 Du-IN。

局限与展望

  • 自监督学习高度依赖数据吞吐规模:作者在消融中坦承,MOJO 存在“数据饥渴”属性。在单会话小样本场景下从头训练时,若无标注脉冲数据不够充分,欠拟合的 SSL 辅助任务反而会对监督解码带来负迁移干扰。
  • 掩码策略仍局限于时域截断:当前隐变量掩码机制主要沿时间维度执行随机丢弃,尚未显式构建空域掩码(如掩盖特定神经元子集或未见脑区通道),导致跨脑区或缺失通道的推断泛化能力未能充分挖掘。
  • 下游跨会话迁移需重新学习单元嵌入:继承自 POYO 家族的固有局限,当模型迁移到电极排布未知的新记录会话时,新单元的 Embedding 仍必须在少量下游试次中从头初始化并优化,未能实现完全零样本即插即用。未来可结合摊销推断(Amortized Inference)或离散神经码本技术进一步解耦神经元身份学习。

相关工作与启发

  • vs POYO / POSSM:POYO 首次确立了以单脉冲为单位的 Token 化解码方案,POSSM 则引入 SSM 与局部交叉注意力降低了推理延迟;但二者均严格局限在监督学习框架下,无法使用未标注数据。MOJO 在保留二者灵活分词优势的同时,引入隐空间 MAE,突破了标注瓶颈并大幅提高了少样本适应性。
  • vs NDT-2 / NDT-3 / NEDS:这类方法基于传统固定时间 bin 统计放电次数并做多通道矩阵切片,跨会话迁移必须依赖复杂的通道对齐或缝合层(Stitching Layers)。MOJO 沿用点过程脉冲分词,规避了刚性网格假设,在跨物种跨动物解码中表现出明显更优的泛化上限与自监督几何结构。
  • vs OmniMouse / POCO:近期工作探索了针对钙成像(Calcium Imaging)等连续荧光追踪信号的自监督预测模型,采用连续时间 1D 卷积与大范围插值。MOJO 则针对电生理离散点过程设计,在数学上用泊松点过程似然取代连续回归损失,补齐了脉冲领域自监督大模型的关键拼图。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙以后移隐变量时间掩码解决脉冲点过程自监督建模难题,架构设计优雅且通用。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖猴运动皮层、小鼠视觉与决策、人脑 ECoG 语音共 4 类跨物种多模态数据集,包含系统的少样本消融、生理可解释性探测与探针拓扑几何分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑主线极为清晰,图表完备,方法与动机高度自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为推动神经电生理基础模型(Neuro-Foundation Models)吸收海量无标注数据奠定了核心技术范式。