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

In plain words

Brain signals recorded with EEG are expensive to label, so most models that read them learn from one small, labelled dataset. This paper borrows an approach from speech and language AI. First, a model is trained on a very large pile of unlabelled EEG recordings, using a guessing task about hidden pieces of the signal. Then the same model is reused for specific tasks, including sleep staging and several BCI datasets.

The model first shrinks the raw signal into a short sequence of learned vectors. A transformer then reads that sequence. During pre-training, some vectors are hidden, and the model has to pick the correct vector out of a set of wrong ones. The authors report that the learned features carried over to recordings from new people, new hardware and new tasks, and that a simple classifier on top often did best.

TL;DR

BENDR adapts the wav2vec 2.0 speech design to multi-channel EEG. A convolutional encoder turns raw signals into 512-dimensional vectors at about 2.67 Hz, and an 8-layer transformer reads them §2.3. Pre-training on 60-second TUEG windows masks spans of vectors and scores the true vector against 20 distractors with a contrastive loss §2.4.1 eq. 2. Fine-tuning on five public tasks gave the best author-reported means with a linear head on pooled features in four of five tasks §3.2 tab. 2. The official code exists, but it departs from the paper in several places (see Paper vs code).

Problem

Deep networks for brain-computer interfaces (BCI) usually learn from one subject or session, and features do not transfer well across subjects, sessions or tasks §1. Labelling BCI trials is costly, so large labelled pre-training sets are impractical, unlike ImageNet in vision §1. Self-supervised learning could use unlabelled EEG instead §1.

Applying masked-token pre-training from language to raw EEG directly is likely to teach interpolation, because neighbouring samples are highly correlated §1. Reconstructing masked spans with mean squared error or mean absolute error is also problematic for time series, because those losses treat errors at neighbouring points as independent §1. The authors therefore use a contrastive objective, following contrastive predictive coding §1.

The earlier contrastive EEG work for sleep staging used windows of two minutes or more, which the authors consider too long for BCI trials that last seconds §1. This paper aims at shorter time scales §1.

Method

BENDR has two stages §2.3. The first is a convolutional encoder that produces the BENDR sequence. The second is a transformer encoder that maps that sequence to a task representation.

1. Preprocess. Each sequence is resampled to 256 Hz. The authors keep 19 channels of the unambiguous 10/20 set, add one channel carrying relative amplitude, and set missing channels to zero, giving 20 channels §2.2. Each sequence is also scaled to the range minus one to one §2.2.
2. Encode. Six 1D convolution blocks, each followed by GroupNorm and GELU, use kernel widths and strides of 3, 2, 2, 2, 2, 2 and 512 filters. The total downsampling is 96, which gives about 2.67 Hz §2.3.
3. Mask. For each sample, the chance of starting a masked span is 0.065, and each span covers 10 contiguous vectors. The code applies this to BENDR positions (see Pitfalls). Spans may overlap. A single learned mask vector of the same length as a BENDR vector replaces the masked vectors §2.4.1.
4. Contextualize. A grouped position convolution with receptive field 25 and 16 groups comes before the transformer. The transformer has 8 layers, 8 heads, model dimension 1536 and feed-forward dimension 3076, with GELU, T-Fixup initialisation and no internal batch normalisation §2.3. A start vector filled with the value minus five is prepended §2.3.
5. Score. At each masked position t, the transformer output c_t is compared by cosine similarity with the true BENDR vector b_t and with 20 distractors drawn uniformly from the same sequence, using temperature 0.1 §2.4.1 eq. 2. The loss adds the mean squared BENDR activation with weight 1 §2.4.1.
6. Evaluate pre-training. Contrastive accuracy is measured with half the masking rate, with masked spans spread evenly through the sequence §2.4.1.
7. Fine-tune. Six configurations are compared §2.4.2:
   • (1) a new softmax layer on the start-token output, with the whole model trained;
   • (2) average pooling of the BENDR sequence to four vectors and a new classifier, with the transformer skipped;
   • (3) configuration 1 without pre-training;
   • (4) configuration 1 with the BENDR sequence held fixed while the transformer keeps training;
   • (5) configuration 2 without pre-training;
   • (6) configuration 2 with the first stage frozen, training only the classifier.

