PSDNORM: TEST-TIME TEMPORAL
NORMALIZATION FOR DEEP LEARNING IN
SLEEP STAGING

Theo Gnassounou, Antoine Collas, Rémi Flamary, Alexandre Gramfort

► To cite this version:

Theo Gnassounou, Antoine Collas, Rémi Flamary, Alexandre Gramfort. PSDNORM: TEST-TIME TEMPORAL NORMALIZATION FOR DEEP LEARNING IN SLEEP STAGING. 2025. hal-05384934

HAL Id: hal-05384934

https://hal.science/hal-05384934v1

Preprint submitted on 27 Nov 2025

**HAL** is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# PSDNORM: TEST-TIME TEMPORAL NORMALIZATION FOR DEEP LEARNING IN SLEEP STAGING

**Théo Gnassounou**

Université Paris-Saclay,
Inria, CEA,
91120 Palaiseau, France
theo.gnassounou@inria.fr

**Antoine Collas**

Université Paris-Saclay,
Inria, CEA,
91120 Palaiseau, France
antoine.collas@inria.fr

**Rémi Flamary**

École Polytechnique,
IP Paris, CMAP, UMR 7641,
91120 Palaiseau, France
remi.flamary@polytechnique.edu

**Alexandre Gramfort***

Université Paris-Saclay,
Inria, CEA,
91120 Palaiseau, France
alexandre.gramfort@inria.fr

## ABSTRACT

Distribution shift poses a significant challenge in machine learning, particularly in biomedical applications using data collected across different subjects, institutions, and recording devices, such as sleep data. While existing normalization layers, BatchNorm, LayerNorm and InstanceNorm, help mitigate distribution shifts, when applied over the time dimension they ignore the dependencies and auto-correlation inherent to the vector coefficients they normalize. In this paper, we propose PSDNorm that leverages Monge mapping and temporal context to normalize feature maps in deep learning models for signals. Notably, the proposed method operates as a test-time domain adaptation technique, addressing distribution shifts without additional training. Evaluations with architectures based on U-Net or transformer backbones trained on 10K subjects across 10 datasets, show that PSDNorm achieves state-of-the-art performance on unseen left-out datasets while being 4-times more data-efficient than BatchNorm.

# 1 Introduction

**Data Shift in Physiological Signals** Machine learning techniques have achieved remarkable success in various domains, including computer vision, biology, audio processing, and language understanding. However, these methods face significant challenges when there are distribution shifts between training and evaluation datasets [1]. For example, in biological data, such as electroencephalography (EEG) signals, the distribution of the data can vary significantly. Indeed, data is collected from different subjects, electrode positions, and recording conditions. This paper focuses on sleep staging, a clinical task that consists in classifying periods of sleep in different stages based on EEG signals [2]. Depending on the dataset, the cohort can be composed of different age groups, sex repartition, health conditions, and recording conditions [3, 4, 5]. Such variability brings shift in the distribution making it challenging for the model to generalize to unseen datasets.

**Normalization to Address Data Shift** Normalization layers are widely used in deep learning to improve training stability and generalization. Common layers include BatchNorm [6], LayerNorm [7], and InstanceNorm [8], which respectively compute statistics across the batch, normalize across all features within each sample, and normalize each channel independently within a sample. Some normalization methods target specific tasks, such as EEG covariance matrices [9] or time-series forecasting [10], but they do not fully address spectral distribution shifts reflected in the temporal auto-correlations of signals. In sleep staging, a simple normalization is often applied as preprocessing, e.g.,

<footnote>*A. Gramfort joined Meta and can be reached at agramfort@meta.com</footnote>
<figure><img src="image_2.png" /><figcaption>Figure 1: **Description of normalization layers.** The input shape is (N, c, ℓ) with batch size N, channels c, and signal length ℓ. BatchNorm estimates the mean $\hat{\mu}$ and variance $\hat{\sigma}^2$ over batch and time, and learns parameters $(\gamma, \beta)$ to normalize the input. PSDNorm estimates PSDs $\hat{\mathbf{P}}$ over time and accounts for local temporal correlations. It computes the barycenter PSD $\hat{\mathbf{P}}$, updates it via a running Riemannian barycenter (6), and applies the filter $\hat{\mathbf{H}}$ to normalize the input. The hyperparameter F controls the extent of temporal correlation considered, thereby adjusting the strength of the normalization.</figcaption></figure>

standardizing signals over entire nights [11] or short temporal windows [12]. Recent studies [13, 14] highlight the importance of considering temporal correlation and spectral content in normalization, proposing Temporal Monge Alignment (TMA), which aligns Power Spectral Density (PSD) to a common reference using Monge mapping, going beyond simple z-score normalization. However, these methods remain preprocessing steps that cannot be inserted as layers in the network architecture as it is done with BatchNorm, LayerNorm or InstanceNorm.

**Domain Generalization** Sleep staging has been addressed by various neural network architectures, which process raw signals [12, 15, 16], spectrograms [17, 18], or both [19]. More recent approaches involve transformer-based models that handle multimodal [20], spectrogram [21], or heterogeneous inputs [22], offering improved modeling of temporal dependencies. However, most existing models are trained on relatively small cohorts, typically consisting of only a few hundred subjects, which limits their ability to generalize to diverse clinical settings. Notable exceptions include U-Sleep [15], which was trained on a large-scale dataset and incorporates BatchNorm layers to mitigate data variability, and foundational models [23, 24, 25] that achieve strong generalization from vast amount of data but require significant computational resources and are challenging to adapt without fine-tuning. Our focus is on developing smaller, efficient models that balance good generalization with ease of training and deployment in clinical practice.

**Test-time Domain Adaptation** Domain generalization involves training models on data from multiple domains without using domain labels, with the goal of achieving broad generalization. In contrast, domain adaptation (DA) utilizes domain information to improve performance by aligning source and target distributions during training [26, 27]. While DA methods have been applied to address domain shifts in sleep staging [28], they typically require access to source data for each new target, which can be limiting. Test-time DA offers a more flexible approach, adapting the model during inference without requiring access to source data [13]. This makes test-time DA particularly well-suited for clinical applications where retraining or data sharing is impractical.

**Contributions** In this work, we introduce the PSDNorm deep learning layer, a novel approach to address distribution shifts in machine learning for signals. PSDNorm leverages Monge Mapping to incorporate temporal context and normalize feature maps effectively. As a test-time DA method, it adapts to distribution shifts during inference without extra training or source domain access. Compared to a normalization like LayerNorm or InstanceNorm, PSDNorm exploits the sequential nature of the intermediate layer representations as illustrated in Figure 1. We evaluate PSDNorm through extensive experiments on 10 sleep datasets. This evaluation covers 10M of samples across 10K subjects, using a leave-one-dataset-out (LODO) protocol with 3 different random seeds. To the best of our knowledge, such a large-scale and systematic evaluation has never been conducted before. PSDNorm achieves state-of-the-art performance and requires 4 times fewer labeled data to match the accuracy of the best baseline. Results highlight the potential of
PSDNorm as a practical and efficient solution for tackling domain shifts in signals.

The paper is structured as follows: Section 2 discusses existing normalization layers and pre-processing. Section 3 introduces PSDNorm, followed by numerical results in Section 4.

**Notations** Vectors are denoted by small cap boldface letters (e.g., $\mathbf{x}$), matrices by large cap boldface letters (e.g., $\mathbf{X}$). The element-wise product, power of $n$ and division are denoted $\odot$, $\odot^n$ and $\ominus$, respectively. $\llbracket 1, K \rrbracket$ denotes $\{1, \dots, K\}$. The absolute value is $|\cdot|$. The discrete circular convolution along the temporal axis operates row-wise as, $*: \mathbb{R}^{c \times l} \times \mathbb{R}^{c \times F} \to \mathbb{R}^{c \times l}$ for $l \ge F$. $\text{vec}: \mathbb{R}^{c \times l} \to \mathbb{R}^{c l}$ concatenates rows of a time series into a vector. $x_l = [\mathbf{x}]_l$ refers to the $l^{th}$ element of $\mathbf{x}$, and $X_{l,m} = [\mathbf{X}]_{l,m}$ denotes the element of $\mathbf{X}$ at the $l^{th}$ row and $m^{th}$ column. $\mathbf{X}^*$ and $\mathbf{X}^\top$ are the conjugate and the transpose of $\mathbf{X}$, respectively. diag puts the elements of a vector on the diagonal of a matrix. $\otimes$ is the Kronecker product. $\mathbf{1}_c$ is the vector of ones of size $c$.

## 2 Related Works

In this section, we first review fundamental concepts of normalization layers. Then, we recall the Temporal Monge Alignment (TMA) method [13] that aligns the PSD of signals using optimal transport.

**Normalization Layers** Normalization layers enhance training and robustness in deep neural networks. The most common are BatchNorm [6], InstanceNorm [8], and LayerNorm [7]. BatchNorm normalizes feature maps using batch and time statistics, ensuring zero mean and unit variance. The output is adjusted with learnable parameters. InstanceNorm normalizes each channel per sample using its own statistics, independent of the batch (see Fig. 1). Popular in time-series forecasting, it is used in RevIN [10], which reverses normalization after decoding. LayerNorm normalizes across all channels and time steps within each sample, with learnable scaling and shifting. While these normalization layers are widely employed, they operate on vectors ignoring statistical dependence and autocorrelation between their coefficients, which are prevalent when operating on time-series. To address this limitation, the Temporal Monge Alignment (TMA) [13, 14] was introduced as a pre-processing step to align temporal correlations by leveraging the Power Spectral density (PSD) of multivariate signals using Monge Optimal Transport mapping.

**Gaussian Periodic Signals** Consider a multivariate signal $\mathbf{X} \triangleq [\mathbf{x}_1, \dots, \mathbf{x}_c]^{\top} \in \mathbb{R}^{c \times l}$ of sufficient length. A standard assumption is that this signal follows a centered Gaussian distribution where sensors are uncorrelated and signals are periodic. This periodicity and uncorrelation structure implies that the signal's covariance matrix is block diagonal, with each block having a circulant structure. A fundamental property of symmetric positive definite circulant matrices is their diagonalization [29] with real and positive eigenvalues in the Fourier basis $\mathbf{F}_\ell \in \mathbb{C}^{l \times l}$ of elements

$$[\mathbf{F}_\ell]_{l,l'} \triangleq \frac{1}{\sqrt{\ell}} \exp \left( -2i\pi \frac{(l-1)(l'-1)}{\ell} \right), \quad (1)$$

where $l, l' \in \llbracket 1, \ell \rrbracket$. Hence, we have $\text{vec}(\mathbf{X}) \sim \mathcal{N}(0, \Sigma)$ with $\Sigma$ block-diagonal,

$$\Sigma = (\mathbf{I}_c \otimes \mathbf{F}_\ell) \text{diag}(\text{vec}(\mathbf{P})) (\mathbf{I}_c \otimes \mathbf{F}_\ell^*) \in \mathbb{R}^{c\ell \times c\ell}, \quad (2)$$

