Leveraging unlabelled data for generalizable neural population decoding

Introduction

The introduction discusses the importance of neural decoders in brain-computer interfaces and highlights the limitations of current spike-based models that rely solely on supervised learning. It emphasizes the potential of advanced architectures like Transformers and State-Space Models for improved decoding performance.

Methods

This section outlines the MOJO framework, detailing its dual pathways for self-supervised and supervised learning, and introduces the key components necessary for its implementation.

Spike Tokenization and Input Cross-Attention

The spike tokenization scheme assigns tokens to individual spikes, facilitating the processing of neural data. MOJO employs a specific encoder that computes cross-attention on time chunks, differing from previous models.

Backbones

MOJO utilizes sequence models as backbones, comparing attention-based and recurrence-based architectures for processing spike tokens, detailing how each architecture updates hidden states.

Masked Autoencoder

In the self-supervised learning pathway, a masked autoencoder is employed, applying temporal masking to latent outputs to enhance learning from unlabelled data.

Figures Explained

The paper’s visual material highlights the workflow and the main system components.

• Figure 1 :: Figure 1: Model and task schematics. (a) Schematic showing a POYO-style model augmented with MOJO. Latent representations extracted from tokenized neural data are simultaneously used for supervised learning (SL) and self-supervised learning (SSL). The former is carried out by minimizing error in predicted behaviour while the latter is carried out by reconstructing spike counts from masked latents. (b-e) Schematics describing the monkey reaching tasks that MOJO is evaluated on. (f) Schematic of the mouse visual and decision tasks. (g) Schematic of the human speech decoding task.
• Figure 2 :: Figure 2: Leveraging unlabelled data for finetuning and pretraining. (a) MOJO improves fewshot finetuning performance over standard SL and leverages additional unlabelled data to improve performance further. (b) MOJO can improve decoding performance by exploiting all available unlabelled data even when little labelled data is available.
• Figure 3 :: Figure 3: Mouse brain region classification and human speech decoding results. (a-b) Confusion matrices for (a) 3-class and (b) 7-class brain region classification of neurons in mice from the Allen visual coding dataset. (c) Classification accuracy of syllables, consonants, and vowels for speech decoding from human electrocorticography. For (a-c), we report the mean accuracy over sessions averaged across 5 seeds. (**,***): p < 0.01, 0.001 resp. for a paired t-test across all sessions. MOJO vs. Du-IN is n.s. for the t-test.
• Figure 4: are the projected queries, keys, and values, respectively. The process is repeated for T c time chunks to obtain the latent token sequence Z = [z (1) 1 , …, z (1) Nc , …, z (Tc) 1 , …, z (Tc) Nc ].
• Figure 4 :: Figure 4: Additional results on pretraining with unlabelled data. (a) Finetuning performance with additional unlabelled data during pretraining. (b) Brain region classification performance with unlabelled data during pretraining.