PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2310.16046v1cs.LGq-bio.NC

TL;DR

Neural recordings are difficult to combine because neuron sets differ across individuals and sessions. POYO tokenizes individual spikes and uses a PerceiverIO-based framework to train across recordings; its pretrained models improve decoding and support transfer with few samples.

  • Problem

    Integrating neural recordings is challenging because different individuals and sessions contain neuron sets without straightforward correspondence.

    Source: §1 · p. 1

  • Method

    POYO tokenizes individual spikes and combines this representation with a PerceiverIO-based architecture for scalable training across neural recordings.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    Pretraining on large amounts of data improves overall brain decoding and supports few-shot transfer to new neural datasets.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    The pretrained models offer a potential way to compare datasets and investigate common activity motifs across sessions and individuals.

    Source: §5 · p. 10

    The study does not reprocess or standardize data across labs or tasks despite differences in preprocessing, controllers, and behavioral sampling rates.

    Source: §3.1 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

Neural recordings are difficult to unify because neuron sets differ across individuals and sessions without straightforward alignment. POYO combines spike-level tokenization with a scalable multi-session framework, evaluated on recordings from seven nonhuman primates.

1 Introduction

  • Different individuals and recording days yield distinct neuron sets, complicating the integration of recordings into a unified population-level account.Electrode/tissue variability can change neuron sets across days, even within the same individual.Source: §1 · pp. 1–2
  • POYO tokenizes individual spikes and uses a PerceiverIO-based architecture to learn interactions across spikes in time and neurons.The framework is designed for training across multiple sessions, including recordings with no known neuron correspondence.Source: §1 · p. 2
  • Across over 158 sessions from seven non-human primates, spanning over 27,373 units and 100 hours of recordings, pretraining on large amounts of data enables few-shot transfer and improves overall brain decoding performance.The evaluation uses open electrophysiology datasets.Source: §1 · p. 2
  • The work introduces an end-to-end multi-session training framework, spike-based tokenization, and two pretrained models, POYO-1 and POYO-mp.The authors state that model weights and code will be made available as a community resource.Source: §1 · p. 2

2 Approach

The approach represents spikes as timestamped tokens, compresses them into a latent sequence for attention, and queries that sequence to predict outputs. It supports recordings with varying neuron sets and output sampling rates, then adapts to new sessions by learning new unit embeddings while freezing the existing model weights.

2.1 Tokenizing neural population activity

  • Each spike becomes a token defined by its unit embedding and event time, allowing arbitrary unit populations without fixed time bins.The number of input tokens varies with the number of units and their firing rates.Source: §2.1 · p. 3

2.2 Building the latent space

  • Cross-attention compresses the input sequence into a shorter latent sequence, where self-attention costs O(N^2) rather than O(M^2).The learned latent tokens provide queries, while input tokens provide keys and values.Source: §2.2 · p. 4

2.3 Encoding relative timing information

  • Rotary position encoding uses relative positions, with latent tokens assigned timestamps spread evenly across the context window.RoPE applies rotations in query and key embedding space across attention layers.Source: §2.3 · p. 4

2.4 Querying the latent space

  • Cross-attention queries the latent sequence to predict output sequences of varying lengths and sampling rates, using timestamps to retrieve temporally relevant context.Learned session embeddings are injected into output queries to capture experimental variables.Source: §2.4 · pp. 4–5

2.5 Unit identification in new sessions

  • For a new session, unit identification learns embeddings for its units and adds a session embedding while freezing the model’s existing weights.The model that maps neural activity to behavior is transferred unchanged to the new dataset.Source: §2.5 · p. 5

3 Experiments

The experiments assess neural decoding across diverse datasets, then test whether joint pretraining, model and data scaling, and adaptation can improve performance across sessions and recording conditions. POYO models improve on single-session baselines and transfer to held-out sessions and datasets, with POYO-1 extending the evaluation across labs and tasks.

3.1 Datasets and experiment setup

  • The evaluation spans 178 sessions, 29,453 units, and nine nonhuman primates; tasks and recording conditions varied across datasets, which the authors did not reprocess or standardize.The CO task was stereotyped, whereas RT involved continuous, self-paced reaches to randomly appearing targets.Source: §3.1 · pp. 5–6, Table 1 · p. 5