where $\mathbf{P} \in \mathbb{R}^{c \times l}$ contains positive entries corresponding to the Power Spectral Density of each sensor. In practice, since we only have access to a single realization of the signal, the PSD is estimated with only $F \ll l$ frequencies, i.e., $\mathbf{P} \in \mathbb{R}^{c \times F}$. This amounts to considering the local correlation of the signal and neglecting the long-range correlations.

**Power Spectral Density Estimation** The Welch estimator [30] computes the PSD of a signal by averaging the squared Fourier transform of overlapping segments of the signal. Hence, the realization of the signal $\mathbf{X} \in \mathbb{R}^{c \times l}$ is decimated into overlapping segments $\{\mathbf{X}^{(1)}, \dots, \mathbf{X}^{(L)}\} \subset \mathbb{R}^{c \times F}$ to estimate the PSD. The Welch estimator is defined as

$$\hat{\mathbf{P}} \triangleq \frac{1}{L} \sum_{l=1}^{L} \left| \left( (\mathbf{1}_c \mathbf{w}^{\top}) \odot \mathbf{X}^{(l)} \right) \mathbf{F}_F^* \right|^{\odot 2} \in \mathbb{R}^{c \times F}, \quad (3)$$

where $\mathbf{w} \in \mathbb{R}^F$ is the window function such that $\|\mathbf{w}\|_2 = 1$.

**F-Monge Mapping** Let $\mathcal{N}(0, \Sigma^{(s)})$ and $\mathcal{N}(0, \Sigma^{(t)})$ be source and target centered Gaussian distributions respectively with covariance matrices following the structure (2) and PSDs denoted by $\mathbf{P}^{(s)}$ and $\mathbf{P}^{(t)} \in \mathbb{R}^{c \times F}$. Given a signal
$\mathbf{X} \in \mathbb{R}^{c \times l}$ such that $\text{vec}(\mathbf{X}) \sim \mathcal{N}(\mathbf{0}, \mathbf{\Sigma}^{(s)})$, the $F$-Monge mapping as defined by [13, 14] is

$$m_F(\mathbf{X}, \mathbf{P}^{(t)}) \triangleq \mathbf{X} * \mathbf{H} \in \mathbb{R}^{c \times l}, \quad \text{where} \quad \mathbf{H} \triangleq \frac{1}{\sqrt{F}} (\mathbf{P}^{(t)} \otimes \mathbf{P}^{(s)})^{\odot \frac{1}{2}} \mathbf{F}_F^* \in \mathbb{R}^{c \times F}. \quad (4)$$

In this case, $F$ controls the alignment between the source and target distributions. Indeed, if $F = l$, then the $F$-Monge mapping is the classical Monge mapping between Gaussian distributions and the source signal has its covariance matrix equal to $\mathbf{\Sigma}^{(t)}$ after the mapping. If $F = 1$, then each sensor is only multiplied by a scalar.

**Gaussian Wasserstein Barycenter** For Gaussian distributions admitting the decomposition (2), the Wasserstein barycenter [31] admits an elegant closed-form solution. Consider $K$ centered Gaussian distributions admitting the decomposition (2) of PSDs $\mathbf{P}^{(1)}, \dots, \mathbf{P}^{(K)}$. Their barycenter is also a centered Gaussian distribution $\mathcal{N}(\mathbf{0}, \bar{\mathbf{\Sigma}})$ admitting the decomposition (2) with PSD

$$\bar{\mathbf{P}} \triangleq \left( \frac{1}{K} \sum_{k=1}^{K} \mathbf{P}^{(k)} \odot \frac{1}{2} \right)^{\odot 2} \in \mathbb{R}^{c \times F}. \quad (5)$$

**Temporal Monge Alignment** TMA is a pre-processing method that aligns the PSD of multivariate signals using the $F$-Monge mapping. Given a source signal $\mathbf{X}_s$ and a set of target signals $\mathbf{X}_t = \{\mathbf{X}_t^{(1)}, \dots, \mathbf{X}_t^{(K)}\}$, the TMA method uses the $F$-Monge mapping between the source and the Wasserstein barycenter of the target signals. Hence, it simply consists of 1) estimating the PSD of all the signals, 2) computing the Wasserstein barycenter of the target signals, and 3) applying the $F$-Monge mapping to the source signal. TMA, as a preprocessing method, is inherently limited to handling PSD shifts in the raw signals and cannot address more complex distributional changes in the data. This limitation highlights the need for a layer that can effectively capture and adapt to these complex variations during learning and inside deep learning models.

## 3 PSDNorm Layer

The classical normalization layers, such as BatchNorm or InstanceNorm do not take into account the temporal autocorrelation structure of signals. They treat each time sample in the intermediate representations independently. In this section, we introduce the PSDNorm layer that aligns the PSD of each signal onto a barycenter PSD within the architecture of a deep learning model.

**PSDNorm as a Drop-in Replacement** As discussed in Section 2, BatchNorm and InstanceNorm layers are commonly applied after each non-linearity in deep learning models, such as U-Time, to normalize feature maps. In contrast, PSDNorm offers a novel approach by aligning the PSD of feature maps to a barycenter PSD, providing a powerful alternative to classical normalization layers. Designed as a drop-in replacement, PSDNorm is optimized for modern hardware accelerators like GPUs, ensuring efficient execution. Once the deep-learning model is trained, PSDNorm operates as a test-time domain adaptation technique, allowing it to adapt to new data without additional training or access to training data. We define the normalized feature map as $\tilde{\mathbf{G}} \triangleq \text{PSDNorm}(\mathbf{G})$. The following sections introduce the core components of PSDNorm and its implementation.

### 3.1 Core Components of the layer

In the following, we formally define PSDNorm and present each of its three main components: 1) PSD estimation, 2) running Riemannian barycenter estimation, and 3) $F$-Monge mapping computation. Given a batch $\mathcal{B} = \{\mathbf{G}^{(1)}, \dots, \mathbf{G}^{(N)}\}$ of $N$ pre-normalization feature maps, PSDNorm outputs a normalized batch $\tilde{\mathcal{B}} = \{\tilde{\mathbf{G}}^{(1)}, \dots, \tilde{\mathbf{G}}^{(N)}\}$ with normalized PSD. Those three steps are detailed in the following and illustrated in the right part of Figure 1.

**PSD Estimation** First, the estimation of the PSD of each feature map is performed using the Welch method. The per-channel mean $\hat{\mu}^{(j)}$ is computed for each feature map $\mathbf{G}^{(j)}$ as $\hat{\mu}^{(j)} \triangleq \frac{1}{l} \sum_{l=1}^{l} [\mathbf{G}^{(j)}]_{:,l} \in \mathbb{R}^c$. Then, the PSD of the centered feature map $\mathbf{G}^{(j)} - \hat{\mu}^{(j)} \mathbf{1}_l^\top$, denoted $\hat{\mathbf{P}}^{(j)}$, is estimated as described in Equation (3). This centering step is required as feature maps are typically non-centered due to activation functions and convolution biases but they are assumed to have a stationary mean. The Welch estimation involves segmenting the centered feature map into overlapping windows, computing the Fourier transform of each window and then averaging them.
**Geodesic and Running Riemannian Barycenter** The PSD-Norm aligns the PSD of each feature map to a barycenter PSD. This barycenter is computed during training by interpolating between the batch Wasserstein barycenter and the current running Riemannian barycenter using the geodesic associated with the Bures metric [32]. The batch barycenter is first computed from the current batch PSDs {$\hat{\mathbf{P}}^{(1)}, \dots, \hat{\mathbf{P}}^{(N)}$} using Equation (5). To ensure gradual adaptation, the running barycenter is updated via an exponential geodesic average with $\alpha \in [0, 1]$:

$$ \hat{\mathbf{P}} \leftarrow \left( (1 - \alpha) \hat{\mathbf{P}}^{\odot \frac{1}{2}} + \alpha \hat{\mathbf{P}}_{\mathrm{B}}^{\odot \frac{1}{2}} \right)^{\odot 2} \in \mathbb{R}^{c \times F}. \quad (6) $$

A proof of the geodesic is provided in Appendix A.1.

**PSD Adaptation with F-Monge Mapping** The final step of the PSDNorm is the application of the F-Monge mapping to each feature map after subtracting the per-channel mean. Indeed, for all $j \in [1, N]$, it is defined as

$$ \tilde{\mathbf{G}}^{(j)} = m_F \left( \mathbf{G}^{(j)} - \hat{\mu}^{(j)} \mathbf{1}_{\ell}^{\top}, \hat{\mathbf{P}} \right) = \left( \left( \mathbf{G}^{(j)} - \hat{\mu}^{(j)} \mathbf{1}_{\ell}^{\top} \right) * \hat{\mathbf{H}}^{(j)} \right) \in \mathbb{R}^{c \times \ell} \quad (7) $$

where $\hat{\mathbf{H}}^{(j)}$ is the Monge mapping filter computed as

$$ \hat{\mathbf{H}}^{(j)} \triangleq \frac{1}{\sqrt{F}} \left( \hat{\mathbf{P}} \otimes \hat{\mathbf{P}}^{(j)} \right)^{\odot \frac{1}{2}} \mathbf{F}_{F}^{*} \in \mathbb{R}^{c \times f} \quad (8) $$

where $\hat{\mathbf{P}}^{(j)}$ is the estimated PSD of $\mathbf{G}^{(j)} - \hat{\mu}^{(j)} \mathbf{1}_{\ell}^{\top}$.

## 3.2 Implementation details

**Overall Algorithm** The forward computation of the proposed layer is outlined in Algorithm 1. At train time, the PSDNorm performs three main operations: 1) PSD estimation, 2) running Riemannian barycenter update, and 3) Monge mapping application. At test time, the PSDNorm operates similarly, except it does not update the running barycenter. The PSDNorm is fully differentiable and can be integrated into any deep learning model. Similarly to classical normalization layers, a stop gradient operation is applied to the running barycenter to prevent the backpropagation of the gradient computation through the barycenter. PSDNorm has a unique additional hyperparameter $F$ which is the filter size. It controls the alignment between each feature map and the running barycenter PSD and it is typically chosen in our experiments between 1 and 17. In practice,

the Fourier transforms are efficiently computed using the Fast Fourier Transform (FFT) algorithm. Because of the estimation of PSDs, the complexity of the PSDNorm, both at train and test times, is $O(Nc\ell F \log(F))$, where $N$ is the batch size, $c$ the number of channels, $\ell$ the signal length, and $F$ the filter size.

**PSDNorm as a generalization of InstanceNorm** InstanceNorm applies a per-channel z-score over time, subtracting the mean and dividing by the standard deviation—equivalent to whitening under an i.i.d. assumption over time. In contrast, PSDNorm explicitly accounts for temporal structure by estimating the PSD and whitening/re-coloring in the frequency domain. InstanceNorm is recovered as a special case of PSDNorm by setting the filter size to $F = 1$ and using the identity as the re-coloring transform instead of the barycentric PSD.

<figure><img src="image_3.png" /><figcaption>Figure 2: Description of the running Riemannian barycenter. The barycenter of the batch $\hat{\mathbf{P}}_{\mathrm{B}}$ is estimated from the PSD of each batch sample. Then the running Riemannian barycenter is updated through an exponential average along the geodesic (- -), parameterized by $\alpha \in [0, 1]$.</figcaption></figure>

