MEASURE: Multi-scale Minimal Sufficient Representation Learning for Domain Generalization in Sleep Staging

TL;DR

Sleep staging models often fail to generalize to unseen subjects due to domain shifts across individuals. MEASURE introduces a theoretically grounded framework that minimizes excess domain-relevant information via minimal sufficient representation learning while preserving multi-scale temporal and spectral features, integrated with contrastive learning. The method combines multi-scale alignment, conditional entropy maximization, and a transformer-based sleep-staging phase, achieving state-of-the-art DG performance on SleepEDF-20 and MASS and outperforming existing DG and contrastive approaches. This approach provides a principled, non-adversarial route to robust sleep staging across diverse subjects, with broader implications for domain-general biosignal analysis.

Abstract

Deep learning-based automatic sleep staging has significantly advanced in performance and plays a crucial role in the diagnosis of sleep disorders. However, those models often struggle to generalize on unseen subjects due to variability in physiological signals, resulting in degraded performance in out-of-distribution scenarios. To address this issue, domain generalization approaches have recently been studied to ensure generalized performance on unseen domains during training. Among those techniques, contrastive learning has proven its validity in learning domain-invariant features by aligning samples of the same class across different domains. Despite its potential, many existing methods are insufficient to extract adequately domain-invariant representations, as they do not explicitly address domain characteristics embedded within the unshared information across samples. In this paper, we posit that mitigating such domain-relevant attributes-referred to as excess domain-relevant information-is key to bridging the domain gap. However, the direct strategy to mitigate the domain-relevant attributes often overfits features at the high-level information, limiting their ability to leverage the diverse temporal and spectral information encoded in the multiple feature levels. To address these limitations, we propose a novel MEASURE (Multi-scalE minimAl SUfficient Representation lEarning) framework, which effectively reduces domain-relevant information while preserving essential temporal and spectral features for sleep stage classification. In our exhaustive experiments on publicly available sleep staging benchmark datasets, SleepEDF-20 and MASS, our proposed method consistently outperformed state-of-the-art methods. Our code is available at : https://github.com/ku-milab/Measure

Figures (7)

• Figure 1: Comparison between (a) sufficient representation and (b) minimal sufficient representation. In conventional contrastive learning, $\boldsymbol{z}_i$ denotes the normalized feature of the $i$-th sample $\boldsymbol{v}_i$, while $\boldsymbol{z}_p$ represents the feature of a positive sample $\boldsymbol{v}_p$ that shares the same label as $\boldsymbol{v}_i$. The domain factor $D$ denotes the set of attributes that contribute to the domain gap. (a) Sufficient Representation Learning: This approach seeks to maximize the shared information between feature and positive samples $I(\boldsymbol{z}_i ; \boldsymbol{v}_p)$, while simultaneously introducing the superfluous information $I(\boldsymbol{z}_i; \boldsymbol{v}_i | \boldsymbol{v}_p)$, which corresponds to the information present in $\boldsymbol{v}_i$ but absent in $\boldsymbol{v}_p$. Among these, excess domain-relevant information$I(\boldsymbol{z}_i ; d_i | \boldsymbol{v}_p)$ caused by domain attributes hinders the learning of domain-invariant features, where $d_i$ refers to the domain label of $\boldsymbol{v}_i$. (b) Minimal Sufficient Representation Learning: This approach aims to reduce the superfluous information $I(\boldsymbol{z}_i; \boldsymbol{v}_i | \boldsymbol{v}_p)$, thereby diminishing the excess domain-relevant information and enabling the learning of more domain-invariant features.
• Figure 2: Overview of MEASURE. Our method consists of two stages: (a) Multi-scale minimal sufficient representation learning and (b) Sleep staging. In stage (a), multi-scale features are extracted from various layers of the encoder, capturing diverse frequency and temporal information. These features are subsequently projected into a shared feature space and optimized using the proposed objective. In the sleep staging (b), the encoder learned in the early stage is frozen, and the extracted multi-scale features are fed into a transformer to produce level-specific predictions. The final predicted sleep stage label $\hat{y}$ is obtained by aggregating these predictions using an argmax operation, as described in lee2024sleepyco.
• Figure 3: Comparison of hypnograms generated by the baseline model (a) SleePyCo and (b) the proposed MEASURE model on SleepEDF-20. True sleep stages (blue) and predicted stages (red dashed) are visualized over time.
• Figure 4: Performance comparison across varying the $\alpha$ on SleepEDF-20 and MASS datasets.
• Figure 5: Visualization of correlation between the superfluous information and domain-relevant information.

Theorems & Definitions (3)

• Definition 1: Sufficient representation for contrastive learning
• Definition 2: Minimal sufficient representation
• Theorem 1