1. Introduciton

MDD(Major Depsresssion Disorder)는 흔한 정신장애이며 비이상적인 행동 패턴과 생각을 가지고 계속해서 슬픔, 무의미함, 희망없음 등의 기분을 느낀다. 이는 자해, 자살까지 이어지고 WHO에 의하면 매해 백만이 목숨을 앗아간다고 한다. 그래서 MDD의 조기 진단은 중요하다. 전형적인 MDD 진단 방법에는 행동이나 정신 상태를 외부에서 관찰하거나 DSM, BDI, PHQ 같은 설문지가 있다. 하지만 이런 방법들은 의료 전문가에 의존하기 때문에 주관적인 진단을 할 수 있다. 사실, 외부 현상과는 반대로 MDD 환자들은 비이상적인 뇌활동패턴을 보인다. 예를 들어, MDD환자는 더 낮은 뇌활동을 보인다. 뇌활동은 일반적으로 비침습 방법인 EEG같은 방법으로 전극활동을 기록하여 측정된다. EEG는 시간에 따른 뇌파의 진폭 변화를 기록한다.

그림 1은 MDD환자와 건강한 사람 간의 다중채널 EEG를 ICA 방법으로 분리한 결과이다. MDD 환자의 뇌파는 건강한 사람보다 덜 자주 극적으로 변하는것으로 보아 뇌활동이 보다 덜 활동적인 걸 알 수 있다. EEG를 더 분석하면, 다른 주파수 밴드에서의 에너지 분포, 기능적인 영역, 반구 사이의 관계 등 많은 뇌 활동이 환자와 건강한 사람 사이에서 다르게 나타난다는 것을 볼 수 있다.

하지만, 데이터의 양과 질이 EEG로 MDD를 진단하는데 큰 문제가 된다.
(1) EEG 데이터 수집에서 많은 양의 노이즈는 raw EEG로 부터의 특징 추출이 어렵다. EEG 데이터 수집에서 환경(눈 깜빡임, 심장박동) 및 기계 노이즈는 피할 수 없다. 그림 1에서 볼 수 있듯이 두 서브젝트는 노이즈를 가지고 있다(b). 윗부분은 사인파 noise를 가지고 있고 아래 부분은 급작스러운 움직임의 노이즈가 있다. noise를 제거하는 기법인 ICA와 밴드필터링이 사용되지만 사람의 검사에 의존해 시간이 많이 걸리고 일반화 되지 않으며 노이즈를 완전히 제거할 수 없다. 이 문제를 해결하기 위해서 automatic한 방법이 제안되어야 모델의 노이즈에 대한 강건성이 늘어날것이다.

(2) 모델 학습을 위한 충분하고 다양한 데이터를 수집하기 위해 많은 수의 피험자를 구하기 어렵다. 사생활과 데이터 클리닝의 고비용 문제로 MDD 진단에 데이터셋의 양은 작다. 같은 피험자의 clipped EEG 샘플들은 밀접하게 관련있고 다양성은 부족하다. 기존 연구들은 타당한 결론을 내었지만, 여전히 적은양 데이터는 모델의 성능을 억제한다.

언급한 두 문제는 오버피팅을 야기할 수 있는 문제점도 있다. 기존 MDD 진단을 위한 딥러닝 방법들은 직접 학습에 초점을 두고 성능은 한정적이다.

위 문제들을 해결하기 위해, 우리는 diffusion 기반 MDD 진단 프레임워크를 제시한다. diffusion의 forward와 reverse는 MDD 진단 성능을 향상시키기 위해 사용되고 통합되었다.
우리의 컨트리뷰션은 아래와 같다.

• 우리느 3-step diffusion 프레임워크를 제시한다. 두 개의 실제 공공데이터셋에서 subject-independent 방법으로 10-fold cross validation 성능을 얻었다.
• 우리는 forward diffusion 노이즈 훈련 모듈을 설계하여 모델의 강건성을 향상시키기 위해 더 많은 노이즈와 상관없는 특징을 추출한다.
• 우리는 reverse diffusion 데이터 증분 모듈을 설계하여 모델이 더 generalized features를 배울 수 있게 데이터의 양과 다양성을 향상시다.

