SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Convolutional Monge Mapping Normalization for learning on sleep data
Theo Gnassounou
Université Paris-Saclay, Inria, CEA
Palaiseau 91120, France
theo.gnassounou@inria.fr
Rémi Flamary
IP Paris, CMAP, UMR 7641
remi.flamary@polytechnique.edu
Alexandre Gramfort*
alexandre.gramfort@inria.fr
In many machine learning applications on signals and biomedical data, especially electroencephalogram (EEG), one major challenge is the variability of the data across subjects, sessions, and hardware devices. In this work, we propose a new method called Convolutional Monge Mapping Normalization (CMMN), which consists in filtering the signals in order to adapt their power spectrum density (PSD) to a Wasserstein barycenter estimated on training data. CMMN relies on novel closed-form solutions for optimal transport mappings and barycenters and provides individual test time adaptation to new data without needing to retrain a prediction model. Numerical experiments on sleep EEG data show that CMMN leads to significant and consistent performance gains independent from the neural network architecture when adapting between subjects, sessions, and even datasets collected with different hardware. Notably our performance gain is on par with much more numerically intensive Domain Adaptation (DA) methods and can be used in conjunction with those for even better performances.
Data shift in biological signals Biological signals, such as electroencephalograms (EEG), often exhibit a significant degree of variability. This variability arises from various factors, including the recording setup (hardware specifications, number of electrodes), individual human subjects (variations in anatomies and brain activities), and the recording session itself (electrode impedance, positioning). In this paper, we focus on the problem of sleep staging which consists of measuring the activities of the brain and body during a session (here one session is done over a night of sleep) with electroencephalograms (EEG), electrooculograms (EOG), and electromyograms (EMG) [1] to classify the sleep stages. Depending on the dataset, the populations studied may vary from healthy cohorts to cohorts suffering from disease [2, 3, 4]. Different electrodes positioned in the front/back [3] or the side of the head [4] can be employed. Sleep staging is a perfect problem for studying the need to adapt to this variability that is also commonly denoted as data shift between domains (that can be datasets, subjects, or even sessions).
Normalization for data shift A traditional approach in machine learning to address data shifts between domains is to apply data normalization. Different approaches exist in the literature to
$^{*}$ A. Gramfort joined Meta and can be reached at agramfort@meta.com
normalize data. One can normalize the data per Session which allows to keep more within-session variability in the data [5]. If the variability during the session is too large, one can normalize independently each window of data (e.g., 30 s on sleep EEG) [6], denoted as Sample in the following. It is also possible not to normalize the data, and let a neural network learn to discard non-pertinent variabilities helped with batch normalization [7, 8]. More recently, a number of works have explored the possibility of learning a normalization layer that is domain specific in order to better adapt their specificities [9, 10, 11, 12, 13, 14, 15]. However, this line of work usually requires to have labeled data from all domains (subjects) for learning which might not be the case in practice when the objective is to automatically label a new domain without an expert.
Domain Adaptation (DA) for data shift Domain Adaptation is a field in machine learning that aims at adapting a predictor in the presence of data shift but in the absence of labeled data in the target domain. The goal of DA is to find an estimator using the labeled source domains which generalizes well on the shifted target domain [16]. In the context of biological signals, DA is especially relevant as it has the ability to fine-tune the predictor for each new target domain by leveraging unlabeled data. Modern DA methods, inspired by successes in computer vision, usually try to reduce the shift between the embeddings of domains learned by the feature extractor. To do that, a majority of the methods try to minimize the divergence between the features of the source and the target data. Several divergences can be considered for this task such as correlation distance [17], adversarial method [18], Maximum Mean Discrepancy (MMD) distance [19] or optimal transport [20, 21]. Another adaptation strategy consists in learning domain-specific batch normalization [9] in the embedding. Interestingly those methods share an objective with the normalization methods: they aim to reduce the shift between datasets. To this end, DA learns a complex invariant feature representation whereas normalization remains in the original data space. Finally, test-time DA aims at adapting a predictor to the target data without access to the source data [22], which might not be available in practice due to privacy concerns or the memory limit of devices.
Contributions In this work we propose a novel and efficient normalization approach that can compensate at test-time for the spectral variability of the domain signals. Our approach called Convolutional Monge Mapping Normalization (CMMN), illustrated in Figure 1, uses a novel closed-form to estimate a meaningful barycenter from the source domains. Then CMMN uses a closed-form solution for the optimal transport Monge mapping between Gaussian random signals to align the power spectrum density of each domain (source and target) to the barycenter. We emphasize that CMMN is, to the best of our knowledge, the first approach that can adapt to complex spectral shifts in the data without the need to access target datasets at training time or train a new estimator for each new domain (which are the limits of DA).
We first introduce in section 2 the problem of optimal transport (OT) between Gaussian distributions, and then propose a novel closed-form solution for the Wasserstein barycenter for stationary Gaussian random signals. We then use this result to propose our convolutional normalization procedure in section 3, where implementation details and related works are also discussed. Finally section 4 reports a number of numerical experiments on sleep EEG data, demonstrating the interest of CMMN for adapting to new subjects, sessions, and even datasets, but also study its interaction with DA methods.
Notations Vectors are denoted by small cap boldface letters (e.g.,  $\mathbf{x}$ ), and matrices are denoted by large cap boldface letters (e.g.,  $\mathbf{X}$ ). The element-wise product is denoted by  $\odot$ . The element-wise power of  $n$  is denoted by  $\cdot^{(n)}$ .  $[K]$  denotes the set  $\{1, \dots, K\}$ .  $|\cdot|$  denotes the absolute value. The discrete convolution operator between two signals is denoted as  $*$ . Any parameters written with a small  $k$  ( $e.g., \mathbf{X}_k$  or  $\mathbf{x}_i^k$ ) is related to the source domain  $k$  with  $1 \leq k \leq K$ . Any parameters written with a small  $t$  ( $e.g., \mathbf{X}_t$  or  $\mathbf{x}_i^t$ ) is related to the target domain.
2 Signal adaptation with Optimal Transport
In this section, we first provide a short introduction to optimal transport between Gaussian Distributions, and then discuss how those solutions can be computed efficiently on stationary Gaussian signals, exhibiting a new closed-form solution for Wasserstein barycenters.
Figure 1: Illustration of the CMMN approach. At train-time the Wasserstein barycenter is estimated from 3 subjects/domains. The model is learned on normalized data. At test time the same barycenter is used to normalize test data and predict.
2.1 Optimal Transport between Gaussian distributions
Monge mapping for Gaussian distributions Let two Gaussian distributions  $\mu_s = \mathcal{N}(\mathbf{m}_s, \boldsymbol{\Sigma}_s)$  and  $\mu_t = \mathcal{N}(\mathbf{m}_t, \boldsymbol{\Sigma}_t)$ , where  $\boldsymbol{\Sigma}_s$  and  $\boldsymbol{\Sigma}_t$  are symmetric positive definite covariances matrices. OT between Gaussian distributions is one of the rare cases where there exists a closed-form solution. The OT cost, also called the Bures-Wasserstein distance when using a quadratic ground metric, is [23, 24]
$$
\mathcal {W} _ {2} ^ {2} \left(\mu_ {s}, \mu_ {t}\right) = \| \mathbf {m} _ {s} - \mathbf {m} _ {t} \| _ {2} ^ {2} + \operatorname {T r} \left(\boldsymbol {\Sigma} _ {s} + \boldsymbol {\Sigma} _ {t} - 2 \left(\boldsymbol {\Sigma} _ {t} ^ {\frac {1}{2}} \boldsymbol {\Sigma} _ {s} \boldsymbol {\Sigma} _ {t} ^ {\frac {1}{2}}\right) ^ {\frac {1}{2}}\right), \tag {1}
$$
where the second term is called the Bures metric [25] between positive definite matrices. The OT mapping, also called Monge mapping, can be expressed as the following affine function:
$$
m (\mathbf {x}) = \mathbf {A} \left(\mathbf {x} - \mathbf {m} _ {s}\right) + \mathbf {m} _ {t}, \quad \text {w i t h} \quad \mathbf {A} = \boldsymbol {\Sigma} _ {s} ^ {- \frac {1}{2}} \left(\boldsymbol {\Sigma} _ {s} ^ {\frac {1}{2}} \boldsymbol {\Sigma} _ {t} \boldsymbol {\Sigma} _ {s} ^ {\frac {1}{2}}\right) ^ {\frac {1}{2}} \boldsymbol {\Sigma} _ {s} ^ {- \frac {1}{2}} = \mathbf {A} ^ {\top}. \tag {2}
$$
In practical applications, one can estimate empirically the means and covariances of the two distributions and plug them into the equations above. Interestingly, in this case, the concentration of the estimators has been shown to be in  $O(N^{-1/2})$ , where  $N$  is the number of samples, for the divergence estimation [26] and for the mapping estimation [27]. This is particularly interesting because optimal transport in the general case is known to be very sensitive to the curse of dimensionality with usual concentrations in  $O(N^{-1/D})$  where  $D$  is the dimensionality of the data [28].
Wasserstein barycenter between Gaussian distributions The Wasserstein barycenter that searches for an average distribution can also be estimated between multiple Gaussian distributions  $\mu_{k}$ . This barycenter  $\bar{\mu}$  is expressed as
$$
\bar {\mu} = \arg \min  _ {\mu} \frac {1}{K} \sum_ {k = 1} ^ {K} \mathcal {W} _ {2} ^ {2} (\mu , \mu_ {k}). \tag {3}
$$
Interestingly, the barycenter is still a Gaussian distribution  $\bar{\mu} = \mathcal{N}(\bar{\mathbf{m}},\bar{\Sigma})$  [29]. Its mean  $\bar{\mathbf{m}} = \frac{1}{K}\sum_{k}\mathbf{m}_{k}$  can be computed as an average of the means of the Gaussians, yet there is no closed-form for computing the covariance  $\bar{\Sigma}$ . In practical applications, practitioners often use the following optimality condition from [29]
$$
\bar {\boldsymbol {\Sigma}} = \frac {1}{K} \sum_ {k = 1} ^ {K} \left(\bar {\boldsymbol {\Sigma}} ^ {\frac {1}{2}} \boldsymbol {\Sigma} _ {k} \bar {\boldsymbol {\Sigma}} ^ {\frac {1}{2}}\right) ^ {\frac {1}{2}}, \tag {4}
$$
in a fixed point algorithm that consists in updating the covariance [30] using equation (4) above until convergence. Similarly to the distance estimation and mapping estimation, statistical estimation of the barycenter from sampled distribution has been shown to have a concentration in  $O(N^{-1/2})$  [31].
2.2 Optimal Transport between Gaussian stationary signals
We now discuss the special case of OT between stationary Gaussian random signals. In this case, the covariance matrices are Toeplitz matrices. A classical assumption in signal processing is that for a long enough signal, one can assume that the signal is periodic, and therefore the covariance matrix is a Toeplitz circulant matrix. The circulant matrix can be diagonalized by the Discrete Fourier Transform (DFT)  $\boldsymbol{\Sigma} = \mathbf{F}\mathrm{diag}(\mathbf{p})\mathbf{F}^{*}$ , with  $\mathbf{F}$  and  $\mathbf{F}^*$  the Fourier transform operator and its inverse, and  $\mathbf{p}$  the Power Spectral Density (PSD) of the signal.
Welch PSD estimation As discussed above, there is a direct relation between the correlation matrix of a signal and its PSD. In practice, one has access to a matrix  $\mathbf{X} \in \mathbb{R}^{N \times T}$  containing  $N$  signals of length  $T$  that are samples of Gaussian random signals. In this case, the PSD  $\mathbf{p}$  of one random signal can be estimated using the Welch periodogram method [32] with  $\hat{\mathbf{p}} = \frac{1}{N} \sum_{i=1}^{N} |\mathbf{F} \mathbf{x}_i|^{\odot 2}$  where  $|\cdot|$  is the element-wise magnitude of the complex vector.
Monge mapping between two Gaussian signals The optimal transport mapping between Gaussian random signals can be computed from (2) and simplified by using the Fourier-based eigenfactorization of the covariance matrices. The mapping between two stationary Gaussian signals of PSD respectively  $\mathbf{p}_s$  and  $\mathbf{p}_t$  can be expressed with the following convolution [27]:
$$
m (\mathbf {x}) = \mathbf {h} * \mathbf {x}, \quad \text {w i t h} \quad \mathbf {h} = \mathbf {F} ^ {*} \left(\mathbf {p} _ {t} ^ {\odot \frac {1}{2}} \odot \mathbf {p} _ {s} ^ {\odot - \frac {1}{2}}\right). \tag {5}
$$
Note that the bias terms  $\mathbf{m}_s$  and  $\mathbf{m}_t$  do not appear above because one can suppose in practice that the signals are centered (or have been high-pass filtered to be centered), which means that the processes are zero-mean. The Monge mapping is a convolution with a filter  $\mathbf{h}$  that can be efficiently estimated from the PSD of the two signals. It was suggested in [27] as a Domain Adaptation method to compensate for convolutional shifts between datasets. For instance, it enables compensation for variations such as differing impedances, which can be physically modeled as convolutions [33]. Nevertheless, this paper focused on theoretical results, and no evaluation on real signals is reported. Moreover, the method proposed in [27] cannot be used between multiple domains (as explored here). This is why in the following we propose a novel closed-form for estimating a barycenter of Gaussian signals that we will use for the normalization of our CMMN method.
Wasserstein barycenter between Gaussian signals As discussed above, there is no known closed-form solution for a Wasserstein barycenter between Gaussian distributions. Nevertheless, in the case of stationary Gaussian signals, one can exploit the structure of the covariances to derive a closed-form solution that we propose below.
Lemma 1 Consider  $K$  centered stationary Gaussian signals of PSD  $\mathbf{p}_k$  with  $k \in [K]$ , the Wasserstein barycenter of the  $K$  signals is a centered stationary Gaussian signal of PSD  $\bar{\mathbf{p}}$  with:
$$
\bar {\mathbf {p}} = \left(\frac {1}{K} \sum_ {k = 1} ^ {K} \mathbf {p} _ {k} ^ {\odot \frac {1}{2}}\right) ^ {\odot 2}. \tag {6}
$$
Proof. The proof is a direct application of the optimality condition (4) of the barycenter. The factorized covariances in (4), the matrix square root and the inverse can be simplified as element-wise square root and inverse, recovering equation (11). We provide a detailed proof in the appendix.
The closed-form solution is notable for several reasons. First, it is a novel closed-form solution that avoids the need for iterative fixed-point algorithms and costly computations of matrix square roots. Secondly, the utilization of the Wasserstein space introduces an alternative approach to the conventional  $\ell_2$  averaging of Power Spectral Density (PSD). This approach involves employing the square root, like the Hellinger distance [34], potentially enhancing robustness to outliers. Note that while other estimators for PSD averaging could be used this choice is motivated here by the fact that we use OT mappings and that the barycenter above is optimal w.r.t. those OT mappings.
3 Multi-source DA with Convolutional Monge Mapping Normalization
We now introduce the core contribution of the paper, that is an efficient method that allows to adapt to the specificities of multiple domains and train a predictor that can generalize to new domains at test time without the need to train a new model. We recall here that we have access to  $K$  labeled source domains  $(\mathbf{X}_k,\mathbf{y}_k)_k$ . We assume that each domain contains  $N_{k}$  centered signals  $\mathbf{x}_i^k$  of size  $T$ .
CMMN at train time The proposed approach, illustrated in Figure 1 and detailed in Algorithm 1, consists of the following steps:
1. Compute the PSD  $\hat{\mathbf{p}}_k$  for each source domain and use them to estimate the barycenter  $\hat{\mathbf{p}}$  with (11).
2. Compute the convolutional mapping  $\mathbf{h}_k$  (5) between each source domain and the barycenter  $\bar{\mathbf{p}}$ .
3. Train a predictor on the normalized source data using the mappings  $\mathbf{h}_k$ :
$$
\min  _ {f} \quad \sum_ {k = 1} ^ {K} \sum_ {i = 1} ^ {N _ {k}} L \left(y _ {i} ^ {k}, f \left(\mathbf {h} _ {k} * \mathbf {x} _ {i} ^ {k}\right)\right). \tag {7}
$$
In order to keep notations simple, we consider here the case for a single sensor, but CMMN can be extended to multi-sensor data by computing independently the Monge mapping for each sensor. Note that steps 1 and 2 can be seen as a pre-processing and are independent of the training of the final predictor, so CMMN can be used as pre-processing on any already existing learning framework.
CMMN at test time At test time, one has access to a new unlabeled target domain  $(\mathbf{X}_t)$  and the procedure, detailed in Algorithm 2, is very simple. One can estimate the PSD  $\hat{\mathbf{p}}_t$  from the target domain unlabeled data and compute the mapping  $\mathbf{h}_t$  to the barycenter  $\bar{\mathbf{p}}$ . Then the final predictor for the new domain is  $f^t (\mathbf{x}_t) = f(\mathbf{h}_t*\mathbf{x}_t)$ , that is the composition of the domain-specific mapping to the barycenter of the training data, and the already trained predictor  $f$ . This is a very efficient test-time adaptation approach that only requires an estimation of the target domain PSD that can be done with few unlabeled target samples. Yet, it allows for a final predictor to adapt to the spectral specificities of new domains thanks to the convolutional Monge normalization.
Numerical complexity and filter size The numerical complexity of the method is very low as it only requires to compute the PSD of the domains and the barycenter in  $O\left(\sum_{k}^{K}N_{k}T\log (T)\right)$ . It is also im
Algorithm 1: Train-Time CMMN
Input:  $f,F,\{\mathbf{X}_k\} _k^K$    
for  $k = 1\to K$  do  $\hat{\mathbf{p}}_k\gets$  Welch PSD estimation of  $\mathbf{X}_k$    
end   
 $\bar{\mathbf{p}}\gets$  Compute barycenter with (11)   
