Deep invariant networks with differentiable augmentation layers
Authors: Cédric Rommel
Thomas Moreau
Alexandre Gramfort
Venue: NeurIPS 2022 Accept (2022)
Keywords: invariance learning
data augmentation
automatic data augmentation

Abstract: Designing learning systems which are invariant to certain data transformations is critical in machine learning. Practitioners can typically enforce a desired invariance on the trained model through the choice of a network architecture, e.g. using convolutions for translations, or using data augmentation. Yet, enforcing true invariance in the network can be difficult, and data invariances are not always known a piori. State-of-the-art methods for learning data augmentation policies require held-out data and are based on bilevel optimization problems, which are complex to solve and often computationally demanding. In this work we investigate new ways of learning invariances only from the training data. Using learnable augmentation layers 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 broad class of learning problems beyond computer vision. We provide empirical evidence showing that our approach is easier and faster to train than modern automatic data augmentation techniques based on bilevel optimization, while achieving comparable results. Experiments show that while the invariances transferred to a model through automatic data augmentation are limited by the model expressivity, the invariance yielded by our approach is insensitive to it by design.

Deep invariant networks with differentiable augmentation layers
Cédric Rommel, Thomas Moreau & Alexandre Gramfort Université Paris-Saclay, Inria, CEA, Palaiseau, 91 91120, France (firstname., lastname)@inria...
Abstract
Designing learning systems which are invariant to certain data transformations is critical in machine learning. Practitioners can typically enforce desired invariance on the trained model through the choice of a network architecture, e.g. using convolutions for translations, invariance in the network can be difficult, and data invariances are not always known piori. State-of-the-an methods for learning data augmentation policies require held-out data and are based on bilevel optimization problems, which are complex to solve and often computationally demanding. In this work investigate new ways of learning invariances only from the training data. Using learnable augmentation layers 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 broad class of learning problems beyond computer vision. We provide empirical evidence showing that our approach is easier and faster to train than modern automatic data augmentation techniques based on bilevel optimization, while achieving comparable results. Experiments show that while the invariances transferred to model through automatic data augmentation are limited by the model expressivity, the invariance yielded by our approach is insensitive to it by design.
using data augmentation. Yet, enforcing true
Introduction
Inductive biases encoding known data symmetries are key to make deep learning models gener- alize in high-dimensional settings such as computer vision, speech processing and computational neuroscience, just to name few. Convolutional layers [1] are the perfect illustration of this, as their translation equivariant property allowed to reduce dramatically the size of the hypothesis space compared to fully-connected layers, opening the way for modern computer vision achievements [2]. This illustrates one way to encode desired symmetries deep learning models by hard-coding them directly in the network architecture. An alternative way is to use data augmentation, where transfor- mations encoding the desired symmetries are applied to the training examples, thus adding additional cost when such symmetry is not recognized by the model. While in the first invariances are built in the network by design and are therefore a hard constraint, data augmentation promotes certain invariances more softly. As opposed to being built-in, desired invariances here trained-in [3].
In both cases, the invariances present in the data are not always known beforehand. While relevant invariances are intuitive for some tasks such as object recognition (e.g. a slightly tilted or horizontally flipped picture of a mug, still represents a mug), the same cannot be said for many important predictive tasks such as classifying brain signals into different sleep stages [4, 5]. In order to be able to tackle this problem of learning optimal systems from complex data two strategies are pursued the literature: neural architecture search (NAS) which aims to find the best architectural elements from the training data, and automatic data augmentation (ADA) which aims to learn augmentation policies automatically from given dataset. Both fields tackle the problem very similarly, by parametrizing the
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

