SRI-EEG: State-Based Recurrent Imputation for EEG Artifact Correction

1. Introduction

Electroencephalography (EEG) is a non-invasive and widely-adopted approach to capture surface electrical activity in the human brain with electrodes placed on the scalp. EEG data has been widely used in domains including neuroscience, cognitive science, cognitive psychology, and mental health (Henry, 2006). One line of research has explored EEG as a potential interaction method, e.g., brain computer interfaces (BCIs). Specifically, by extracting context or intention from EEG data, hands-free BCI applications have been tested. Those applications include playing games (van de Laar et al., 2013), controlling a mouse cursor (Aydemir and Kayikcioglu, 2014), controlling robotic arms (Latif et al., 2017), managing smart home appliances (Anindya et al., 2016), and using smartphones (Kumar et al., 2017; Rashid et al., 2018). EEG has also been used as an accessible input modality when traditional input methods are not a feasible option. For example, brain-to-text communication (Willett et al., 2021) and hands-free wheelchair control (Singla et al., 2014) have enabled paralyzed or disabled users to communicate and move without assistance. Moreover, there is a growing number of EEG enabled BCI devices for consumers. Galea¹, Emotiv², Advanced Brain Monitoring³, and Muse⁴ are some representative devices that allow integration of various brain signals into a single headset for use in daily life and virtual reality applications. Research-grade EEG caps such as those developed by OpenBCI⁵ are being used to support both lab-based researchers as well biosensing enthusiasts. In addition, less bulky devices, such as the InEar BioFeed Controller (Matthies, 2013) or the Kokoon EEG-based sleep headphones⁶ are part of a new set of devices called hearables designed to track neural and physiological function of a user over long periods of time (Goverdovsky et al., 2017).

Although EEG-aided BCIs have seen significant interest and development since the 1970s, an open and challenging research question is to build robust user experiences under unconstrained (i.e., general and everyday use) conditions, given that EEG signals are easily contaminated by artifacts from users and their environment (Aricò et al., 2018). To deal with artifacts, the recorded EEG signals usually need to be heavily processed by sophisticated software, proprietary algorithms, or manual inspection before use. For example, independent component analysis (ICA) is a widely used approach that plays an important role in the EEG pre-processing phase for artifact removal. Though ICA variants and other advanced methods (e.g., canonical correlation analysis, wavelet transform algorithm, empirical-mode decomposition) have been introduced over the years, the EEG artifact correction procedure still requires manual inspection of the signal data to pick out noise . Furthermore, random and irregularly-occurring artifacts make signal correction increasingly challenging. The complex, time consuming, and labor intensive pre-processing requirements thus limit the robustness and generalization of EEG supported BCIs for many consumer use cases, and hinder the development of ubiquitous EEG-based applications. Therefore, developing an automatic EEG artifact correction approach that can significantly reduce manual tagging time and expense is imperative to making BCI broadly useful as a general purpose input method.

Automated artifact correction methods using deep-learning-based signal processing have achieved state-of-the-art performance on a variety of time series data. One line of research has explored the use of recurrent neural networks (RNNs) for time series data imputation to predict noisy or missing portions from the remaining segments of data. A large set of relevant works [e.g., GRU-D (Che et al., 2018a), MR-HDMM (Che et al., 2018b), MRNN (Yoon et al., 2017), GAIN (Yoon et al., 2018), and BRITS (Cao et al., 2018)] have shown that RNNs can effectively encode the temporal relationships in time series data and are, thus, able to achieve satisfactory performance on the data imputation task. Additionally, transformer-based neural networks have achieved impressive performance in various research domains including natural language processing and computer vision. Motivated by their success, recent work has started to tackle time series data in a sequence-to-sequence manner by applying the self-attention mechanism common in transformer-based models. For example, Zerveas et al. (2021) and Tran et al. (2021) present two representative works to adapt transformer-based encoders on the classification, regression and imputation problems in time series data and achieve superior performance.

While EEG is a type of time series data, the automatic imputation designed for correcting artifacts, caused by EEG's susceptibility to various types of noise during data collection, remains a relatively unexplored and challenging area. To address this challenge and take advantage of recent success of deep learning methods for time series data imputation, we propose a state-based recurrent neural network (RNN) for EEG artifact correction (SRI-EEG) in an automatic manner. Our network is built upon an RNN as its backbone structure. It takes into account the state information encoded in EEG signal representing features of stimuli applied during EEG data recording. Added to the state information is a spatial decay matrix to model EEG channel dependencies. We evaluate the performance of our proposed network on three publicly available datasets: Bike (Bullock et al., 2015), Kaggle (Margaux et al., 2012), and SMR (Tangermann et al., 2012) (Section 6).

In summary, our main contributions include: (1) an automatic imputation method to correct irregularly-occurring artifacts in EEG data leading to significantly reduced manual effort that could encourage broader usage of EEG based interaction modalities; (2) a novel bidirectional long short-term memory (LSTM) model for imputing artifacts in EEG signals. This approach captures the temporal and spatial dependencies of recorded data, and leverages EEG state features to enable automatic artifact correction through data imputation; (3) results from a qualitative and quantitative evaluation of SRI-EEG on three publicly available EEG datasets. Experimental results show that our algorithm achieves state-of-the-art performance on EEG imputation.

2. Related Work

2.1. Time Series Imputation

Time series data imputation is defined as replacing data gaps with predicted values computed from the remaining data. Simple methods replace the missing data with the mean or median of non-empty values, or the last observed value. Such methods offer a fast and easy way to impute missing portions from known data. However, these approaches are based on simple statistical rules and can easily introduce large errors, e.g., imputing long-range missing segments from limited known data.

Non-deep-learning-based techniques have been commonly adopted as they outperform the aforementioned simple imputation methods and do not demand dense computation. K nearest neighbors (KNNs)-based imputation (Zhang and Zhou, 2007) utilizes the statistical dependency in neighboring data, and imputes missing parts using the mean of K nearest neighbors. Factorization-based techniques (Friedman et al., 2001), decompose the non-missing signal into basis vectors to approximate values at missing points. Adaptations of factorization-based algorithms, such as multivariate imputation by chained equations (MICE) (Azur et al., 2011) and SoftImpute (Mazumder et al., 2010), have also been explored with varying levels of success. However, one of the main limitations of these approaches is the inability to capture long-term temporal dependencies in time series data.

Recent advancements in deep learning have shown desirable capability to encode long-range temporal correlations and have achieved state-of-the-art imputation performance on time series datasets (Moritz and Bartz-Beielstein, 2017; Luo et al., 2018; Miller et al., 2018; Yuan et al., 2018; Suo et al., 2019; Zhang and Yin, 2019; Tang et al., 2020; Miao et al., 2021; Zerveas et al., 2021). Among these works, BRITS (Cao et al., 2018) and MRNN (Yoon et al., 2017) use bidirectional RNNs, which impute missing values based on hidden states updated in both the forward and backward directions. These two approaches have achieved state-of-the-art performance on one of the most popular multivariate time series datasets, PhysioNet Challenge 2012 (Goldberger et al., 2000). MRNN takes the values to-be-imputed as constraints and does not encode correlations between the missing segments. BRITS does not impose strong assumptions on the imputation setup such as linear dynamics in the hidden states, and enables imputing correlated missing segments.

Motivated by the successful adoption of bidirectional RNNs on the time series imputation problem, we build our network's backbone using a bidirectional LSTM. The LSTM processes time series data in both the forward and backward directions with two separate sequences of hidden states. The two sequences are then averaged to predict imputed values. Such a backbone network has mostly been used for general time series data in prior work (e.g., air quality, human activity, and health care datasets). It fails to impute EEG signals containing a large amount of irregularly occurring artifacts. To tackle this limitation, we introduce encoded EEG state constraints into the network along with taking as input the spatial correlations between signal collected by all the electrodes in an EEG cap. Through our evaluation, reported in Section 6, we demonstrate the effectiveness of our approach targeted at the EEG artifact imputation problem.

2.2. EEG Artifact Correction

Artifacts in EEG are undesired noise mainly originating from two types of sources: (1) extrinsic artifacts include environment noise and experimental errors, and (2) intrinsic artifacts include physiological artifacts (Jiang et al., 2019) such as body motions or eye movements. Extrinsic artifacts can be eliminated by filtering recorded signal or following proper experimental procedures. However, preventing intrinsic artifacts is more challenging and their removal requires particular algorithms (Anderer et al., 1999). In this work, we focus on eliminating physiological artifacts stemming from the human body, specifically body motions, eye movements, and cardiac activities. Therefore, in this section we focus on work most closely related to handling physiological artifacts in EEG data.

Blind source separation (BSS) approaches, including principal component analysis (PCA) and independent component analysis (ICA), are widely used for EEG artifact removal. PCA (Berg and Scherg, 1991) is a widely used BSS technique. It decomposes EEG signals into uncorrelated variables, called principal components (PCs), through an orthogonal transformation. The PCs representing artifacts are removed to denoise the signal, while the remaining PCs are used to reconstruct clean EEG data. PCA typically fails to unravel signal dependencies and can mistakenly discard non-artifact signal (Casarotto et al., 2004). ICA (Somers and Bertrand, 2016) is a flexible BSS method. It assumes that collected EEG signals are linear mixtures of artifacts and non-artifacts. The first step involves decomposing recorded signals into independent components (ICs). Sejnowski (1996) introduces artifactual and non-artifactual ICs as independent ICs, such that obvious artifactual ICs can be segregated easily. By eliminating ICs representing artifacts, ICA reconstructs clean signal with the remaining ICs. Built upon the conception of ICA, recent research has studied EEG artifact removal with various ICA variants and demonstrated their effectiveness (Flexer et al., 2005; Bian et al., 2006; Li et al., 2006; Ting et al., 2006; Zhou and Gotman, 2009; Winkler et al., 2011; Rejer and Górski, 2015, 2019; Dimigen, 2020; Klug and Gramann, 2020). Although a large number of ICA variants have been developed over the last twenty years to control EEG artifacts in different scenarios, the inherent idea of how ICA works imposes assumptions on the recorded signal but the assumptions are not always satisfied. For example, ICA can estimate non-Gaussian signals (one of the premises of ICA), but recorded signals are usually unknown to be Gaussian or non-Gaussian (Jiang et al., 2019).

