Original: A UTOMATIC D ATA A UGMENTATION VIA I NVARIANCEC ONSTRAINED L EARNING
Revised: A UTOMATIC D ATA A UGMENTATION VIA I NVARIANCEC ONSTRAINED L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Underlying data structures, such as symmetries or invariances to transformations, are often exploited to improve the solution of learning tasks.
Revised: Underlying data structures, such as symmetries or invariances to transformations, are often exploited to improve the solution of learning tasks.

Original: However, embedding these properties in models or learning algorithms can be challenging and computationally intensive.
Revised: However, embedding these properties in models or learning algorithms can be challenging and computationally intensive.

Original: Data augmentation, on the other hand, induces these symmetries during training by applying multiple transformations to the input data.
Revised: Data augmentation, on the other hand, induces these symmetries during training by applying multiple transformations to the input data.

Original: Despite its ubiquity, its effectiveness depends on the choices of which transformations to apply, when to do so, and how often.
Revised: Despite its ubiquity, its effectiveness depends on the choices of which transformations to apply, when to do so, and how often.

Original: In fact, there is both empirical and theoretical evidence that the indiscriminate use of data augmentation can introduce biases that outweigh its benefits.
Revised: In fact, there is both empirical and theoretical evidence that the indiscriminate use of data augmentation can introduce biases that outweigh its benefits.

Original: This work tackles these issues by automatically adapting the data augmentation while solving the learning task.
Revised: This work tackles these issues by automatically adapting the data augmentation while solving the learning task.

Original: To do so, it formulates data augmentation as an invariance-constrained learning problem and leverages
Revised: To do so, it formulates data augmentation as an invariance-constrained learning problem and leverages

Original: Monte Carlo Markov Chain (MCMC) sampling to solve it.
Revised: Monte Carlo Markov Chain (MCMC) sampling to solve it.

Original: The result is a practical algorithm that not only does away with a priori searches for augmentation distributions, but also dynamically controls if and when data augmentation is applied.
Revised: The result is a practical algorithm that not only does away with a priori searches for augmentation distributions, but also dynamically controls if and when data augmentation is applied.

Original: Our experiments illustrate the performance of this method, which achieves stateof-the-art results in automatic data augmentation benchmarks for CIFAR datasets.
Revised: Our experiments illustrate the performance of this method, which achieves stateof-the-art results in automatic data augmentation benchmarks for CIFAR datasets.

Original: Furthermore, this approach can be used to gather insights on the actual symmetries underlying a learning task.
Revised: Furthermore, this approach can be used to gather insights on the actual symmetries underlying a learning task.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Exploiting the underlying structure of data has always been a key principle in data analysis.
Revised: Exploiting the underlying structure of data has always been a key principle in data analysis.

Original: Its use has been fundamental to the success of machine learning solutions, from the translational equivariance of convolutional neural networks (Fukushima and Miyake, 1982) to the invariant attention mechanism in Alphafold (Jumper et al., 2021).
Revised: Its use has been fundamental to the success of machine learning solutions, from the translational equivariance of convolutional neural networks (Fukushima and Miyake, 1982) to the invariant attention mechanism in Alphafold (Jumper et al., 2021).

Original: However, embedding invariances and symmetries in model architectures is hard in general and when possible, often incurs a high computational cost.
Revised: However, embedding invariances and symmetries in model architectures is hard in general and when possible, often incurs a high computational cost.

Original: This is the case, of rotation invariant neural network architectures that rely on group convolutions, which are feasible only for small, discrete transformation spaces or require coarse undersampling due to their high computational complexity (Cohen and Welling, 2016; Finzi et al., 2020).
Revised: This is the case, of rotation invariant neural network architectures that rely on group convolutions, which are feasible only for small, discrete transformation spaces or require coarse undersampling due to their high computational complexity (Cohen and Welling, 2016; Finzi et al., 2020).

Original: A widely used alternative consists of modifying the data rather than the model.
Revised: A widely used alternative consists of modifying the data rather than the model.

Original: That is, to augment the dataset by applying transformations to samples in order to induce the desired symmetries or invariances during training.
Revised: That is, to augment the dataset by applying transformations to samples in order to induce the desired symmetries or invariances during training.

Original: Data augmentation, as it is commonly known, is used to train virtually all state-of-the-art models in a variety of domains (Shorten and Khoshgoftaar, 2019).
Revised: Data augmentation, as it is commonly known, is used to train virtually all state-of-the-art models in a variety of domains (Shorten and Khoshgoftaar, 2019).

Original: This empirical success is supported by theoretical results showing that, when the underlying data distribution is invariant to the applied transformations, data augmentation provides a better estimation of the statistical risk (Chen et al., 2019; Sannai et al., 2019; Lyle et al., 2020; Shao et al., 2022).
Revised: This empirical success is supported by theoretical results showing that, when the underlying data distribution is invariant to the applied transformations, data augmentation provides a better estimation of the statistical risk (Chen et al., 2019; Sannai et al., 2019; Lyle et al., 2020; Shao et al., 2022).

Original: On the other hand, applying the wrong transformations can introduce biases that may outweigh these benefits (Chen et al., 2019;
Revised: On the other hand, applying the wrong transformations can introduce biases that may outweigh these benefits (Chen et al., 2019;

Original: Shao et al., 2022).
Revised: Shao et al., 2022).

Original: Choosing which transformations to apply, when to do so, and how often, is thus paramount to achieving good results.
Revised: Choosing which transformations to apply, when to do so, and how often, is thus paramount to achieving good results.

Original: However, it requires knowledge about the underlying distribution of the data that is typically unavailable.
Revised: However, it requires knowledge about the underlying distribution of the data that is typically unavailable.

Original: Several approaches to learning an augmentation policy or distribution over a fixed set of transformations exist, such as reinforcement learning (Cubuk et al., 2018), genetic algorithms (Ho et al., 2019), density matching
Revised: Several approaches to learning an augmentation policy or distribution over a fixed set of transformations exist, such as reinforcement learning (Cubuk et al., 2018), genetic algorithms (Ho et al., 2019), density matching

Original: (Lim et al., 2019; Cubuk et al., 2020; Hataya et al., 2020), gradient matching (Zheng et al., 2022), bi-level optimization (Li et al., 2020b; Liu et al., 2021), jointly optimizing over transformations using regularised objectives (Benton et al., 2020) and variational bayesian inference (Chatzipantazis et al., 2021).
Revised: (Lim et al., 2019; Cubuk et al., 2020; Hataya et al., 2020), gradient matching (Zheng et al., 2022), bi-level optimization (Li et al., 2020b; Liu et al., 2021), jointly optimizing over transformations using regularised objectives (Benton et al., 2020), variational bayesian inference (Chatzipantazis et al., 2021) , bayesian model selection (Immer et al., 2022) and alignment regularization (Wang et al., 2022).

Original: Optimization based methods often require computing gradients with respect to transformations (Chatzipantazis et al., 2021;
Revised: Optimization based methods often require computing gradients with respect to transformations (Chatzipantazis et al., 2021;

Original: Li et al., 2020b).
Revised: Li et al., 2020b).

Original: Moreover, several methods resort to computationally intensive search phases, the optimization of auxiliary models, or additional data, while failing to outperform fixed user-defined augmentation distributions (Müller and Hutter, 2021).
Revised: Moreover, several methods resort to computationally intensive search phases, the optimization of auxiliary models, or additional data, while failing to outperform fixed user-defined augmentation distributions (Müller and Hutter, 2021).

Original: In this work, we formulate data augmentation as an invariance-constrained learning problem.
Revised: In this work, we formulate data augmentation as an invariance-constrained learning problem.

Original: The use of invariance overcomes the need to search for a specific distribution over transformations.
Revised: That is, we specify a set of transformations and a desired level of invariance, and recover an augmentation distribution that enables imposing this requirement on the learned model, without explicitly parametrising the distribution over transformations.

Original: In addition, the constrained learning formulation mitigates the potential biases introduced by data augmentation whithout doing away with its potential benefits.
Revised: In addition, the constrained learning formulation mitigates the potential biases introduced by data augmentation without doing away with its potential benefits.

Original: More specifically, we rely on an approximate notion of invariance that is weighted by the probability of each data point.
Revised: More specifically, we rely on an approximate notion of invariance that is weighted by the probability of each data point.

Original: Hence, we require the output of our model to be stable only on the support of the underlying data distribution, and more so on common samples.
Revised: Hence, we require the output of our model to be stable only on the support of the underlying data distribution, and more so on common samples.

Original: By imposing this requirement as a constraint on the learning task and leveraging recent duality results, the amount of data augmentation can be automatically adjusted during training.
Revised: By imposing this requirement as a constraint on the learning task and leveraging recent duality results, the amount of data augmentation can be automatically adjusted during training.

Original: We propose an algorithm that combines stochastic primal-dual methods and MCMC sampling to do away with the need for transformations to be differentiable.
Revised: We propose an algorithm that combines stochastic primal-dual methods and MCMC sampling to do away with the need for transformations to be differentiable.

Original: Our experiments show that it leads to state-of-the-art results in automatic data augmentation benchmarks in CIFAR datasets.
Revised: Our experiments show that it leads to state-of-the-art results in automatic data augmentation benchmarks in CIFAR datasets.

Original:  D ATA A UGMENTATION IN S UPERVISED L EARNING
Revised:  D ATA A UGMENTATION IN S UPERVISED L EARNING

Original: As in the standard supervised learning setting, let x ∈ X ⊆
Revised: As in the standard supervised learning setting, let x ∈ X ⊆

Original: R
Revised: R

Original: d denote a feature vector and y ∈ Y ⊆
Revised: d denote a feature vector and y ∈ Y ⊆

Original: R its associated label or measurement.
Revised: R its associated label or measurement.

Original: For classification tasks, we take Y ⊆ N .
Revised: For classification tasks, we take Y ⊆ N .

Original: Let D denote a probability distribution over the data pairs ( x , y ) and ℓ : Y × Y
Revised: Let D denote a probability distribution over the data pairs ( x , y ) and ℓ : Y × Y

Original: → R + be a non-negative, convex loss function, e.g., the cross entropy loss.
Revised: → R + be a non-negative, convex loss function, e.g., the cross entropy loss.

