Deep invariant networks with differentiable augmentation layers
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Designing learning systems which are invariant to certain data transformations iscritical in machine learning.
Practitioners can typically enforce a desired invarianceon the trained model through the choice of a network architecture, e.g. usingconvolutions for translations, or using data augmentation.
Yet, enforcing trueinvariance in the network can be difﬁcult, and data invariances are not always knowna piori .
State-of-the-art methods for learning data augmentation policies requireheld-out data and are based on bilevel optimization problems, which are complex tosolve and often computationally demanding.
In this work we investigate new waysof learning invariances only from the training data.
Using learnable augmentationlayers built directly in the network, we demonstrate that our method is very versatile.
It can incorporate any type of differentiable augmentation and be applied to a broadclass of learning problems beyond computer vision.
We provide empirical evidenceshowing that our approach is easier and faster to train than modern automaticdata augmentation techniques based on bilevel optimization, while achievingcomparable results.
Experiments show that while the invariances transferred to amodel through automatic data augmentation are limited by the model expressivity,the invariance yielded by our approach is insensitive to it by design.
Introduction
Inductive biases encoding known data symmetries are key to make deep learning models generalize in high-dimensional settings such as computer vision, speech processing and computationalneuroscience, just to name a few.
Convolutional layers [1] are the perfect illustration of this, astheir translation equivariant property allowed to reduce dramatically the size of the hypothesis spacecompared to fully-connected layers, opening the way for modern computer vision achievements [2].
This illustrates one way to encode desired symmetries in deep learning models by hard-coding themdirectly in the network architecture.
An alternative way is to use data augmentation, where transformations encoding the desired symmetries are applied to the training examples, thus adding additionalcost when such symmetry is not recognized by the model.
While in the ﬁrst case invariances are builtin the network by design and are therefore a hard constraint, data augmentation promotes certaininvariances more softly.
As opposed to being built-in , desired invariances are here trained-in [3].
In both cases, the invariances present in the data are not always known beforehand.
While relevantinvariances are intuitive for some tasks such as object recognition (e.g. a slightly tilted or horizontallyﬂipped picture of a mug, still represents a mug), the same cannot be said for many importantpredictive tasks such as classifying brain signals into different sleep stages [4, 5].
In order to be ableto tackle this problem of learning optimal systems from complex data two strategies are pursued in theliterature: neural architecture search (NAS) which aims to ﬁnd the best architectural elements fromthe training data, and automatic data augmentation (ADA) which aims to learn augmentation policies automatically from a given dataset.
Both ﬁelds tackle the problem very similarly, by parametrizing thenetwork architecture or the augmentation policies, leading to a bilevel hyperparameters optimizationproblem [6].
While these techniques allowed to ﬁnd architectures and augmentations capable ofoutperforming the state-of-the-art in some cases, solving such bilevel optimization problems is oftendifﬁcult and computationally demanding [7, 8, 9].
This work investigates how to learn data invariances directly from the training data, avoiding thecomplex bilevel structure of ADA.
To this end, we propose to integrate learnable data augmentationlayers within the network, and train them together with other model parameters using a novelinvariance promoting regularizer.
Our approach extends previous works into a very general frameworkthat goes beyond computer vision tasks, as it can incorporate any type of differentiable augmentations.
We demonstrate the versatility of our method on two well controlled simulated settings, as well as animage recognition and a sleep stage classiﬁcation dataset.
We show that our approach can correctlyselect the transformations to which the data is invariant, and learn the true range of invariance, evenfor nonlinear operations.
Our experiments also demonstrate that the data augmentation layer proposedhere leads to almost perfect built-in invariances irrespective to the complexity of the original networkit is added to.
Moreover, we are able to achieve comparable performance and speed as state-of-the-art
ADA approaches on our sleep staging experiment with a completely end-to-end approach, avoidingtedious bilevel optimization parameters.
Related work
Automatic data augmentation Automatic data augmentation aims to learn relevant data invarianceswhich increase generalization power.
More precisely, ADA is about searching augmentationsthat, when applied during the model training, will minimize its validation loss.
This objective issummarized in the following bilevel optimization problem:
[Equation]
where T is an augmentation policy, ✓ denotes the parameters of some predictive model, and L ( ✓ | D )its loss over the set D .
Initial ADA approaches such as AutoAugment [6] and PBA [10] use discretesearch algorithms to approximately solve (1).
Despite the impressive results obtained, they aretremendously costly in computation time, which makes them impractical in many realistic settings.
Inan attempt to alleviate this limitation, Fast AutoAugment [11] proposes to solve a surrogate densitymatching problem, which breaks the bilevel structure of (1).
It is hence substantially faster to solve,since it does not require to train the model multiple times.
However, this method needs a pre-trainedmodel and its success highly depend on whether the latter was able to learn relevant data invarianceson its own.
Another way of carrying ADA efﬁciently is by using gradient-based algorithms, asproposed in Faster AutoAugment [12], DADA [13] and ADDA [14].
While Faster AutoAugmentalso tackles a surrogate density matching problem, DADA and ADDA solve the bilevel problem (1)directly.
These ADA methods are the most related to our work since we rely on the same differentiablerelaxations of standard augmentation transformations.
However, we are mostly interested in buildingthe learned invariances into the model, which is out of the scope of these methods.
Indeed, theyrequire substantially more overhead than our approach since they learn augmentations from thevalidation set and need to retrain the model after the augmentation search is completed.
Moreover,
DADA and ADDA are based on an alternating optimization of the inner and outer problems [15],which suffer from noisy hypergradients on the outer level due to the stochastic inner problem, andrequire careful tuning of the outer learning rate [7, 8, 9].
Embedding invariances within neural network architecture A vast literature has focused onencoding predeﬁned invariances or equivariances into neural network architectures.
For instance,group convolutions allow to extend traditional convolutional layers to groups of afﬁne transformationsother than translations [16].
More related to our work, DeepSets
Augmentation layer
T q Selected transform
T q
Unselected transform
T
T
T Q …
T
T
T Q …
T
T
T Q …
Augmentation module ∑
Aggregation module Trunk model f f f suggests to discover data symmetries and build them into neural networks by learning weight sharingpatterns in a meta-learning framework.
Most related to our work is Augerino [19], which allows tolearn invariances to afﬁne transformations from the training data, by adding a learnable sampling brickto an existing model and averaging its predictions.
They propose to learn the range of distributionsfrom which afﬁne transformations are sampled, by parametrizing the corresponding Lie group interms of its Lie algebra.
As it is unclear how to extend Augerino beyond Lie groups such as afﬁnetransformations, its scope of application is rather limited, making it mostly taylored for computervision tasks.
In this work we build on their ideas, extending Augerino to more diverse applicationswith hierarchical differentiable data augmentation layers built into the networks.
Moreover, inaddition to learning sampling distribution ranges, we also learn to select transformations whichencode the strongest data invariance.
AugNet: A general framework to embed data augmentation into neuralnetworks3.1 Preliminaries
Let us consider a dataset ( x 1 , y 1 ) , . . .
, ( x N , y N ) of observations sampled from an unknown distribution P X,Y over X ⇥ Y .
In a supervised setting, one wants to use this data to train a model 102 f :
X ! Y (e.g. a neural network) so that it can predict y from new observations x sampled fromthe marginal P X .
Now, suppose that the data joint distribution P X,Y is invariant to a certain group
G of transformations acting on X , i.e. for any g 2 G the distribution P X,Y is close (in some sense)to the transformed distribution P gX,Y ( cf.
[20] for further details).
This means for example that, if
Y =
y is the class "dog", than the probability of sampling an image X = x of a dog is close to theprobability of sampling the transformed image gx .
In this situation, one would like the model f
tohave this same invariance by design, so that it does not have to learn it from scratch and can bettergeneralize to new observations of P X,Y .
In other words, we would like f ( gx ) = f
( x ) for any g 2 G .
One way of achieving this is to average f over G endowed with a uniform distribution ⌫ G :
Proposition 3.1.
For a given model f : X ! Y and a group of transformations G acting on X ,¯ f : X ! Y deﬁned as
[Equation]
where ⌫ G is a uniform distribution on G , is invariant through the action of G .
The proof follows from the assumption that ⌫ G is uniform and G is a group, hence stable bycomposition and inverse.
So for any u 2 G , gu  1 is also in G and
[Equation]
where h = gu .
In practice, we will consider sets of transformations G which do not necessarily forma group (e.g. rotations within a given range).
Because of this, ¯ f is only approximately invariant,although it can get very close to perfect invariance, as shown in our experiments from Section 4.3.
3.2 Architecture: augment, forward and aggregate
As the averaged model (3.1) is intractable, one can approximate it with an empirical average
[Equation]
where g c ’s are sampled from ⌫ G .
In practice, using a large number C of sampled transformationswould be prohibitive as well.
Fortunately, even when C is small, ˜ f is an unbiased estimator of ¯ f .
Hence stochastic gradient descent can be used to train exactly ¯ f by minimizing the loss of ˜ f .
Based on these observations, we propose to create nearly invariant neural networks made of threeblocks:
• an augmentation module , which takes an input x 2 X , samples C transformations from G ,
and outputs C transformed copies of x ;
• a trunk model f , which can be any neural network mapping transformed inputs to predictions;
• and an aggregation module , which is responsible for averaging the C predictions.
This general architecture is illustrated on Figure 1 and referred to as AugNet hereafter.
Note thatunlike standard data augmentation, the transformations distributions is part of the model and kept atinference.
3.3 Augmentation layers
There are two missing elements from the architecture described in the previous section, namely thechoice of the set of transformations G and how transformations are sampled from it.
Indeed, the trueinvariances present in a dataset are often unknown, which is why we propose to learn both the set oftransformations and a parametrized distribution used to sample from them.
In this paper, we are mostly interested in transformations T : X ! X deﬁning a data augmentation,such as random rotations in image recognition.
Most often, such operations can transform the datawith more or less intensity depending on a parameter µ , called magnitude hereafter.
If we takethe same example of random rotations, the magnitude can be the maximum angle by which weare allowed to rotate the images.
In order to have an homogeneous scale for all transformationsconsidered, we assume without loss of generality that magnitudes lie in the interval [0 , 1] , with µ = 0being equivalent to the identity (i.e. no augmentation) and µ = 1 being the maximal transformationstrength considered.
By using the reparametrization trick [21] or some other type of gradient estimatorsuch as straight-through [22] or relax [23], these transformations can be made differentiable withrespect to µ as shown for example by [12] for image augmentations and by [14] for augmentations ofelectroencephalography signals (EEG).
One can hence learn the right magnitude to use from the databy backpropagating gradients through T .
But the correct transformation T describing a data invariance is also supposed to be unknownand needs to be learned in addition to its magnitude.
Let T =
{ T 1 , . . .
, T Q } be a discrete set oftransformations possibly describing a data invariance.
We propose to use backpropagation to learnwhich transformation to pick from this set by using a layer consisting of a weighted average of alltransformations:
[Equation]
where the weights w q sum to 1.
In practice, we optimize some other hidden weights w 0 which passthrough a softmax activation w =  ( w 0 ) .
As explained in Section 3.4 and demonstrated in ourexperiments, this layer architecture favors a single transformation describing a correct data invarianceand tune its magnitude.
As the data might be invariant to more than one transformation, these layers can be stacked on top ofeach other, so that each one learns a different transformation (see Figure 1).
This is justiﬁed by thestability property of invariance for composition.
Indeed, if a function h : X ! Y is invariant to actionsu and g , then it must be invariant to their composition: h (( u  g ) x )
=
h ( u ( gx ))
=
h ( gx ) =
h ( x ) .
Hence, we build augmentation modules with sequences of augmentation layers (2) and learn theirparameters w, µ from the training set together with the parameters of the trunk model.
3.4 Selective regularizer
As illustrated in our experiments from Section 4.1, if we train AugNet with a standard loss suchas the cross-entropy, the model tends to ﬁnd the most unconstrained model: the one with µ = 0 ,which never augments the data.
The same was described for Augerino [19], which is why its authorsproposed to add a regularization R pushing towards broader distributions:
[Equation]
In the previous equation, ✓ denotes the parameters of thetrunk neural network f and ` is a loss function.
With ourformalism, the regularization proposed by [19] is equivalent to penalizing the negative L 2 -norm of the magnitudevector: R ( µ ) =
k µ k 2 .
As illustrated in the ablation experiments of Section 4.1, this regularizer is not sufﬁcient toensure that our augmentation layers do not converge to theidentity transform because we have an additional degree offreedom than Augerino: the weights w .
Hence, the modelmay reach low loss values by just maximizing the weightof a single transform w q = 1 with magnitude µ q = 0 ,while maximizing all other magnitudes µ q 6 = q 0 = 1 .
Be- 182 cause of this, we propose instead to penalize the normof the element-wise product of weights and magnitudes:
R ( w, µ ) =
k w  µ k 2 .
This has the effect of tying together w q and µ q of the same transformation T q , avoidingthe problem described before.
Another property of this regularization is that it promotesthe selection of a single transformation per augmentationlayer, as illustrated on Figure 2 for the simple case ofn = 2 transformations.
Because P q w q = w 1 + w 2 = 1 ,the vector w  µ = ( x 1 , x 2 ) 2 [0 , 1] 2 is bound to movewithin the line of equation x 2 = µ [-
-]2 (1  x 1 µ 1 ) when themagnitudes µ 1 and µ 2 are ﬁxed.
At the beginning of thetraining, with all weights initialized at 0 . 5 and magnitudesclose to 0 , if neither T 1 or T 2 harm the training, the vector’s norm can grow following the bissectorby increasing the magnitudes equally.
Once one of the transformations T 1 reaches a magnitude ¯ µ 1 , itwill start harming the training loss and its gradient will converge to 0.
This introduces an imbalancebetween the transformations as µ 2 keeps increasing, until it reaches its maximal value ¯ µ 2 > ¯ µ 1 .
Asw  µ is still bound to move within the line ¯ µ 2 (1  x 1 ¯ µ 1 ) , the only way to keep increasing its norm isby pointing more and more vertically, until it gets to w 2 = 1 and w 1 = 0 .
This can also be seen fromthe expression of the gradients of the regularizer:
[Equation]
which implies that:
[Equation]
Each augmentation layer hence promotes the transformation with the highest admissible magnitude.
Experiments with synthetic data
In this section we present experimental results of AugNet in controlled settings in order to verifyempirically some of its properties and compare it to Augerino and standard data augmentations.
[Figure]
4.1 Comparison to Augerino
Learning the correct invariance First, we reproduced the Mario-Iggy experiment by [19] to showthat AugNet can also learn afﬁne invariances from image datasets.
The data is generated from twoinitial images, which are rotated by a random angle either between [  ⇡ / 4 , ⇡ / 4] (labels 0 and 2)or between [  ⇡ ,  3 ⇡ / 4] S [3 ⇡ / 4 , ⇡ ] (labels 1 and 3).
This procedure is illustrated in Figure A.1.
Hence, labels depend both on the initial image and on whether it has its head pointing up or down.
Bydesign, the data is invariant to inputs rotations between [  ⇡ / 4 , ⇡ / 4] , and the experiments objectiveis to verify whether AugNet can learn this invariance.
For this, we use a single augmentationlayer containing 5 geometric augmentations: translate-x , translate-y , rotate , shear-x andshear-y .
Note that while these augmentations are also encoded in Augerino, we don’t use the Liealgebra and exponential maps to implement them here.
Maximal translations are set to half the imagewidth/height, maximal rotations correspond to an angle of ± ⇡ and maximal shearing coefﬁcients are0 .
The trunk model used is a simple 5-layer convolutional network, whose architecture is describedin Section A.1 together with other experimental details.
As shown on Figure 3a, both Augerino (blue)and AugNet (orange) are capable of learning the correct angle of invariance from the data.
Ablation study As explained in Section 3.4, the regularizer plays a crucial role both in AugNet and
Augerino.
It can be seen on Figure 3a that without regularization the rotation angle converges to 0 .
Indeed, our model naturally tends to nullify any transformations of the input, similarly to what isshown by [19] for Augerino.
Furthermore, we also see that when we replace AugNet’s regularizerby Augerino’s one ( k µ k 2 ), the model does not learn the correct angle.
This happens because thispenalty is not sufﬁcient to prevent the model from converging to the identity transform, as shownby the learned weights and magnitudes depicted at the bottom of Figure 4.
We see indeed that themodel maximizes the weight for translate-x and minimizes its magnitude, while maximizing themagnitude of other transformations which have very small weights.
In contrast, we can see at the topof Figure 4 that the regularizer we propose allows to both select the rotate operation and correctlytune its magnitude.
4.2 Application to problems beyond computer vision
In this experiment we demonstrate that AugNet is also applicable to other types of data and canlearn non-linear transformations which go beyond the scope of Augerino.
For this, we create adataset consisting of four classes corresponding to 4 generating frequencies: ! = 2 , 4 , 6 and 8 Hz.
For each of these frequencies, 10 seconds long sinusoidal waves of unit amplitude are created withfrequencies sampled uniformly between ! ± 0 .
5 Hz.
Moreover, the phase of these waves are sampleduniformly and they are corrupted with additive Gaussian white noise with a standard-deviation of0 .
5 ( cf.
Figure A.4).
We generate 400 training examples and 200 test examples.
The learning taskconsists in predicting the correct generating frequency ! from the noisy signals.
Note that, becauseclasses are separated by 1 Hz, the dataset presents an invariance to frequency shifts of the inputsbetween ± 0 .
5 Hz.
Augerino’s regularizer, the may still converge to the identity map by maximizing the probability weight of any (irrelevant) transformation and minimizing its magnitude towards
In order to learn this invariance, we use a single augmentation layer implementing three augmentationsfrom the EEG literature: frequency shift [14], FT surrogate [24] and Gaussian noise [25].
The trunk model used is a simple 3-layer convolutional network described in Section A.2 togetherwith other experimental settings.
As shown in Figure 3b, AugNet is able to learn the correct range ofinvariance, and select the correct transformation.
This experiment was repeated in Section B.1 withmultiple augmentation layers to show that AugNet is robust to the chosen number of layers.
4.3 Insensitivity to model capacity
[Figure]
We report the median values across the test set, with a 75% conﬁdence interval.
Data augmentation helps to increase the model invariance, but is limited by its expressivity, even when data invariances are known.
be almost invariant
In this experiment we aim to demonstrate thatmodels trained with our proposed methodology are invariant regardless of the capacity ofthe trunk model f used, which represents anadvantage over (automatic) data augmentation.
For this, we re-use the sinusoids datasetfrom Section 4.2 and compare three methods:a baseline model trained directly on theraw data, the same model trained with anoracle augmentation T ( frequency shiftwith bounds of ± 0 .
5 Hz) and an AugNetusing the baseline as trunk model.
Unlikein Section 4.2, we use here a very smallmulti-layer perceptron with a varying numberof neurons and layers, in order to be able toassess the impact of the model’s expressivity.
In addition to the generalization power of thetrained models, evaluated through their testaccuracy, we are interested in how invariantthey are to the true symmetry T encoded inthe data.
To gauge this property for a modelf at some input x 2 X , we use the following 273 metric adapted from [26]:
[Equation]
where d is the cosine distance andb =
d ( f ( x i ) , f ( x j )) is a baseline distance between randomly shufﬂed inputs 1 .
As applying T to the inputs of an invariant model shouldleave its outputs unchanged, the closer Inv T ( f ( x )) is to 1 for various inputs x , the more f is invariantto T .
Figure 5 reports the median invariance across the test set of the three models for different number oflayers and neurons.
It can be seen that data augmentation helps to increase the level of invariance ofthe baseline model considerably.
However, the level of invariance reached for the smallest modelsis lower than for larger models, suggesting that the former cannot encode perfectly the invariancetaught to them through data augmentation.
This problem would be worse in a realistic setting whereaugmentations are more complex and imperfectly learned by ADA approaches rather than known apriori .
On the contrary, results show that AugNet is capable of learning the correct augmentation fromscratch and make the baseline model almost perfectly invariant to it, regardless of its expressivity.
Furthermore, the plot at the bottom of Figure 5 suggests that the level of invariance of each model istightly related to its predictive performance.
Experiments with real datasets
5.1 Image recognition
In this experiment we showcase AugNet on a standard image recognition task using the CIFAR10dataset [27].
All models considered in this experiment are trained for 300 epochs over 5 differentseeds on a random 80% fraction of the ofﬁcial CIFAR10 training set.
The remaining 20% is used as avalidation set for early-stopping and choosing hyperparameters, and the ofﬁcial test set is used forreporting performance metrics.
The trunk model used here is a pre-activated ResNet18 [28, 29].
We consider ﬁve baselines in this experiment: Augerino [19], AutoAugment [6] (pretrained for 5000h), RandAugment [30], as well as the trunk model trained with and without ﬁxed data augmentations( horizontal-flip with probability 0.5 and random-crop of 32x32 pixels with a padding of 4pixels [2]).
For Augerino, AugNet and the baseline with ﬁxed augmentations, we apply a standardnormalization of RGB channels to the augmented examples before feeding them to the trunk model.
AugNet’s augmentation module is made of two layers containing the ﬁve augmentations detailed in Section 4.1 and a third layer with the following four different augmentations: horizontal-flip ,sample-pairing [31], contrast , brightness jitter [32].
Both AugNet and Augerino weretrained with a number of copies set to C = 1 and tested with C = 20 to take advantage of theinvariant averaging architecture.
The reader is referred to Section A.3 for further experimental details,and to Section B.2 for a sensitivity analysis on hyperparameters C and  .
Figure 6a reports the performance across epochs.
For a given epoch, we report the test accuracycorresponding to the best validation accuracy up to that point in training (equivalent to having earlystopped the model).
The performance of the baseline model with no augmentation ( 86 . 6 ± 0 %) isconsiderably improved by Augerino ( 91 . 6 ± 0 . 1 %), AutoAugment ( 92 . 0 ± 0 . 1 %), RandAugment( 92 . 4 ± 0 . 2 %) and AugNet ( 93 . 2 ± 0 .
AugNet consistently outperforms all three baselines overall seeds within a 90% conﬁdence interval.
This is probably due to AugNet selecting transformationswhich cannot be learned by Augerino, such as the horizontal flip ( cf.
Figure B.7), and to thebetter invariance of AugNet compared to models trained with data augmentation ( cf. ...).
Concerningthe baseline with ﬁxed augmentations ( 93 . 6 ± 0 . 2 %), AugNet reaches a comparable performance.
Thisdemonstrates that in a realistic scenario with data from a real-world application where augmentationshave not been explored as extensively as for CIFAR10, AugNet can be trained end-to-end andlead to a performance comparable to what we would get after manually trying many augmentationcombinations.
5.2 Sleep stage classiﬁcation
Similarly to Section 4.2, this experiment aims to demonstrate the practical usefulness of AugNet in arealistic setting beyond afﬁne transformations.
For this, a sleep stage classiﬁcation task is considered.
As most commonly done [33], it consists in assigning to windows of 30 s of EEG signals a labelamong ﬁve: Wake (W), Rapid Eye Movement (REM) and Non-REM of depth 1, 2 or 3 (N1, N2,
N3).
The public dataset MASS - Session 3 [34] is used for this purpose (more details in Section A.4).
As done by [14], both the training and validation sets consist of 24 nights each, and the test setcontained 12 nights.
The trunk model used for this experiment is the convolutional network proposedin [4], whose architecture is detailed in section Section A.4 together with other experimental settings.
In this experiment, we compare the proposed approach to ADA methods.
For this, two discretesearch methods were considered: AutoAugment [6] and Fast AutoAugment [11].
Additionally, threegradient-based methods were also tested: Faster AutoAugment [12], DADA [13] and ADDA [14].
Allthese methods shared the same policy architecture, consisting in 5 subpolicies made of 2 augmentationtransformations ( cf.
Section A.4).
As for the previous experiment, we used two types of augmentation
[Figure]
layers to implement AugNet: two layers capable of sampling time-frequency transformations, andone layer made of sensors transformations ( cf.
Table A.4).
Figure 6b presents the balanced accuracy over the test set as a function of the computation time.
For
ADA methods, we stopped the search every given number of steps (2 epochs or 5 samplings), usedthe learned policy to retrain the model from scratch (leading to a point in Figure 6b) and resumed thesearch from where it stopped.
For each run, we report the test accuracy of the best retrained modelaccording to the validation accuracy.
Also note that some methods (Fast AutoAugment and ADDA)require a pretraining of the model.
Hence, the reported GPU hours for ADA methods correspond tothe sum of pretraining, search and retraining time, which explains why they start with some horizontaloffset.
We see that, given a budget of 12 hours, AugNet is able to outperform four out of the ﬁvestate-of-the-art approaches both in speed and accuracy.
It reaches a ﬁnal performance comparableto the recently proposed ADDA, while requiring considerably less efforts in parameter ﬁne tuning.
AugNet only requires to set two hyperparameters: the regularization parameter  which is onlyone-dimensional, and the number augmentation layers stacked in the network.
In contrast bilevelapproaches require to carefully tune jointly the learning rate and the batch size of the validation set forthe outer problem.
Moreover, while AugNet is trained end-to-end once on the training set, ADDA requires pre-training the model and a ﬁnal retraining of the model with the learned augmentation policy.
Conclusion
In this paper we propose a new method coined AugNet to learn data invariances from the training data,thanks to differentiable data augmentation layers embedded in a deep neural network.
Our methodcan incorporate any type of differentiable augmentations and is applicable to a broad class of learningproblems.
We show that our approach can correctly select the transformation to which the data isinvariant, and learn the true range of invariance, even for nonlinear operations.
While (automatic)data augmentation is limited by the capacity of the model to encode symmetries, our approach leadsto almost perfect invariance regardless of the model size.
On computer vision tasks, this advantageallows our method to reach high generalization power without hand-crafted data augmentations.
Promising results are also obtained for sleep stage classiﬁcation, where AugNet outperforms most
ADA approaches both in speed and ﬁnal performance with end-to-end training, avoiding tediousbilevel setup.
Limitations In order to be able reach top performance, a large number of copies C might benecessary, increasing AugNet’s computational complexity at test time ( cf.
Section B.2).
This meansthat a proper trade-off between performance and speed has to be determined for each use-case.
Moreover, although AugNet is able to learn augmentations in a broader scope than previous end-toend methods, it is still limited to augmentations that can be relaxed into differentiable surrogates and requires the availability of a set of possible augmentations T .
These limitations are not speciﬁc to
AugNet however, as it is a standard requirement in ADA [6, 12, 13, 11].