S$^3$Net: Stage-Aware Sleep Staging Network

Xiaomao Fan, Kaiwen Xiong, Cheng Lin, Xiaojiang Peng, Ye Li, Ruxin Wang, Yumeng Liu

Keywords: sleep staging, dual-tower network, mixture of expert

Abstract: Automated sleep staging is a critical component in the diagnosis of sleep disorders and the analysis of sleep architecture. While deep learning approaches that leverage time-frequency representations have shown promise, their performance remains suboptimal, primarily due to two fundamental limitations: (1) the inability to effectively model the subtle distinctions of transitional sleep stages (N1 and N2), which exhibit ambiguous electrophysiological patterns, and (2) the inefficient fusion of complementary information from time-domain and frequency-domain representations. To this end, we propose S$^3$Net, a novel Stage-Aware Sleep Staging Network that introduces two dedicated components to address these challenges. First, a Stage-Aware Experts (SAE) module explicitly partitions the sleep stages into easy- and hard-to-separate groups, processing them through separate expert network branches. This allows for specialized feature refinement, particularly for the challenging transitional stages. Second, to foster a cohesive representation, we design a Time Alignment Module (t-ALN) that projects frequency-derived features onto the time axis, effectively bridging the domain gap and enabling synergistic integration of multi-view features. We evaluate S$^3$Net on three public polysomnography datasets (ISRUC-S1, ISRUC-S3, and Sleep-EDF-153). Our model consistently sets a new state-of-the-art, achieving an overall accuracy of 85.6\%, 86.6\%, and 86.9\%, respectively, and demonstrates a \minew{noticeable} improvement in classifying the N1 and N2 stages. The results validate the efficacy of our stage-aware design and structural alignment strategy, offering a more robust framework for clinical and portable sleep staging. Source code is available at \url{https://anonymous.4open.science/r/S3Net/}.

Primary Area: applications to neuroscience & cognitive science

Submission Number: 11065