network architecture or the augmentation policies, leading to bilevel hyperparameters optimization problem [6]. While these techniques allowed to find architectures and augmentations capable of outperforming the tate-of-the-an in some cases, solving such bilevel optimization problems is often difficult and computationally demanding [7. 8, 9].
This work investigates how to learn data invariances directly from the training data, avoiding the complex bilevel structure of ADA. To this end, we propose to integrate learnable data augmentation layers within the network, and train them together with other model parameters using a novel invariance promoting regularizer. Our approach extends previous works into a very general framework that 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 an image recognition and a sleep stage classification dataset. We show that our approach can correctly select the transformations to which the data invariant, and learn the true range of invariance, even for nonlinear operations. Our experiments also demonstrate that the data augmentation layer proposed here leads to almost perfect built-in invariances irrespective to the complexity of the original network it 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, avoiding tedious bilevel optimization parameters. The accompanying code can be found at https:/ /gi thub. com/ cedri cwwwmel/ /augnet.
2 Related work
Automatic data augmentation Automatic data augmentation aims to learn relevant data invariances which increase generalization power. More precisely, ADA is about searching augmentations that, when applied during the model training. will minimize its validation loss. This objective summarized in the following bilevel optimization problem:
min T
C(0")
S.L. 0" E Era min (())i()))))))
where T an augmentation policy, denotes the parameters of some predictive model, and C()) D) its loss over the set D. Initial ADA approaches such as AutoAugment [6] and PBA [10] use discrete search algorithms to approximately solve (1). Despite the impressive results obtained, they are tremendously costly in computation time, which makes them impractical in many realistic settings. In an attempt to alleviate this limitation, Fast AutoAugment [1 [11] proposes to solve a surrogate density matching 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- pre-trained model and its success highly depend whether the latter was able to learn relevant data invariances on its own. Another way of carrying ADA efficiently is by using gradient-based algorithms, as proposed in Faster Augmmnt [12], DADA [13] and ADDA [14]. While Faster Augment also 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 differentiable relaxations of standard augmentation transformations. However, we are mostly interested in building the learned invariances into the model, which is out of the scope of these methods. Indeed, they require substantially more overhead than our approach since they learn augmentations from the validation 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, and require careful tuning of the outer learning rate [7, 8, 9].
Embedding invariances within neural network architecture A vast literature has focused on encoding predefined invariances or equivariances into neural network architectures. For instance, group convolutions allow to extend traditional convolutional layers to groups of affine transformations other than translations [16]. More related to work, DeepSets [17] encode permutation invariance by summing networks predictions. Although related, these methods are not designed for learning symmetries from the data, which is the objective of this study.
Prior work on this matter from [18] proposes to learn invariances using the marginal likehood in the context of Gaussian processes. In contrast, we are mostly focused on deep neural networks. [3]

---

Augmentation module
Trunk model
Aggregation module
Ti
Z
Ti
mm
f
1mm
T2
T2
T
mmgl
f
Σ
f
To Ta
T
Augmentation layer
T+
Selected transform
T Unselected transform
Figure I: General architecture of AugNet. Input data is copied times and randomiy augmented by the augmentation layers forming the augmentation module. Each copy is then mapped by the trunk model f. whose predictions are averaged by the aggregation module. Parameters both f and the augmentation layers are learned together from the training set.
suggests discover data symmetries and build them into neural networks by learning weight sharing patterns in meta-learning framework. Most related to our work is Augerino [19], which allows to learn invariances affine transformations from the training data, by adding learnable sampling brick to an existing model and averaging its predictions. They propose to learn the range of distributions from which affine transformations are sampled, by parametrizing the corresponding Lie group in terms of its Lie algebra. As it is unclear how to extend Augerino beyond Lie groups such as affine transformations, its scope of application is rather limited, making it mostly taylored for computer vision tasks. In this work we build on their ideas, extending Augerino to more diverse applications with hierarchical differentiable data augmentation layers built into the networks. Moreover, in addition to learning sampling distribution ranges, we also learn to select transformations which encode the strongest data invariance.
AugNet: general framework to embed data augmentation into neural networks
3.1 Preliminaries
Let consider a dataset
DN) of observations sampled from an unknown dis- tribution Px,y over X' y. In a supervised setting, one wants to use this data train model f: + +) (e.g. neural network) so that it can predict from new observations sampled from marginal Px. Now, suppose that the data joint distribution Px invariant to certain group G of transformations acting on X, i.e. for any
G the distribution IP P is close (in some sense) to the transformed distribution 9x.Y (ef. [20] for further details). This means for example that, Y is the class "dog". than the probability of sampling
image X of dog is close the probability of sampling the transformed image gr. this situation, one would like the model f to have this same invariance by design, that it does not have to learn it from scratch and can better generalize to new observations Px. Px,y- In other words, we would like fgz) = f fz) ( for any g E G. One way of achieving this is to average f over G endowed with uniform distribution ve: Proposition 3.1. For a given model f X - y and a compact group of transformations G acting on X, f X - defined as
f(x)) Egrue [f(gx)]
where VG is uniform distribution G. is invariant through the action of G.
The proof follows from the assumption that PG uniform and G group. hence stable by composition and inverse. So for any G, -u is also in G and
f(ux) Ey [f(gux)] Enu- ~wn [f Ex-re [f(hx)] f(x)

---

where h = gu. In practice, we will consider sets of transformations G which do not necessarily form group (e.g. rotations within 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
1 - E f(oor) c=l
f(x)
where ge's 's are sampled from Vg. In practice, using large number C' of sampled transformations would be prohibitive as well. Fortunately, even when C' is small, f is unbiased estimator of Hence stochastic gradient descent can be used to train exactly by minimizing the loss of f.
Based on these observations, we propose to create nearly invariant neural networks made of three blocks:
an augmentation module, which takes an input x: E X, samples C transformations from G, and outputs C transformed copies of x;;
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 AugNet hereafter. Note that unlike standard data augmentation, the transformations distributions is part of the model and kept at inference.
3.3
Augmentation layers
There are two missing elements from the architecture described in the previous section, namely the choice of the of transformations G and how transformations are sampled from it. Indeed, the true invariances present in dataset are often unknown, which is why we propose to learn both the set of transformations and parametrized distribution used sample from them.
In this paper, we are mostly interested in transformations T X + X' defining a data augmentation, such random rotations in image recognition. Most often, such operations transform the data with more or less intensity depending on a parameter called magnitule hereafter. If we take the same example of random rotations, the magnitude can be the maximum angle by which are allowed to rotate the images. In order to have an homogeneous scale for all transformations considered, we assume without loss of generality that magnitudes lie in the interval [O, 1]. with μ being equivalent to the identity (i.c. no augmentation) and μc = 1 being the maximal transformation strength considered. By using the reparametrization trick [21] or some other type of gradient estimator such as straight-through [22] or relax [23]. these transformations can be made differentiable with respect to 1 as shown for example by [12] for image augmentations and by [14] for augmentations of electreencephalggrapyy signals (EEG). One can hence learn the right magnitude to use from the data by backpropagating gradients through
But the correct transformation T describing a data invariance is also supposed to be unknown and needs to be learned in addition its magnitude. Let T (Ti...., To} be discrete set of transformations possibly describing data invariance. We propose to use backpropagation to learn which transformation pick from this set by using a layer consisting of weighted average of all transformations:
AugLayer(z;w,, w.
Σ w. (TTT(
where the weights w, sum to I. practice, we optimize some other hidden weights w' which pass through softmax activation = aw'). As explained in Section 3.4 and demonstrated in our experiments, this layer architecture favors a single transformation describing correct data invariance and tune its magnitude.

---

As the data might be invariant to more than one transformation, these layers can be stacked on top of each other, so that each one learns different transformation (see Figure 1). This is justified by the stability property of invariance for composition. Indeed, if a function h X > 3 is invariant to actions u and g. then must be invariant to their composition: h((u o g)ar) = h( h(((gz))
h(gr)) = h(x)). Hence, we build augmentation modules with sequences of augmentation layers (2) and learn their parameters 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 AugNer with a standard loss such as the cross-entropy. the model tends to find the most unconstrained model: the one with μ 0, which never augments the data. The same was described for Augerino [19), which is why its authors proposed to add regularization R pushing towards broader distributions:
min ((f(z; w. p. p, y) + AR(p) w.p.0
(3)
In the previous equation, 0 denotes the parameters of the trunk neural network f and is loss function. With our
formalism, the regularization proposed by [19] is equiva- lent to penalizing the negative Lo-norm of the magnitude vector: R(u) - |1/4|12. As illustrated in the ablation ex- periments of Section 1. 1, this regularizer is not sufficient to ensure that our augmentation layers do not converge to the identity transform because we have an additional degree of freedom than Augerino: the weights w. Hence, the model may reach low loss values by just maximizing the weight of a single transform w, = with magnitude 14g = 0, while maximizing all other magnitudes Maztn' = 1. Be- cause of this, we propose instead to penalize the norm of the element-wise product of weights and magnitudes: R(w, (4) = - || @ slll2- This has the effect of tying to- gether w. and B&g of the same transformation Tg. avoiding the problem described before.
1.0
Walla
A
W 0
0
1.0
W1
Figure 2: Illustration of the selective property of AugNet's regularizer. As the norm or the vector w O μ increases and J41 reaches the maximal value allowed by the training loss fi, A22 starts getting greater than fl, creating an imbalance in the weights gradients. When #2 reaches Aig the only way to keep increasing the vector's norm is by decreasing wi to- wards O and increasing Wy towards 1.
Another property of this regularization is that it promotes the selection of a single transformation per augmentation layer, as illustrated on Figure 2 for the simple case of n = transformations. Because E, W/g = Wi + W2 = 1, the vector w @ μ (Z2) E [0, 1]2 is bound to move within the line of equation x2 = (12(1
# when the magnitudes j1 and μl2 are fixed. At the beginning of the training, with all weights initialized at 0.5 and magnitudes close to O, neither Ti or T2 harm the training, the vector's norm can grow following the bissector by increasing the magnitudes equally. Once one of the transformations Ti reaches a magnitude e it will start harming the training loss and its gradient will converge to O. This introduces an imbalance between the transformations as U2 keeps increasing, until it reaches its maximal value juy Pı. As w 3 μ is still bound to move within the line 7d2(1 N1 ), the only way to keep increasing its norm is by pointing more and more vertically, until it gets to w22 = 1 and wh = O. This can also be seen from the expression of the gradients of the regularizer:
RR
(w, ( = 2n - 2w (p (1) + m3)
Dw,
which implies that:
V Vw R (w. πi) S Each augmentation layer hence promotes the transformation with the highest admissible magnitude.
4 Experiments with synthetic data
In this section we present experimental results of AugiNez in controlled settings in order to verify empirically some of its properties and compare it to Augerino and standard data augmentations.
5

---

m//4
Rocation Whath
Augerino Aldl
Shift range (Hz]
0.5
AugNet- Alp o w/p
AugNet Alpb
AugNet no reg
-0.5
-104
Iterations
Iterations
(a) Mario&Iggy experiment
(b) Sinusoidal waves experiment
Figure 3: (a) Learned rotation angle during training. Both Augerino [19] (blue) and AugNet (orange) able to learn the correct level of rotation invariance in the data. When the regularization is removed from AugNet (red), it converges to the identity transform (no rotation) as observed in [19]. If we replace the regularizer in AugNet by Augerino's one, AugNet misses the correct rotation angle (green). (b) Learned frequency shift invariance. AugNet is able to learn nonlinear invariances from various types of data.
4.1 Comparison to Augerino
Learning the correct invariance First, reproduced the Mario-Iggy experiment by [19] to show that AugNet can also learn affine invariances from image datasets. The data is generated from two initial images, which
rotated by random angle either between |- /4, π r/4] (labels 0 and 2) or between [-т, - 3T /4] U[3T// /4, T] (labels and 3). This procedure is illustrated
Figure A.1. Hence, labels depend both on the initial image and whether has its head pointing up or down. By design, the data invariant to inputs rotations between [-#/4, /4, и /4), and the experiments objective to verify whether AugNet can learn this invariance. For this, we layer containing 5 geometric augmentations: translate-x, translate-y, rotate, shear-x and shear-y. Note that while these augmentations are also encoded in Augerino, we don't 't use the Lie algebra and exponential maps to implement them here. Maximal translations are set to half the image width/height, maximal rotations correspond to angle of = and maximal shearing coefficients are 0.3. The trunk model used is simple 5-layer convolutional network, whose architecture is described 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.
single augmentation
Ablation study As explained in Section 3.4, the regularizer plays a crucial role both in AugNet and Augerino. can be seen on Figure 3a that without regularization the rotation angle converges Indeed, our model naturally tends to nullify any transformations of the input, similarly to what is shown by [19] for Augerino. Furthermore, we also see that when we replace AugNet's regularizer by Augerino's 's one (-1|x(|2), the model does not learn the correct angle. This happens because this penalty is not sufficient to prevent the model from converging to the identity transform, as shown by the learned weights and magnitudes depicted the bottom of Figure 4. We see indeed that the model maximizes the weight for aranslate- x and minimizes its magnitude, while maximizing the magnitude of other transformations which have very small weights. In contrast, of Figure 4 that the regularizer we propose allows to both select the rotate operation and correctly tune its magnitude.
at the top
4.2 Application to problems beyond computer vision
In this experiment we demonstrate that AugNet is also applicable to other types of data and can learn non-linear transformations which go beyond the scope of Augerino. For this, we create dataset consisting of four classes corresponding to 4 generating frequencies: w 2, 4, and Hz. For each of these frequencies, 10 seconds long sinusoidal waves of unit amplitude are created with frequencies sampled uniformly between w 0.5 Hz. Moreover, the phase of these waves are sampled uniformly and they are corrupted with additive Gaussian white noise with a standard-devittio 0.5 (ef. Figure A.4). We generate 400 training examples and 200 test examples. The learning task consists in predicting the correct generating frequency from the noisy signals. Note that, because classes are separated by Hz, the dataset presents an invariance frequency shifts of the inputs between +0.5 Hz.
of

---

Weights
Magnitudes
TranslateX
1.0
1.0
TranslateY
Rotate
0.5
0.5
ShearX
ShearY
TranslateX
1.0
1.0
Translatey
Rotate
0.5
0.5
ShearX
ShearY
Iterations
Iterations
Figure 4: Top: With its standard regularizer, AugNet quickly learns to maximize the probability weight of rotation and adjusts the corresponding magnitude, ignoring other irrelevant transformations. Bottom: With Augerino's 's regularizer, the model may still converge the identity map by maximizing the probability weight of any (irrelevant) transformation and minimizing its magnitude towards O.
In order to learn this invariance, from the EEG literature: frequency shift [14], FT surrogate [24] and Gaussian noise [25]. The trunk model used is simple 3-layer convolutional network described in Section A.2 together with other experimental settings. As shown Figure 3b, AugNet is able to learn the correct range of invariance, and select the correct transformation. This experiment multiple augmentation layers to show that AugNet is robust to the chosen number of layers.
single augmentation layer implementing three augmentations
repeated in Section B.1 with
4.3 Insensitivity model capacity
In this experiment aim demonstrate that models trained with our proposed methodol- ogy are invariant regardless of the capacity of the trunk model f used, which represents an advantage over (automatic) data augmentation. For this, we re-use the sinusoids dataset from Section 4.2 and compare three methods: a baseline model trained directly on the raw data, the same model trained with oracle augmentation T (frequency shift with bounds of +0.5 Hz) and an AugNet using the baseline as trunk model. Unlike in Section 4.2, we use here very small multi-layer perceptron with a varying number of neurons and layers, in order be able assess the impact of the model's 's expressivity.
1.00
Invariance
Baseline
0.25
Oracle augmentation
AugNet
0.00
Test accuracy
0.8
0.7
In addition the generalization power of the trained models, evaluated through their test accuracy,
neurons, layers
neurons.
neurons,
2 neurons, 3 layers
layers
2 layers
interested in how invariant they are to the true symmetry T encoded in the data. To gauge this property for model ff at some input E X. we use the following metric adapted from [26]:
Figure 5: Top: Model invariance (4) the true fre- quency shift in the data across different architectures. We report the median values across the test set, with 75% confidence interval. Data augmentation helps to increase the model invariance, but is limited by its expressivity, even when data invariances are known. AugNet learns be almost perfectly invariant and is insensitive to model capacity. Bottom: AugNet outperforms model trained with oracle augmenta- tion with 3x fewer parameters.
Invr(f(z)) d(f(o), b
(4)
where is the cosine distance and b = d(f(x:) is a bascline distance between randomly shuffled inputs' As applying T to the inputs of an invariant model should leave its outputs unchanged, closer invr(f(x)) is to for various inputs x, the more f is invariant to T.
'eere f refers to the activations before the final softmax.

---

Figure 5 reports the median invariance across the test set of the three models for different number of layers and neurons. It can be seen that data augmentation helps to increase the level of invariance of the baseline model considerably. However, the level of invariance reached for the smallest models lower than for larger models, suggesting that the former cannot encode perfectly the invariance taught to them through data augmentation. This problem would be worse in a realistic setting where augmentations are more complex and imperfectly learned by ADA approaches rather than known a priori. On the contrary, results show that AugNet is capable of learning the correct augmentation from scratch and make the baseline model almost perfectly invariant to it, regardless of its expressivity. Furthermore, the plot at the bottom of Figure suggests that the level of invariance of each model is tightly related to its predictive performance.
5 Experiments with real datasets
5.1 Image recognition
In this experiment we showcase AugNet on a standard image recognition task using the CIFARIO dataset [27]. All models considered in this experiment are trained for 300 epochs over 5 different seeds on a random 80% fraction of the official CIFAR1O training set. The remaining 20% is used as a validation set for early-stopping and choosing hyperparameterr, and the official test set is used for reporting performance metrics. The trunk model used here is pre-activated R [28, 29].
We consider five baselines in this experiment: Augerino [19], Auto [6] (pretrained for 5000h), Ragdent [30], as well as the trunk model trained with and without fixed data augmentations
(horizontal-flip with probability 0.5 and random- crop of 32x32 pixels with padding of 4 pixels [2]). For Augerino, AugNet and the baseline with fixed augmentations, we apply a standard normalization of RGB channels to the augmented examples before feeding them to the trunk model. AugNet's s augmentation module is made of two layers containing the five augmentations detailed in Section 4.1 and third layer with the following four different augmentations: horizontal-flin sample-pairing [31], contrast, br brightness jitter [32]. Both AugNet and Augerino were trained with number of copies set to C = 1 and tested with C 20 take advantage of the invariant averaging architecture. Contrary to [6, 30), we do not add the learned augmentations on top of the fixed baseline transformations. 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 A.
Figure 6a reports the performance across epochs. For a given epoch, we report the test accuracy corre- sponding to the best validation accuracy up to that point in training (equivalent to having stopped the model). The performance of the baseline model with augmentation (86.6
considerably improved by Augerino (91.6 = 0.1%), AutoAugment (92.0 = 0.1%), RandAugment (92.4 = 0.2%) and AugNet (93.2 + 0.4%).. Net consistently outperforms all three baselines over all seeds within a 90% confidence interval. This is probably due to AugNet selecting transformations which cannot be learned by Augerino, such as the horizontal flip (ef. Figure B.7), and to the better invariance of AugNet compared to models trained with data augmentation (cf. Figure B.9). Concerning the baseline with fixed augmentations (93.6 = 0.2%), AugNet reaches a comparable performance. This demon- strates that in a realistic scenario with data from real-world application where augmentations have not been explored as extensively as for CIFARIO. AugNet can be trained end-to-end and lead to per- formance comparable to what we would get after manually trying many augmentation combinations.
5.2 Sleep stage classification
Similarly to Section 4.2, this experiment aims to demonstrate the practical usefulness of AugNet in a realistic setting beyond affine transformations. For this, sleep stage classification task is considered. As most commonly done [33], it consists in assigning to windows of 30s of EEG signals a label among five: Wake (W), Rapid Eye Movement (REM) and Non-REM of depth 1, 2 or (NI, 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 set contained 12 nights. The trunk model used for this experiment is the convolutional network proposed in [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 discrete search methods were considered: AutoAugment [6] and Fast AutoAugment [11]. Additionally, three gradient-basse based methods were also tested: Faster AutoAugment [12], DADA [13] and ADDA [14]. All these methods shared the same policy architecture, consisting in 5 subpolicies made of augmentation transformations (cf. Section A.4). As for the previous experiment, we used two types of augmentation
8

---

0.925
0.836
Test balanced accuracy
0.834
Test accuracy
0.875
0.850
augmentation Fixed augmentation
0.830
0.825
Augerino
0.828
AutsAugment
Fast AA
0.800
RandAugment
0.826
Faster AA
AutoAugment
0.775
0.824
AugNet
AugNett
0.750
0.822
50
100
150
200
250
300
Epoch
GPU
(a) CIFAR10
(b) MASS
Figure 6: (a) Median accuracy CIFARIO, over seeds, where shades represent 90% confidence. AugNet learns invariances from seratch, outperforms automatic search approaches and achieves a performance comparable state-of-the-an fixed data augmentations. (b) Median performance (over folds) on the MASS dataset, as function of the computation time. AugNet comparable in speed and accuracy ADDA, while being trained end-to-end. This figure is also available with error bars in the appendix Figure B.10
layers to implement NetNet: two layers capable of sampling time-frequency transformations, and one layer made of sensors transformations (ef. Table A.4).
Figure 6b presents the balanced accuracy over the test set as function of the computation time. For ADA methods, stopped the search every given number of steps (2 epochs or 5 samplings), used the learned policy to retrain the model from scratch (leading to point Figure 6b) and resumed the search from where it stopped. For each we report the test accuracy of the best retrained model according to the validation accuracy. Also note that some methods (Fast AutoAugment and ADDA) require pretraining of the model. Hence, the reported GPU hours for ADA methods correspond to the sum of pretraining, search and retraining time, which explains why they start with some horizontal offset. We that, given budget of 12 hours, AugNet is able outperform four out of the five state-of-the-art approaches both in speed and accuracy. reaches final performance comparable to the recently proposed ADDA, while requiring considerably less efforts in parameter fine tuning. AugNet only requires to set two hyperparameters: the regularization parameter which only one-dimensional, and the number augmentation layers stacked in the network. In contrast bilevel approaches require to carefully tune jointly the learning rate and the batch size of the validation set for the outer problem. Moreover, while AugNet trained end-to-end once the training set, ADDA re- quires pre-training the model and a final retraining of the model with the learned augmentation policy.
Conclusion
In this paper we propose 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 method incorporate any type of differentiable augmentations and is applicable problems. We show that our approach can correctly select the transformation to which the data is invariant, 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 leads to al- most perfect invariance regardless the model size. computer vision tasks, this advantage allows our method to reach high generalization power without hand-crafted data augmentations. Promising results are also obtained for sleep stage classification, where AugNet outperforms most ADA ap- proaches both in speed and final performance with end-to-end training, avoiding tedious bilevel setup.
broad class of learning
Limitations In order to be able reach top performance, a large number of copies C might be necessary, increasing AugNet's computational complexity at test time (cf. Section B.2). This means that 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-to- end methods, it is still limited to augmentations that can be relaxed into differentiable surrogates and requires the availability AugNet however, as it is standard requirement in ADA [6, 12, 13, 11].
set of possible augmentations T. These limitations are not specific to

---

Acknowledgments and Disclosure of Funding
This work was supported by the BrAIN grant LANR-20-CHIA-MAN and ANR AI-Cog grants (ANR- 20-1ADJ-0002). It was also granted access to the HPC resources of IDRIS under the allocation 2021-AD011012284RR and 2022-AD01101117228 made by GENCI.
References
[1] Y. LeCun, Boser, S. Denker, D. Henderson, R. Howard, Hubbard, and L. Jackel. Backpropa- gation applied to handwritten zip code recognition. Neural Computation, 1:541-551, 1989.
[2] Alex Krizhevsky, Ilya Sutskever, and Geoffre E. Hinton. ImageNet classification with deep convolutional neural networks. Advances neural information processing systems (NeurIPS), 2012.
[3] Allan Zhou, Tom Knowles, and Chelsea Finn. Meta-Learning Symmetries by Reparameterization. In International Conference Learning Representations (ICLR), 2021.
[4] Stanislas Chambon, Mathieu Galtier, Pierrick Arnal, Gilles Wainrib, and Alexandre Gramfort. A deep learning architecture for temporal sleep stage classification using multivariate and multimodal time series. IEEE Transactions on Neural Systems and Rehabilitation Engineering. 26(4):758- 769, 2018
[5] Huy Phan, Oliver Y. Chen. Minh C. Tran, Philipp Koch, Alfred Mertins, and Maarten De Vos. Xsleepnet: Multi-view sequential model for automatic sleep staging. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
[6] Ekin D. Cubuk, Barret Zoph. Dandelion Mane, Vijay Vasudevan, and Quoe Le. AutoAugment: Learning Augmentation Strategies From Data. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, 2019.
[7] Arber Zela, Thomas Elsken, Tonmoy Saikia, Yassine Marrakchi, Thomas Brox, and Frank Hutter. Under- standing and Robustifying Differentiable Architecture Search. In International Conference on Learning Representations (ICLR), 2019.
[8] Xiangning Chen and Cho-Jui Hsieh. Stabilizing Differentiable Architecture Search via Perturbation-banl Regularization. In International Conference on Machine Learning (ICML), 2020.
[9] Miao Zhang, Steven W. Su. Shirui Pan, Xiaojun Chang, Ehsan M Abbasnejad. and Reza Haffari. iDARTS: Differentiable Architecture Search with Stochastic Implicit Gradients. In International Conference on Machine Learning (ICML), 2021.
[10] Daniel Ho, Eric Liang. lon Stoica, Pieter Abbeel, and Xi Chen. Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules. In International Conference on Machine Learning (ICML), 2019.
[11] Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. Fast Automugment. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[12] Ryuichiro Hataya. Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Faster AutoAugment: Learning Augmentation Strategies Using Backpropagation. In Computer Vision ECCV 2020. Springer International Publishing. 2020.
[13] Yonggang Li. Guosheng Hu, Yongtao Wang, Timothy Hospedales, Neil M. Robertson, and Yongxin Yang. DADA: Differentiable Automatic Data Augmentation. In ECCV, 2020.
[14] Cédric Rommel, Thomas Moreau, Joseph Paillard, and Alexandre Gramfort. CADDA: Class- Automatic Differentiable Data Augmentation for EEG Signals. In International Conference Learning Representattons (ICLR), 2022.
[15] Hanxiao Liu. Karen Simonyan, and Yiming Yang. DARTS: Differentiable Architecture Search. In
International Conference Learning Representations (ICLR), 2019.
[16] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning (ICML). PMLR, 2016.
[17] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnahas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[18] Mark van der Wilk, Matthias Bauer, S. John, and James Hensman. Learning Invariances using the Marginal Likelihood. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
[19] Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning Invariances in Neural Networks. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[20] Shuxiao Chen, Edgar Dobriban, and Jane H. Lee. A Theorethe Framework for Data Augmentation. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
10

---

[21] John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient Estimation Using Stochastic Computation Graphs. In Advances in Neural Information Processing Systems (NeurIPS), 2015.
[22] Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. CoRR, 2013.
[23] Will Grathwoll. Dami Choi, Yuhuai Wu. Geoff Roeder, and David Duvenaud. Backpropagation through the Void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations (ICLR), 2018.
[24] Justus C. Schwabedal, John C. Snyder, Ayse Cakmak, Shamim Nemati, and Gari D. Clifford. Addressing Class Imbalance in Classification Problems of Noisy Signals by using Fourier Transform Surrogates. arXiv:180.006673. 2019.
[25] Fang Wang, Sheng-hua Zhong, Jianfeng Peng, Jianmin Jiang, and Yan Liu. Data Augmentation for EFG-Based Emotion Recognition with Deep Convolutional Neural Networks. In MulriMedia Modeling, volume 10705, pages 82-93. Springer International Publishing, 2018. Series Title: Lecture Notes in Computer Science.
[26] Diane Bouchacourt, Mark Ibrahim, and Ari S. Morcos. Grounding inductive biases natural images: invariance stems from variations in data. In Advances in Newral Information Processing Systems (NeurIPS), 2021.
[27] Alex Krizhevsky, Geoffrey Hinton, et Learning multiple layers features from tiny images. 2009.
[28] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identity mappings in deep residual networks. In Computer Vision - ECCV 2016, 2016.
[29] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
[30] Ekin Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V. Le. Randaugment: Practical automated data augmentation with reduced search space. In JEFF/CVV Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pages 3008-3017. IEEF, 2020.
[31] Hiroshi Inoue. Data augmentation by pairing samples for images classification. In IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), 2018.
[32] Andrew G. Howard. Some improvements deep convolutional neural network hased image classification. CoRR, abs/1312.5402, 2014.
[33] C. Iber, S. Ancoli-Israel, A. Chesson, and S. F Quan. The AASM Manual for the Scoring of Sleep and Associated Events: Rules, Terminology and Technical Specification, 2007.
[34] Christian O'reilly, Nadia Gosselin, Julie Carrier, and Tore Nielsen. Montreal archive of sleep studies: an open-access resource for instrument benchmarking and exploratory research. Journal of sleep research, 23(6):628 635, 2014.
[35] Alexandre Gramfort, Martin Luessi, Eric Larson, Denis A. Engemann, Daniel Strohmeier, Christian Brodbeck, Roman Goj. Mainak Jas, Teon Brooks, Lauri Parkkonen, and Matti S. Hämäläinen. MEG and EEG data analysis with MNE-Python. Frontiers in Neuroscience, 7(267):1-13, 2013.
[36] Robin Tibor Schirmeesstrr, Jost Tobias Springenberg, Lukas Dominique Josef Fiederer, Martin Glasstetter, Katharina Eggensperger, Michael Tangermann, Frank Hutter, Wolfram Burgard, and Tonio Ball. Deep lcarning with convolutional neural networks for EEG decoding and visualization. Human Brain Mapping, 2017.
[37] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on learning Representations (ICLR), 2015.
[38] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Reprexentations (ICLR), 2019.
[39] E. Riba, D. Mishkin, F. Rublee D. Ponsa, and G. Bradski. Komia: an Open Source Differentiable Computer Vision Library for PyTorch. In Winter Conference on Applications of Computer Vision, 2020.
[40] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta and Masanori Koyama. Optuna: A Next-generation Hyperparameter Optimization Framework. In Proceedings of the 25th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, 2019.
11

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work existing assets, did you cite the creators? [Yes] Our work based code from [19] and [12], as well as open-source libraries like MNE-PYTHON [35] and BRAINDECODE [36]. All of them are cited the paper and in our open-sourced code.
(b) Did you mention the license of the assets? [No]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're
using/curating? [No]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [No] Note however that only use public datasets and that the EEG dataset used has been anonymized by their creators.
5. If you used crowdsoureing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
12


Reviews:
[{"review_id": "plJubC1tAQ5", "reviewer": "Reviewer_7J9r", "summary": "The author proposed an augmentation layer to learn data invariance from the training data. The proposed method is applicable to any differentiable transformation, and the augmentation layer can learn a weighted sum of different transformations that can best captures the data invariance. Sufficient experiments are conducted to verify the proposed arguments.", "questions": "see Strengths And Weaknesses", "limitations": "There is no negative social impact of this work.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n\nThe paper proposed an idea that the data invariance can be learnt in an end to  end fashion. Greatly eased the learning process comparing to bi-level optimization methods.\nGreat amount of experiments are conducted on both synthetic data and benchmark datasets. Experiments are well designed to demonstrate the effectiveness of proposed methods.\n\nCons:\nOverall, I think the paper shares very similar idea comparing to the Augerino paper.  The main difference is that the proposed method applied learnable weights to multiple transformations. Thus I consider the novelty of this paper as an incremental variation of existing works.\n\nThe author proposed a selective regularizer that is different from the one used in the Augerino work. But the difference of the proposed regularizer is basically injecting the newly proposed learnable weight parameters to the regularizer proposed in Augerino, making the proposed method more like a mirror design of the existing work.\n\nThe author explored the proposed augmentation method beyond computer vision tasks and emphasized the novelty of this extension. I do not see this point as something unique for this proposed method. In fact, there shouldn’t be obstacles for other existing methods to be applied to tasks beyond computer vision. It would be great if the author can further explain on this point if I missed/misunderstand something.\n\nsome minor points:\nin line 123, the author mentioned : Fortunately, even when C is small, f˜ is an unbiased estimator of f. Can the author further explain on the reason?\nFigure 4 could be a bit hard to read, the position of each label could be a bit confusing. Might be better to reformat the figure.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TVL0YmPskKA", "reviewer": "Reviewer_Km4K", "summary": "The paper proposes a method to learn data invariance along with model training. The method avoids architectural modification and bilevel optimization, which makes it easy to use in many scenarios.", "questions": "My primary concern is the claim that the method can recover the *true* data invariance. It ideally would require two steps: (1) all possible types of invariance are considered in the augmentation module, and (2) the distribution is properly learned for each type of invariance. However, these two steps can hardly be satisfied. For the first step, (1a) it is impossible to enumerate all types of invariance; (1b) not all types of invariance are differentiable (e.g., cutoff); (1c) the paper does not discuss how to select the invariances among all possible combinations (e.g., using validation). For the second step, (2a) the scalar amplitude itself is insufficient to characterize an unparameterized distribution, and (2b) there is no theory in the paper that the recovered distribution matches the true distribution (Indeed, the learned distribution degenerates to identical mapping without regularizer). In summary, I can hardly agree that the learned augmentation matches the *true* invariance.\n\nIt also makes the motivation of the paper unclear. One reason to use augmentation is to boost performance --- however, the proposed method still falls behind the fixed augmentation. Another reason to use augmentation is to boost robustness against invariance attack --- however, such robustness is not systematically evaluated in the paper.\n\nMy last minor concern regards computational complexity. As mentioned in the paper, the model in inference requires more than one sample per example. I wonder if the authors could provide an analysis of the tradeoff between computational complexity and model accuracy/uncertainty. ", "limitations": "Not applicable.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The technical details are generally good and easy to follow; however, I am concerned about the motivation and theoretical foundation of the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FTL99AJUi5v", "reviewer": "Reviewer_kEKu", "summary": "The authors propose a module that learns the invariances (and their magnitudes) present in a dataset while training a neural network. The module applies a convex combination of augmentations before passing to a trunk network, where the outputs are aggregated. To prevent trivial augmentations (such as identity), the authors apply regularization to the magnitude predictions of their module. The authors test their module on synthetic and real-world data, comparing to other learned invariance approaches.", "questions": "* Why do the authors not compare to Augerino for the sinusoidal dataset?\n* Is there any way for the authors to apply their method on a dataset where a set of plausible transformations T_Q are not known beforehand? This would be closer to the problem setup of a priori learning of augmentations. \n* How does AugNet compare to related methods on a larger scale image dataset, like CIFAR-10 or ImageNet?\n* Why does AugNet learn to apply a single augmentation at each layer? Wouldn’t it be preferable to capture all invariances using fewer parameters?", "limitations": "* The problem set-up states that the authors are interested in “data invariances [that] are not always known a priori”, but the authors still rely on pre-specifying which augmentations they will apply (denoted as a set T_Q). The propose methodology seems to be limited in scenarios where a set of transformations is not known a priori.\n* Without more experimental comparison and validation, it is difficult to assess the efficacy of the proposed method", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n* The paper is well written and clearly presents its methodology and experiments.\n* The synthetic experiment shows the ability to uncover ground truth invariances from a set, and the Mario&Iggy experiment demonstrates the benefit of the selected magnitude prediction regularization. \n\nWeaknesses\n* The proposed approach is an incremental improvement to that proposed by (Augerino, Benton et al. 2020). Both rely on aggregating the output of a network across augmentation samples to approximate an expectation, and both require some form of regularization to prevent learning identity for the transformation magnitude. \n* The paper is limited in experimental validation. Some experiments do not adequately compare to related methods (e.g., Augerino on the EEG dataset)\n* The authors should reference (“Learning Invariances using the Marginal Likelihood,” van der Wilk et al. 2018). It seems the aggregation module proposed by the author is similar to the Bayesian approach in the mentioned work.\n* Since the authors are comparing to learned augmentation approaches, they are missing many related works that should be cited from that literature. Just a few, for example: (RandAugment, Cubuk et al. 2019), (Faster AutoAugment, Hataya et al. 2019), (Cutout, DeVries et al. 2017), (Adversarial AutoAugment, Zhang et al. 2019), (AugMix, Hendrycks et al. 2020). The authors should at least include a comparison to RandAugment in their experiments, such as Figure 6b.\n* Regarding test accuracy, many recent methods that have been proposed to learn an augmentation policy perform comparable or worse than the simple and computationally quick approach offered by RandAugment. Justification of learning the dataset invariances would be strengthened by also comparing to a network that implicitly learns invariances by training on RandAugment.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "B85nwj8yqQW", "reviewer": "Reviewer_GmuX", "summary": "The paper recognizes the difficulties of current bilevel optimization of automatic data augmentation (ADA) methods. \nThe authors propose AugNet, by adding learnable augmentation layers and an aggregation module to a base trunk model, and making network training and augmentation learning a single level end-to-end optimization problem.\nExperiments across two simulated tasks, one image classification task, and one physiological data classification task are conducted to show the empirical versatility of AugNet.\n", "questions": "While itemized suggestions are laid out in the “weaknesses” section, here’s some higher level suggestions.\n\nOverall I think this work is well motivated and has potential to make an impact if the experimental validation was much more convincing.\n\nTo me the first thing the authors could do without adding additional tasks and models is getting more solid results on existing tasks. As I mentioned in “weaknesses” error bars should be reported in all results to determine whether any effects are statistically significant, this should be the basis of experimental validation. It will also be nice to see stronger empirical performance across the existing tasks. As is, the improvements of AugNet over prior art is marginal and qualitative at best, with significant performance drop for Cifar. Important experiments such as performance vs lambda are missing.\n\nThe paper would also benefit a lot by going beyond the toyish task/model combinations. It is important to see whether AugNet generalizes to larger tasks and models for it to be more useful practically.\n", "limitations": "The authors discussed various limitations of previous work such as the bilevel optimization nature of ADAs. While there are mentions such as how practical transforms deviate from the group assumption, the authors did not have a clear discussion of the limitations of AugNet.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Learning data augmentation policies is an important topic in training generalizable deep neural networks. The authors provide a good review of related work and identify some important limitations of previous methods such as AA, FAA, and ADDA. Thus the topic and motivation of this work are well justified.\n2. Some mathematical intuitions are discussed in the paper in Section 3. Practical approximations are mentioned for Proposition 3.1.\n3. Basic experimental details such as number of epochs, learning rates, and weight decay are generally well documented, choices of hyperparameters are mostly sensible (see weaknesses for exceptions).\n4. Error estimates are provided for Figure 5 Top and Figure 6 (a).\n5. A range of synthetic and real classification tasks are tested.\n\nWeaknesses:\nMajor:\nGenerally, empirical evaluation is weak and unconvincing at several levels which makes it hard to believe the proposed method has much practical advantage over existing methods.\n\n1. Although the tasks are somewhat diverse, they are all small or toy problems. It is a well known trend that data augmentation tends to work better when training data is limited and its performance boost may diminish when the training data is abundant. It is important to empirically show how well AugNet scales to larger (both in size and complexity) datasets and tasks.\n\n2. Similar to tasks, the models used in the experiments have very low capacities (especially for Section 4.3). Not testing AugNet with larger and more realistic models would significantly limit its potential usefulness.\n\n3. For the Mario&Iggy experiments:\n\n  3.1. Qualitatively, both mu regularization and mu-omega regularization work well for Augerino and AugNet as they all roughly learn the optimal Pi/4 rotation.\n\n  3.2. It is unclear how sensitive AugNet is to lambda. Experiments with different lambda are missing.\n\n  3.3. Since C is set to 10 for Section 4.3 and 4 for other experiments, it is reasonable C is at least a somewhat important hyperparameter. \nExperiments to show how C affects performance are missing. \n\n4. The Cifar-10 experiments are very problematic in many ways. \n\n  4.1. AugNet simply performs poorly by falling behind even the base augmentation that was developed about a decade ago (from the AlexNet paper) which has 0 GPU hours of computation cost since it was handcrafted. Spending 1 extra GPU hour to learn an inferior augmentation policy is hardly a win.\n\n  4.2. The gap in top-1 accuracy between AugNet (91.9%) vs Fixed Augment (93.6%) is significant as it could correspond to a 2x or 3x model size difference, this shouldn’t be understated.\n\n  4.3. Since AugNet is supposed to improve upon methods like AA, it is natural to show results with AA and RA, since Pytorch has official implementations of searched policies on Cifar and ImageNet. These methods should significantly boost performance further.\n\n  4.4. Figure B.5 shows that layer 2 learns TraslateY but with magnitude 0, which means identity transform, which suggests the failure of the proposed regularization. The same issue can be found in Figure B.6 layer 1.\n\n  4.5. Additionally, other learned invariances with non-zero weights have magnitudes of 1, which also suggests that AugNet is only learning binary choices for both weights and magnitudes in this task. That also contradicts with what the authors propose regarding AugNet learning both weights and (continuous) magnitude.\n\n5. Most accuracies in the MASS task differ by less than 1%, it is hard to evaluate if any median trends observed are statistically significant without error bars.\n\n6. The limitations of this work are not properly discussed. \n\nMinor:\nOverall I think some math can be more concise or moved to the appendix, as this paper does not focus on proposing new theories. This will free up some space to discuss the limitations of this work and potentially more experimental evaluations of AugNet.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept