EEGDM: EEG Representation Learning via Generative Diffusion Model

This work moves away from the typical masked-reconstruction pretraining employed in EEG foundation models (FMs), thereby avoiding the introduction of discontinuities—such as patching or tokenization—into fundamentally continuous EEG signals. Instead, a score-matching objective from diffusion models is utilized for pretraining, and the latent activities of the pretrained diffusion backbone are leveraged for EEG representation.

Two unique challenges were encountered and resolved:

• ⚠️ Noisy, dynamic EEG signals and Gaussian noise injection lead to unstable variance.
  💡 A cosine noise schedule and μ-law companding were proposed for variance control.
• ⚠️ Emergent representations are scattered across the diffusion backbone.
  💡 A Latent Fusion Transformer was proposed for effective representation extraction from the diffusion backbone.

As a result, our proposed method demonstrates state-of-the-art performance alongside strong cross-domain and cross-task generalization in epilepsy-related tasks. It outperforms existing foundation models while using a smaller model size and only a single pretraining dataset.

Benchmarking on a single Nvidia RTX 4090 demonstrates the high efficiency of EEGDM:

• Pretraining: 23 hours of training on 250 hours of EEG data.
• Finetuning: ~1 minute per epoch, enabled by latent caching with a frozen backbone.
• Inference: 30 ms latency per sample without parallelization.

In contrast to FMs, our approach does not require high-grade hardware or tremendous training data, while still achieving high performance and strong generalization.

This component is responsible for learning robust representations from raw EEG signals. The method draws inspiration from Generative Diffusion Models, which are a class of self-supervised learning models. Instead of the typical masking-based methods, SSMDP learns to reverse a process of noise injection.

The core of SSMDP is a Structured State-Space Model (SSM) architecture. This architecture is particularly well-suited for capturing the temporal dynamics of EEG signals. During pretraining, the model is tasked with progressively "denoising" the EEG signal to restore the original data. The resulting latent activities and representations from this process are then used for downstream tasks.

Unlike autoencoders, which are optimized to output representations, our diffusion-based pretraining encourages representations to spontaneously emerge across the entire backbone. Consequently, this requires special care to collect and condense these scattered representations.

LFT is a transformer-based architecture designed to take the learned latent representations from the SSMDP and fuse them for classification tasks. Because the latent representations are high-dimensional, they are first pooled to reduce computational complexity.

The LFT has a "latent fusion module" that uses a decoder-only transformer to aggregate the information from different EEG channels and layers. This module uses "fusion tokens" to create a context-aware feature for each time window of the signal. The "classification module" then takes these fused representations and processes them through an encoder-only transformer, finally using a linear classification head to make predictions. This fine-tuning phase is trained in a supervised manner on a specific task.

In this section, we present the main results of the proposed EEGDM framework in downstream tasks. EEGDM achieved new state-of-the-art results in multi-event classification, highlighting the effectiveness of its diffusion-based representation learning. These findings demonstrate that EEGDM offers a promising and computationally efficient alternative, outperforming existing methods while being approximately 19x more lightweight.