Generalizable, real-time neural decoding with hybrid state-space models
Avery Hee-Woon Ryoo, Nanda H Krishna, Ximeng Mao, Mehdi Azabou, Eva Dyer, Matthew Perich, Guillaume Lajoie
Real-time decoding of neural spiking data is a core aspect of neurotechnology applications such as brain-computer interfaces, where models are subject to strict latency constraints. Traditional methods, including simple recurrent neural networks, are fast and lightweight but are less equipped for generalization to unseen data. In contrast, recent Transformer-based approaches leverage large-scale neural datasets to attain strong generalization performance. However, these models typically have much larger computational requirements and are not suitable for settings requiring low latency or limited memory. To address these shortcomings, we present POSSM, a novel architecture that combines individual spike tokenization and an input cross-attention module with a recurrent state-space model (SSM) backbone, thereby enabling (1) fast and causal online prediction on neural activity and (2) efficient generalization to new sessions, individuals, and tasks through multi-dataset pre-training. We evaluate our model's performance in terms of decoding accuracy and inference speed on monkey reaching datasets, and show that it extends to clinical applications, namely handwriting and speech decoding. Notably, we demonstrate that pre-training on monkey motor-cortical recordings improves decoding performance on the human handwriting task, highlighting the exciting potential for cross-species transfer. In all of these tasks, we find that POSSM achieves comparable decoding accuracy with state-of-the-art Transformers, at a fraction of the inference cost. These results suggest that hybrid SSMs may be the key to bridging the gap between accuracy, inference speed, and generalization when training neural decoders for real-time, closed-loop applications.