**Algorithm 1 Forward pass of PSDNorm**
1: **Input:** Batch $\mathcal{B} = \{\mathbf{G}^{(1)}, \dots, \mathbf{G}^{(N)}\}$, running barycenter $\hat{\mathbf{P}}$, filter-size $F$, momentum $\alpha$, training flag
2: **Output:** Normalized batch $\{\tilde{\mathbf{G}}^{(1)}, \dots, \tilde{\mathbf{G}}^{(N)}\}$
3: **for** $j = 1$ to $N$ **do**
4: &nbsp;&nbsp;&nbsp;&nbsp; $\hat{\mu}^{(j)} \leftarrow \text{Mean estimation}$
5: &nbsp;&nbsp;&nbsp;&nbsp; $\hat{\mathbf{P}}^{(j)} \leftarrow \text{PSD est. from } \tilde{\mathbf{G}}^{(j)} - \hat{\mu}^{(j)} \mathbf{1}_{\ell}^{\top}$ with eq. (3)
6: **end for**
7: **if** training **then**
8: &nbsp;&nbsp;&nbsp;&nbsp; $\hat{\mathbf{P}}_{\mathrm{B}} \leftarrow \text{Batch bary. from } \{\hat{\mathbf{P}}^{(j)}\}_{j}$ with eq. (5)
9: &nbsp;&nbsp;&nbsp;&nbsp; $\hat{\mathbf{P}} \leftarrow \text{Running bary. up. from } \hat{\mathbf{P}}, \hat{\mathbf{P}}_{\mathrm{B}}$ with eq. (6)
10: **end if**
11: **for** $j = 1$ to $N$ **do**
12: &nbsp;&nbsp;&nbsp;&nbsp; $\hat{\mathbf{H}}^{(j)} \leftarrow \text{Filter estimation from } \hat{\mathbf{P}}^{(j)}, \hat{\mathbf{P}}$ with eq. (8)
13: &nbsp;&nbsp;&nbsp;&nbsp; $\tilde{\mathbf{G}}^{(j)} \leftarrow \text{F-Monge mapping with eq. (7)}$
14: **end for**
# 4 Numerical Experiments

In this section, we evaluate the proposed method through a series of experiments designed to highlight its effectiveness and robustness on the clinically relevant task of sleep staging. We first describe the datasets and training setup employed, followed by a performance comparison with existing normalization techniques. Next, we assess the efficiency of PSDNorm by training over varying numbers of subjects per dataset. Finally, we analyze the robustness of PSDNorm against domain shift by focusing on subject-wise performance and different architectures. The code will be available on GitHub upon acceptance. The anonymized code is available in the supplementary material. All numerical experiments were conducted using a total of 1500 GPU hours on NVIDIA H100 GPUs.

## 4.1 Experimental Setup

**Datasets** To evaluate the effect of normalization layers, we use ten datasets of sleep staging described in Table 1. ABC [33], CCSHS [34], CFS [35], HPAP [36], MROS [37], SHHS [4], CHAT [5], and SOF [38] are publicly available sleep datasets with restricted access from National Sleep Research Resource (NSRR) [39]. PHYS [40] and MASS [3] are two other datasets publicly available. Every 30 s epoch is labeled with one of the five sleep stages: Wake, N1, N2, N3, and REM. These datasets are unbalanced in terms of age, sex, number of subjects, and have been recorded with different sensors in different institutions which makes the sleep staging task challenging. We now describe the pre-processing steps and splits of the datasets.

<table><caption>Table 1: Characteristics of the datasets.</caption><thead><tr><th>Dataset</th><th>Subj.</th><th>Rec.</th><th>Age ± std</th><th>Sex (F/M)</th></tr></thead><tbody><tr><td>ABC</td><td>44</td><td>117</td><td>48.8 ± 9.8</td><td>43%/57%</td></tr><tr><td>CCSHS</td><td>515</td><td>515</td><td>17.7 ± 0.4</td><td>50%/50%</td></tr><tr><td>CFS</td><td>681</td><td>681</td><td>41.7 ± 20.0</td><td>55%/45%</td></tr><tr><td>HPAP</td><td>166</td><td>166</td><td>46.5 ± 11.9</td><td>43%/57%</td></tr><tr><td>MROS</td><td>2101</td><td>2698</td><td>76.4 ± 5.5</td><td>0%/100%</td></tr><tr><td>PHYS</td><td>70</td><td>132</td><td>58.8 ± 22.0</td><td>33%/67%</td></tr><tr><td>SHHS</td><td>5730</td><td>8271</td><td>63.1 ± 11.2</td><td>52%/48%</td></tr><tr><td>MASS</td><td>61</td><td>61</td><td>42.5 ± 18.9</td><td>55%/45%</td></tr><tr><td>CHAT</td><td>1230</td><td>1635</td><td>6.6 ± 1.4</td><td>52%/48%</td></tr><tr><td>SOF</td><td>434</td><td>434</td><td>82.8 ± 3.1</td><td>100%/0%</td></tr><tr><td>Total</td><td>11032</td><td>14710</td><td>—</td><td>—</td></tr></tbody></table>

**Data Pre-processing** We follow a standard pre-processing pipeline used in the field [41, 42]. The datasets vary in the number and type of available EEG and electrooculogram (EOG) channels. To ensure consistency, we use two bipolar EEG channels, as some datasets lack additional channels. For dataset from NSRR, we select the channels C3-A2 and C4-A1. For signals from Physionet and MASS, we use the only available channels Fpz-Cz and Pz-Oz. The EEG signals are low-pass filtered with a 30 Hz cutoff frequency and resampled to 100 Hz. All data extraction and pre-processing steps are implemented using MNE-BIDS [43] and MNE-Python [44].

**Leave-One-Dataset-Out (LODO) Setup and Balancing** We evaluate model performance using a leave-one-dataset-out (LODO) protocol: in each fold, one dataset is held out for testing, and the model is trained on the union of the remaining datasets. From the training data, 80% of subjects are used for training and 20% for validation, which is used for early stopping. The full held-out dataset is used for testing. To assess performance in low-data regimes, we also evaluate a variant in which we subsample at most *N* subjects per dataset, promoting balanced contributions across training sources. We refer to this configuration as balanced@N, with *N* ranging from 40 to 400. The exact number of subjects per dataset in each case is listed in Appendix Table 3.

**Architecture and Training** Sleep staging has inspired a variety of neural architectures, from early CNN-based models [41, 42, 19] to recent attention-based approaches [21, 17, 20]. We evaluate two architectures: **U-Sleep** [45, 15], a state-of-the-art temporal CNN model designed for robustness and large-scale training, and a newly introduced architecture, **CNNTransformer**. CNNTransformer combines a lightweight convolutional encoder with a Transformer applied to epoch-level embeddings. It is specifically tailored for two-channel EEG and designed to scale efficiently to large datasets, while remaining minimal in implementation (under 100 lines of code) and training cost (Appendix A.3). Its design draws inspiration from recent transformer-based models for time series [46], with an emphasis on simplicity and practicality.

We use the Adam optimizer [47] with a learning rate of $10^{-3}$ to minimize the weighted cross-entropy loss, where class weights are computed from the training set distribution. Training is performed with a batch size of 64, and early stopping is applied based on validation loss with a patience of 3 epochs. Each input corresponds to a sequence of 17'30s, with a stride of 10'30s between sequences along the full-night recording.

**Evaluation** At test time, the model similarly processes sequences of 17'30s with a stride of 10'30s. Performance is evaluated using the balanced accuracy score (BACC), computed on the central 10'30s of each prediction window.
<table><caption>Table 2: **Balanced Accuracy (BACC) scores on the left-out datasets.** The top section reports results in the **large-scale** setting (using all available subjects), while the bottom section presents results in the **medium-scale** setting (balanced@400). For each row, the best score is highlighted in **bold**, and standard deviations reflect training variability across 3 random seeds.</caption><thead><tr><th></th><th></th><th>BatchNorm</th><th>LayerNorm</th><th>InstanceNorm</th><th>PSDNorm(F=5)</th><th>PSDNorm(F=9)</th><th>PSDNorm(F=17)</th></tr></thead><tbody><tr><td rowspan="10">All subjects</td><td>ABC</td><td>78.49 ± 0.42</td><td>77.94 ± 0.31</td><td><strong>78.83 ± 0.59</strong></td><td>78.56 ± 0.67</td><td>78.73 ± 0.32</td><td>78.60 ± 0.28</td></tr><tr><td>CCSHS</td><td><strong>88.79 ± 0.21</strong></td><td>87.51 ± 0.77</td><td>88.75 ± 0.04</td><td>88.56 ± 0.36</td><td>88.48 ± 0.07</td><td>88.52 ± 0.19</td></tr><tr><td>CFS</td><td>84.97 ± 0.37</td><td>84.29 ± 0.67</td><td><strong>85.73 ± 0.29</strong></td><td>85.42 ± 0.09</td><td>85.25 ± 0.19</td><td>85.28 ± 0.12</td></tr><tr><td>CHAT</td><td>64.72 ± 3.94</td><td>64.36 ± 0.40</td><td>68.86 ± 2.49</td><td>70.57 ± 1.24</td><td>70.36 ± 2.22</td><td><strong>70.71 ± 2.15</strong></td></tr><tr><td>HOMEPAP</td><td>76.39 ± 0.29</td><td>75.23 ± 0.78</td><td>76.70 ± 0.35</td><td>76.72 ± 0.27</td><td>76.93 ± 0.10</td><td><strong>77.02 ± 0.32</strong></td></tr><tr><td>MASS</td><td><strong>73.71 ± 0.62</strong></td><td>71.39 ± 3.00</td><td>72.12 ± 0.70</td><td>72.51 ± 1.68</td><td>71.34 ± 2.68</td><td>71.61 ± 0.71</td></tr><tr><td>MROS</td><td>81.30 ± 0.25</td><td>80.44 ± 0.29</td><td>81.49 ± 0.18</td><td><strong>81.57 ± 0.34</strong></td><td>81.35 ± 0.17</td><td>81.30 ± 0.13</td></tr><tr><td>PhysioNet</td><td>76.13 ± 0.57</td><td>75.12 ± 0.22</td><td>76.15 ± 0.52</td><td>75.96 ± 1.02</td><td><strong>76.35 ± 0.27</strong></td><td>76.02 ± 0.47</td></tr><tr><td>SHHS</td><td>77.97 ± 1.46</td><td>75.98 ± 0.48</td><td>79.05 ± 0.89</td><td>79.14 ± 1.01</td><td><strong>79.33 ± 0.62</strong></td><td>79.12 ± 0.11</td></tr><tr><td>SOF</td><td>81.33 ± 0.54</td><td>81.82 ± 0.79</td><td>81.98 ± 0.22</td><td><strong>82.50 ± 0.34</strong></td><td>82.15 ± 1.00</td><td>81.94 ± 0.65</td></tr><tr><td>Mean</td><td>78.38 ± 0.47</td><td>77.41 ± 0.28</td><td>78.97 ± 0.11</td><td><strong>79.15 ± 0.14</strong></td><td>79.03 ± 0.10</td><td>79.01 ± 0.36</td></tr><tr><td rowspan="10">Balanced@400</td><td>ABC</td><td>78.26 ± 1.33</td><td>75.29 ± 0.81</td><td><strong>78.73 ± 0.42</strong></td><td>78.18 ± 0.68</td><td>78.18 ± 0.91</td><td>77.76 ± 1.00</td></tr><tr><td>CCSHS</td><td>87.42 ± 0.16</td><td>85.20 ± 0.48</td><td><strong>87.62 ± 0.42</strong></td><td>87.58 ± 0.30</td><td>87.35 ± 0.52</td><td><strong>87.62 ± 0.48</strong></td></tr><tr><td>CFS</td><td>84.32 ± 0.57</td><td>81.66 ± 1.36</td><td><strong>84.72 ± 0.33</strong></td><td>84.29 ± 0.36</td><td>84.06 ± 0.10</td><td>84.46 ± 0.06</td></tr><tr><td>CHAT</td><td>66.55 ± 0.88</td><td>61.19 ± 1.16</td><td>64.43 ± 4.41</td><td><strong>70.28 ± 1.70</strong></td><td>68.11 ± 3.94</td><td>69.88 ± 0.46</td></tr><tr><td>HOMEPAP</td><td>75.25 ± 0.50</td><td>74.86 ± 0.25</td><td>76.47 ± 0.63</td><td><strong>76.83 ± 0.61</strong></td><td>76.61 ± 0.74</td><td>76.49 ± 0.45</td></tr><tr><td>MASS</td><td>70.00 ± 1.91</td><td>68.56 ± 3.33</td><td>71.52 ± 1.13</td><td>72.77 ± 1.09</td><td><strong>73.07 ± 1.30</strong></td><td>72.23 ± 2.40</td></tr><tr><td>MROS</td><td>80.37 ± 0.20</td><td>78.05 ± 0.22</td><td>80.28 ± 0.21</td><td>80.26 ± 0.11</td><td>80.32 ± 0.22</td><td><strong>80.70 ± 0.42</strong></td></tr><tr><td>PhysioNet</td><td><strong>75.81 ± 0.13</strong></td><td>71.82 ± 2.12</td><td>74.68 ± 0.55</td><td>74.82 ± 2.11</td><td>73.77 ± 1.73</td><td>75.09 ± 0.97</td></tr><tr><td>SHHS</td><td>76.44 ± 0.92</td><td>75.12 ± 0.39</td><td>78.68 ± 0.37</td><td><strong>78.88 ± 0.68</strong></td><td>77.28 ± 0.91</td><td>78.41 ± 0.49</td></tr><tr><td>SOF</td><td>81.08 ± 1.14</td><td>78.70 ± 0.50</td><td>80.68 ± 1.38</td><td>79.49 ± 0.41</td><td><strong>81.44 ± 0.97</strong></td><td>81.07 ± 0.66</td></tr><tr><td>Mean</td><td>77.55 ± 0.34</td><td>75.05 ± 0.28</td><td>77.78 ± 0.46</td><td>78.34 ± 0.42</td><td>78.02 ± 0.67</td><td><strong>78.37 ± 0.38</strong></td></tr></tbody></table>

