[논문 리뷰] An Attention-Based Deep Learning Approach for Sleep Stage Classification With Single-Channel EEG

 

1. Introduction

sleep stage는 인간의 면역체계, 기억력, 대사작용 등과 관련하여 중요한 요소이기 때문에 수면을 모니터링하고, sleep stage를 분류, 분석하여 sleep quality를 높이는 것이 필요하다. 현재 sleep stage를 측정하는 기본적인 방법으로는 수면 다원검사(PSG : polysomnography)가 있으며, 이는 EEG(Electroencephalogram), EOG(Electrooculogram), EMG(Electromyogram), ECG(Electrocardiogram)을 측정하여 수면에 대한 전반적인것들을 검사한다. 하지만 수면 다원검사는 비용적인 부담과 검사 방법의 불편함 등 여러 문제점이 있기 때문에, 보다 쉽고 간편하며 저렴한 방법으로 수면 관련 검사를 진행하고 sleep stage를 측정하는 것이 필요한 상황이다. 따라서 이 논문에서는 보다 간편하게 sleep monitoring을 수행하기 위한 방법으로 AttnSleep model(novel attention-based deep learning architecture)을 제안한다. AttnSleep model은 single channel EEG signals를 이용하여 sleep stage를 분류는것을 목표로 고안된 모델이다.

 

 

 

해당 모델은 크게 2개의 부분으로 구성되어있다.

1) Feature extraction

2) Sleep stage Classification

 

overview of AttnSleep Model

 

기존에 연구들에서 CNN으로 feature extraction를 진행한뒤 RNN으로 time dependency를 반영하는 방법이 제시되었지만, 기본 RNN모델은 병렬적으로 학습하기 어렵다는 제한점이 있다. 따라서 이 논문에서는 Attention mechanism 사용한다. EEG를 30초 단위로 분할한 데이터에 self-attention을 사용하여 intra-epoch feature와 inter-epoch temporal feature를 학습하게 되는 것이다.

 

 

 

---

2) Proposed Method

2-1) Feature Extraction 

AttnSleep Model에서 raw single channel EEG signals의 feature extraction을 위해

1) MRCNN(Multi-Resolution Convolution Neural Network) + 2) AFR(Adaptive Feature Recalibration)을 사용한다.

 

 

 

 

1️⃣ MRCNN(Multi-Resolution CNN)

MRCNN layer는 다른 kernel size에 따라 두 개의 갈래로 나눠진다. Kernel size는 EEG signals의 샘플링(sampling) 비율과 관련 있으며 이러한 구조를 통해 서로 다른 주파수 대역을 추출할 수 있게 된다. 서로 다른 사이즈의 커널을 사용하여 다양한 범위의 timestep을 캡처하고 다양한 수면 관련 주파수 대역의 기능을 처리한다.

 

예를 들어, sampling rate가 100Hz(1초에 100개의 timestep이 sampling 되는)인 데이터에 400 크기의 커널을 적용하면,  ~0.25까지 사인파 신호 주기가 추출된다. 반면, 커널 사이즈를 50으로 줄이게 되면 ~2Hz까지 사인파 신호 주기를 추출할 수 있다. 즉, 커널 사이즈가 더 작을수록 고주파 특성을 추출하고, 커널 사이즈가 클수록 저주파 특성 추출하는 것이다. EEG signal처럼 non- stationary 한 특성을 가지고 있는 데이터의 경우, 이렇게 다양한 방식으로 추출된 feature를 조합하는 것이 중요하다.

 

 

2️⃣ AFR(Adaptive Feature Recalibration)

AFR는 성능 향상을 위해 MRCNN에서 학습한 feature들을 재보정하는 역할을 한다.

특히, AFR은 feature들과 residual squeeze & excitation(residual SE) block을 통해 특징적으로 구별되는 feature들을 알맞게 선택한다.

 

 

 

3️⃣ Activation function - GELU, ReLU

보편적으로 사용되는 activation function 중 하나인 ReLU는 모든 음수 가중치를 0으로 만들기 때문에 AFR모듈에서 해당 가중치를 더 이상 사용할 수 없게 된다. 하지만 음수 가중치들이 AFR module에 중요할 수 있기 때문에 MRCNN에서는 GELU를 활성화 함수로 사용하여 input의 일부 음력 가중치가 통과할 수 있도록 한다. 하지만 ReLU는 쉽고 빠르게 수렴하고 gradient exploding/vanishing을 방지할 수 있기 때문에 AFR에서는 ReLU를 사용한다.

 

 

 

