Resource Efficient Sleep Staging via Multi-Level Masking and Prompt Learning

TL;DR

This work tackles resource-constrained sleep staging by learning from partially observed EEG using Mask-Aware Sleep Staging (MASS). MASS introduces a multi-level masking scheme with epoch- and patch-level masking ratios $r_e$ and $r_a$, and a global prompt token to guide patch- and epoch-level encoding via a two-tier Transformer/GRU pipeline. The method leverages spectral-domain features via PSD within 1-second patches and optimizes a joint objective including cross-entropy, cosine similarity, and transition losses. Across DREAMS-SUB, Sleep-EDF-20/78, and SHHS, MASS achieves state-of-the-art accuracy and macro-F1 under various data availability levels, while reducing data acquisition and inference costs, making it suitable for on-device wearable sleep monitoring.

Abstract

Automatic sleep staging plays a vital role in assessing sleep quality and diagnosing sleep disorders. Most existing methods rely heavily on long and continuous EEG recordings, which poses significant challenges for data acquisition in resource-constrained systems, such as wearable or home-based monitoring systems. In this paper, we propose the task of resource-efficient sleep staging, which aims to reduce the amount of signal collected per sleep epoch while maintaining reliable classification performance. To solve this task, we adopt the masking and prompt learning strategy and propose a novel framework called Mask-Aware Sleep Staging (MASS). Specifically, we design a multi-level masking strategy to promote effective feature modeling under partial and irregular observations. To mitigate the loss of contextual information introduced by masking, we further propose a hierarchical prompt learning mechanism that aggregates unmasked data into a global prompt, serving as a semantic anchor for guiding both patch-level and epoch-level feature modeling. MASS is evaluated on four datasets, demonstrating state-of-the-art performance, especially when the amount of data is very limited. This result highlights its potential for efficient and scalable deployment in real-world low-resource sleep monitoring environments.

Figures (4)

• Figure 1: Comparison of macro-F1 scores in resource-efficient scenarios. When the signal mask ratio increases, performance of models rely on complete training data decrease significantly. Even though there is redundancy in the complete signal, they cannot directly utilize such a small amount of signal for inference. Our MASS framework based on multi-level masking and prompt learning, which can focus on features under a small amount of information and achieve reliable and resource-efficient sleep staging.
• Figure 2: Complete Structure of the Proposed Mask-Aware Sleep Staging Framework.
• Figure 3: Comparison of Model Parameter Sizes and Inference Time on DREAMS-SUB dataset. The results are calculated in 10% signal integrity (80% patch-level masking and 50% epoch-level masking). MASS achieves highest $n_p$ and $n_t$, especially compared with latest models such as LGSleepNet(2023) and NeuroNet(2024).
• Figure 4: Confusion Matrix of Different Mask Ratios on Four Datasets.