Pre-training loss, as the paper defines it. Citations are to the paper; the code differs (Paper vs code).

# Pre-training loss for one batch  [§2.3] [§2.4.1] [eq. 2]
# Input:  eeg (B, 20, 15360), 60 s at 256 Hz  [§2.2]
# Output: scalar loss

def pretraining_loss(eeg):
    # Encode raw EEG into BENDR vectors, one per 96 samples  [§2.3]
    bendr = encoder(eeg)                              # (B, 512, 160)

    # Choose masked spans: start rate 0.065, span of 10 vectors  [§2.4.1]
    mask = sample_spans(bendr.shape, rate=0.065, span=10)

    # Replace masked vectors with the one learned mask vector  [§2.4.1]
    masked = where(mask, mask_vector, bendr)

    # Position convolution, projection, then transformer  [§2.3]
    context = transformer(masked)                     # (B, 512, 160)

    # Contrastive loss at each masked position t  [§2.4.1] [eq. 2]
    losses = [
        contrastive_loss(context[t], bendr[t], sample_distractors(t))
        for t in masked_positions(mask)
    ]

    # Aggregation over masked positions is not stated  [§2.4.1]
    contrastive = aggregate(losses)

    # Add the mean squared BENDR activation with weight 1  [§2.4.1]
    return contrastive + 1.0 * mean(bendr ** 2)

# Contrastive loss at one masked position  [§2.4.1] [eq. 2]
# Temperature 0.1; 20 distractors drawn from the same sequence  [§2.4.1]

def contrastive_loss(predicted, true_vector, distractors):
    # Cosine similarity to the true vector, divided by the temperature  [eq. 2]
    positive = cosine(predicted, true_vector) / 0.1
    # The same for each distractor  [eq. 2]
    negatives = [cosine(predicted, d) / 0.1 for d in distractors]
    # Negative log of the softmax weight on the true vector  [eq. 2]
    partition = exp(positive) + sum(exp(n) for n in negatives)
    return -log(exp(positive) / partition)

Downstream readouts, as the paper describes them:

# Downstream classifier, configuration 1  [§2.4.2]
# Input: trial (B, 20, T)

def classify_with_transformer(trial):
    # Encode the trial into BENDR vectors  [§2.3]
    bendr = encoder(trial)
    # Run the pre-trained transformer with the start vector prepended  [§2.3]
    context = transformer(bendr)
    # Read the classifier input from the start-vector output  [§2.4.2]
    return linear(context[:, :, start_position])

# Downstream classifier, configuration 2  [§2.4.2]
# Input: trial (B, 20, T)

def classify_pooled(trial):
    # Encode the trial into BENDR vectors; the transformer is skipped  [§2.4.2]
    bendr = encoder(trial)                            # (B, 512, T')
    # Average-pool the sequence to four vectors and concatenate  [§2.4.2]
    pooled = adaptive_avg_pool(bendr, outputs=4)      # (B, 512, 4)
    return linear(flatten(pooled))                    # (B, 2048)

Hyperparameters

Pre-training values come from the paper unless the row says otherwise.

Per-dataset fine-tuning settings. The paper's values are in tab. 2 (windows), tab. 1 (folds) and app. A (batch, epochs, learning rate). Code values are in configs/downstream_datasets.yml.

Windows are shown as start to start plus length. All datasets match except P300, whose code window is -0.1 s to 0.6 s, not -0.7 s to 1.3 s.

Results