Though ICA is an effective standard for EEG artifact correction in general, using ICA alone still requires careful manual effort and time to inspect the source signal, which may not have satisfactory properties to fit ICA. We use ICA to detect artifacts followed by automatically imputing them using our proposed method. Combining ICA with the proposed automatic imputation approach achieves improved performance compared to using ICA alone as shown in Section 6.1.

2.3. Brain Computer Interfaces and Brain Signals

Depending on the features of interest, EEG based BCIs can be categorized into two classes (Lotte, 2015): (1) event-related potential (ERP) based BCI, and (2) oscillatory activity based BCI. ERP BCIs can detect high-amplitude and low-frequency brain responses to known stimuli. They usually contain well-stereotyped waveforms and are robust across subjects (Fazel-Rezai et al., 2012). The robustness of ERPs thus enables efficient learning of ERP features using machine learning based algorithms. For example, the P300 response, one type of ERP, is evoked by visual stimuli, and forms identifiable electrical waveforms. Other types of ERPs, such as motor imagery, also induce recognizable waveforms. Prior deep-learning-based methods have deployed various classifiers to distinguish between different ERP types (Lawhern et al., 2018; Santamaŕıa-Vázquez et al., 2019, 2020; Wen et al., 2019; Zhao et al., 2020; Zang et al., 2021). Oscillatory BCIs can detect the signal power of EEG frequency bands. Due to low signal-to-noise ratio and great variations across subjects (Pfurtscheller and Neuper, 2001), oscillatory BCIs are typically challenging to model using machine learning models. An example of oscillatory BCIs is the sensorimotor rhythm (SMR). Subjects produce lower SMR amplitudes when the corresponding sensorimotor regions are active, and higher amplitudes otherwise (Arroyo et al., 1993). SMR data is greatly variable across and even within subjects, and thus demands subject training for data collection, and longer signal calibration session for practical use (Krusienski et al., 2008).

To assess the generality and robustness of the proposed EEG artifact imputation method under different EEG paradigms, we evaluate SRI-EEG on two ERP datasets: Bike (Bullock et al., 2015) and Kaggle (Margaux et al., 2012), and one oscillatory dataset: SMR (Tangermann et al., 2012).

3. Method

Similar to the standard bidirectional LSTM for time series data imputation, our network takes as input a data matrix, a mask matrix and a time gap matrix processed from raw data. However, unlike using a bidirectional LSTM designed for general time series data, we also take into account a spatial gap matrix for modeling EEG channel dependency. Added to that is a state vector of encoded EEG event features to handle irregularly occurring artifacts that form a typical but challenging problem for EEG based BCI modalities. Section 3.1 introduces the preliminaries covering notation and definitions. Section 3.2 presents the backbone bidirectional LSTM for general time series imputation and discusses its limitations. Sections 3.3 and 3.4 present our design based on the RNN backbone network to specifically tackle the EEG imputation problem. Section 3.5 discusses the learning objectives. Figure 1 shows an overview of the SRI-EEG architecture.

3.1. Preliminary

We denote the raw EEG data as a two-dimensional matrix where the first dimension represents time: ${{X}}^i=\left\{x_1,x_2,...,x_t,...,x_T\right\}$ is a sequence of data with T timestamp observations collected at trial i. The t-th observation $x_t\in {ℝ}^D$ consists of values from D EEG channels $\left\{x_t^1,x_t^2,...,x_t^d,...,x_t^D\right\}$ at timestamp t. Note that the time range 1 to T and the channels 1 to D in this section are from any single trial of data, where trial i is defined as a segment/epoch of interest during EEG data collection.

The time difference between successive valid (non-artifact) observations varies due to artifacts of random lengths. To model the changing time gaps between neighboring valid observations, we define a temporal decay matrix: ${\delta }_t^d$ at time t and channel d similar to Cao et al. (2018):

where we denote a binary mask matrix ${M}=\left\{m_1^d,m_2^d,...,m_t^d,...,m_T^d\right\}$ along the time domain from channel d∈D to mask out the artifacts in the data matrices. Specifically,

Additionally, we denote by y_i the EEG state vector representing the state of the i-th trial. It is a learned feature vector using EEGNet (Lawhern et al., 2018), one of the state-of-the-art EEG classifiers. EEGNet allows us to extract a state feature vector to represent stimuli applied to subjects at trial i (Section 3.4). The state feature vector is fed into our network as supervised information together with the data matrix, mask matrix, temporal decay matrix, and spatial decay matrix (defined in Section 3.3). Figure 1 depicts all the input matrices.

3.2. Bidirectional LSTM for Time Series Imputation

In this section, we briefly introduce unidirectional LSTMs, and our backbone network which consists of a bidirectional LSTM with input mask and temporal decay matrices. Following the background information, we discuss the backbone structure's limitations and introduce our improvements.

3.2.1. Unidirectional LSTM

In a unidirectional LSTM, the hidden layer receives an input vector x_t, and predicts an output vector ${\widehat{x}}_t$. At each timestamp t, the hidden layer maintains a hidden state h_t, and updates the hidden state based on the current input x_t and the previous hidden state h_t−1. We formulate this procedure as:

where σ_h is the hidden layer's activation function, W_xh is the weight matrix from the input layer to the hidden layer, W_hh is the weight matrix between two temporally consecutive hidden states, and b_h is the hidden layer's bias vector. Consequently, we compute the output vector ${\widehat{x}}_t$ as:

