BENDR: using transformers and a contrastive self-supervised learning task to learn from massive amounts of EEG data

Demetres Kostas, Stephane Aroca-Ouellette, Frank Rudzicz

Seminal January 2021

TL;DR BENDR adapts wav2vec 2.0 to EEG, pretraining a convolutional encoder and a transformer with a masked contrastive task on about 1.5 TB of unlabelled clinical recordings. The single pretrained model transfers to five downstream EEG tasks recorded with other hardware and subjects, reaching 86.7 balanced accuracy on left/right motor imagery.

What they do

EEG classifiers generalize poorly across subjects, sessions and tasks, and labelled EEG is scarce. BENDR pretrains one network on a large unlabelled clinical EEG corpus with a self-supervised task taken from speech recognition (wav2vec 2.0), then fine-tunes it on a battery of BCI and sleep-staging datasets recorded with different hardware, subjects and tasks.

How it works

Data. Pretraining uses the Temple University Hospital EEG Corpus: over 10,000 people, about 1.5 TB of recordings, cut into 60 s sequences. Every dataset is resampled to 256 Hz and mapped to the same 19 channels of the 10/20 layout plus one relative-amplitude channel.
Encoder. Six 1D convolution blocks downsample raw EEG by 96x into a sequence of 512-dimensional vectors at about 2.67 Hz. These vectors are the BENDR.
Sequence model. An 8-layer, 8-head transformer takes the BENDR with contiguous spans of 10 masked. Its output at each masked position must be more similar to the true BENDR vector than to 20 distractors drawn from the same sequence.
Fine-tuning. Six variants are compared: the full model or a linear classifier on average-pooled BENDR, each either pretrained, randomly initialized, or pretrained with the encoder frozen.

Results

Accuracy on the contrastive task for unseen downstream data is nearly uniform across datasets and subjects, and rises as the evaluation sequence grows from 20 to 60 s.
The fine-tuned linear classifier on pretrained BENDR is best on four of five datasets: 86.7 balanced accuracy on MMI motor imagery, 42.6 accuracy on 4-class BCIC, 0.65 AUROC on ERN and 0.72 balanced accuracy on sleep staging. The full pretrained model is best on P300 (0.72 AUROC). The two always have overlapping confidence intervals.
The full model without pretraining is generally ineffective, except on the larger sleep dataset.
On sleep staging, both fine-tuned models reach higher mean balanced accuracy than the two task-specific contrastive schemes of Banville et al.

Limitations

Most results are not state of the art. The 1D convolutions mix channels and only 19 are used, so it is unclear how much spatial information BENDR captures. The rise in contrastive accuracy with sequence length suggests neighbouring BENDR vectors still hold redundant information. The convolutional position encoder spans a little over 9 s, longer than every downstream trial except sleep staging.

Figures

From Figure 1: Violin plot of test-subject-wise accuracy on the contrastive task for each downstream dataset, i.e. ho... — From Figure 1: Violin plot of test-subject-wise accuracy on the contrastive task for each downstream dataset, i.e. how often the transformer stage creates a representation closest to the correct one at masked sequence positions. P300 is lower, likely due to its shorter evaluation context; otherwise there is minimal test-subject-wise variation.

From Figure 2: Contrastive accuracy versus evaluation length in seconds (x-axis logarithmic). Performance is similar ... — From Figure 2: Contrastive accuracy versus evaluation length in seconds (x-axis logarithmic). Performance is similar for all datasets, rising for longer sequences.

From Figure 3: Performance of all downstream datasets for each of the six model configurations, normalized to range f... — From Figure 3: Performance of all downstream datasets for each of the six model configurations, normalized to range from chance (0) to perfect (1). Translucent points are single subjects, solid diamonds are the mean across all subjects and folds, with .95 confidence intervals. Configurations (1.) and (2.) were consistently among the best performing, yet both remained within the confidence levels of each other.

Abstract

Deep neural networks (DNNs) used for brain-computer-interface (BCI) classification are commonly expected to learn general features when trained across a variety of contexts, such that these features could be fine-tuned to specific contexts. While some success is found in such an approach, we suggest that this interpretation is limited and an alternative would better leverage the newly (publicly) available massive EEG datasets. We consider how to adapt techniques and architectures used for language modelling (LM), that appear capable of ingesting awesome amounts of data, towards the development of encephalography modelling (EM) with DNNs in the same vein. We specifically adapt an approach effectively used for automatic speech recognition, which similarly (to LMs) uses a self-supervised training objective to learn compressed representations of raw data signals. After adaptation to EEG, we find that a single pre-trained model is capable of modelling completely novel raw EEG sequences recorded with differing hardware, and different subjects performing different tasks. Furthermore, both the internal representations of this model and the entire architecture can be fine-tuned to a variety of downstream BCI and EEG classification tasks, outperforming prior work in more task-specific (sleep stage classification) self-supervision.

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