BrainBERT: Self-supervised representation learning for intracranial recordings

Christopher Wang, Vighnesh Subramaniam, Adam Uri Yaari, Gabriel Kreiman, Boris Katz, Ignacio Cases, Andrei Barbu

Seminal February 2023

TL;DR A Transformer pretrained by masked spectrogram reconstruction on 43.7 hours of unannotated intracranial recordings from 10 subjects. A linear decoder on its fine-tuned embeddings reaches 0.83 average ROC-AUC on four audio and speech tasks, against 0.59 for a linear decoder on the raw signal.

What they do

Linear decoders of neural recordings are explainable but weak; deep decoders are stronger but may solve the task themselves instead of reading it from the brain. BrainBERT sits in between: a Transformer is pretrained without labels on stereo-EEG (SEEG) recordings, one electrode at a time, and a linear decoder is then trained on its embeddings. Because it is subject- and electrode-agnostic, it can be used off the shelf on new subjects.

How it works

Data. 10 epilepsy patients (aged 4-19) watched feature-length movies while 1,688 electrodes were recorded at 2 kHz. After cleaning, 1,249 electrodes (4,551 electrode-hours) are cut into 5 s segments and pooled for pretraining.
Input. Each segment becomes a spectrogram, either STFT or the superlet transform, z-scored per frequency bin so that no band dominates.
Pretraining. A 6-layer Transformer encoder (12 heads, hidden size 768) reconstructs randomly masked time and frequency bands. For superlets the mask width adapts to frequency, and a content-aware loss adds weight on spectrogram elements far from zero.
Decoding. A linear classifier reads the mean embedding over about 244 ms, with BrainBERT either frozen or fine-tuned.

Results

Four tasks on held-out sessions (sentence onset, speech vs. non-speech, pitch, volume), on the 10 electrodes where the raw linear decoder does best: fine-tuned BrainBERT (STFT) averages 0.83 ROC-AUC, against 0.63 for a 5-layer deep network and 0.59 for the linear decoder. Random initialization gives 0.60.
Frozen BrainBERT (superlet) averages 0.63, on par with the deep network.
Holding a subject out of pretraining changes that subject’s decoding performance negligibly.
On one superior temporal gyrus electrode, BrainBERT matches with 150 examples what the best baseline reaches with 1,000.

Limitations

The released weights were pretrained only on passive movie viewing; the authors recommend pretraining from scratch if the downstream setting departs significantly from it. Decoding is evaluated on a subset of electrodes per task to limit compute, and all four tasks are binary classifications of the movie audio. The content-aware loss and adaptive mask each add only about 0.01 AUC.

Figures

From Figure 1: (a) Intracranial electrode locations (yellow dots) across all subjects. (b) Subjects watched movies wh... — From Figure 1: (a) Intracranial electrode locations (yellow dots) across all subjects. (b) Subjects watched movies while neural data was recorded. (c) Recordings are converted to spectrograms and embedded with BrainBERT for downstream tasks such as sample-efficient classification. (d) In pretraining, BrainBERT learns embeddings from which a masked spectrogram can be reconstructed.

From Figure 2: Spectrograms of one electrode over a 5 s interval, computed with the short-time Fourier transform (top... — From Figure 2: Spectrograms of one electrode over a 5 s interval, computed with the short-time Fourier transform (top left) or the superlet transform (bottom left). Random frequency bands and time intervals are masked (red rectangles); for superlets the mask adapts to the temporal resolution, which falls off as the inverse of frequency.

From Figure 5: Sentence onset classification on one electrode in the superior temporal gyrus (red dot) as a function ... — From Figure 5: Sentence onset classification on one electrode in the superior temporal gyrus (red dot) as a function of training examples. Linear decoders saturate quickly; BrainBERT without fine-tuning matches the deep network, and with fine-tuning it exceeds the deep network peak at 1,000 examples with only 150 examples.

From Figure 3: Sentence onset classification with a linear decoder on the neural recordings directly (left) or on Bra... — From Figure 3: Sentence onset classification with a linear decoder on the neural recordings directly (left) or on BrainBERT superlet embeddings (right). Each circle is one of the 947 held-out electrodes, colored by ROC-AUC (chance is 0.5). BrainBERT highlights far more relevant electrodes, in superior temporal and frontal regions.

Abstract

We create a reusable Transformer, BrainBERT, for intracranial recordings bringing modern representation learning approaches to neuroscience. Much like in NLP and speech recognition, this Transformer enables classifying complex concepts, i.e., decoding neural data, with higher accuracy and with much less data by being pretrained in an unsupervised manner on a large corpus of unannotated neural recordings. Our approach generalizes to new subjects with electrodes in new positions and to unrelated tasks showing that the representations robustly disentangle the neural signal. Just like in NLP where one can study language by investigating what a language model learns, this approach opens the door to investigating the brain by what a model of the brain learns. As a first step along this path, we demonstrate a new analysis of the intrinsic dimensionality of the computations in different areas of the brain. To construct these representations, we combine a technique for producing super-resolution spectrograms of neural data with an approach designed for generating contextual representations of audio by masking. In the future, far more concepts will be decodable from neural recordings by using representation learning, potentially unlocking the brain like language models unlocked language.

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