RPNT: Robust Pre-trained Neural Transformer -- A Pathway for Generalized Motor Decoding

TL;DR

RPNT tackles robust generalized motor decoding from neural activity by pretraining a neural transformer with three innovations: Multidimensional rotary positional embedding (MRoPE) that encodes recording configurations, a context-based attention mechanism that learns local temporal structure to manage non-stationarity, and a robust self-supervised objective using uniform random masking with Poisson reconstruction. The model is pretrained on two diverse datasets (public microelectrode benchmark and Neuropixel recordings) and evaluated on cross-session, cross-task, cross-subject, and cross-site decoding tasks, showing consistent improvements over strong baselines, including NDT/POYO variants. Key findings include superior $R^2$ scores under various generalization settings, data-efficient few-shot finetuning, and interpretable spatial-attention maps that reveal functional connectivity patterns. The work advances neural foundation modeling for neural decoding and has implications for robust, scalable brain-computer interfaces, while acknowledging limitations such as focus on motor cortex and single-modality data with plans for multimodal extensions.

Abstract

Brain decoding aims to interpret and translate neural activity into behaviors. As such, it is imperative that decoding models are able to generalize across variations, such as recordings from different brain sites, distinct sessions, different types of behavior, and a variety of subjects. Current models can only partially address these challenges and warrant the development of pretrained neural transformer models capable to adapt and generalize. In this work, we propose RPNT - Robust Pretrained Neural Transformer, designed to achieve robust generalization through pretraining, which in turn enables effective finetuning given a downstream task. In particular, RPNT unique components include 1) Multidimensional rotary positional embedding (MRoPE) to aggregate experimental metadata such as site coordinates, session name and behavior types; 2) Context-based attention mechanism via convolution kernels operating on global attention to learn local temporal structures for handling non-stationarity of neural population activity; 3) Robust self-supervised learning (SSL) objective with uniform causal masking strategies and contrastive representations. We pretrained two separate versions of RPNT on distinct datasets a) Multi-session, multi-task, and multi-subject microelectrode benchmark; b) Multi-site recordings using high-density Neuropixel 1.0 probes. The datasets include recordings from the dorsal premotor cortex (PMd) and from the primary motor cortex (M1) regions of nonhuman primates (NHPs) as they performed reaching tasks. After pretraining, we evaluated the generalization of RPNT in cross-session, cross-type, cross-subject, and cross-site downstream behavior decoding tasks. Our results show that RPNT consistently achieves and surpasses the decoding performance of existing decoding models in all tasks.

Figures (6)

• Figure 1: Overall illustration of the pretraining and finetuning workflow for generalized motor decoding. (A) Experimental setup for data collection while NHPs performed reaching tasks. (B) Preparation of pretaining data. (C) and (D) overall schemes for SSL and SFT, respectively. (E) Illustration of model adaptation during pretraining and finetuning.
• Figure 1: Spatial attention map provided data-driven brain network insights for motor behaviors.
• Figure 2: A schematic of components in RPNT. Components in black indicate standard transformer signal flow (i.e, no masking and standard attention mechanism). Our novel proposed components include MRoPE (green), context-based attention (cyan), and uniform random masking strategy (pink). MRoPE incorporates experimental metadata into the transformer positional embeddings. The context-based attention captures local temporal structures. The input spike activity is randomly masked according to a uniform distribution for causal masked neural modeling. The masked input spike activity is projected into neural embeddings, forming Q, K, and V representations. MRoPe is applied to get the rotary queries and keys. Attention scores are modulated by context-based convolutional kernels and combined with causal masking before aggregating with values.
• Figure 2: Neuropixel data collection setup. (A) Chamber location and neuropixel probe. (B) Neuropixels probe insertion locations (black dots) on the cortical surface. Each black dot stands for a different recording site.
• Figure 3: Decoding performance comparison across sites (B-CS) on the neuropixel dataset. The mean and standard deviation were calculated across all successful trials.