where σ_x is the output layer's activation function, W_x is the weight matrix from the hidden layer to the output layer, and b_x is the output layer's bias vector.

3.2.2. Unidirectional LSTM With Input Mask

The unidirectional LSTM takes continuous input without data gaps to update hidden states progressively. However, for the imputation task, where the input data matrix ${{X}}^i$ consists of both artifacts and non-artifacts, the raw data cannot be fed into the network to update the hidden states. The artifacts will induce distracting features and thus introduce undetermined bias in the imputation process. To resolve this issue, we leverage a mask matrix ${M}$ to filter out artifacts and pass only the valid data into the network, similar to Cao et al. (2018) and Yoon et al. (2018). Accordingly, the input matrix ${{X}}^i=\left\{x_1^d,x_2^d,...,x_t^d,...,x_T^d\right\}$ is computed by:

where ⊙ is element-wise multiplication.

3.2.3. Unidirectional LSTM With Input Mask and Temporal Decay

Although the hidden states of LSTMs enable the modeling of decayed influence along the time axis (i.e., distant values have less impact than close ones), the length of effective time window is automatically learned and thus not easily controllable by hyper-parameter tuning. To address this limitation, we utilize the defined temporal decay matrix ${\delta }_t^d$ to control the temporal influence on hidden state update, similar to Cao et al. (2018). Specifically, the temporal decay factor ${\beta }_t^d$ is computed by:

where W_β is the time gap's weight matrix, and b_β is the bias vector. Consequently, the hidden state is updated by applying an element-wise multiplication between the previous hidden state (at timestamp t−1) and the current temporal decay factor (at timestamp t):

where β_t represents the temporal decay factor at time t from all the D channels.

3.2.4. Bidirectional LSTM With Input Mask and Temporal Decay

We denote by $\vec{h}$ the hidden state sequence of the forward-directional LSTM, and we iteratively update the sequence based on inputs from timestamps 1 to T. Conversely, the backward-direction LSTM updates the hidden state sequence $\stackrel{⃖}{h}$ by taking in reversed inputs from timestamps T to 1. Combining both the forward and backward hidden states, we compute the estimated values ${\widehat{x}}_t$ by:

where σ is an average function to combine the forward and backward hidden states. $W_x^{\mathrm{\text{forward}}}$ and $b_x^{\mathrm{\text{forward}}}$ are the forward-directional LSTM's weight matrix and bias vector from the hidden layer to the output layer, respectively. It is similar for $W_x^{\mathrm{\text{backward}}}$ and $b_x^{\mathrm{\text{backward}}}$ in the backward-direction LSTM.

3.2.5. Our Improvements

Introducing the mask matrix and temporal decay into the bidirectional LSTM proposed by prior work forms an effective backbone network to impute multiple time series datasets as evaluated by Yoon et al. (2017) and Cao et al. (2018). Nevertheless, EEG is a special type of time series data with unique characteristics such as spatial dependency between electrodes and state-related global waveforms. For example, when a region of the human brain is stimulated by a certain event, the imputation is likely to benefit from assigning higher weights to the EEG signals recorded on the scalp above that active brain region, e.g., all electrodes above the motor cortex. Additionally, the imputation performance can be potentially improved by capturing EEG waveform patterns from EEG sequences in a different trial that are temporally similar, i.e., in the same EEG state. Such EEG state correlation introduces favorable inductive biases. We take into account the aforementioned EEG properties and design the imputation algorithm by exploring the channel/electrode spatial correlations (Section 3.3), including EEG state modeling (Section 3.4).

3.3. Learning the Spatial Correlations of EEG Channels

As shown in Figure 2, the 32 electrodes are positioned on a cap in a particular layout (the International 10-20 system). The electrode topology presents varying proximity between different pairs of electrodes. For example, when certain brain regions are active, triggered by some stimuli (red regions in Figure 2), all the electrodes within that region collect related signals. These signals, thus, become important neighboring values to impute contaminated channels affected by artifacts. To model the spatial correlations between the electrodes, which prior work has not extensively studied, we propose a novel spatial decay matrix ${\delta }_d^{p,q}$. It consists of the Euclidean distance between any EEG channel pair: p and q. We position the channels in a 3D Euclidean coordinate system, and denote the coordinates of any two channels as p = (p_x, p_y, p_z) and q = (q_x, q_y, q_z), respectively. Accordingly, the distance between any EEG electrode pair is calculated by:

Similar to the temporal decay factor β_t, we introduce a spatial decay factor η_d to model the decay dependency between EEG channel pairs as their distance increases:

where δ_d denotes the distances between any channel d and all the remaining channels, W_η is the weight matrix, and b_η is the bias vector.

Combining the temporal and spatial decay factors with the weight hyper-parameter λ∈[0, 1], we formulate the temporal-spatial decay matrix as:

We update the bidirectional LSTM's hidden states by element-wise multiplying the temporal-spatial decay matrix and the last hidden state (forward and backward):

3.4. Modeling the State of EEG Trials