Original: Our goal is to learn a predictor f θ : X → Y in some hypothesis class H θ
Revised: Our goal is to learn a predictor f θ : X → Y in some hypothesis class H θ

Original: = { f θ | θ ∈ Θ ⊆ R p } that minimizes the expected loss, namelyminimize
Revised: = { f θ | θ ∈ Θ ⊆ R p } that minimizes the expected loss, namely

Original: [Equation]
Revised: [Equation]

Original: Θ [Equation] We consider the distribution D to be unknown, except for the dataset { ( x i , y i ) , i = 1 , . . . , n } of n i.i.d.
Revised: We consider the distribution D to be unknown, except for the dataset { ( x i , y i ) , i = 1 , . . . , n } of n i.i.d.

Original: samples from D .
Revised: samples from D .

Original: Therefore, we rely on the empirical approximation of the objective of (SRM), explicitly
Revised: Therefore, we rely on the empirical approximation of the objective of (SRM), explicitly

Original: [Equation]
Revised: [Equation]

Original: The aim of data augmentation is to improve the approximation ˆ R of the statistical risk R when dealing with a dataset that is not sufficiently representative of the data distribution.
Revised: One of the aims of data augmentation is to improve the approximation ˆ R of the statistical risk R when dealing with a dataset that is not sufficiently representative of the data distribution.

Original: To do so, we consider transformations of the feature vector g : X → X , taken from the (possibly infinite) transformation set G .
Revised: To do so, we consider transformations of the feature vector g : X → X , taken from the (possibly infinite) transformation set G .

Original: Common examples include rotations and translations in images.
Revised: Common examples include rotations and translations in images.

Original: Data augmentation leverages these transformations to generate new data pairs ( g x , y ) by sampling transformations according to a probability distribution G over G , leading to the learning problem
Revised: Data augmentation leverages these transformations to generate new data pairs ( g x , y ) by sampling transformations according to a probability distribution G over G , leading to the learning problem

Original: [Equation]
Revised: [Equation]

Original: Note that the empirical risk approximation ˆ R in (1) can be interpreted as an approximation of the data distribution D by a discrete distribution that places atoms on each data point.
Revised: Note that the empirical risk approximation ˆ R in (1) can be interpreted as an approximation of the data distribution D by a discrete distribution that places atoms on each data point.

Original: In that sense, ˆ R aug in (2) can be thought of as the Vicinal Risk Minimization (Chapelle et al., 2000) counterpart of (1), in which the atoms on x i are replaced by a local distribution over the transformed samples g x i , i.e.,
Revised: In that sense, ˆ R aug in (2) can be thought of as the Vicinal Risk Minimization (Chapelle et al., 2000) counterpart of (1), in which the atoms on x i are replaced by a local distribution over the transformed samples g x i , i.e.,

Original: [Equation]
Revised: [Equation]

Original:  where the distribution P over X is induced by the distribution G over G .
Revised: ) where the distribution P over X is induced by the distribution G over G .

Original: As it can be seen from (3) if G is not chosen adequately, ˆ R aug can be a poor estimate of R , introducing biases that outweigh the benefits of data augmentation (Chen et al
Revised: As it can be seen from (3) if G is not chosen adequately, ˆ R aug can be a poor estimate of R , introducing biases that outweigh the benefits of data augmentation (Chen et al., 2019; Shao et al., 2022).

Original:  D ATA A UGMENTATION VIA I NVARIANCE C ONSTRAINTS
Revised: On the other hand, if the data distribution D is statistically invariant under the action of G , invariant solutions have provable advantages in terms of sample complexity (Chen et al., 2019; Bietti and Mairal, 2019; Sannai et al., 2019; Lyle et al., 2020). In this work, given a set of transformations G , we tackle the choice of G , i.e., how to sample transformations so that the solution of the learning problem is sufficiently invariant with respect to the learning task, as defined on the next section. Note that invariance to transformations in G may hold for the true distribution or may also be a desirable property of the solution, for example, to achieve robustness (Kanbak et al., 2017; Volpi et al., 2018; Joshi et al., 2019). Unlike invariance learning (Jebara, 2003;

Original:  Finding a transformation distribution G that leads toa good approximation of the underlying distribution of the data D in (3) can be challenging.
Revised: Zhou et al., 2021a; Benton et al., 2020; Immer et al., 2022), we do not seek to learn the transformations G from the data.  D ATA A UGMENTATION VIA I NVARIANCE C ONSTRAINTS Finding a transformation distribution G that leads to the desired properties in the solution can be challenging .

Original: What is more, using a fixed G as in (2) prevents us from controlling when and how much augmentation is used during training, running the risk of biasing the final solution.
Revised: What is more, using a fixed G as in (2) prevents us from controlling when and how much augmentation is used during training, running the risk of biasing the final solution.

Original: On the other hand, it is straightforward to specify a set of transformations G to which the solution should be approximately invariant (e.g., image rotations and translations).
Revised: On the other hand, it is straightforward to specify a set of transformations G to which the solution should be approximately invariant (e.g., image rotations and translations).

Original: In the next sections, we explain how a data augmentation distribution can be obtained from such an invariance requirement.
Revised: In the next sections, we explain how a data augmentation distribution can be obtained from such an invariance requirement.

Original: We first show how invariance can be interpreted as an augmentation distribution (Section 3.1).
Revised: We first show how invariance can be interpreted as an augmentation distribution (Section 3.1).

Original: We then incorporate this invariance in a constrained learning problem (Section 3.2).
Revised: We then incorporate this invariance in a constrained learning problem (Section 3.2).

Original: 3. F ROM INVARIANCE TO AN A UGMENTATION D ISTRIBUTION
Revised: 3. F ROM INVARIANCE TO AN A UGMENTATION D ISTRIBUTION

Original: The goal of this section is to replace the task of choosing an augmentation distribution by the task of choosing a set of transformations we wish to be (approximately) invariant to.
Revised: The goal of this section is to replace the task of choosing an augmentation distribution by the task of choosing a set of transformations we wish to be (approximately) invariant to.

Original: To do so, we we will show how invariance can be used to implicitly determine the augmentation distribution.
Revised: To do so, we we will show how invariance can be used to implicitly determine the augmentation distribution.

Original: However, rather than requiring the output of the model to be invariant, i.e., f θ ( x ) =
Revised: However, rather than requiring the output of the model to be invariant, i.e., f θ ( x ) =

Original: f θ ( g x ) for all g ∈ G , we will consider invariance in terms of quality of its predictions as evaluated by the loss function.
Revised: f θ ( g x ) for all g ∈ G , we will consider invariance in terms of quality of its predictions as evaluated by the loss function.

Original: Namely,
Revised: Namely,

Original: [Equation]
Revised: [Equation]

Original: This notion of invariance explicitly incorporates the structure of the learning task by using the loss to identify which changes in the output of the model would lead to a significant change in prediction performance.
Revised: This notion of invariance explicitly incorporates the structure of the learning task by using the loss to identify which changes in the output of the model would lead to a significant change in prediction performance.

Original: More precisely, we want to limit the difference in our model’s performance on a sample and its transformed versions, i.e., we wish to have | ℓ ( f θ ( x ) , y ))
Revised: More precisely, we want to limit the difference in our model’s performance on a sample and its transformed versions, i.e., we wish to have | ℓ ( f θ ( x ) , y ))

Original: − ℓ ( f θ ( g x ) , y )) | small for all g in
Revised: − ℓ ( f θ ( g x ) , y ) [Equation] | small for all g in

Original: G .
Revised: G .

Original: Equivalently, we wish to control the magnitude of
Revised: Equivalently, we wish to control the magnitude of

Original: [Equation]
Revised: [Equation]

Original: However, rather controlling (4) for all x ∈ X , which may be overly conservative, we want to restrict our attention to the support of the data distribution.
Revised: However, rather controlling (4) for all x ∈ X , which may be overly conservative, we want to restrict our attention to the support of the data distribution.

Original: Furthermore, we wish to weight different inputs depending on their probability, in order to reduce the importance of unlikely or pathological cases.
Revised: Furthermore, we wish to weight different inputs depending on their probability, in order to reduce the importance of unlikely or pathological cases.

Original: We therefore average (4) over the data distribution to obtain
Revised: We therefore average (4) over the data distribution to obtain

Original: [Equation]
Revised: [Equation]

Original: To connect the invariant risk R inv in (5) to the data augmentation formulation in (2), notice that it can be bounded, using triangle inequality and the monotonicity of the expectation, by
Revised: To connect the invariant risk R inv in (5) to the data augmentation formulation in (2), notice that it can be bounded, using triangle inequality and the monotonicity of the expectation, by

Original: [Equation]
Revised: [Equation]

Original: The first term is simply the statistical risk R that is usually minimized in order to tackle the learning task.
Revised: The first term is simply the statistical risk R that is usually minimized in order to tackle the learning task.

Original: The second term resembles the objective typically found in adversarial learning problems, e.g., (Madry et al., 2017).
Revised: The second term resembles the objective typically found in adversarial learning problems, e.g., (Madry et al., 2017).

Original: This term, it turns out, can be interpreted as an augmentation distribution.
Revised: This term, it turns out, can be interpreted as an augmentation distribution.

Original: As shown by Robey et al. (2021a), the maximisation of the loss over transformations can be written as the semi-infinite constrained optimization problem
Revised: As shown by Robey et al. (2021a), the maximisation of the loss over transformations can be written as the semi-infinite constrained optimization problem

Original: [Equation]
Revised: [Equation]

Original: Notice that the solution of this optimization problem λ ⋆ ( g ) is a non-negative, normalized function and can therefore be interpreted as a distribution over transformations that depends on the sample point ( x , y ) as well as the model f θ .
Revised: Notice that the solution of this optimization problem λ ⋆ ( g ) is a non-negative, normalized function and can therefore be interpreted as a distribution over transformations that depends on the sample point ( x , y ) as well as the model f θ .

Original: This allows us to re-interpret the maximization over G as an expectation, i.e.,
Revised: This allows us to re-interpret the maximization over G as an expectation, i.e.,

Original: [Equation]
Revised: [Equation]

Original: Observe that the right-hand side of (8) resembles the statistical form of the data augmentation objective in (2).
Revised: Observe that the right-hand side of (8) resembles the statistical form of the data augmentation objective in (2).