2. Related Work

A. MDD Diagnosis by EEG Using Deep Learning

초기단계의 딥러닝 기반 MDD 진단 방법은 MLP, PNN과 같은 수공예 방법에 의존한다. 최근의 CNN기반 모델들은 짧은 기간의 특징을 추출합니다. 예를 들어 DeprNet은 ConvNet 구조를 사용하고 5겹의 CNN 레이어로 특징을 포착한다. InceptionNet은 다른 사이즈의 커널을 사용하여 다른 시간대의 특징을 포착하고 채널의 중요성을 학습하기 위해 높은 단계의 레이어에 channel-wise attention을 사용한다. [Ref 8]은 spectral feature를 먼저 추출하고 time-frequency map을 2D-CNN에 넣는다. 더 최근에는, LSTM이 긴 기간의 특징을 포착하는데 좋기 때문에 하이브리드 구조를 쓰기도 한다. [Ref 18, 24]는 CNN-LSTM 구조를 사용하여 단기간과 장기간의 특징을 모두 포착하기 위한 모델을 만들었다. [Ref 25]는 GCN과 GRU 모델을 쌓아 뇌파의 공간-시각적 특징을 포착하고 adpative 방식으로 뇌 연결 그래프가 그려졌다. 하지만, 기존 연구들은 원본 데이터를 직접 학습하고 성능이 적은 양의 데이터 때문에 억제된다. 위의 딥러닝 모델들은 학습중에 오버피팅되기 쉽고 학습하지 않은 피험자에 대해 일반화에실패한다. 우리는 CNN-Transformer 분류기로 특징을 장단기 특징을 인코딩한다. 하지만 이전 연구들과 다른점은 Diffusion 기반 모델을 사용하여 데이터의 성능과 양 문제를 해결한다는 것이다.

B. Denoising Diffusion Probabilities Models

이 작업에서, 데이터의 품질과 희소성 문제를 해결하기 위해 forward and reverse diffusion을 각각 noisy 학습, 데이터 증분으로 사용한다. 이들은 일반적인 DDPM의 과정이다. DDPM은 forward, reverse diffusion chain으로 이루어진 생성 모델의 한 종류이고 다양한 분야에서 SOTA의 성능을 보인다. DDPM은 GAN보다 안정적이고 더 높은 품질을 생성한다.

하지만 원본 reverse diffusion은 step by step 으로 동작하고 많은 시간을 소비한다. DDIM과 같은 몇 연구는 reverse process를 조건부확률 공식이나 noise parameter를 수정하여 빠르게 하였다. DDPM을 사용한 Guided, 조건부 데이터 생성에 초점을 둔 연구도 있다. 분류기를 사용하여 reverse process를 guide하는 것과 학습중에 기울기를 수정하는 것이 수학적으로 동등하다는 것이 증명되었다.

우리는 forward diffusion 을 노이즈 학습에 사용하고 DDIM의 유도 버전의 reverse diffusion을 data 증분에 사용한다. 일반적인 classifier-guided diffusion과 다른 점은, 우리는 학습된 guide and condition diffusion model을 사용한다. 이 방법에서, 두 diffusion 모듈을 통합되어 있고 분류기의 출력 특징은 reverse diffusion에 더 많은 정보를 줄 수 있다. 게다가, 일반적인 U-Net을 사용하는 것이 아닌 공간 트랜스포머 블록을 설계하여 EEG 채널의 공간 관계를 포착합니다. Diffusion 기반 모델과의 결합으로, 모델은 noise에 상관없이 더 일반화된 특징을 학습하고 보지 못한 피험자들에 대한 일반솨 성능이 향상됩니다.

3. Methodology

A. Problem Definition