3.3 POYO-mp: Building a large pretrained across-animal, multi-session model

  • Multi-session POYO-mp outperformed single-session models across sessions, with larger gains on sessions with fewer trials and on the more complex random-target task.With increasing model depth and joint training data, RT decoding improved by 4% over the single-session baseline when training used 134M tokens across CO sessions.Source: Figure 2 · p. 6, §3.3 · p. 7

3.2 Testing the model on single sessions

  • Across 100 single-session models, average R2 was 0.9347 on CO and 0.8402 on RT, consistently outperforming Wiener Filter, GRU, and MLP baselines.The sessions came from three nonhuman primates and used the same task setup, behavioral recording apparatus, and spike-sorting procedure.Source: §3.2 · p. 6

3.3 POYO-mp: Building a large pretrained across-animal, multi-session model

  • POYO-mp was trained on 100 sessions, and scaling analyses found improvements with greater model depth and joint training data on CO and RT tasks.Figure 3 reports test-split R2 averaged across both tasks, with single-session performance averaged over 100 models.Source: §3.3 · pp. 6–7, Figure 3 · p. 7

3.4 Transferring to new sessions

  • On held-out sessions from trained animals, unit identification reached 0.9675 versus 0.9682 for single-session models trained from scratch; fine-tuning raised accuracy to 0.9708.For a new animal, fine-tuning reached 0.9379 on CO and 0.7916 on RT, using a few minutes of labeled data.Source: §3.4 · pp. 7–8, Figure 4 · p. 8

3.5 POYO-1: A multi-lab, multi-task model for neural decoding

  • POYO-1 improved across variable datasets; NLB-Maze unit identification reached R2 0.9329, almost 4% above POYO-mp's unit-identification result.The models also showed different strengths: POYO-mp scored 0.8788 on same-lab RT sessions versus POYO-1's 0.8664, while POYO-1 had an edge in diverse transfer, especially with unit identification.Source: Figure 5 · p. 9, §3.5 · p. 9

4 Related Work

Prior transformer approaches to time series and neural activity typically discretize data into bins, while earlier multi-session methods rely on narrower recording and task assumptions. The paper situates its approach against these limits in modeling heterogeneous neural recordings.

Transformer architectures for time-series data

  • Earlier time-series transformers commonly bin continuous data, while Perceiver uses cross-attention to process inputs directly and reduce attention complexity.Many event-stream models extend only to univariate streams or assume fixed, known channels for multivariate data.Source: §4 · p. 9

Transformers applied to neural data

  • Neural-data transformers tokenize binned activity and train on a single session with a fixed neuron set.NDT embeds all neurons jointly per time bin; EIT and STNDT use separate neuron/time tokenizations or representations.Source: §4 · p. 10

Multi-session training and alignment

  • Prior multi-session methods generally assume the same subject, brain region, implant, and stereotyped task, or align sessions after training.LFADS requires per-session alignment initialization and has been demonstrated for a single task in one subject.Source: §4 · p. 10

5 Discussion

The discussion presents POYO as a framework for unified multi-session training and efficient fine-tuning, while identifying interpretability and self-supervised training as directions for further development. The appendix describes its relative-time encoding, fixed-duration context windows, and tokenization details.

5 Discussion

  • POYO provides a framework for training transformers on large multi-session neural datasets, with pretrained models that can be fine-tuned on new datasets and are available to the community.The paper reports that unified multi-recording training improves decoding performance.Source: §5 · p. 10
  • Further work should probe learned activity motifs and embedding organization, while scaling across more brain regions and tasks is tied to moving toward self-supervised objectives.The framework could also be extended to BMI and disease-modeling applications.Source: §5 · p. 10

A.1 Rotary position encoding

  • RoPE encodes relative timing in attention scores, using rotation periods that span 1 ms to 4 s in these models.The rotation matrix is incorporated into attention scores, which depend on relative query-key timing.Source: §A.1 · p. 14
  • The model also rotates values to encode relative timing in attention outputs, using a factorization to avoid computing rotations for all N^2 time differences.The authors report improved performance with value rotation in input cross-attention and self-attention, and with only half the head dimensions rotated.Source: §A.1 · pp. 14–15

A.2 Time-based context window

  • A fixed 1-second context window gives sequences a variable number of spike tokens, so batches are padded and attention-masked, with load balancing for variable-length inputs.Training does not use trial structure; inference uses 500 ms sliding windows and averages overlapping predictions.Source: §A.2 · p. 15

A.3 Further details on tokenization

  • Unit delimiters mark the start and end of each window, making even units with no spikes in that window visible to the model.The delimiter tokens receive the unit embedding plus a learned start or end embedding.Source: §A.3 · p. 15
  • The 256 latent tokens are arranged in groups of eight sharing timestamp-specific embeddings, so each group queries spike tokens at the same time in the first cross-attention layer.After information is drawn from the inputs, those tokens evolve independently.Source: §A.3 · p. 15

A.4 Training details

Training uses LAMB with weight decay and a cosine learning-rate decay near the end of training. Large and single-session models use different batch sizes and GPU setups.

A.4 Training details

  • The models use the LAMB optimizer with weight decay and a constant learning rate followed by cosine decay over the last 25% of epochs.Source: §A.4 · p. 15
  • Single-session models use one GPU and batch size 128, while large models use eight GPUs and total batch size 1,400.Source: §A.4 · p. 15
  • Increasing batch size did not benefit single-session model training.Source: §A.4 · p. 15

A.4.1 Compute

POYO-mp and POYO-1 were trained for 400 epochs on eight-GPU hardware, while single-session training and unit identification required substantially less compute time.

A.4.1 Compute

  • POYO-mp trained for 2 days and POYO-1 for 5 days, each for 400 epochs.Source: §A.4.1 · p. 16
  • The large models used eight Nvidia A40 GPUs, while single-session models used one Nvidia GeForce RTX 3090 GPU and took less than an hour to train.Source: §A.4.1 · p. 16
  • Unit identification took less than a minute on one GPU or a few minutes on the CPU.Source: §A.4.1 · p. 16

A.4.2 Data augmentation

Training uses unit dropout, a loss designed for mixed behavioral sampling rates and task quality, and length-aware GPU batching to manage variable input sequences.

A.4.2 Data augmentation

  • Unit dropout randomly samples neural units, with a minimum augmented population of 30; training itself accepts any population size, and other augmentations were not explored.The authors identify other augmentation strategies as a possible direction for improving model capabilities.Source: §A.4.2 · p. 16

A.4.3 Loss

  • MSE trains hand-velocity sequences after lab-level z-scoring; randomly sampled 100 timepoints address unequal sampling rates, and reaching-segment loss is weighted by a factor of 5.The sampled timepoints are resampled at each pass; the increased weight is applied to center-out reaching segments.Source: §A.4.3 · p. 16

A.4.4 Load Balancing

  • Sequences of similar length are assigned to the same GPU node; buckets and local batch sizes are set to use compute capacity while minimizing padding.GPUs assigned longer sequences receive smaller local batches.Source: §A.4.4 · p. 16

A.5 Evaluation Details

Evaluation uses task-specific scoring segments, while fine-tuning follows a shared gradual-unfreezing procedure across experiments and the recordings span diverse platforms and brain regions.

A.5 Evaluation Details

  • Evaluation excludes training-seen segments and scores successful center-out reaches, reaches in trialized random-target tasks, and all segments in continuous random-target datasets.During training, the model uses any 1-second recording segment.Source: §A.5 · p. 16

B Datasets

  • The curated datasets span neural and behavioral recording platforms; POYO-mp covers PMd and M1, while POYO-1 spans PMd, M1, and S1 across both hemispheres and seven animals.The figure identifies four animals for POYO-mp and seven for POYO-1.Source: Figure 6 · p. 17, B Datasets · p. 17

A.6 Fine-tuning Details

  • Fine-tuning first updates unit and session embeddings before unfreezing the remaining weights; the same hyperparameters are used across experiments.The authors say this consistency suggests adaptation without expert neural-network training knowledge or expensive hyperparameter search.Source: §A.6 · p. 17

B.1 Datasets used for training

The training datasets combine movement recordings collected across animals, tasks, and sites, while retaining substantial differences in recording and preprocessing practices.

B.1 Datasets used for training

  • POYO-mp uses 100 sessions from three nonhuman primates performing two movement tasks with a custom 2-D planar manipulandum.The datasets came from four previously published studies.Source: §B.1 · p. 17
  • The tasks were center-out reaches among eight radial targets and random-target movements among four randomly placed targets.The center-out task includes a variable hold period before the go cue.Source: §B.1 · p. 17
  • POYO-1 adds Churchland, Makin, and Flint datasets while holding out 12 Monkey T sessions, two Monkey C sessions, and two NLB datasets.The added recordings include touch-screen tasks with different behavioral sampling rates.Source: §B.1 · p. 18
  • The datasets are not standardized and velocity is not resampled; behavior is clipped only during periods of extremely high acceleration using a threshold heuristic.The authors describe minimizing processing as important for integrating more datasets and making the model accessible out of the box.Source: §B.1 · p. 18
  • Inputs retain collection differences, including threshold-crossing and spike-sorted units, without filtering by presence ratio or rejecting multi-units.Spike-sorting algorithms were tuned differently across labs.Source: §B.1 · p. 18
  • Minimizing processing is intended to ease integration of more datasets and provide an out-of-the-box model accessible to the wider community.Source: §B.1 · p. 18

B.2 Datasets for evaluation and fine-tuning

Evaluation holds out unseen data, includes two NLB datasets with distinct reach structures, and compares POYO with established decoding baselines using dataset-specific binning.

B.2 Datasets for evaluation and fine-tuning

  • Evaluation holds out 20% of each session, all six CO and six RT sessions from Monkey T, and the MC-Maze and RTT NLB datasets.Reported NLB performance uses the same data segments as the NLB benchmark.Source: §B.2 · p. 18
  • MC-Maze has instructed-delay reaches with varied barriers and trajectories, whereas RTT has continuous, variable-length point-to-point reaches with few repetitions.MC-Maze recordings support varied behavioral configurations and separate preparation from execution.Source: §B.2 · p. 18

C.1 Comparison with single-session baselines

  • Comparisons include Wiener Filter, MLP, and GRU baselines; NLB comparisons additionally include Auto-LFADS rates followed by a linear behavior decoder.Wiener Filter, MLP, and GRU predict behavior using a history window of neural activity.Source: §C.1 · pp. 18–19
  • NLB uses 5 ms neural binning, while other MP datasets use 10 ms bins matching the behavior sampling rate.Source: §C.1 · p. 19
  • Table 4 lists the hyperparameters used to train all single-session POYO models.Source: §C.1 · p. 19, Table 4 · p. 19

C.2 Robustness of Unit-identification

The unit-identification analysis assesses whether finetuning can recover pretrained unit embeddings, and the hyperparameter analysis reports stable single-session performance across a range of settings.

C.2 Robustness of Unit-identification

  • On sessions from two animals, the analysis tracks finetuned embeddings from random initialization against pretrained embeddings and evaluates cosine similarity and nearest-neighbor identification accuracy.The reported cosine similarity is averaged over units; the accuracy is the fraction whose true embedding is the nearest neighbor.Source: §C.2 · p. 19, Table 5 · p. 19, Figure 7 · p. 20
  • The authors conclude that unit-identification is robust and reliable for identifying units.They identify understanding the functional similarities of nearby units in embedding space as future work.Source: §C.2 · p. 19

C.3 Hyperparameter sensitivity

  • The authors report that single-session models were stable across a wide range of hyperparameters, based on validation-set performance.Figure 8 shows validation R2 scores for two datasets.Source: §C.3 · p. 19, Figure 8 · p. 20
  • They ran 300 trials per session, averaged performance by hyperparameter, and used the selected settings to train all 100 single-session models studied in Figure 2.The settings were selected using a subset of datasets.Source: §C.3 · pp. 19–20

Related papers

Loading 2310.16046v1…