Original: 3. A C ONSTRAINED L EARNING A PPROACH
Revised: 3. A C ONSTRAINED L EARNING A PPROACH

Original: Returning to the invariant risk bound in (6), notice that it is composed of two parts, namely, the statistical risk R from (SRM) and what we have shown in (8) to be a intance/model-dependent data augmentation.
Revised: Returning to the invariant risk bound in (6), notice that it is composed of two parts, namely, the statistical risk R from (SRM) and what we have shown in (8) to be a intance/model-dependent data augmentation.

Original: However, in order to address the potential biases introduced by transformations, rather than modifying the objective as in standard data augmentation, we propose to combine these two terms in a constrained learning problem.
Revised: However, in order to address the potential biases introduced by transformations, rather than modifying the objective as in standard data augmentation, we propose to combine these two terms in a constrained learning problem.

Original: Explicitly,
Revised: Explicitly,

Original: [Equation]
Revised: [Equation]

Revised: ) , y )]

Original:  Notice that this formulation tackles the two terms forming the invariant risk bound in (6), but instead of combining them directly, it incorporates the data augmentation term as a constraint in the typical statistical risk minimization problem (SRM).
Revised: [Equation] s. to E ( x ,y ) ∼ D [ E g ∼ λ ⋆ [ ℓ ( f θ ( g x ) , y )]] ≤ ϵ . Notice that this formulation tackles the two terms forming the invariant risk bound in (6), but instead of combining them directly, it incorporates the data augmentation term as a constraint in the typical statistical risk minimization problem (SRM).

Original: This formulation has the advantage that if a solution to the unconstrained problem is feasible, i.e., satisfies the invariance constraint in (CSRM), the presence of that constraint has no effect on the statistical problem.
Revised: This formulation has the advantage that if a solution to the unconstrained problem is feasible, i.e., satisfies the invariance constraint in (CSRM), the presence of that constraint has no effect on the statistical problem.

Original: Yet, it can be beneficial when approximating the solution of (CSRM) empirically.
Revised: Yet, it can be beneficial when approximating the solution of (CSRM) empirically.

Original: We will explore this fact in the next section, where we tackle the practical challenges involved in solving (CSRM).
Revised: We will explore this fact in the next section, where we tackle the practical challenges involved in solving (CSRM).

Revised: For conciseness in (CSRM) we have included only one invariance constraint associated with a single set of transformations G .

Revised: However, our approach can be extended to an arbitrary number of constraints defined by transformation sets G i , i = 1 , . . .

Revised: , m (each inducing an augmentation distribution λ ⋆i ), and constraint levels

Revised: ϵ i .

Revised: All of the following derivations still hold, regardless of the number of constraints.

Original:  A LGORITHM D EVELOPMENT
Revised:  A LGORITHM D EVELOPMENT

Original: Solving (CSRM) presents two challenges.
Revised: Solving (CSRM) presents two challenges.

Original: First, it is a constrained statistical learning problem, which involves the unknown data distribution D .
Revised: First, it is a constrained statistical learning problem, which involves the unknown data distribution D .

Original: We address this by resorting to an empirical dual problem as explained on Section 4.1.
Revised: We address this by resorting to an empirical dual problem as explained on Section 4.1.

Original: Second, it can be hard to sample from λ ⋆ .
Revised: Second, it can be hard to sample from λ ⋆ .

Original: We address this by introducing a smooth approximation that leverages MCMC methods on Section 4.2.
Revised: We address this by introducing a smooth approximation that leverages MCMC methods on Section 4.2.

Original: 4. E MPIRICAL D UAL C ONSTRAINED L EARNING
Revised: 4. E MPIRICAL D UAL C ONSTRAINED L EARNING

Original: To tackle the invariance-constrained statistical risk minimization problem we leverage recent duality results in constrained learning theory (Chamon et al., 2021), that approximate the problem by its empirical dual
Revised: To tackle the invariance-constrained statistical risk minimization problem we leverage recent duality results in constrained learning theory (Chamon et al., 2021), that approximate the problem by its [Table] [Equation] [Equation]

Revised: D

Original: [Equation]
Revised: [Equation]

Original:  The advantage of (D-CERM) is that it is an unconstrained problem that, provided we have enough samples and the parametrization is rich enough, can approximate the constrained statistical problem ( CSRM).
Revised: CERM) The advantage of (D-CERM) is that it is an unconstrained problem that, provided we have enough samples and the parametrization is rich enough, can approximate the constrained statistical problem ( CSRM).

Original: Namely, the difference between the optimal value of the empirical dual D ⋆ emp and the statistical primal P ⋆ , i.e., the empirical duality gap is bounded (Chamon et al., 2021).
Revised: Namely, the difference between the optimal value of the empirical dual D ⋆ emp and the statistical primal P ⋆ , i.e., the empirical duality gap is bounded (Chamon et al., 2021).

Original: As in regular data augmentation, we will also approximate the expectation over λ ⋆ by sampling transformations as discussed in Section 4.2.
Revised: As in regular data augmentation, we will also approximate the expectation over λ ⋆ by sampling transformations as discussed in Section 4.2.

Original: Then, the problem D-CERM becomes an unconstrained deterministic problem, which can be solved using the algorithm described in Section 4.3.
Revised: Then, the problem D-CERM becomes an unconstrained deterministic problem, which can be solved using the algorithm described in Section 4.3.

Original: Note that finding a Lagrangian minimizer for a fixed value of the dual variable ( γ ) is equivalent to minimising the risk under a fixed mixture augmentation distribution 1 , where the value of γ controls the probability of sampling the identity.
Revised: Note that finding a Lagrangian minimizer for a fixed value of the dual variable ( γ ) is equivalent to minimising the risk under a fixed mixture augmentation distribution 1 , where the value of γ controls the probability of sampling the identity.

Original: We can also interpret this as optimising a penalised or regularised learning objective.
Revised: We can also interpret this as optimising a penalised or regularised learning objective.

Original: However, solving the constrained problem, namely maximising over γ , has fundamental differences.
Revised: However, solving the constrained problem, namely maximising over γ , has fundamental differences.

Original: ∼ λ ∗ .

Original: First, constraints explicit the requirement they represent.
Revised: First, constraints explicit the requirement they represent.

Original: While the degree of invariance imposed should depend only on the statistical problem at hand, the value of γ needed to achieve it will depend on the sample size, the parametrization and the learning algorithm.
Revised: While the degree of invariance imposed should depend only on the statistical problem at hand, the value of γ needed to achieve it will depend on the sample size, the parametrization and the learning algorithm.

Original: In contrast, constrained learning dynamically adjusts the amount of augmentation — dictated by γ — to a particular learning setup.
Revised: In contrast, constrained learning dynamically adjusts the amount of augmentation — dictated by γ — to a particular learning setup.

Original: Second, the optimal dual variable can give information about the trade-off minimising the loss over training samples and satisfying the invariance constraint.
Revised: Second, the optimal dual variable can give information about the trade-off minimising the loss over training samples and satisfying the invariance constraint.

Original: In penalised approaches, on the contrary, this trade-off is fixed.
Revised: In penalised approaches, on the contrary, this trade-off is fixed.

Original: Lastly, the aforementioned informativeness and interpretability can facilitate hyper-parameter tuning.
Revised: Lastly, the aforementioned informativeness and interpretability can facilitate hyper-parameter tuning.

Original: The insights gathered from optimal dual variables can be leveraged a posteriori, for instance, to manually choose appropriate transformations, relax the invariance constraint levels, or change the learning setup (e.g. increase the capacity of the model class).
Revised: The insights gathered from optimal dual variables can be leveraged a posteriori, for instance, to manually choose appropriate transformations, relax the invariance constraint levels, or change the learning setup (e.g. increase the capacity of the model class).

Original: 4. S AMPLING TOWARDS INVARIANCE
Revised: 4. S AMPLING TOWARDS INVARIANCE

Original: Sampling the augmentation distribution λ ⋆ can be difficult when G is not finite and f θ is a deep neural network.
Revised: Sampling the augmentation distribution λ ⋆ can be difficult when G is not finite and f θ is a deep neural network.

Original: Even when the transformation space is low dimensional, as in the case of translations and rotations, the highly non-convex loss landscape of these models makes the maximization over G challenging (Engstrom et al., 2017).
Revised: Even when the transformation space is low dimensional, as in the case of translations and rotations, the highly non-convex loss landscape of these models makes the maximization over G challenging (Engstrom et al., 2017).

Original: If the optimal distribution λ ⋆ is not smooth, it is challenging to sample from it with sufficient accuracy (Homem-de Mello and Bayraksan, 2014).
Revised: If the optimal distribution λ ⋆ is not smooth, it is challenging to sample from it with sufficient accuracy (Homem-de Mello and Bayraksan, 2014).

Original: Consequently, obtaining an unbiased estimator of E g ∼ λ ⋆ [ ℓ ( f θ ( g x ) , y )] may not be possible.
Revised: Consequently, obtaining an unbiased estimator of E g ∼ λ ⋆ [ ℓ ( f θ ( g x ) , y )] may not be possible.

Original: Therefore, we add an
Revised: Therefore, we add an

Original: L 2 norm penalisation, which promotes smoothness, to leverage MCMC methods.
Revised: L 2 norm penalisation, which promotes smoothness, to leverage MCMC methods.

Original: We then define the c-smoothed distribution λ
Revised: We then define the c-smoothed distribution λ

Original: ⋆c as a solution to the regularised problem
Revised: ⋆c as a solution to the regularised problem

Original: [Equation]
Revised: [Equation]

Original: The regularization term introduces an optimality gap with respect to worst case perturbations, i.e., E g ∼ λ
Revised: The regularization term introduces an optimality gap with respect to worst case perturbations, i.e., E g ∼ λ

Original: ⋆c
Revised: ⋆c

Original: [ ℓ ( f θ ( g x ) , y )] ≤ max g ∈G ℓ ( f θ ( g x ) , y ) .
Revised: [ ℓ ( f θ ( g x ) , y )] ≤ max g ∈G ℓ ( f θ ( g x ) , y ) .