Author-reported. Setup: pre-training on TUEG v1.1 and v1.2, about 1.5 TB of EDF before preprocessing §2.1. Downstream tasks are the five datasets in tab. 1. The baseline is the same architecture with random initialisation (configurations 3 and 5) §2.4.2 §3.2.

The table reports means over subjects and folds for each dataset's best configuration tab. 2. Configuration 2 was the highest on four of five datasets. Configuration 1 was close, with overlapping confidence intervals §3.2 fig. 3. The full model without pre-training (configuration 3) was generally ineffective, except on sleep staging §3.2. For sleep staging, fine-tuning the pre-trained model was roughly on par with random initialisation §3.2. The authors say their sleep-staging results with configurations 1 and 2 beat the two contrastive schemes of Banville et al., but this paper does not give those numbers §3.2.

Pre-training generalization: contrastive accuracy varied little across subjects and was similar across the four datasets other than P300 §3.1 fig. 1. Accuracy rose with longer evaluation sequences, and P300 was lower because its evaluation context was shorter, at 20 s §3.1 fig. 2. The authors read this as evidence that the features generalise to new subjects, hardware and tasks, while noting that their use downstream remains to be seen §3.1.

Paper vs code

Code: https://github.com/SPOClab-ca/BENDR at ac918abaec111d15fcaa2a8fcd2bd3d8b0d81a10. The repository has no LICENSE file and no license statement in its README, so code is cited by path and line only, with no verbatim snippets. Downstream training uses DN3 (the paper's own library, recommended at version 0.2 in the README). DN3 is not in the repository, so behaviour attributed to it is marked as unverified.

• Loss scope. Paper: the loss is defined for masked positions §2.4.1 eq. 2. Code: the contrastive loss is computed at every position and averaged, and the mask affects only the input dn3_ext.py dn3_ext.py dn3_ext.py. The mask is returned only for the Mask_pct metric dn3_ext.py. The paper's figures may use a different accuracy; this is not verified.
• Distractors. Paper: 20 distractors uniformly sampled from the same sequence §2.4.1. Code: indices are drawn with randint(0, full_len - 1), so the last position is never a distractor dn3_ext.py. Candidate exclusion of the target position is commented out dn3_ext.py dn3_ext.py. The guard that masks equal distractors compares the prediction with the distractors, not with the target vector dn3_ext.py.
• Possible in-place aliasing. The contextualizer writes the mask vector into its input tensor in place dn3_ext.py. Negatives are then drawn from that same tensor dn3_ext.py. Whether masked positions are affected depends on whether DataParallel copies its input, which multi_gpu: True may trigger pretraining.yml. Not run; check before reproducing.
• Start vector placement. Paper: the start vector is prepended to the BENDR sequence §2.3. Code: it is added after the position convolution and projection, in the 1536-dimensional transformer space dn3_ext.py. Pre-training drops it before the loss dn3_ext.py.
• Readout for configuration 1. Paper: the classifier reads the first output position, the start token §2.4.2. Code: it reads the last sequence position dn3_ext.py, while the start token sits at index 0 dn3_ext.py. The paper's configuration 1 therefore differs from the code.
• Encoder gradient factor. Config sets encoder_grad_frac: 0.1 pretraining.yml. It is applied through a backward hook on the encoder dn3_ext.py. The paper does not mention this factor. A module backward hook scales gradients with respect to the module's inputs, not its weights, so the factor may not reach encoder parameters. This follows from PyTorch hook semantics and was not run.
• Pooling head. Paper: configuration 2 pools the BENDR sequence §2.4.2. Code: the encoder output first passes through the augmentation module, which is a LayerNorm, a 1x1 convolution to 1536 channels and the position convolution dn3_ext.py dn3_ext.py. The pooled tensor is therefore 1536 channels, not 512. The classifier is a Linear(6144 to 2048) with dropout 0.4, ReLU and BatchNorm dn3_ext.py. None of this is in the paper.
• Frozen augmentation module. When pre-trained weights are loaded into the linear head, all augmentation-module parameters are set to not require gradients dn3_ext.py. The paper says configuration 2 trains "only the first stage and new layer" §2.4.2. The code therefore freezes the input projection and position convolution, which the paper does not state.
• Mask span in configuration 2. Paper: 10% of the BENDR sequence for all six configurations §2.4.2. Code: the linear head uses 5% dn3_ext.py dn3_ext.py. The BENDR classifier uses 10% dn3_ext.py.
• Pre-training hyperparameters not in the paper. Learning rate 2e-5, betas (0.9, 0.98), batch 64, one epoch, and the 0.05 random crop pretraining.yml pretrain.py. The paper states no pre-training learning rate or batch size.
• Validation during pre-training. BCI IV 2a is used as the validation set and is not trained on pretrain.py pretraining.yml. The paper does not mention a validation set. The config sets limit: 10000 for TUEG v1.1 pretraining.yml; how DN3 applies this is not verified. The paper does not state how many recordings it used.
• Fine-tuning batch sizes and P300 window. Paper app. A gives batch 80 for P300 and 64 for sleep staging. Code uses 64 for P300 downstream_datasets.yml and 80 for sleep staging downstream_datasets.yml. The paper's tab. 2 window for P300 is -0.7 s for 2 s. The code uses -0.1 s for 0.7 s downstream_datasets.yml.
• Scaling. Paper: each sequence is scaled to the range minus one to one §2.2. Code: fixed data_min and data_max values per dataset are set pretraining.yml downstream_datasets.yml. The scaling is applied by DN3, which is not in this repo. Not verified.
• P300 filtering. Paper: P300 is low-pass filtered below 120 Hz §2.2. Code: the filter was applied offline, and the in-code option is commented out downstream_datasets.yml.

