SPD domain-specific batch normalization to crack interpretable unsupervised domain adaptation in EEG
Abstract
Electroencephalography (EEG) provides access to neuronal dynamics noninvasively with millisecond resolution, rendering it a viable method in neuroscience and healthcare.
However, its utility is limited as current EEG technology does not generalize well across domains (i.e., sessions and subjects) without expensive supervised re-calibration.
Contemporary methods cast this transfer learning (TL) problem as a multi-source/-target unsupervised domain adaptation (UDA) problem and address it with deep learning or shallow, Riemannian geometry aware alignment methods.
Both directions have, so far, failed to consistently close the performance gap to state-of-the-art domain-specific methods based on tangent space mapping (TSM) on the symmetric, positive definite (SPD) manifold.
Here, we propose a machine learning framework that enables, for the first time, learning domain-invariant TSM models in an end-to-end fashion.
To achieve this, we propose a new building block for geometric deep learning, which we denote SPD domain-specific momentum batch normalization (SPDDSMBN).
A SPDDSMBN layer can transform domain-specific SPD inputs into domain-invariant SPD outputs, and can be readily applied to multi-source/-target and online UDA scenarios.
In extensive experiments with 6 diverse EEG brain-computer interface (BCI) datasets, we obtain state-of-the-art performance in inter-session and -subject TL with a simple, intrinsically interpretable network architecture, which we denote TSMNet.
Introduction
Electroencephalography (EEG) measures multi-channel electric brain activity from the human scalp with millisecond precision [1].
Transient modulations in the rhythmic brain activity can reveal cognitive processes [2], affective states [3] and a person’s health status [4].
Unfortunately, these modulations exhibit low signal-to-noise ratio (SNR), domain shifts (i.e., changes in the data distribution) and have low specificity, rendering statistical learning a challenging task - particularly in the context of brain-computer interfaces (BCI) [5] where the goal is to predict a target from a short segment of multi-channel EEG data in real-time.
Under domain shifts, domain adaptation (DA), defined as learning a model from a source domain that performs well on a related target domain, offers principled statistical learning approaches with theoretical guarantees [6, 7].
DA in the BCI field mainly distinguishes inter-session and -subject transfer learning (TL) [8].
In inter-session TL, domain shifts, are expected across sessions mainly due to mental drifts (low specificity) as well as differences in the relative positioning of the electrodes and their impedances.
Inter-subject TL is more difficult, as domain shifts are additionally driven by structural and functional differences in brain networks as well as variations in the performed task [9].
These domain shifts are traditionally circumvented by recording labeled calibration data and fitting domain-specific models [10, 11].
As recording calibration data is costly, models that are robust to scarce data with low SNR perform well in practice.
Currently, tangent space mapping (TSM) models [12, 13] operating with symmetric, positive definite (SPD) covariance descriptors of preprocessed data are considered state-of-the-art (SoA) [10, 14, 15].
They are well suited for EEG data as they exhibit invariances to linear mixing of latent sources [16], and are consistent [13] and intrinsically interpretable [17] estimators for generative models that encode label information with a log-linear relationship in source power modulations.
Competitive, supervised calibration-free methods are one of the long-lasting grand challenges in
EEG neurotechnology research [5, 10, 18, 19, 15].
Among the applied transfer learning techniques, including multi-task learning [20] and domain-invariant learning [21–23], unsupervised domain adaptation (UDA) [24] is considered as key to overcome this challenge [10, 19].
Contemporary methods cast the problem as a multi source and target UDA problem and address it with deep learning [25–28] or shallow, Riemannian geometry aware alignment methods [29–32].
Successful methods must cope with notoriously small and heterogeneous datasets (i.e., dozens of domains with a few dozens observations per domain and class).
In a recent, relatively large scale inter-subject and -dataset
TL competition with few labeled examples per target domain [19], deep learning approaches that aligned the first and second order statistics either in input [33, 27] or latent space [34] obtained the highest scores.
Whereas, in a pure UDA competition [15] with a smaller dataset, Riemannian geometry aware approaches dominated.
With the increasing popularity of geometric deep learning [35], Ju and Guan [36] proposed an architecture based on SPD neural networks [37] to align SPD features in latent space and attained SoA scores.
Despite the tremendous advances in recent years, the field still lacks methods that can consistently close the performance gap to state-of-the-art domain-specific methods.
To close this gap, we propose a machine learning framework around domain-specific batch normalization on the SPD manifold (Figure 1).
The proposed framework is used to implement domain-specific
TSM (Figure 1a), which requires tracking the domains’ Fréchet means in latent space as they are changing during training a typical TSM model in an end-to-end fashion (Figure 1b).
After reviewing some preliminaries in section 2, we extend momentum batch normalization (MBN) [32] to
SPDMBN that controls the Fréchet mean and variance of SPD data in section 3.
In a theoretical analysis, we show under reasonable assumptions that SPDMBN can track and converge to the data’s true Fréchet mean, enabling, for the first time, end-to-end learning of feature extractors, TSM and tangent space classifiers.
Building upon this insight, we combine SPDMBN with domain-specific batch normalization (DSBN) [38] to form SPDDSMBN (Figure 1a).
A SPDDSMBN layer can transform domain-specific SPD inputs into domain-invariant SPD outputs (Figure 1c).
Like DSBN,
SPDDSMBN easily extends to multi-source, multi-target and online UDA scenarios.
In section 4, we briefly review the generative model of EEG, before the proposed methods are combined in a simple, intrinsically interpretable network architecture, denoted TSMNet (Figure 1b).
We obtain state-of-theart performance in inter-session and -subject UDA on small and large scale EEG BCI datasets, and show in an ablation study that the performance increase is primarily driven by performing DSBN on the SPD manifold.
Preliminaries
Multi-source multi-target unsupervised domain adaptation Let X denote the space of input features, Y a label space, and I d ⊂ N an index set that contains unique domain identifiers.
In the multi-source, multi-target unsupervised domain adaptation scenario considered here, we are given a set T source =
{T i | i ∈ I sourced ⊂ I d } with |I sourced | = N domains.
Each domain T i = { ( X ij , y ij ) } Mj =1 ∼ P iXY contains M observations of feature ( X ∈ X ) and label ( y ∈ Y ) tuples sampled from a joint distribution P iXY .
1 While the joint distributions can be different (but related) across domains, we assume that the class priors are the same (i.e., P iY =
P Y ).
The goal is to learn a predictive function h :
X × I d
→ Y that, once fitted to T source , can generalize to unseen target domains T target =
{T l | l ∈ I targetd ⊂ I
d , I targetd ∩ I sourced = ∅} merely based on unsupervised adaptation of h to each target domain T l once its label l and features { X lj } Mj =1 ∼ P lX are revealed.
Riemannian geometry on S + D
We start with recalling notions of geometry on the space of real
D × D symmetric positive definite (SPD) matrices S
+ D = { Z ∈ R D
× D : Z T =
Z , Z ≻ 0 } .
The space S + D forms a cone shaped Riemannian manifold in R D × D [39].
A Riemannian manifold M is a smooth manifold equipped with an inner product on the tangent space T Z M at each point Z ∈ M .
Tangent spaces have Euclidean structure with easy to compute distances T Z M × T Z M → R + which locally approximate Riemannian distances on M induced by an inner product [40].
Logarithmic
Log Z : M → T Z M and exponential Exp Z : T Z M → M mappings project points to and from tangent spaces.
Using the inner product ⟨ S 1 , S 2 ⟩ Z =
Tr( Z − 1 S 1 Z − 1 S 2 ) for points S 1 , S 2 in the tangent space T Z S + D
(i.e., the space of real symmetric D × D matrices) results in a globally defined affine invariant
Riemannian metric on S + D [41, 39], which can be computed in closed form:
[Equation]
where Z 1 and Z 2 are two SPD matrices, log( · ) denotes the matrix logarithm 2 , || · || F the Frobenius norm, and Tr( · ) in the inner product the trace operator.
Due to affine invariance, we have δ ( AZ 1 A T , AZ 2 A T ) =
δ ( Z 1 , Z 2 ) for any invertible D × D transformation matrix A .
The exponential and logarithmic mapping are also globally defined in closed form as
[Equation]
For a set of SPD points Z = { Z j ∈ S + D } j ≤ M , we will use the notion of Fréchet mean G Z ∈ S + D and Fréchet variance ν 2 Z ∈ R + .
The Fréchet mean is defined as the minimizer of the average squared distances
[Equation]
For M = 2 , there is a closed form solution expressed as
[Equation]
with weight γ = 0 . 5 .
Choosing γ ∈ [0 , 1] computes weighted means along the geodesic (i.e., the shortest curve) that connects both points.
For M > 2 , (4) can be solved using the Karcher flow algorithm [42], which iterates between projecting the data to the tangent space (2) at the current estimate, arithmetic averaging, and projecting the result back (3) to obtain a new estimate.
The
Fréchet variance ν 2 Z is defined as the attained value at the minimizer G Z :
[Equation]
To shift a set of tangent space points to vary around a parametrized mean G ϕ , parallel transport on
S + D can be used [43]:
[Equation]
While, parallel transport is generally defined for tangent space vectors S [40], on S + D the same operations also apply directly to points on the manifold (i.e., Z ∈ Z ) [31, 44].
[Table]
Domain-specific batch normalization on S + D
In this section, we review relevant batch normalization (BN) [45] variants with a focus on S + D .
We then present SPDMBN and show in a theoretical analysis that the running estimate converges to the true Fréchet mean under reasonable assumptions.
At last, we combine the idea of domain-specific batch normalization (DSBN) [38] with SPDMBN to form a SPDDSMBN layer.
Table 1 provides a brief overview of related and proposed methods.
Batch normalization Batch normalization (BN) [45] is a widely used training technqiue in deep learning as BN layers speed up convergence and improve generalization via smoothing of the engery landscape [47, 32].
A standard BN layer applies slightly different transformations during training and testing to independent and identically distributed (iid) observations x j ∈ R d within the k -th minibatch B k of size M drawn from a dataset T .
During training, the data are normalized using the batch mean b k and variance s 2 k , and then scaled and shifted to have a parametrized mean g ϕ and variance σ 2 ϕ .
Internally, the layer updates running estimates of the dataset’s statistics ( g k , σ 2 k ) during each training step k ; the updates are computed via exponential smoothing with momentum parameter γ .
During testing, the running estimates are used.
Using batch statistics to normalize data during training rather than running estimates introduces noise whose level depends on the batch size [32]; smaller batch sizes raise the noise level.
The introduced noise regularizes the training process, which can help to escape poor local minima in initial learning but also lead to underfitting.
Momentum BN (MBN) [32] allows small batch sizes while avoiding underfitting.
Like batch renormalization [48], MBN uses running estimates during training and testing.
The key difference is that MBN keeps two sets of running statistics; one for training and one for testing.
The latter are updated conventionally, while the former are updated with momentum parameter γ train ( k ) that decays over training steps k .
MBN can, therefore, quickly escape poor local minima during initial learning and avoid underfitting at later stages [32].
Batch normalization on S + D
It is intractable to compute the Fréchet mean G B k for each minibatch
B k = { Z j ∈ S + D } Mj =1 , as there is no efficient algorithm to solve (4).
Brooks et al.
[44] proposed
Riemannian Batch Normalization (RBN) as a tractable approximation.
RBN approximateley solves (4) by aborting the iterative Karcher flow algorithm after one iteration.
To transform Z j ∈ B k with estimated mean B k to vary around G ϕ , parallel transport (7) is used.
The RBN input output transformation is then expressed as
[Equation]
Using (5), the running estimate of the dataset’s Fréchet mean can be updated in closed form
[Equation]
In [46] we proposed an extension to RBN, denoted SPD batch renormalization (SPDBN) that controls both Fréchet mean and variance.
Like batch renormalization [48], SPDBN uses running estimates G k and ν 2 k during training and testing.
To transform Z j ∈ B k to vary around G ϕ with variance ν 2 ϕ , each observation is first transported to vary around the identity matrix I , rescaled via computing matrix powers and finally transported to vary around G ϕ .
The sequence of operations can be expressed as
[Equation]
The standard backpropagation framework with extensions for structured matrices [49] and manifoldconstrained gradients [40] can be used to propagate gradients through RBN and SPDBN layers and learn the parameters ( G ϕ , ν ϕ ).
Input : batch B training , running mean ¯ G k − 1 , ¯ G {+
+} =
I and variance ¯ ν 2 k − 1 , ¯ ν 20 = training, running mean ˜
G k − 1 , ˜ G
= I and variance ˜
ν
2 k − 1 , ˜ ν 20 = 1 for parameters ( G ϕ , ν ϕ ) , and momentum for training and testing γ train ( k ) , γ
Output: normalized batch } if training then B k
= karcher_flow
B
=
); // approx.
solve problem (4) ¯
G k = ¯ G k − 1 # γ train ( k ) B k ; // update running stats for training ¯ ν 2 k
= (1 − γ train ( k ))
¯ ν 2 k − 1 + γ train ( k )Var B k
( ¯ G k ) ˜
G k = ˜
G k − 1 # γ B k ; // update running stats for testing ˜ ν
2 k = ( − γ )˜ ν
2 k − 1 + γ Var B k
( ˜ G k ) end ( =
ν if else ν
˜
Z j = Γ I
→
G ϕ ◦ Γ G k
→ I ( j ) νϕνk + ε ; // use ( to whiten, rescale and rebias
Momentum batch normalization on S + D
SPDBN [46] suffers from the same limitations as batch renormalization [48].
Consequently, we propose to extend MBN [32] to S + D .
We list the pseudocode of our proposed extension, which we denote SPDMBN, in algorithm 1.
SPDMBN uses approximations of batch-specific Fréchet means to update two sets of running estimates of the dataset’s Fréchet mean.
As MBN [32], we decay γ train ( k ) with a clamped exponential decay schedule
[Equation]
where K defines the training step at which γ min ∈ [0 , 1] should be attained.
The running mean in SPDMBN converges to the Fréchet mean Here, we consider models that apply a SPDMBN layer to latent representations generated by a feature extractor f
θ : X → S + D with learnable parameters θ .
We define a dataset that contains the latent representations generated with feature set θ k as Z θ
k
=
{ f θ k ( x ) | x ∈ T } , and a minibatch of M iid samples at training step k as B k .
We denote the Fréchet mean of Z θ k as G θ k , the estimated Fréchet mean, defined in (9), as G k , and the estimated batch mean as B k .
Since the batches are drawn randomly, we consider the batch and running means as random variables.
We assume that the variance Var θ k
( B k − 1 ) =
E B k − 1 { δ 2 ( B k − 1 , G θ k ) } of the previous batch mean B k − 1 with respect to the current Fréchet mean G θ k is bounded by the current variance Var θ k
( B k ) and the norm of the difference in the parameters
[Equation]
[Equation]
That is, across training steps k the parameter updates are required to change the first and second order moments of the distribution of Z θ k gradually so that the expected distance between G θ k and the previous batch mean B k − 1 is bounded.
We conjecture that this is the case for feature extractors f θ that are smooth in the parameters and small learning rates, but leave the proof for future work.
Proposition 1 (Error bound for G k ) .
Consider the setting defined above, and assumption (12) holds true.
Then, the variance of the running mean Var θ k ( G k ) is bounded by
[Equation]
(13) over training steps k if
[Equation]
holds true.
The proof is provided in appendix A.1 of the supplementary material and relies on the proof of the geometric law of large numbers [50].
Proposition 1 states that if (12) and (14) are met, the expected distance between the true Fréchet mean and the running mean is less or equal to the one of the batch mean.
Consequently, the introduced noise level of SPDBN (equation 10) and SPDMBN (algorithm 1), which use G k to normalize batches during training, is smaller or equal to RBN (equation 8), which uses B k .
Since γ controls the adaptation speed of G k , proposition 1 also states that if γ converges to zero (=no adaptation), the parameter updates are required to converge to zero as well (=no learning).
Hence, for a fixed γ ∈ (0 , 1) , as in the case of SPDBN (equation 10), proposition 1 is fulfilled, if the learning rate for the parameters θ is chosen sufficiently small.
This can substantially slow down initial learning for standard choices of γ (e.g., 0.1 or 0.01).
As a remedy, SPDMBN (algorithm 1) uses an adaptive momentum parameter, which allows larger parameter updates during initial training steps.
If we consider a late stage of learning, and in particular assume that after a certain number of iterations κ the parameters stay in a small ball with radius ρ around θ ∗ (i.e., || θ k − θ ∗ || ≤ ρ ∀ k > κ ) and the feature extractor is L -smooth in the parameters (i.e., δ ( f θ ( x ) , f ˜ θ ( x )) ≤ L || θ − ˜ θ
|| ∀ x ∈ T , ∀ θ, ˜ θ )
then the distances are bounded δ ( f θ k ( x ) , f θ ∗ ( x )) ≤ ρL .
Remark 1 (Convergence of G k for SPDMBN) .
If ρL is neglibile compared to the dataset’s variance, then the Fréchet mean and variance can be considered fixed, and the theorem of large numbers on
S + D [50] applies directly.
That is, if the momentum parameter is decayed exponentially ∀ k > κ the running mean G k converges to the Fréchet mean G θ ∗ in probability as k → ∞ .
Taken together, Proposition 1 and Remark 1 provide guidelines to update G k in SPDMBN so that the introduced estimation error is bounded during initial fast learning (large γ ) and decays towards zero in late learning (small γ ).
SPDMBN to learn tangent space mapping at Fréchet means Typical TSM models for classification [12] and regression [13] first use (2) to project Z ∈ T ⊂ S + D to the tangent space at the Fréchet mean G T , then use (7) to transport the result to vary around I , and finally extract elements in the upper triangular part 3 to reduce feature redundancy.
The invertible mapping P G T : S + D → R D ( D +1) /is expressed as:
[Equation]
We propose to use a SPDMBN layer followed by a LogEig layer [37] to compute a similar mapping m ϕ (Figure 1a).
A LogEig layer simply computes the matrix logarithm and vectorizes the result so that the norm is preserved.
If the parametrized mean of SPDMBN is fixed to the identify matrix ( G ϕ = I ), the composition computes
[Equation]
where ( G k , ν 2 k ) are the estimated Fréchet mean and variance of the dataset T at training step k , and ϕ = { ν ϕ } the learnable parameters.
According to remark 1 G k converges to G T and, in turn, m ϕ to a scaled version of P G T , since upper is linear.
The mapping m ϕ offers several advantageous properties.
First, the features are projected to a
Euclidean vector space where standard layers can be applied and distances are cheap to compute.
Second, distances between the projected features locally approximate δ and, therefore, inherit its invariance properties (e.g., affine mixing) [41].
This improves upon a LogEig layer [37] which projects features to the tangent space at the identity matrix.
As a result, distances between LogEig projected features correspond to distances measured with the log-Euclidean Riemannian metric (LERM) [51] which is not invariant to affine mixing.
Third, controlling the Fréchet variance in (16) empirically speeds up learning and improves generalization [46].
Domain-specific batch normalization on S + D
Considering a multi-source UDA scenario, Chang et al.
[38] proposed a domain-specific BN (DSBN) layer which simply keeps multiple parallel BN layers and distributes observations according to the associated domains.
Formally, we consider minibatches B k that form the union of N B k ≤ |I
d | domain-specific minibatches B ik drawn from distinct domains i ∈ I B k ⊆
I d .
As before, each B i k contains j = 1 , ..., M/N B k iid observations x j .
A DSBN layer mapping R d × I d → R d can then be expressed as
[Equation]
In practice, the batch size M is typically fixed.
The particular choice is influenced by resource availability and the desired noise level introduced by minibatch based stochastic gradient descent.
A drawback of DSBN is that for a fixed batch size M and an increasing number of source domains
N B k , the effective batch size declines for the BN layers within DSBN.
Since small batch sizes increase the noise level introduced by BN, increasing the number of domains per batch can lead to underfitting [32].
To alleviate this effect, we use the previously introduced SPDMBN layer.
The proposed domain-specific BN layer on S + D is then formally defined as
[Equation]
The layer can be readily adapted to new domains, as new SPDMBN layers can be added on the fly.
If the entire data of a domain becomes available, the domain-specific Fréchet mean and variance can be estimated by solving (4), otherwise, the update rules in algorithm 1 can be used.
SPDDSMBN to crack interpretable multi-source/-target UDA for EEG data
With SPDDSMBN introduced in the previous section, we focus on a specific application domain, namely, multi-source/-target UDA for EEG-based BCIs and propose an intrinsically interpretable architecture which we denote TSMNet.
Generative model of EEG EEG signals x ( t ) ∈ R P capture voltage fluctuations on P channels.
An EEG record (=domain) is uniquely identified by a subject and session identifier.
After standard pre-processing steps, each domain i contains j = 1 , ..., M labeled observations with features
X ij ∈ X
⊂ R P
×
T where T is the number of temporal samples.
Due to linearity of Maxwell’s equations and Ohmic conductivity of tissue layers in the frequency ranges relevant for EEG [52], a domain-specific linear instantaneous mixture of sources model is a valid generative model:
[Equation]
(19) where S ij ∈ R Q × T represents the activity of Q latent sources, A i ∈ R P × Q a domain-specific mixing matrix and N ij ∈ R P × T additive noise.
Both A i and S ij are unknown which demands making assumptions on A i (e.g., anatomical prior knowledge [53]) and/or S ij (e.g., statistical independence [54]) to extract interesting sources.
Interpretable multi-source/-target UDA for EEG data As label information is available for the source domains, our goal is to identify discriminative oscillatory sources shared across domains.
Our approach relies on TSM models with linear classifiers [12], as they are consistent [13] and intrinsically interpretable [17] estimators for generative models with log-linear relationships between the target y ij and variance Var { s ( k ) ij ( t ) } of k = 1 , ..., K ≤ Q discriminative sources:
[Equation]
where b k ∈ R summarizes the coupling between the target y ij and the variance of the encoding source, and ε ij additive noise.
In [17] we showed that the encoding sources’ coupling and their patterns 4 (columns of A i ) can be recovered via solving a generalized eigenvalue problem between the Fréchet mean G T i and classifier patterns [55] that were back projected to S + D with P − 1 G T i .
The resulting eigenvectors are the patterns and the eigenvalues λ k reflect the relative source contribution c k
[Equation]
To benefit from the intrinsic interpretability of TSM models, we constrain our hypothesis class
H to functions h :
X × I d
→ Y that can be decomposed into a composition of a shared linear feature extractor with covariance pooling f θ : X → S + D , domain-specific tangent space mapping m ϕ : S + D × I d → R D
( D +1) / 2 , and a shared linear classifier g ψ : R D ( D +1) / 2 → Y with parameters Θ = { θ, ϕ, ψ } .
TSMNet with SPDDSMBN Unlike previous approaches which learn f θ , m ϕ , g ψ sequentially [29, 31, 13, 30], we parametrize h
= g ψ ◦ m ϕ ◦
f θ as a neural network and learn the entire model in an end-to-end fashion (Figure 1b).
Details of the proposed architecture, denoted TSMNet, are provided in appendix A.2.5.
In a nutshell, we parametrize f θ as the composition of the first two linear convolutional layers of ShConvNet [56], covariance pooling [57], BiMap [37], and ReEig [37] layers.
A BiMap layer applies a linear subspace projection, and a ReEig layer thresholds eigenvalues of symmetric matrices so that the output is SPD.
We used the default threshold ( 10 − 4 ) and found that it was never active in the trained models.
Hence, after training, f θ fulfilled the hypothesis class constraints.
In order for m ϕ to align the domain data and compute TSM, we use SPDDSMBN (18) with shared parameters (i.e., G ϕ i =
G ϕ
=
I , ν ϕ i = ν
ϕ ) in (16).
Finally, the classifier g ψ was parametrized as a linear layer with softmax activations.
We use the standard-cross entropy loss as training objective, and optimized the parameters with the Riemannian ADAM optimizer [58].
Experiments with EEG data
In the following, we apply our method to classify target labels from short segments of EEG data.
We consider two BCI applications, namely, mental imagery [59, 5] and mental workload estimation [60].
Both applications have high potential to aid society in rehabilitation and healthcare [61, 62, 18] but have, currently, limited practical value because of poor generalization across sessions and subjects [19, 15].
Datasets and preprocessing The considered mental imagery datasets were BNCI2014001 [63] ( subjects/2 sessions/4 classes), BNCI
[64] (12/2-3/2), Lee2019 [65] (54/2/2), Lehner [66] (1/7/2), Stieger2021 [67] (62/4-8/4) and Hehnberger2021 [68] (1/26/4).
For mental workload estimation, we used a recent competition dataset [69] (12/2/3).
A detailed description of the datasets is provided in appendix A.2.1.
Altogether, we analyzed a total of 603 sessions of 158 human subjects whose data was acquired in previous studies that obtained the subjects’ informed consent and the right to share anonymized data.
The python packages moabb [14] and mne [70] were used to preprocess the datasets.
The applied steps comprise resampling the EEG signals to 250/256 Hz, applying temporal filters to extract oscillatory EEG activity in the 4 to 36 Hz range (spectrally resolved if required by a method) and finally extract short segments ( ≤ 3 s ) associated to a class label (details provided in appendix A.2.2).
Evaluation We evaluated TSMNet against several baseline methods implementing direct transfer or multi-source (-target) UDA strategies.
They can be broadly categorized as component based [71, 68],
Riemannian geometry aware [12, 17, 30, 72] or deep learning [56, 73, 25].
All models were fit and evaluated with a randomized leave 5% of the sessions (inter-session TL) or subjects (inter-subject
TL)
out cross-validation (CV) scheme.
For inter-session TL, the models were only provided with data of the associated subject.
When required, inner train/test splits (neural nets) or CV (shallow methods) were used to optimize hyper parameters (e.g., early stopping, regularization parameters).
The dataset
Hehenberger2021 was used to fit the hyper parameters of TSMNet, and is, therefore, omitted in the presented results.
Balanced accuracy (i.e., the average recall across classes) was used as scoring metric.
As the discriminability of the data varies considerably across subjects, we decided to report the results in the figures relative to the score of a SoA domain-specific Riemannian geometry aware method [17], which was fitted and evaluated in a 80%/20% chronological train/test split (for details see appendix A.2.4).
Soft- and hardware We either used publicly available python code or implemented the methods in python using the packages torch [74], scikit-learn [75], skorch [76], geoopt [77], mne [70], pyriemann [78], pymanopt [79].
We ran the experiments on standard computation PCs equipped with 32 core
CPUs with 128 GB of RAM and used up to 1 GPU (24 GRAM).
Depending on the dataset size, fitting and evaluating TSMNet varied from a few seconds to minutes.
5.1 Mental imagery
TSMNet closes the gap to domain-specific methods Figure 2 summarizes the mental imagery results.
It displays test set scores of the considered TL methods relative to the score of a SoA domain-specific reference method.
Combining the results of all subjects across datasets (Figure 2a), it becomes apparent that TSMNet is the only method that can significantly reduce the gap to the reference method (inter-subject) or even exceed its performance (inter-session).
Figure 2b displays the results resolved across datasets (for details see appendix A.3.1).
We make two important observations.
First, concerning inter-session TL, TSMNet meets or exceeds the score of the reference method consistently across datasets.
Second, concerning inter-subject TL, we found that all considered methods tend to reduce the performance gap as the dataset size (# subjects) increases and that
TSMNet is consistently the top or among the top methods.
As a fitted TSMNet corresponds to a typical TSM model with a linear classifier, we can transform the fitted parameters into interpretable patterns [17].
Figure 3a displays extracted patterns for the BNCI2015001 dataset (inter-subject TL).
It is clearly visible that TSMNet infers the target label from neurophysiologically plausible sources (rows in Figure 3a).
As expected [2], the source with highest contribution has spectral peaks in the alpha and beta bands, and originates in contralateral and central sensorimotor cortex.
DSBN on S + D drives the success of TSMNet Since TSMNet combines several advances, we present the results of an ablation study in Table 2.
It summarizes the grand average inter-session TL test scores relative to TSMNet with SPDDSMBN for n = 138 subjects.
We observed three significant effects.
The largest effect can be attributed to S + D , as we observed the largest performance decline if the architecture would be modified 5 to omit the SPD manifold (4.5% with DSBN, 3% w/o DSBN).
[Table]
The performance gain comes at the cost of a 2.6x longer time to fit the parameters.
The second largest effect can be attributed to DSBN; without DSBN the performance dropped by 3.9% (with S + D ) and 2.4% (w/o S + D ).
The smallest, yet significant effect can be attributed to SPDMBN.
5.2 Mental workload estimation
Compared to the baseline methods [Table] TSMNet obtained the highest average scores of 54.7% (7.3%) and 52.4% (8.8%) in inter-session and -subject TL (for details see appendix A.3.1).
Interestingly, the inter-session TL score of TSMNet matches the score (54.3%) of the winning method in last year’s competition [15].
To shed light on the sources utilized by TSMNet, we show patterns for a fitted model in Figure 3b.
For the low mental workload class, the top contributing source’s activity peaked in the theta band and originated in pre-frontal areas.
The second source’s activity originated in occipital cortex with non-focal spectral profile.
Our results agree with the findings of previous research, as both areas and the theta band have been implicated in mind wandering and effort withdrawal [60].
Discussion
In this contribution, we proposed a machine learning framework around (domain-specific) momentum batch normalization on S + D to learn tangent space mapping (TSM) and feature extractors in an end-to-end fashion.
In a theoretical analysis, we provided error bounds for the running estimate of the Fréchet mean as well as convergence guarantees under reasonable assumptions.
We then applied the framework, to a multi-source multi-target unsupervised domain adaptation problem, namely, inter-session and -subject transfer learning for EEG data and obtained or attained state-of-the art performance with a simple, intrinsically interpretable model, denoted TSMNet, in a total of 6 diverse
BCI datasets (138 human subjects, 573 sessions).
In the case of mental imagery, we found that
TSMNet significantly reduced (inter-subject TL) or even exceeded (inter-session TL)
the performance gap to a SoA domain-specific method.
Although our framework could be readily extended to online UDA for unseen target domains, we limited this study to offline evaluations and leave actual BCI studies to future work.
A limitation of our framework, and also any other method that involves eigen decomposition, is the computational complexity, which limits its application to high-dimensional SPD features (e.g., fMRI connectivity matrices with fine spatial granularity).
Altogether, the presented results demonstrate the utility of our framework and in particular TSMNet as it not only achieves highly competitive results but is also intrinsically interpretable.
While we do not foresee any immediate negative societal impacts, we provide direct contributions towards the scalability and acceptability of EEG-based healthcare [1, 5] and consumer [18, 60] technologies.
We expect future works to evaluate the impact of the proposed methods in clinical applications of EEG like sleep staging [80, 81], seizure [82] or pathology detection [83, 84].
Acknowledgments and Disclosure of Funding
This work was supported by the JSPS KAKENHI (Grants-in-Aid for Scientific Research) grants 20H04249, 20H04208, 21H03516 and 21K12055.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See section
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See section (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See paragraph 4 in section 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See appendix A.1 in the supplementary material.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Supplementary Material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See paragraph 2 in section 5 for a brief summary and appendix A.2 for details.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] We provide error bars with respect to the variability across sessions and subjects.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See paragraph 3 in section 5. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See paragraph 1 in section 5.
(b) Did you mention the license of the assets?
[Yes] See appendix A.2.1 in the Supplementary material.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[Yes] See paragraph 1 in section 5.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[Yes] See paragraph 1 in section 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Supplementary Material A.1 Proof of proposition
In the proof, we will use Theorem 2 of [50] which relates the distance between interpolated points along the geodesic R # γ S , connecting points R and S , and point T to the distances between R , S and T .
Formally, for all R , S , T ∈ S
+ D we have
[Equation]
As a last ingredient for the proof, we use Proposition 1 of [50], which states that for a random variable
U , following distribution P U defined on S + D with Fréchet mean G U , we have for any point V ∈ S + D :
[Equation] [Equation]
which means that the expected distance between V and U is bounded from below by the variance of
U and the distance between V and its Fréchet mean G U .
Let us quickly repeat the definitions, assumptions and proposition 1.
We define a dataset containing the latent representations generated with feature set θ k as Z θ k
= { f θ k ( x ) | x ∈ T } , and a minibatch of M iid samples drawn from Z θ k at training step k as B k .
We denote the Fréchet mean of Z θ k as
G θ k , the estimated Fréchet mean, defined in (9), as G k , and the estimated batch mean as B k .
Since the batches are drawn randomly, we consider B k and G k as random variables.
We assume that the variance Var θ k
( B k − 1 ) =
E B k − 1 { δ 2 ( B k − 1 , G θ k ) } of the previous batch mean B k − 1 with respect to the current Fréchet mean G θ k is bounded by the current variance Var θ k
( B k ) and the norm of the difference in the parameters
[Equation]
Proposition 1 states then that the variance of the running mean Var θ k
[Equation]
G k ) is bounded by
[Equation]
25) over training steps k , if
[Equation]
holds true.
Proof of Proposition 1.
We prove Proposition 1 via induction.
We assume that the variance
Var θ k
( G k − 1 ) , that is the expected distance between the running mean G k − 1 and the Fréchet mean
G θ k , is bounded by the variance of the batch mean B k − 1 :
[Equation]
(27) and show that this also holds for G k and B k .
The assumption is trivially satisfied for G 0 = B 0 .
We start the induction step with (22) and apply it to the update rule for the running estimate G k .
As a result, we have
[Equation]
[Equation]
where we used G k
= G k − 1 # γ B k , as defined in algorithm 1.
Taking the expectations for the random variables G k , G k − 1 and B k , we get
[Equation]
Using (23) to simplify the last term [Equation]
we obtain
[Equation]
Applying assumptions (27) and (24) [Equation] we get
[Equation]
The resulting inequality holds true, for
[Equation]
[Equation]
and, in turn, results in feasible bounds for the parameter updates for fixed γ ∈ (0 , 1) .
This concludes the proof.
A.2 Supplementary Methods A.2.1 Datasets
A summary of the datasets’ key attributes is listed in Table 3.
The datasets contain a diverse sample of 154 human subjects, whose data was acquired in Europe (38 subjects; BNCI2014001,BNCI2015001,Lehner2021,Hehenberger2021,Hinss2021), Asia (54 subjects;
Lee2019)
and North America (62 subjects; Stieger2021).
A.2.2 Preprocessing
Depending on the dataset, either all or a subset of EEG channels was selected, and then resampled along the temporal dimension to a sampling rate of either 250 or 256 Hz.
(see Table 3).
Thereafter, an infinite impulse response (IIR) bandpass filter was used to extract EEG activity in the 4 to 36 Hz range (4th order Butterworth filter, 4 and 36 Hz cut-off frequencies, zero-phase).
Some baseline methods required spectrally resolved input data.
For these, we applied a bank of 8 filters with similar parameters except for the cut-off frequencies ([4, 8], [8, 12], ..., [32, 36] Hz).
Finally, short epochs (=segments) were extracted (see Table 3) relative to the task cues (=labels).
The labeled data were then extended with a domain index (= unique integer associated to one session of one subject).
A.2.3 Baseline methods
We considered several established and SoA baseline methods which were previously applied to intersession/-subject TL.
They can be broadly categorized as component based, Riemannian geometry aware or deep learning which we denote component, geometric and end-to-end, respectively.
For the component category, we considered the popular filter-bank common spatial patterns (FBCSP+SVM) [71] and a variant [68], designed for MSMTUDA, that applies domain-specific standardization (DSS) to features before classification, denoted FBCSP+DSS+LDA.
The geometric category was represented by TSM+SVM [12], a spectrally resolved variant [17] denoted FB+TSM+LR (which was also used as domain-specific baseline method).
We additionally considered two MSUDA methods, denoted
URPA+MDM and SPDOT+TSM+SVM here, that align SPD observations (=spatial covariance matrices) of different domains.
The former uses Riemannian Procrustes Analysis (RPA) [30] to align domains, and the latter optimal transport (OT) on S + D [72].
The end-to-end category was represented by EEGNet [73] and ShConvNet [56] two convolutional neural network architectures specifically designed for EEG data.
We additionally considered variants [25] that use domain-adversarial neural networks (DANN) [85] to learn domain-invariant latent representations.
A.2.4 Domain-specific reference method
Due to the success of TSM models [14, 10], we considered a spectrally resolved model [13, 17] which consisted of a filter-bank to separate activity of canonical frequency bands.
For each frequency band, PCA was used to reduce the spatial dimensionality and TSM to project the SPD features to the Euclidean vector space.
Finally, all features were pooled and submitted to a penalized logistic regression classifier.
For further details, see [17].
A.2.5 TSMNet
Architecture The architecture of TSMNet is outlined in Figure 4 and detailed in Table 4.
The feature extractor f θ comprises two convolutional layers, followed by covariance pooling [57], BiMap [37] and
ReEig [37] layers.
The first convolutional layer performs convolution along the temporal dimension; implementing a finite impulse response (FIR) filter bank (4 filters) with learnable parameters.
The second convolutional layer applies spatio-spectral filters (40 filters) along the spatial and convolutional channel dimensions.
Covariance pooling is then applied along the temporal dimension.
A subsequent
BiMap layer projects covariance matrices to a subspace via a bilinear mapping (i.e, BiMap( Z ) =
W Tθ ZW θ ) where the parameter matrix W θ is constrained to have orthogonal rows (i.e., W θ ∈ { U ∈
R
I × O : U T U
=
I
O , I ≥ O } ).
Next, a ReEig layer rectifies all eigenvalues, lower than a threshold ϵ = 10 − 4 (i.e, ReEig( Z ) =
U max( Σ , ϵ I )
U T with [ Σ , U ] =
eig
( Z ) ).
Domain-specific tangent space mapping m ϕ is implemented via combining SPDDSMBN and LogEig layers.
In order for m ϕ to align the domain data and compute TSM, we use SPDDSMBN (18) with shared parameters (i.e., G ϕ i =
G ϕ
=
I , ν ϕ i = ν
ϕ ) in (16).
The classifier g ψ was parametrized as a linear layer with softmax activations.
Parameter estimation We used the cross entropy loss as training objective, and the standard backpropagation framework with extensions for structured matrices [49] and manifold-constrained gradients [40] to propagate gradients through the layers of TSMNet.
Gradients were estimated with mini-batches of fixed size (50 observations; 10 per domain; 5 distinct domains) and converted into parameter updates with the Riemannian ADAM optimizer [58] ( 10 − 3 learning rate, 10 − 4 weight decay applied to unconstrained parameters; β 1 = 0 .
9 , β 2 = 0 .
For every MSMTUDA problem, comprising source and target domain sets, we split the source domains’ data into training and validation sets (80%/20% splits; randomized; stratified across domains and labels) and repeatedly iterated through the training set for 50 epochs via exhaustive minibatch sampling.
As required by SPDMBN, we implemented a decaying schedule (over epochs) for the training momentum parameter γ train ( k ) , defined in (11), with γ train (0)
= 1 and γ min = 0 .attained at epoch K = 40 .
During training, we monitored the loss on the validation data (at the end of every epoch).
Post training, the model with minimal loss on the validation data was selected.
For each target domain, the associated data was then passed through this model to estimate the labels.
During the forward pass, the domain’s normalization statistics within the SPDMBN layers were computed by solving (4) with the Karcher flow algorithm [42].
A.3 Supplementary results A.3.1 EEG data
The test set score for each considered EEG dataset is summarized in in Table 5.
Significant differences between the proposed method (TSMNet with SPDDSMBN) and baseline methods are highlighted.
[Table]
[Table] [Table]
[Table]
TempConv 1 x P x
[Table]
4 x P
[Table]
4 x P
[Table]
SPDDSMBN
[Table]
C
[Table] [Table]