Each experiment is repeated three times with different random seeds, and we report the mean and standard deviation of BACC.

**Normalization Strategies** We compare the proposed PSDNorm with three normalization strategies: BatchNorm, LayerNorm, and InstanceNorm. Note that InstanceNorm corresponds to a special case of PSDNorm with $F = 1$ and a fixed identity mapping instead of a learned running barycenter. In the following experiments, the BatchNorm layers in the first three convolutional layers are replaced with either PyTorch's default implementations of LayerNorm, InstanceNorm [48], or PSDNorm. To preserve the receptive field, the filter size $F$ of PSDNorm is used in the first layer and progressively halved in the following ones. We fix the momentum $\alpha$ to $10^{-2}$.

## 4.2 Numerical Results

This section presents results from large-scale sleep stage classification experiments. The analysis begins with a comparison of PSDNorm against standard normalization layers—BatchNorm, LayerNorm, and InstanceNorm—on the full datasets. Then, the data efficiency of each method is evaluated under limited training data regimes. Finally, robustness to distribution shift is assessed via subject-wise performance across multiple neural network architectures.

**Performance Comparison on Full Datasets** Table 2 (top) reports the LODO BACC of U-Sleep across all datasets, averaged over three random seeds. PSDNorm consistently outperforms all baseline normalization layers—BatchNorm, LayerNorm, and InstanceNorm—achieving the highest mean BACC of 79.15%, which exceeds BatchNorm (78.38%), InstanceNorm (78.97%), and LayerNorm (77.41%) by at least one standard deviation. This performance holds across tested filter sizes, with PSDNorm reaching 79.03% for $F=9$ and 79.01% for $F=17$. In total, PSDNorm ranks first on 6 out of the 10 datasets. On the challenging CHAT dataset, where all methods struggle, PSDNorm with $F=5$ outperforms all other normalizations by more than 2 percentage points, highlighting its robustness under strong distribution shifts. Although InstanceNorm is a strong baseline—outperforming BatchNorm and LayerNorm by at least one standard deviation on average—it is consistently surpassed by PSDNorm in average performance. In contrast, LayerNorm underperforms across the board, achieving the lowest average BACC and never ranking first, confirming its limited suitability for this task.
**Efficiency: Performance with 4× Less Data** The PSDNorm layer improves model performance when trained on the full dataset (~10000 subjects), but such large-scale data availability is not always the case. In many real-world scenarios—such as rare disease studies, pediatric populations, or data collected in constrained clinical settings—labeled recordings are scarce, expensive to annotate, or restricted due to privacy concerns. Evaluating model robustness under these constraints is therefore essential. To this end, we train all models using the balanced@400 setup, which reduces the training data by a factor of 4 compared to the full-data setting. In this lower-data regime, PSDNorm continues to outperform all baseline normalization strategies on 6 out of 10 datasets and achieves

higher average BACC across all tested values of F. The performance improvement of PSDNorm over the best baseline is more pronounced in this setting: for F=5, the BACC gain reaches +0.56%, compared to +0.18% in the full-data setting. Again, the gains exceed one standard deviation. To assess statistical significance, we conducted a critical difference (CD) test [49]. Figure 3 (top) reports the average rank of each method and the corresponding statistical comparisons. The results confirm that PSDNorm significantly outperforms the baselines, underscoring the value of incorporating temporal structure into normalization for robust and data-efficient generalization. The same trend is observed for U-Sleep trained on all subjects (see in Appendix Figure 6). The following experiments focus on the balanced@400 setup.

**Robustness Across Architectures** PSDNorm is a plug-and-play normalization layer that can be seamlessly integrated into various neural network architectures. To demonstrate this flexibility, we evaluate its performance on both the U-Sleep and CNNTransformer models. Figure 3 reports the average rank of each normalization method across datasets and subjects for both architectures using datasets balanced@400. In both architectures, PSDNorm with *F*=5 achieves the best overall ranking and demonstrates statistically significant improvements over both BatchNorm, LayerNorm and InstanceNorm. Notably, it is also the only method to rank among the top two across both architectures. The results confirm that PSDNorm generalizes well beyond a single architecture and can provide consistent improvements in diverse modeling setups. InstanceNorm performs competitively in some cases but is never significantly better than PSDNorm. The figure also illustrates that PSDNorm with *F*=17 does not yield additional benefit in this medium-scale setting, further supporting the conclusion that moderate temporal context is sufficient at this data scale. Detailed numerical scores for CNNTransformer are reported in the supplementary material (Table 7).

<figure><img src="image_4.png" /><figcaption>Figure 3: Critical Difference (CD) diagram for two architectures on datasets balanced @400. Average ranks across datasets and subjects for USleep and CNNTransformer. Black lines connect methods that are not significantly different.</figcaption></figure>

<figure><img src="image_5.png" /><figcaption>Figure 4: Performance of PSDNorm and BatchNorm with varying training set sizes. The BACC score is plotted against the number of training subjects used with U-Sleep.</figcaption></figure>

**Sensitivity to Filter Size** The choice of *F* in PSDNorm controls the intensity of the normalization: larger *F* provide stronger normalization, while smaller *F* allow more flexibility in the model. In Figure 4, we evaluate its impact across different training set sizes and observe a clear trend: when trained on fewer subjects, larger filter sizes yield better performance (i.e., *F* = 17), whereas smaller filter sizes are more effective with larger datasets (i.e., *F* = 5). This suggests that with limited data, stronger normalization helps prevent overfitting, while with more data, a more flexible model is preferred. On average, PSDNorm with *F* = 5 offers a good compromise, achieving one of the best performances across all training set sizes.
**Performance on the most challenging subjects** Performance variability across subjects is a key challenge in biomedical applications where ensuring consistently high performance—even for the most challenging subjects—is critical. To highlight the robustness of PSDNorm, Figure 5 presents a scatter plot of subject-wise BACC scores comparing BatchNorm or InstanceNorm vs. PSDNorm across two selected target datasets. CHAT and MASS are two challenging datasets, where the prediction performance is significantly lower than the other datasets. For CHAT, most of the dots are below the diagonal, indicating that PSDNorm improves performance for 91% of subjects against BatchNorm and 99% of subjects against InstanceNorm, with the largest gains observed for the hardest subjects, reinforcing its ability to handle challenging cases. For MASS, PSDNorm improves performance for 75% of subjects against BatchNorm and 69% against InstanceNorm. This demonstrates that PSDNorm is not only effective in improving overall performance but also excels in enhancing the performance of the most challenging subjects.

<figure><img src="image_6.png" /><figcaption>Figure 5: **Subject-wise BACC comparison on MASS and CHAT (balanced @400).** Blue dot means improvement with PSDNorm.</figcaption></figure>

# 5 Conclusion, Limitations, and Future Work

This paper introduced PSDNorm, a normalization layer that aligns the power spectral density (PSD) of each signal to a geodesic barycenter. By leveraging temporal correlations, PSDNorm offers a principled alternative to standard normalization layers. Experiments on large-scale sleep staging datasets show that PSDNorm consistently improves performance, robustness, and data efficiency, especially under domain shift and limited-data settings—outperforming BatchNorm, LayerNorm, and InstanceNorm across architectures.

While the results are promising, some limitations remain. PSDNorm introduces a filter size hyperparameter ($F$) that controls normalization strength; although we provide default values that perform well across datasets, selecting it automatically in adaptive settings could be challenging.

Despite these limitations, PSDNorm is flexible and easy to integrate into existing models. Future work includes extending it to other signals such as audio and other biomedical applications.

# Acknowledgements

This work was supported by the grants ANR-22-PESN-0012 to AC under the France 2030 program, ANR-20-CHIA-0016 and ANR-20-IADJ-0002 to AG while at Inria, and ANR-23-ERCC-0006 to RF, all from Agence nationale de la recherche (ANR). This project has also received funding from the European Union's Horizon Europe research and innovation programme under grant agreement 101120237 (ELIAS).

This project received funding from the Fondation de l'École polytechnique

