DualStabSleepNet: A Dual-Domain Diffusion Stabilization Network for Robust Sleep Staging

Chongjian Wang; Chen Liu; Junjie Gao; Xiaofang Zhong; Shiyuan Han; Tong Zhang

cs.CVAug 17, 2026

DualStabSleepNet: A Dual-Domain Diffusion Stabilization Network for Robust Sleep Staging

Authors: Chongjian Wang, Chen Liu, Junjie Gao, Xiaofang Zhong, Shiyuan Han, Tong Zhang

Organizations: School of Mathematics and Systems Science, Shandong University of Science and Technology, Qingdao, 266590, Shandong, China · School of Artificial Intelligence, Shandong Women’s University, Jinan, 250352, Shandong, China · School of Computer Science and Engineering, South China University of Technology, Guangzhou, 510641, Guangdong, China

Abstract

Existing deep learning approaches for automatic sleep staging suffer from limited robustness under heterogeneous recording conditions, where non-stationary noise, inter-subject differences and cross-dataset distribution shifts cause unstable features and poor generalization. This work proposes DualStabSleepNet (DSSNet), a dual-domain diffusion stabilization network for robust sleep staging, which improves robustness in both data and feature domains. After preprocessing multi-channel polysomnography (PSG), a continuous-scale diffusion-based stabilization module suppresses noise while preserving physiological signal structures. Stabilized signals are converted to time-frequency representations and fed into a Vision Transformer backbone. A teacher-student guided diffusion feature stabilization module further mitigates feature drift and enforces multi-level feature consistency. Evaluated on four public PSG datasets SleepEDF-20, SleepEDF-78, SHHS and ISRUC-S3, DSSNet achieves state-of-the-art accuracy of 89.2%, 88.0%, 89.7%, 86.7% with improved macro-F1 and Cohen's kappa. It obtains notable improvements on hard transitional stages (e.g., 12.5% gain for N1 on SHHS) and boosts N2/REM recognition. Under cross-dataset settings, DSSNet is robust to distribution shift and performs on par with or superior to target-dataset trained baselines, demonstrating its practical potential for real-world sleep staging across heterogeneous cohorts.

Figures & tables

Fig. 1 : (a) Data domain. Raw multi-channel PSG signals (EEG, EOG, and EMG) are normalized and processed by a continuous-scale diffusion process (blue), where a U-Net–based denoiser (green) progressively suppresses noise while preserving physiological signal structures. (b) Feature domain. The enhanced 1D signals are transformed into 2D time–frequency representations via STFT, encoded by a Vision Transformer (green), and further stabilized by feature-domain diffusion modules (purple) before feature fusion and sleep stage classification.

Fig. 2 : The forward process adds noise to PSG signals along predefined scales, and the reverse process uses a U-Net–based denoising function conditioned on the noise level to reconstruct clean and physiologically consistent signals [32].

Fig. 3 : Multi-level features extracted by the Vision Transformer are refined through feature-domain diffusion modules (EDM), where diffusion-based correction is applied at each feature level to reduce representation instability. The stabilized features are then fed into fusion layers for supervised finetuning and final sleep stage classification.

Dataset	Subjects	Channels	Freq.	Epoch numbers					Total
Dataset	Subjects	Channels	Freq.	Wake	N1	N2	N3	REM	Total
Sleepedf-20	20	3	100 Hz	8285	2804	17799	5703	7717	42308
				19.6%	6.6%	42.1%	13.5%	18.2%
Sleepedf-78	78	3	100 Hz	63802	21229	68645	12883	25651	192210
				33.2%	11.0%	35.71%	6.7%	13.34%
SHHS	329	3	125 Hz	46369	10304	142125	60153	65953	324854

TABLE I : Statistics of Sleep Stage Distribution in Different Datasets

Symbol	Value
$[\sigma^{data}_{min},\sigma^{data}_{max}]$	[0.05, 0.2]
$\sigma_{data}$	0.5
$N_{data}$	8 / 12 / 16
$\sigma^{feat}_{max}$	0.2
$K$	4
$\lambda$	0.2

TABLE II : Key Hyperparameters Used in Experiments