우리는 MDD 진단 문제를 2진 분류로 공식화한다. $x_{sub}\isin R^{C\times N}$을 원본 EEG 기록이라 하면, C는 EEG 채널의 수이고 C=19이다. N은 샘플링 포인트의 수이다. 각 $x_{sub}$는 0(=건강 ),1(=MDD)의 정답 라벨($\hat{y}$)이 있다. 이후, 우리는 각 원본 EEG를 같은 크기로 n=1280 개의 샘플링 포인트로 자르고 $x_{sub}\isin R^{C\times n}$이라 표시한다. 각 샘플은 정답 라벨이 있다($\hat{y}_{sub}$). 우리의 MDD 진단 작업은 sample x가 0 or 1 중 어느것에 가장 가까운지 매핑하는 분류기 f를 학습하는 것으로 정의할 수 있다.

B. Overview

forward와 reverse diffusion을 통합한 오버뷰이다. 세 가지 모듈로 이루어져 있다.
(1) Forward Diffusion Noisy Training, (2) Reverse Diffusion Data Augmentation, (3) Re-training and Classification.

첫 번째로 (1)은 데이터 품질 문제를 해결하기 위해 설계되었다. 우리는 모델이 noise와 상관 없이 특징을 학습하도록 우리는 forward diffusion 과정을 적용하고 Gaussian noise를 주입하여 분류기를 초기에 학습했다.
두 번째로 (2)는 새로운 EEG 샘플을 생성하여 데이터의 희소성 문제를 해결하기 위해 설계되었다. 생성 품질과 학습 과정이 안정적이게 하기 위해 초기 훈련된 분류기의 출력 기울기로 확산 모델을 안내하는 것 외에도 분류기의 출력 뇌파 임베딩으로 확산 모델을 조건한다. 특히, Diffusion에서 1D-CNN-트랜스포머를 사용하여 시간적 특징을 포착하고, 다른 채널 간의 관계를 포착하는 반구 및 기능적 임베딩 영역을 가지는 공간 트랜스포머를 사용한다. 높은 품질과 다양성을 가진 생성 데이터는 분류기가 더 일반화되게 학습하도록 도울 수 있다.
마지막으로, 우리는 증분된 EEG 데이터로 분류기를 재학습하고 MDD 진단 결과를 얻는다. 각 모듈들은 서로 밀접하게 연관되어 있으며 MDD 진단의 오버피팅 문제를 완화한다. 모듈들은 분류기가 noise-irrelavant하고 generalized된 특징을 학습하도록 하고 unseen 피험자들에 대해 더 나은 성능을 하도록 regularize한다.

C. Forward Diffusion Noisy Training

EEG는 수집중 불가피한 환경 소음과 안구 운동 때문에 신호 대 노이즈 비율이 낮은 특성을 가집니다. 기존의 EEG를 직접 학습한 MDD 진단 모델은 소음에 민감하고 raw EEG에서 특징을 추출하는데 실패했을 것이다. 그러모르, 우리는 forward diffusion 모듈이 분류기를 regularize 하도록 하고 noisy EEG에 대해서 noise-irrelevant 한 특징을 학습하도록 돕는다.

우리의 (1)에서 각 샘플에 대해 먼저 가우시안 noise를 timestamp t에 대하여 점진적으로 넣는다. time-stamp의 최대치는 1000이다. 우리가 가우시안 noise를 주입하기로 한 이유는 이론적으로 reverse 과정에서 guide하기 더 좋기 때문이다. 게다가, EEG 수집중의 노이즈의 일반적인 유형이기 때문이다.

이후, 우리는 가우시안 noise가 부여된 EEG 샘플 $x_t$를 사용하여 우리의 분류기$f_{noisy}$를 학습하는데 사용한다. Fig.3 은 재학습 후 분류하는 과정이다.

Fig. 3.

D. Reverse Diffusion Data Augmentation

