[Paper Review] A Unified, Scalable Framework for Neural Population Decoding

Mehdi Azabou, Vinam Arora+8 more|arXiv (Cornell University)|Oct 24, 2023

Neural dynamics and brain function21 citations

TL;DR

The paper introduces POYO, a spike-tokenization and PerceiverIO-based framework that trains large-scale pretrained models across multi-session neural recordings, enabling few-shot adaptation to unseen sessions without fixed neuron correspondence.

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.

Motivation & Objective

• Motivate scaling neural decoding with large datasets and diverse sessions across individuals.
• Address the lack of a shared neuron correspondence by developing a token-based neural population representation.
• Introduce a scalable training framework that supports multi-session, cross-individual electrophysiology data.
• Demonstrate pretrained models that can transfer to new sessions with minimal labeled data.

Proposed method

• Tokenize neural population activity by representing each spike as a token with a learnable unit embedding and a timestamp.
• Compress the input tokens into a latent sequence using a PerceiverIO cross-attention encoder to enable scalable processing.
• Apply time-aware rotary position encoding and sparse, event-based representations to preserve temporal structure.
• Pretrain large models on data from multiple sessions/animals and enable fine-tuning for new tasks and recording contexts.
• Provide pretrained models (POYO-1 and POYO-mp) and code to the community for cross-session neural decoding.

Experimental results

Research questions

• RQ1 Can a single model decode neural population activity across heterogeneous sessions with no known neuron correspondence?
• RQ2 Does spike-based tokenization enable scalable, cross-session learning without fixed input units across datasets?
• RQ3 Can large pretrained POYO models adapt quickly (few-shot) to unseen sessions and tasks with minimal labels?
• RQ4 What scaling laws emerge for neural decoding when training on multi-session neural data across multiple animals?

Key findings

• The pretrained POYO models can be rapidly adapted to unseen sessions with very few labeled examples.
• The approach scales to data from 158 sessions, seven non-human primates, 27,373 units, and over 100 hours of recordings.
• Spike-based tokenization preserves fine temporal structure while enabling a sparse, scalable input representation.
• The framework enables end-to-end training of transformer-based models on multi-session neural data.
• Two large pretrained models, POYO-1 and POYO-mp, are developed to support cross-recording transfer and task generalization.