Dataset	Method	Architecture	Acc	MF1	Kappa	W	N1	N2	N3	R
ISRUC-S3	MVF-SleepNet [21]	CNN + LSTM	84.1	82.8	0.795	90.0	62.5	83.3	91.1	87.3
	SVM [14]	Traditional	71.4	67.2	0.626	82.4	42.8	72.4	81.5	56.9
	RF [15]	Traditional	70.2	68.5	0.616	83.8	47.0	67.1	76.3	68.4
	MixSleepNet [33]	CNN + Transformer	83.0	82.1	0.782	89.9	62.5	81.9	89.9	86.0
	XSleepNet1 [12]	CNN + LSTM	82.5	80.8	0.774	90.1	58.6	82.5	88.7	84.3
	XSleepNet2 [12]	CNN + LSTM	82.6	81.0	0.774	89.9	59.0	82.6	88.4	84.9

TABLE III : Comparison with State-of-the-Art Methods on Five Public Datasets

SleepEDF-78
Stage	Pre.	Rec.	F1	Spec.	Samples
Wake	0.9583	0.9278	0.9428	0.9800	63,802
N1	0.6297	0.6215	0.6256	0.9546	21,229
N2	0.8985	0.9083	0.9034	0.9430	68,645
N3	0.8435	0.8476	0.8455	0.9887	12,883
REM	0.8677	0.9185	0.8924	0.9784	25,651

TABLE IV : Per-Class Performance Across Datasets

Panel of Figure 8 — Fig. 4 : Hyperparameter sensitivity analysis under different settings. Each curve represents performance variation with respect to one hyperparameter while keeping others fixed. The proposed model shows limited performance fluctuation, suggesting strong robustness.

Panel of Figure 9 — Fig. 5 : Confusion matrices of sleep stage classification obtained from a single multi-channel model. Rows denote ground-truth labels and columns denote predicted labels. Subfigures (A–D) correspond to evaluations on four individual datasets, while (E–F) present cross-dataset results. The matrices indicate consistent classification behavior across datasets, with most confusions occurring between neighboring sleep stages.

Panel of Figure 10 — Fig. 6 : t-SNE visualizations of feature representations learned by a single multi-channel model. Each point corresponds to one EEG epoch and is colored according to the sleep stage. Subfigures (A–D) show results on four individual datasets, while (E–F) present cross-dataset settings. The learned representations exhibit compact intra-class clustering and consistent inter-class separation across datasets, indicating stable and generalizable feature learning.

Data Prepare	SleepEDF-78			SleepEDF-20
Data Prepare	ACC	MF1	$\boldsymbol{\kappa}$	ACC	MF1	$\boldsymbol{\kappa}$
Wavelet + STFT	86.6	83.2	0.827	88.5	84.8	0.841
STFT	86.4	83.1	0.818	88.3	84.4	0.841
Diffusion + STFT	88.0	84.2	0.836	89.2	85.5	0.852

TABLE V : Ablation Study of Data-Domain Filtering

EEGDenoise
Method	RRMSE $\downarrow$	CC $\uparrow$	SNR $\uparrow$	WSNR $\uparrow$	WCC $\uparrow$
1D-ResCNN [39]	0.302	0.931	12.746	16.897	0.906
EEGdenoiseNet [40]	0.344	0.904	12.121	15.094	0.888
EEGDnet [41]	0.301	0.945	13.251	17.584	0.924
DuoCL [42]	0.315	0.919	13.236	16.374	0.902
GCTNet [43]	0.262	0.948	14.534	18.702	0.926

TABLE VI : Denoising Performance Comparison on EEG Datasets

Baseline	Feature	Data	SleepEDF-78			SleepEDF-20
Baseline	Feature	Data	ACC	MF1	$\boldsymbol{\kappa}$	ACC	MF1	$\boldsymbol{\kappa}$
✓	✓	✓	88.0	84.2	0.836	89.2	85.5	0.852
✓	✓	$\times$	87.4	83.5	0.824	88.5	85.1	0.849
✓	$\times$	✓	86.3	82.4	0.806	87.1	84.3	0.836
✓	$\times$	$\times$	84.6	81.1	0.789	85.9	82.8	0.819

TABLE VII : Overall Ablation of the Two-Stage Strategy Diffusion

Layer Ablation	SleepEDF-78			SleepEDF-20
Layer Ablation	ACC	MF1	$\boldsymbol{\kappa}$	ACC	MF1	$\boldsymbol{\kappa}$
L1 + L2 + L3 + L4	88.0	84.2	0.836	89.2	85.5	0.852
L2 + L4	87.4	83.6	0.814	88.6	84.9	0.848
L4	86.9	83.1	0.807	88.2	84.5	0.843