Studies on EEG classification (Lotte et al., 2007; Fazel-Rezai et al., 2012; Lawhern et al., 2018) have demonstrated that different stimuli applied to subjects during each trial can lead to different features in the collected EEG data. These features can be differentiated to enable classifying the EEG signals, e.g., classifying ERPs. Motivated by these findings and to take advantage of the recognizable features of EEG waveforms, we utilize the state features to provide supervised information on the EEG signal imputation task. Specifically, we obtain the features learned from the state-of-the-art EEG classifier EEGNet (Lawhern et al., 2018). We subsequently transform the features into a 1-dimensional vector (state vector) to serve as supervised information for our network's input. More formally, we denote the state of EEG trial i as y_i, and concatenate y_i with the masked data matrix to pass them together into the network.

In summary, the update of hidden states considering temporal-spatial decay and the EEG state is computed by:

where ⊕ signifies concatenation.

3.5. Learning Objective

We train our model using the Mean Squared Error (MSE) loss in the back-propagation process:

where n denotes the total number of observations in each trial. The MSE is calculated as an averaged sum of the difference between the observed and imputed EEG values along the time domain and across all EEG channels.

4. EEG Dataset

For evaluation, we use three publicly available datasets, namely Bike, Kaggle⁷, and SMR⁸. The feature type, the number of subjects and samples, and the number of trials per subject in all three datasets are summarized in Table 1.

4.1. Bike

The Bike dataset was introduced by Bullock et al. (2015) and provides P300 event related potentials (ERPs). Twelve subjects performed two different versions of three-stimuli oddball tasks while seated on a stationary bike. Participants were asked to respond to the target stimuli and ignore the remaining two types of distracting stimuli. The stimuli were presented at different rates: 1 Hz [200 ms stimulus presentation with 800 ms inter-stimulus interval (ISI)], and 2 Hz (200 ms stimulus presentation with 300 ms ISI). Participants completed the two-version tasks at rest (sat on the bike but not pedaling), and during exercise (sat on the bike and pedaling).

We followed the train-test split protocol and pre-processing procedure used in Ding et al. (2019). We used both the 1 and 2 Hz datasets (pedaling or noisy, and not pedaling or denoised). Specifically, the Noisy_1 Hz and Noisy_2 Hz datasets are the “noisy” versions of the Bike dataset since pedaling introduced a large amount of artifacts in the collected EEG signals. Additionally, we analyzed “denoised” versions of the Bike dataset, named Denoised_1 Hz and Denoised_2 Hz, where the subjects were not pedaling during EEG collection. For the data pre-processing, we applied a bandpass filter and kept the 1–40 Hz frequency band and lower the sampling rate from 512 to 128 Hz.

4.2. Kaggle

For the Kaggle dataset (Margaux et al., 2012), 26 participants engaged in a P300 speller task. The system presented a random sequence of flashing letters and numbers, arranged in a 6 ×6 grid, to elicit the P300 responses (Krusienski et al., 2008). The goal of this task was to determine whether the displayed item is the target item by analyzing recorded EEG signals. Following the same pre-processing procedure as the Bike dataset, the EEG signal was band-pass filtered to 1–40 Hz and down-sampled to 128 Hz. The training and testing sets were split as suggested in the official release.

4.3. SMR

We use the BCI Competition IV Dataset 2A (Tangermann et al., 2012) involving nine subjects for oscillatory EEG data analysis. The SMR data consists of four classes of imagined movements of left and right hands, feet and tongue. The EEG data was originally recorded using 22 electrodes, sampled at 250 Hz and bandpass filtered between 0.5 and 100 Hz (Schirrmeister et al., 2017). We re-sampled the data to 128 Hz, and band-pass filtered it to 1–40 Hz to keep consistent pre-processing setups. The training and testing sets were split as in the official release.

5. Experimental Setup

5.1. Evaluation Setup and Metrics

We train our imputation model using the Adam optimizer (Kingma and Ba, 2015) with a learning rate of 0.001 and batch size of 32. The weight λ between temporal and spatial decay factors is set to 0.8 which was determined to be a suitable value through experimentation. We adopt the early stopping strategy by randomly selecting 15% of the training sets to form the validation dataset. Training is stopped when the lowest validation error is obtained. The training weights are saved for subsequent evaluation on the testing sets.

We adopt two metrics for performance evaluation: mean absolute error (MAE) and root mean squared error (RMSE). We define ${{G}}^i=\left\{g_1,g_2,...,g_n\right\}$ as the observed data matrix, and ${{P}}^i=\left\{p_1,p_2,...,p_n\right\}$ as the corresponding imputed data matrix. n is the total number of observations in each trial i. Accordingly, the metrics are defined as:

5.2. Baseline and Comparison Methods

We compare the proposed imputation model with the following non-deep-learning algorithms: Mean, KNN, SoftImpute, and ICA, and with two deep learning algorithms: MRNN and BRITS. The non-deep-learning approaches compute comparatively fast and are widely adopted as baselines in prior studies. The deep-learning methods have achieved great improvements in imputation performance on time series datasets allowing us comparisons with the state-of-the-art.

• Mean: Mean imputation replaces artifacts with the global average along the time domain for each trial.

• KNN⁹: KNN uses the k-nearest neighbors to find similar samples, and imputes the artifacts with weighted average of the k neighbors.

