PSDNorm: Temporal Normalization for Deep Learning in Sleep Staging
Authors: Theo Gnassounou
Antoine Collas
Rémi Flamary
Alexandre Gramfort
Venue: ICLR 2026 Poster (2026)
Keywords: Normalization Layer
Sleep Staging
Optimal Transport
Primary area: applications to neuroscience & cognitive science

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.
    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.

Published as a conference paper at ICLR 2026
PSDNORM: TEMPORAL NORMALIZATION FOR DEEP LEARNING IN SLEEP STAGING
Théo Gnassounou Université Paris-Saclay, Inria, CEA, 91120 Palaiseau, France
Antoine Collas Université Paris-Saclay, Inria, CEA, 91120 Palaiseau, France antoines wow.sneooollaeiniiaa..
theo wwr..nnnsssuuuunnnnnn. inria. fr
Rémi Flamary
Alexandre Gramfort Université Paris-Saclay Inria, CEA, 11020 Palaiseau, France
École Polytechnique. IP Paris, CMAP, UMR 7641, 91120 Palaiseau, France
remi. Llaaarypollyeechrique...
a alexandre. e granfort@innia.
ABSTRACT
Distribution shift poses 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. 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 unseen left-out datasets while being more robust to data scarcity.
INTRODUCTION
Data Shift in Physiological Signals. Machine learning techniques have achieved remarkable success 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 (Moreno- Torres al., (012). For example, in biological data, such as electreencephalograpyy (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 sleep staging, clinical task that consists in classifying periods of sleep in different stages based on EEG signals (Stevens & Clark 2004). Depending on the dataset, the repartition, health conditions, and recording
cohort can be composed of different age groups, conditions (O'Reilly et al., 2014 Quan et al. 1998 Marcus et al. 2013). Such variability brings shift in the distribution making challenging for the model generalize 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 (loffe & Szegedy 2015), LayerNorm (Ba et al. 2016), and InstanceNorm (Ulyanoy 2016). which respectively compute statisties across the batch, normalize across all features within each sample, and normalize each channel independently within sample. Some normalization methods target specific tasks, such as EEG covariance matrices (Kobler al.) 2022) or time-series forecasting (Kim al. 2021). but they do not fully address spectral distribution shifts reflected in the temporal auto-corre citto-ccorllationss of signals. Other papers have proposed to adapt layer statistics to new domains (Li et al., 2016) [Chang et al.) 2019). In sleep staging, a simple normalization often applied preprocessing, e.g., standardizing
A. Gramfort joined Meta and he reached a agramfort @neta.com

---

Published as a conference paper at ICLR 2026
BateiaNom
IrstanceNorm
PSDNNomm
Shape reduced
Statiatios
Batch estimation
Lesmable
Porametors
N
Normaliaing with Tamporal Convolution
(0,8-=
(7.8%)
p
11
N
&
Vigiina
bv
P F
8
Figure Description of normalization layers. The input shape is (N, c, with batch size N, channels c, and signal length t. BatchNorm estimates the mean and variance 2² over batch and time, and learns parameters (7,B) normalize the input. PSDNorm estimates PSDs P over time and accounts for local temporal correlations. It computes the barycenter PSD P. updates it via running Riemannian barycenter (6), and applies the filter H to normalize the input. The hyperparameter f controls the extent of temporal correlation considered, thereby adjusting the strength of the normalization.
signals over entire nights (Apicella et al. 2023) or short temporal windows (Chambon et al. 2018).
Recent studies (Gnassounou et al., 2023 2024 highlight the importance of considering temporal correlation and spectral content in normal ization, proposing Temporal Monge Alignment (TMA), which aligns Power Spectral Density (PSD) to 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 it is done with BatchNorm, LayerNorm InstanceNorm.
Deep Learning for Sleep Staging. Sleep staging has been addressed by various neural network architectures, which process raw signals (Chambon et al. 2018 Perslev et al. 202 Guillot & Thorey 2021), spectrograms (Phan et al. 2023 2019), or both (Phan al. 2022a). More recent approaches
involve transformer-bant models that handle multimodal (Wang et al.). spectrogram (Phan et al.
2022b), heterogeneous inputs (Guo al. 2024), offering improved modeling of temporal depen- dencies. However, most existing models are trained on relatively small cohorts, typically consisting of only few hundred subjects, which limits their ability to generalize to diverse clinical settings. Notable exceptions include U-Sleep (Perslev et al. 2021), which was trained large-scale dataset and incorporates BatchNorm layers to mitigate data variability, and foundational models (Thapa al. Fox et al.) Deng et al.) 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 developing smaller, efficient models that balance good generalization with of training and deployment clinical practice.
Contributions. In this work, introduce the PSDNorm deep learning layer, novel approach to address distribution shifts in machine learning for signals. PSDNorm leverages Monge Map- ping incorporate temporal context and normalize feature maps effectively. This layer enhances model robustness to new subjects at inference time. Unlike standard normalization layers such as LayerNorm or InstanceNorm, PSDNorm leverages the sequential nature intermediate feature maps, as illustrated in Figure [ We evaluate PSDNorm through extensive experiments on 10 sleep datasets. This evaluation covers 10M of samples across 10K subjects, using leave-one-dataett- (LODO) protocol with 3 different random seeds. To the best of our knowledge, such large-scale and systematic evaluation has never been conducted before. PSDNorm achieves state-of-the-ee 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 introduces PSDNorm, followed by numerical results in Section
2

---

Published as a conference paper at ICLR 2026
Notations. Vectors are denoted by small cap boldface letters (e.g., x), matrices by large cap boldface letters (e.g., X). The element-wise product, power of n and division are denoted o .0m and @. respectively. [1, K] denotes {1, {I....,KT K}. The absolute value is |.]. The discrete circular convolution along the temporal axis operates row-wise as, * Rext x IRexf - Rext for E 2 f. vec : RRx - c concatenates rows of time series into vector. 27 = [x] refers to the 1thh element of x, and Xt,m = [X]m. denotes the element of X at the (th row and m column. X and X are the conjugate and the transpose of X, respectively. diag puts the elements of a vector on the diagonal of matrix. @ is the Kronecker product. 1, is the vector of ones of size
2 RELATED WORKS
In this section, we first review classical architectures for sleep staging and fundamental concepts of normalization layers. Then, we recall the Temporal Monge Alignment (TMA) method (Gnassoonnou
al. 2023) that aligns the PSI) of signals using optimal transport.
Deep Learning for Sleep Staging, Numerous neural network architectures have been proposed for sleep staging, processing data in different formats as introduced in Section [ Different types
of architectures have been explored, such as convolutional neural networks (CNNs) (Chambon
al. 2018), recurrent neural networks (RNNs) (Supratak et al.) 2017: Phan et al., 2019). and more recently transformers (Phan et al., 20226 Wang et al.) Guo et al. (024), which have shown promise modeling temporal dependencies in sleep data. While many models are typically evaluated on limited number of datasets, the work by (Perslev et al.] 2021 introduced U-Sleep, model trained on a large-scale dataset of sleep recordings. Their architecture, based on U-Time (Perslev al. 2019), incorporates BatchNorm layers to mitigate data variability, and they employ domain generalization approach: training a single model on a sufficiently diverse set of domains to ensure it generalizes to unseen datasets without additional adaptation. This architecture is composed of encoder-decoder blocks with skip connections, allowing the model to capture both local and global features of the sleep signals effectively. Each encoder and decoder block consists of convolutional layers followed by BarchNorm and non-linear activation functions, enabling the model to learn robust representations of the input data. A more detailed description of U-Time is provided in Section A.3 Normalization Layers. Normalization layers enhance training and robustness in deep neural networks. The most common are BatchNorm (Toffe & Szegedy 2015), InstanceNorm (Ulyanov) 2016), and LayerNorm (Ba et al., 2016). BatchNorm normalizes feature maps using batch and time statisties, ensuring zero mean and unit variance. The output adjusted with learnable parameters. InstanceNorm normalizes each channel per sample using its own statistics, independent of the batch (see Fig. D. Popular in time-series forecasting, it is used in RevIN (Kim et al.) 2021], 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 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) (Cnassounouu et al.) 2023] 2024) was introduced temporal correlations by leveraging the Power Spectral density (PSD) of multivariate signals using Monge Optimal Transport mapping.
pre-processing step to align
Gaussian Periodic Signals. Consider multivariate signal X Xi,-.., x.]T Rext 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 's covariance matrix is block diagonal, with each block having a circulant structure. A fundamental property of symmetric positive definite eirculant matrices is their diagonalization (Gray 2006) with real and positive eigenvalues in the Fourier basis F2 E CFX of elements
exp (-2im (L 1)(1' -1)
Fely A
where 1,1' [1 Hence, we have vec(X) (X) NO, Σ) with Σ block-diagonal,
Σ = (Ie & F2) diag (v))(P)) (I. @ Fj) E Retxel
(2)
where P E x contains positive entries corresponding to the Power Spectral Density of each sensor. In practice, since we only have access to single realization of the signal, the PSD is estimated with
3

---

Published as a conference paper at ICLR 2026
only f < l frequencies, i.e.. P R This amounts to considering the local correlation of the signal and neglecting the long-range correlations.
Power Spectral Density Estimation. The Welch estimator [Welch] 1967) computes the PSD of signal by averaging the squared Fourier transform of overlapping segments of the signal. Hence, the realization of the signal E _ decimated into overlapping segments Rex/ to estimate the PSD. The Welch estimator is defined
x C
P
1 L
((1o
02
E ERxx
(3)
where w R/ the window function such that |w/l-
M-Monge Mapping. Let (0, ) and (0, ())) source and target centered Gaussian distri-
butions respectively with covariance matrices following the structure 2 and PSDs denoted by p(a)
and p()) E Rex Given a signal E RexE such that vec(X)
NOO, (*)), the f-Monge mapping
as defined by (Gnassounou et al. 2023 2024 is
pl.)) 35 Fý E ROx
my (x, x= H gext
where
H
P 2
(4)
In this case, f controls the alignment between the source and target distributions. Indeed, if then the f-Monge mapping is the classical Monge mapping between Gaussian distributions and the source signal has its covariance matrix equal to (+) 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 (Agueh & Carlier] 20TT) admits an elegant closed-form solution. Consider K centered Gaussian distributions admitting the decomposition (2) of PSDs p(1) ,P()) Their barycenter is also a centered Gaussian distribution (0, Σ) admitting the decomposition with PSD
K
82
O 1
P e
R
(5)
K
Temporal Monge Alignement. TMA is a pre-processing method that aligns the PSD of multi- variate signals using the f-Monge mapping. Given source signal X, and a set of target signals
X, = [x!") (! ("), 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, preprocessing method, is inherently limited to handling PSD shifts the raw signals and cannot address more complex distributional changes the data. This limitation highlights the need for layer that effectively capture and adapt these complex variations during learning and inside deep learning models.
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 barycenter PSD within the architecture of a deep learning model.
PSDNorm is novel normalization layer that can be used as drop-in replacement for layers like NatchNorm or InstanceNorm. Instead of simple standardization, it aligns the Power Spectral Density (PSD) of feature maps to running barycenter PSD. This approach, optimized for modern hardware, enhances model robustness to new subjects at inference time without retraining. We define the normalized feature map as á PSDNorm G). The following sections introduce the core components of PSDNorm and its implementation.
4

---