TABLE VIII : Layer Ablation Study in the Feature Domain

EEG	EOG	EMG	SleepEDF-78			SleepEDF-20
EEG	EOG	EMG	ACC	MF1	$\boldsymbol{\kappa}$	ACC	MF1	$\boldsymbol{\kappa}$
✓	✓	✓	88.0	84.2	0.836	89.2	85.5	0.852
✓	✓	$\times$	87.1	82.8	0.803	88.8	84.2	0.837
✓	$\times$	$\times$	85.4	81.4	0.782	87.7	82.6	0.818

TABLE IX : Overall Ablation of Data Channels

L1+L2+L3+L4
Stage	Params (M)	FLOPs (G)	FPS (samples/s)
1	9.276	1.883	1265.03
2	9.276	2.508	1015.05
3	9.276	3.133	726.62
4	9.276	3.758	672.28
L2+L4

TABLE X : Computational Cost under Different Feature-Layer Configurations

Explore similar work

STDA-Net: Spectrogram-Based Domain Adaptation for cross-dataset Sleep Stage Classification
May 7, 2026Unaza Tallal, Shruti Kshirsagar, Ankita ShuklaElectroencephalography Sleep Stage Classification
School of Computing, Wichita State University, 1845 N. Fairmount, Wichita, KS 67260, USA · Computer Science and Engineering Department, University of Nevada, Reno, NV, USA
StageGuard: Physiologically Constrained Sleep Staging
Jul 25, 2026Juntang Wang, Yihan Wang, Hao Wu +3Constrained Decoding Sleep Stage Classification
Zu Chongzhi Center, Duke Kunshan University Kunshan, China
SWINSleepNet: A Hierarchical Context-Aware Framework for Sleep Staging (v2)
Aug 3, 2026Chongjian Wang, Junjie GaoSleep Stage Classification Hybrid CNN-Transformer Architectures
School of Mathematics and Systems Science, Shandong University of Science and Technology, Qingdao 266590, Shandong, China · School of Artificial Intelligence, Shandong Women’s University, Jinan 250352, Shandong, China

May 7, 2026, cs.LGJ/K move · Enter open · S save

STDA-Net: Spectrogram-Based Domain Adaptation for cross-dataset Sleep Stage Classification

Unaza Tallal, Shruti Kshirsagar, Ankita Shukla

School of Computing, Wichita State University, 1845 N. Fairmount, Wichita, KS 67260, USA · Computer Science and Engineering Department, University of Nevada, Reno, NV, USA

Electroencephalography Sleep Stage Classification Unsupervised Domain Adaptation Time Series Classification Cross-Dataset Generalization

Open paper AskPDF arXiv

Accurate sleep stage classification across datasets remains challenging due to variability in EEG channel montages, sampling rates, recording environments, and subject populations. Although deep learning has shown considerable promise for automated sleep staging, most existing cross-dataset methods rely on one-dimensional EEG signal representations, whereas the use of two-dimensional spectrogram-based inputs within an unsupervised domain adaptation framework has remained largely unexplored. Here, we propose STDA-Net (Spectrogram-based Temporal Domain Adaptation Network), a framework that combines a convolutional neural network (CNN) for spectrogram-based feature extraction, a bidirectional long short-term memory (BiLSTM) module for temporal modeling of sleep dynamics, and a domain-adversarial neural network (DANN) for source-to-target feature alignment without requiring any labeled target-domain data during training. Experiments are conducted on three publicly available datasets Sleep-EDF, SHHS-1, and SHHS-2 under six cross-dataset transfer settings. Results show that the proposed framework achieves an average accuracy of 89.03% and an average macro F1-score of 87.64%, consistently outperforming existing 1D baseline methods in terms of balanced classification performance, with substantially lower variance across five independent runs, indicating improved stability and reproducibility. Overall, these findings demonstrate that 2D spectrogram-based representations, combined with temporal modeling and adversarial domain adaptation, provide a robust and competitive alternative to conventional 1D EEG inputs for cross-dataset sleep staging.

Abstract

Explore similar work

STDA-Net: Spectrogram-Based Domain Adaptation for cross-dataset Sleep Stage Classification

StageGuard: Physiologically Constrained Sleep Staging

SWINSleepNet: A Hierarchical Context-Aware Framework for Sleep Staging (v2)