SSF-SET: A Discrete EEG Token-Based Framework for Sleep Stage Forecasting

초록

Sleep is important to human health. To improve sleep quality, recorded EEG signals have been utilized for automated sleep staging either in real time during sleep or after sleep. However, because previous approaches classify events that have already occurred rather than forecasting the future, their effectiveness is limited for personalized sleep management. This study proposes the sleep stage forecaster with sleep EEG tokenizer (SSF-SET) framework, which predicts the future sleep-stage sequence using only earlier EEG of the current sleep. SET combines a multi-branch transformer for epoch-level representations, an LSTM-based sequence encoder-decoder, and quantization to convert continuous EEG features into sleep-informative tokens. This quantization makes an information bottleneck that reveals the latent transition structure and suppresses artifacts, enabling reliable next-stage prediction. The decoder-only transformer SSF is first pretrained for next-token prediction with causal attention, then fine-tuned via reinforcement learning that uses sequence-level macro-F1 and token-consistency rewards; throughout it does not access future EEG at inference. With subject-wise cross-validation on the SleepEDF20 and SleepEDF78 datasets, SSF-SET consistently outperformed direct forecasting and forecasting with predicted sleep stages. On SleepEDF20, accuracy was 0.596 and macro-F1 was 0.516. In addition, we achieved an accuracy of 0.611 and a macro-F1 of 0.537 on SleepEDF78. These results show that quantized sleep EEG tokens are effective for autoregressive prediction and demonstrate that future sleep stages can be predicted without future EEG. We believe that SSF-SET is an important component for developing closed-loop and personalized sleep interventions that can act before disruptive transitions occur, and we expect it to improve sleep quality.

키워드