This project was provided with computer and storage resources by GENCI at IDRIS thanks to the grant 2025-AD011016052 and 2025-AD011016067 on the supercomputer Jean Zay's the V100 & H100 partitions.

This work was conducted at Inria, AG is presently employed by Meta Platforms. All the datasets used for this work were accessed and processed on the Inria compute infrastructure.

All the datasets used for this work were accessed and processed on the Inria compute infrastructures. Numerical computation was enabled by the scientific Python ecosystem: Matplotlib [50], Scikit-learn [51], Numpy [52], Scipy [53], PyTorch [48] and MNE [44].
## References

[1] J. G. Moreno-Torres, T. Raeder, R. Alaiz-Rodríguez, N. V. Chawla, and F. Herrera, "A unifying view on dataset shift in classification," *Pattern Recognition* **45** no. 1, (2012) 521–530.

[2] S. Stevens and G. Clark, "Chapter 6 - polysomnography," in *Sleep Medicine Secrets*, D. STEVENS, ed., pp. 45–63. Hanley & Belfus, 2004.

[3] C. O'Reilly, N. Gosselin, and J. Carrier, "Montreal archive of sleep studies: an open-access resource for instrument benchmarking and exploratory research," *Journal of sleep research* **23** (06, 2014).

[4] S. Quan, B. Howard, et al., "The sleep heart health study: Design, rationale, and methods," *Sleep* **20** (01, 1998) 1077–85.

[5] C. L. Marcus, R. H. Moore, et al., "A randomized trial of adenotonsillectomy for childhood sleep apnea," *The New England Journal of Medicine* **368** no. 25, (June, 2013) 2366–2376.

[6] S. Ioffe and C. Szegedy, "Batch normalization: accelerating deep network training by reducing internal covariate shift," in *Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37*, ICML'15, p. 448–456. JMLR.org, 2015.

[7] J. L. Ba, J. R. Kiros, and G. E. Hinton, "Layer normalization." 2016. http://arxiv.org/abs/1607.06450.

[8] D. Ulyanov, "Instance normalization: The missing ingredient for fast stylization," *arXiv preprint arXiv:1607.08022* (2016).

[9] R. Kobler, J.-i. Hirayama, Q. Zhao, and M. Kawanabe, "Spd domain-specific batch normalization to crack interpretable unsupervised domain adaptation in EEG," in *Advances in Neural Information Processing Systems*, S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, eds., vol. 35, pp. 6219–6235. Curran Associates, Inc., 2022.

[10] T. Kim, J. Kim, Y. Tae, C. Park, J.-H. Choi, and J. Choo, "Reversible instance normalization for accurate time-series forecasting against distribution shift," in *International Conference on Learning Representations*. 2021.

[11] A. Apicella, F. Isgrò, A. Pollastro, and R. Prevete, "On the effects of data normalization for domain adaptation on EEG data," *Engineering Applications of Artificial Intelligence* **123** (2023) 106205.

[12] S. Chambon, M. N. Galtier, P. J. Arnal, G. Wainrib, and A. Gramfort, "A deep learning architecture for temporal sleep stage classification using multivariate and multimodal time series," *IEEE Transactions on Neural Systems and Rehabilitation Engineering* **26** no. 4, (2018) 758–769.

[13] T. Gnassounou, R. Flamary, and A. Gramfort, "Convolutional monge mapping normalization for learning on biosignals," in *Neural Information Processing Systems (NeurIPS)*. 2023.

[14] T. Gnassounou, A. Collas, R. Flamary, K. Lounici, and A. Gramfort, "Multi-source and test-time domain adaptation on multivariate signals using spatio-temporal monge alignment," *arXiv preprint arXiv:2407.14303* (2024).

[15] M. Perslev, S. Darkner, L. Kempfner, M. Nikolic, P. Jennum, and C. Igel, "U-Sleep: resilient high-frequency sleep staging," *npj Digital Medicine* **4** (04, 2021) 72.

[16] A. Guillot and V. Thorey, "RobustSleepNet: Transfer learning for automated sleep staging at scale," *arXiv:2101.02452* [stat.ML].

[17] H. Phan, K. P. Lorenzen, E. Heremans, O. Y. Chén, M. C. Tran, P. Koch, A. Mertins, M. Baumert, K. B. Mikkelsen, and M. De Vos, "L-SeqSleepNet: Whole-cycle Long Sequence Modeling for Automatic Sleep Staging," *IEEE Journal of Biomedical and Health Informatics* **27** no. 10, (Oct., 2023) 4748–4757. https://ieeexplore.ieee.org/document/10210638/?arnumber=10210638. Conference Name: IEEE Journal of Biomedical and Health Informatics.

[18] H. Phan, F. Andreotti, N. Cooray, O. Y. Chén, and M. D. Vos, "SeqSleepNet: End-to-End Hierarchical Recurrent Neural Network for Sequence-to-Sequence Automatic Sleep Staging." Feb., 2019. http://arxiv.org/abs/1809.10932. arXiv:1809.10932.

[19] H. Phan, O. Y. Chen, M. C. Tran, P. Koch, A. Mertins, and M. D. Vos, "XSleepNet: Multi-view sequential model for automatic sleep staging," *IEEE Transactions on Pattern Analysis and Machine Intelligence* **44** no. 09, (Sep, 2022) 5903–5915.

[20] J. Wang, S. Zhao, H. Jiang, Y. Zhou, Z. Yu, T. Li, S. Li, and G. Pan, "CareSleepNet: A hybrid deep learning network for automatic sleep staging," https://ieeexplore.ieee.org/document/10595067/.
[21] H. Phan, K. Mikkelsen, O. Y. Chén, P. Koch, A. Mertins, and M. De Vos, “Sleeptransformer: Automatic sleep staging with interpretability and uncertainty quantification,” *IEEE Transactions on Biomedical Engineering* **69** no. 8, (2022) 2456–2467.

[22] Y. Guo, M. Nowakowski, and W. Dai, “Flexsleeptransformer: a transformer-based sleep staging model with flexible input channel configurations,” *Scientific Reports* **14** (11, 2024).

[23] R. Thapa, M. R. Kjær, et al., “A multimodal sleep foundation model developed with 500k hours of sleep recordings for disease predictions,” https://www.ncbi.nlm.nih.gov/pmc/articles/PMC11838666/.

[24] B. Fox, J. Jiang, S. Wickramaratne, P. Kovatch, M. Suarez-Farinas, N. A. Shah, A. Parekh, and G. N. Nadkarni, “A foundational transformer leveraging full night, multichannel sleep study data accurately classifies sleep stages,”.

[25] G. Deng, M. Niu, et al., “A unified flexible large polysomnography model for sleep staging and mental disorder diagnosis,” https://www.ncbi.nlm.nih.gov/pmc/articles/PMC11661386/.

[26] B. Sun and K. Saenko, “Deep CORAL: Correlation Alignment for Deep Domain Adaptation.” July, 2016. http://arxiv.org/abs/1607.01719. arXiv:1607.01719 [cs].

[27] B. B. Damodaran, B. Kellenberger, R. Flamary, D. Tuia, and N. Courty, “Deepjdot: Deep joint distribution optimal transport for unsupervised domain adaptation,” in *Proceedings of the European conference on computer vision (ECCV)*, pp. 447–463. 2018.

[28] E. Eldele, M. Ragab, Z. Chen, M. Wu, C. Kwoh, X. Li, and C. Guan, “Adast: Attentive cross-domain EEG-based sleep staging framework with iterative self-training,” *IEEE Transactions on Emerging Topics in Computational Intelligence* **7** (2021) 210–221.

[29] R. M. Gray, “Toeplitz and circulant matrices: A review,” *Foundations and Trends® in Communications and Information Theory* **2** no. 3, (2006) 155–239. http://dx.doi.org/10.1561/0100000006.

[30] P. Welch, “The use of fast fourier transform for the estimation of power spectra: a method based on time averaging over short, modified periodograms,” *IEEE Transactions on audio and electroacoustics* **15** no. 2, (1967) 70–73.

[31] M. Agueh and G. Carlier, “Barycenters in the wasserstein space,” *SIAM Journal on Mathematical Analysis* **43** no. 2, (2011) 904–924.

[32] R. Bhatia, T. Jain, and Y. Lim, “On the bures–wasserstein distance between positive definite matrices,” *Expositiones Mathematicae* **37** no. 2, (2019) 165–191.

[33] B. Jessie P., T. Ali, R. Michael, W. Wei, A. Robert, M. Atul, O. Robert L., A. Amit, D. Katherine, and P. Sanya R., “Gastric Banding Surgery versus Continuous Positive Airway Pressure for Obstructive Sleep Apnea: A Randomized Controlled Trial,” *American journal of respiratory and critical care medicine* **197** no. 8, (Apr., 2018). https://pubmed.ncbi.nlm.nih.gov/29035093/. Publisher: Am J Respir Crit Care Med.

[34] C. L. Rosen, E. K. Larkin, H. L. Kirchner, J. L. Emancipator, S. F. Bivins, S. A. Surovec, R. J. Martin, and S. Redline, “Prevalence and risk factors for sleep-disordered breathing in 8- to 11-year-old children: association with race and prematurity,” *The Journal of Pediatrics* **142** no. 4, (Apr., 2003) 383–389.

[35] S. Redline, P. V. Tishler, T. D. Tosteson, J. Williamson, K. Kump, I. Browner, V. Ferrette, and P. Krejci, “The familial aggregation of obstructive sleep apnea,” *American Journal of Respiratory and Critical Care Medicine* **151** no. 3 Pt 1, (Mar., 1995) 682–687.

[36] C. L. Rosen, D. Auckley, R. Benca, N. Foldvary-Schaefer, C. Iber, V. Kapur, M. Rueschman, P. Zee, and S. Redline, “A multisite randomized trial of portable sleep studies and positive airway pressure autotitration versus laboratory-based polysomnography for the diagnosis and treatment of obstructive sleep apnea: the HomePAP study,” *Sleep* **35** no. 6, (June, 2012) 757–767.

[37] T. Blackwell, K. Yaffe, S. Ancoli-Israel, S. Redline, K. E. Ensrud, M. L. Stefanick, A. Laffan, K. L. Stone, and Osteoporotic Fractures in Men Study Group, “Associations between sleep architecture and sleep-disordered breathing and cognition in older community-dwelling men: the Osteoporotic Fractures in Men Sleep Study,” *Journal of the American Geriatrics Society* **59** no. 12, (Dec., 2011) 2217–2225.

[38] A. P. Spira, T. Blackwell, K. L. Stone, S. Redline, J. A. Cauley, S. Ancoli-Israel, and K. Yaffe, “Sleep-disordered breathing and cognition in older women,” *Journal of the American Geriatrics Society* **56** no. 1, (Jan., 2008) 45–50.

[39] G.-Q. Zhang, L. Cui, R. Mueller, S. Tao, M. Kim, M. Rueschman, S. Mariani, D. Mobley, and S. Redline, “The National Sleep Research Resource: towards a sleep data commons,” *Journal of the American Medical Informatics Association: JAMIA* **25** no. 10, (Oct., 2018) 1351–1358.

