MEG-Mamba: A Scalable State-Space Foundation Model for Magnetoencephalography

Abstract

Magnetoencephalography (MEG) is an imaging technique that offers a non-invasive, millisecond-resolution view of human brain activity. The increasing availability of MEG data presents an opportunity to take advantage of a recent advance in artificial intelligence, namely self-supervised foundation models. Existing foundation models for MEG (and electroencephalography) have been built on a transformer architecture. Here, we introduce MEG-Mamba: a generative foundation model for neural activity (source reconstructed, parcellated MEG) built on a Mamba architecture. MEG-Mamba is trained to predict the next token of a discretised MEG signal, conditioned on a brain region and recording session. MEG-Mamba surpasses the generative fidelity of a transformer-based alternative (MEG-GPT) while pre-training in less time (22 vs 400 GPU-hours) and modelling a longer context (4 vs 0.32 s). We evaluate MEG-Mamba by examining the spatio-spectral characteristics of the data it generates and by interpreting its learned embeddings. Furthermore, we demonstrate that lightweight stimulus conditioning (LoRA) can be used to generate realistic task-evoked responses that are not included in pre-training. Our results suggest that Mamba is a promising backbone for scaling neural foundation models.