CADDA: Class-wise Automatic Differentiable Data Augmentation for EEG Signals

1 Introduction . The interest in using deep learning for EEG related tasks has been rapidly growing in the last years , specially for applications in sleep staging , seizure detection and prediction , and brain-computer interfaces ( BCI ) ( Roy et al. , 2019 ) . Data augmentation is a well-known regularization technique , widely used to improve the generalization power of large models , specially in deep learning ( Krizhevsky et al. , 2012 ; Yaeger et al. , 1996 ; Simard et al. , 2003 ) . Not only does it help by synthetically increasing the size of the dataset used for training , it also creates useful inductive biases , as it encodes invariances of the data and the underlying decision function which the model does not have to learn from scratch ( Chen et al. , 2020a ; van der Maaten et al. , 2013 ) . Such invariant transforms are also a key ingredient for stateof-the-art self-supervised learning ( Chen et al. , 2020b ) . Unfortunately , these transforms have to be known a priori and the best augmentations to use often highly depend on the model architecture , the task , the dataset and even the training stage ( Ho et al. , 2019 ; Cubuk et al. , 2020 ) . Manually finding what augmentation to use for a new problem is a cumbersome task , and this motivated the proposal of several automatic data augmentation search algorithms ( Cubuk et al. , 2019 ) . The existing automatic data augmentation literature often focuses on computer vision problems only , and its application to other scientific domains such as neuroscience has been under-explored . Data augmentation is all the more important in this field , as brain data , be it functional MRI ( fMRI ) or electroencephalography ( EEG ) signals , is very scarce either because its acquisition is complicated and costly or because expert knowledge is required for labelling it , or both . Furthermore , while atomic transformations encoding suitable invariances for images are intuitive ( if you flip the picture of a cat horizontally it is still a cat ) , the same can not be said about functional brain signals such as EEG . Hence , automatic data augmentation search could be helpful not only to improve the performance of predictive models on EEG data , but also to discover interesting invariances present in brain signals . Another interesting aspect of data augmentation that has gotten little attention is the fact that suitable invariances often depend on the class considered . When doing object recognition on images , using color transformations during training can help the model to better recognize cars or lamps , which are invariant to it , but will probably hurt the performance for classes which are strongly defined by their color , such as apples or oranges . This also applies to neuroscience tasks , such as sleep staging which is part of a clinical exam conducted to characterize sleep disorders . As most commonly done ( Iber et al. , 2007 ) , it consists in assigning to windows of 30 s of signals a label among five : Wake ( W ) , Rapid Eye Movement ( REM ) and Non-REM of depth 1 , 2 or 3 ( N1 , N2 , N3 ) . While some sleep stages are strongly characterized by the presence of waves with a particular shape , such as spindles and Kcomplexes in the N2 stage , others are defined by the dominating frequencies in the signal , such as alpha and theta rhythms in W and N1 stages respectively ( Rosenberg & Van Hout , 2013 ) . This means that while randomly setting some small portion of a signal to zero might work to augment W or N1 signals , it might wash out important waves in N2 stages and slow down the learning for this class . This motivates the study of augmentations depending on the class . Of course , as this greatly increases the number of operations and parameters to set , handcrafting such augmentations is not conceivable and efficient automatic searching strategies are required , which is the central topic of this paper . Using black-box optimization algorithms as most automatic data augmentation papers suggest seemed unsuitable given the exponential increase in complexity of the problem when separate augmentations for each class are considered . In this paper , we extend the bilevel framework of AutoAugment ( Cubuk et al. , 2019 ) in order to search for class-wise ( CW ) data augmentation policies . First , Section 3 introduces three novel augmentation operations for EEG signals , and Section 4 quantifies on sleep staging and digit classification tasks how CW augmentations can enable gains in prediction performance by exploiting interesting invariances . Then , Section 5 introduces a novel differentiable relaxation of this extended problem which enables gradient-based policy learning . Finally , in Section 6 , we use the EEG sleep staging task in the class-agnostic setting to evaluate our approach against previously proposed gradient-based methods . In the class-wise setting , the CADDA method is compared against gradient-free methods that can suffer significantly from the dimension of policy learning problem . Furthermore , we carry an ablation study which clarifies the impact of each architecture choices that we propose . Our experiments also investigate density matching-based approaches ( Lim et al. , 2019 ; Hataya et al. , 2020 ) in low or medium data regimes . 2 Related Work . EEG Data Augmentation Given the relatively small size of available EEG datasets , part of the community has explored ways of generating more data from existing ones , e.g. , using generative models ( Hartmann et al. , 2018 ; Bouallegue & Djemal , 2020 ) or data augmentation strategies ( e.g. , Roy et al . 2019 ; Yin & Zhang 2017 ; Wang et al . 2018 ) . Here , we give a succinct review which is completed in Appendix B . The reader is referred to Roy et al . ( 2019 ) for a more detailed discussion on previous EEG data augmentation papers . Noise addition is the most straight-forward data augmentation that can be applied to either raw EEG signals ( Wang et al. , 2018 ) or to derived features ( Yin & Zhang , 2017 ) . Adding such transformed samples forces the estimator to learn a decision function that is invariant to the added noise . Other transforms have also been proposed to account for other sources of noise , such as label misalignment with the time shift ( Mohsenvand et al. , 2020 ) , positional noise for the sensors with sensor rotations ( Krell & Kim , 2017 ) or corrupted sensors with channel dropout ( Saeed et al. , 2021 ) . Other data augmentations aim at promoting some global properties in the model . While masking strategies such as time masking , bandstop filter ( Mohsenvand et al. , 2020 ) or sensors cutout ( Cheng et al. , 2020 ) ensure that the model does not rely on specific time segments , frequency bands or sensor , channel symmetry ( Deiss et al. , 2018 ) encourages the model to account for the brain bilateral symmetry . Likewise , the Fourier Transform ( FT ) surrogate ( Schwabedal et al. , 2019 ) consists in replacing the phases of Fourier coefficients by random numbers sampled uniformly from [ 0 , 2π ) . The authors of this transformation argue that EEG signals can be approximated by linear stationary processes , which are uniquely characterized by their Fourier amplitudes . Automatic Data Augmentation Automatic data augmentation ( ADA ) is about searching augmentations that , when applied during the model training , will minimize its validation loss , leading to greater generalization . Let Dtrain and Dvalid denote a training and validation set respectively , and let T be an augmentation policy , as defined in more detail in Section 4 . ADA is about finding algorithms solving the following bilevel optimization problem : min T L ( θ∗|Dvalid ) s.t . θ∗ ∈ arg min θ L ( θ|T ( Dtrain ) ) , ( 1 ) where θ denotes the parameters of some predictive model , and L ( θ|D ) its loss over set D. One of the first influential works in this area is AutoAugment ( Cubuk et al. , 2019 ) , where problem ( 1 ) is solved by fully training multiple times a smaller model on a subset of the training set with different augmentation policies and using the validation loss as a reward function in a reinforcement learning setting . The main drawback of the method is its enormous computation cost . Many alternative methods have been proposed since , differing mainly in terms of search space , search algorithm , and metric used to assess each policy . The first attempts to make AutoAugment more efficient consisted in carrying model and policy trainings jointly , as done with a genetic algorithm in Population-Based Augmentation ( Ho et al. , 2019 ) . A different way of alleviating the computation burden of AutoAugment is proposed in Tian et al . ( 2020 ) . Observing that data augmentation is mostly useful at the end of training , the authors propose to pre-train a shared model close to convergence with augmentations sampled uniformly , and then to use it to warmstart AutoAugment . The previous methods ( Cubuk et al. , 2019 ; Ho et al. , 2019 ; Lim et al. , 2019 ) use proxy tasks with small models and training subsets to carry the search . This idea is challenged in RandAugment ( Cubuk et al. , 2020 ) , where it is shown that optimal augmentations highly depend on the dataset size and model . RandAugment simply samples augmentations uniformly with the same shared magnitude , which can be tuned with a grid-search . Competitive results are obtained on computer vision tasks with this naive policy . A similar approach is proposed in Fons et al . ( 2021 ) , where all possible augmentations are weighted with learnable parameters and used to derive enlarged batches . Density Matching While all previously cited ADA methods try to solve in some sense the original problem ( 1 ) , Fast AutoAugment ( Lim et al. , 2019 ) suggests to solve a surrogate problem , by moving the policy T into the upper-level : min T L ( θ∗|T ( Dvalid ) ) s.t . θ∗ ∈ arg min θ L ( θ|Dtrain ) . ( 2 ) Problem ( 2 ) can be seen as a form of density matching ( Lim et al. , 2019 ) , where we look for augmentation policies such that the augmented validation set has the same distribution as the training set , as evaluated through the lens of the trained model . This greatly simplifies problem ( 1 ) which is no longer bilevel , allowing to train the model only once without augmentation . Computation is massively reduced , yet this simplification assumes the trained model has already captured meaningful invariances . This density matching objective has been later reused in Faster AutoAugment ( Hataya et al. , 2020 ) , where a Wasserstein GAN network ( Arjovsky et al. , 2017 ) is used instead of the trained classifier to assess the closeness between augmented and original data distributions . Gradient-based Automatic Data Augmentation Further efficiency improvements in ADA were obtained by exploring gradient-based optimization . In Online Hyper-parameter Learning ( Lin et al. , 2019 ) and Adversarial AutoAugment ( Zhang et al. , 2019 ) , policies are 0.0 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 Frequency ( Hz ) 10 20 30 40 V 2 /H z ( dB ) Sleep stage N2 Subject 1 Subject 2 Subject 1 w/ 0.5 Hz Freq . shift Figure 1 : Averaged power spectral density of N2 windows from one night sleep of two different subjects from the sleep Physionet dataset ( channel Pz-Oz used here ) ( Goldberger et al. , 2000 ) . We notice that peak frequencies are shifted . Applying a 0.5Hz frequency shift transform to subject 1 leads a power spectrum density more similar to subject 2. modeled as parametrized discrete distributions over possible transformations and updated using the REINFORCE gradient estimator ( Williams , 1992 ) . As such estimators are quite noisy , Faster AutoAugment ( Faster AA ) and DADA ( Li et al. , 2020 ) derive full continuous relaxations of the discrete original formalism of AutoAugment , allowing them to backpropagate directly through policies . We have revisited this idea in this work . Class-dependent Data Augmentation While class-dependent data generation has been studied in the GAN literature ( Mirza & Osindero , 2014 ) , to our knowledge , Hauberg et al . ( 2016 ) is the only work that has explored class-dependent data augmentations . It presents a method for learning a distribution of CW spatial distortions , by training a model to find the C1-diffeomorphism allowing to transform one example into another within the same class . Although the authors state it is applicable to other domains , it is only demonstrated for digit classification and its extension to other frameworks seems non-trivial . The main difference with our work is that Hauberg et al . ( 2016 ) learns transformations from scratch , while we try to learn which one to pick from a pool of existing operations and how to aggregate them .

Summary: The paper proposes an automatic differentiable data augmentation algorithm for EEG data that outperforms existing methods. They also propose novel augmentations for EEG that help the model to train better in low-labeled data regimes. They also show preliminary results showcasing that class-wise augmentation can be better than class agnostic augmentations for EEG data.