A Unified, Scalable Framework for Neural Population Decoding

Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael J. Mendelson, Blake Richards, Matthew G. Perich, Guillaume Lajoie, Eva L. Dyer

★ Must read October 2023

TL;DR POYO turns every spike into a token with a learned unit embedding and a timestamp, so one PerceiverIO-style transformer can train on recordings with different neurons. Trained on 27,373 units from seven monkeys, it adapts to unseen sessions and animals by learning only new unit embeddings or by finetuning, reaching R2 0.9482 on NLB-Maze against 0.9062 for the best prior baseline.

What they do

Neural recordings share no vocabulary: each session, day and animal gives a different set of neurons with no known correspondence, so decoders are usually trained per session. POYO (Pre-training On manY neurOns) is a transformer that trains across sessions, animals and labs to decode hand velocity from motor cortex spikes, and can then be adapted to new recordings with little labeled data.

How it works

Spike tokens. Each spike is one token: a learned embedding for the unit that fired plus the spike’s timestamp. There is no binning and no fixed number of units, so populations of any size can be ingested.
Latent compression. Cross-attention pulls a 1 s window of spikes (1k to 20k tokens) into a short sequence of latent tokens, followed by self-attention in the latent space. Rotary position encoding carries relative timing in every attention layer.
Querying. Output tokens with timestamps and a learned session embedding query the latents, so behavior sampled at any rate (100 Hz to 1 kHz) can be predicted.
Unit identification. For a new session, all weights are frozen and only new unit and session embeddings are learned by gradient descent; full finetuning can follow.

Results

POYO-mp, trained on 100 sessions from three monkeys (9,789 units), reaches average R2 0.9512 on center-out and 0.8738 on random target, against 0.9347 and 0.8402 for single-session models.
On a new animal (Monkey T), finetuned POYO-mp gets 0.9379 on center-out versus 0.9194 for a single-session model and 0.8577 for the best baseline (MLP). Unit identification alone gets 0.9012.
POYO-1 (seven monkeys, three labs, 27,373 units) reaches 0.9329 on NLB-Maze with unit identification alone and 0.9482 finetuned; AutoLFADS + Linear gets 0.9062. On NLB-RTT it gets 0.7378 versus 0.6953 for the MLP.
Unit identification takes less than a minute on one GPU.

Limitations

Training is supervised on hand velocity, and the authors say scaling to more brain regions and tasks will need a self-supervised objective. All data are motor, premotor and somatosensory cortex recordings from monkeys during reaching tasks. On the new animal, unit identification alone falls slightly below a single-session model, so finetuning is needed there. What the latent tokens and unit embeddings encode is left for future work.

Figures

From Figure 1: Overview of the approach. Input spike tokens are compressed into a smaller set of latent tokens which ... — From Figure 1: Overview of the approach. Input spike tokens are compressed into a smaller set of latent tokens which are processed through multiple self-attention blocks, finally time-varying outputs are predicted by querying the latent space. All tokens in this model are assigned a timestamp, which are used towards rotary position encoding.

From Figure 2: Building a multi-session model spanning multiple animals and tasks. (A) Center-out reach and (B) Rando... — From Figure 2: Building a multi-session model spanning multiple animals and tasks. (A) Center-out reach and (B) Random target task, along with examples of true and predicted behavior (x,y velocities). In (C-D), we show the decoding performance for single-session models (gray) and the POYO-mp multi-session model (green).

From Figure 3: Scaling curves. R2 is evaluated and averaged over the test splits on both CO and RT tasks. The single-... — From Figure 3: Scaling curves. R2 is evaluated and averaged over the test splits on both CO and RT tasks. The single-session performance is the mean over 100 different models.

From Figure 4: Sample and compute efficiency for training, unit identification, and fine-tuning approaches. On the le... — From Figure 4: Sample and compute efficiency for training, unit identification, and fine-tuning approaches. On the left, we show the sample and compute efficiency for a heldout CO session. On the right, we plot the sample and compute efficiency for a new animal not seen during training.

Abstract

Our ability to use deep learning approaches to decipher neural activity would likely benefit from greater scale, in terms of both model size and datasets. However, the integration of many neural recordings into one unified model is challenging, as each recording contains the activity of different neurons from different individual animals. In this paper, we introduce a training framework and architecture designed to model the population dynamics of neural activity across diverse, large-scale neural recordings. Our method first tokenizes individual spikes within the dataset to build an efficient representation of neural events that captures the fine temporal structure of neural activity. We then employ cross-attention and a PerceiverIO backbone to further construct a latent tokenization of neural population activities. Utilizing this architecture and training framework, we construct a large-scale multi-session model trained on large datasets from seven nonhuman primates, spanning over 158 different sessions of recording from over 27,373 neural units and over 100 hours of recordings. In a number of different tasks, we demonstrate that our pretrained model can be rapidly adapted to new, unseen sessions with unspecified neuron correspondence, enabling few-shot performance with minimal labels. This work presents a powerful new approach for building deep learning tools to analyze neural data and stakes out a clear path to training at scale.

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong