Representation learning for neural population activity with Neural Data Transformers

Joel Ye, Chethan Pandarinath

Seminal August 2021

TL;DR The Neural Data Transformer (NDT) replaces the recurrent network in LFADS-style models of neural spiking with a BERT-style Transformer encoder trained by masking timesteps. It matches AutoLFADS on monkey reaching data (0.918 vs 0.915 R2 for hand-velocity decoding) while inferring in 3.9 ms, 6.7x faster.

What they do

Models that turn noisy spike trains into single-trial firing rates, such as LFADS, rely on recurrent networks. These process a window of data step by step, which slows real-time uses like brain-computer interfaces. The Neural Data Transformer (NDT) is a non-recurrent alternative: a Transformer encoder that reads the whole window of binned spikes in parallel and outputs a firing rate for every channel and time bin.

How it works

Architecture. A BERT-style Transformer encoder (6 layers in most experiments, 1 or 2 attention heads) with learned position embeddings. A linear layer outputs log firing rates, which are exponentiated and scored against the observed spikes with a Poisson likelihood.
Masked training. A random subset of timesteps is masked and the model must reproduce the spikes there from the unmasked context. Masked inputs are zeroed out instead of being replaced by a special mask token.
Heavy regularization. Dropout ratios are swept between 0.2 and 0.6, which matters most for small datasets.
Evaluation. NDT is compared with AutoLFADS on two synthetic autonomous systems (Lorenz, chaotic RNN) and on recordings of 202 motor and premotor neurons from a monkey doing a delayed reaching task (2296 trials).

Results

Synthetic data: inferred rates match ground truth with R2 of 0.934 (Lorenz) and 0.846 (chaotic RNN), against 0.921 and 0.869 for AutoLFADS.
Monkey reaching: linear decoding of hand velocity from inferred rates reaches 0.918 R2 for NDT and 0.915 for AutoLFADS.
Speed: a forward pass takes 3.9 ms on a CPU and grows only slightly with sequence length, against 26 ms for LFADS on 70 bins (6.7x faster).
Smaller models: a 2-layer NDT trains in 45 minutes, infers in 2.2 ms and decodes at 0.921 R2; trained on only 92 trials it reaches 0.866 R2, ahead of AutoLFADS. The 6-layer model took 9.4 hours to train.

Limitations

NDT cannot model non-autonomous dynamics, where unpredictable inputs from unrecorded areas perturb the population: LFADS outperforms it by over 20% R2 in a preliminary test on a synthetic dataset with inputs. Timing was measured offline, not in a closed-loop interface. Without log-rate outputs and strong dropout, performance is much worse and more variable. The Transformer also lacks the recurrent structure that dynamical-systems analyses of RNNs rely on.

Figures

From Figure 1: Sequential vs. parallel models. (a) Unsupervised models of sequential spiking activity take in binned ... — From Figure 1: Sequential vs. parallel models. (a) Unsupervised models of sequential spiking activity take in binned spikes and output inferred rates. A likelihood loss trains the network to output the most likely rates. (b) A Transformer architecture (top) performs parallel modeling, contrasting with RNNs (bottom) which use sequential processing. (c) Spike input and inferred rate examples for the NDT applied to the reaching dataset.

From Figure 2: Transformer architecture. (a) A single Transformer layer; the full encoder stacks several of these lay... — From Figure 2: Transformer architecture. (a) A single Transformer layer; the full encoder stacks several of these layers. (b) Inputs to Transformer layers are normalized (Norm blocks), enriched through contextual information (Self-Attention blocks), and passed through a feedforward module (MLP blocks). (c) Inputs at each time step are multiplied by three learned weight matrices to create the queries, keys, and values. (d) A single timestep output is a weighted sum of value vectors, with weights from the similarity between that timestep query and every key.

From Figure 3: Transformer training. The model is trained with masked modeling, that is, model outputs are optimized ... — From Figure 3: Transformer training. The model is trained with masked modeling, that is, model outputs are optimized to maximize likelihood of the masked activity given the context provided by unmasked activity.

Abstract

Neural population activity is theorized to reflect an underlying dynamical structure. This structure can be accurately captured using state space models with explicit dynamics, such as those based on recurrent neural networks (RNNs). However, using recurrence to explicitly model dynamics necessitates sequential processing of data, slowing real-time applications such as brain-computer interfaces. Here we introduce the Neural Data Transformer (NDT), a non-recurrent alternative. We test the NDT's ability to capture autonomous dynamical systems by applying it to synthetic datasets with known dynamics and data from monkey motor cortex during a reaching task well-modeled by RNNs. The NDT models these datasets as well as state-of-the-art recurrent models. Further, its non-recurrence enables 3.9ms inference, well within the loop time of real-time applications and more than 6 times faster than recurrent baselines on the monkey reaching dataset. These results suggest that an explicit dynamics model is not necessary to model autonomous neural population dynamics. Code: https://github.com/snel-repo/neural-data-transformers

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