Retrospect - Lightweight Neural Network for EEG Source Imaging
Sep 2, 2024
MoNET Lab에서 수행했던 딥러닝 프로젝트인 “Lightweight Neural Network for EEG Source Imaging”에 대한 회고이다. 개인으로 진행하였으며, 코드는 이 레포지토리에서 확인 가능하다.
문제 정의
측정된 EEG를 바탕으로 source region 및 spike timing을 추정하는 task(EEG inverse problem)이다.
더 간단하게 표현하자면 EEG를 바탕으로 뇌의 활동을 추정하는 task라고 할 수 있고, 더 포멀하게 표현하자면 intracranial EEG(invasive)를 scalp EEG(non-invasive)로부터 추정하는 task이다.
수행 과정
재구현 및 데이터 준비
DeepSIF의 모델 아키텍쳐를 발전시키려 했고, 로컬에서 재구현부터 시도해 보았다.
(iEEG, EEG) labelled dataset은 매우 희귀하므로, neural mass model로 생성하였다. 이 과정은 매우 CPU-heavy함과 동시에 디스크 용량을 많이 사용하였으므로, 후처리 파이프라인도 함께 리팩토링한 후에 병렬화하여서 약 16개의 job으로 분리하고, 후처리 후 raw data는 지우는 방식으로 작업하였다.
모델 성능 개선
원 논문의 아키텍쳐(residual linear block as spatial encoder, 3x bi-LSTM)를 구현해 보았고, 약 22M 파라미터 정도가 있었다.
모델 성능 개선에 있어 해본 생각들은 다음과 같다.
iEEG Discretization
channel 수나 위치가 비교적 표준화되어있는 scalp EEG와 달리 iEEG는 그렇지 않다. source region의 신호를 어떻게 잘 나타내야 할까?
20484-region fsaverage5 템플릿을 우선 사용하였다. 하지만 20484 채널은 neural network가 직접 다루기에는 VRAM의 한계가 명확하므로, 원래 방법처럼 994채널 매핑을 사용하였다.
하지만, 원래 제공되었던 forward matrix가 EEG 채널 정보를 포함해 모든 것이 모호했기에, forward calculation부터 직접 수행하여 학습 과정에 그 결과를 참조할 수 있게끔 하였다. 물론, 이 과정에서 매핑 정보는 그대로 사용하였다.
Parameter Reduction
사실 이 task는 이미 수학적으로 잘 풀린 task이다. 꼴의 간단한 linear relationship(forward calculation; with leadfield matrix)을 가지고 있다. 이러한 종류의 relationship은 상당히 잘 연구되어 있으며, 예를 들어 digital communication에서도 사용된다.
한정된 데이터셋에서 DeepSIF의 방법인 994채널 bidirectional LSTM 3개가 그 파라미터 수에 비해 잘 학습될까? 라는 의문이 들었다.
Importing Time-series Models
neural network의 관점에서, 다른 일반적인 multichannel time-series task의 접근을 적용해 볼 수 있을까?
하지만 다른 task와 달리, 이 task는 채널 수가 매우 많이 변한다. 60채널 내외에서 994채널로. 채널 수가 매우 많으므로, attention 기반 모델은 코스트가 매우 크다. 특히, 994채널의 경우 상당히 폭력적인 메모리 사용량을 보일 것이다.
또한, time-series task의 경우 vision 쪽과는 달리 transformer 기반 모델들이 우위를 보이지 않는다.
그리고, 데이터셋 자체가 최신 SOTA들의 트렌드? 라고 볼 수 있는 long-term dependency가 잘 나타나지 않는다. 그도 그럴 것이, 전처리 과정에서 spike 단위로 잘라내어서, 각 spike는 1초 내외의 길이를 가지고 있다. 이는 대부분의 task에서는 short-term dependency로 볼 수 있다.
즉, 채널 수, 시간 길이, 파라미터 수에 대한 제약이 매우 크다.
결론
결론적으로 모델 아키텍쳐는 다음과 같이 변경하였다. 파라미터 수는 5M 정도로 줄었다.
- spatial encoder는 동일한 구조(residual linear block)을 사용한다.
- temporal encoder로 GRU를 사용한다. 이는 LSTM보다 적은 파라미터를 가지고 있으며, 채널 수를 확장하는데 적합하다.
- temporal translator로 LightTS를 사용하였다. 이 구조는 linear 블록으로만 이루어져 있으면서도 time-series task에서 SOTA에 근접한 성능을 보이며, patchmixer 등등의 long-term dependency에 과도하게 집중하는 모델들보다 범용적이다.
실험
사실, 생성한 데이터셋에 대한 정량적 비교분석을 진행하려 했지만, 직접 만든, 그것도 실제 데이터셋이 아니라 무엇을 반영할지도 모르는 데이터셋에서 메트릭을 추출해내는 것은 큰 의미가 없다고 판단하였다. 물론, DeepSIF에 비해서는 개선점이 있기 때문에 기록하진 않았지만, validation loss 상에서의 향상은 있었다.
그렇다고 (iEEG, EEG) label이 있는 실제 데이터셋도 없는 노릇이기에, 특정 상황에서 측정된 EEG를 바탕으로 source region이 잘 추정되는지 정성적으로 확인해보기로 하였다. 사용한 데이터셋은 Auditory Oddball Task를 수행한 환자의 EEG 데이터셋이며, auditory cortex에서 활성이 나타남을 확인할 수 있었다.
넋두리
spatial encoder 구조 역시 개선이 가능하지 않을까? 사실 이 태스크 자체가 forward calculation에서는 temporal dimension이랑 아예 무관하게 계산이 되는 건 맞는데, 그것을 역산해내기 위해 temporal context를 추정한다…는 느낌이라서 정보의 많은 부분이 spatial encoder에서 계산된다 하더라도 아주 틀린 말은 아니라고 생각한다.
또한, 데이터셋과 메트릭이 명확해서 리더보드를 만들 수도 있는 대부분의 태스크들과 다르게, 이 태스크는 어떤 모델이 좋은지 판단하기가 상당히 어려웠다.
그리고 살짝 장난식으로 말하자면, 컴퓨터 성능은 항상 부족하다는 것을 항상 깨닫는다.