[Paper Review] Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion

Dongyang Li, Wei Chen+4 more|arXiv (Cornell University)|Mar 12, 2024

Blind Source Separation Techniques4 citations

TL;DR

This paper proposes a novel EEG-based visual reconstruction framework using a plug-and-play EEG encoder, Adaptive Thinking Mapper (ATM), and a two-stage diffusion-based image generator. By aligning EEG embeddings with image features via contrastive learning and leveraging guided diffusion for image generation, the method achieves state-of-the-art performance in EEG-based image classification, retrieval, and reconstruction, demonstrating high temporal resolution and low cost for brain-computer interface applications.

ABSTRACT

How to decode human vision through neural signals has attracted a long-standing interest in neuroscience and machine learning. Modern contrastive learning and generative models improved the performance of visual decoding and reconstruction based on functional Magnetic Resonance Imaging (fMRI). However, the high cost and low temporal resolution of fMRI limit their applications in brain-computer interfaces (BCIs), prompting a high need for visual decoding based on electroencephalography (EEG). In this study, we present an end-to-end EEG-based visual reconstruction zero-shot framework, consisting of a tailored brain encoder, called the Adaptive Thinking Mapper (ATM), which projects neural signals from different sources into the shared subspace as the clip embedding, and a two-stage multi-pipe EEG-to-image generation strategy. In stage one, EEG is embedded to align the high-level clip embedding, and then the prior diffusion model refines EEG embedding into image priors. A blurry image also decoded from EEG for maintaining the low-level feature. In stage two, we input both the high-level clip embedding, the blurry image and caption from EEG latent to a pre-trained diffusion model. Furthermore, we analyzed the impacts of different time windows and brain regions on decoding and reconstruction. The versatility of our framework is demonstrated in the magnetoencephalogram (MEG) data modality. The experimental results indicate that our EEG-based visual zero-shot framework achieves SOTA performance in classification, retrieval and reconstruction, highlighting the portability, low cost, and high temporal resolution of EEG, enabling a wide range of BCI applications. Our code is available at https://github.com/ncclab-sustech/EEG_Image_decode.

Motivation & Objective

• To address the limitations of fMRI in brain-computer interfaces (BCIs), such as high cost and low temporal resolution, by developing a low-cost, portable, and high-temporal-resolution EEG-based visual decoding framework.
• To improve EEG-based visual reconstruction by introducing a novel EEG encoder, Adaptive Thinking Mapper (ATM), that achieves superior alignment with image embeddings.
• To enable high-fidelity image reconstruction from EEG signals using a two-stage diffusion-based generative model conditioned on EEG embeddings.
• To investigate the impact of different brain regions and time windows on EEG-based visual decoding performance.
• To demonstrate the versatility and generalization of the framework across EEG and MEG data modalities.

Proposed method

• The proposed framework uses a plug-and-play EEG encoder, Adaptive Thinking Mapper (ATM), which integrates an attention module and a spatiotemporal convolutional module to extract real-time EEG and MEG representations.
• The ATM encoder is trained via a self-supervised contrastive learning framework to align EEG embeddings with CLIP-based image embeddings, enhancing representation quality.
• A two-stage image generation strategy is employed: first, a prior diffusion model generates image embeddings conditioned on EEG features; second, an enhanced Stable Diffusion model reconstructs the final image from these EEG-conditioned priors.
• The method enables zero-shot image retrieval and classification using EEG embeddings without fine-tuning on downstream tasks.
• The framework is evaluated on both EEG (THINGS-EEG) and MEG (THINGS-MEG) datasets, demonstrating robustness across modalities and subjects.
• Ablation studies validate the contribution of each component, including the attention and spatiotemporal modules in ATM.

Experimental results

Research questions

• RQ1 Can EEG-based visual reconstruction achieve state-of-the-art performance comparable to fMRI-based methods while maintaining low cost and high temporal resolution?
• RQ2 How does the performance of EEG-based visual decoding vary across different brain regions and time windows during stimulus processing?
• RQ3 To what extent can a plug-and-play EEG encoder like ATM improve alignment between EEG and image embeddings for downstream tasks?
• RQ4 Can a two-stage diffusion-based generation strategy effectively reconstruct natural images from EEG signals with high fidelity?
• RQ5 Does the proposed framework generalize across EEG and MEG data modalities?

Key findings

• The proposed ATM-S model achieves a top-1 accuracy of 28.50% and top-5 accuracy of 60.39% on zero-shot retrieval for even-numbered subjects on the THINGS-EEG dataset, outperforming all baselines.
• On subject-independent retrieval, ATM-S achieves an overall top-1 accuracy of 13.65% and top-5 of 37.34%, significantly exceeding prior methods like EEGNetV4 (8.69% top-1) and NICE (8.81% top-1).
• For subject-dependent retrieval, ATM-S achieves 25.60% top-1 accuracy on Subject 1 and 30.50% on Subject 7, surpassing the next best method (NICE) by over 10 percentage points.
• The two-stage diffusion generation strategy enables vivid and perceptually accurate image reconstruction from EEG signals, with qualitative results showing high fidelity to original stimuli.
• Ablation studies confirm that both the attention module and spatiotemporal convolution in ATM are essential for optimal performance, with performance dropping significantly when either is removed.
• The framework generalizes to MEG data, achieving SOTA performance on the THINGS-MEG dataset, confirming its versatility across neuroimaging modalities.