Original: However, for particular values of c the regularized problem has a closed form solution (Robey et al., 2021a) that allows us to sample from it easily.
Revised: However, for particular values of c the regularized problem has a closed form solution (Robey et al., 2021a) that allows us to sample from it easily.

Original: Namely, there exists a constant c ≥ 0 such that λ ⋆c ( x , y, g ) =
Revised: Namely, there exists a constant c ≥ 0 such that λ ⋆c ( x , y, g ) =

Original: ℓ ( f θ ( g x ) ,y ) c .
Revised: ℓ ( f θ ( g x ) ,y ) c .

Original: Since λ ⋆c is a smooth probability distribution, we do not need to estimate the multiplicative factor c to sample from it by leveraging Monte Carlo Markov Chain methods (MCMC).
Revised: Since λ ⋆ c is a smooth probability distribution, we do not need to estimate the multiplicative factor c to sample from it by leveraging Monte Carlo Markov Chain methods (MCMC).

Original: MCMC methods (Hastings, 1970) are based on constructing a Markov chain that has the target distribution as an equilibrium distribution.
Revised: MCMC methods (Hastings, 1970) are based on constructing a Markov chain that has the target distribution as an equilibrium distribution.

Original: Independent Metropolis Hastings uses a state independent - usually fixed - proposal for each step.
Revised: Independent Metropolis Hastings uses a state independent - usually fixed - proposal for each step.

Original: In our case, it only requires applying a transformation and computing a forward pass of the neural network to evaluate the loss.
Revised: In our case, it only requires applying a transformation and computing a forward pass of the neural network to evaluate the loss.

Original: This enables the use of non-differentiable transformations, and has the advantage that the density at consecutive proposals can be evaluated in parallel allowing speedups in the sampling step.
Revised: This enables the use of non-differentiable transformations, and has the advantage that the density at consecutive proposals can be evaluated in parallel allowing speedups in the sampling step.

Original: Although MH methods thus allow to sample the proposal distribution with low computational cost, they exhibit random walk behaviour, which leads to slow convergence in high dimensional settings (Dellaportas and Roberts, 2003; Holden et al., 2009).
Revised: Although MH methods thus allow to sample the proposal distribution with low computational cost, they exhibit random walk behaviour, which leads to slow convergence in high dimensional settings (Dellaportas and Roberts, 2003; Holden et al., 2009).

Original: However, a uniform augmentation distribution has been shown to give good results (Müller and Hutter, 2021) and sampling λ ∗ c accurately is not a concern given that we are interested in promoting approximate invariance.

Original: We can then approximate E g ∼ λ
Revised: We can then approximate E g ∼ λ

Original: ⋆c
Revised: ⋆c

Original: [ ℓ ( f θ ( g x i ) , y i )] by sampling transformations according to the loss.
Revised: [ ℓ ( f θ ( g x i ) , y i )] by sampling transformations according to the loss.

Original: Namely, we can obtain a set of m samples drawn from λ ⋆c and approximate the expectation over the group by the sample mean
Revised: Namely, we can obtain a set of m samples drawn from λ ⋆c and approximate the expectation over the group by the sample mean

Original: [Equation]
Revised: [Equation]

Original: where g 1 , . . . , g m i.i.d.
Revised: where g 1 , . . . , g m i.i.d.

Original: ∼ ℓ ( f θ ( g x i ) , y i ) /c are m transformations sampled from the smoothed distribution λ
Revised: ∼ ℓ ( f θ ( g x i ) , y i ) /c are m transformations sampled from the smoothed distribution λ

Original: ⋆c ( f θ , x i , y i ) .
Revised: ⋆c ( f θ , x i , y i ) .

Original: In the next section, the implementation of independent-MH with a uniform proposal is described in
Revised: In the next section, the implementation of independent-MH with a uniform proposal is described in

Original: Algorithm 2, together with the primal-dual augmentation algorithm.
Revised: Algorithm 2, together with the primal-dual augmentation algorithm.

Original: 4. P RIMAL -D UAL A LGORITHM
Revised: 4. P RIMAL -D UAL A LGORITHM

Original: Since the cost of the inner minimization, i.e. training the model, can be high, we adopt an alternating update scheme (K. J. Arrow and Uzawa., 1958) for the primal and dual variables, as in (Chamon et al., 2021; Fioretto et al., 2020).
Revised: Since the cost of the inner minimization, i.e. training the model, can be high, we adopt an alternating update scheme (K. J. Arrow and Uzawa., 1958) for the primal and dual variables, as in (Chamon et al., 2021; Fioretto et al., 2020).

Original: A bounded empirical duality gap does not guarantee that the primal variables obtained after running the alternating primal-dual algorithm 1 and solving the saddle point problem approximately are near optimal or approximately feasible.
Revised: A bounded empirical duality gap does not guarantee that the primal variables obtained after running the alternating primal-dual algorithm 1 and solving the saddle point problem approximately are near optimal or approximately feasible.

Original: Although stronger primal recovery guarantees can be obtained by randomizing the learning algorithm (Chamon et al., 2021), it requires storing model parameters θ at each iteration, and there is empirical evidence (Chamon et al., 2021;
Revised: Although stronger primal recovery guarantees can be obtained by randomizing the learning algorithm (Chamon et al., 2021), it requires storing model parameters θ at each iteration, and there is empirical evidence (Chamon et al., 2021;

Original: Robey et al., 2021a;
Revised: Robey et al., 2021a;

Original: Elenter et al., 2022;
Revised: Elenter et al., 2022;

Original: Shen et al., 2022; Cervino et al., 2022; Zhang et al., 2022) that good solutions can still be obtained without randomization.
Revised: Shen et al., 2022; Cervino et al., 2022; Zhang et al., 2022) that good solutions can still be obtained without randomization.

Original: Algorithm 2 describes transformation sampling.
Revised: Algorithm 2 describes transformation sampling.

Original: By keeping only one sample ( m = 1 ) we recover the usual augmentation setting, that yields one augmentation per sample in the training batch.
Revised: By keeping only one sample ( m = 1 ) we recover the usual augmentation setting, that yields one augmentation per sample in the training batch.

Original: In our experiments we address this setting, since it has lower computational costs.
Revised: In our experiments we address this setting, since it has lower computational costs.

Original: However, simply keeping more samples from the chain ( m > 1 ) allows to extend the method to the batch augmentation setting (Hoffer et al., 2020), which creates several augmented samples from the same instance in each batch.
Revised: However, simply keeping more samples from the chain ( m > 1 ) allows to extend the method to the batch augmentation setting (Hoffer et al., 2020), which creates several augmented samples from the same instance in each batch.

Original: Although several steps of the chain may be required to deviate enough from the proposal distribution, we show in the experimental section that sampling the constraint approximately by using few sampling steps suffices.
Revised: Although several steps of the chain may be required to deviate enough from the proposal distribution, we show in the experimental section that sampling the constraint approximately by using few sampling steps suffices.

Original: 1: λ = 0 , θ = θ 0 . 2: for Sample = 1 | Batch
Revised: Algorithm

Original: ) ∈ Batch (cid:104) 1 m (cid:80) mj =1 ℓ ( f θ ( g ij x i ) , y i ) (cid:105) − ϵ ▷Evaluate constraint slack 6: ℓ = 1 | Batch | (cid:80) ( x i ,y i ) ∈ Batch ℓ ( f θ ( x i ) , y i )
Revised: ▷

Original: ▷ Evaluate Loss 7:ˆ L = ℓ + γs ▷ Compute Lagrangian 8: θ = θ − η p ∇ θ ˆ L ▷ Primal update 9: γ=
Revised: ▷ Evaluate Loss 7: L = ℓ + γs ▷ Compute Lagrangian 8: θ = θ − η p ∇ θ ˆ L ▷ Primal update 9: γ

Original: [ γ+ η d s ] + ▷ Dual Update
Revised: [ γ η d s ] + ▷ Dual Update

Original: Algorithm

Original: 1: g (0) ∼ U ( G )
Revised: 1: g (0) ∼ U ( G )

Original: ▷ Sample initial State 2: ℓ (0) =
Revised: ▷ Sample initial State 2: ℓ (0) =

Original: ℓ (cid:0) f θ (cid:0) g (0) x (cid:1) , y (cid:1)
Revised: ℓ (cid:0) f θ (cid:0) g (0) x (cid:1) , y (cid:1)

Original: ▷ Evaluate loss 3: for t = 1 , . .
Revised: ▷ Evaluate loss 3: for t = 1 , . .

Original: . , n steps do 4: g prop ∼ U ( G )
Revised: . , n steps do 4: g prop ∼ U ( G )

Original: ▷ Sample next proposal 5: ℓ prop = ℓ ( f θ ( g prop x ) , y )
Revised: ▷ Sample next proposal 5: ℓ prop = ℓ ( f θ ( g prop x ) , y )

Original: ▷ Evaluate Loss 6: p = min (cid:16) 1 , ℓ prop ℓ ( t − 1) (cid:17) ▷ Acceptance Prob 7: w.p. p : ▷ Accept/Reject 8: g ( t ) =
Revised: ▷ Evaluate Loss 6: p = min (cid:16) 1 , ℓ prop ℓ ( t − 1) (cid:17) ▷ Acceptance Prob 7: w.p. p : ▷ Accept/Reject 8: g ( t ) =

Original: g prop , ℓ ( t ) =
Revised: g prop , ℓ ( t ) =

Original: ℓ prop 9: else: 10: g ( t ) =
Revised: ℓ prop 9: else: 10: g ( t ) =

Original: g ( t − 1) , ℓ ( t ) =
Revised: g ( t − 1) , ℓ ( t ) =

Original: ℓ ( t − 1)
Revised: ℓ ( t − 1)

Original:  A UTOMATIC D ATA A UGMENTATION E XPERIMENTS
Revised:  A UTOMATIC D ATA A UGMENTATION E XPERIMENTS

Original: This section showcases Algorithm 1 in common image classification benchmarks.
Revised: This section showcases Algorithm 1 in common image classification benchmarks.

Original: We compare it to state-of-the-art data augmentation methods in terms of classification accuracy.
Revised: We compare it to state-of-the-art data augmentation methods in terms of classification accuracy.

Original: Furthermore, we discuss other advantageous properties of our method through ablations.
Revised: Furthermore, we discuss other advantageous properties of our method through ablations.