• SoftImpute⁹[]: SoftImpute is a matrix completion based technique that works by applying iterative soft thresholding of SVD decompositions.

• ICA¹⁰: ICA decomposes the original signal into independent components, and then removes the components consisting of artifacts for denoising the signal. ICA is a popular EEG artifact correction approach widely adopted by researchers.

• MRNN¹¹ (Yoon et al., 2017): MRNN is a bidirectional RNN based algorithm, and treats the imputed values as constants without modeling correlations between channels.

• BRITS (see textfootnote 11) (Cao et al., 2018): BRITS is an extension of MRNN. It considers the correlations between different imputed values.

5.3. Dataset Preparation

We evaluate our proposed method on two setups: (1) we detect artifacts in EEG with ICA and manual inspection, and impute the detected artifacts using our proposed imputation method, and (2) we synthesize EEG artifacts by removing varying percentages of the recorded EEG data, and impute the synthesized artifacts using our proposed imputation method.

5.3.1. EEG Artifact Detection

We extract when and where the artifacts occur during EEG recording via ICA (Winkler et al., 2011). Additionally, we detect obvious missing values caused by bad channels using functions in EEGLAB (Delorme and Makeig, 2004) by manually investigating kurtosis in the channel and looking at how well each channel correlates with the surrounding channels.

Although this setup needs manual effort to detect artifacts in the recorded EEG signal, it lays a practical basis to evaluate our network's effectiveness. This setup also incurs complexity in extracted artifacts, such as different artifact quantities across trials and different artifact patterns caused by various sources of noise.

5.3.2. Synthesized EEG Artifacts

We synthesize EEG noise by randomly removing certain portions of the observed data, and replacing the removed parts with NaNs. The synthesized artifact rates in our experiment are: 2.5, 5, 7.5, 10, 12.5, 15, 17.5, 20%. This setup offers controllable artifact rates to compare our approach's performance with other imputation methods at fixed artifact rates. The artifact ratios help to analyze an imputation method's performance on varying artifact rates (studied in Section 6.2).

The motivation for utilizing the synthesized artifacts is to prevent the network from getting scoped by certain types of artifacts, such as those caused by heart beats. This goal is achieved by introducing randomness and increasing complexity of the input data. Specifically, the synthesized artifacts may consist of clean EEG, real artifacts or both. The synthesized random and complex artifacts prevent the model from simply “remembering” artifact patterns but failing to effectively capture the temporal and spatial relationships between data points. By evaluating the imputation methods based on the synthesized artifacts setup, we aim to test the algorithm's robustness and further boost EEG usage for scenarios that introduce different artifact types. The synthesized artifacts are obtained automatically, as opposed to the manual effort in detecting real artifacts using ICA. This setup is therefore, a step toward the potential goal of building automatic EEG artifact correction pipelines where people can detect real artifacts without heavily relying on human effort. Following automatic artifact detection, our proposed imputation algorithm can be applied in the artifact correction process to generate denoised EEG signal.

6. Results

6.1. Detected Artifacts Imputation

We qualitatively compare the observed EEG data and the corresponding corrected EEG data output by our proposed imputation method, shown in Figure 3. In the Bike-Denoised_1 Hz dataset, we randomly select EEG data collected by 10 out of 32 channels from a random trial. This is done to avoid crowding in the plots while still ensuring generality. Each channel is plotted in a unique color. Marked by the boxes, signals from timestamps 30–40 are detected as contaminated data. The artifacts are visible as mostly extreme values in this time window. After applying our proposed imputation method, as shown in the right plot, the contaminated data has been corrected by taking into account values from neighboring channels and timestamps, and by considering the trial state feature. The imputed values eliminate the waveforms consisting of extreme large or small values to maintain the waveform trend as a whole.

We also conduct a quantitative study and Tables 2, 3 show the detected EEG artifact imputation performance comparisons. The results show great variance across different datasets as the occurrence of detected artifacts differs considerably. The two denoised Bike datasets present lower MAE and RMSE compared with the corresponding noisy sets. For each dataset, we can see an improvement in imputation performance from the top to the bottom. Our method performs well on all three datasets and achieves lower errors compared to the competing approaches in general. Overall, the results demonstrate our proposed method's ability to correct real EEG artifacts.

To measure imputation performance significance on data reported in Table 2, we performed a Wilcoxon signed rank test over all the imputed values averaged over 10 runs on the Denoised-1 Hz, Denoised-2 Hz, Noisy-1 Hz, Noisy-2Hz, Kaggle and SMR test sets. We report the full results in Supplementary Tables 1–6 and Supplementary Figures 1–6. Here we summarize the main findings.

Denoised_1 Hz dataset. Absolute error for the imputed values was significantly different for the compared imputation methods as determined using the Friedman test, ${{X}}_{(6)}^2=440.15$, p < 0.0001, n = 1, 980. All pairwise differences between SRI-EEG and the other methods were significant to at least a level of p < 0.001. Exception is SRI-EEG vs ICA, which was significant to a level of p < 0.05 instead.

Denoised_2 Hz dataset. Absolute error for the imputed values was significantly different for the compared methods as determined using the Friedman test, ${{X}}_{(6)}^2=369.44$, p < 0.0001, n = 1, 980. All pairwise differences between SRI-EEG and the other methods were significant to at least a level of p < 0.001.

