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: This paper studies gradient-based automatic data augmentation algorithms amenable to class-wise policies with exponentially larger search spaces. It presents a method, called  CADDA, to address the problem of automatic data augmentation with application on EEG signals. The proposed method achieves good results for EEG sleep stage classification.