Original: Namely, we focus on the ability to control the effect of data augmentation by modifying the constraint level, the informativeness of dual variables, and the benefits of adapting the augmentation distribution throughout training.
Revised: Namely, we focus on the ability to control the effect of data augmentation by modifying the constraint level, the informativeness of dual variables, and the benefits of adapting the augmentation distribution throughout training.

Original: We follow the setup (including the transformation sets) used in recent automatic data augmentation literature (Müller and Hutter, 2021).
Revised: We follow the setup (including the transformation sets) used in recent automatic data augmentation literature (Müller and Hutter, 2021).

Original: A complete list of transformations together with other hyperparameters and training settings can be found on Appendix B. Throughout these experiments
Revised: A complete list of transformations together with other hyperparameters and training settings can be found on Appendix B. Note that four out of the sixteen transformations used are non-differentiable.

Original: [Table]
Revised: Whereas other works have introduced gradient approximations for transformation operations with discrete parameters (Li et al., 2020b; Hataya et al., 2020), our approach does not require such approximations.

Original:  we fixed the number of steps of the MH sampler (Algorithm 2) to two, which has the added advantage of reducing the computational cost of evaluating proposals.
Revised: Throughout these experiments, we fixed the number of steps of the MH sampler (Algorithm 2) to two, which has the added advantage of reducing the computational cost of evaluating proposals.

Original: The constraint level was determined by a grid search targeting cross-validation accuracy.
Revised: The constraint level was determined by a grid search targeting cross-validation accuracy.

Original: As shown in Table 1, in both transformation sets considered, we find that our approach improves or closely matches existing approaches.
Revised: As shown in Table 1, in both transformation sets considered, we find that our approach improves or closely matches existing approaches.

Original: The failure to achieve large improvements in accuracy over baselines, which has been attributed to a stagnation in data augmentation research (Müller and Hutter, 2021), can also reflect the limits of the benchmarking setup
Revised: The failure to achieve large improvements in accuracy over baselines, which has been attributed to a stagnation in data augmentation research (Müller and Hutter, 2021), can also reflect the limits of the benchmarking setup.

Original: [Table]
Revised: [Table]

Original: our method. In CIFAR datasets our approach leads to an improvement in accuracy, whereas in SVHN (Netzer et al., 2011) it leads to a degradation in performance in the wide setup.
Revised: accuracy, whereas in SVHN (Netzer et al., 2011) it leads to a degradation in performance in the wide setup.

Original: Regardless, our approach yields improvements in test accuracy over a baseline model without augmentation for a wide range of constraint levels (Figure 1).
Revised: Regardless, our approach yields improvements in test accuracy over a baseline model without augmentation for a wide range of constraint levels (Figure 1).

Original: This illustrates the robustness of the solution to this hyperparameter.
Revised: This illustrates the robustness of the solution to this hyperparameter.

Original: Observe also that as the constraint is relaxed (by increasing ϵ ), the training error decreases while the generalization gap, i.e., the difference between train and test errors, increases.
Revised: Observe also that as the constraint is relaxed (by increasing ϵ ), the training error decreases while the generalization gap, i.e., the difference between train and test errors, increases.

Original: In other words, by loosening the invariance requirement the model can fit better to training samples at the cost of worse generalization.
Revised: In other words, by loosening the invariance requirement the model can fit better to training samples at the cost of worse generalization.

Original: Eventually, only the generalization gap increases while the training error stagnates ( ϵ > 2 . 1 for CIFAR100 and ϵ > 0 .
Revised: Eventually, only the generalization gap increases while the training error stagnates ( ϵ > 2 . 1 for CIFAR100 and ϵ > 0 .

Original: 8 for CIFAR10).
Revised: 8 for CIFAR10).

Original: This transition occurs at the same point at which the final value of the dual variable γ from (D-CERM) essentially vanishes (darker color).
Revised: This transition occurs at the same point at which the final value of the dual variable γ from (D-CERM) essentially vanishes (darker color).

Original: This showcases the infromativeness of the dual variable.
Revised: This showcases the infromativeness of the dual variable.

Original: In the case of
Revised: In the case of

Original: CIFAR10, even the training error begins to increase at that point, suggesting that the invariance requirement need not be at odds with accuracy.
Revised: CIFAR10, even the training error begins to increase at that point, suggesting that the invariance requirement need not be at odds with accuracy.

Original:  Not only does Algorithm 1 tune the effect of data augmentation on the solution (by adapting γ in step 9), but also modifies the distribution over transformations during training (step 4).
Revised: [Figure] Not only does Algorithm 1 tune the effect of data augmentation on the solution (by adapting γ in step 9), but also modifies the distribution over transformations during training (step 4).

Original: To showcase the benefits of this over the use of a fixed distribution, Figure 2 compares the results obtained using our approach (sampling according to λ ⋆c ) and one where step 4 is replaced by a uniform sampling over transformations.
Revised: To showcase the benefits of this over the use of a fixed distribution, Figure 2 compares the results obtained using our approach (sampling according to λ ⋆c ) and one where step 4 is replaced by a uniform sampling over transformations.

Original: For the same constraint levels, lower test errors are obtained by sampling transformations according to λ ⋆c , i.e., promoting invariance.
Revised: For the same constraint levels, lower test errors are obtained by sampling transformations according to λ ⋆c , i.e., promoting invariance.

Original: Note also that for ϵ = 2 .
Revised: Note also that for ϵ = 2 .

Original: in CIFAR and ϵ = 0 . 

Revised: in CIFAR and ϵ = 0 .

Original: for CIFAR10, the performance gap is quite large.
Revised: for CIFAR10, the performance gap is quite large.

Original: Once again, this occurs at the point in which γ vanishes (darker color) for the uniform distribution, i.e., no data augmentation occurs by the end of training.
Revised: Once again, this occurs at the point in which γ vanishes (darker color) for the uniform distribution, i.e., no data augmentation occurs by the end of training.

Original: At this point, however, there is still value in promoting invariance by sampling from λ ⋆c as evidenced by the positive value of the dual variable (lighter color) in this approach.
Revised: At this point, however, there is still value in promoting invariance by sampling from λ ⋆c as evidenced by the positive value of the dual variable (lighter color) in this approach.

Original: [Figure]
Revised: [Figure]

Original: CIFAR [Figure] [Figure] Dual variable at epoch
Revised: r

Original: [Figure]
Revised: [Figure]

Original: [Figure]  R ELATED W ORK Zhang et al. (2020) have shown that using adversarial transformations - which as already mentioned is related to promoting invariance - can give competitive results with respect to other automatic augmentation methods in image classification.
Revised: Uniform c

Original: However, Blaas et al. (2021) have since evidenced the importance oftwo factors: the implicit learning curricula and the suboptimality of the adversarial used by Zhang et al.
Revised: Furthermore, to assess the impact of the sampling approximation on the performance of the solution we conduct an ablation on the number of steps of the MH sampler (Algorithm 2), keeping the constraint level ( ϵ ) fixed.

Original: (2020), which mitigate the biases introduced by worst-case transformations. Furthermore, Blaas et al. (2021) report that an explicit cyclic curricula in whichaugmentations are mild at first, then get harder as training progresses, and finally revert to milder augmentations at the end of training, performs better empirically. We note some interesting commonalities with our approach and experimental results.
Revised: Using more steps of the chain allows samples to deviate further from the uniform distribution, which as shown in Figure 3 is reflected on the lower entropy of sampled transformations (first row).

Original: First, the dynamics of our primal-dual algorithm resemble the aforementioned heuristically defined curricula.
Revised: As shown in the second row, this results in higher classification error rates on the training set (represented by blue circular markers).

Original: Second, the suboptimality with respect to worst case perturbations can be related to the smoothed approximation used in our approach.
Revised: However, we find that increasing the number of steps does not affect significantly test error (denoted by orange crosses).

Original: Lastly, Xu et al.
Revised: In Appendix C.1.

Original: (2021) have motivated a two-stage data augmentation policy, as learning subject to a constraint in the risk under a fixed augmentation distribution.
Revised: we show how the number of sampling steps affects the evolution of dual variables.

Original: Unlike our approach, this distribution still needs to be user defined.
Revised: [Figure]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this paper we have proposed a constrained learning approach for automatic data augmentation, which instead of using augmented samples as a modified learning objective, imposes an invarianceconstraint.
Revised: In this paper, we have proposed a constrained learning approach for automatic data augmentation, which instead of using augmented samples as a modified learning objective, imposes an invarianceconstraint.

Original: We have shown that this yields an augmentation distribution that adapts during training, and found that coarse sampling approximations based on MCMC methods can improve generalization in automatic data augmentation benchmarks.
Revised: We have shown that this yields an augmentation distribution that adapts during training, and found that coarse sampling approximations based on MCMC methods can improve generalization in automatic data augmentation benchmarks.

Original: Furthermore, our experiments showed that the dual variable can give insights about the resulting augmentation distribution.
Revised: Furthermore, our experiments showed that the dual variable can give insights about the resulting augmentation distribution.

Original: We also found that strictly feasible solutions were obtained for a wide range of constraint levels, with notably different generalization gaps and that in some cases tightening the constraint even led to a lower training error.
Revised: We also found that strictly feasible solutions were obtained for a wide range of constraint levels, with notably different generalization gaps, and that in some cases tightening the constraint even led to a lower training error.

Original: Thus, analysing the interplay between the learning problem and the optimization algorithm is a promising future work direction.
Revised: Thus, analysing the interplay between the learning problem and the optimization algorithm is a promising future work direction

Original: A A DDITIONAL R ELATED WORK A. C ONSTRAINED L EARNING AND D ATA AUGMENTATION
Revised: A A DDITIONAL R ELATED WORK A. C ONSTRAINED L EARNING AND D ATA AUGMENTATION

Original: Xu et al.
Revised: Xu et al.

Original: (2021) also formulate data augmentation as a constrained learning problem.
Revised: (2021) also formulate data augmentation as a constrained learning problem.

Original: They impose a constraint on the excess risk, i.e. the difference between the statistical risk and its optimal value, on augmented data.
Revised: They impose a constraint on the excess risk, i.e. the difference between the statistical risk and its optimal value, on augmented data.