[40] A. Goldberger, L. Amaral, et al., “Components of a new research resource for complex physiologic signals,” *PhysioNet* **101** (01, 2000).
[41] S. Chambon, M. Galtier, P. Arnal, G. Wainrib, and A. Gramfort, "A deep learning architecture for temporal sleep stage classification using multivariate and multimodal time series," *arXiv:1707.03321 [stat.ML]*.

[42] J. B. Stephansen, A. N. Olesen, et al., "Neural network analysis of sleep stages enables efficient diagnosis of narcolepsy," *Nature Communications* **9** no. 1, (Dec, 2018).

[43] S. Appelhoff, M. Sanderson, et al., "MNE-BIDS: Organizing electrophysiological data into the BIDS format and facilitating their analysis," *Journal of Open Source Software* **4** no. 44, (2019) 1896.

[44] A. Gramfort, M. Luessi, et al., "MEG and EEG data analysis with MNE-Python," *Frontiers in Neuroscience* **7** no. 267, (2013) 1-13.

[45] M. Perslev, M. Jensen, S. Darkner, P. J. Jennum, and C. Igel, "U-time: A fully convolutional network for time series segmentation applied to sleep staging," *Advances in Neural Information Processing Systems* **32** (2019).

[46] C. Yang, M. Westover, and J. Sun, "Biot: Biosignal transformer for cross-data learning in the wild," in *Advances in Neural Information Processing Systems*, A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, eds., vol. 36, pp. 78240-78260. Curran Associates, Inc., 2023. https://proceedings.neurips.cc/paper_files/paper/2023/file/f6b30f3e2dd9cb53bbf2024402d02295-Paper-Conference.pdf.

[47] D. P. Kingma, "Adam: A method for stochastic optimization," *arXiv preprint arXiv:1412.6980* (2014).

[48] A. Paszke, S. Gross, et al., "Pytorch: An imperative style, high-performance deep learning library," in *Advances in Neural Information Processing Systems (NeurIPS)*, pp. 8024-8035. Curran Associates, Inc., 2019.

[49] J. Demšar, "Statistical comparisons of classifiers over multiple data sets," *Journal of Machine Learning Research* **7** no. 1, (2006) 1-30. http://jmlr.org/papers/v7/demsar06a.html.

[50] J. D. Hunter, "Matplotlib: A 2d graphics environment," *Computing in science & engineering* **9** no. 3, (2007) 90-95.

[51] F. Pedregosa, G. Varoquaux, et al., "Scikit-learn: Machine Learning in Python," *Journal of Machine Learning Research* **12** (2011) 2825-2830.

[52] C. Harris, K. Millman, et al., "Array programming with NumPy," *Nature* **585** no. 7825, (2020) 357-362.

[53] P. Virtanen, R. Gommers, et al., "SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python," *Nature Methods* **17** (2020) 261-272.

[54] O. Ronneberger, P. Fischer, and T. Brox, "U-net: Convolutional networks for biomedical image segmentation," in *Medical image computing and computer-assisted intervention-MICCAI 2015: 18th international conference*, Munich, Germany, October 5-9, 2015, *proceedings, part III 18*, pp. 234-241, Springer. 2015.

[55] R. T. Schirrmeister, J. T. Springenberg, L. D. J. Fiederer, M. Glasstetter, K. Eggensperger, M. Tangermann, F. Hutter, W. Burgard, and T. Ball, "Deep learning with convolutional neural networks for EEG decoding and visualization," *Human Brain Mapping* **38** no. 11, (Nov., 2017) 5391-5420. http://arxiv.org/abs/1703.05051. arXiv:1703.05051 [cs].
# A Appendix

## A.1 Proof of the Bures-Wasserstein geodesic (6) between covariance matrices of structure (2)

**Proposition A.1.** Let $\Sigma^{(s)}$ and $\Sigma^{(t)}$ be two covariance matrices in $\mathbb{R}^{c_F \times c_F}$ following (2). Let us denote $\mathbf{P}^{(s)}$ and $\mathbf{P}^{(t)}$ the corresponding PSD matrices. The geodesic associated with the Bures-Wasserstein metric between $\Sigma^{(s)}$ and $\Sigma^{(t)}$ and parameterized by $\alpha \in [0, 1]$ is $\Sigma(\alpha)$ following (2) of PSD

$$ \mathbf{P}(\alpha) = \left( (1 - \alpha) \mathbf{P}^{(s)} \odot \frac{1}{2} + \alpha \mathbf{P}^{(t)} \odot \frac{1}{2} \right)^{\odot 2}. $$

*Proof.* From [32], the geodesic associated with the Bures-Wasserstein metric between two covariance matrices $\Sigma^{(s)}$ and $\Sigma^{(t)}$ is given by

$$ \gamma(\alpha) = (1 - \alpha)^2 \Sigma^{(s)} + \alpha^2 \Sigma^{(t)} + \alpha(1 - \alpha) \left[ (\Sigma^{(s)} \Sigma^{(t)})^{\frac{1}{2}} + (\Sigma^{(t)} \Sigma^{(s)})^{\frac{1}{2}} \right]. \quad (9) $$

where

$$ (\Sigma^{(s)} \Sigma^{(t)})^{\frac{1}{2}} = \Sigma^{(s)}^{\frac{1}{2}} \left( \Sigma^{(s)}^{\frac{1}{2}} \Sigma^{(t)} \Sigma^{(s)}^{\frac{1}{2}} \right)^{\frac{1}{2}} \Sigma^{(s)}^{-\frac{1}{2}}. \quad (10) $$

Since $\Sigma^{(s)}$ and $\Sigma^{(t)}$ diagonalize in the unitary basis $\mathbf{I}_c \otimes \mathbf{F}_F$, $\gamma(\alpha)$ also diagonalizes in this basis. Thus, we only have to compute the geodesic between the PSD matrices $\mathbf{P}^{(s)}$ and $\mathbf{P}^{(t)}$ and from now on, all operations are element-wise. Let $\mathbf{P}(\alpha)$ be the PSD of $\gamma(\alpha)$, we have

$$ \mathbf{P}(\alpha) = (1 - \alpha)^2 \mathbf{P}^{(s)} + \alpha^2 \mathbf{P}^{(t)} + \alpha(1 - \alpha) \left[ (\mathbf{P}^{(s)} \odot \mathbf{P}^{(t)})^{\odot \frac{1}{2}} + (\mathbf{P}^{(t)} \odot \mathbf{P}^{(s)})^{\odot \frac{1}{2}} \right] \quad (11) $$

$$ = (1 - \alpha)^2 \mathbf{P}^{(s)} + \alpha^2 \mathbf{P}^{(t)} + 2\alpha(1 - \alpha)(\mathbf{P}^{(s)} \odot \mathbf{P}^{(t)})^{\odot \frac{1}{2}} \quad (12) $$

$$ = \left( (1 - \alpha) \mathbf{P}^{(s)} \odot \frac{1}{2} + \alpha \mathbf{P}^{(t)} \odot \frac{1}{2} \right)^{\odot 2}. \quad (13) $$

This concludes the proof.

<figure><img src="image_7.png" /></figure>

## A.2 Balanced datasets

In the main paper, we report results across different training set sizes. Since the datasets are highly imbalanced (e.g., ABC has 44 subjects, SHHS has 5,730), we create balanced subsets by randomly selecting up to *N* subjects per dataset. This avoids over-representing the largest dataset and ensures greater diversity in the training data. We consider four values of *N*: 40, 100, 200, and 400. The average number of subjects in each balanced set is shown in Table 3. Notably, the balanced set with 400 subjects contains roughly four times less data than the full dataset.

<table><caption>Table 3: Number of samples in the balanced datasets. Average and standard deviation (across LODO) are computed over 10 datasets left-out from the training set.</caption><thead><tr><td>Balanced datasets</td><td>Number of subjects</td></tr></thead><tbody><tr><td>Balanced@40</td><td>360 ± 0</td></tr><tr><td>Balanced@100</td><td>787 ± 19</td></tr><tr><td>Balanced@200</td><td>1387 ± 63</td></tr><tr><td>Balanced@400</td><td>2466 ± 157</td></tr><tr><td>All subjects</td><td>9929 ± 1659</td></tr></tbody></table>

## A.3 U-Time: CNN for time series segmentation

U-Time [45, 15] is a convolutional neural network (CNN) inspired by the U-Net architecture [54], designed for segmenting temporal sequences. U-Time maps sequential inputs of arbitrary length to sequences of class labels on a freely chosen temporal scale. The architecture is composed of several encoder and decoder blocks, with skip connections between them.

**Encoder blocks** A single encoder block is composed of a convolutional layer, an activation function, a BatchNorm layer, and a max pooling layer. First, the convolution is applied to the input signal, followed by the activation function and the BatchNorm layer. Finally, the max pooling layer downsamples the temporal dimension. In the following, the pre-BatchNorm feature map is denoted $\mathbf{G}$ and the post-BatchNorm feature map $\tilde{\mathbf{G}}$, i.e., $\tilde{\mathbf{G}} \triangleq \text{BatchNorm}(\mathbf{G})$. Each encoder block downsamples by 2 the signal length but increases the number of channels.
**Decoder blocks and Segmentation Head** The decoding part of U-Time is symmetrical to the encoding part. Each decoder block doubles the signal length and decreases the number of channels. It is composed of a convolutional layer, an activation function, a BatchNorm layer, an upsampling layer and a concatenation layer of the skip connection of the corresponding encoding block. Finally, the segmentation head applies two convolutional layers with an activation function in between to output the final segmentation. It should be noted that U-Time employs BatchNorm layers but other normalization layers, such as LayerNorm [7] or InstanceNorm [8] are possible.

**Implementation** The architecture is inspired from Braindecode [55]. The implementation is improved to make it more efficient and faster. One epoch of training takes about 30 min on a single H100 GPU.

## Architecture: CNNTransformer

The CNNTransformer is a hybrid architecture designed for multichannel time series classification inspired by transformers for EEG-Data [20, 21, 46, 23]. It combines convolutional feature extraction with long-range temporal modeling via a Transformer encoder at epoch-level. The model processes an input tensor of shape (B, S, C, T), where B is the batch size, S is the number of temporal segments, C is the number of input channels, and T is the number of time samples per segment. It outputs a tensor of shape (B, n_classes, S), where n_classes is the number of classes and S is the number of epochs.

The architecture consists of the following components:

* **Reshaping:** The input is first permuted and reshaped to a 3D tensor of shape (B, C, S, T) to be compatible with 1D convolutional layers applied along the temporal dimension.

* **CNN-based Feature Extractor:** A stack of 10 Conv1D layers, each followed by ELU activation and Batch Normalization. Some layers use a stride greater than 1 to progressively reduce the temporal resolution. This block extracts local temporal patterns and increases the representational capacity up to a dimensionality of $d_{model}$.

* **Adaptive Pooling:** An AdaptiveAvgPool1D layer reduces the temporal length to a fixed number of steps (S), independent of the input sequence length. This step ensures a consistent temporal resolution before the Transformer.