Noisy_1 Hz dataset. Absolute error for the imputed values was significantly different for the compared imputation methods as determined using the Friedman test, ${{X}}_{(6)}^2=656.41$, p < 0.0001, n = 1, 980. All pairwise differences between SRI-EEG and the other methods were significant to at least a level of p < 0.001. Exceptions are SRI-EEG vs MRNN and SRI-EEG vs BRITS, both of which were non-significant (p >0.05).

Noisy_2 Hz dataset. Absolute error for the imputed values was significantly different for the compared imputation methods as determined using the Friedman test, ${{X}}_{(6)}^2=349.25$, p < 0.0001, n = 1, 980. All pairwise differences between SRI-EEG and the other methods were significant to at least a level of p < 0.001.

Kaggle dataset. Absolute error for the imputed values was significantly different for the compared imputation methods as determined using the Friedman test, ${{X}}_{(6)}^2=2317.41$, p < 0.0001, n = 1, 980.

SMR dataset. Absolute error for the imputed values was significantly different for the different compared methods as determined using the Friedman test, ${{X}}_{(6)}^2=1967.3$, p < 0.0001, n = 1, 980. All pairwise differences between SRI-EEG and the other methods were significant to at least a level of p < 0.001.

6.2. Synthesized Artifacts Imputation

6.2.1. Imputation With the Same Synthesized Artifact Rate

Tables 4, 5 show the performance of different imputation methods evaluated on all the datasets with 10% artifact rate. Evidently, simply replacing artifacts with the mean of observed data is highly inaccurate. KNN and SoftImpute perform much better than simple averaging imputation. Using ICA to detect and remove artifacts outperforms the aforementioned three methods. In contrast, the two RNN based methods, BRITS and MRNN, demonstrate significantly improved performance compared to the non-deep-learning-based approaches. SRI-EEG shows comparable performance to BRITS and MRNN.

6.2.2. Imputation With Varying Synthesized Artifact Rates

On the Bike-Denoised_1Hz dataset, we evaluate the imputation performance at various artifact rates. Figure 4 reveals the MAE (left) and RMSE (right) using different imputation methods at artifact rates from 2.5 to 20%. As is clear, the two errors increase with increase in artifact rates. In particular, simple mean imputation is very sensitive to artifacts, where the errors grow rapidly as the number of artifacts increases. KNN, SoftImpute, and ICA perform much better than the mean imputation. The two RNN based methods (BRITS and MRNN) and our method achieve further lower imputation errors than the non-RNN based approaches at all artifact rates. This study demonstrates the robustness of our model's ability to handle varying proportions of artifacts in the signal, and the effectiveness of deep-learning-based methods in general to tackle the imputation problem at different artifact rates.

6.3. Ablation Study

6.3.1. Influence of EEG Channel Spatial Dependency and EEG Trial State

We perform an ablation study to evaluate the effectiveness of including EEG state and channel spatial correlation as the network input. As shown in Tables 6, 7, after removing either or both the trial state supervision and spatial decay in SRI-EEG, the imputation performance degenerates on all the datasets evaluated by MAE and RMSE. Based on these results, both factors appear to be important components contributing to our algorithm's performance demonstrating effectiveness for the EEG imputation task compared to other methods.

6.4. Post-imputation Analysis

6.4.1. Goal

To evaluate SRI-EEG on a practical task, we conducted post-imputation analysis in the form of an EEG classification problem. Our aim was to analyze if our model can correct artifacts in EEG and further boost the EEG classification performance. We designed this study as an indirect reflection of the imputation performance and a first step toward using imputation based EEG artifact correction to support EEG-based applications.

6.4.2. Dataset and Metrics

We performed analysis on the Bike dataset, i.e., Denoised_1 Hz, Denoised_2 Hz, Noisy_1 Hz, and Noisy_2 Hz. The Bike dataset consists of P300 ERPs from three-stimuli oddball tasks (introduced in Section 4.1). Therefore, we conduct the classification experiments to evaluate the three-class classification performance.

We use precision, recall and F1 score as the classification metrics:

where TP, FP, and FN stand for “true positive,” “false positive,” and “false negative,” respectively.

6.4.3. Setup

We take the classification performance without imputation as a baseline. Specifically, based on the recorded EEG data, i.e., no imputation applied, we perform baseline classification using xDAWN+RG, which is a multi-class EEG classifier (Barachant et al., 2011). Compared against this baseline, we classify the imputed EEG data using BRITS, MRNN, and SRI-EEG as the imputation methods, respectively.

6.4.4. Results and Analysis

We present the classification comparisons in Table 8 with the best performance in bold. Higher precision and recall correspond to better performance. We find that by classifying imputed EEG signals, both the precision and recall are improved in general on the denoised and noisy Bike sets compared with directly classifying the raw data without artifact correction. Importantly, compared with the two RNN-based imputation methods, SRI-EEG achieves improved precision and recall on the classification task evaluated on Denoised_1Hz and Denoised_2 Hz datasets. Although the precision on Noisy_1 Hz and Noisy_2 Hz datasets is less satisfactory than that using MRNN based post-imputation analysis, the recall shows improved performance.