---

2-2) Temporal Context Encoder(TCE)

TCE layer에서는 input feature들의 temporal dependency(시간 종속성)를 기록하는 것을 목표로 하며,

이 layer는  1) Multi-head Attention(MHA), 2) normalization layer, 3) FC layer로 구성되어있다. 

 

 

1️⃣ Multi-head Attention

MHA는 2가지 측면에서 slef-attention기능을 향상하는데,

1) 각 head의 encoding이 다른 head의 encoding 또한 알고 있으며 각 head가 서로 다른 위치에 주목함으로써 모델의 기능이 향상된다. 특히 시간적 의존성을 학습해야 하는 해당 모델의 능력이 향상될 수 있다.

2) input feature들을 분할하여 넣으면 표현 부분 공간(representation subsapaces)이 증가한다. 따라서 각 subspace에 서 생성된 attention weight는 각 부분의 중요성을 더 잘 대표할 수 있게 되며, 이러한 representation을 연결하면 전체적으로 더 나은 representation이 된다.

 

뿐만 아니라, Attnsleep모델에서 MHA는 인과관계를 활용하여 입력 기능의 위치정보를 인코딩하고, 시간적 관계를 기록하는데, 이때 인과관계 컨볼루션(caual convolution)은 빠르고, 병렬 처리가 가능하기 때문에 RNN에 비해 모델 훈련시간을 크게 단축시킬 수 있다.

 

 

 

---

2-3) Class-Aware Loss Function & Optimization

 

Wk: 클래스 K에 할당된 가중치

μk : 조정 가능한 매개변수

Mk : 클래스 k의 샘플수

 

사람의 수면 단계에서 각 클래스 간의 데이터는 불균형이 있기 때문에 이를 해결하기 위해 class aware loss function을 사용한다. 만약 기본적인 multi-class cross-entropy 사용하면 모든 클래스의 오분류에 동등하게 페널티를 주기 때문에 훈련 모델은 데이터가 많은 클래스에 편향될 수 있다. 따라서 위의 문제를 해결하기 위해 가중 교차 엔트로피 손실을 사용하는 class - aware loss function을 사용하는 것이다.

 

 

---

 

Sleep data를 분석한 결과 두 가지 결론에 도달할 수 있는데,

 

1) sleep stage중, 클래스 N2는 수많은 샘플이 있는 반면, N1, N3 클래스에 해당하는 샘플의 수가 현저히 적었다.

2) N3의 신호는 다른 클래스보다 Amplitude의 크기가 훨씬 더 크다. 따라서 해당 모델은 N3샘플에 대해서는 정확하게 예측할 수 있다. 반면, N1과 N2, REM은 명확하게 구분하기 힘들다. 따라서 가장 높은 μk를 N1에 할당하고, 가장 낮은 값을 N3에 할당하며, 나머지 세 개의 클래스는 W(wake)와 유사한 값을 할당하기로 한다.

 

 

 

 

---

3. Experimental Results

3-1) Datasets and Evaluation Metrics

 

해당 논문에서 사용된 dataset은 총 3가지이다.

1) Sleep-EDF-20 

2) Sleep-EDF-78

3) Sleep Heart Health Study(SHHS) 

 

 

논문에 사용된 데이터

 

 

1️⃣  Data Preprocessing

1) UNKNOWN stages -> 제외시킨다.

2) AASM의 기준에 따라, N3와 N4를 N3으로 통합시킨다.

3) sleep stage에 조금 더 집중하기 위해 sleep periods 전, 후의 30분 동안 Wake periods를 포함한다.

 

 

 

2️⃣ Model evaluate

 

해당 연구에서는 모델의 성능을 평가하기 위해 다음과 같은 지표를 사용한다.

 

Accuracy(ACC)

macro-averaged F1-score(MF1)

Cohen Kappa

macro-averaged G-mean(MGm)

 

 

 

---

3-2) Scoring Performance of AttnSleep

〰️ Confusion-matrix : Sleep-EDF-20 Dataset

 

 

 

 