* **Positional Encoding:** Learnable positional embeddings of shape (1, S, $d_{model}$) are added to the feature representations to preserve temporal ordering before passing through the Transformer encoder.

* **Transformer Encoder:** A standard Transformer encoder composed of L layers, each consisting of multi-head self-attention and a feedforward sublayer. This module models global temporal dependencies across the S steps.

* **Classification Head:** After transposing the data to shape (B, $d_{model}$, S), a final 1D convolution with a kernel size of 1 projects the output to $n_{classes}$, yielding predictions for each epoch segment.

The model is trained end-to-end using standard optimization techniques. The use of adaptive pooling and self-attention enables it to generalize across variable-length inputs while maintaining temporal resolution. A full summary of the architecture is provided in Table 4.

## A.4 Equation for BatchNorm and InstanceNorm

**BatchNorm** The BatchNorm layer [6] normalizes features maps in a neural network to have zero mean and unit variance. At train time, given a batch $B = \{G^{(1)}, \dots, G^{(N)}\} \subset \mathbb{R}^{c \times l}$ of N pre-BatchNorm feature maps and for all $j, m, l \in [1, N] \times [1, c] \times [1, l]$, the BatchNorm layer is computed as

$$ \tilde{G}_{m,l}^{(j)} = \gamma_m \frac{G_{m,l}^{(j)} - \hat{\mu}_m}{\sqrt{\hat{\sigma}_m^2 + \varepsilon}} + \beta_m, \quad (14) $$

where $\gamma, \beta \in \mathbb{R}^c$ are learnable parameters. The mean and standard deviation $\hat{\mu} \in \mathbb{R}^c$ and $\hat{\sigma} \in \mathbb{R}^c$ are computed across the time and the batch,

