Automatic Differentiable Data Augmentation for EEG Signals in a Class-wise Setting
Anonymous authors Paper under double-blind review
Abstract
Data augmentation is a key element of deep learning pipelines, as it informs the network during training about transformations of the input data that keep the label unchanged.
Manually ﬁnding adequate augmentation methods and parameters for a given pipeline is however rapidly cumbersome.
In particular, while intuition can guide this decision for images, the design and choice of augmentation policies remains unclear for more complex types of data, such as neuroscience signals.
Besides, class-dependent augmentation strategies have been surprisingly unexplored in the literature, although it is quite intuitive: changing the color of a car image does not change the object class to be predicted, but doing the same to the picture of an orange does.
This paper investigates gradient-based automatic data augmentation algorithms amenable to class-wise policies with exponentially larger search spaces.
Motivated by supervised learning applications using EEG signals for which good augmentation policies are mostly unknown, we propose a new diﬀerentiable relaxation of the problem.
In the class-agnostic setting, results show that our new relaxation leads to optimal performance with faster training than competing gradient-based methods, while also outperforming gradient-free methods in the class-wise setting.
This work proposes also novel diﬀerentiable augmentation operations relevant for sleep stage classiﬁcation.
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 ﬁnding 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 scientiﬁc domains such as neuroscience has been under-explored.
Data augmentation is all the more important in this ﬁeld, 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 ﬂip the picture of a cat horizontally it is still a cat), the same cannot 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 deﬁned 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 ﬁve: 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 deﬁned 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 eﬃcient 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 quantiﬁes on sleep staging and digit classiﬁcation tasks how CW augmentations can enable gains in prediction performance by exploiting interesting invariances.
Then, Section 5 introduces a novel diﬀerentiable 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 suﬀer signiﬁcantly from the dimension of policy learning problem.
Furthermore, we carry an ablation study which clariﬁes 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.
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.
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.
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 speciﬁc 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 coeﬃcients 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 D train and D valid denote a training and validation set respectively, and let T be an augmentation policy, as deﬁned in more detail in Section 4.
ADA is about ﬁnding algorithms solving the following bilevel optimization problem:
[Equation]
where θ denotes the parameters of some predictive model, and L ( θ | D ) its loss over set D .
One of the ﬁrst inﬂuential 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 diﬀerent 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, diﬀering mainly in terms of search space, search algorithm, and metric used to assess each policy.
The ﬁrst attempts to make AutoAugment more eﬃcient consisted in carrying model and policy trainings jointly, as done with a genetic algorithm in Population-Based Augmentation (Ho et al., 2019).
A diﬀerent way of alleviating the computation burden of AutoAugment is proposed in Tian et al.
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:
[Equation]
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 simpliﬁes problem (1) which is no longer bilevel, allowing to train the model only once without augmentation.
Computation is massively reduced, yet this simpliﬁcation 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 classiﬁer to assess the closeness between augmented and original data distributions.
Gradient-based Automatic Data Augmentation Further eﬃciency 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
[Figure]
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 ﬁnd the
C 1 -diﬀeomorphism 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 classiﬁcation and its extension to other frameworks seems non-trivial.
The main diﬀerence 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.
New EEG data augmentations
In addition to the data augmentations described in Section 2, we investigate in this paper three novel operations acting on the time, space and frequency domains.
Time Reverse As a new time domain transformation, we propose to randomly reverse time in certain input examples ( i.e., ﬂip the time axis of the signal).
Our motivation is the belief that most of the sleep stage information encoded in EEG signals resides in relative proportions of frequencies and in the presence of certain prototypical short waves.
Yet, on this last point, it is possible that not all sleep stages are invariant to this transform, as some important waves ( e.g., K-complexes) are asymmetric and are therefore potentially altered by the transformation.
Sign Flip
In the spatial transformations category, we argue that the information encoded in the electrical potentials captured by EEG sensors are likely to be invariant to the polarity of the electric ﬁeld.
Given the physics of EEG, the polarity of the ﬁeld is deﬁned by the direction of the ﬂow of electric charges along the neuron dendrites: are charges moving towards deeper layers of the cortex or towards superﬁcial ones?
As both are likely to happen in a brain region, we propose to augment EEG data by randomly ﬂipping the signals sign (multiplying all channels’ outputs by − 1 ).
This can be interpreted as a spatial transformation, since in some conﬁgurations it corresponds to inverting the main and the reference electrode (e.g. in the Sleep Physionet dataset).
Frequency Shift Reading the sleep scoring manual (Rosenberg & Van Hout, 2013), it is clear that dominant frequencies in EEG signals play an important role in characterising diﬀerent sleep stages.
For instance, it is said that the N2 stage can be recognized by the presence of sleep spindles between 11-16 Hz.
Such frequency range is quite broad, and when looking at the averaged power spectral density of windows in this stage for diﬀerent individuals, one can notice that the frequency peaks can be slightly shifted ( cf.
Figure 1).
To mimic this phenomena we propose to shift the frequencies of signals by an oﬀset ∆ f sampled uniformly from a small range (See Appendix B for details).
Class-wise Data Augmentation
Background on auto augmentation framework We adopt the same framework of deﬁnitions from AutoAugment (Cubuk et al., 2019), which was also reused in Lim et al.
(2019) and Hataya et al.
Let X denote the inputs space.
We deﬁne an augmentation operation O as a mapping from X to X , depending on a magnitude parameter µ and a probability parameter p .
While µ speciﬁes how strongly the inputs should be transformed, p sets the probability of actually transforming them:
[Equation]
Operations can be chained, forming what we call an augmentation subpolicy of length K :
[Equation]
where µ τ and p τ denote the concatenation of the K operations parameters.
To increase the randomness of the augmentation strategy, L subpolicies are grouped into a set called an augmentation policy T , and sampled with uniform probability for each new batch of inputs:
[Equation]
Note that some operations, such as time reverse and sign flip augmentations described in Section 3, may not depend on a magnitude parameter, which reduces the search space.
Novel class-wise subpolicies and policies We introduce in this paper augmentation subpolicies which are conditioned on the input class.
Hence, we deﬁne a class-wise subpolicy as a mapping between inputs X and labels y to an augmentation subpolicy
τ y :
[Equation]
(5) where { τ y : y ∈ Y} is a set of subpolicies for each possible class in the output space Y .
These can be grouped into class-wise policies as in (4).
Manual selection for class-wise augmentations In order to illustrate the interest of CW augmentations, we explore their performances with a manual search on a sleep staging task using the sleep Physionet dataset (Goldberger et al., 2000).
Following the standardization and low-pass ﬁltering of the the data, 30-seconds labelled time windows were extracted from the recordings.
To avoid observing the eﬀects due to class imbalance, we sub-sampled the windows to obtain a balanced dataset.
Indeed, as data augmentation yields more signiﬁcant enhancements in low data regime, it tends to beneﬁt underrepresented classes more than others.
Out of 83 subjects
, 8 were left out for testing and the remaining ones were then split in training and validation sets, with respective proportions of 0.8 and 0.2.
The convolutional neural network from Chambon et al.
(2017) 1 was then trained with Gaussian noise and sign flip augmentations on 250 randomly selected windows (50 per class) from the training set, and evaluated on the test set.
The experiment was repeated 10 times using diﬀerent folds.
Figure 2 (a) reports the improvement of the per-class F1 score for Gaussian noise and sign flip augmentations on the diﬀerent folds.
The Gaussian noise augmentation increases the F1 score for stages W, N1 and N3, while it degrades the median classiﬁcation power for N2 and REM stages.
The sign flip augmentation appears to be less helpful than
Gaussian noise for W, N1 and N3, but it helps on classifying N2 and REM stages.
With this observation in mind, we have carried out the same experiment using a CW subpolicy with
K = 1 (5) which augments W, N1 and N3 stages with Gaussian noise , and N2 and REM stages with sign flip .
The relative improvement in overall balanced accuracy compared to only augmenting with the class-agnostic augmentations is depicted on Figure 2 (b), which shows that a better performance is reached with the CW augmentation.
This manual search experiment suggests that looking for augmentation strategies depending on the label of the data being transformed can be relevant to improve prediction performance, but also to help discovering some interesting invariances for the task considered.
A similar toy experiment with MNIST dataset (LeCun et al., 1998) can be found in Appendix C with automated CW-policy selection using Random Search.
Efficient Automatic Data Augmentation using gradients
Although the previous experiment in Figure 2 motivates the interest in CW data augmentation, it also illustrates how impractical it can be to search such augmentations manually for each new task, dataset and model.
Moreover, considering N µ possible magnitudes, N O operations,
N p probability values and |Y| ( |Y| = 1 in the class-agnostic case) possible classes leads to ( N p × N µ × N O )
L × K ×|Y| possible policies, which becomes very rapidly a huge space to explore, especially with gradient-free algorithms.
This is typically illustrated in Appendix C where even though the setting is quite simple, Random Search requires a very large number of draws to ﬁnd interesting policies.
Therefore, it is unlikely that such algorithms can scale well when considering CW policies, which considerably increase the search space size.
This motivates the exploration of gradient-based search approaches, which can select policies more eﬃciently in such a huge search space.
5. Differentiable Augmentation Policies
Most ADA approaches are based on gradient-free algorithms, mainly because of the discrete structure of augmentation policies (Section 4).
For this reason, Faster AA (Hataya et al., 2020) and DADA (Li et al., 2020) propose continuous relaxations of the AutoAugment policies, inspired by a recent Neural Architecture Search (NAS) method: DARTS (Liu et al., 2019).
Hereafter we build on these relaxation ideas in the EEG setting.
Background on probabilities and magnitudes relaxation Concerning probabilities, equation (3) is relaxed using a Relaxed Bernoulli random variable (Maddison et al., 2017) as in Hataya et al.
(2020) and Li et al.
(2020) ( cf.
Appendix A for further details).
As for magnitudes, while some operations don’t depend on any magnitude parameter, most of the augmentations described in section Section 2 do and are not diﬀerentiable with respect to it.
It is argued in Hataya et al.
(2020) that using a simple straight-through gradient estimator (Bengio et al., 2013) is enough to face this diﬃculty ( i.e., approximating the gradient of O ( X ; µ ) w.r.t µ by 1 ).
In practice, we found that this approach alone did not allow to properly transmit the gradient loss information to tune the magnitudes.
For this reason, we preferred to carry out a case-by-case relaxation of each operations considered, which is detailed in Appendix B.
Novel operations relaxation Operations composing subpolicies also need to be selected over a discrete set.
For this, Faster AA replaces the sequence of operations by a sequence of stages { ˜ O k : k = 1 , . . .
, K } , which are a convex combination of all operations:
[Equation]
6) where w k denote the weights vector of stage k .
These weights are passed through a softmax activation σ η , so that, when parameter η is small, σ η ( w k ) becomes a onehot-like vector. Algorithm
Input : ξ, (cid:15) > 0 , Datasets D train , D valid , Trainable policy T α , Model θ
Result: Policy parameters α while not converged do // compute the unrolled model g θ
=
L ( θ |T α
( D train )) .
backward ( θ ) θ (cid:48) :
= θ − ξg θ // Estimate ∇ α L ( θ (cid:48) | D valid ) g (cid:48) θ
=
L ( θ (cid:48) | D valid ) .
backward ( θ ) g + α = L ( θ + (cid:15)g (cid:48)
θ |T
α ( D train )) .
backward ( α ) g − α
=
L ( θ − (cid:15)g (cid:48) θ |T
α ( D train )) .
backward ( α ) g α = 12 (cid:15) ( g + α − g − α ) //
Update Policy parameters α α = α − ξg α // Update the model parameters θ g θ
=
L ( θ |T α
( D train )) .
backward ( θ ) θ = θ − ξg θ end
We used the same type of architecture, but replaced the softmax σ η by a straight-through Gumbel-softmax distribution (Jang et al., 2017) parametrized by { w k } .
As shown in experiments of Section 6, this allows to gain in eﬃciency, as we only sample one operation at each forward pass and do not need to evaluate all operations each time.
Furthermore, given that the original Gumbel-softmax gradients are biased, we use the unbiased RELAX gradient estimator (Grathwohl et al., 2018).
The same idea is used in DADA (Li et al., 2020), except that they consider a diﬀerent policy architecture where they sample whole subpolicies instead of operations within subpolicies.
These architecture choices are compared in our ablation study in Section 6 and
Appendix F.
5. Policy optimization
While Faster AA casts a surrogate density matching problem (2), which is easier and faster to solve then (1), Tian et al.
(2020) presents empirical evidence suggesting that these metrics are not really correlated.
Hence, we propose to tackle directly the bilevel optimization from (1) as done in DADA and DARTS.
Let α be the vector grouping all the parameters of a continuously relaxed policy model T α
(Section 5.1).
We carry alternating optimization steps of the policy parameters α and the model parameters θ as described in Algorithm 1.
In order to estimate the augmentation policy gradient, we approximate the lower-level of (1) by a single optimization step:
[Equation]
where ξ denotes the learning rate.
By applying the chain rule in Equation (7) we get:
[Equation]
where (cid:15) is a small scalar and θ ±
= θ ± (cid:15) ∇ θ L ( θ (cid:48)
Experiments on EEG data during sleep
Datasets We used the public dataset MASS - Session 3 (O’reilly et al., 2014).
It corresponds to 62 nights, each one coming from a diﬀerent subject.
Out of the 20 available EEG channels, referenced with respect to the A2 electrode, we used 6 (C3, C4, F3, F4, O1, O2).
We also used the standard sleep Physionet data (Goldberger et al., 2000) which contains 153 recordings from 83 subjects.
Here two EEG derivations are available (FPz-Cz and Pz-Oz).
For both datasets, sleep stages were annotated according to the AASM rules (Iber et al., 2007).
The
EEG time series were ﬁrst lowpass ﬁltered at 30 Hz (with 7 Hz transition bandwidth) and
Training set fraction 0.
R e l a t i v e i m p
r
o v e
m e n t
Frequency transforms
Bandstop Filter
FT Surrogate
Frequency Shift 22 -
Training set fraction 0.
Time transforms
Time Masking
Time Reverse
White Noise 22 -
Training set fraction 0.
Sensors transforms
Channel Dropout
Channel Shuffle
Sign Flip standardized.
For MASS they were resampled from 256 Hz to 128 Hz, while the Physionet dataset was kept at 100 Hz.
Further experimental details can be found in Appendix D.
Manual exploration All EEG data augmentation techniques were ﬁrst tested individually on the two datasets without an automatic search strategy, using a batch size of 128.
The
Physionet dataset was split in 5 folds with 16 subjects for Physionet and 12 subjects for
MASS in a test set.
Among the training folds, data were split in 80/20 randomly to obtain a training and a validation set.
To assess the eﬀect of augmentation at diﬀerent data regimes, the training set was sequentially subset using a log2 scale (with stratiﬁed splits).
All augmentations tested had 0.5 probability and a manually ﬁne-tuned magnitude ( cf.
Section H.2).
Results on Physionet are presented on Figure 3, while similar plots for
MASS can be found in the appendix (Figure H.8).
As expected, augmentation plays a major role in the low data regime, where it teaches invariances to the model which struggles to learn them from the small amount of data alone.
Transforms like FT surrogate , frequency shift and time reverse lead to up to 20% performance improvement on the test set for
Physionet and 60% for MASS in extremely low data regimes.
Surprisingly, rotations and time masking do not seem to help a lot.
More eﬃcient gradient-based search As commonly done (Cubuk et al., 2019; Ho et al., 2019;
Lim et al., 2019; Hataya et al., 2020), we considered policies of size L = 5 , made of subpolicies of length K = 2 containing N O = 12 operations ( cf.
Appendix D).
The MASS dataset was used for this experiment.
Both the training and validation sets consisted of 24 nights each, and the test set contained 12 nights.
For each method, we stopped the search every given number of steps (2 epochs or 5 samplings), used the learned policy to retrain from scratch the model (leading to a point in Figure 4) and resumed the search from where it stopped.
For each run, when the ﬁnal retraining validation accuracy improves compared to previous retraining, we report the new test accuracy obtained by the retrained model (otherwise, we keep the previous value).
Figure 4 (a) shows that our novel automatic diﬀerentiable data augmentation method (ADDA) outperforms existing gradient-based approaches both in speed and ﬁnal accuracy.
We also see that the density matching approach (Faster AA) is outperformed by the other solutions that solve directly the bilevel problem (1).
This is also evidenced in Appendix G, where it is demonstrated that density matching approaches suﬀer in low data regimes as the trained classiﬁer is not able to learn relevant invariances on its own (Figure G.7).
In Figure G.4, a comparison with other gradient-free approaches in a class-agnostic setting is also provided, where ADDA reaches comparative predictive performance 15 hours earlier than AutoAugment.
Ablation study of diﬀerentiable architecture in the class-agnostic setting In order to better understand previous results, we carried out in the class-agnostic setting an ablation study where we sequentially removed from ADDA (green): the RELAX gradient estimator (red), followed by the Gumbel-softmax sampling (light blue).
We see in Figure 4 (a) that
GPU Hours of search 0.
T e s t b a l a n c e d a cc u r a c y (a)
[Figure]
[Figure]
w/
[Figure]
RELAX
ADDA w/o GS sampling 0.
[Figure]
[Figure] [Figure] [Figure]
[Figure]
GPU Hours of search (
[Figure]
CADDA*
RELAX is mainly responsible for the ﬁnal accuracy increase, as it removes the gradients’ bias induced by the Gumbel-softmax sampling.
We also clearly see from curves in red and light blue that replacing the softmax σ η from the subpolicy stages (6) by a Gumbel-softmax sampling is responsible for the considerable 4x speed-up.
Note that the core diﬀerence between ADDA and DADA is the fact that the former samples operations within each subpolicy, and the latter samples whole subpolicies.
Curves in green and dark blue show that
ADDA converges faster, which means the ﬁrst choice is more eﬃcient.
This might be due to the fact that our architecture expresses the same possibilities with L × K × N O = 120 weights w k , which in our case is less than the ( N O ) K = 144 weights in DADA.
Most importantly, note that during the training DADA will sample and update the probabilities and magnitudes of subpolicies with higher weights more often, while ADDA keeps updating all its subpolicies with equal probability.
This allows ADDA to converge faster to 5 good subpolicies while
DADA can get stuck with only one or two strong ones.
Eﬃcient search in class-wise setting Figure 4 (b) compares the best performing gradientbased approach in the class-wise setting, namely CADDA, against state-of-the-art gradientfree methods (AutoAugment and Fast AutoAugment) optimized with TPE (Bergstra et al., 2011) as implemented in optuna (Akiba et al., 2019).
One can see that CADDA achieves the best performance while being signiﬁcantly faster than AutoAugment.
However, one can observe on this CW setting that CADDA does not improve over ADDA in the class-agnostic case.
This illustrates the diﬃculty of learning CW policies due to the dimensionality of the search space.
We also hypothesize that this could reﬂect the diﬃculty of the simple CNN model considered here to encode the complex invariances promoted by the CW augmentations.
Concluding remarks
Our work explores the problem of automatic data augmentation beyond images which is still rarely considered in the literature.
We provided novel transforms for EEG data, and a state-of-the-art search algorithm – called (C)ADDA – to address this question.
This method yields very promising results for sleep stage classiﬁcation, with up to 40% speed up and >1.4% accuracy increase compared to existing gradient-based methods in a standard setting.
Motivated by the possible impact of CW policies that we illustrated empirically, we proposed a framework for the automatic search of such augmentations.
Our work shows that gradient-based automatic augmentation approaches are possible and necessary in the CW setting due to the size of the search space.
Finally, while CW does not provide the expected performance boost yet, we believe that this framework opens a novel avenue to improve the ﬁeld of data augmentation, in particular for neuroscience applications.
A Details on operations relaxation
Hereafter are some general explanations concerning the continuous relaxation of augmentation operations.
A. Relaxation regarding probabilities
Concerning the probabilities, equation (3) can be rewritten as follows:
[Equation]
where b is sampled from a Bernoulli distribution with probability p .
As the Bernoulli distribution is discrete, equation (10) can be made diﬀerentiable with respect to p by using a Relaxed Bernoulli random variable (Maddison et al., 2017), i.e., by reparametrizing the distribution as a continuous function depending on p and on a uniform random variable (Schulman et al., 2015).
A. Relaxation regarding magnitudes
As for magnitudes, while some operations don’t depend on any magnitude parameter, most of the augmentations described in section Section 2 do and are not diﬀerentiable with respect to it.
It is argued in Hataya et al.
(2020) that using a simple straight-through gradient estimator (Bengio et al., 2013) is enough to face this diﬃculty ( i.e., approximating the gradient of
O ( X ; µ ) w.r.t µ by 1 ).
In practice, we found that this approach alone did not allow to properly transmit the gradient loss information to tune the magnitudes ( cf.
Section A.3).
For this reason, we preferred to carry out a case-by-case relaxation of each operations considered ( cf.
Appendix B).
Overall, the same relaxation principles were used in most of them:
when µ is used to set some sampling operation, a pathwise gradient estimator is used (Schulman et al., 2015), just as for p in Eq.
when some masking is necessary, vector indexing is replaced by element-wise multiplications between the transformed vector and a smooth mask built using sigmoid functions; 3.
the straight-through estimator is only used to propagate the gradient through the permutation operation in channel shuffle .
A. Further explanation concerning Straight Through magnitude gradient estimator
In order to validate the continuous relaxation of each augmentation operation described in
Appendix B, we tried to ﬁt the identity, i.e., we minimized the mean squared error between augmented and unchanged batches of EEG signals:
[Equation]
This allowed us to realize that using a simple straight-through estimator as suggested by
Hataya et al.
(2020) was not enough to learn good magnitude parameters.
Indeed, we observed with many diﬀerent learning rates that µ would be modiﬁed but would not converge to 0 when p was ﬁxed.
This motivated us to propose our own relaxation described in
Section 5.1 and Appendix B.
B Implementation of EEG augmentations
In this section we review the EEG augmentation that have been proposed and describe their precise implementation, including their case-by-case relaxation mentioned in Appendix A.
Their implementation in python is provided in the supplementary material ( braindecode-wip folder).
The most basic form of data augmentation proposed is the addition of Gaussian white noise to EEG signals (Wang et al., 2018) or to derived features (Yin & Zhang, 2017).
This transformation modiﬁes the waveform shapes in the time domain, while moderately distorting the proportions of diﬀerent frequencies in the signals spectra.
It simply adds the same power to all frequencies equally.
It should hence work as an augmentation for tasks where the predictive information is well captured by frequency-bands power ratios, as used for pathology detection in Gemein et al.
Another type of augmentation is time related transformations, such as time shifting and time masking (Mohsenvand et al., 2020) (a.k.a. time cutout (Cheng et al., 2020)).
Both aim to make the predictions more robust, the latter supposing that the label of an example must depend on the overall signal, and the former trying to teach the model that misalignment between human-annotations and events should be tolerated up to some extent.
Similarly, transformations acting purely in the frequency domain have been proposed.
Just as time masking zeros-out a small portion of the signal in the time domain, narrow bandstop ﬁltering at random spectra positions (Cheng et al., 2020; Mohsenvand et al., 2020) attempts to prevent the model from relying on a single frequency band.
Another interesting frequency data augmentation is the FT-surrogate transform (Schwabedal et al., 2019).
It consists in replacing the phases of Fourier coeﬃcients 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.
Another widely explored type of EEG data augmentation are spatial transformations, acting on sensors positions.
For example, sensors are rotated and shifted in Krell & Kim (2017), to simulate small perturbations on how the sensors cap is placed on the head.
Likewise, the brain bilateral symmetry is exploited in Deiss et al.
(2018), where the left and right-side signals are switched.
Sensors cutout ( i.e., zeroing-out signals of sensors in a given zone) is studied in Cheng et al.
(2020), while Saeed et al.
(2021) propose to randomly drop or shuﬄe channels to increase the model robustness to diﬀerences in the experimental setting.
Mixing signals from diﬀerent examples has also been suggested in Mohsenvand et al.
B. Frequency domain transforms
Frequency Shift To shift the frequencies in EEG signals, we carry a time domain modulation of the following form:
[Equation]
with x a = x + j H [ x ] being the analytic signal corresponding to x , where H denotes the
Hibert transform.
At each call, a new shift ∆ f is sampled from a range linearly set by the magnitude µ , where µ = 1 corresponds to the range [0 − 5 Hz).
This value was chosen after carrying data exploration, based on the observed shifts between subjects in the Physionet dataset.
Given that Equation 11 is completely diﬀerentiable on ∆ f , we only had to relax the sampling (using the pathwise derivatives trick (Schulman et al., 2015)) to make it diﬀerentiable w.r.t µ .
More precisely, we deﬁne ∆ f = f max · u , where f max = 5 µ is the frequencies range upper-bound and u is sampled from a uniform distribution over [0 , 1) .
FT Surrogate
In this transform, we compute the Fourier transform of the signal F ( x ) (using fft ) and shift its phase using a random number ∆ ϕ :
[Equation]
We then transform it back to the time domain.
This was reproduced from the authors (Schwabedal et al., 2019) original code.
However, in our implementation we added a magnitude parameter setting the range in which random phases are sampled [0 , ϕ max ) , with ϕ max = 2 π when µ = 1 .
The procedure used to make this transform diﬀerentiable is very similar to frequency shift , 3 with ∆ ϕ parametrized as ϕ max u and ϕ max = 2 πµ .
It requires using Pytorch version
8, which now supports fft diﬀerentiation.
Bandstop Filter This transform was implemented using the FIR notch ﬁlter from MNEPython package (Gramfort et al., 2013) with default parameters.
The center of the band ﬁltered out is uniformly sampled between
Hz and the Nyquist frequency of the signal.
Here, the magnitude µ is used to set the size of the band, with µ = 1 corresponding to 2 Hz.
This value was chosen after some small experiments, where we observed that larger bands degraded systematically the predictive performance on Physionet.
This transformation was not relaxed (and hence not available to the gradient-based automatic data augmentation searchers).
B. Temporal domain transforms
Gaussian Noise This transform adds white Gaussian noise with a standard deviation set by the magnitude µ .
When µ = 1 , the corresponding standard deviation was 0 .
Larger standard deviations would degrade systematically the predictive performance in our manual exploration on Physionet.
The Gaussian noise implementation is straight forward.
It’s relaxation only included the use of pathwise derivatives for the sampling part: we sample the noise from a unit normal distribution and multiply it by the standard deviation σ = 0 .
2 ∗ µ .
Time Masking
In this operation, we sample a central masking time t cut with uniform probability (shared by all channels) and smoothly set to zero a window between t cut ± ∆ t/ 2 , where ∆ t is the masking length.
The latter is set by the magnitude µ , where µ =corresponds to ∆ t = 1 s .
This value was chosen because it corresponds roughly to the length of important sleep-related events.
The sampling part of the operation was made diﬀerentiable as above.
Masking was computed by multiplying the signal by a function valued in [0 , 1] , built with two opposing steep sigmoid functions
[Equation]
where λ was arbitrarily set to 1000.
B. Spatial domain transforms Rotations
In this transform, we use standard sensors positions from a 10-20 montage (using the mne.channels.make_standard_montage function from MNE-Python package (Gramfort et al., 2013)).
The latter are multiplied by a rotation matrix whose angle is uniformly sampled between ± ψ max .
The value of the maximum angle ψ max is determined by the magnitude µ , where µ = 1 corresponds to π 6 radian (value used in Krell & Kim (2017); Cheng et al.
Signals corresponding to each channel are then interpolated towards the rotated sensors positions.
While Krell & Kim (2017); Cheng et al.
(2020) used radial basis functions to carry the interpolation, we obtained better results using a spherical interpolation, which also made more sense in our opinion.
The only part that needed to be relaxed to allow automatic diﬀerentiation was the rotation angle sampling, which was done as before, with the angle ∆ ψ parametrized as ψ max (2 u − 1) , u a uniform random variable in [0 , 1] and ψ max =
π 6 µ .
Channel Dropout
Here, each channel is multiplied by a random variable sampled from a relaxed Bernoulli distribution (Maddison et al., 2017) with probability 1 − µ , where µ is the magnitude.
This allows to set the channel to 0 with probability µ and to have a diﬀerentiable transform.
Channel Shuﬄe
In this operation, channels to be permuted are selected using relaxed Bernoulli variables as in the channel dropout , but with probability µ instead of 1 − µ .
The selected channels are then randomly permuted.
As the permutation operation necessarily uses the channels indices, it is not straight-forward to diﬀerentiate it, which is why we used a straight-through estimator here to allow gradients to ﬂow from the loss to the magnitude parameter µ .
C Illustration of the usefulness of class-wise augmentation on MNIST
Motivation In order to illustrate the potential of CW augmentation, we present in this section a simple example using the MNIST dataset (LeCun et al., 1998).
Intuitively, some common image augmentation operations such as horizontal ﬂips or rotations with large angles should not be helpful for handwritten digits recognition: if you rotate a 9 too much, it might be diﬃcult to distinguish it from a 6.
However, some digits have more invariances than others.
For example, 0’s and 8’s are relatively invariant to vertical and horizontal ﬂips, as well as to 180 degrees rotations.
We used this framework to compare a naive approach of automatic data augmentation search both in a standard and in a CW setting.
Algorithm For simplicity, we used a random search algorithm to look for augmentation policies, with operations’ parameters discretized over a grid.
More precisely, each new candidate subpolicy sampled by the search algorithm is used to train the model from scratch, before evaluating it on the validation set.
In terms of search space and search metric, this is equivalent to what is done in AutoAugment (Cubuk et al., 2019), although the reinforcement learning algorithm is replaced by random sampling.
Search space Let N p , N µ and N O denote the number of possible probability, magnitude values and operations in our setting.
The search space of both standard and CW policies is potentially huge:
( N p × N µ × N O )
L × K ×|Y| , where the number of classes |Y| is replaced by 1 in the standard case.
Given the simplicity of the algorithm used, we tried to reduce the search space as much as possible.
Hence, we only considered subpolicies of length K =(a single operation) and restrained our problem to the classiﬁcation of 4 digits only: 4 , 6 ,and 9 .
A pool of four transformations that only depend on a probability parameter (no magnitude) was used in the search: horizontal and vertical flip , as well as 90 anddegrees rotation (counter-clockwise).
Despite this signiﬁcant simpliﬁcation, the search space size for policies made of L =subpolicies is still immense: around 10 6 for standard policies and 1 .
2 × 10 24 for CW policies.
For this reason, 20,000 trials were sampled in both cases.
As an additional baseline, we also trained the same model without any type of data augmentation.
Experimental setting After the 4-class reduction, the MNIST training set was randomly split into a training set containing 1000 images per class and a validation set of 3000 images per class.
The reduced MNIST test set used to evaluate the ﬁnal performance hasimages per class, as the original one.
All transforms were implemented using torchvision (Marcel & Rodriguez, 2010).
We used the LeNet-1 architecture (LeCun et al., 1998) for the classiﬁer.
It was trained using batches of 64 images and Adam optimizer (Kingma & Ba, 2015) with an initial learning rate of 2 × 10 − 3 , β 1 = 0 .
9 , β 2 = 0 .
999 and no weight decay.
Trainings had a maximum number of 50 epochs and were early-stopped with a patience of 5 epochs.
The code that was used to generate this plots are part of the supplementary material ( mnist folder).
Results As seen on Figure C.1, random search converges to the identity policy (no augmentation).
This was expected as none of the available operations describe an invariance shared by all four digits.
From Figure C.1-right we see that the CW random search does not ﬁnd a good augmentation for digits 4 and 6 as well.
Yet, it is able to recover that digit 8 is invariant to horizontal flip .
Interestingly, 90 degree rotation of digit 9 is also selected with a small probability.
This can be understood, as some people might draw the 9’s leg more horizontally than others, so that rotating up to 90 degrees counter-clockwise still preserves in some cases the picture semantics.
Figure C.1-left shows that the CW algorithm can ﬁnd interesting policies in such a constrained setting, but also that it can outperform the baselines.
Yet, this happens only after 10,000 trials.
[Figure]
sampled
[Figure]
e s
[Figure]
a cc
[Figure]
a c
[Figure]
Standard setting Class-wise setting
No augmentation
H o r .
F li p
V e r .
F li p 90 R o t.
180 R o t.
I den t i t y
All classes
C l a ss               0.25 0 0.
This toy experiment suggests that looking for augmentation strategies depending on the label of the data being transformed can be relevant to improve prediction performance, but also to help discovering some interesting invariances for the task considered.
However, despite the extremely simpliﬁed setting used, the CW search took 14.4 GPU days (Tesla V100 SXM2).
It becomes clear that a more eﬃcient search strategy is required for real-world applications.
D Experimental setting shared across EEG experiments
In all EEG experiments, learning was carried using the convolutional network proposed in Chambon et al.
The architecture can be found on Table 1 and was chosen as it seemed like a good compromise as a simple yet deep and relevant sleep staging architecture.
The ﬁrst layers (1-4) implements a spatial ﬁlter, computing virtual channels through a linear combination of the original input channels.
Then, layers 5 to 9 correspond to a standard convolutional feature extractor and last layers implement a simple classiﬁer.
More details can be found in Chambon et al. Other relevant and recent sleep staging architectures are Perslev et al.
; Jia et al. [Table] [Table] The optimizer used was Adam with a learning rate of 10 − 3 , β 1 = 0 .
and β 2 = 0 .
At most 300 epochs were used for training.
Early stopping was implemented with a patience of 30 epochs.
Balanced accuracy was used as performance metric using the inverse of original class frequencies as balancing weights.
The MNE-Python (Gramfort et al., 2013) and Braindecode software (Schirrmeister et al., 2017) were used to preprocess and learn on the EEG data.
Training was carried on single Tesla V100 GPUs.
The N O =operations considered were: time reverse, sing flip, FT surrogate, frequency shift, bandstop filtering, time masking, Gaussian noise, channel dropout, channel shuffle, channel symmetry and rotations around each cartesian axis.
In automatic search experiments, bandstop filter was not included in the diﬀerentiable strategies (Faster AA, DADA, ADDA and CADDA) because we did not implement a diﬀerentiable relaxation of it.
E Implementation
One will ﬁnd included in this supplementary material the complete code of:
• the MNIST experiment from Appendix C in folder mnist , • the new eeg augmentations from Section 5.1 and Appendix B in folder braindecode-wip , • the experiment from section Section 4 in folder cw-sleep-experiment • the automatic data augmentation implementations from Section 6 and Appendix G in folder eeg-augment .
It will be made publicly available upon publication, and we plan on releasing a package compatible with braindecode (Schirrmeister et al., 2017) and PyTorch .
F Architecture details and comparison between ADDA, DADA and Faster AA
[Figure] sample L uniform probability parametrized operations
DADA their are stages operations is computed
ADDA however operations are sampled with learnable log
As depicted on Figure F.2, DADA policies will sample a subpolicy according to a Gumbelsoftmax distribution parametrized by learnable weights w :
T DADA ( X ) =
τ i ( X ) , with i ∼ Gumbel-Softmax ( { 1 , . . .
, ( N O ) K } ; w ) .
Each subpolicy τ i is just a sequence of operations diﬀerentiable wrt its parameters ( cf.
Appendix A).
However, in ADDA and Faster AA, only a predetermined number of L subpolicies are considered and sampled uniformly:
[Equation]
with i ∼ U ( { 1 , . . .
, L } ) .
In this case, subpolicies ˜ τ
i are sequences of stages ˜
O k , as described in equation (6).
While these stages compute a convex combination of all possible operations in Faster AA, using learnable weights w mapped through a softmax function σ η , ADDA stages carry a diﬀerentiable sample (i.e. σ η is exactly equal to 1 for exactly one out of N O in (6)).
In this case, the weights correspond to the log probabilities of sampling each operation.
G Complementary automatic data augmentation results on EEG sleep staging
Class-wise policies selected by CADDA and other approaches
t i m e r e v e r s
e
s i g n f l i p
F T s u r r o g a t e f r e q u e n
c y s h
i f t b a n d s t
o p f i l t e r t i m e m a s k
i n g
G a u s s i a n n o i s e c h a n n e l d r o p o u t c h a n n e l s h u f f l e c h a n n e l s y m m e t r y r o t a t i o n z a x i s r o t a t i o n y a x
i s r o t a t i o n x a x
i s i d e n t i
[Figure]
Probability of applying each operation in the best policies obtained by the diﬀerent search strategies.
For each method, probabilities of all subpolicies are grouped.
Methods based on diﬀerentiable augmentation policies discover relevant augmentation strategies such as time reverse , sign flip , FT surrogate etc.
CADDA suggests that each class beneﬁts diﬀerently from the diﬀerent strategies.
Complete comparison of gradient-based and gradient-free methods in the classagnostic setting
[Figure]
T e s t b a l a n c e d a cc u r a
[Figure]
[Figure]
a
[Figure]
a n c e d [Figure] cc [Figure]
[Figure]
Impact of data regime on density matching approaches
[Figure]
to capture relevant invariances.
H Further manual exploration details H. Manual search results on MASS dataset
22 - 2Training set fraction 0
R e l a t i v e i m p
r
o v e
m e n t
Frequency transforms
Bandstop Filter
FT Surrogate
Frequency Shift 22 - 2Training set fraction 0
Time transforms
Time Masking
Time Reverse
White Noise 22 - 2Training set fraction 0
Sensors transforms
Channel Dropout
Rotation Y-axis
Rotation Z-axis
Sign Flip
In this section we explain the ﬁne-tuning of magnitudes used for the manual exploration results presented in Section 6.
A ﬁrst training with probability p = 0 .
5 and magnitude µ = 0 .
5 was carried over the same training subsets as in Figure 3 and Figure H.8.
Then we selected the subset where the eﬀects of augmentation were maximized ( 2 − 11 times the original training set for both Physionet and
MASS datasets).
Finally, for each operation, we carried a grid-search over the magnitude parameter (with ﬁxed probability p = 0 . 5 ) with the selected training subsets (over 5-folds, as described in Section 6).
The result of the grid-search can be seen on Figure H.9 and
Figure H.10.
[Figure]
B a l a n c e d t e s t a cc u r a c y
i m p
r
o v e
m e n t
Frequency transforms transform
Bandstop Filter
FT Surrogate
Frequency
[Figure]
[Figure] [Figure] [Figure]
[Figure]
B a l a n c e d t e s t a cc u r a c y
i m p
r
o v e
m e n t
[Figure]
Time
[Figure]
u
[Figure]
[Figure]
[Figure]
[Figure]
B a l a n c e d t e s t a cc u r a c y
i m p
r
o v e
m e n t
[Figure]
transforms
[Figure]
Bandstop
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
B a l a n c e d t e s t a cc u r a c y
i m p
r
o v e
m e n t
[Figure]
t e s t a cc u r a c y
i m p
r
o v e m e n
[Figure]