Pitfalls

• Units of the mask. Span and rate are counted in BENDR positions, not raw samples, in both paper and code. A span of 10 is about 3.75 s at 2.67 Hz §2.3 §2.4.1. The paper says "contiguous sequences of 10" without units.
• Sequence length. A 60-s window gives 160 encoder positions §2.3 §2.4.1. This is derived from the 96x downsampling: 15,360 / 96 = 160.
• Positional encoding. The position convolution covers about 9 s of encoded input. The paper says this may weaken position information for short sequences §2.3 §4.
• Losses and metrics. Accuracy and loss in code are computed over all positions, not just masked ones dn3_ext.py. Expect loss values that differ from a masked-only definition.
• Aggregation. The paper does not say how per-position losses combine. The pseudocode leaves this open.
• Data. TUEG is not included. Reproduction needs access to TUEG v1.1 and v1.2 or the pre-trained weights from the v0.1-alpha release [README.md#L16-L17]. The README recommends DN3 version 0.2 [README.md#L12].
• Model selection. Fine-tuning retains the best validation checkpoint on the first metric downstream.py utils.py. The paper does not describe validation or early stopping.
• Variability. Results are averages over subjects, and per-subject variability is large for classifiers §3.2 fig. 3. The MMI results are discrete because each subject has few trials fig. 3.
• Hidden costs. The reference config sets multi_gpu: True pretraining.yml. The transformer runs 160 steps of width 1536 per window §2.3.

Lineage

• Builds on: wav2vec 2.0 (Baevski et al., 2020), whose architecture and masking this paper follows §2.3 §2.4.1. The source names it but does not give an arXiv id, and it has no card in this wiki.
• Also draws on: contrastive predictive coding (Oord et al., 2018) §1; BERT's masked-language objective (Devlin et al., 2018) §1; T-Fixup (Huang et al., 2020) §2.3; LayerDrop (Fan et al., 2019) §2.3. None of these has a card.
• Compared against: contrastive sleep-staging self-supervision (Banville et al., 2019) §1 §3.2; fully supervised and randomly initialised models of the same architecture §3.2.
• Superseded by: not known from this wiki.