[论文解读] A Deep Knowledge Distillation framework for EEG assisted enhancement of single-lead ECG based sleep staging

Vaibhav Joshi, Sricharan Vijayarangan+2 more|arXiv (Cornell University)|Dec 14, 2021

EEG and Brain-Computer Interfaces被引用 5

一句话总结

本文提出了一种跨模态知识蒸馏框架，通过基于响应和基于特征的蒸馏，将来自脑电图（EEG）的睡眠分期知识迁移至单导联心电图（ECG），显著提升了基于ECG的睡眠分期准确率。该方法在4类分期任务中使加权-F1分数相对提升了14.3%（达到0.61），在3类分期任务中提升了13.4%（达到0.66），相较于基线ECG模型，证明了知识蒸馏在提升无创睡眠监测性能方面的可行性。

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迁移知识的手段，且无需在推理阶段使用多传感器数据。
• 开发一种跨模态蒸馏框架，在保持单模态推理的同时提升基于ECG的睡眠分期性能。
• 评估各个蒸馏组件——基于响应的蒸馏、基于注意力的蒸馏和对比学习——对模型性能的贡献。
• 分析瓶颈层中的特征表示，以理解性能提升背后的机制。

提出的方法

• 在EEG信号上训练教师模型，以学习高层次的睡眠阶段表征。
• 在单导联ECG数据上训练学生模型，通过基于响应的蒸馏（从Softmax层蒸馏软标签）从教师模型迁移知识。
• 应用基于注意力的蒸馏，将教师模型中间特征图的知识迁移至学生模型。
• 使用对比学习（CL）对瓶颈层中ECG与EEG特征的潜在表征进行对齐。
• 完整框架结合了基于响应的蒸馏（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在所有类别上均表现出相对均衡的性能，表明对类别不平衡具有鲁棒性。
• 特征可视化显示，成功蒸馏可产生更具判别性的瓶颈特征，而失败案例则因类别不平衡或蒸馏不优导致表征错位。