To measure any significance in performance based on data in Table 8, we performed a McNemar's test over all the classified labels (1,200 labels) on the Denoised-1 Hz, Denoised-2 Hz, Noisy-1 Hz, and Noisy-2 Hz test sets. McNemar's test has been proven to be a suitable and effective way to compare machine learning classifiers (Dietterich, 1998). We report the full results in Supplementary Tables 7–10. Here we summarize the main findings.

(1) On the Denoised_1Hz dataset, all pairwise differences between imputation with SRI-EEG and other methods were significant to at least a level of p < 0.01. Exception is SRI-EEG vs MRNN, which was significant to a level of p < 0.05 instead. (2) On the Denoised_2 Hz dataset, all pairwise differences between imputation with SRI-EEG and other methods were significant to at least a level of p < 0.05. (3) On the Noisy_1Hz dataset, all pairwise differences between imputation with SRI-EEG and other methods were significant to at least a level of p < 0.01. (4) On the Noisy_2 Hz dataset, all pairwise differences between imputation with SRI-EEG and other methods were significant to at least a level of p < 0.01. Exception is SRI-EEG vs MRNN, which was significant to a level of p < 0.05 instead.

In summary, this study demonstrates that applying imputation for artifact correction can enable using EEG for practical scenarios such as classification. Our proposed EEG imputation technique can effectively recover EEG signals through data imputation and help further improve the post-imputation classification performance.

7. Limitations and Future Work

From the assessments presented in Section 6.1, our method achieves improved performance on imputing detected artifacts and thus demonstrates potential for use in real EEG artifact correction. However, in some scenarios when subjects perform long-term body motions during EEG data collection, the artifacts caused by muscle movements and cardiac activities can exist in a long time window or for the entire trial. These long-range artifacts significantly limit the availability of enough clean data which is required for RNN based imputation approaches. We can observe this issue from the experiments on synthesized artifacts (Section 6.2.2): when the artifact rate reaches 17.5% and beyond, the imputation error increases rapidly. Thus, RNN based imputation methods rely on having enough valid signal to predict/impute artifacts confidently. In future work, we plan to tackle this limitation inherent to RNNs by exploring the use of reconstruction based denoising approaches that can aid RNN imputation.

We model the spatial dependency of EEG signals as 3D distances between pairs of electrodes. This setup has been proven effective to handle the linearly decayed influence on the EEG cap across channels. This linear decay dependency holds under the assumption that a single event/state is applied to the subject at each EEG trial and the applied stimulus activates a specific region of human brain. When multiple tasks are taken into account, e.g., click the mouse when you see a certain letter on the screen, multiple cortical regions (e.g., the motor cortex and the visual cortex) are activated. In such cases, a simple distance based decay model may fail to capture the complex spatial relationships between electrodes placed above those cortical regions. Future research should address this limitation by using EEG state to support the modeling of EEG channel relationships knowing that channel state and active brain regions are dependent. New work could also explore dynamic capture and computation of electrode relevance given different stimuli.

8. Conclusion

In this paper, we introduced an imputation supported artifact removal algorithm to correct irregular and non-uniformly distributed artifacts that ubiquitously exist in EEG signal data. Our imputation approach enables automatic EEG artifact correction with significantly reduced human effort compared to other methods. Our method is based on a bidirectional LSTM as the backbone network, which has been proven effective by prior work, to learn the temporal dependencies in EEG signal. Into this backbone, we introduce EEG trial state data and spatial correlations between EEG electrodes as supervised information to update the bidirectional hidden states. We evaluated our proposed method via qualitative and quantitative analyses. From the evaluations, we demonstrated that SRI-EEG can greatly improve EEG imputation performance under multiple setups on all the tested datasets. Furthermore, we examined our network design with ablation studies and conducted post-imputation analysis to discuss potential applications that can be supported by imputation based EEG artifact correction. We believe that imputation approaches can improve the effectiveness, robustness and generalization of EEG enabled BCI applications, and see our approach as an important step toward that goal.

Data Availability Statement

The original contributions presented in the study are included in the article/Supplementary Material, further inquiries can be directed to the corresponding author/s.

Author Contributions

The ideas presented in the paper were formulated through a series of meetings to which all authors contributed. All authors contributed to the conception and design of the evaluation. YL organized the database, performed the experiments and analyses, and wrote the first draft of the manuscript. All authors contributed to manuscript revision, read, and approved the submitted version.

Funding

This work was supported in part by NSF award IIS-1845587 as well as ONR awards N00014-19-1-2553 and N00174-19-1-0024.

Conflict of Interest

The authors declare that the research was conducted in the absence of any commercial or financial relationships that could be construed as a potential conflict of interest.

Publisher's Note

All claims expressed in this article are solely those of the authors and do not necessarily represent those of their affiliated organizations, or those of the publisher, the editors and the reviewers. Any product that may be evaluated in this article, or claim that may be made by its manufacturer, is not guaranteed or endorsed by the publisher.

Acknowledgments

We would like to thank Barry Giesbrecht and Tom Bullock for providing and contextualizing the bike dataset and also Yi Ding for valuable discussions. We would also like to thank the reviewers for providing insightful and constructive feedback.

Supplementary Material

The Supplementary Material for this article can be found online at: https://www.frontiersin.org/articles/10.3389/fncom.2022.803384/full#supplementary-material

Footnotes

References