Original: Thus the constraint level on the augmented risk is also determined by the data distribution, augmentations considered, and model class, and the existence of a strictly feasible point is guaranteed.
Revised: Thus the constraint level on the augmented risk is also determined by the data distribution, augmentations considered, and model class, and the existence of a strictly feasible point is guaranteed.

Original: [Equation]
Revised: [Equation]

Original: where R and R aug are the statistical risk under the original and augmented distribution, exlipicity
Revised: where R and R aug are the statistical risk under the original and augmented distribution, exlipicity

Original: [Equation]
Revised: [Equation]

Original: Unlike our formulation, this formulation assumes a fixed distribution of augmentations G is given.
Revised: Unlike our formulation, this formulation assumes a fixed distribution of augmentations G is given.

Original: By formulating it as a constrained problem, they aim to avoid introducing a bias when the data distribution is not invariant to augmentations.
Revised: By formulating it as a constrained problem, they aim to avoid introducing a bias when the data distribution is not invariant to augmentations.

Original: Two types of biases induced by augmentation are explicitly addressed, covariate shift (i.e. label-preserving augmentations) and concept shift (i.e. label mixing augmentations).
Revised: Two types of biases induced by augmentation are explicitly addressed, covariate shift (i.e. label-preserving augmentations) and concept shift (i.e. label mixing augmentations).

Original: Interestingly, they show that under some conditions on the risk, augmented risk and constraint level, by utilizing the augmented data to constrain the solution to a small region SGD can achieve lower error (Xu et al., 2021, Proposition 1).
Revised: Interestingly, they show that under some conditions on the risk, augmented risk and constraint level, by utilizing the augmented data to constrain the solution to a small region SGD can achieve lower error (Xu et al., 2021, Proposition 1).

Original: Instead of resorting to constrained optimization algorithms, they propose a two stage algorithm that consists of first finding an approximate minimizer of the augmented risk and then using that solution as an initialisation to the (unconstrained) statistical risk minimization problem.
Revised: Instead of resorting to constrained optimization algorithms, they propose a two stage algorithm that consists of first finding an approximate minimizer of the augmented risk and then using that solution as an initialisation to the (unconstrained) statistical risk minimization problem.

Original: The first stage obtains a feasible point, and then under some conditions the SGD iterates obtained when solving the second problem remain feasible (Xu et al., 2021, Theorem 1).
Revised: The first stage obtains a feasible point, and then under some conditions the SGD iterates obtained when solving the second problem remain feasible (Xu et al., 2021, Theorem 1).

Original: A. C ONSTRAINED L EARNING AND D OMAIN G ENERALIZATION
Revised: A. A DVERSARIAL D ATA A UGMENTATION

Revised: Zhang et al.

Revised: (2020) have shown that using adversarial transformations - which as already mentioned is related to promoting invariance - can give competitive results with respect to other automatic augmentation methods in image classification.

Revised: However, Blaas et al.

Revised: (2021) have since evidenced the importance of two factors: the implicit learning curricula and the suboptimality of the adversarial used by Zhang et al.

Revised: (2020), which mitigates the biases introduced by worst-case transformations.

Revised: Furthermore, Blaas et al.

Revised: (2021) report that an explicit cyclic curricula in which augmentations are mild at first, then get harder as training progresses, and finally revert to milder augmentations at the end of training, performs better empirically.

Revised: We note some interesting commonalities with our approach and experimental results.

Revised: First, the dynamics of our primal-dual algorithm resemble the aforementioned heuristically defined curricula.

Revised: Second, the suboptimality with respect to worst case perturbations can be related to the smoothed approximation used in our approach.

Revised: A. C ONSTRAINED L EARNING AND D OMAIN G ENERALIZATION

Original: Domain Generalization (DG) involves training the model in different but related data distributions, and evaluating in an unseen domain .
Revised: Domain Generalization (DG) involves training the model in different but related data distributions, and evaluating in an unseen domain .

Original: For example, a common benchmark consists of domains created by rotating images in the MNIST dataset by different angles.
Revised: For example, a common benchmark consists of domains created by rotating images in the MNIST dataset by different angles.

Original: Constrained formulations have been proposed in this context (Robey et al., 2021b; Zhang et al., 2022), that enforce invariance under learnt domain translation transformations.
Revised: Constrained formulations have been proposed in this context (Robey et al., 2021b; Zhang et al., 2022), that enforce invariance under learnt domain translation transformations.

Original: In contrast, our approach addresses pre-defined transformations that are commonly used in data augmentation pipelines in order to improve generalization in the same domain used to train the model.
Revised: In contrast, our approach addresses pre-defined transformations that are commonly used in data augmentation pipelines in order to improve generalization in the same domain used to train the model.

Original: A. G ROUP INVARIANCE
Revised: A. G ROUP INVARIANCE

Original: The relationship between convolutional structure and equivariance has been long known in algebraic signal processing theory (Püschel and Moura, 2006).
Revised: The relationship between convolutional structure and equivariance has been long known in algebraic signal processing theory (Püschel and Moura, 2006).

Original: Recently, necessity results have re-gained attention in the context of neural network architecture design (Kondor and Trivedi, 2018).
Revised: Recently, necessity results have re-gained attention in the context of neural network architecture design (Kondor and Trivedi, 2018).

Original: Several
Revised: Several

Original: Group Convolutional neural network architectures that generalize CNNs to different groups by leveraging group convolutions have been proposed (Cohen and Welling, 2016; Esteves et al., 2017;
Revised: Group Convolutional neural network architectures that generalize CNNs to different groups by leveraging group convolutions have been proposed (Cohen and Welling, 2016; Esteves et al., 2017;

Original: Finzi et al., 2020).
Revised: Finzi et al., 2020).

Original: In the case of images, it has been shown that that SE2 equivariant layers can be implemented efficiently using regular 2D convolutions (Weiler and Cesa, 2019).
Revised: In the case of images, it has been shown that that SE2 equivariant layers can be implemented efficiently using regular 2D convolutions (Weiler and Cesa, 2019).

Original: This approach allows to derive a parametrization for CNN filters under finite subgroups of SE2.
Revised: This approach allows to derive a parametrization for CNN filters under finite subgroups of SE2.

Original: Among other works (Bietti and Mairal, 2019; Sannai et al., 2019) give theoretical analyses of the benefits of group invariance in learning settings.
Revised: Among other works (Bietti and Mairal, 2019; Sannai et al., 2019) give theoretical analyses of the benefits of group invariance in learning settings.

Original: As already mentioned, achieving invariance through architecture design is both challenging and limited in the sense that it relies on transformations having a specific structure (e.g: a group).
Revised: As already mentioned, achieving invariance through architecture design is both challenging and limited in the sense that it relies on transformations having a specific structure (e.g: a group).

Original: The goal of exact invariance over the whole input space is more strict than the approximate invariance notion that our work addresses.
Revised: The goal of exact invariance over the whole input space is more strict than the approximate invariance notion that our work addresses.

Original: As argued by Mallat (2016), CNNs can learn locally invariant features with respect to arbitrary transformation groups, which could explain their generalization properties.
Revised: As argued by Mallat (2016), CNNs can learn locally invariant features with respect to arbitrary transformation groups, which could explain their generalization properties.

Original: Furthermore, empirical studies evidence modern CNN architectures learn approximately equivariant features to transformations such as scaling and rotations (Olah et al., 2020), or diffeomorphisms (Petrini et al., 2021), even when trained without direct augmentation.
Revised: Furthermore, empirical studies evidence modern CNN architectures learn approximately equivariant features to transformations such as scaling and rotations (Olah et al., 2020), or diffeomorphisms (Petrini et al., 2021), even when trained without direct augmentation.

Original: However, commonly used augmentations do not form a group.
Revised: However, commonly used augmentations do not form a group.

Original: Our approach does not require this structure.
Revised: Our approach does not require this structure.

Original: B E XPERIMENTAL S ETUP D ETAILS B. T RANSFORMATIONS
Revised: B E XPERIMENTAL S ETUP D ETAILS B. T RANSFORMATIONS

Original: As in recent automatic augmentation literature (Ho et al., 2019; Lim et al., 2019; Hataya et al., 2020; Zhang et al., 2020; Cubuk et al., 2020; LingChen et al., 2020; Zhang et al., 2020; Müller and Hutter, 2021;
Revised: As in recent automatic augmentation literature (Ho et al., 2019; Lim et al., 2019; Hataya et al., 2020; Zhang et al., 2020; Cubuk et al., 2020; LingChen et al., 2020; Zhang et al., 2020; Müller and Hutter, 2021;

Original: Zhou et al., 2021; Zheng et al., 2022; Cheung and Yeung, 2022), we focus on image classification datasets and employ a transformation search space comprising 14 operations, introduced by Cubuk et al.
Revised: Zhou et al., 2021b; Zheng et al., 2022; Cheung and Yeung, 2022), we focus on image classification datasets and employ a transformation search space comprising 14 operations, introduced by Cubuk et al.

Original: For those that have parameters, their magnitudes are discretized in thirty levels, which does not compromise performance and greatly reduces the search space (Cubuk et al., 2020).
Revised: For those that have parameters, their magnitudes are discretized in thirty levels, which does not compromise performance and greatly reduces the search space (Cubuk et al., 2020).

Original: Most approaches compose transformations, i.e. applying more than one transformation to the same image.
Revised: Most approaches compose transformations, i.e. applying more than one transformation to the same image.

Original: However, recently Müller and Hutter (2021) have shown that applying only one transformation at a time, defined over a wider magnitude space (noted Wide in Table 2) can outperform other approaches.
Revised: However, recently Müller and Hutter (2021) have shown that applying only one transformation at a time, defined over a wider magnitude space (noted Wide in Table 2) can outperform other approaches.

Original: We thus use the same transformation space as (Müller and Hutter, 2021).
Revised: We thus use the same transformation space as (Müller and Hutter, 2021).

Original: Table 2 from Müller and Hutter (2021) lists the operations and their magnitude ranges.
Revised: Table 2 from Müller and Hutter (2021) lists the operations and their magnitude ranges.

Original: In our experiments we used the wide search space, the standard ranges from Cubuk et al.
Revised: In our experiments we used the wide search space, the standard ranges from Cubuk et al.

