[논문 리뷰] A Deep Knowledge Distillation framework for EEG assisted enhancement of single-lead ECG based sleep staging

Vaibhav Joshi, Sricharan Vijayarangan외 2명|arXiv (Cornell University)|2021. 12. 14.

EEG and Brain-Computer Interfaces인용 수 5

한 줄 요약

이 논문은 EEG에서 단일 유도 ECG로 수면 단계 분류 지식을 응답 기반 및 특징 기반 distillation을 통해 전이하는 크로스모달 지식 정련 프레임워크를 제안한다. 이는 ECG 기반 수면 단계 분류 정확도를 크게 향상시킨다. 기준 ECG 모델 대비 4단계 분류에서는 가중 F1 점수에서 14.3% 상대적 향상(0.61로)을, 3단계 분류에서는 13.4% 상대적 향상(0.66로)을 기록하며, 비침습적 수면 모니터링 향상을 위한 정련 기법의 타당성을 입증한다.

ABSTRACT

Automatic Sleep Staging study is presently done with the help of Electroencephalogram (EEG) signals. Recently, Deep Learning (DL) based approaches have enabled significant progress in this area, allowing for near-human accuracy in automated sleep staging. However, EEG based sleep staging requires an extensive as well as an expensive clinical setup. Moreover, the requirement of an expert for setup and the added inconvenience to the subject under study renders it unfavourable in a point of care context. Electrocardiogram (ECG), an unobtrusive alternative to EEG, is more suitable, but its performance, unsurprisingly, remains sub-par compared to EEG-based sleep staging. Naturally, it would be helpful to transfer knowledge from EEG to ECG, ultimately enhancing the model's performance on ECG based inputs. Knowledge Distillation (KD) is a renowned concept in DL that looks to transfer knowledge from a better but potentially more cumbersome teacher model to a compact student model. Building on this concept, we propose a cross-modal KD framework to improve ECG-based sleep staging performance with assistance from features learned through models trained on EEG. Additionally, we also conducted multiple experiments on the individual components of the proposed model to get better insight into the distillation approach. Data of 200 subjects from the Montreal Archive of Sleep Studies (MASS) was utilized for our study. The proposed model showed a 14.3\% and 13.4\% increase in weighted-F1-score in 4-class and 3-class sleep staging, respectively. This demonstrates the viability of KD for performance improvement of single-channel ECG based sleep staging in 4-class(W-L-D-R) and 3-class(W-N-R) classification.

연구 동기 및 목표

• EEG 기반 방법에 비해 정확도가 낮은 단일 유도 ECG 기반 수면 단계 분류의 문제를 해결하기 위해.
• 다중 센서 추론이 필요 없이 EEG에서 ECG로 지식을 전이하기 위한 지식 정련의 가능성을 탐색하기 위해.
• 단일 모odal 추론를 유지하면서도 ECG 기반 수면 단계 분류 성능을 향상시키는 크로스모달 정련 프레임워크를 개발하기 위해.
• 응답 기반, 주의 기반, 대비 학습과 같은 개별 정련 구성 요소가 모델 성능에 기여하는 방식을 평가하기 위해.
• 성능 향상의 메커니즘을 이해하기 위해 버티브 레이어의 특징 표현을 분석하기 위해.

제안 방법

• 고수준 수면 단계 표현을 학습하기 위해 교사 모델이 EEG 신호에서 훈련된다.
• 학생 모델은 단일 유도 ECG 데이터에서 훈련되며, 교사의 소프트 레이블(소프트맥스 레이어에서 유도)을 통한 응답 기반 정련을 통해 지식을 전달받는다.
• 중간 특징 맵의 지식을 교사에서 학생으로 전이하기 위해 주의 기반 정련이 적용된다.
• 잠재 표현을 버티브 레이어에서 ECG 및 EEG 특징 간에 정렬하기 위해 대비 학습(Contrastive Learning, CL)이 사용된다.
• 완전한 프레임워크는 응답 기반 정련(SD), 주의 기반 정련(AT), 대비 학습(CL)을 결합하여 공동 최적화를 수행한다.
• 모델은 MASS 데이터셋의 200명의 환자에서 엔드 투 엔드로 훈련되며, 배포 시에는 추론을 위해 오직 ECG에만 의존한다.

실험 결과

연구 질문

• RQ1 지식 정련이 EEG에서 단일 유도 ECG로 수면 단계 분류 지식을 효과적으로 전이하여, 추론 시 EEG가 필요 없이 성능을 향상시킬 수 있는가?
• RQ2 응답 기반 정련, 주의 정련, 대비 학습과 같은 개별 구성 요소가 ECG 기반 수면 단계 분류 성능 향상에 기여하는 방식은 무엇인가?
• RQ3 다양한 정련 방식의 조합이 개별 방식보다 더 나은 성능을 내는가, 아니면 간섭을 유발하는가?
• RQ4 버티브 레이어에서 성공적인 정련과 실패한 정련의 경우, 학습된 특징 표현은 어떻게 다를까? 이는 특징 품질에 대한 통찰을 제공하는가?
• RQ5 클래스 불균형이 정련 성능에 미치는 영향은 어느 정도이며, 프레임워크는 불균형한 수면 단계 분포에 대해 얼마나 강건한가?

주요 결과

• 제안된 AT+SD+CL 정련 프레임워크는 4단계 수면 단계 분류(W-L-D-R)에서 가중 F1 점수에서 14.3% 상대적 향상(0.45에서 0.51로)을 기록했다.
• 3단계 분류(W-N-R)에서 가장 우수한 아블레이션(AT+CL)은 가중 F1 점수에서 13.4% 상대적 향상(0.58에서 0.66로)을 달성했다.
• SD+CL 방법은 4단계 분류에서 경도 수면(L) 단계에서 가장 높은 향상을 보였으며, F1 점수를 0.47에서 0.61로 상승시켰다.
• AT+CL 방법은 3단계 분류에서 NREM(N) 단계에서 가장 높은 향상을 기록했으며, F1 점수를 0.65에서 0.77로 상승시켰다.
• 완전한 AT+SD+CL 모델은 모든 클래스에서 상대적으로 균형 잡힌 성능을 보이며, 클래스 불균형에 대한 강건성을 나타냈다.
• 특징 시각화 결과, 성공적인 정련은 더 구분력 있는 버티브 레이어 특징을 이끌어내었고, 실패한 경우는 클래스 불균형이나 최적화되지 않은 정련으로 인해 표현이 정렬되지 않은 것을 확인할 수 있었다.