for  $k = 1\rightarrow K$  do  $\mid \mathbf{h}_k\gets$  Compute mapping from (5)   
end   
 $\hat{f}\gets$  Train on adapted data with (7)   
return  $\hat{f},\bar{\mathbf{p}}$
Algorithm 2: Test-Time CMMN
Input:  $\hat{f},\bar{\mathbf{p}},\mathbf{X}_t$ $\hat{\mathbf{p}}_t\gets$  Welch PSD estimation of  $\mathbf{X}_t$ $\mathbf{h}_t\gets$  compute mapping from 5   
return  $\hat{\mathbf{y}}_t = \hat{f} (\mathbf{h}_k^t *\mathbf{X}_k^t)$
important to note that in practice, the method allows for a size of normalization filter  $F$  that is different (smaller) than the size  $T$  of the signals. This consists in practice in estimating the PSD using Welch periodogram on signal windows of size  $F \leq T$  that can be extracted from the raw signal or from already extracted fixed sized source training samples. Indeed, if we use  $F = T$  then the estimated average PSD can be perfectly adapted by the mapping, yet using many parameters can lead to overfitting which can be limited using a smaller filter size  $F \leq T$ . In fact, it is interesting to note that the special case  $F = 1$  boils down to a scaling of the whole signal similar to what is done with a simple z-score operation. This means that the filter size  $F$  is an hyperparameter that can be tuned on the data. From an implementation point of view, one can use the Fast Fourier Transform (FFT) to compute the convolution (for large filters) or the direct convolution, which both have very efficient implementation on modern hardware (CPU and GPU).
Related Works CMMN is a computationally efficient approach that benefits from a wide array of recent results in optimal transport and domain adaptation. The idea of using optimal transport to adapt distributions was first proposed in [35]. The idea to compute a Wasserstein barycenter of distributions
from multiple domains and use it for adaptation was introduced in [21]. Both of those approaches have shown encouraging performances but were strongly limited by the numerical complexity of solving the OT problems (mapping and barycenters) on large datasets  $(O(\sum_{k=1}^{K} N_k^3 \log(N_k))$  or at least quadratic in  $N_k$  for entropic OT). CMMN does not suffer from this limitation as it relies on both Gaussian and stationary signals assumptions that allow to estimate all the parameters for a complexity  $O(\sum_{k=1}^{K} N_k \log(N_k))$  linear with the number of samples  $N_k$ , and quasi-linear in dimensionality of the signals  $T$ . The use of Gaussian modeling and convolutional Monge mapping for DA was first proposed in [27] but the paper was mostly theoretical and only focused on the standard 2-domain DA problem whereas CMMN handles multi-source and provides test-time adaptation.
Finally, CMMN also bears resemblance with the convolutional normalization layer proposed in [11] that also uses the FFT for fast implementation. Yet, it needs to be trained using labeled source and target data, which prevents its use on DA at test time on new unseen domains.
In this section, we evaluate CMMN on the clinical application of sleep stage classification from EEG signals with [6, 36]. In the following one domain can be a session of one subject (i.e., in the domain-specific experiment in section 4.2) or one subject (i.e., all other experiments). We first compare CMMN to classical normalization methods and to subject-specific normalizations. Next, we illustrate the behavior of CMMN when used with different neural network architectures, and study the effect of CMMN on low-performing subjects. Finally, we study the use of CMMN in conjunction with domain adaptation approaches. In order to promote research reproducibility, code is available on github  $^{2}$ , and the datasets used are publicly available.
Sleep staging datasets We use three publicly available datasets: Physionet (a.k.a SleepEDF) [3], SHHS [4, 37] and MASS [2]. On all datasets, we want to perform sleep staging from 2-channels EEG signals. The considered EEG channels are 2 bipolar channels, Fpz-Cz and Pz-Cz that have been known to provide discriminant information for sleep staging. Note that those channels were not available on the SHHS dataset, and we used the C3-A2 and C4-A1 instead. This brings another level of variability in the data. More details about the datasets are available as supplementary material.
Pre-processing For all experiments, we keep 60 subjects of each dataset and the two EEG channels. The same pre-processing is applied to all sensors. First, the recordings are low-pass filtered with a  $30\mathrm{Hz}$  cutoff frequency, then the signals are resampled to  $100\mathrm{Hz}$ . Then we extract 30 s epochs having a unique class. This pre-processing is common in the field [6, 38]. All the data extraction and the pre-processing steps are done with MNE-BIDS [39] and MNE-Python [40].
Neural network architectures and training Many neural network architectures dedicated to sleep staging have been proposed [8, 7, 41]. In the following, we choose to focus on two architectures: Chambon [6] and DeepSleepNet [7]. For both architectures, we use the implementation from braindecode [42]. Chambon is an end-to-end neural network proposed to deal with multivariate time series and is composed of two convolutional layers with non-linear activation functions. DeepSleepNet is a more complex model with convolutional layers, non-linear activation functions, and a Bi-LSTM to model temporal sequences.