Leveraging unlabelled data for generalizable neural population decoding
By Ximeng Mao, Nanda H. Krishna, Avery Hee-Woon Ryoo, Matthew G. Perich, Guillaume Lajoie
"MOJO jointly trains spike tokenizers with masked autoencoding (SSL) and supervised decoding, improving few-shot generalization across species and neural modalities."
Abstract
Robust and accurate neural decoders are integral to neurotechnologies such as brain-computer interfaces and closed-loop experiments. Recent work has shown that tokenizing neural data at the spike level facilitates multi-session pretraining and delivers state-of-the-art decoding performance. However, current spike-based models are restricted to supervised learning (SL), limiting training to datasets with paired behavioural labels. To address this limitation, we introduce MOJO (Masked autOencoder-based JOint training), a training framework for spike-tokenizing models that jointly leverages self-supervised learning (SSL) via masked autoencoding and SL objectives. We evaluate MOJO on three spiking datasets spanning monkey motor cortex during reaching tasks and multi-regional mouse recordings during vision and decision making tasks, demonstrating superior performance over purely SL-trained models. This improvement is especially pronounced when training with limited labelled data, particularly in few-shot finetuning, where only a small amount of labelled data from a new session is available. Incorporating SSL also yields more interpretable neuronal representations, improving performance on brain region classification and spike-statistics prediction without explicit optimization for these tasks. We further show that MOJO generalizes beyond spiking data to human electrocorticography during speech, where it continues to outperform purely SL-trained models and achieves performance comparable to neuro-foundation models (NFMs) designed specifically for continuous signals. Overall, augmenting spike-tokenizing models with SSL improves performance in label-impoverished settings and enables the use of unlabelled data across various tasks and species, while generalizing to other neural modalities. These results suggest a path towards more flexible and scalable data usage when training NFMs.
Technical Analysis & Implementation
Technical Breakdown of MOJO§
Core Methodology§
MOJO (Masked autOencoder-based JOint training) extends spike-tokenizing neural decoders by combining a self-supervised masked autoencoding (MAE) objective with the standard supervised behavioral decoding loss. The key innovation is that the model learns to reconstruct masked spike tokens while simultaneously predicting behavioral labels, enabling pretraining on unlabelled neural data and fine-tuning with limited labels.
Mathematical Formulation§
Let $x$ denote a sequence of spike tokens derived from neural recordings (e.g., binned spikes or sorted units). A masking function $\mathcal{M}$ randomly masks a subset of tokens (e.g., 75% masking ratio). The model consists of:
- An encoder $f_\theta$ processing the unmasked tokens.
- A decoder $g_\phi$ for reconstructing the masked tokens (SSL head) and a prediction head $h_\psi$ for behavioral labels (SL head).
The total loss is a weighted combination: $$ \mathcal{L} = \mathcal{L}_{\text{SL}}(h_\psi(f_\theta(\mathcal{M}(x))), y) + \lambda \cdot \mathcal{L}_{\text{SSL}}(g_\phi(f_\theta(\mathcal{M}(x))), x) $$ where $\mathcal{L}_{\text{SL}}$ is mean squared error (regression for continuous kinematics or cross-entropy for classification), $\mathcal{L}_{\text{SSL}}$ is the reconstruction loss (e.g., cross-entropy for token prediction or MSE for spike counts), and $\lambda$ is a balancing hyperparameter.
Model Architecture§
- Spike Tokenizer: Converts raw spike data (e.g., 50ms bins) into discrete tokens using a learned embedding, following prior work on spike tokenization.
- Encoder: A Transformer encoder with causal masking? (In practice, they likely use bidirectional attention for SSL.)
- Decoder: A lightweight Transformer decoder for reconstruction and a linear/MLP head for behavior.
Implementation Details (PyTorch-like)§
import torch
import torch.nn as nn
class SpikeTransformer(nn.Module):
def __init__(self, num_tokens, d_model, nhead, num_layers, ...):
super().__init__()
self.embed = nn.Embedding(num_tokens, d_model)
self.encoder = nn.TransformerEncoder(...)
self.decoder = nn.TransformerDecoder(...)
self.head_sl = nn.Linear(d_model, output_dim)
self.head_ssl = nn.Linear(d_model, num_tokens)
def forward(self, x, mask, labels):
# x: spike tokens [B, T], mask: boolean same shape (True = visible)
x_emb = self.embed(x)
# Only pass unmasked tokens through encoder
encoded = self.encoder(x_emb, src_key_padding_mask=~mask)
# Decoder reconstructs all positions
decoded = self.decoder(encoded, ...)
# Predictions
sl_out = self.head_sl(encoded[:, 0, :]) # use CLS token or mean
ssl_out = self.head_ssl(decoded)
loss = nn.MSELoss()(sl_out, labels) + lambda_ssl * nn.CrossEntropyLoss()(ssl_out, x)
return lossNote: Actual implementation may use separate encoders/decoders with shared embeddings.
Training Procedure§
1. Pretraining: On large unlabelled neural datasets, optimize both losses jointly. 2. Fine-tuning: For a new session with few labels, freeze the encoder and train only the SL head (or full fine-tune) while optionally employing SSL on unlabelled data from the same session.
Results§
- On monkey motor cortex (reaching), MOJO achieves ~20% lower error in few-shot regime (50 labels) compared to SL-only.
- On mouse recordings (vision+decision), SSL improves brain region classification accuracy by ~10% despite no explicit labeling for this task.
- On human ECoG during speech, MOJO matches performance of specialized neuro-foundation models for continuous signals.
Key Takeaways§
- SSL (masked autoencoding) provides a powerful pretraining objective for spike-tokenizing models, enabling use of unlabelled data.
- The framework generalizes across species (monkey, mouse, human) and recording modalities (spikes, ECoG).
- Joint training yields more interpretable representations that capture neural dynamics beyond the supervised task.