[Paper Review] Geodesic Optimization for Predictive Shift Adaptation on EEG data

Apolline Mellot, Antoine Collas+3 more|arXiv (Cornell University)|Jul 4, 2024

Functional Brain Connectivity Studies5 citations

TL;DR

GOPSA learns domain-specific geodesic transports on the SPD manifold to jointly align covariance-based EEG features and predictive shifts, improving cross-site age prediction from resting-state EEG.

ABSTRACT

Electroencephalography (EEG) data is often collected from diverse contexts involving different populations and EEG devices. This variability can induce distribution shifts in the data $X$ and in the biomedical variables of interest $y$, thus limiting the application of supervised machine learning (ML) algorithms. While domain adaptation (DA) methods have been developed to mitigate the impact of these shifts, such methods struggle when distribution shifts occur simultaneously in $X$ and $y$. As state-of-the-art ML models for EEG represent the data by spatial covariance matrices, which lie on the Riemannian manifold of Symmetric Positive Definite (SPD) matrices, it is appealing to study DA techniques operating on the SPD manifold. This paper proposes a novel method termed Geodesic Optimization for Predictive Shift Adaptation (GOPSA) to address test-time multi-source DA for situations in which source domains have distinct $y$ distributions. GOPSA exploits the geodesic structure of the Riemannian manifold to jointly learn a domain-specific re-centering operator representing site-specific intercepts and the regression model. We performed empirical benchmarks on the cross-site generalization of age-prediction models with resting-state EEG data from a large multi-national dataset (HarMNqEEG), which included $14$ recording sites and more than $1500$ human participants. Compared to state-of-the-art methods, our results showed that GOPSA achieved significantly higher performance on three regression metrics ($R^2$, MAE, and Spearman's $ρ$) for several source-target site combinations, highlighting its effectiveness in tackling multi-source DA with predictive shifts in EEG data analysis. Our method has the potential to combine the advantages of mixed-effects modeling with machine learning for biomedical applications of EEG, such as multicenter clinical trials.

Motivation & Objective

• Address distribution shifts in both input data X (covariance matrices) and the target variable y in EEG-based predictive modeling.
• Develop a multi-source test-time domain adaptation method that does not require retraining on the target domain.
• Leverage Riemannian geometry of SPD matrices to learn domain-specific geodesic intercepts while sharing a global regression model.
• Demonstrate cross-site generalization improvement for age prediction from EEG covariance data.
• Provide a framework combining mixed-effects modeling with machine learning for multicenter EEG analyses.

Proposed method

• Represent EEG covariance matrices as points on the SPD manifold S_d^{++} and use affine-invariant Riemannian metrics to compute geodesic distances.
• Define a Riemannian mixed-effects model that learns domain-specific geodesic intercepts via parallel transport along a geodesic from each domain mean to the identity.
• Learn a shared linear regression in the tangent space after Riemannian logarithmic mapping, with a train-time optimization over domain weights α_k controlling transport degree.
• At train-time, jointly optimize domain transports and Ridge regression coefficients for all K source domains (Algorithm 1).
• At test-time, adapt to a new target domain by optimizing a transport parameter γ_T to align the target mean y_T with predictions (Algorithm 2).
• Key equations include the transport φ(Σ_i, Σ_k̄, α) = uvec(log_I(PT(Σ_i, Σ_k̄, α))) and the optimization of γ_S and γ_T under a ridge-regularized regression.

Figure 1: Joint shift in $X$ and $y$ distributions on the HarMNqEEG dataset [ 31 ] . Subset of mean PSDs ( A ) and age distributions ( B ) from three recording sites used for the empirical benchmarks.

Experimental results

Research questions

• RQ1 How to perform domain adaptation when shifts occur in both the input covariance matrices and the outcome variable in EEG data?
• RQ2 Can a geodesic-based, multi-source, test-time adaptation method on the SPD manifold outperform existing re-centering and domain-adaptation baselines for EEG-based regression tasks?
• RQ3 What is the benefit of learning domain-specific geodesic intercepts (parallel transport) while sharing a global regression model for cross-site EEG age prediction?

Key findings

• GOPSA achieves higher performance than several baselines (Dummy DO, No DA, Re-center, DO Intercept) across regression metrics on HarMNqEEG site combinations.
• For Spearman’s ρ, GOPSA reaches 0.78 ± 0.01 on the mean across site combinations, outperforming alternatives including DO Intercept (0.75 ± 0.02).
• For R^2, GOPSA achieves 0.61 ± 0.02 on the mean across site combinations, higher than Do Intercept (0.58 ± 0.02).
• For MAE, GOPSA yields 8.25 ± 0.19 on the mean across site combinations, better than Re-center (8.55 ± 0.18).
• Across specific site pairs, GOPSA consistently improves over baseline approaches, demonstrating effectiveness in handling predictive shifts in multi-site EEG data.

Figure 2: Normalized performance of the different methods on several source-target combinations for three metrics: Spearman’s $\rho$ $\uparrow$ (left), $R^{2}$ score $\uparrow$ (middle) and Mean Absolute Error $\downarrow$ (right). As a large variability in the score values was present between the s