Original: (2020) are included for comparison.
Revised: (2020) are included for comparison.

Original: We extend the codebase provided by Müller and Hutter (2021), which uses the
Revised: We extend the codebase provided by Müller and Hutter (2021), which uses the

Original: Pillow 2 implementation of all transformations except for cutout, and refer to its documentation for further details about image operations.
Revised: Pillow 2 implementation of all transformations except for cutout, and refer to its documentation for further details about image operations.

Original: [Table]
Revised: [Table]

Original: B. T RAINING S ETUP
Revised: B. T RAINING S ETUP

Original: In order to enable comparisons and reproducibility we use the same training pipeline as in previous works Müller and Hutter (2021) .
Revised: In order to enable comparisons and reproducibility we use the same training pipeline as in previous works Müller and Hutter (2021) .

Original: We apply the vertical flip and the pad-and-crop augmentations and a 16 pixel cutout (DeVries and Taylor, 2017) after any augmentation method.
Revised: We apply the vertical flip and the pad-and-crop augmentations and a 16 pixel cutout (DeVries and Taylor, 2017) after any augmentation method.

Original: We trained
Revised: We trained

Original: Wide-ResNet (Zagoruyko and Komodakis, 2016) models in the 40-2 and 28-10 configurations.
Revised: Wide-ResNet (Zagoruyko and Komodakis, 2016) models in the 40-2 and 28-10 configurations.

Original: Except for epoch ablation experiments [Table] we use SGD with Nesterov Momentum and a learning rate of 0.1, a batch size of 128, a 5e-4 weight decay [Table]
Revised: Except for epoch ablation experiments, we use SGD with Nesterov Momentum and a learning rate of 0.1, a batch size of 128, a 5e-4 weight decay, and cosine learning rate decay.

Original: and cosine learning rate decay.

Original: In ablation experiments we also trained for 600 epochs and used a custom learning rate schedule.
Revised: In ablation experiments we also trained for 600 epochs and used a custom learning rate schedule.

Original: For the first 185 epochs we followed the same cosine learning rate decay schedule, and then switching to a custom step learning rate scheduler detailed on Table 3.
Revised: For the first 185 epochs we followed the same cosine learning rate decay schedule, and then switching to a custom step learning rate scheduler detailed on Table 3.

Original: This schedule was implemented after observing that just scaling the cosine learning rate schedule to 600 epochs resulted in slow convergence, thus yielding solutions similar to training for 600 epochs, and failing reflect the effects of early stopping which this experiment addressed.
Revised: This schedule was implemented after observing that just scaling the cosine learning rate schedule to 600 epochs resulted in slow convergence, thus yielding solutions similar to training for 600 epochs, and failing reflect the effects of early stopping which this experiment addressed.

Original: We used 2 MH steps for the smoothed adversarial unless stated otherwise, and a learning rate of 10e-3 for the dual ascent step.
Revised: We used 2 MH steps for the smoothed adversarial unless stated otherwise, and a learning rate of 10 −for the dual ascent step.

Original: No other hyperparameters were tuned or modified with respect to standard settings.
Revised: No other hyperparameters were tuned or modified with respect to standard settings.

Original: The constraint level were set for Wide-ResNet-40-2 in the wide augmentation space by maximising three fold cross validation over the grid specified on Table 4.
Revised: The constraint levels were set for Wide-ResNet-40-2 in the wide augmentation space by maximising three fold cross validation over the grid specified on Table 4.

