[논문 리뷰] Visual Decoding and Reconstruction via EEG Embeddings with Guided Diffusion

Dongyang Li, Wei Chen외 4명|arXiv (Cornell University)|2024. 03. 12.

Blind Source Separation Techniques인용 수 4

한 줄 요약

이 논문은 플러그 앤 플레이형 EEG 인코더인 적응형 사고 매핑 기반(Adaptive Thinking Mapper, ATM)과 이중 단계의 확산 기반 이미지 생성기를 사용한 새로운 EEG 기반 시각 재구성 프레임워크를 제안한다. 대조적 학습을 통해 EEG 임베딩을 이미지 특징과 정렬하고, 유도된 확산을 이용한 이미지 생성을 통해, EEG 기반 이미지 분류, 검색 및 재구성에서 최신 기술 수준(SOTA) 성능을 달성하며, 뇌-컴퓨터 인터페이스(BCI) 응용 분야에서 높은 시간 해상도와 낮은 비용을 실현한다.

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.

연구 동기 및 목표

• fMRI가 뇌-컴퓨터 인터페이스(BCI)에서 가지는 높은 비용과 낮은 시간 해상도 등의 한계를 해결하기 위해, 저비용, 이식 가능, 고시간 해상도의 EEG 기반 시각 해독 프레임워크를 개발한다.
• 새로운 EEG 인코더인 적응형 사고 매핑 기반(Adaptive Thinking Mapper, ATM)을 도입하여 EEG 임베딩과 이미지 임베딩 간의 우수한 정렬을 달성함으로써 EEG 기반 시각 재구성 성능을 향상시킨다.
• EEG 임베딩을 조건으로 하는 이중 단계의 확산 기반 생성 모델을 활용하여 EEG 신호로부터 고해상도 이미지 재구성을 실현한다.
• 자극 처리 중 다양한 뇌 영역과 시간 창이 EEG 기반 시각 해독 성능에 미치는 영향을 조사한다.
• 프레임워크가 EEG 및 MEG 데이터 모odalities 간에 어떻게 일반화되는지 확인한다.

제안 방법

• 제안된 프레임워크는 플러그 앤 플레이형 EEG 인코더인 적응형 사고 매핑 기반(Adaptive Thinking Mapper, ATM)을 사용하며, 이는 주의 메커니즘 모듈과 시공간 컨볼루션 모듈을 통합하여 실시간 EEG 및 MEG 표현을 추출한다.
• ATM 인코더는 자기지도 학습 기반의 대조적 학습 프레임워크를 통해 EEG 임베딩을 CLIP 기반 이미지 임베딩과 정렬함으로써 표현 품질을 향상시킨다.
• 이중 단계의 이미지 생성 전략을 적용한다: 첫 번째 단계에서는 사전 확산 모델이 EEG 특징에 조건부로 이미지 임베딩을 생성하고, 두 번째 단계에서는 향상된 Stable Diffusion 모델이 이러한 EEG 조건부 사전 임베딩에서 최종 이미지를 재구성한다.
• 이 방법은 후행 작업에 대한 미세조정 없이도 EEG 임베딩을 사용해 제로샷 이미지 검색 및 분류를 가능하게 한다.
• 프레임워크는 EEG(THINGS-EEG) 및 MEG(THINGS-MEG) 데이터셋 모두에서 평가되어 모달리티 및 피험자 간에 강건성을 입증한다.
• 제거 실험을 통해 ATM 내 주의 모듈과 시공간 컨볼루션 모듈의 기여도를 검증한다.

실험 결과

연구 질문

• RQ1 fMRI 기반 방법과 유사한 최신 기술 수준의 성능을 달성하면서도 저비용과 고시간 해상도를 유지할 수 있는가?
• RQ2 자극 처리 중 다양한 뇌 영역과 시간 창에서 EEG 기반 시각 해독 성능는 어떻게 변화하는가?
• RQ3 ATM과 같은 플러그 앤 플레이형 EEG 인코더가 후행 작업에서 EEG와 이미지 임베딩 간 정렬을 얼마나 향상시킬 수 있는가?
• RQ4 이중 단계의 확산 기반 생성 전략이 EEG 신호로부터 고해상도 자연 이미지를 효과적으로 재구성할 수 있는가?
• RQ5 제안된 프레임워크는 EEG 및 MEG 데이터 모달리티 간에 일반화되는가?

주요 결과

• 제안된 ATM-S 모델은 THINGS-EEG 데이터셋에서 짝수 번호의 피험자에 대해 제로샷 검색에서 상위 1위 정확도 28.50%와 상위 5위 정확도 60.39%를 기록하여 모든 기준선을 초월한다.
• 피험자 간 검색에서 ATM-S는 전체적으로 상위 1위 정확도 13.65%와 상위 5위 정확도 37.34%를 달성했으며, 이는 이전 방법들인 EEGNetV4(8.69% 상위 1위)와 NICE(8.81% 상위 1위)를 크게 뛰어넘는 성능이다.
• 피험자 의존적 검색에서 ATM-S는 피험자 1에서 25.60%의 상위 1위 정확도, 피험자 7에서 30.50%의 상위 1위 정확도를 기록했으며, 이는 다음으로 우수한 방법(NICE)보다 10个百分点 이상 높다.
• 이중 단계의 확산 생성 전략은 EEG 신호로부터 생생하고 인지적으로 정확한 이미지 재구성을 가능하게 하였으며, 정성적 결과는 원본 자극과의 높은 유사도를 보여준다.
• 제거 실험을 통해 ATM 내 주의 모듈과 시공간 컨볼루션 모듈이 최적 성능을 내기 위해 필수적임을 확인하였으며, 둘 중 하나를 제거할 경우 성능이 크게 떨어진다.
• 프레임워크는 MEG 데이터로도 일반화 가능하며, THINGS-MEG 데이터셋에서 최신 기술 수준의 성능을 달성하여 다양한 신경영상 모달리티 간의 유연성을 확인한다.