Published as a conference paper at ICLR 2026
3.1 CORE COMPONENTS OF THE LAYER
In the following, we formally define PSDNorm and present each of its three main compo- nents: 1) PSD estimation, 2) running Riemannian barycenter estimation, and 3) f-Monge map- ping computation. Given batch B {G ,.... ,...,,(N)) of N pre-normalization feature maps, PSDNorm outputs a normalized batch B (G ..... G(M) with normalized PSD. Those three steps are detailed in the following and illustrated in the right part of Figure
PSD Estimation. First, the estimation of the PSD of each feature map is performed using the Welch method. The per-channel mean () is computed for each feature map G ju A Ei [G6] R-
GR
ap
P
P p P
((2)
p(O)
Then, the PSD of the centered feature map G 1 denoted p()), is estimated as described in Equation (3). This centering step is required as fea- ture maps typically non-centered due to activa- tion functions and convolution biases but they are assumed to have a stationary mean. The Welch es- timation involves segmenting the centered feature map into overlapping windows, computing the Fourier transform of each window and then averaging them.
Figure 2: Description of the running Rie- manian barycenter. The barycenter of the batch P5 is estimated from the PSD of each batch sample.
Geodesic and Running Riemanian Barycenter.
The PSDNorm aligns the PSD of each feature map barycenter PSD. This barycenter is computed during training by interpolating between the batch Wasserstein barycenter and the current running Riemanian barycenter using the geodesic associated with the Bures metric (Bhatia al., 2019). The batch barycenter is first computed from the current batch PSDs /P(2),... ,..., ((N) using Equation 5).
To ensure gradual adaptation, the running barycenter is updated via exponential geodesic average with a E [0,1]:
62
P (a-ap) 27
+aP (
E E
(6)
proof of the geodesic provided in Appendix A.I
PSD Adaptation with/-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 E [1, [1, it defined as
( my (
(c in0))
gext
(7)
where () is the Monge mapping filter computed as
V7 P O (O)) o
Fý E Rex
(8)
where p() is the estimated PSD of G6) - 1T.
3.2 IMPLEMENTATION DETAILS
Overall Algorithm The forward computation of the proposed layer is outlined in Algorithm At train time, the PSDNorm performs three main operations: 1) PSD estimation. 2) running Riemannian barycenter update, and 3) Monge mapping application. At inference, 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, stop gradient operation is applied to the running barycenter to prevent the backpropagation of the gradient computation through the barycenter. PSDNNrm has unique additional hyperparameter / which is
5

---

Published as a conference paper at ICLR 2026
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 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 inference times, is O(Nel / log(f)), where N is the batch size, e the number of channels, l the signal length, and the filter size.
3.3 DISCUSSION AND CONNECTIONS TO RELATED METHIODS
PSDNorm generalization of InstanceNorm. InstanceNorm applies a per-channel z-score over time, subtracting the mean and dividing by the standard devi- eoon--quivllent 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-colorngg in the frequency domain. InstanceNorm is recovered as a special case of PSDNorm by setting the filter size to f = 1 and using the uniform
Algorithm Forward pass of PSDNorm
1: Input: Batch = (G G(V), rumning barycenter
P. filter-size f. momentum a, training flag
2: Output: Normalized batch {
{G (
3: for to N do
4: ñ (f) Mean estimation
5: P +- PSD est. from G -p-") 1 with eq.
6: end for
7: if training then
8: Ps + Batch bary. from (PO)}, with eq. 5
9:
< Running bary. up. from P, Ps with eq. (6)
10: end if
11: for j = to N do
PSD barycenter as P 1. as the re-coloring transform instead of the barycentric PSD.
H <- Filter estimation from p()) with eq. (8)
13: GO) < f-Monge mapping with eq. 7
14: end for
Similarity with Test-time Do- main Adaptation. PSDNorm inspired by Temporal Monge Alignment (TMA) (Gassounouu et al. 2023). pre-processing technique that can be used for test- time adaptation. Test-time Domain Adaptation methods adjust pre-trained model new target domain during inference, without requiring access to the original training data (Wang et al.) 2021] Yang et al. 2021). While PSDNorm must be integrated into the model during training and is not post-hoc adaptation method, it provides similar benefit at inference time. Designing new modern architectures that incorporate PSDNorm can enhance robustness to domain shifts without the need for retraining access to source data.
Discussion of Gaussian and Stationarity Assumptions. PSDNorm relies on the Gaussian ap- proximation of OT for compensation variability but does not assume that the signals are Gaussian. This allows for efficient alignment of second-order statistics (covariance structure), but also allows preserving higher-order discriminative information. This approach is computationally tractable and targets the most prominent sources of domain shift without over-distorting the signal, a strategy also used in successful methods like Deep CORAL (Sun & Saenko 2016). Like BatchNorm, PSDNorm assumes shifts are captured by low-order statistics, but it provides richer alignment by incorporating temporal context.
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 is available at httpth:: // supplementary material. All numerical experiments were conducted using a total of 1500 GPU hours on NVIDIA H100 GPUs.
com/tgnasse P The anonymized code is available in the

---

Published as a conference paper at ICLR 2026
4.1 EXPERIMENTAL SETUP
Table 1: Characteristics of the datasets.
Datasets To evaluate the effect of normaliza- tion layers, we use ten datasets of sleep stag- ing described in Table 1 ABC (Jessie P. et al.)
Dataset
Subj.
Rec.
Age + std
Sex (F/M)
ABC
44
117
48.8 + 9.8
43%6/57%
CCSHS
515
515
17.7 = 0.4
50%
2018), CCSHS (Rosen et al. 2003), CFS (Red- line et al. 1 HPAP (Rosen et al. 2012).
CFS
681
681
41.7 + 20.0
55%/455
HPAP
166
166
46.5 + 11.9
43%b/57%
MROS (B Blaacweell et al.] (011T) SHHS (Qun
MROS
2101
2698
76.4 5.5
0%/100%
PHYS
70
132
58.8 +: 22.0
33%6/67%
et al. 1998), CHAT (Marcus et al. 2013). and
SHHS
5730
8271
63.1 +: 11.2
52%//48%
SOF Spira al.) 2008) are publicly available sleep datasets with restricted access from Na-
MASS
61
61
42.5 : 18.9
55%/45%
CHAT
1230
1635
6.6 + 1.4
52%/44%
tional Sleep Research Resource (NSRR) (Zhang ct al. 2018). PHYS GGoldbergrr et al. 2000
SOF
434
434
82.8 + 3.1
100%/0%
Total
11032
14710
and MASS (O'Reilly al. 2014) 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- processsss steps and splits of the datasets.
Data Pre-processing, We follow standard pre- rre-processing pipeline used in the field (Chambon ctt al., 2017 Stephansen et al. 2018). The datasets vary in the number and type of available EEG and clectrooculogram (EOG) channels. To ensure consistency, 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 (Appelhoff et al. 2019 and MNE-Python (Gramfort et al. 2013).
use two bipolar EEG channels,
Seave-OneDattaeetOOn (LODO) Setup and Balancing. We evaluate model performance using a leave-one-dataset (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 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
Architecture and Training. Sleep staging has inspired a variety of neural architectures, from early CNN-based models (Chambon et al., 2017 Stephansen et al., (2018: Phan et al., 2022a) to recent attention- based approaches (Phan et al. 2022b 2023] Wang et al.). We evaluate two architectures: U-Sleep (Perslev et al.) 2019 2021), otate-offie of-te-o art temporal CNN model designed for robustness and large-scale training, and a newly introduced architecture, CNNTransformer. CNNTransformer combines a lightweight convolutional encoder with Transformer applied 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 (Section A. Its design draws inspiration from recent transformer-bat based models for time series (Yang et al. 2023 with an emphasis on simplicity and practicality.
We use the Adam optimizer (Kingma 2014) 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 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. The filter size f of PSDNorm is set to 5. A sensitivity analysis of f is provided in Section A.6 in the appendix, and show that the performance is stable across a range of values from 5 to 11.
Evaluation. At inference, the model similarly processes sequences of 17'30s with a stride of 10' 30s. Performance is evaluated using the balanced accuracy score (BACC), computed the central 10' 30s of each prediction window. Each experiment is repeated three times with different random seeds, and we report the mean and standard deviation of BACC.
7

---

Published as a conference paper at ICLR 2026
Table 2:
Balanced Accuracy (BACC) scores on the left-out datasets with USleep. 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 highlighted in bold, and standard deviations reflect training variability across 3 random seeds. The mean BACC reports the average over all the subjects.
Dataset
BatchNorm
LayerNorm
InstanceNorm
TMA
PSDNorm
ABC
78.49 =0.00
77. 77.9440.31
78.83 78.83+0.55
78.33
78.56
CCSHS
88. 88.79=0.21
87.51
88. 88.75+0.04
88.61=0.10
88.5640.36
CFS
80.97=0.37
81.29+0.67
85. 85.73+00.2
84.85+0.13
85.42+0.09
All subjects
CHAT
64.72_3.94
64.36
68.86
69.76
70.57 ::
HOMEPAP
76.39_a.28
75.23
76.70.
76.77
76.72+..27
MASS
73.71
71.3943.00
72. 12.12
73.90
72.51 4
MROS
81.30=0.25
80.4440.29
81.49 81.49+0.18
80.91
=0.57
PhysioNet
76.13_0.57
75. 75.12+0.22
76.15+0.52
76.48
75.9641.08
SHHS
77.97=1.16
75.98+0.48
79. 79.05+0..89
78.21
79. 14:1.01
SOF
81.33=0.54
81.8240.79
81.98+0.22
81.84
82.50.0.34
Mcan(Dataset)
78.38-00477
77.41
78.97
78.98
79. 79.15.0.14
Mean(Subject)
78.14=1.01
76.78+0.118
79.26 79.26+:048
78.77
79.51
ABC
78.2
75.29
78.73+0.42
78.04
78.1840.65
CCSHS
87.42_0..16
85.20.1..4
87. 87.62.10.42
87.57
87 87.5849.30
CFS
84.32_0.57
81.66+1.36
84.72
84.58
84.29 44.3940.36
Balanced (annced@e440 (e:400
CHAT
66.55_0.88
61.1941.16
64.43
68.73-2248
70.28 28
HOMEPAP
75.25
74.86.40.25
76.47
76.10
76. 83. 0.61
MASS
70.00=1.911
68.5643.35
71.52+1.15
71.63
72.77 1.09
MROS
=0.27
7..05
80.28
80.09
$0.26.10.11
PhysioNct
=5.81
71.82:2.12
74.68
75.31
74.82+2.11
SHHS
76.44_0.92
75. 75.1240.39
78 78.6840.37
77.00
78. 18.88
SOF
81.08
78. 00.5.0
80.6841.38
81.25
79. 79.4944.41
Mean(Dataset)
77.5
75.05 4..5.....
77.78
78.34 78.03
Mean(Subject)
77.22
75.04
78.17
77.74
78.85
Normalization Strategies. We compare the proposed PSDNorm with three normalization strategies: BatchNorm, LayerNorm, and InstanceNorm. Note that InstanceNorm corresponds to special case of PSDNorm with f and 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 s default implementations of LayerNorm, InstanceNorm (Paszke et al. 2211). or PSDNorm. To preserve the receptive field, the filter size / of PSDNorm is used in the first layer and progressively halved in the following ones. We the momentum a to 10-2.
4.2 NUMERICAL. RESULTS
This section presents results from large-scale sleep stage classification experiments. The analysis be- gins with comparison PSDNorm against standard normalization layers- BatchNorm, LayerNorm, and InstanceNorm- the full datasets. Then, the data efficiency of each method is evaluated under limited training data regimes. Finally, robustness distribution shift is assessed via subject-wise performance across multiple neural network architectures.
Performance Comparison on Full Datasets. Table (top) reports the LODO BACC U-Sleep across all datasets, averaged over three random seeds. PSDNorm consistently outperforms all baseline normalization layers BatchNorm, LayerNorm, InstanceNorm, and TMA aahieviig the highest mean BACC of 79.51% over subjects, which exceeds BatchNorm (78.38%). InstanceNorm (78.97%), LayerNorm (77.41% 1%) and TMA (78. (78.77%) On the challenging CHAT dataset, where all methods struggle, PSDNorm outperforms all other normalizations by more than I percentage points, highlighting its robustness under strong distribution shifts. Although InstanceNorm is a strong base- line-outie outperforming BatchNorm and LayerNorm by at least standard deviation average-it it is consistently surpassed by PSDNorm in average performance. contrast, LayerNorm underperforms across the board, achieving the lowest average BACC and never ranking first, confirming its limited
8

---

Published as a conference paper at ICLR 2026
suitability for this task. PSDNorm also improves score by almost 1 1% over TMA, showing that using Monge Alignment inside the network allows for better adaptation.
Performance with 4x Less Data. The PSDNorm layer improves model per- formance when trained on the full dataset (~10000 subjects), but such large-scale data availability is always the case. In many real-world scenarios-suce as rare disease studies, pediatric populations, or data col- lected in constrained clinical settings caatings-labeled recordings are scarce, expensive to annotate, or restricted due to privacy concerns. Evaluat- ing model robustness under these constraints therefore essential. To this end, we train all models using the balanced @440 setup, which reduces the training data by a factor of 4 com- pared to the full-data setting. this lower- data regime, PSDNorm continues outper- form all baseline normalization strategies and achieves higher average BACC. The performance improvement PSDNorm over the best baseline is more pronounced in this setting: the BACC gain reaches +0.67% compared to +0.25% in the full-data setting. The gains exceed one standard devia- tion. To assess statistical significance, we conducted a critical difference (CD) test (Demšar 2006).
Efficiency:
USleep
2.0
2.2
2.4
2.6
2.8
PSDNorm
BatchNorm
InstanceNorm
TMA
CNNTrnstoomer
2.4
PSDNorm
TMA
InstanceNorm
BatchNorm
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.
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 Appendix Figure Sb. The following experiments focus the balanced @400 setup.
Target: MASS (61 subj.)
Target: CHAT (1230 subj.)
Performance on the most challenging subjects. Performance variability across subjects is key challenge biomedical applications where en-
BACC BatchNorm
0.9
25%
9%
0.7
suring consistently high performance- even for the most challenging subjects- is critical. To
0.5
75%
91%
highlight the robustness of PSDNorm, Figure 4 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 the dots are below the diagonal, indicating that PS- DNorm improves performance for of subjects against BatchNorm and 99% of subjects against InstanceNorm, with the largest gains observed for the hardest subjects, reinforcing its ability han- dle 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.
BACC InstanceNorm
0.9
331%
1%
0.7
0.5
69%
99%
0.5
0.7
0.9
0.5
0.7
0.9
BACC PSDNorm
BACC PSDNorm
Figure Subject-wise BACC comparison on MASS and CHAT (balanced @ @400). Blue dot means improvement with PSDNorm.
Robustness Across Architectures. PSDNorm is plug-and-pla 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 reports the average rank of each normalization method across datasets and subjects for both architectures using datasets balanced @ 400. In both architectures, PSDNorm achieves the best overall ranking and demonstrates statistically significant improvements over both BatchNorm, InstanceNorm, and TMA. The results confirm that PSDNorm generalizes well beyond single architecture and can provide consistent improvements in diverse modeling setups which is not the case of TMA that is ranked the worst with CNNTTrasfformer InstanceNorm performs competitively in some cases but is never

---

Published as a conference paper at ICLR 2026
significantly better than PSDNorm. Detailed numerical scores for CNNTransformer reported in the supplementary material (Table (
It is important to highlight that PSDNorm brings improvements without too much additional compu- tational cost. In appendix Section we provide a detailed comparison of the computational time of PSDNorm with other normalization layers. The results show that PSDNNom is only slightly slower than BatchNorm and InstanceNorm, with negligible increase in training time (less than 10%) and no significant impact on inference speed.
4.3 ILLUSTRATION OF PSD NORMALIZATION
Figure shows how different normalization layers affect the PSD of signals at several stages of the network. The input signals display limited variability, which explains why applying TMA as pre-processing step provides only marginal benefit. In the first row, corresponding to BatchNorm, the PSD variability increases with depth, a behavior that is undesirable for generalization. TMA exhibits normalization is applied within the network to counteract this accumulation of variance. In contrast, both InstanceNorm and PSDNorm reduce PSD variability across samples. However, InstanceNorm does not fully align the PSDs, and noticeable differences remain between samples. PSDNorm, on the other hand, achieves strong alignment PSDs across samples, indicating its ability to normalize the underlying temporal correlations. This alignment essential for improving robustness and generalization, particularly in settings involving distribution shifts.
similar pattern,
BatchNorm
InstanceNorm
TMA
PSDNorm
Freq. (Hz) Input
Freq, (Hzz)
Freq, (H2)
Freq, (Hz)
After first norm
After second
After third
Figure 5: Illustration of PSD normalization with different normalization layers. The figure shows the PSD of different segment of 17min from subject batch as input, and after 3 encoders using different normalization layers.
5 CONCLUSION, LIMITATIONS, AND FUTURE WORK
This paper introduced PSDNorm, a normalization layer that aligns the power spectral density (PSD) of each signal geodesic barycenter. By leveraging temporal correlations, PSDNorm offers 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 settings-outpefeee outperforming BatchNorm, LayerNorm, and InstanceNorm across architectures.
While the results are promising, some limitations remain. PSDNorm introduces filter size hyperpa- rameter (f) that controls normalization strength; although we provide default values that perform well across datasets, selecting 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 other signals such audio and other biomedical applications.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work was supported by the grants ANN-22-PESN--00) to AC under the France 2030 program, and ANR-20-IADJ-0002 to AG while Inria, and ANR-23-ERCC-000 and ANR-25-PELA-OU0 to RE, all from Agence nationale de la recherche (ANR). This work is supported by Hi! PARIS and ANR/France 2030 program (ANR-22 (ANR-23-IACE 23-IACL-0005), This project has also received funding from the European Union's 's Horizon Europe research and innovation programme under grant agreement 101 120237 (ELIAS).
This project received funding from the Fondation de I' l'École polytechnique
This project was provided with computer and storage resources by GENCI at IDRIS thanks to the grant 0025-AD01101605 and 2025-AD01 1016067 on the supercomputer Jean Zay's the V100 & H100 partitions.
This work 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 the Inria compute infrastructures.
Numerical computation was enabled by the scientific Python ecosystem: Matplotlib Hunter (2007). Scikit-learn [Pedregosa et al. (2011), Numpy Harris al. (2020), Scipy VVrttnnnn al. (2020)
PyTorch Paszke et (201)) and MNE Gramfort et al. (2013)
REFERENCES
Martial Agueh and Guillaume Carlier. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2):904-925 924, 2011.
Andrea Apicella, Francesco Isgrò, Andrea Pollastro, and Roberto Prevete. the effects of data nor- malization for domain adaptation on EEG data, Engineering Applications of Artificial Intelligence, 123:106205, 2023.
Stefan Appelhoff, Matthew Sanderson, Teon L. Brooks, Marijn Vliet, Romain Quentin, Chris Holdgraf, Maximilien Chaumon. Ezequiel Mikulan, Kambiz Tavabi, Richard Höchenberger, Dominik Welke, Clemens Brunner, Alexander P. Rockhill, Eric Larson, Alexandre Gramfort, and Mainak Jas. MNE-BIDS: Organizing electrophyssological data into the BIDS format and facilitating their analysis. Journal of Open Source Software, 4(44):1896, 2019. doi: 10.21105/joss. 01896.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016. URL http: wrw.obrg///oo..ooooo. abtp/bbss//11 .6650
Rajendra Bhatia, Tanvi Jain, and Yongdo Lim. On the wasserstenn distance between positive definite matrices. Expositiones Mathematicae, 1652)1165-191, 2019.
Terri Blackwell, Kristine Yaffe, Sonia Ancoli-Israel, Susan Redline, Kristine E. Ensrud, Mar- cia L. Stefanick, Alison Laffan, Katie L. Stone, and Osteoporotic Fractures in Men Study Group. Associations between sleep architecture and disopddid breathing and cognition in older community-dwelling men: the Osteoporotic Fractures in Men Sleep Study. Journal of the American Geriatrics Society, 59(12);2217-22) December 2011. ISSN 1532-5415. doi: 1533-5555532-
Stanislas Chambon, Mathieu Galtier, Pierrick Arnal, Gilles Wainrib, and Alexandre Gramfort. A deep learning architecture for temporal sleep stage classification using multivariate and multimodal time series. 2017.
Stanislas Chambon, Mathieu N Galtier, Pierrick J Arnal, Gilles Wainrib, and Alexandre 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(4):758-769, 2018.
Woong-Gr Gi Chang, Tackgeun You, Seonguk Seo, Suha Kwak, and Bohyung Han. Domain-specific batch normalization for unsupervised domain adaptation, 2019. URL https: stttps://arxiie //tpxxx. abs 1906. 03950)
11

---

Published as a conference paper at ICLR 2026
Janez Demšar. Statistical comparisons of classifiers over multiple data sets. Journal of Machine Learning Research, 7(1):1-30, 2006. BRLhttt::/jjmr http://trr.oaoaaaeee.oeooeeeeeae htt..///paapeee////aaaaaaaa
html
Guifeng Deng, Mengfan Niu, Yuxi Luo, Shuying Rao, Junyi Xie, Zhenghe Yu, Wenjuan Liu, Sha Zhao, Gang Pan. Xiaojing Li, Wei Deng, Wanjun Guo, Tao Li, and Haiteng Jiang. A unified flexible large polysomnography model for sleep staging and mental disorder diagnosis. pp- 2024.12.11.24318888 doi: 10.11011/222.12111.2233333 243318815. URL httpt::://// www ncbi. nlm. nih. hutp://ww/mmm/arrrcccee////LLL/1
Benjamin Fox, Joy Jiang. Sajila Wickramaratne, Patricia Kovatch, Mayte Suarez-Farinar Neomi A. Shah, Ankit Parekh, and Girish N. Nadkarni. A foundational transformer leveraging full night, multichannel sleep study data accurately classifies sleep stages. pp. zsaf061. ISSN 1550-9109. doi: 10. 10.1093/slccp//ssf061.
Théo Gnassounou, Antoine Collas, Rémi Flamary, Karim Lounici, and Alexandre Gramfort. Multi- source and test-time domain adaptation on multivariate signals using apatio-tempora monge alignment. arXiv preprint arXiv:2407. 14303, 2024.
Théo Gnassounou, Rémi Flamary, and Alexandre Gramfort. Convolutional monge mapping nor- malization for learning on biosignals. In Neural Information Processing Systems (NeurIPS), 2023.
Ary Goldberger, Luís Amaral, L. Glass, Shlomo Havlin, J. Hausdorg, Plamen Ivanov, R. Mark, J. Mietus, G. Moody, Chung-Kang Peng, H. Stanley, and Physiotoolkit Physiobank. Components of a new research resource for complex physiologic signals. PhysioNer, 101, 01 2000.
Alexandre Gramfort, Martin Luessi, Eric Larson, Denis A. Engemann, Daniel Strohmeier, Christian Brodbeck, Roman Goj, Mainak Jas, Teon Brooks, Lauri Parkkonen, and Matti S. Hämäläinen. MEG and EEG data analysis with MNE-Python. Frontiers in Neuroscience, 7(267):1-13, 2013. doi: 10.3389/fmins..02 001.00066.
Robert M. Gray. Toeplitz and circulant matrices: A review. Foundations and Trends@ in Communica- fions and Information Theory, 2(3): 2006. ISSN 1567-2190. doi: 10.1561 1010000000 URL http::/t/aatt doi. wwtp/r/......... 0ttp///0000011111000000000000000
Antoine Guillot and Valentin Thorey. RobustSleepNet: Transfer learning for automated sleep staging at scale. 2021.
Yanchen Guo, Maciej Nowakowski, and Weiying Dai. Flexsleeprransfomme:: a transformer based sleep staging model with flexible input channel configurations. Scientific Reports, 14, 11 2024. doi: 10. 111103//415988-04-76977779
C.R. Harris, K.J. Millman, S.J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, Taylor, S. Berg, N.J. Smith, R. Kern, M. Picus, S. Hoyer, M.H. van Kerkwijk, M. Brett, A. Haldane, J. Fernández del Río, M. Wiebe, P. Peterson, P. G' G'erard-Marcham K. Sheppard, T. Reddy, W. Weekesser, H. Abbasi, C. Gohlke, and T.E. Oliphant. Array programming with NumPy. Nature, 585(7825)3357336, 2020.
John D. Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering, 9(3):
90-95, 2007.
Sergey loffe and Christian Szegedy. Batch normalization: accelerating deep network training by reducing internal covariate shift. Proceedings of the 32nd International Conference on International Conference on Machine Learning- Volume 37, ICML'15, pp. 2015.
JMLR.org,
Bakker Jessie P., Tavakkoli Ali, Rueschman Michael, Wang Wei, Andrews Robert, Malhotra Atul, Owens Robert L., Anand Amit, Dudley Katherine, and Patel Sanya R. Gastric Band- ing Surgery versus Continuous Positive Airway Pressure for Obstructive Sleep Apnea: A Randomized Controlled Trial. American journal of respiratory and critical care medicine, 197(8). April 2018 ISSN 1535-4970.
doi: 10. 11166///...2170000 1637LE. URL https Publisher: Am J Respir Crit Care Med.
//pubmed ncbi mmm. nih. gov/29635093/
12

---

Published as a conference paper at ICLR 2026
Taesung Kim, Jinhee Kim, Yunwon Tae, Cheonbok Park, Ho Choi, and Jaegul Choo. Re- versible instance normalization for accurate time-series forecasting against distribution shift. In International Conference Learning Representations, 2021.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv: 141.v9141..6980. 2014.
Reinmar Kobler, Jun-ichiro Hirayama, Qibin Zhao, and Motoaki Kawanabe. Spd domain-specific batch normalization crack interpretable unsupervised domain adaptation in EEG. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35. pp. 6219 6235.. Curran Associates, Inc., 2022.
Yanghao Li, Naiyan Wang, Jianping Shi, Jiaying Liu, and Xiaodi Hou. Revisiting batch normalization for practical domain adaptation, 2016. URL https: /ttppp::/aaaxi..oot org/ abs abs/1603. 01t.//103300777999
Carole L. Marcus, Reneé H. Moore. Carol L. Rosen, Bruno Giordani, Susan L. Garetz, H. Gerry Taylor, Ron B. Mitchell, Raouf Amin, Eliot S. Katz, Raanan Arens, Shalini Paruthi, Hiren Muzumdar, David Gozal, Nina Hattiangadi Thomas, Janice Ware, Dean Beebe, Karen Snyder, Lisa Elden, Robert C. Sprecher, Paul Willging. Dwight Jones, John P. Bent, Timothy Hoban, Ronald D. Chervin, Susan S. Ellenberg, Susan Redline, and Childhood Adenotonsillectomy Trial (CHAT), A randomized trial of adenotonsillectomy for childhood sleep apnea. The New England Journal of Medicine, 368(25)22366-233333 2376, June 2013. ISSN 1533-4406. doi: 10.1056/NEJMoa 1218881.
Jose G. Moreno- Torres, Troy Raeder, Rocío Alaiz-Rodríguez, Nitesh V. Chawla, and Francisco Herrera. A unifying view on dataset shift in classification. Pattern Recognition, 45(1):521 2012. ISSN 0031-3203. doi: https://doi.ooo http:://doi.or/1100 10. 1016//j.patcog.22220 patcog.2011.06.00
Christian O' O'Reilly, Nadia Gosselin, and Julie Carrier. Montreal archive of sleep studies: an open- access resource for instrument benchmarking and exploratory research. Journal of sleep research, 23, 06 2014. doi: 10.1111111 10js1 1/jrr.12169. 12169.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), pp. 8024-8035.
Curran Associates, Inc., 2019,
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Pretten- hofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine Learning in Python Journal of Machine Learning Research, 12:2825 -830,, 2011.
Mathias Perslev, Michael Jensen, Sune Darkner, Poul Jørgen Jennum, and Christian Igel. U-time: fully convolutional network for time series segmentation applied to sleep staging. Advances in Neural Information Processing Systems, 32, 2019.
Mathias Perslev, Sune Darkner, Lykke Kempfner, Miki Nikolic, Poul Jennum, and Christian Igel. U-Sleep: resilient high-frequency sleep staging. npj Digital Medicine, 4:72, 04 2021. doi: 10.100884444446-00144444000
Huy Phan, Fernando Andreotti, Navin Cooray, Oliver Y. Chén, and Maarten De Vos. SeqSleepNet: End-to-End Hierarchical Recurrent Neural Network for wewpuuenneett-SSqqeeeee Automatic Sleep Staging, February 2019. URL h.Rp:tttp/aaxx...oo
arxiv. org /109 10932 arXiv:1809.1099
Huy Phan, Oliver Y. Chen, Minh C. Tran, Philipp Koch, Alfred Mertins, and Maarten De Vos. XSleepNet: Multi-view sequential model for automatic sleep staging. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(09))5903-5915 sep 2022a. ISSN 1939-3539. doi:
Huy Phan, Kaare Mikkelsen, Oliver Y. Chén, Philipp Koch, Alfred Mertins, and Maarten De Vos. Sleeptransformer: Automatic sleep staging with interpretability and uncertainty quantification. IEEE Transactions on Biomedical Engineering. 69(8):4456224 2467, 2022b. doi: 10.
2022.3147187.
13

---

Published as a conference paper at ICLR 2026
Huy Phan, Kristian P. Lorenzen, Elisabeth Heremans, Oliver Y. Chén, Minh C. Tran, Philipp Koch, Alfred Mertins, Mathias Baumert, Kaare B. Mikkelsen, and Maarten De Vos. L- SeqSleepNet: Whole-cycle Long Sequence Modeling for Automatic Sleep Staging. IEEE Journal of Biomedical and Health Informatics, 7(10))4748-47511 October 2023. ISSN 2168-2208. doi: 10. 111////....2.3333000
URL htt..https:.. /tttps:/ieeexpl..o. eeee org. document./
Conference Name: IEEE Journal of Biomedical and
Health Informatics.
Stuart Quan, Barbara Howard, Conrad lber, James Kiley, Nieto, George O'Connor, David Rapoport, Susan Redline, John Robbins, Jonathan Samet, and Patricia Wahl. The sleep heart health study: Design, rationale, and methods. Sleep, 20: 01 1998. doi: 10. 10.10933sleep/22..12000 12.177777.
S. Redline, P. V. Tishler, T. D. Tosteson, J. Williamson, K. Kump, I. Browner, V. Ferrette, and P. Krejei. The familial aggregation of obstructive sleep apnea. American Journal of Respiratory and Critical Care Medicine, (31 Pt 1):682-687, March 1995. ISSN 1073-449X. doi: 10.1164/
ajreem/151.3 1.682.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MIC 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234 241. Springer, 2015.
Carol Rosen, Emma K. Larkin, H. Lester Kirchner, Judith I.. Emancipator, Sarah F. Bivins, Susan A. Surovec, Richard J. Martin, and Susan Redline. Prevalence and risk factors for sleep-disordered breathing in 8- I--ear-olld children: association with race and prematurity. The Journal of Pediatrics, 142(4):383-38) April 2003. ISSN 0022-3476. doi: 10. 20.1067mmd.2003.22
Carol L.. Rosen, Dennis Auckley, Ruth Benca, Nancy Foldvary-Schaefer, Conrad Iber, Vishesh Kapur, Michael Rueschman, Phyllis Zee, and Susan 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(6): doi: 10.5665/sleep.1877
757-767, June 2012. ISSN 1550-9109
Robin Tibor Schirrmeister, Jost Tobias Springenberg, Lukas Dominique Josef Fiederer, Martin Glasstetter, Katharina Eggensperger, Michael Tangermann, Frank Hutter, Wolfram Burgard, and Tonio Ball. Deep learning with convolutional neural networks for EEG decoding and visualization. Human Brain Mapping, 38(1 ()::3313351000 5420, November 2017. ISSN 10. 10.1002//bm.23733 URL http:/:://a/ww..or////abaaa 703.. 05051
1097-0193. doi: arXiv: 703.05051 [cs].
Adam P. Spira, Terri Blackwell, Katie L. Stone, Susan Redline, Jane A. Cauley, Sonia Ancoli-Israel, and Kristine Yaffe. Sleep-disordered breathing and cognition older women. Journal of the American Geriatrics Society, 56(1):45-50, January 2008. ISSN 1532-5415. 1532-5415.22 00.44415520..000 0157..00......
doi: 10.1111/j.
Jens B. Stephansen, Alexander N. Olesen, Mads Olsen, Aditya Ambati, Eileen B. Leary, Hyatt E. Moore, Oscar Carrillo, Ling Lin, Fang Han, Han Yan, Yun L. Sun, Yves Dauvilliers, Sabine Scholz, Lucie Barateau, Birgit Hogl, Ambra Stefani, Seung Chul Hong, Tae Won Kim, Fabio Pizza, Giuseppe Plazzi, Stefano Vandi, Elena Antelmi. Dimitri Perrin, Samuel T. Kuna, Paula K. Schweitzer, Clete Kushida, Paul E. Peppard, Helge B. D. Sorensen, Poul Jennum, and Emmanuel Mignot. Neural network analysis of sleep stages enables efficient diagnosis of narcolepsy. Nature Communications, 9(1), dec 2018. doi: 10. (0.10388/41167-018---77 --2022---
Suzanne Stevens and Glenn Clark. Chapter 6 polysomnography. In DAMIEN STEVENS (ed.), Sleep Medicine Secrets, pp. 45-63. Hanley & Belfus, 2004. ISBN 978-1-56053-59248 doi: https://do..rr/100 0. 1.100.11111110 1016/
Baochen Sun and Kate Sacnko. Deep CORAL: Correlation Alignment for Deep Domain Adaptation, July 2016. URL
abs ab////0/s///./.../////// 01 01719 1101:1607.01000 [cs).
Akara Supratak, Hao Dong, Chao Wu, and Yike Guo. DeepSleepNet: A model for automatic sleep stage scoring based on raw single-channel EEG. IEEE Transactions on Neural Systems and Rehabilitation Engineering, 25(11):1998-200 2017.
14

---

Published as a conference paper at ICLR 2026
Rahul Thapa, Magnus Ruud Kjær, Bryan He. lan Covert, Hyatt Moore, Umaer Hanif, Gauri Ganjoo, M. Brandon Westover, Poul Jennum, Andreas Brink- Brink-Kjærr Emmanuel Mignot, and James Zou. A multimodal sleep foundation model developed with 500k hours of sleep recordings for disease predictions. pp. 2025.02.04..5332177.
doi: 10. 02000/2222.022.04225551111
URL https
//ww. ncbi n mmm htt.p//go..m./rrraiceeeeeee//eee /1
Ulyanov. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016.
P. Virtanen, R. Gommers, T.E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, Bright, S.J. van der Walt, M. Brett, J. Wilson, J.K. Millman, N. Mayorov, A.R.J. Nelson, E. Jones, R. Kern, E. Larson, C.J. Carey, I. Polat, Y. Feng, E.W. Moore, VanderPlas, D. Laxalde, J. Perktold, R. Cimrman, Henriksen, E.A. Quintero, C.R. Harris, A.M. Archibald, A.H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020.
Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Dureell. Tent: Fully Test-time Adaptation by Entropy Minimization, March 2021. htLlhttpp//xxxi./
/bs/
2006. 10 10726 arXiv:2006. 10726 [cs, stat].
Jiquan Wang, Sha Zhao, Haiteng Jiang, Yangxuan Zhou, Zhenghe Yu, Tao Li, Shijian Li, and Gang Pan. CareSleepNet: A hybrid deep learning network for automatic sleep staging. (8(1))333 URL
7405. ISSN 2168-2208. doi: 10.
Peter Welch. The use of fast fourier transform for the estimation of power spectra: method based on time averaging over short, modified periodograms. IEEE Transactions on audio and electroacoustics, 15(2):70-73, 1967.
Chaoqi Yang, M Westover, and Jimeng Sun. Biot: Biosignal transformer for cross-data learning in the wild. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 78240-78260. Curran Asso- ciates, Inc., 2023. URL hta.hcppee/pproeednnss.eeuuuddip cepppp http://ppppppeleeppppppeeeeeee 2023 f: www.ffllefffee 202.:2211/553333addee22222222220
pdf
Shiqi Yang. Yaxing Wang, Joost Van De Weijer, Luis Herranz, and Shangling Jui. Generalized Source- free Domain Adaptation. In 202/ IEEE/CVF International Conference on Computer Vision (ICCV), pp. 8958-8967, Montreal, QC, Canada, October 2021. IEEE. ISBN 978- 178-11666 1-6--1--6-2----55 doi: 10. 10..109//CCC449922222..0000 9710764/]
URL https: htt:///eeeeeeereeeee// i ieee org/ document/ /
Guo-Qiang Zhang, Licong Cui, Remo Mueller, Shiqiang Tao, Matthew Kim, Michael Rueschman, Sara Mariani, Daniel Mobley, and Susan Redline. The National Sleep Research Resource: towards a sleep data commons. Journal of the American Medical Informatics Association: JAMIA, 25(10): October 2018 8. ISSN 1527-974X. doi: 10. 1033///aaaaa/oy0044.
15

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 PROOF OF THE BURES- WASSERSTEIN GEODESIC (6) BETWEEN COVARIANCE MATRICES OF STRUCTURE (2)
Proposition A.1. Let (x)) and E be two covariance matrices in Refxef following 2). Let us denote p(s) and p(t) the corresponding PSD matrices. The geodesic associated with the Bures- Wasserstein metric between ((t) and (() and parameterized by E [0, 1] is ((x)) following 2 of
PSD
P(a)= (1- + +ap(00z
Proof. From Bhatia et al. (2019). the geodesic associated with the Wares-Wasserstein metric between two covariance matrices Σ and (1) is given by
7(a) = (1 a) ( + (1) + a(1 - aa( (1(0) (0)}
(9)
where
E ( (
{ (1)) ((0) (+))
(10)
Since (+) and (()) diagonalize in the unitary basis I F, (a) also diagonalizes in this basis. Thus, we only have compute the geodesic between the PSD matrices (()) and p(t) and from now all operations are element- wise. Let P(a) be the PSD of y(a), we have
P(a) = (1 pp)) + (())
+ a(1 - a) (P()) p( + (P(9 @ p
(11)
(1 a) pp()) + ap
+ 2a(1 - a) @ p(
(12)
(1- p tap 000)
(13)
This concludes the proof.
A.2 BALANCED DATASETS
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.
Balanced datasets
Number of subjects
Balanced @ 40
360 + 0
Balanced @ 100
787 +
Balanced @ 200
1387 + 63
Balanced 4400
2466 = 157
All subjects
9929 + 1659
In the main paper, we report results across different training set sizes. Since the datasets 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 each balanced is shown in Table 3 Notably, the balanced set with 400 subjects contains roughly four times less data than the full dataset.
A.3 U-TIME: CNN FOR TIME SERIES SEGMENTATION
U-Time Perslev et al. (2019 202[ architecture Ronneherger et al. 2015), designed for segmenting temporal sequences. U-Time maps
| is a convolutional neural network (CNN) inspired by the U-Net
16

---

Published as a conference paper at ICLR 2026
sequential imputs of arbitrary length to sequences of class labels on 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 convolutional layer, an activation function, BatchNorm layer, and 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 G and the post-BatchNorm feature map G, i.e., G & BatchNorm (G). Each encoder block downsamples by the signal length but increases the number of channels.
Decoder blocks and Segmentation Head The decoding part of U-Time symmetrical to the encoding part. Each decoder block doubles the signal length and decreases the number of channels. composed of convolutional layer, an activation function, 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 Ba et al. (2016) or InstanceNorm possible.
Ulyanov (2016) are
The architecture is inspired from Braindecode |Schirmeister et al. (2017). The implementation is improved to make it more efficient and faster. One epoch of training takes about 30 min on a single H100 GPU.
Implementation
A.4 ARCHITECTURE: CNNTRANSFORMER
The CNNTransformer is a hybrid architecture designed for multichannel time series classification inspired by transformers for EEG-Data Wang et al.) Phan et al. (2022b); Yang et al. (2023): Thapa et al. It combines convolutional feature extraction with long-range temporal modeling via Transformer encoder at epoch-level. The model processes an input tensor of shape (B, S. C, T). where
the batch size, S is the number of temporal segments, C' is the number of input channels, and T the
number of time samples per segment. It outputs tensor of shape (B, Telassess S), where Tlllssss is the number of classes and S is the number of epochs.
The architecture consists of the following components:
The input is first permuted and reshaped 3D tensor of shape (B, C, T)) to be compatible with ID convolutional layers applied along the temporal dimension.
Reshaping:
CNN-based Feature Extractor: stack of 10 Conv ID layers, each followed by ELU activation and Batch Normalization. Some layers use stride greater than I to progressively reduce the temporal resolution. This block extracts local temporal patterns and increases the representational capacity up to a dimensionality of dmodet-
Adaptive Pooling: An AdaptiveAvgPoolll layer reduces the temporal length to a fixed number of steps (S), independent of the input sequence length. This step ensures consistent temporal resolution before the Transformer.
Positional Encoding: Learnable positional embeddings of shape (1, S, dmmode) are added to the feature representations to preserve temporal ordering before passing through the 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.
Transformer Encoder:
Classification Head: After transposing the data to shape (B, dmodel. S), a final ID convolu- tion with a kernel size of 1 projects the output to Nelasses, 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
17

---

Published as a conference paper at ICLR 2026
Table Architecture overview of the CNNTransformer model. In pratice, dmodel is 768, Nheend to 8, and S is 35.
Stage
Operation
Details
Output Shape
Input
Raw signal
Multichannel EEG signal with S segments and (,
time samples per segment
Reshape
Permute & flatten
Rearranged as (B, C, S T) to process with ID (B, (T) convolutions
Feature Extractor
ID CNN stack
10-layer sequence of ConvID + ELU + Batch- Norm; includes temporal downsampling via stride
(B, dmodess T)
Temporal Pooling
AdaptiveAvgPoolID
Downsamples fixed temporal resolution defined by S
(B, duustes S)
Positional Eincoding
Learnable embeddings
Added to temporal dimension to encode temporal order before transformer layers
(B, dmadees S)
Transformer Encoder
Multi-head attention
2 Transformer layers with dmudel embedding di- mension, Phead heads, and feedforward sublayers
(B, dmodets
Classifier
Linear projection
Projects feature vectors class logits each epoch time step
(B, Ficlusess S)
A.5 EQUATION FOR BATCHÑORM AND INSTANCENORM
NarchNorm The BatchNorm layer loffe Szegedy (2015) normalizes features maps in neural network to have zero mean and unit variance, At train time, given batch { ..... ,.... G(M} c
Rexe of N ne-BBathNNorm feature maps and for all j, m, [1, M1 x [1,c] cl x [1, [] the BatchNorm layer is computed
(ii)
Pm
ワ m.
7m
mnd
8m
(14)
V
02 +
where 7,3 E TRC are learnable parameters. The mean and standard deviation ğu E RR and E R computed across the time and the batch,
(
Tim
NO
m,
j=1
(15)
N
4
NO
m.!
pm
At test time, the mean and variance ğá and are replaced by their running mean and variance, also called exponential moving average, estimated during training.
InstanceNorm layer Ulyanov (2016). During training. InstanceNorm operates similarly to (14), but the mean and variance are computed per sample instead of across the batch dimension, i.e., (. and a(fe are computed for each sample j.
InstanceNorm Another popular normalization
P: 4
(
m,
(16)
( 2
(ali)-
Hence, each sensor of each sample is normalized independently of the others. At test time, Instan- ceNorm behaves identically to its training phase and therefore does not rely on running statistics contrary to the BatchNorm.
A.6 SENSITIVITY TO FILTER SIZE
The filter size f in PSDNorm controls the temporal context used for normalization, influencing the strength of adaptation temporal variations. shows the impact of different f values on the
18

---

Published as a conference paper at ICLR 2026
Balanced Accuracy Score
PSDNorm
0.79
InstanceNorm
0.78
BatchNorm.
0.77
11
13
15
17
21
Filter size
Figure 6: Performance of PSDNorm with varying filter sizes. The BACC score plotted against the filter size used with U-Sleep.
BACC score across datasets using U-Sleep trained on balanced 400. This experiments shows that for any f. PSDNorm consistently improves performance compared to other normalization techniques. Taking a f between and yields the best results, with peak at f = 5. Smaller values (e.g.. - 1, equivalent to InstanceNorm) provide less adaptation, while larger values (e.g., ʃ = 21) may over-smooth temporal variations, leading to diminished performance. Overall, the results shows that is not sensitive yielding good performance for wide range of values.
A.7 F1 SCORE BALANCED ACCURACY
In the main paper, report Balanced Accuracy scores, which account for class imbalance in sleep stage classification. Prior work, such as the U-Time paper Perslev et al. (2019), uses the F1 score evaluate performance. In Table 5 report F1 scores 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 datasets, while InstanceNorm is never the top performer.
These results confirm that our implementation achieves state-of-the offth------e-rtt performance in sleep stage classification. Moreover, PSDNorm maintains its advantage even in data-limited settings
A.8 IMPACT OF WHITENING AND TARGET COVARIANCE
As explained in the main paper, InstanceNorm is 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.c., colorization).
In this section, we evaluate the impact of whitening on the performance of PSDNorm, to assess the benefit of using barycenter the target covariance matrix. Table reports results 10 datasets (alancee@4000) with and without whitening.
Whitening improves performance only one dataset (CCSHS), while projecting to the barycenter yields the best results on datasets.
This suggests that, while whitening may help when F = 1, it is less effective when F 1. Using barycenter leads to a more robust and stable target covariance matrix.
A.9 GENERALIZATION OF PSDNORM IN CNNTRANSFORMER
The CNNTransformer architecture a hybrid model that combines convolutional and transformer layers for time series classification.
19

---

Published as a conference paper at ICLR 2026
Table Fl 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.
Dataset
BatchNorm
LayerNorm
InstanceNorm
TMA
PSDNorm(F=5)
ABC
$1.00 11
79.50
80.56+0.39
=0.89
81.12=0033
CCSHS
89.83 0.19
89.01 10.43
89.39
89.37
89.13
CFS
88.30
87 87.39+0..0
88.45+0.17
88. 08.28
88. 80.52
All subjects
CHAT
65.77=44.06
65.25.43.9
71.35+2.75
71.80
72. 72.16+2.21
HOMEPAP
77.06
76.62
77.50
77 77.82
77 77.30 77.30
MASS
77.27
74.21+2.05
75.12+2.08
77. 77.74#1.05
76. 76.00+3.00
MROS
85.53=0.48 48
84.02++0.95
80.12
0..13
85. 00.022
PhysioNet
74.98
74.29
75.07 11.05
76. 10.73
75. 75.29 1.21
SHHS
78.95
18.04
80.30+1.29
78. 08.84
80.32=0.99
SOF
86.30
85.82++022
86.57 +0.60
86.31
86. 86.9940.55
Mean(Dataset) 80.50
19.41
$0.95
81. $1.22: 1.020
81.19
Mean(Subject)
80.05
79.094+0.90
81.31+0.53
80.59
81.39+0.60
ABC
79.80
77.86
78.36 1..20
79.49
78 78.08
CCSHS
88.32
87 87.22+0.51
88. 00.552
88.47 20.62
88.79 88.7940.99
CFS
87.01 10.18
85.61 10.16
87.62 10.27
87.37
87 87.06
Balanced Balanced@4400
CHAT
66.5641.42
61 61.3242.22
64.19
69. 69.90
71.86 20095
HOMEPAР
76,20+1.25
76.1541.11
77.66+0.58
76.83
77,85+11 +
MASS
76.0611.00
73.95
76.94
76.32
1.773
MROS
83.6940.30
82.22
83.95
84.15
83.51 +0.84
PhysioNet
76.26
70.40
73.84
75.24
73.51+3.05
SHHS
76.98
75.98
79. 79.12 10.96
0..0
79.26611.0
SOF
85.4940.58
84.23
85.50
85.56 1..90
84. 84.14
Mcan(Datasct)
79.64
77.5740.73
80.15 30.15.0.22 0.
80. 30.12.40.55
Mean Mcan(Subject)
78.57
76.86
79.53
79.70=..ee
80. 40.29.2m 10.68
Table 6: Impact of the whitening on the performance of PSDNorm on the 10 datasets balanced @ 400.
PSDNorm
Dataset
BatchNorm
InstanceNorm
Barycenter
Whitening
ABC
78.26 1.33
78.73
78.1840.08
77.86
CCSHS
87.42
87.62-0.42
87 87.5840.30
87.80+0.29
CFS
84.32
84.72+0.33
84.29+0.36
84.01
CHAT
66.5510.46
64.43
70.28 1.70
69.07 13.73
HOMEPAP
15.25
76.47
76. 76.83... 0.61
76. 10.93
MASS
70.00+1.91
71 71.52+1.13
72.77
69. 69.11. +1.51
MROS
80.37
80.28+0.21
80.26+0.11
80. =0.0775
PhysioNet
75.81
74.68
74.82+2.11
74.58
SHHS
76.44
78.68+0.57
78.88=0.68
78. 78.77
SOF
81.08
80.68
79.494t.41
80.1044.62 10.62
Mean
77.55+0.34
77.78+0.46
78.34
77. 77.79+0.30
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 pest-performing normalization layer.
In Table 7 we report the results of different normalization layers used in the CNNTransformer architecture datasets balanced@400..
20

---

Published as a conference paper at ICLR 2026
Table 7: Different normalization layers used in the CNNTransformer architecture for datasets
balanced 4400.
Dataset
BatchNorm
InstanceNorm
TMA
PSDNorm
ABC
76.99
75.40=0.30
77.50
76.31
CCSHS
86.75 =0.4
87.00
86.73 10.25
86. 10.92
CFS
83.32 0.33
83.7710.34
83.16 10.38
83. 63.71 10.29
CHAT
66.44
66.40
66.47=1.55
10.34
HOMEPAP
74.81
75.92.0.44
74.7641088
75. 25.26
MASS
71.51
71-70
70.57
72.55.10881
MROS
79.77=0.31
79.74
79.85=0.08
79. 70.77
PhysioNet
72.54-0.31
74.36=0.81
71.39+1.38
74.95=0.41
SHHS
75.34_0.54
76.55.0..2
75. 75.15.1..0 10.98
77. 26 10557
SOF
80.63 0.60
80.78 10.58
81.03 00.48
80.31 10.90
Mean
76.38_0.17
77.07 20.28
76.30 0..57
77. 77.83
Balanced Accuracy Score
0.78
Normalization
0.76
BatchNorm
InstanceNorm
PSDNorm(f=5)
0.74
All subjects
Number of subjects
Figure 7: 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.
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 outperforms BatchNorm by margin of 0.9 and InstanceNorm by 0.54 in average score.
These results highlight that PSDNorm is plug-andd-layy normalization layer that can be seamlessly integrated into various architectures to reduce feature space variability.
A.10 EVOLUTION OF PERFORMANCE WITH TRAINING SET 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. Figure we evaluate impact across different training set sizes and observe 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., 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 offers good compromise, achieving one of the best performances across all training sizes.
A.11 CRITICAL DIFFERENCE DIAGRAM FOR U-SLEEP ON ALL SUBIECTS
The main paper presents the critical difference diagram for U-Sleep the dataset balanced @ 400. Figure & extends this analysis to all subjects across datasets. The conclusion remains consistent: PSDNNom with F = 5 is the best-perfomming normalization layer, while NatchNorm performs the
21

---

Published as a conference paper at ICLR 2026
worst. Interestingly, PSDNorm with F 17 ranks second to last, suggesting that overly strong adaptation can hurt performance when the dataset is large.
USleep
2.2
2.4
2.6
2.8
3.0
PSDNorm
BatchNorm
InstanceNorm
TMA
Figure 8: Critical difference diagram for U-Sleep all subjects.
A.12
COMPUTATIONAL TIME OF PSDNORM
Table 8: Computational time of PSDNorm compared BatchNorm and InstanceNorm for USleep and CNNTransformer. The time is done for leave out the dataset CHAT and with the dataset balanced @ 4400. The time averaged over 3 runs and reported in seconds.
Model
Normalization
Time per epoch (sec)
Time of inference (sec)
USleep
BatchNorm
161.94 10.18
95.63 + 4.85
USleep
InstanceNorm
258.71 + 2x15(*)
93.40 + 8.18
USleep
PSDNorm(f
172.85 + 4.05
98.72 12.09
CNNTransformer
BatchNorm
130.88 + 2.67
93.03 + 5.25
CNNTransformer
InstanceNorm
127.47 + 5.14
92.70 + 4.11
CNNTransformer
PSDNorm(f = 152.83 + 2.29
92.57 2.64
One important aspect of normalization layers is their computational cost, which can impact training and inference times. Table compares the computational time PSDNorm with BatchNorm and InstanceNorm in both U-Sleep and CNNTransformer architectures. In U-Sleep, PSDNorm takes 172.85 seconds per epoch, which slightly higher than BatchNorm (161.94 seconds) but significantly lower than InstanceNorm (258.71 seconds). The high cost of InstanceNorm is due to the fact that the torch.compile was not working for Usleep and InstanceNorm. For inference, PSDNorm takes 98.72 seconds, which is comparable to BatchNorm (95.63 seconds) but slightly higher than InstanceNorm (93.40 seconds).
In CNNTransformer. PSDNorm takes 152.83 seconds per epoch, which is higher than BatchNorm (130.88 seconds) and InstanceNorm (127.47 seconds). However, for inference, PSDNorm is equiv- alent to both BatchNorm and InstanceNorm. The modest computational overhead introduced by PSDNorm is worthwhile trade-off for its superior performance. This efficiency is enabled by the highly optimized implementation of the Fast Fourier Transform (FFT) on GPUs.
A.13 CLASS- WISE PERFORMANCE
The Tables show that the most complicated sleep stage to classify is N1, with F1 scores consistently lower than other stages across all normalization methods. This is likely due to the inherent di difficulty of distinguishing NI from other stages, shares characteristics with both wakefulness and deeper sleep stages. In contrast, stages like Wake and REM tend to have higher Fl scores, indicating that they easier to classify accurately.
Figure iillustrates the class-wise score differences between normalization layers score against BatchNorm score. For almost all the classes, other nonmalization increase the performance compared to BatchNorm except for where InstanceNorm shows decrease in performance. PSDNorm is consistently the best performing normalization across all classes, highlighting its effectiveness in
22

---

Published as a conference paper at ICLR 2026
Table 9: Class-wise Fl scores for BatchNorm layer on datasets balanced @ 400.
Dataset
Wake
N1
N2
N3
REM
F1
ABC
86.07 10.68
53.97 10.36
80.05
71.35
88.43 10.13
79.80 10.34
CCSHS
95.22+0.46
48.08+2.87
84.88
86.13 86.13+1001
88.5440.93
88.32++0.49
CFS
94.6440.10
42.69+0.39
82.27+0.75
77. 77.20+0.30
86.9440.34
87.01+0.18
CHAT
18.3
35.57+2.64
52.634:2.77
72.67+1.25
76.12+1.50
66.5641.42
HOMEPAP
84.51+1.01
41.45+t071
73.88+-2.00
57.41+1.72
82.52+1..3
76.20+1:25
MASS
66.93
40.02422.0
78.7340.67
67.0110.39
75.8045.65
76.06+1.69
MROS
94.63
41.7311.11
73.7410.72
47.8610.33
82.7010.46
83.6 13.69910.39
PhysioNet
89.22.0.49
46.01
73.93
55.04
77.30.0.61
76. 16.26411.27
SHHS
85.68+1.57
32.56+1.31
72.48+22.00
61.91+1.11
79.09+1.07
76.98+0.70
SOF
93.91+0.15
38.29+1.09
79.15+0.57
71.86+3.17
86.49 18
85.49+0
Mean
86.91
42.04+1.34
75.17+1.39
66.84+1.25
82.39+1.22
79.6440.88
Table 10: Class-wise scores for LayerNorm layer with f = 5 on datasets balanced 400.
Dataset
Wake
N1
N2
N3
REM
F1
ABC
83.29
52.93
78.07+1100
68.91 +2.99
85.05
77. 10.80
CCSHS
93.66.10.40
40.68
84.80++1.26
86.46.1.88
0..22.82
85.39
87.2240.51
CFS
93.7840.50
38.5642.07
81.15
75.39
83.6840.09
85.61
CHAT
72.07
30.11
46.82.16.98
70.45
68. 18.077
61.32
HOMEPAP
83.58+1.72
43.85
73.97
57 57.54
79. 79.1641.11
76.1541.13
MASS
60.73+41.2
40.5640.87
77.39+6.28
65.91+1.01
69 69.92++1.55
73.95
MROS
94.12+0010
38.4142.21
71.4843.66
47.85 47.85+0.37
79.05+00.6
82.22+1.27
PhysioNet
88.04
44.7212.03
64.5311.93
48.49 10.54
68.17
70.40.10.14
SHHS
84.61 1:2.49
32.89 11.92
72.02
61.13
75.63
75.98
SOF
93. 93.2440.27
35.61
77.4942.15
70.83
83.52+0.01
84.23
Mean
84.71
39.83 1.58
72.77
65.30
77.7742.00
77.4941.36
Table 11: Class-wise Fl scores for Instancenorm layer on datasets balanced @ 400.
Dataset
Wake
NI
N2
N3
REM
F1
ABC
86.46
54.45
75.6842.16
70.7540.51
88.60+0.75
78.36
CCSHS
95.6940.14
49. 19.23422.2
85.29 $5.2940.89
85.9840.68
89. 89.2344099
88. 13.7334..0
CFS
94.95+0.19
44.97+1.31
83.30+0.60
77.1740.39
87.4540.38
87 87.62++0.2
CHAT
77.67
29.47-8.15
48.78+6.21
71.34+261
74.58
64.19
HOMEPAP
86.27
43.1911.45
75.1911.07
58.3911.03
83.44
77. 77.6610.58
MASS
67.25
13.19422.0
78.64
65.7811.37
78.34
76.94
MROS
94.80+0.18
41.4640.71
74.42+
48.8942.31
82.1140.10
83. 03.95440.58
PhysioNet
89.43+0.41
44.35++0.2
68.91422.2
51.0541.32
77.3540.95
73.8440.93
SHHS
88.62+0.30
33.02+2.26
74.31
64.28
80.32+t.42
79. 19.1244000
SOF
94.42+0..0
37. 37.18
78.43+1.88
72.3941.56
86.82++070
85.50+0.50
Mean
87.55+1.20
42.05
74.29+2..6
66.60+1.25
82.83+0.77
79.5941.16
23

---

Published as a conference paper at ICLR 2026
Table Class-wise F1 scores for TMA preprocessing with f = 5 on datasets balanced 400.
Dataset
Wake
NI
N2
N3
REM
F1
ABC
80.50
54.7640.88
18.9
71.40+1.23
88.43+0.07
79.4940.66
CCSHS
95.51 10.31
48.72 18.7222.0 12.07
85.0211.09
85.37 1.33
89.33..0.29
88.4710.68
CFS
94.7510.26
43.28
83.0610.79
77-38
87.17.0.33
87.37.10.4
CHAT
80.70+4.88
37.51
58.89+41.08
75.71+2.09
75.9542.35
69.90+2.74
HOMEPAP
84.13+0.98
43.92+0.51
74.83+
57.94+1.19
81.58400.1
76. 76.8340099
MASS
70.45+7.07
41.83+3.63
77.30+
65.04+2.01
79.50+2.42
76.32+0036
MROS
94.5040.32
41.7441.18
75.20+4..26
48.92+2.07
82.35-10.70
84. 15:10.4
PhysioNet
89.40+0.50
44.65+-2.19
71.13++5.36
51.6143.94
80.06+1.02
75. 75.244..72
SHHS
88.30
32.95
73.2342.58
62.33
79.27
78. 18.1940.00
SOF
93.6940.34
37.8542.16
79.3341.49
72.4611.92
86.8310.29
85.564.0.0
Mean
87.6911.63
42.72
75.6912.08
66.8211.64
83.05
80. 15.1..0
Table 13: Class-wise Fl scores for PSDNorm layer with f 5 on datasets balanced @ 400.
Dataset
Wake
N1
N2
N3
REM
Fl
ABC
14.57
54.4640.59
75.94+ 1.14
70. 70.73+1.00
88. 88.1940.
78.08+0.78
CCSHS
95.7640.21
17.9
15.3411.76
86.1742.24
89.71
88. 18.799
CFS
95.01
42.92+1.06
82.08
76.98+0.96
87.31
87. 87.06.1077
CHAT
82.93
36.5946.79
61.8442.98
77.0641.32
78.01+1.62
71.8640.95
HOMEPAP
85.75+1.45
44.4740.78
75.54+1.72
58.7041.40
83.2440.55
77.85+1.22
MASS
72.74
42.20
78.5642..98
66.1342.37
78.2343.21
77. 77.1641.73
MROS
94.63+0.31
41.40+1.61
73.33++180
47.5642.25
82.52+0.47
83.5140.84
PhysioNet
89.48+0.41
44.33+1.43
67.67+6.16
49.37+4.91
79.23+1.21
73.51 05
SHHS
80.09
33.7842.53
74.15
64.4010.29
80.2441.03
19.26
SOF
93.74
34.70+2.45
76.584.2.97
70.2042.02
85.5041.63
84. 84.1441.05
Mean
88.37 +0.90
42.284:2.01
75.104:2.61
66. 66.73+1.88
83.22+1.04
80.12
Normalization InstanceNorm
Difference from BatchNorm
0.03
TMA
PSDNorm(F=5 (5)
0.02
0.00
-0.01
Figure 9: Class-wise Fl score differences between normalization layers. The variance is giving by the seeds.
24

---

Published as a conference paper at ICLR 2026
0.85
Balanced Accuracy
0.80
0.75
0.70
0.65
1
5
7
11
Number of layers with Temporal Normalization
Figure 10: Impact of the number of layers in U-Sleep using PSDNorm with f 5. The BACC score plotted against the number of layers. The variance is one the Datasets.
Table 14: Performance of different normalization layers in U-Sleep in very low data regime on datasets balanced @ 40.
Dataset
BatchNorm
LayerNorm
InstanceNorm
PSDNorm(F=5)
PSDNorm(F=15)
ABC
73.83
64.5443.46
12.29+1.50
71.35
72.61
CCSHS
83.58.00.5
7..91
85.33_0.69
85. 85.10.0.17
85.0040.25
CFS
81.1340.85
76.57
81.59
81.7940.89
80.9340.13
CHAT
55.7442.38
58. 28.44+1..1
63.3845.26
59.6640.89
67.86+3.59
HOMEPAP
1..52
72.19 1.665
76.03
76.01 10.32
76.14 1.61
MASS
70.1513.09
64.79
66.5810.50
69.49
68.21 16.25
MROS
77.12.:0.03
71.76
76.59
77.19 10.38
76.77 1:1.30
PhysioNet
71.68+1.61
69.59
72.6843.09
73.67
72.08 72.08
SHHS
73.74+1.11
71.50+0.97
75.5641.60
75.43+1.38
76. 76.004.0.6
SOF
75.84+2.16
73.50
76. 76.54+1.59
75.14+1.79
76. 76.00
Mean
73.35
70.72+1.2
75.19 1.03
74.79
75. 75.88
improving sleep stage classification. But have to note that for N1 and N2 the variance over the seed is big showing the instability of the training for these classes.
A.14 STUDY OF IMPACT NUMBER LAYER U-SLEEP USING PSDNORM
In the main paper, we apply PSDNorm in layers of U-Sleep. Here, investigate the impact of varying the number of layers that utilize PSDNorm. Figure I0 shows the BACC score as a function of the number of layers with PSDNorm. The results indicate that increasing the number of layers with PSDNorm reach a plateau after 3 layers, but does reduce the variance across datasets. This suggests that while adding more layers with PSDNorm can enhance performance, there are diminishing returns beyond certain point. Thus, using PSDNorm in 3 layers strikes good balance between performance, good variance, and computational efficiency.
A.15 STUDY VERY LOW DATA REGIME
In this section, explore the performance of different normalization layers in U-Sleep when trained on a very limited dataset, specifically balanced @ 40 subjects. The results indicate that in this low data regime, PSDNorm with small filter size (F 5) struggle to outperform InstanceNorm while still outperforming BatchNorm and LayerNorm. However, PSDNorm with a larger filter size (F = 15) gives the best average performance across datasets with an increase of more than 10% in BACC
25

---

Published as a conference paper at ICLR 2026
Table 15: Comparison of PSDNorm with AdaBN on datasets balanced @400 using U-Sleep.
Dataset
BatchNorm
LayerNorm
InstanceNorm
AdaBN(3)
AdaBN(12)
AdaBN(full)
TMA
PSDNorm
ABC
78. 26.1..3
26.29
78. :0.43
78.25=1.30
77.21 56
26.80
78. 04.00....5
78.18
CCSHS
87 12
85.20
87 62 :0.42
87.38
NaN
86,90 86.99
87. 57.57 10.50
87.58, 87.58.0.30
54.32
81.66g1.36
84.72.0.a
20.00
NaN
83.61ammm NaN
61.5820.20
84.29
CHAT
60.55
61.19=1.16
64.43=4.41
66.49-0.89
NaN
68. 23..2.2
70.28
HOMEPAP
75.25+0.50
74.86.0.25
76.47=0.68
76.148
24.30
74.46
26.10
76.8 76.83=0.61
MASS
10.00
68.56.3.as
71.52:1.18
69.68_1.05
68. 16
68. (8.31
11:63:1.22
72.77
MROS
80.37
78.05
80.28go.m
80.34_0.20
NaN
NaN
80.09
80.266.0.1
PhysioNet
75.81
71.82.2.12
74.08=0.60
75.27
74.01
14:01
75.31
74.82
SHHS
76.4420.92
76.1220.a9
78. =8
76.43=0.92
NaN
NaNN
77.00+0.39
78. 78.88&0.68
SOF
81.08
78.70g0.a0
80.68
81.05g1.15
NaN
81.17 amam
81.25.:0.7
79.19:0.44
Mean
77. 77.2220.54
75.0420.42
78. 28.17
77.18
74. 21.29
76.59
77.74g0.56
78.85=0.50
compared to other normalization layers for CHAT dataset. This suggests that in scenarios with very limited data, stronger normalization (larger filter size) is beneficial to prevent overfitting and enhance generalization.
A.16 COMPARISON WITH ADABN
Table 15 presents comparison between PSDNorm and AdaBN using the U-Sleep architecture on datasets balanced @ @400. AdaBN adapts the BatchNorm statistics separately for each subject. In original paper, all BN layers are replaced (AdaBN(full)), but for fair comparison we also evaluate two additional settings: AdaBN(3), which adapts only the first three BN layers, and AdaBN(12), which adapts only the first BN layers of the encoders.
As expected, AdaBN struggles to achieve strong performance on sleep staging. consistently underperforms compared to TMA and, in some cases, even performs worse than standard BatchNorm. Notably, increasing the number adapted BN layers further degrades performance, highlighting the importance of not adapting too many layers within the model. In contrast, PSDNorm consistently outperforms AdaBN across all datasets, demonstrating its effectiveness in normalizing features for sleep staging tasks.
26


Reviews:
[{"review_id": "yAvNvIllZF", "reviewer": "Reviewer_Q7SG", "summary": "The paper proposes PSDNorm, a normalization layer for deep learning models on time-series signals, particularly EEG-based sleep staging. Traditional normalization methods (BatchNorm, LayerNorm, InstanceNorm) do not properly handle spectral distribution shifts arising from temporal autocorrelations and subject variability. PSDNorm addresses this issue by integrating Temporal Monge Alignment (TMA) within the network, allowing temporal context to be leveraged during normalization rather than as a separate preprocessing step. The method consists of three stages: 1) PSD estimation, 2) Running Riemannian barycenter update, and 3) Temporal Monge alignment. Experiments are conducted on 10 large-scale sleep datasets (10M samples, 10K subjects) under a Leave-One-Dataset-Out (LODO) setup. Experimental results show that PSDNorm outperforms all baseline methods, BatchNorm, LayerNorm, InstanceNorm, and TMA, demonstrating improved robustness under domain shifts.", "questions": "Please address the above weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- **S1:** The paper addresses an important problem, domain shift in physiological datasets, where variability across subjects and recording conditions poses a major challenge.\n- **S2:** The paper is theoretically well-grounded, combining Monge mapping, Riemannian barycenters, and temporal alignment within a unified normalization framework. It is mathematically clear.\n- **S3:** Consistent improvements on both USleep and CNNTransformer indicate its strong applicability to other architectures.\n- **S4:** The analysis on subject-wise performance demonstrates the robustness of PSDNorm to inter-subject variability, one of the most difficult problems in physiological data modeling.\n- **S5:** The idea of normalization based on spectral composition shows strong potential for sleep staging and other biomedical applications where temporal dependencies are critical.", "weaknesses": "- **W1:** The motivation behind PSDNorm is not entirely clear. While the paper states that PSDNorm is inspired by Temporal Monge Alignment (TMA), the concepts and roles of data preprocessing (TMA) and in-network normalization (PSDNorm) are fundamentally different. The rationale for integrating TMA into a normalization layer, and the specific benefits of this design are not sufficiently demonstrated.\n- **W2:** Experiments are primarily conducted on large, well-curated datasets. Low-resource or few-subject scenarios are not adequately tested (the balanced@400 setting is still relatively large). For example, can the authors provide the performance of PSDNorm on balanced@40? While the authors argue that PSDNorm is effective in limited-data settings, the result in Table 2 does not fully support their argument (reports only balanced@400, which seems close to medium-scale rather than small-scale).\n- **W3:** The normalization baselines are limited to BN, IN, and LN. Broader domain adaptation methods such as AdaBN, DSBN, or Deep CORAL would strengthen the comparative analysis. \n- **W4:**  While PSDNorm targets dataset shift at test time, its behavior across datasets with highly divergent characteristics (Table 1) remains ambiguous. Performance trends in Table 2 do not clearly confirm improved robustness to severe distribution shifts. \n- **W5:**  Although Figure 4 provides subject-level insight, since sleep datasets are often highly imbalanced, a class-wise performance analysis for poorly predicted subjects would provide more insight. The results from Table 2 and Table 5 indicate a higher F1-score compared to BAAC (balanced accuracy), which appears as certain classes dominating the overall improvements. Could the authors provide class-wise performance results for further clarification?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3gHyxI90Mp", "reviewer": "Reviewer_yQGx", "summary": "The paper introduces a new normalization method called PSDNorm to address the challenges that models applied to sleep staging face. The authors point out challenges in building neural networks for sleep staging, like distribution shift across subjects, devices, and datasets, and they point out limitations in existing normalization layers like BatchNorm/LayerNorm, which ignore temporal autocorrelation in signals. PSDNorm aligns the power spectral density of intermediate representations to a barycenter using an f-Monge mapping. The authors claim that they achieve SOTA generalization under a leave-one-dataset evaluation across 10 datasets, and they find that they are 4x more data efficient.", "questions": "•\tHyperparameters: The paper fixes \\alpha and f, with sensitivity analysis mainly on f. How important is \\alpha? Other aspects like Welch window overlap or FFT length could be useful to discuss as well to get a sense of stability and cost.\n•\tEEG channels: Will PSDNorm actually scale to other channels or modalities?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "•\tPrincipled Formulation: The optimal transport on Gaussian models is clean. I can understand why this normalization would be beneficial for OOD generalization.\n•\tDomain shift: Authors show that PSD has a strong gain, notably with MASS and CHAT, and on the hardest subjects, as shown in Figure 4.\n•\tData Efficiency: If the hyperparameter is well-tuned, it seems that PSDNorm requires 4x fewer labeled subjects than BatchNorm.", "weaknesses": "I think this paper needs stronger comparisons in order to better characterize whether PSDNorm is making a real difference. I also think better scoping the claims will help significantly.\n\n•\tNormalization replacement: In my opinion, the paper makes a very odd choice in the experiments, as stated in line 386, with mixing normalization strategies. Instead of replacing the entire network, the paper takes a CNN Transformer and replaces the first three BatchNorm layers, keeping all other BatchNorm layers. For TMA, no layer is replaced, and this is treated as a preprocessing step, but TMA uses BatchNorm as well. This design makes comparisons confounded by remaining BatchNorm layers and by different amounts of architectural change across baselines. I think a full replacement is required here for clean comparison. \n•\tChannel independence: The method uses per-channel PSDs, from my understanding, as shown in Eq. 2. I think this assumes uncorrelated sensors. For learned feature maps, channel independence is unlikely to hold, so cross-channel correlations are ignored. I think this caps achievable alignment across channels. \n•\tTest-time adaptation: This point is more minor. The paper discusses TTA in Section 3.3. However, I think the comparison overlooks key issues with PSDNorm for TTA that make me uncertain about the claims in line 274. PSDNorm does not update its barycenter at inference (Algorithm 1), so it does not adapt to a novel target domain at test time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GkhfnpGrqN", "reviewer": "Reviewer_vTBu", "summary": "This paper proposes a novel deep learning normalization method—PSDNorm (Temporal Normalization Layer)—aimed at addressing the data distribution bias problem in biomedical data (such as sleep staging). Traditional normalization methods (such as BatchNorm, LayerNorm, and InstanceNorm) typically ignore temporal correlations, while PSDNorm, by employing Monogram mapping and temporal context, can effectively normalize feature maps in deep learning models, thereby improving performance on unseen data.", "questions": "See the section on weaknesses for details.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1) The proposed PSDNorm layer effectively addresses the data distribution shift problem in biomedical signals within a deep learning framework, providing a novel perspective and approach for processing complex time-series signals.\n\n2) Experimental results show that PSDNorm exhibits superior performance on multiple large-scale sleep datasets, particularly in handling data distribution shifts, outperforming existing normalization methods (such as BatchNorm and LayerNorm).\n\n3) PSDNorm is a \"plug-and-play\" normalization layer that can be seamlessly integrated into existing neural network architectures with relatively low computational cost.", "weaknesses": "1) The proposed PSDNorm introduces a filter size hyperparameter (f), which controls the normalization strength. The choice of f can affect model performance across different datasets and tasks. Although the authors provide some default values, automatically selecting this parameter in adaptive settings may still be a challenge.\n\n2) While the experiments cover multiple sleep datasets, the diversity of these datasets remains limited and may not fully represent all challenges in biomedical signals. Future research could consider extending this method to other domains, such as audio processing and other biomedical applications, to further validate its generalization capabilities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "F8Htugvu98", "reviewer": "Reviewer_ovsu", "summary": "This paper introduces PSDNorm, a novel normalization layer designed for deep learning on sequential signals. Unlike standard normalization layers (BatchNorm, LayerNorm, InstanceNorm) that treat time steps independently, PSDNorm leverages the temporal auto-correlation structure of signals by aligning their Power Spectral Density (PSD) to a running Riemannian barycenter via Monge mapping. The method is evaluated extensively on sleep stage classification using 10 datasets comprising ~10K subjects. Results demonstrate consistent improvements over baseline normalization techniques, with particular benefits in data-limited regimes and under domain shift.", "questions": "1. Can you empirically verify that intermediate feature maps exhibit the assumed covariance structure (Eq. 2)? Visualizations of PSD estimates from actual feature maps would strengthen the claims.\n2. Under what conditions (signal properties, network architecture, data distribution) should we expect PSDNorm to outperform alternatives? Can you provide theoretical guidance?\n3. What do the learned running barycenter PSDs look like? Are they interpretable? Do they differ meaningfully across layers or datasets?\n4. Have you evaluated PSDNorm on other time series tasks (e.g., ECG, audio, speech)? Results on even 1-2 additional domains would strengthen generality claims.\n5. Have you tried alternatives to the geodesic update (Eq. 6)? What about learnable momentum α or adaptive filter sizes f?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. LODO setup with 10 datasets, proper statistical testing, multiple seeds, and comprehensive ablations set a high standard. The transparency in reporting (including failed torch.compile) is commendable.\n\n2. The use of optimal transport, Wasserstein barycenters, and Riemannian geometry provides solid theoretical foundation, even if the novelty is limited.\n\n3. Data efficiency results (Figure 6, Table 2 bottom) demonstrate real value for resource-constrained scenarios common in medical applications.\n\n4. Algorithm 1, architecture details (Table 4), implementation notes, and comprehensive appendix facilitate reproducibility.", "weaknesses": "1. Limited theoretical justification:\n- When do we expect PSDNorm to outperform alternatives? No formal analysis provided.\n- Why is the geodesic update (Eq. 6) optimal? Other update schemes not explored.\n- Gaussian assumption (Eq. 2) likely violated for deep network features - no empirical validation.\n- No analysis of what happens when periodicity assumption fails.\n\n2. Narrow evaluation scope:\n- Only sleep staging evaluated - no other time series tasks\n- Both architectures are CNN-based - what about pure transformers or RNNs?\n- No evaluation on non-biomedical sequential data\n- Limits ability to assess general applicability\n\n\n3. This paper might overclaimed contributions, \"4× more data-efficient\" is misleading, should say \"more robust to data scarcity\"; \"State-of-the-art\" improvements are marginal (often <1%); Title suggests \"temporal normalization\" but method is specific to signals with spectral content\n\n4. Something missed\n- No visualization of learned barycenter PSDs - are they interpretable?\n- No analysis of what spectral patterns PSDNorm captures\n- Missing failure case analysis\n- No discussion of when PSDNorm might be worse than alternatives\n\n\n5. Computational: \n- Training time overhead 7-17% (Table 8) is non-negligible\n- Requires FFT operations which may not be hardware-optimized on all platforms\n- No analysis of memory overhead\n- The InstanceNorm timing anomaly (torch.compile issue) undermines the comparison", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WB5mqWJaHU", "reviewer": "Reviewer_6ujs", "summary": "The paper proposes PSDNorm (Power Spectral Density Normalization), a normalization layer that can be directly inserted into a network for time-series learning. Its core idea is to perform \"power spectrum alignment\" on feature maps in the time-frequency domain, allowing the model to learn domain-invariant temporal statistical structures. Experiments on a large-scale sleep dataset demonstrate the effectiveness of the framework.", "questions": "See in weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The structure of this paper is clear and reasonable.\n\n2. This paper is easy to follow and generally well presented.\n\n3. Extensive experiments have been conducted to evaluate the performance of the proposed framework.", "weaknesses": "1. The current related work section is more like a preliminary. There should be several subsections introducing works with similar topics, such as methods for training sleep data and normalization in time series.\n\n2. The feature map of the middle layer of the deep network may not meet the prior conditions in the article after nonlinearity and convolution, a Gaussian periodic signal, a covariance block diagonal, and diagonalized according to the Fourier basis. Should we consider the failure mode?\n\n3. The comparison is mainly limited to normalization layers, but the core idea of PSDNorm is based on Monge/OT/PSD alignment. It could be better to compare PSDNorm with existing TMA/STMA or general test-time adaptation methods. It could also compare PSDNorm with TMA, which uses PSD as preprocessing (or TMA embedded in the network).\n\n4. It could be better to consider the parameter sensitivity of the Welch estimator.\n\n5. It should show more results on different testing windows.\n\n6. What about the computational cost of this normalization?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)