Original: Then the constraint levels were adjusted for other architectures and search spaces so that the dual variables at the end of training were small but not zero (of the order of 10 − 1 , which empirically showed good results.
Revised: Then the constraint levels were adjusted for other architectures and search spaces so that the dual variables at the end of training were small but not zero (of the order of 10 − 1 ), which empirically showed good results.

Original: The resulting constraint levels corresponding to the results in 1 are detailed in table 5.
Revised: The resulting constraint levels corresponding to the results in 1 are detailed in table 5.

Original: [Table]
Revised: We used the provided code and reported hyperparameters for running TrivialAugment (Müller and Hutter, 2021) and DeepAA (Zheng et al., 2022). For the latter, unlike the results reported in (Zheng et al., 2022), we kept all evaluation hyperparameters (including weight decay) consistent with that of other methods. Results for the wide augmentation space were not included in (Zheng et al., 2022).

Original:  [Table]
Revised: We thus performed the augmentation policy search for the wide search space using the same hyperparameters (except for the augmentation space) reported in (Zheng et al., 2022) for CIFAR datasets. We also run the search for both augmentation spaces in SVHN datasets, and evaluated the policy with the same setup as other methods. [Table] [Table]

Original: [Table]
Revised: [Table]

Original: C A DDITIONAL RESULTS
Revised: C A DDITIONAL E XPERIMENTS C. A UTOMATIC D ATA AUGMENTATION

Original: The following section contains additional ablations and discussions about our algorithm.
Revised: The following section contains additional ablations and discussions about our algorithm.

Original: As in previous sections, we use the wide Müller and Hutter (2021) augmentation space and CIFAR image classification benchmarks.
Revised: As in previous sections, we use the wide Müller and Hutter (2021) augmentation space and CIFAR image classification benchmarks.

Original: Our main motivation is to analyse how the different hyperparameter choices and the learning algorithm affect the generalization and invariance of the obtained solutions.
Revised: Our main motivation is to analyse how the different hyperparameter choices and the learning algorithm affect the generalization and invariance of the obtained solutions.

Original: First, we show how the dual variables adapts to different constraint levels during training, and link its dynamics to heuristically defined learning curricula.
Revised: First, we show how the dual variables adapts to different constraint levels during training, and link its dynamics to heuristically defined learning curricula.

Original: We then evaluate the effect of training for more epochs and link our observations to the known properties of early stopping in unconstrained learning, in section C.2.
Revised: We then evaluate the effect of training for more epochs and link our observations to the known properties of early stopping in unconstrained learning, in section C.1.2.

Original: In section C.3, we analyse how the sampling approximation affects regularisation, by performing an ablation on the number of MH steps.
Revised: In section C.1.3, we analyse how the sampling approximation affects regularisation, by performing an ablation on the number of MH steps.

Original: Lastly, we include the observed frequencies of sampled transformations for different setups, so as to obtain further insights in how the distribution adapts throughout training (Section C.4).
Revised: Lastly, we include the observed frequencies of sampled transformations for different setups, so as to obtain further insights in how the distribution adapts throughout training (Section C.1.4).

Original: C. D UAL V ARIABLE D YNAMICS
Revised: C.1. D UAL V ARIABLE D YNAMICS

Original: As already mentioned, the dual variables control the weight of augmented samples during training, thus balancing the trade-off between fitting the primal objective (i.e. loss over training samples) and satisfying the constraints (i.e. loss over transformed samples).
Revised: As already mentioned, the dual variables control the weight of augmented samples during training, thus balancing the trade-off between fitting the primal objective (i.e. loss over training samples) and satisfying the constraints (i.e. loss over transformed samples).

Original: In penalised approaches, on the contrary, this trade-off is fixed.
Revised: In penalised approaches, on the contrary, this trade-off is fixed.

Original: In Figure 3 we show how the dual variable adapts to different constraint levels, for Wide-ResNet-40-2 in CIFAR datasets using the standard setup.
Revised: In Figure 4 we show how the dual variable adapts to different constraint levels, for Wide-ResNet-40-2 in CIFAR datasets using the standard setup.

Original: Note that for stricter constraint levels, the algorithm has not converged when it reaches 200 epochs.
Revised: Note that for stricter constraint levels, the algorithm has not converged when it reaches 200 epochs.

Original: We also observe that the dynamics of our primal-dual algorithm resembles the augmentation learning curricula proposed by Blaas et al. 
Revised: We also observe that the dynamics of our primal-dual algorithm resembles the augmentation learning curricula proposed by Blaas et al. [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Original: C. E ARLY S TOPPING
Revised: C.1. E ARLY S TOPPING

Original: Our approach has slower convergence than unconstrained approaches and stopping training arbitrarily - after a fixed number of epochs - can result in solutions that are unfeasible or sub-optimal.
Revised: Our approach has slower convergence than unconstrained approaches and stopping training arbitrarily - after a fixed number of epochs - can result in solutions that are unfeasible or sub-optimal.

Original: However, early stopping is a popular regulariser, particularly for neural networks trained through gradient descent.
Revised: However, early stopping is a popular regulariser, particularly for neural networks trained through gradient descent.

Original: Several empirical (Caruana et al., 2000; Rice et al., 2020) and theoretical (Ji et al., 2021;
Revised: Several empirical (Caruana et al., 2000; Rice et al., 2020) and theoretical (Ji et al., 2021;

Original: Li et al., 2020a; Duvenaud et al., 2016) results show its advantages in terms of generalisation and robustness to noisy training data.
Revised: Li et al., 2020a; Duvenaud et al., 2016) results show its advantages in terms of generalisation and robustness to noisy training data.

Original: In general, the advantages of early stopping do not lie in the sub-optimality of the solution in terms of training error, but on nature of the regularization or prior imposed, which leads to non-uniform model selection among models with a given training error (Cataltepe et al., 1999).
Revised: In general, the advantages of early stopping do not lie in the sub-optimality of the solution in terms of training error, but on nature of the regularization or prior imposed, which leads to non-uniform model selection among models with a given training error (Cataltepe et al., 1999).

Original: To the best of our knowledge, there is no literature that explicitly addresses early stopping in empirical primal-dual learning.
Revised: To the best of our knowledge, there is no literature that explicitly addresses early stopping in empirical primal-dual learning.

Original: Whether the generalisation gap in constraint satisfaction can be reduced by early stopping regularisation, in the same manner early stopping regularisation can reduce the generalisation gap in unconstrained learning, thus remains unclear.
Revised: Whether the generalisation gap in constraint satisfaction can be reduced by early stopping regularisation, in the same manner early stopping regularisation can reduce the generalisation gap in unconstrained learning, thus remains unclear.

Original: Figure 4 shows an ablation on the number of epochs.
Revised: Figure 5 shows an ablation on the number of epochs.

Original: Non-zero dual variables and strict feasiblility show the algorithm has not yet converged at 200 epochs.
Revised: Non-zero dual variables and strict feasiblility show the algorithm has not yet converged at 200 epochs.

Original: At 600 epochs whereas constraint satisfaction shows little change, training loss decreases and the generalization gap increases.
Revised: At 600 epochs whereas constraint satisfaction shows little change, training loss decreases and the generalization gap increases.

Original: Thus, we observe early stopping has a larger impact on the primal objective than on the constraint.
Revised: Thus, we observe early stopping has a larger impact on the primal objective than on the constraint.

Original: that although for stricter levels of the constraint the algorithm has not converged when training is stopped atepochs, it can yield solutions that are still feasible and have a smaller generalization gap.
Revised: that although for stricter levels of the constraint the algorithm has not converged when training is stopped atepochs, it can yield solutions that are still feasible and have a smaller generalization gap.

Original: [Figure]
Revised: [Figure]

Original: C.  S AMPLING S TEPS A BLATION
Revised: C.1.  S AMPLING S TEPS A BLATION

Original: As already mentioned in Section 4.3, we used the Metropolis-Hastings algorithm with independent uniform proposals.
Revised: As already mentioned in Section 4.3, we used the Metropolis-Hastings algorithm with independent uniform proposals.

Original: As shown in Figure 5, using more steps of the chain allows samples to deviate further from the uniform distribution, as measured by the decrease in entropy.
Revised: As shown in Figure 6, using more steps of the chain allows samples to deviate further from the uniform distribution, as measured by the decrease in entropy.

Original: For a fixed constraint level, augmentation distributions with higher entropy result in weaker regularization.

Original: As already mentioned, dual variables give information of the trade-off between fitting clean and augmented data.
Revised: As already mentioned, dual variables give information of the trade-off between fitting clean and augmented data.

Original: We observe the final value of the dual variable is highly correlated with the entropy of sampled  transformations and final training loss.
Revised: We observe the final value of the dual variable is highly correlated with the entropy of sampled transformations. That is, transformations that deviate further from uniformity result in larger dual variables at the end of training.

Original: In 6 we show the evolution of dual variables for different sampling steps.
Revised: Furthermore, in 7 we show the evolution of dual variables for different sampling steps.

Original: Sampling distributions that are closer to worst case perturbations result in more stringent requirements, and thus dual variables adapt accordingly.
Revised: Sampling distributions that are closer to worst case perturbations results in more stringent requirements, and thus dual variables grow more rapidly. E n t r op y CIFAR     2. CIFAR

Original: [Figure]
Revised: [Figure]

Original: opy

Original: [Figure]
Revised: [Figure]

Original:  u a l
Revised: D u a l

Original: V a r i a b l e ( )
Revised: V a r i a b l e ( )

Original: [Figure]
Revised: [Figure]

Original: The constraint level is fixed for each dataset (0.8 in CIFAR10 and 2.1 for CIFAR100).
Revised: The constraint level is fixed for each dataset (0.8 in CIFAR10 and 2.1 for CIFAR100).

Original: We compute the entropy of the augmentations sampled at the last epoch of training.
Revised: We compute the entropy of the augmentations sampled at the last epoch of training.

Original: Thegap between train and test loss and the final value of the dual variable, and the augmentation distribution entropy increase with the number of steps.
Revised: The value of the dual variable at epoch 200 increases with the number of steps, whereas the augmentation distribution entropy decreases. Each point represents an independent run. Experiments were repeated four times.

Original: [Figure]
Revised: [Figure]

Original: C. S AMPLED TRANSFORMATIONS
Revised: C.1. S [Figure] T

Original:  We now include plots of the empirical transformation distributions.
Revised: [Figure] We now include plots of the empirical transformation distributions.

Original: As shown in Figure 7, the frequency with which transformations are sampled varies depending on the dataset, which is a desirable property.
Revised: As shown in Figure 8, the frequency with which transformations are sampled varies depending on the dataset, which is a desirable property.

Original: Similar to Zhang et al.
Revised: Similar to Zhang et al.

Original: (2020), we observe a prevalence of geometric transformations, unlike Cubuk et al.
Revised: (2020), we observe a prevalence of geometric transformations, unlike Cubuk et al.

Original: In SVHN color transformations are less frequent than in CIFAR datasets, and the frequency of geometric transformations increases, which is interesting due to the perceptual importance of shape in the digit recognition task.
Revised: In SVHN color transformations are less frequent than in CIFAR datasets, and the frequency of geometric transformations increases, which is interesting due to the perceptual importance of shape in the digit recognition task.

Original: As already noted, using more steps allows the chain to deviate further from the uniform proposal distribution.
Revised: As already noted, using more steps allows the chain to deviate further from the uniform proposal distribution.

Original: We thus plot the frequency of sampled transformations accross training epochs on
Revised: We thus plot the frequency of sampled transformations across training epochs on

Original: Figure 8.
Revised: Figure 9.

Original: We observe that as training progresses, the observed frequencies also deviate further from uniformity.
Revised: We observe that as training progresses, the observed frequencies also deviate further from uniformity.

Original: This suggests that, due to the dataset and architecture, some transformations may be harder to fit than others.
Revised: This suggests that, due to the dataset and architecture, some transformations may be harder to fit than others.

Original: Similarly, we include histograms for the sampled transformation levels in Figure 9.
Revised: Similarly, we include histograms for the sampled transformation levels in Figure 10.

Original: Extreme levels are sampled more frequently, but the empirical distributions vary depending on the transformation.
Revised: Extreme levels are sampled more frequently, but the empirical distributions vary depending on the transformation.

Original: The deviation from uniformity and the differences between transformations are accentuated as sampling steps increase (Figure 10).
Revised: The deviation from uniformity and the differences between transformations are accentuated as sampling steps increase (Figure 11). [Figure] [Figure] [Figure]

Original: [Figure]
Revised: [Figure]

Revised: C.1. R UNTIME ANALYSIS

Revised: As mentioned in section 4.3, the added computational cost of our algorithm is that of computing a forward pass for each mcmc step.

Revised: As a result, trade-off between sampling λ ⋆c accurately and computation arises.

Revised: Table 6 provide an empirical runtime analysis for our method for different numbers of MH steps, and compare it with the training time of baseline methods.

Revised: These times correspond solely to training, and it is difficult to account for the time taken to tune the hyperparameters of each method, which hinders direct comparisons.

Revised: In the case of DeepAA (Zheng et al., 2022), it requires running a data augmentation policy search that takes 11 hours (more than 5 × training time) using the same hardware.

Revised: [Table]

Revised: D ATA AUGMENTATION CAN HINDER PERFORMANCE

Revised: As already mentioned in section A.2 there is empirical evidence that certain distributions over commonly used transformations can introduce biases that are detrimental to model performance and generalization (Blaas et al., 2021).

Revised: We provide a simple experiment to show that there exist transformation distributions that make data augmentation deteriorate model performance, and that in such cases balancing the amount of data augmentation (e.g. by including the original data) is important to mitigate and overcome this issue.

Revised: We restrict the transformations in the wide augmentation space (Müller and Hutter, 2021) to their maximum magnitude.

Revised: We sample transformations according to λ

Revised: ⋆c using two MH steps as previously described.

Revised: We compare against training without augmentation, and training using both augmentation and the original data (i.e. adding the identity) equally weighted.

Revised: While sampling transformations according to λ ⋆c deteriorates performance with respect to the model without augmentation, including both the identity and the augmented samples achieves a superior performance, as shown in Table 7.

Revised: [Table]

Revised: C. S YNTHETIC I NVARIANCES

Revised: Although our approach does not aim to learn the set of symmetries or invariant transformations from the data, but rather to impose it on the predictor, dual variables can be used to gather insights on the actual invariances underlying a learning task.

Revised: We showcase this on datasets with artificial invariances, following the setup of Immer et al.

Revised: Namely, we apply transformations independently drawn from the distributions specified in Table 8, to each sample in the datasets MNIST (LeCun et al., 2010) and FashionMNIST (Xiao et al., 2017).

Revised: We use the same MLP and CNN architectures and hyperparameters as (Immer et al., 2022).

Revised: We run our algorithm constraining the loss on the transformation spaces detailed in Table 9, which (except for the fully rotated dataset) are larger than the true transformation range used to construct the synthetic dataset.

Revised: Note that we use the same transformation sets and constraint levels ( ϵ ) for all synthetic datasets.

Revised: As shown in table 10, except for scalings in FashionMNIST, the dual variables ( γ ) associated with transformations corresponding to the true synthetic invariances in the dataset are considerably smaller.

Revised: This indicates that when the transformations in the constraint correspond to a true invariance of the dataset, the constraint is easier to satisfy.

Revised: Since the transformation ranges in the constraint and those used to construct the dataset do not match exactly (except for the fully rotated dataset), the degree of invariance to the transformation sets in the constraints is only approximate.

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: D R ELATED N OTIONS OF INVARIANCE

Revised: There are several definitions of invariance that capture different properties of the solution or data distribution under the action of transformations in G .

Revised: In the context of supervised learning, the data distribution is said to be exactly invariant iff it does not change when transformations are applied to the covariates, i.e.,

Original: [Figure]
Revised: [Equation]

Revised: , where = d denotes equality in distribution.

Revised: Note that this can be relaxed to approximate equality in distribution, i.e. that the distribution shift induced is small in an appropriate metric, e.g. Wassertain distance (Chen et al., 2019).

Revised: This definition implies that the probability of a given label is - exactly or approximately - the same for a feature vector x and its transformed versions g x :

Original: [Figure]
Revised: [Equation]

Revised: On the other hand, a predictor can satisfy hard invariance, explicitly,

Original: [Figure]
Revised: [Equation]

Revised: which needs to hold point-wise over the transformation set and input space.

Revised: This can be relaxed by taking the mean over the data distribution to obtain a definition of statistical invariance , i.e.,

Revised: [Equation]

Revised: which still needs to hold for every transformation in G .

Revised: This notions of invariance do not explicitly contemplate the task at hand, that is, not all changes in y equally affect performance.

Revised: Thus, we use the loss to encode meaningful differences in labels with respect to the task, as described in section 3.1.

Revised: Throughout this work, we thus refer to the invariant risk R inv defined on equation 5 as the degree of invariance, unless otherwise noted.