$$ \begin{aligned} \hat{\mu}_m &\triangleq \frac{1}{Nl} \sum_{j=1}^{N} \sum_{l=1}^{l} G_{m,l}^{(j)}, \\ \hat{\sigma}_m^2 &\triangleq \frac{1}{Nl} \sum_{j=1}^{N} \sum_{l=1}^{l} \left( G_{m,l}^{(j)} - \hat{\mu}_m \right)^2. \end{aligned} \quad (15) $$
<table><caption>Table 4: Architecture overview of the CNNTransformer model. In pratice, $d_{model}$ is set to 768, $n_{head}$ to 8, and $S$ is 35.</caption><thead><tr><th>Stage</th><th>Operation</th><th>Details</th><th>Output Shape</th></tr></thead><tbody><tr><td>Input</td><td>Raw signal</td><td>Multichannel EEG signal with S segments and T time samples per segment</td><td>(B, S, C, T)</td></tr><tr><td>Reshape</td><td>Permute & flatten</td><td>Rearranged as (B, C, S · T) to process with 1D convolutions</td><td>(B, C, S · T)</td></tr><tr><td>Feature Extractor</td><td>1D CNN stack</td><td>10-layer sequence of Conv1D → ELU → Batch-Norm; includes temporal downsampling via stride</td><td>(B, d<sub>model</sub>, T')</td></tr><tr><td>Temporal Pooling</td><td>AdaptiveAvgPool1D</td><td>Downsamples to fixed temporal resolution defined by S</td><td>(B, d<sub>model</sub>, S)</td></tr><tr><td>Positional Encoding</td><td>Learnable embeddings</td><td>Added to temporal dimension to encode temporal order before transformer layers</td><td>(B, d<sub>model</sub>, S)</td></tr><tr><td>Transformer Encoder</td><td>Multi-head attention</td><td>2 Transformer layers with d<sub>model</sub> embedding dimension, n<sub>head</sub> heads, and feedforward sublayers</td><td>(B, d<sub>model</sub>, S)</td></tr><tr><td>Classifier</td><td>Linear projection</td><td>Projects feature vectors to class logits at each epoch time step</td><td>(B, n<sub>classes</sub>, S)</td></tr></tbody></table>

At test time, the mean and variance $\hat{\mu}$ and $\hat{\sigma}$ are replaced by their running mean and variance, also called exponential moving average, estimated during training.

**InstanceNorm** Another popular normalization is the InstanceNorm layer [8]. During training, InstanceNorm operates similarly to (14), but the mean and variance are computed per sample instead of across the batch dimension, i.e., $\hat{\mu}_m^{(j)}$ and $\hat{\sigma}_m^{(j)}$ are computed for each sample $j$,

$$ \hat{\mu}_m^{(j)} \triangleq \frac{1}{\ell} \sum_{l=1}^{\ell} G_{m,l}^{(j)}, \\ (\hat{\sigma}_m^{(j)})^2 \triangleq \frac{1}{\ell} \sum_{l=1}^{\ell} (G_{m,l}^{(j)} - \hat{\mu}_m^{(j)})^2. \qquad (16) $$

Hence, each sensor of each sample is normalized independently of the others. At test time, InstanceNorm behaves identically to its training phase and therefore does not rely on running statistics contrary to the BatchNorm.

## A.5 F1 score vs. Balanced Accuracy

In the main paper, we report Balanced Accuracy scores, which account for class imbalance in sleep stage classification. Prior work, such as the U-Time paper [45], uses the F1 score to evaluate performance. In Table 5, we report F1 scores on the left-out datasets. These scores are slightly higher than the Balanced Accuracy scores and are comparable to those reported in the U-Time paper.

Our main findings remain consistent: BatchNorm and InstanceNorm are the strongest baselines and achieve the best performance on 3 out of 10 datasets. PSDNorm outperforms all other methods on 7 out of 10 datasets. The same trend holds for the balanced@400 setup, where PSDNorm again outperforms all baselines on 7 datasets, while InstanceNorm is never the top performer.

These results confirm that our implementation achieves state-of-the-art performance in sleep stage classification. Moreover, PSDNorm maintains its advantage even in data-limited settings

## A.6 Impact of Whitening and Target Covariance

As explained in the main paper, InstanceNorm is a special case of PSDNorm with $F = 1$ and an identity target covariance matrix (i.e., whitening). PSDNorm extends this by (i) using temporal context with $F > 1$, and (ii) mapping the PSD to a target covariance matrix, such as a barycenter (i.e., colorization).

In this section, we evaluate the impact of whitening on the performance of PSDNorm, to assess the benefit of using a barycenter as the target covariance matrix. Table 6 reports results on 10 datasets (balanced@400), with and without whitening.

Whitening improves performance on only one dataset (CCSHS), while projecting to the barycenter yields the best results on 6 datasets.
<table><caption>Table 5: F1 scores of different methods on the left-out datasets. The lower section displays results for training over datasets balanced @400 i.e., small-scale dataset, while the upper section presents results for training over all subjects i.e., large-scale dataset. The best scores are highlighted in bold. The reported standard deviations indicate performance variability across 3 seeds.</caption><thead><tr><th></th><th></th><th>BatchNorm</th><th>LayerNorm</th><th>InstanceNorm</th><th>PSDNorm(F=5)</th><th>PSDNorm(F=9)</th><th>PSDNorm(F=17)</th></tr></thead><tbody><tr><td rowspan="10">All subjects</td><td>ABC</td><td><strong>79.80 ± 0.34</strong></td><td>77.86 ± 0.80</td><td>78.36 ± 1.20</td><td>78.08 ± 0.78</td><td>76.93 ± 1.29</td><td>77.44 ± 0.58</td></tr><tr><td>CCSHS</td><td>88.32 ± 0.49</td><td>87.22 ± 0.51</td><td>88.73 ± 0.52</td><td>88.79 ± 0.99</td><td>87.86 ± 1.08</td><td><strong>88.91 ± 0.28</strong></td></tr><tr><td>CFS</td><td>87.01 ± 0.18</td><td>85.61 ± 0.16</td><td><strong>87.62 ± 0.27</strong></td><td>87.06 ± 0.77</td><td>87.26 ± 0.42</td><td>86.52 ± 0.99</td></tr><tr><td>CHAT</td><td>66.56 ± 1.42</td><td>61.32 ± 2.25</td><td>64.19 ± 4.63</td><td><strong>71.86 ± 0.95</strong></td><td>68.61 ± 5.36</td><td>71.12 ± 1.09</td></tr><tr><td>HOMEPAP</td><td>76.20 ± 1.25</td><td>76.15 ± 1.13</td><td>77.66 ± 0.58</td><td><strong>77.85 ± 1.29</strong></td><td>77.57 ± 0.68</td><td>77.18 ± 1.25</td></tr><tr><td>MASS</td><td>76.06 ± 1.69</td><td>73.95 ± 5.80</td><td>76.94 ± 1.12</td><td>77.16 ± 1.73</td><td><strong>78.50 ± 1.82</strong></td><td>76.86 ± 4.06</td></tr><tr><td>MROS</td><td>83.69 ± 0.39</td><td>82.22 ± 1.27</td><td>83.95 ± 0.53</td><td>83.51 ± 0.84</td><td><strong>84.63 ± 1.02</strong></td><td>83.96 ± 0.87</td></tr><tr><td>PhysioNet</td><td><strong>76.26 ± 1.27</strong></td><td>70.40 ± 0.14</td><td>73.84 ± 0.93</td><td>73.51 ± 3.05</td><td>73.80 ± 1.24</td><td>74.76 ± 0.50</td></tr><tr><td>SHHS</td><td>76.98 ± 0.70</td><td>75.98 ± 0.22</td><td>79.12 ± 0.96</td><td><strong>79.26 ± 1.35</strong></td><td>78.65 ± 1.04</td><td>77.69 ± 0.88</td></tr><tr><td>SOF</td><td>85.49 ± 0.58</td><td>84.23 ± 1.30</td><td>85.50 ± 0.86</td><td>84.14 ± 1.05</td><td><strong>85.54 ± 0.16</strong></td><td>85.33 ± 0.16</td></tr><tr><td>Mean</td><td>79.64 ± 0.41</td><td>77.57 ± 0.73</td><td>79.59 ± 0.25</td><td><strong>80.12 ± 0.57</strong></td><td>79.94 ± 0.80</td><td>79.98 ± 0.47</td></tr><tr><td rowspan="10">Balanced@400</td><td>ABC</td><td>81.00 ± 0.11</td><td>79.50 ± 0.49</td><td>80.56 ± 0.39</td><td><strong>81.12 ± 0.37</strong></td><td>80.80 ± 0.22</td><td>80.90 ± 0.31</td></tr><tr><td>CCSHS</td><td><strong>89.83 ± 0.19</strong></td><td>89.01 ± 0.43</td><td>89.39 ± 0.16</td><td>89.13 ± 0.17</td><td>89.22 ± 0.17</td><td>89.45 ± 0.57</td></tr><tr><td>CFS</td><td>88.30 ± 0.52</td><td>87.39 ± 0.06</td><td>88.45 ± 0.17</td><td><strong>88.52 ± 0.15</strong></td><td>88.32 ± 0.34</td><td>88.23 ± 0.47</td></tr><tr><td>CHAT</td><td>65.77 ± 4.06</td><td>65.25 ± 3.96</td><td>71.35 ± 2.75</td><td>72.16 ± 2.21</td><td>71.93 ± 3.01</td><td><strong>72.99 ± 4.01</strong></td></tr><tr><td>HOMEPAP</td><td>77.06 ± 0.14</td><td>76.62 ± 1.06</td><td>77.50 ± 0.46</td><td>77.30 ± 0.24</td><td>77.78 ± 0.62</td><td><strong>77.98 ± 0.90</strong></td></tr><tr><td>MASS</td><td><strong>77.27 ± 1.42</strong></td><td>74.21 ± 2.05</td><td>75.12 ± 2.08</td><td>76.00 ± 3.00</td><td>74.16 ± 4.21</td><td>74.10 ± 0.47</td></tr><tr><td>MROS</td><td><strong>85.53 ± 0.48</strong></td><td>84.02 ± 0.95</td><td>85.22 ± 0.19</td><td>85.02 ± 0.42</td><td>85.37 ± 0.30</td><td>84.93 ± 0.53</td></tr><tr><td>PhysioNet</td><td>74.98 ± 1.84</td><td>74.29 ± 1.50</td><td>75.07 ± 1.05</td><td>75.29 ± 1.21</td><td><strong>75.41 ± 0.68</strong></td><td>74.57 ± 1.16</td></tr><tr><td>SHHS</td><td>78.95 ± 0.92</td><td>78.04 ± 1.21</td><td>80.30 ± 1.29</td><td><strong>80.32 ± 0.91</strong></td><td>79.22 ± 1.27</td><td>80.06 ± 1.30</td></tr><tr><td>SOF</td><td>86.30 ± 0.40</td><td>85.82 ± 0.22</td><td>86.57 ± 0.60</td><td><strong>86.99 ± 0.33</strong></td><td>86.75 ± 0.54</td><td>86.76 ± 0.23</td></tr><tr><td>Mean</td><td>80.50 ± 0.51</td><td>79.41 ± 0.73</td><td>80.95 ± 0.36</td><td><strong>81.19 ± 0.11</strong></td><td>80.89 ± 0.00</td><td>81.00 ± 0.28</td></tr></tbody></table>

This suggests that, while whitening may help when $F = 1$, it is less effective when $F > 1$. Using a barycenter leads to a more robust and stable target covariance matrix.

<table><caption>Table 6: Impact of the whitening on the performance of PSDNorm on the 10 datasets balanced @ 400.</caption><thead><tr><th rowspan="2"></th><th rowspan="2">BatchNorm</th><th rowspan="2">InstanceNorm<br>Whitening</th><th colspan="2">PSDNorm(F=5)</th><th colspan="2">PSDNorm(F=9)</th><th colspan="2">PSDNorm(F=17)</th></tr><tr><th>Barycenter</th><th>Whitening</th><th>Barycenter</th><th>Whitening</th><th>Barycenter</th><th>Whitening</th></tr></thead><tbody><tr><td>ABC</td><td>78.26 ± 1.33</td><td>78.73 ± 0.42</td><td>78.18 ± 0.68</td><td>77.86 ± 1.33</td><td>78.18 ± 0.91</td><td>78.16 ± 0.93</td><td>77.76 ± 1.00</td><td>78.34 ± 0.33</td></tr><tr><td>CCSHS</td><td>87.42 ± 0.16</td><td>87.62 ± 0.42</td><td>87.58 ± 0.30</td><td>87.80 ± 0.23</td><td>87.35 ± 0.52</td><td>87.53 ± 0.21</td><td>87.62 ± 0.48</td><td><strong>87.90 ± 0.57</strong></td></tr><tr><td>CFS</td><td>84.32 ± 0.57</td><td><strong>84.72 ± 0.33</strong></td><td>84.29 ± 0.36</td><td>84.01 ± 0.60</td><td>84.06 ± 0.10</td><td>84.20 ± 0.23</td><td>84.46 ± 0.06</td><td>84.04 ± 0.91</td></tr><tr><td>CHAT</td><td>66.55 ± 0.88</td><td>64.43 ± 4.41</td><td><strong>70.28 ± 1.70</strong></td><td>69.07 ± 3.73</td><td>68.11 ± 3.94</td><td>69.31 ± 2.50</td><td>69.88 ± 0.46</td><td>69.43 ± 1.80</td></tr><tr><td>HOMEPAP</td><td>75.25 ± 0.50</td><td>76.47 ± 0.63</td><td><strong>76.83 ± 0.61</strong></td><td>76.13 ± 0.93</td><td>76.61 ± 0.74</td><td>76.45 ± 0.29</td><td>76.49 ± 0.45</td><td>76.81 ± 0.22</td></tr><tr><td>MASS</td><td>70.00 ± 1.91</td><td>71.52 ± 1.13</td><td>72.77 ± 1.09</td><td>69.11 ± 1.51</td><td><strong>73.07 ± 1.30</strong></td><td>69.01 ± 0.78</td><td>72.23 ± 2.40</td><td>70.41 ± 1.50</td></tr><tr><td>MROS</td><td>80.37 ± 0.20</td><td>80.28 ± 0.21</td><td>80.26 ± 0.11</td><td>80.50 ± 0.75</td><td>80.32 ± 0.22</td><td>80.52 ± 0.15</td><td><strong>80.70 ± 0.42</strong></td><td>79.96 ± 0.53</td></tr><tr><td>PhysioNet</td><td><strong>75.81 ± 0.13</strong></td><td>74.68 ± 0.55</td><td>74.82 ± 2.11</td><td>74.58 ± 1.57</td><td>73.77 ± 1.73</td><td>75.08 ± 2.18</td><td>75.09 ± 0.97</td><td>75.61 ± 0.71</td></tr><tr><td>SHHS</td><td>76.44 ± 0.92</td><td>78.68 ± 0.37</td><td><strong>78.88 ± 0.68</strong></td><td>78.77 ± 0.67</td><td>77.28 ± 0.91</td><td>78.85 ± 0.16</td><td>78.41 ± 0.49</td><td>77.77 ± 1.35</td></tr><tr><td>SOF</td><td>81.08 ± 1.14</td><td>80.68 ± 1.38</td><td>79.49 ± 0.41</td><td>80.10 ± 0.62</td><td><strong>81.44 ± 0.97</strong></td><td>80.53 ± 0.05</td><td>81.07 ± 0.66</td><td>79.63 ± 0.89</td></tr><tr><td>Mean</td><td>77.55 ± 0.34</td><td>77.78 ± 0.46</td><td>78.34 ± 0.42</td><td>77.79 ± 0.30</td><td>78.02 ± 0.67</td><td>77.96 ± 0.12</td><td><strong>78.37 ± 0.38</strong></td><td>77.99 ± 0.35</td></tr></tbody></table>

## A.7 Generalization of PSDNorm in CNNTransformer

The CNNTransformer architecture is a hybrid model that combines convolutional and transformer layers for time series classification.

The main paper presents a critical difference diagram for the CNNTransformer evaluated on datasets balanced@400. It shows that PSDNorm with $F = 5$ is the best-performing normalization layer.

In Table 7, we report the results of different normalization layers used in the CNNTransformer architecture on datasets balanced@400.

First, we observe that CNNTransformer performs slightly below U-Sleep. Second, BatchNorm and InstanceNorm are the best performers on one and two datasets respectively, while PSDNorm achieves the best performance on 7 out of 10 datasets.

PSDNorm with $F = 5$ outperforms BatchNorm by a margin of 0.9 and InstanceNorm by 0.54 in average score.
<table><caption>Table 7: Different normalization layers used in the CNNTransformer architecture for datasets balanced@400.</caption><thead><tr><th></th><th>BatchNorm</th><th>InstanceNorm</th><th>PSDNorm(F=5)</th><th>PSDNorm(F=9)</th><th>PSDNorm(F=17)</th></tr></thead><tbody><tr><td>ABC</td><td><strong>76.99</strong> ± <strong>0.53</strong></td><td>75.40 ± 0.36</td><td>76.31 ± 0.46</td><td>75.55 ± 0.82</td><td>75.12 ± 0.31</td></tr><tr><td>CCSHS</td><td>86.75 ± 0.48</td><td>87.00 ± 0.34</td><td>86.92 ± 0.32</td><td>86.92 ± 0.32</td><td><strong>87.27</strong> ± <strong>0.40</strong></td></tr><tr><td>CFS</td><td>83.32 ± 0.35</td><td><strong>83.77</strong> ± <strong>0.34</strong></td><td>83.71 ± 0.29</td><td>83.47 ± 0.61</td><td>83.68 ± 0.34</td></tr><tr><td>CHAT</td><td>66.44 ± 0.49</td><td>66.40 ± 2.55</td><td><strong>70.04</strong> ± <strong>0.37</strong></td><td>69.70 ± 2.49</td><td>67.60 ± 3.84</td></tr><tr><td>HOMEPAP</td><td>74.81 ± 1.36</td><td><strong>75.92</strong> ± <strong>0.44</strong></td><td>75.26 ± 0.55</td><td>75.55 ± 0.81</td><td>75.14 ± 1.14</td></tr><tr><td>MASS</td><td>71.51 ± 0.47</td><td>71.70 ± 1.17</td><td>72.55 ± 0.81</td><td>72.86 ± 0.11</td><td><strong>73.26</strong> ± <strong>0.26</strong></td></tr><tr><td>MROS</td><td>79.77 ± 0.31</td><td>79.74 ± 0.55</td><td>79.77 ± 0.30</td><td>79.52 ± 0.33</td><td><strong>79.83</strong> ± <strong>0.52</strong></td></tr><tr><td>PhysioNet</td><td>72.54 ± 0.34</td><td>74.36 ± 0.84</td><td>74.95 ± 0.41</td><td><strong>75.32</strong> ± <strong>1.08</strong></td><td>75.19 ± 1.25</td></tr><tr><td>SHHS</td><td>75.34 ± 0.34</td><td>76.55 ± 0.92</td><td><strong>77.26</strong> ± <strong>0.57</strong></td><td>76.44 ± 0.99</td><td>76.64 ± 0.80</td></tr><tr><td>SOF</td><td>80.63 ± 0.60</td><td>80.78 ± 0.54</td><td>80.31 ± 0.90</td><td><strong>80.84</strong> ± <strong>0.30</strong></td><td>80.65 ± 1.00</td></tr><tr><td>Mean</td><td>76.81 ± 0.23</td><td>77.16 ± 0.24</td><td><strong>77.71</strong> ± <strong>0.22</strong></td><td>77.62 ± 0.37</td><td>77.44 ± 0.53</td></tr></tbody></table>

These results highlight that PSDNorm is a plug-and-play normalization layer that can be seamlessly integrated into various architectures to reduce feature space variability.

## A.8 Critical Difference Diagram for U-Sleep on all subjects

The main paper presents the critical difference diagram for U-Sleep on the dataset balanced@400. Figure 6 extends this analysis to all subjects across datasets. The conclusion remains consistent: PSDNorm with $F = 5$ is the best-performing normalization layer, while BatchNorm performs the worst. Interestingly, PSDNorm with $F = 17$ ranks second to last, suggesting that overly strong adaptation can hurt performance when the dataset is large.

<figure><img src="image_8.png" /><figcaption>Figure 6: Critical difference diagram for U-Sleep on all subjects.</figcaption></figure>