소음이 많은 것 외에도 뇌파를 이용한 MDD 진단을 위한 데이터가 부족한 것도 심각한 과적합 문제를 일으키는 중요한 원인입니다. EEG 수집과 clearning의 비용과 건강관리의 사생활 문제 때문에 데이터셋의 크기는 일반적으로 작으며, 특히 강력한 딥러닝 모델일수록 더 그렇다. 그래서 우리는 이 모듈을 데이터 증분을 위해 설계하고, 분류기가 더 일반화된 특징을 학습하도록 한다. 처음 학습된 분류기 $f_{noisy}$는 MDD 진단에 상=대적으로 의존하고 diffusion model을 conditioning과 guiding하기에 유용하다는 것을 발견한다. 이런 관점에서, 우리의 diffusion model은 원본 샘플과 대응되는 새로운 EEG 샘플과 라벨을 생성할 수 있고 guidance는 생성 데이터의 품질을 향상시킨다.

게다가, EEG의 다른 채널과의 관계성이 있고 이 관계성은 뇌의 기능적 패턴을 나타낸다. 게다가 동일한 기능 영ㅇ역 및 반구의 뇌파 채널은 일반적으로 강하게 관련되어 있습니다. 이러한 특징은 MDD 환자와 건강한 사람을 구분하는데 유용하다. channel들의 연관성을 포착하기 위해, 우리는 diffusion model에 시간 1D-CNN-트랜스포머에서 떨어진 공간 트랜스포머를 설계한다.

Fig. 4.

Reverse diffusion은 noise가 부여된 $x_t$로부터 노이즈를 제거하는 것을 예측한다. 우리의 diffusion model $\epsilon_0$은 아래 그림과 같다. diffusion model은 $x_t,t,x_t^{EEG}$를 입력으로 받고 non-guided 예측 noise를 출력한다. 이 아키텍처는 CNN, 시간 및 공간 트랜스포머 계층 및 Residual layer로 이루어진다. 동일한 기능 영역과 반구의 뇌파 채널이 일반적으로 강하게 관련되어 있기 때문애 공간 트랜스포머에 각 채널의 위치 임베딩 표현을 위해서 뇌 기능 영역 임베딩과 반구 임베딩 계층을 연결한다. 우리는 non-guided predicted noise $\epsilon_\theta(x_t,x_t^{EEG},t)$ 와 정답 noise $\epsilon$ 사이의 MSE loss를 최소화한다. K는 noise가 주입된 샘플의 개수이다.

학습 후에, 무작위로 샘플링 된 각 노이즈 주입 뇌파 샘플 $x_t$를 classifier-guided reverse diffusion 을 통해 $\hat{x}_0$를 DDIM을 기반으로 추론한다. 아래는 마지막으로 얻게되는 reverse diffusion process 수식

$\hat\epsilon_\theta$는 $x_t$를 통해 예측된 noise이고, y는 $x_t$의 정답 라벨이다. DDIM은 샘플링 스텝을 스킵할 수 있기 때문에 t의 최대값을 100으로 두었다.

E. Re-Training and Classification

마지막으로, 우리는 분류기를 reverse diffusion module로부터 얻은 데이터를 사용하여 재학습시킨다. 분류기는 Forward diffusion noisy training에서 처음 학습되고, 데이터 입력은 같은 형태이다. 이 모듈에서 t=0일때 원본 이미지, 합성 EEG 샘플은 $x_t$이다. 우리는 CNN을 아래와 같이 정의한다. $Conv^l, BatchNorm^l, MaxPooling^l$은 각각 convolutional layer, batch normalization layer, max-pooling layer로 이루어져 있다. $x_t^l$은 l번째 CNN 레이어의 출력이다.

CNN layer L의 수는 2로 설정되어 있고 $x_t^L\isin R^{19\times32}$이다. 이후 CNN 계층의 출력을 표준 트랜스포머 인코더 계층에 넣고 모든 단계를 임베딩한다. 핵심 모듈은 classic multi-head self-attention 매커니즘이다. $head_i$는 i번째 head이고 $Q_i, K_i, V_i$는 i번째 head의 Query key value이다. 그리고 $W^O$는 head projection matrix이다.