〰️ Confusion-matrix : Sleep-EDF-78 Dataset

 

 

 

〰️ Confusion-matrix : Sleep-EDF-78 Dataset

 

 

 

---

3-3) Baselines and Experimental Setup

본 연구에서 사용된 AttnSleep model과  sleep stage classification에 관련된 다른 모델들을 비교해보았다.

 

DeepSleepNet

SleepEEGNet

ResnetLSTM

MultitaskCNN

SeqSleepNet

 

 

 

 

- 20-fold cross-validation 적용 (각 dataset을 20개의 그룹으로 나눔)

- 모든 모델들을 동일하게 100 epochs -> training time 측정

- Batch size : 128

- Optimizer : Adam

- Learning rate : 1e-3

- Input dimension : d , Output dimension : 120

 

 

 

다른 모델들과의 성능 비교

 

 

 

 

---

3-4) Ablation Study

* Ablation Study : 모델이나 알고리즘을 구성하는 다양한 구성요소(component) 중, feature를 하나씩 제거해보면서 성능(performance)에 어떠한 영향을 미치는지 파악하는 방법 

 

AttnSleep은 class aware loss function과 함께 MRCNN, AFR, TCE module로 구성된다.

이때 AttnSleep의 각 모듈의 효율성을 분석하기 위해 Ablation Study를 진행했다.

총 5가지의 모델 변형을 시도했으며, 1번~4번까지는 class-aware loss function을 사용하지 않았다.

 

 

1)  MRCNN: MRCNN module only.

2)  MRCNN+AFR: MRCNN and AFR without TCE.

3)  MRCNN+TCE: MRCNN and TCE without AFR.

4)  MRCNN+AFR+TCE: training MRCNN, AFR and TCE together, without the class-aware loss function.

5)  AttnSleep: training MRCNN, AFR and TCE together, with the class-aware loss function.

 

 

Ablation study

 

Ablation Study를 통해 몇 가지 결론을 도출해낼 수 있다. 

 

1) AFR는 분류 성능을 향상할 수 있으며, 이는 feature inter-dependencies를 모델링하는 것이 필수적임을 의미한다.

2) MRCNN과 MRCNN_TCE를 비교해보면 temporal dependencies를 포착하는 것을 포함시키는 게 sleep stage classification에 중요하다고 볼 수 있다. 뿐만 아니라, MRCNN + TCE가 MRCNN + AFR보다 성능이 훨씬 우수하기 때문에 TCE가 AFR보다 훨씬 더 중요하다는 것을 의미한다.

3) AttnSleep은 다른 네 가지의 변형 모델보다 훨씬 더 나은 MF1 / MGM을 보인다. 따라서 class-aware cost-sensitive loss function은 데이터 불균형 문제를 효과적으로 해결할 수 있음을 나타낸다.

 

 

 

 

---

3-5) Sensitivity Analysis for the Number of Heads in MHA

MHA는 모델에서 중요한 요소 중 하나이므로 head 수에 대한 고민 또한 필요하다.

해당 연구에서는 head의 수가 5일 때 모델이 가장 높은 성능을 보였다.

 

 

 

 

---

4. Conculsion

- 해당 논문에서는 Sleep stage를 조금 더 쉽게 예측하기 위한 방법으로 Attnsleep model을 제시한다.

- Attnsleep은 single channel EEG signal을 이용하여 sleep stage를 분류하는 모델이다.

- Attnsleep은 MRCNN /  AFR 을 이용하여 feature 추출한다.

- feature 추출 후, 이를 MHA에 연결하여 시간적 정보를 반영한다.

- 사람의 수면단계에서 각 단계에 해당되는 데이터의 수가 불균형하므로, 보완하기 위해 class-aware loss function을 사용하였다.

 

 

 

---

 

🖇 논문 & 코드 

https://github.com/emadeldeen24/AttnSleep

https://ieeexplore.ieee.org/abstract/document/9417097

 

 

🖇 논문에서 사용된 데이터 다운로드

Sleep Heart Health Study(SHHS) https://sleepdata.org/datasets/shhs

Sleep-EDF-78 (https://physionet.org/content/sleep-edfx/1.0.0/)

Sleep-EDF-20 (https://gist.github.com/emadeldeen24/a22691e36759934e53984289a94cb09b)