은닉차원 $d_k=8$, head count $H=4$이고 시간 모델은 마지막 단계에서 많은 정보를 통합하기 때문에 마지막 단계의 추력은 $x_t^{EEG} \isin R^{32}$를 사용한다. 우리는 또한 시간 스텝에 따른 $x_t^{noise}\isin R^{32}$ 를 얻기 위해 임베딩 레이어 $T_{emb}$ 를 사용한다. 우리가 사용하는 cross-entropy loss는 아래와 같다. m은 noise가 주입된 EEG 샘플이고, $\hat{y}_j$는 j 번째 noise 부여 샘플의 정답 라벨이다. 그리고 $y_j$는 재학습된 분류기의 예측 점수 출력이다.

4. Experiments

A. Datasets and Preprocessing

두 개의 public MDD 진단 데이터셋, Mumtaz2016, Arizona2020을 사용함. Mumtaz2016 데이터셋은 34명의 MDD 환자와 30명의 건강한 피험자의 EEG 기록을 포함한다. 피험자는 HUSM 병원에서 고용되었다. Arizona2020 dataset은 총 121명의 참여자로 구성되어 있다. EEG는 New Mexico 대학에서 수집되었고 참여자들은 Arizonsa state 대학의 BDI점수에 대한 설문좃가를 통해 모집되었다. 이 실험에서는 23명의 MDD 환자와 19명의 가장 낮은 BDI 점수의 건강한 참여자들을 데이터의 밸런스와 품질을 위해 선택했다. 두 데이터셋을 위해, 두 데이터셋에서 EEG 데이터 수집중에 피험자는 모두 resting state였다. EEG 전극은 10-20 system에 따라 배치되었고 19개의 채널을 사용하였다. : Fp1, F3, C3, P3, O1, F7, T3, T5, Fz, Fp2, F4, C4, P4, O2, F8, T4, T6, Cz and Pz.

EEG 신호는 256Hz로 resampling되었다. 0.5Hz 고주파필터와 50Hz 저주파 필터가 환경적 noise 삭제에 사영되었다. Arizona2020 데이터셋에서, 낮은 퀄리티의 기간은 EEGLAB을 사용하여 삭제되었다. 모든 EEG 기록은 5분동안 이루어졌고 5초 길이의 샘플로 2.5초의 stride로 잘렸다.
마지막으로, Mumtaz2016 데이터셋은 14,035 sample (7,223 긍정, 6,812 부정) 으로 이루어진다.
Arizona2020 데이터셋은 10,582 sample (5,710 긍정, 4,872 부정) 으로 이루어진다.

B. Experimental Setup

1) Performance Metrics

동일한 피험자의 샘플이 서로 밀접한 관련이 있다는 사실을 고려서 데이터 leakage 문제를 해결하기 위해 10-Fold subject indepent 교차검증법을 사용하였다. (한 피험자는 Train, Test 중 하나만 속한다)

우리는 이진 분류 지표인 sample-wise ACC, F1 score, recall and precision를 사용한다. 또한 정확하게 진단된 피험자의 비율을 나타내는 subject-wise accuracy도 소개한다. 한 피험자의 50%가 잘 분류되었다면, 정확하게 진단되었다고 본다.

2) Implementation

우리는 파이토치를 사용하여 모델을 구현한다. 우리의 모델은 $10^{-3}$
의 학습률로 50 epochs 동안 첫번째와 두번째 모듈이 학습되었고 마지막 모듈에서 $10^{-4}$의 학습률로 변경되었다. 우리는 early stopping 전략을 도입하여 오버피팅을 피했다. 우리는 배치사이즈를 128로 하였고 drop rate을 0.5로 두었다. 분류기와 디퓨전 모델의 트랜스포머 블록은 4개의 헤드와 32개의 은닉층을 가진다. 우리는 모델을 Intel Core i9 10900K CPU and 4 NVIDIA RTX 3090 GPU로 학습한다.

3) Compared Methods

우리는 같은 학습 전략을 사용하여 다른 MDD 진단 방법과 우리 모델을 비교한다.

LR (Logistic Regression), XGBoost, 1D-CNNLSTM, 1D-CNN-Transformer, CWT-1D-CNN, CWT-2D-CNN, EEGNet, DeprNet, GC-GRU, TSception, GRU-Conv