SAFLEX: Self-Adaptive Augmentation via Feature Label Extrapolation
Authors: Mucong Ding
Bang An
Yuancheng Xu
Anirudh Satheesh
Furong Huang
Venue: ICLR 2024 poster (2024)
Keywords: Data Augmentation
Learnable Augmentation
Bilevel Optimization
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We introduce SAFLEX, a bilevel optimization method for data augmentation, achieving promising results across multiple data types and tasks and seamlessly integrating with existing augmentation techniques.

Abstract: Data augmentation, a cornerstone technique in deep learning, is crucial in enhancing model performance, especially with scarce labeled data. While traditional techniques are effective, their reliance on hand-crafted methods limits their applicability across diverse data types and tasks. Although modern learnable augmentation methods offer increased adaptability, they are computationally expensive and challenging to incorporate within prevalent augmentation workflows. In this work, we present a novel, efficient method for data augmentation, effectively bridging the gap between existing augmentation strategies and emerging datasets and learning tasks. We introduce SAFLEX (Self-Adaptive Augmentation via Feature Label EXtrapolation), which learns the sample weights and soft labels of augmented samples provided by any given upstream augmentation pipeline, using a specifically designed efficient bilevel optimization algorithm. Remarkably, SAFLEX effectively reduces the noise and label errors of the upstream augmentation pipeline with a marginal computational cost. As a versatile module, SAFLEX excels across diverse datasets, including natural and medical images and tabular data, showcasing its prowess in few-shot learning and out-of-distribution generalization. SAFLEX seamlessly integrates with common augmentation strategies like RandAug, CutMix, and those from large pre-trained generative models like stable diffusion and is also compatible with frameworks such as CLIP's fine-tuning. Our findings highlight the potential to adapt existing augmentation pipelines for new data types and tasks, signaling a move towards more adaptable and resilient training frameworks.

Published as a conference paper at ICLR 2024
SAFLEX: SELF- AUGMENTATION VIA FEA- TURE LABEL EXTRAPOLATION
Mucong Ding, Bang An, Yuancheng Xu, Anirudh Satheesh, Furong Huang Department of Computer Science
University of Maryland medi ingecs cs. umd.edu
ABSTRACT
Data augmentation, cornerstone technique in deep learning, is crucial in enhanc- ing model performance, especially with scarce labeled data. While traditional methods, such as hand-crafted augmentations, are effective but limited in scope, modern, adaptable techniques often come at the cost of computational complexity and are hard to fit into existing processes. In this work, we unveil efficient approach that universally enhances existing data augmentation techniques by en- abling their adaptation and refinement, thereby providing significant and com- prehensive improvement across all existing methods. We present SAFLEX (Self- Adaptive Augmentation via Feature Label EXtrapolation), approach that uti- lizes an efficient bilevel optimization to learn the sample weights and soft labels of augmented samples. This is applicable augmentations from any source, seam- lessly integrating with existing upstream augmentation pipelines. Remarkably, SAFLEX effectively reduces the noise and label errors of the upstream augmenta- tion pipeline with marginal computational cost. As a versatile module, SAFLEX excels across diverse datasets, including natural, medical images, and tabular data, showcasing its prowess in few-shot learning and out-of distriiiittiittion generaliza- tion. SAFLEX seamlessly integrates with common augmentation strategies like RandAug and CutMix, as well augmentations from large pre-trained gener- ative models like stable diffusion. It is also compatible with contrastive learn- ing frameworks, such as fine-tuning CLIP. Our findings highlight the potential to adapt existing augmentation pipelines for new data types and tasks, signaling a move towards more adaptable and resilient training frameworks.
INTRODUCTION
Data augmentation is cornerstone in improving machine learning models, especially when labeled data is scarce. enhances model performance by introducing varied training samples. Though traditional methods like rotation and cropping are widely used, they operate under ame-size-fits-ae assumption that often falls short in the complexity of real-world data. The key is not just to augment data, but do it in way that does not mislead the learning process.
not.uust
Recent work emphasizes the benefits of learned data augmentation, where techniques such as Au- toAugment (Cubuk al., 2019) and RandAugment (Cubuk al. (2020) adapt to specific datasets and tasks. While promising, this area is still nascent and lacks a comprehensive framework to ad- dress diverse tasks and data nuances. Furthermore, selecting meaningful transformations remains challenge, often relying on heuristics or domain expertise, which is especially problematic in spe- cialized fields. Inappropriate transformations can harm model performance, underscoring the need for systematic selection. Amid the rise of image generation methods, such as diffusion models and other generative AI, an abundance of synthetic data is available but requires discerning use. A recent study, LP-A3 (Yang et al., 2022a). aims to generate "hard positive examples" for aug- mentation but risks introducing false positives that could mislead learning. Another recent work, Loft-Augmentation (Liu et al. 2023), introduces soft learning targets and loss reweighting train on augmented samples but is primarily limited to improving image crop augmentation. The overar- ching need is for smarter, more adaptable data augmentation algorithms.
This paper proposes SAFLEX Adalt-Ade Augmentation via Feature Label Extrapolation). which automatically learns the sample weights and soft labels of augmented samples provided by any given

---

Published as a conference paper at ICLR 2024
Training Samples
Model
Augmented Samples
Upstream Augment
Pipeline
Sample Weights
Soft Labels
Grodient matching
1-+ 0.953
Augmented
Samples
Validation
Figure 1: SAFLEX learns to adjust sample weights and soft labels of augmented samples from an upstream pipeline, aiming to maximize the model's performance on the validation set. While formulated as a bilevel optimization problem, it can be efficiently solved by linear programming with a gradient-matchin objective. SAFLEX is a plug-in to the existing training framework.
upstream augmentation pipeline. Existing learnable augmentation methods that directly learn in the feature space (e.g., image space) often restrict augmentation scope due to differentiability needs and suffer from complicated training in high-dimensional spaces. Contrary to this, we advocate for learning only low-dimensional sample weights and soft labels for each augmented instance sourced from pre-existing upstream augmentation pipeline like synthetic data generation. While upstream augmentation methods can sometimes alter labels or introduce noise, especially when creating sam- ples outside the data distribution, our approach offers mechanism to correct them. By calibrating sample weights and labels after augmentation, we considerably alleviate issues stemming from up- stream augmentation methods. Without the complexity of learning augmentation transformations from scratch, this strategy ensures that augmentation is both diverse and consistent with the inher- ent data distribution, thereby fostering better generalization across various tasks. See Fig. for demonstration of our proposed SAFLEX.
We frame learning sample weights and soft labels as a bilevel optimization problem. This captures the interdependent nature of the model and its augmented data: the model's performance depends on the quality of the augmented data, which in turn is guided by the model itself (Bard 2013). This new perspective advances our understanding of data augmentation, offering theoretical framework that underpins its practical applications. Despite the bilevel nature of the problem, direct solutions are computationally infeasible for large-scale real-world applications. To mitigate this, we propose a streamlined, greedy, online, single-level approximation algorithm, which optimizes a gradient- matching objective to accelerate the learning process.
We conducted extensive empirical evaluations to highlight SAFLEX'S superior performance and adaptability. On eight medical images (Yang et al. 2023), SAFLEX elevates popular augmentation techniques like RandAugment (Cubuk et al. (020)) and Mixup (Zhang et al. 2018). boosting per- formance by up to 3.6%. On seven tabular datasets, SAFLEX shows compatibility with categorical data and effectively enhances CutMix (Yun et al., 2019). Furthermore, SAFLEX improves image augmentations from diffusion models, yielding an average improvement of 1.9% in fine-grained generalization against three diffussion-augmmnnation methods,
classification and harnessing on their pre- pre-trained expertise. We also validate SAFLEX'S 's integration with contrastive learning through CLIP fine-tuning experiment. These findings underline its versatility across var- ied data types and learning tasks.
Our contributions are threefold:
(1) We unveil a novel parametrization for learnable augmentation complemented by an adept bilevel algorithm primed for online optimization.
(2) Our SAFLEX method is distinguished by its universal compatibility, allowing to be effortlessly incorporated into a plethora of supervised learning processes and to collaborate seamlessly with an extensive array of upstream augmentation procedures.
(3) The potency of our approach is corroborated by empirical tests diverse spectrum of datasets and tasks, all underscoring SAFEEX'S efficiency and versatility, boosting performance by1.2% on average over all experiments.
2 PROPOSED METHOD: SAFLEX
Our goal is to refine augmented samples from any upstream pipeline to enhance classifier general- ization. The proposed methodology is founded on two pivotal questions: (1) Which aspects of the
2

---

Published as a conference paper at ICLR 2024
augmented samples should be refined? (2) What approach should be taken to learn these refined samples? We start from these questions and defer the derivation of the algorithm to Section 3
Limitations of Augmentation Methods. Data augmentation is pivotal in enhancing model gener- alization. However, its limitations, particularly the unintentional introduction of noise, can some- times outweigh its benefits. For instance, consider the widespread use of random cropping on nat- ural images. Although largely effective, there are times when this approach inadvertently omits task-relevant information, leading unintended outcomes like false positíves. This inherent noise creates a trade-off: under-augmentation may yield insufficient challenging examples, whereas over- augmentation can flood the dataset with misleading samples. As shown in Fig. 2a reducing the noise in augmentation is the key to resolving the dilemma.
Noise in augmentation primarily arises from two fundamental challenges: (1) the deviation of aug- mented samples from the original data distribution and (2) the potential mislabeling of augmented samples. We shall envision augmentation as a method to harness prior knowledge in capturing the underlying data distribution. This distribution is encapsulated in the joint distribution, P Pxy(2; (az, y). where x E X' are features and y E (I...., }} represents labels, with K indicating the number of classes. Breaking down this joint distribution: Pxy (x. v) (x(x)) Py|x (1(zz), we observe that the primary source of noise is associated with the feature distribution Px (a),, while the secondary source is tied to the conditional distribution y (y)x)) Addressing these challenges, our method- ology is designed to integrate seamlessly with any upstream augmentation process, amending both types of errors post-augmentation, and considering the initial augmentation process as a separate, unchanged entity.
Feature and Label Extrapolation. A key concern in data augmentation pertains to addressing these two types of errors. Some prior works on learning augmentation (c.g., (Yang al. 2022a)) attempted to reduce noise by fine-tuning augmented features, using them as inftializations. Specif- ically, the aim was to derive a modified feature x' that eliminates both error types. Yet, due to the high-dimensionality of feature space X, manipulating æ is computationally burdensome.
A more efficient strategy is to handle the errors individually and abstain from modifying x. When encountering erroneous estimation of the feature distribution Px (2), even if augmented samples lie in low-density areas, we can compensate by modulating the sample weights w [0, 1] in the em- pirical risk minimization loss. Specifically, rarer augmented features are assigned decreased sample weights. For inaccuracies in estimating the conditional distribution yy (y/z), it's advantageous to modify the augmented label y directly. We also propose transitioning from a hard class label to soft one, denoted as simplex E A X representing probability mass across K classes, residing in the K K-diiensionnal The proposed refinement of augmented samples is depicted in Eq. (1. Remark- ably, optimizing these sample weights and soft labels effectively mitigates errors resulting from varied augmentation methods across numerous classification challenges.
(I, v)
Upstream Augment
(nun aug
SAFLEX
ang
yaug soft label A
(1)
sample weight [0, 1]
To elucidate, consider hypothetical example in Fig. 2b Envision training sample from the green class (represented by a pronounced green dot). Upon applying a noise-prone augmentation, such as Gaussian perturbation in 2D setting, the augmented sample could either (1) fall into a region with few validation samples regardless of their class, or (2) be overwhelmingly encompassed by valida- tion samples from different class. In the former case, it is judicious to reduce the sample weights since they might not be pivotal in discerning the conditional distribution. In the latter instance, the label of the augmented sample should be fine-tuned. This can entail a shift to a soft label to rectify or mitigate potential label inconsistencies, informed by patterns in the validation set.
Bilevel Formulation. The remaining question our design is how to learn the sample weights and soft labels for augmented samples. The overarching goal of augmentation is enhancing model generalization. While the test set remains inaccessible, a prevalent approach is to fine-tune perfor- mance using validation set. This methodology aligns with standard practices in hyperparameter optimization and is evidenced in learnable augmentation methods such as AutoAugment (Cubuk) et al., 2019) and RandAugment (Cubuk et al. 22020). Given a neural network f (-) X A (where we assume Softmax is already applied and the outputs are L, normalized) with parameter 8, let us denote the training set, validation set, and the set of augmented samples as Derain, Dent, and
3

---

Published as a conference paper at ICLR 2024
Reduce sample weight
Augmented
1L0523
Coected Genec Pererrnnce
Original
SAFLEX
anu
Adjus soft
Augmented
Storgh Aogneeraeee
Ieagetads,
(a) Mitigating trade-offs.
(b) Addressing two types of error..
Figure 2: (a) Under-augmentation can lead to a scarcity of hard positives, while over-augmentation can introduce an excess of false positives. Reducing the noise in augmentation helps resolve the dilemma. (b) Adjusting sample weights and recalibrating soft labels can address the two types of noises introduced by the augmentation process.
D Daug. respectively. The ambition is to refine Daug such that a model trained on the amalgamation of Dtrain U Daug optimizes performance vall-
min [(Dval, 0) s.t. 0 arg min ((Prrann U D Dang. : 0')
This scenario can be east 0 augmented samples with parametrized by sample weights and soft labels, and the model parameters 0 are learnable. The conventional model training constitutes the inner level, while the quest to iden- tify optimal augmented samples Daug. which minimize the validation loss post-inner level training, establishes the outer problem. Such a paradigm inherently transforms learnable augmentation into bilevel optimization. Intriguingly, much of the existing literature on learnable augmentation eschews this representation. The primary reservations stem from concerns related to efficiency and differen- tiability. Notably, works such as (Mounsaveng et al. 2021 2023) are among the sparse few to apply bilevel optimization for augmentation learning, yet their focus remains constricted to affine transfor- mations. In contrast, approach sidesteps the modification and modeling of feature augmentation, obviating the challenge of differentiability. The low-dimensional nature of sample weights and soft labels potentially simplifies the learning process. In subsequent sections, we demonstrate that, un- der benign approximations, we can adeptly navigate the bilevel problem, determining the apt sample weights and soft labels within singular step for each training iteration.
as bilevel optimization problem as in Eq. ( where Daug. the set of
3 ALGORITHM
We now develop an algorithm for the bilevel problem described in Eq. (
The Greedy Approach. Bilevel optimization is notoriously challenging, often necessitating nested loops, which introduces significant computational overhead. Upon inspecting Eq. (2). it becomes evident that an essential characteristic of the problem the training dynamics of the model has been understated. In standard practice, augmented samples are typically generated during model training for each minibatch across all iterations. Therefore, the actual problem deviates from Eq. ( in two ways: (1) different batches augmentation may influence the learned parameters differently. and the model is not trained on cumulative set of augmented samples, and conversely, (2) the learned parameters are affected differently by the refined augmented samples across batches, imply- ing that augmentation should be optimized with respect to the corresponding model parameters.
To incorporate model optimization dynamics, we should reformulate the problem on a finer scale: Given the model parameter 0.-1 at intermediate training step, how can we determine the batch of refined augmented samples, Dhatch? Through a greedy approach, we posit that the granular objective is to minimize the validation loss after single update, denoted as C(Dval, 0z), where 0, is the model parameter updated from 0.-1.
min L(Dval, 0z) s.t. 0, = 00-1--1 - a- V VoL batch U Dbatch .0.1)) train
(3)
Dbbach
This micro-persppetive of Eq. (2) is represented in Eq. ples Dbatch [w aug .1 aug .y ( ),..,,... .,. (w ang IB
where the batch of augmented sam- is parametrized by the set of sample
Aug B
weights (w) mnn ..... ww aug) and soft labels (y) . y ang
As a direct consequence, if the inner loop uses a first-order optimizer like SGD (as assumed), this significantly eases the optimization task. The emergent problem is no longer bilevel. With the analytical solution of the "inner problem" at our disposal, we can integrate the formula for 6, into the outer objective, L(Deat, 0.), converting into single-level problem.
4

---

Published as a conference paper at ICLR 2024
Efficient Solution. We next derive an algorithm for efficiently addressing Eq. 33) to the linearity of the loss function C(,, 8(-1) with respect to datasets and the inherent linearity of gradient computation, the gradient vector for the combined training and augmentation batch linearly relates to the sample weights and soft labels, assuming the sample-wise loss function, such as the cross-entropy loss, behaves linearly with respect to sample weights and soft labels. Validating this, the cross-entropy loss is indeed linear concerning these variables, typical characteristic based on their definitions.
Crucially, due
By approximating the validation loss ((D 0z) up to the first order around parameter Or, we can recast Eq. (3) as a linear programming problem. The objective seeks to maximize the inner product of the gradient vectors on the combined train and augmented batch and the validation batch, effec- tively yielding gradient-matching loss (Zhao et al. 2020). Here, both the objective and normaliza- tion constraints linearly correspond to our learnable variables: sample weights and soft labels. The derivation provided in Appendix A and we summarize the solution in the subsequent notations and theorem.
Notation 1. Let the Jacobian matrix of logits with respect to the model parameter be Vol (Iam lo=0,_, RxK and the gradient vector on the validation set be VoL(Dvva, 0t-1) E RT, where m is the parameter count and K is the class count. The Jacobian- vector product is 10=0, VaL(Dvat, 0z-1) E RK which can be computed efficiently.
denoted as II = Vol (2a48)
Theorem (Solution of Eq. (3). The approximated soft label solution y
OneHot (arg max; kk), where OneHor(-) denotes one-hot encoding, and the sample weight so- K & 2 O; otherwise, w O.
lution is w if Ek= [
Theorem illustrates that an effective approximation of Eq. (3) is computationally efficient. The gradient inner product, II. a Jacobian-vector product, readily computed alongside standard back-propagation on the combined training and augmentation batch. While determining the val- idation set gradient vector mandates an additional back the gradient vector for the complete validation set, Ve -1), using a minibatch gradient, V VDL
step, we can approximate
val 0,-1). Despite necessitating solution for Eq. (3) at every iteration, our efficient SAFLEX algorithm incurs minimal computational overhead.
A notable takeaway from Theorem is that while we aim to learn continuous sample weights (in [0, 1]) and soft labels (in AK), the derived solutions consistently yield discrete values: either 0 or and one-hot vectors. This consistency does not signify a coarse approximation, especially considering we resolve Eq. (3) with a ((a) tolerance, where is typically small. Nonetheless, this characteristic could potentially impact model generalization in under-parameterizes scenarios.
Generalization Aspects. Let's interpret and examine the solution provided by Theorem ! from generalization standpoint, which is our primary objective. The loss function's lincarity helps un- derstand Theorem [. Given that Eq. (3) is linear program with straightforward normalization constraints, we effectively form a linear combination of K gradient vectors (each pertaining to a logit of the augmented sample), with B representing the augmented batch size, to approximate the m-dimmensionll validation gradient vector. The total constraints sum up to B 1. If these K B gra- dient vectors are linearly independent, we can always align the combined gradient vector with the validation gradient vector when the degree of freedom, B + K (B+1), is greater than or equal to the gradient vector dimension, m. This is represented by the condition KB > m. Such a scenario, exceedingly under-parametrized is rare in deep learning. If the combined gradient vector consis- tently aligns with the validation gradient vector, training with SAFLEX will approximate training on the combined training and validation sets, potentially limiting the generalization improvements.
To enhance generalization, it is essential to circumvent the challenges of the under-parametrizes paradigm, even if we are not closely approaching it. Here, we suggest two modifications to the solution given by Theorem [
1. Encouraging Retention of the Original Label. We can introduce minor constant penalty term to the gradient inner product to incentivize retaining the augmented sample's 's original label. Thus, substitute II with II "us, where uy is one-hot vector with value of 1 at the -th position. If no other entry in II exceeds [I]] y by a margin of at least 3, the learned label remains unaltered. This approach proves especially valuable when the validation set is of limited size.
2. Substituting arg max with Huubel--Softaan Our current solution invariably yields hard labels. This can sometimes manifest as an excessive degree of confidence, particularly when II contains
5

---

Published as a conference paper at ICLR 2024
multiple significant entries. To alleviate this, we can employ the Gumbel- SoftMax function to intro- duce a "softening" effect to the learned labels, adding a measure of stochasticity. Hence, we have y softmax (II + Bey + 8)),), where g consists of i..d. random variables sourced from Gumbel(O, 1). Typically, unless specified otherwise, we opt for relatively low fixed temperature value, T = 0.01.
The pseudo-code of SAFLEX for cross-entropy loss is shown Algorithm
Algorithm 1: SAFLEX (Cross-Entroppy Loss, Single batch).
Input:
Neural network f (-) : X + A (softmax applicd on outputs) with parameters 8, upstream
augmented batch {2] aus ")""),.... ( ") validation batch
D {(xxim] : Vi x J,..., (m U }, penalty coefficient (9, temperature 7.
Compute the gradient vector for the validation batch VoL(Dbatch 0).
for i 1,..,,
vectorized, Determine the gradient inner product II, = Vaf (z,"TK VeL(Diatch val 0) via Jacobian product.
// The actual inp
Apply Gofbel-SofiMax to get y; soft (II, + Se,,
: + g)/+). where eyp E B one-hot at y; aug and consists of i.idd random variables taken from Gumbel(O, 1).
Set wi = 1 if II; y: 2 0, otherwise set Renormalize the sample weights Wi....; BB to sum to aug and soft labels y 7 return Sample weights w ,..., Wa
0.
3ug
SAFLEX for Contrastive Learning. We conclude this section by discussing to encompass the gen- eralization of the proposed method for certain contrastive learning losses, as illustrated in Eq. (12)
and Eq. (3). Notably, the latter is utilized for CLIP training. In the realm of contrastive learning, labels are not conventionally defined. Yet, one can perceive the contrastive training objectives in Eq. (12) and Eq. (13) as proxy classification tasks. Here, we posit that the batch of size B can be construed as containing B classes: one positive example coupled with B-1 1 negative examples. This interpretation paves the way to introduce the notion of (soft) labels over this surrogate classification task with its B distinct classes.
Under this paradigm, the loss function remains linear concerning the soft labels and sample weights, making the methodology in Theorem I applicable. The sole requisite modification pertains to the gradient inner product's definition. Rather than employing gradients from the cross-entropy logits, Vof ), we substitute them with gradients corresponding to the contrastive learning logits.
RELATED WORKS
Traditional data augmentation techniques such as random flipping and cropping (Krizhevsky al., 2017 Simard et al. 2003 Shorten & Khoshgoftaar 2019) are hand-crafted and static, unlike our adaptive SAFLEX method that tunes sample weights based on validation performance. Autonomous approaches like AutoAugment (Cubuk et al. 2019; Lim et al., 2019 Ho et al. 2019; Mounsaveng jet al., [02T 2023) learn transformations but are restricted in scope, primarily focusing on affine transformations. Generative methods employing GANs or diffusion models (Odena et al., 2017]
Sankaranarayanan et al.) 2018: Huang et al.] 2018: He et al., 2022, Shipard et al.J [202]] Dunlap ct al. 2023] Trabucco et al. 2023) can inadvertently alter class-relevant features, which our method avoids by adaptively adjusting sample weights. Research on adversarial perturbations (Goodfellow 2013: Yang al. 2022a by Ho & Nvasconcelos 2020) and noise-robust learning (Han al.;
let al.
2018 Lang et al. 20222 Thulasidasan et al. 2019 Konstantinov & Lampert 2019 Gao et al.) 2022 Ma et al. 2018 Kremer et al. (2018) address similar problems but often suffer from complexity and stability issues, which we mitigate by our principled approach to weight adjustment. Recently Soft-Augmentation (Liu et al. 2023) also proposes to use soft labels and sample weights to train on augmented samples. However, it implements a specific formula to generate them based on the strength parameter of upstream augmentations. This limits the applicability of Soft-Augmentattion mostly to crop augmentation on images. WWang et al. (2023) introduce self-adaptive augmentation within the meta-learning framework, MetaMix, which improves the corruption robustness of con- tinual learning models. BBattarai et al. (2020) propose a progressive sampling strategy for GAN synthetic data, while Caramalau et al. (2021) introduce a sequential graph convolutional network for active learning. Our work extends these findings by developing novel sampling and purifying method for augmented data that is specifically designed to improve the performance of downstream tasks.

---

Published as a conference paper at ICLR 2024
For more detailed discussion of related works, please refer to Appendix
5 EXPERIMENTS
We validate the effectiveness of SAFLEX under four very different learning scenarios: (1) adapt- ing augmentations to medical images, (2) refining augmentations for tabular data, (3) purifying cuftusioonmmodellgeneratte augments, and (4) applying to contrastive fine-tuning. Experimental se- tups and implementation details are provided in Appendix C
Adapting Augmentations to Medical Images. Unlike natural images, medical images often carry quantitative information (e.g., encoded color) and objects without a canonical orientation. While we usually lack the domain knowledge to design effective heuristic augmentation transformations for these images, applying augmentation pipelines designed for natural images, such as RandAug- ment (Cubuk et al.] 2020), can sometimes degrade performance the medical context (Yang et al.) 2022ab. Consequently, investigate whether SAFLEX can adapt these augmentation pipelines for medical images.
We assess multi-class classification across eight medical image datasets from MedMNIST (Yang) iet al., (2023), with each dataset comprising 10K to 236K 28 x8 images and 4 to 11 classes. In line with (Yang et al., 202T). we train a ResNet-18 model (He et al.) 2016) using the Adam opti- mizer (Kingma & Ba, 2014) for 100 epochs. For upstream augmentation, we utilize the widely- adopted RandAugment (Cubuk et al., 2020) and Mixup (Zhang al.) 2018) methods. Test accura- cies are presented in [ highlighting that SAFLEX significantly enhances the performance of both RandAugment and Mixup. It's noteworthy that SALEXX, when combined with basic upstream augmentations as shown in [ achieves better performance than Soft- noft-Augmentation (Liu et al.) 2023), and comparable or superior performance than the adversarial-pertuubauti -basdr augmen- tation, LP-A3 (Yang et al. 2022a). The latter not only takes significantly longer to train but also demands careful hyperparameter tuning. For a comprehensive view, Noft-Augmentation's 's perfor- mance and LP-A3's performance the MedMNIST datasets can be found in Appendix
Method
Tissue
Blood
OCT
OrganA
OrganC
OrganS
No Aug
91.31 = 0.18
76.14
0.00 68.25 +: 0.17 96.81 0.19 78.67 # 1,26 91.21 + 0.09 91.81 = 0.12 81,37 # 0.07
RandAug
93.52 + 0.08 73.71 + 0.33 62.03 0.14 95.00 0.21 76.00 + 0.94 94.18 + 0.20 91.38 + 0.14 60.52 0.32
SAFLEX
96.11 + 0.14 76.00 + 0.33 64.32 0:18 96.91
0.15 79.63 + 0.28 96.32 + 0:29 02.10 = 0.21 82.85 + 0.42
(w/ RandAug)
Mixup
+: 0.19 75.22 0.45 06.62 96.28
0.23 77.08
94.12
90.76
0.28 50.00
SAFLEX
93.71 + 0.37 76.94 # 0.51 68.31 = 0.43 97,21 0.35 79.54 11.04 93.06 + 0.31 92.73 + 0.53 62.14 11:27
(w/ Mixup)
Table 1: On medical images, SAFLEX significantly enhances the performance of RandAugment and Mixup across eight medical image datasets from MedMNIST.
In terms of efficiency, SAFLEX, designed update per iteration. It only extends the average wall-clock time of a training epoch by roughly 42% in this experiment; Appendix [or details.
augmentation plug-in, requires only single-step
Refining Augmentations for Tabular Data. Tabular data typically encompasses heterogeneous features that include a blend of continuous, categorical, and ordinal values. The presence of dis- crete features constrains the space of potential transformations. Furthermore, the domain knowl- edge design invariant, label-preserving transformations is often absent. One of the few traditional augmentation techniques directly applicable to tabular data is CutMix (Yun et al., 2019), which substitutes a portion of continuous or discrete features with values from other randomly chosen samples (see Appendix C for implementation details). However, studies suggest that CutMix, with a relatively small augmentation probability like 0.1, struggles to bolster tabular elassification perfor- mance (Onishi & Meguro 2023). Conversely, higher augmentation probability can introduce ex- cessive noise, potentially downgrading the performance. This leads us to explore whether SAFLEX can mitigate the noise from CutMix and enhance classification performance.
Our experiments span seven tabular datasets varying in size (from 452 to 494K) and feature types (from exclusively continuous features to predominantly discrete ones); detailed dataset information and statistics are available in Appendix [ Except for the Volkert dataset, which involves 10-way classification, all other datasets focus binary classification. Notably, some datasets, like Credit,
7

---

Published as a conference paper at ICLR 2024
exhibit a significantly skewed class distribution (e.g., only 0.17% positive). We consider backbone models such as the sample Multilayer Perceptron (MLP) with two hidden layers and 256 neurons each and tranformer-based models like SAINT (Somepali et al. (022) (without contrastive pre- training). These models undergo training with dropout (Srivastava et al. 2014) and, in certain cases, batch normalization, for 200 epochs.
Method
Model
Appetency
Arrhythmia
Click
Credit
QASR
Shrutime
Volkert
No Aug
MLP
49.03 = 0:01 81.53 0.03
62.54
0:04
66.91 + 0.03 91.84
0.02 86 + 0.04 61.14 0:05
CutMix
MLP
48.98
81.57 = 0.05
62
0.00
73.68 + 0.08 91.87 0.02
86.39
0.05 61:20 0:02
SAFLEX
(w/ CutMix)
MLP
51.04 0.00 83.02
+ 0.06
52.81
= 0.06 74.61 + 0.15 92.69
0.13
86.90
- 0.10 61.51
0.05
No Aug
SAINT 78.90 = 0:03
83,90
+ 0.01
66.72
0.00
79.49 + 0.05
98. 18
0.04
87.53 = 0.01 66.82 = 0.05
CutMix
SAINT 81.05 +: 0.07
85.32 + 0.09
65.77 0.04
78.71 + 0.08 98.61
0.06
87.61
4 0.07 68.23 0.10
SAFLEX
SAINT 81.33 + 0.14 85.27 + 0.14 66.12 : 0.09 79.93 + 0.17 98.38 0.21 87.93 + 0.13 68.91 1 0.17
(w/ CutMix)
Table 2: On tabular data, SAFLEX outperforms the upstream augmentation method, CutMix, across diverse tabular datasets using either MLP or SAINT the backbone models.
Table shows that SAFLEX almost consistently enhances the performance of CutMix across all datasets, regardless of whether the MI.P or SAINT model is used. This improvement especially noticeable with the MLP backbone, which is typically more intricate to train, and on datasets abun- dant in discrete features, such as Click and Shrutime, where CutMix tends to inject more noise. Notably, the Volkert dataset demonstrates considerable performance impovement, potentially at- tributed to the fact that has 10 classes where soft labels might be more useful.
Augments. Recent research (Dunlap et al. 2023 Trabucco ct al. 2023) has advocated the application of diffusion models for image editing via text prompts. Compared to traditional augmentation techniques, images produced by pretrained diffusion models maintain task-specific details while offering enhanced domain diversity, as dictated by the prompts. Diffusion-modee mooel-ggenerated augmentations have been found particularly efficacious in fine-grained classification and out-of outoof-distribution (OOD) generalization tasks (Dunlap et al., [023). However, these diffusion models occasionally generate subtle image alterations, potentially corrupting class- essential information, thus underscoring the necessity for noise reduction (Dunlap et al. (023)). In this context, we probe the capability of SAFEEX to enhance the purity of diffusion model -genrrate augmentations, aiming for improved classification outcomes.
Purifying Serfusioon-Moee-Geeneee
In our experimentation, we adhere to the setups in (Dunlap et al. 2023). We assess AAFFE using images derived from two distinct approaches: The lmg2lmg approach involves an image encoder that first converts a given image into a latent representation. Subse- quently, employing diffusion model (specifically, Stable Diffusion v1.5 (Rombach et al. 2022) for this experiment), this latent representation undergoes a series of prompt-conditioned transforma- tions. Ultimately, the altered representation is decoded, yielding augmented image reflecting the modifications stipulated in the prompt. Notably, the diffusion model may or may not undergo fine- tuning (w/ and w/o finetune) on the dataset in question. (2) The InstruciPix2Pix approach (Brooks) cet al. 2023) accepts an image and an edit instruction sampled (c.g., "position the animals within the forest") ") and outputs a correspondingly modified image. InstructPix2Pix is conditional diffusion model pretrained on a dataset containing paired images and their associated edit instructions.
Our evaluation encompasses two tasks: (1) Fine-grained classification on CUB dataset subset (Wah et al. 201T) (featuring 25 images per category). (2) OOD generalization on an iWildCam subset from the Wilds dataset (Koh et al. 202I) (consisting of over 6,000 images and simplified to 7- way classification). We use a ResNet-50 model (He et al.] 2016) pretrained on ImageNet (Deng lt al. 2009). For comparison, we also consider data generated solely from text (Tex/2Img) and the RandAugment method as baselines.
Results, as depicted in Table 3 affirm that SAFLEX consistently elevates the performance of all three certtuiion-mddl-ggeee and OOD generalization tasks. Notably, the performance boost is more prominent within the OOD generalization task, where feature and label distortions are particularly detrimental. We confirm that SAFLEX is useful to refine diffusion model- -enerreed augmentations, leading to enhanced classifi- cation accuracy.
augmentation techniques, across both fine-grained classification
8

---

Published as a conference paper at ICLR 2024
Aug
RandAug
Text2lmg
InstructPssPPis
Img2lin finetune)
Img2ling (w/ finetuue)
SARRIX
WSSARES
wis SAHUiX
wi SARLx
Fine-Gminnd
63.60 0.16 71,2 1 0.52 69.68 1.0.7 7L.38 1. 0.91
72.34 at
71.25 0.30
73.22 de 0.63 72.01 1.21
73.01 0.78
Clissification
OOD
a 272 04.53 67,29 199
72.61 +: 1.44
Generalization
Table 3: For diffussoon-model-geeeeaa augmentations, SAFLEX enhances the fine-grained clas- sification and OOD generalization performance for various diffusion diffusion-generaaion methods.
Applying to Contrastive Fine-Tuning, We next shift focus from the empirical risk minimiza- tion (ERM) framework utilizing cross-entropy loss, as demonstrated in the prior scenarios. To test the adaptability and compatibility of SAFLEX with contrastive loss, turn to contrastive fine- tuning paradigm termed "Finetune Like You Pretrain" (FLyP)(Goya et al.) 2023). This methodol- ogy offers straightforward yet potent means to fine-tune pretrained image-text models, including notable ones like CLIP (Radford et al. 2021). Remarkably, by simply fine-tuning classifiers through the initial pretraining contrastive loss (refer to Eq. (13), FLYP achieves uniformly better classifica- tion performance. This entails constructing prompts from class labels and subsequently minimizing the contrastive loss between these prompts and the image embeddings within the fine-tuning set.
Our experimentation adopts the framework presented in (Goyal et al. 2023 Specifically, we fine- tune CLIP model equipped with a V encoder on the full Wllld dataset from Wilds (Koh) et al. 2021). Post fine-tuning, we adopt a strategy from (Goyal et al.) (023) that linearly interpo- lates weights between the pretrained and the fine-tuned checkpoints to optimize in-distribution (ID) performance. As our upstream augmentation technique, we employ RandAugment, following hy- perparameter setups as described in (Koh et al.) 202I). For handling the CLIP contrastive loss, apply tailored algorithm, detailed in Section For an in-depth understanding, please refer to Appendix A and Appendix C
Zhro-Shot
FLYP
FLYP-RundAng FLYP+SAFLEX RundAug)
8.7
49.7
52.4
52.7
who Esscrabling
0OD
11.02:000
31.7
1:
35.6 1.2
36.3114 14
36.9 1.5
ID
8.7.100
50.2
0.5
52.5 0.5
52.6 4 1.0
53.0 1.0:::0 0.7
000
Emseubling
11.0:2000
80.7 04 37.1 111:1.
37.6 0.9
37.8 + 1.1
Table 4: Applied to contrastive fine-tuning of CLIP using FLYP (Goyal al., 2023), SAFLEX also enhances the performance of standard image augmentations like RandAugment.
As evidenced in Table incorporating RandAugment alongside FLYP yields favorable outcomes. Moreover, the introduction of SAFLEX amplifies performance gains for both ID and OOD tasks, irrespective of whether ensembling is applied. This observation is particularly noteworthy, as demonstrates that SAFLEX is compatible with contrastive loss, which is key component of many training paradigms, including self-supervised learning.
6 CONCLUSIONS
Our study presents SAFLEX, novel solution to current challenges in data augmentation. At its core, SAFLEX offers a paradigm shift from traditional, one-size-fits-all augmentation strategies to more adaptive, data-driven approach. It allows for the learning of low-dimensional sample weights and soft labels for each augmented instance, thereby cireumventing the complexities and limitations inherent in direct augmentation feature learning. Our method demonstrates universal compatibility, underscoring its vast potential for diverse data types in learning scenarios. Extensive empirical eval- uations confirm SAFLEX's 's prowess, proving its adaptability from medical imaging contexts to the nuances of tabular and natural image datasets. While A demonstrates promising results, there are certain factors to consider for optimal performance. A substantial and high-quality validation set is beneficial. A suboptimal set could limit its effectiivnesss. Additionally, the type of upstream augmentation methods selected plays a role, it impacts the overall performance of SAFLEX. Our approach also entails some computational overhead due to frequent gradient evaluations. These considerations will be the focus of future studies to further refine the methodology. In essence, SAFLEX stands as a testament to the advancements in learnable data augmentation, ushering a more adaptive and customized era of data-centric AI.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Ding, An, Xu, Satheesh, and Huang are supported by National Science Foundation NSF-IIS- 2147276 FAI, AFOSR-AV -Air Force Office of Scientific Research under award number FA9550-23-1-0000 DOD- DARPA-Defens, Advanced Research Projects Agency Guaranteeing Al Robustness against Decep- tion (GARD) HR001 2020007, Adobe, Capital One and Morgan faculty fellowships.
of Naval Research under award number N00014-22-1-2335, DOD-
REFERENCES
Sercan Ô Arik and Tomas Pfister. Tabnet: Attentive interpretable tabular learning. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 6679-6687, 2021.
Shekoofeh Azizi, Simon Korblith, Chitwan Saharia, Mohammad Norouzi, and David J Fleet. Synthetic data from diffusion models improves imagenet classification. arXiv preprint arXiv:2304.08466 2023.
Jonathan F Bard. Practical bilevel optimization: algorithms and applicattios, volume 30. Springer Science & Business Media, 2013.
Binod Bhattarai, Seungryul Back, Rumeysa Bodur, and Tae-Kyun Kim. Sampling strategies for gan synthetic data. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 2303-2307. IEEE, 2020.
Tim Brooks, Aleksander Holynski, and Alexei A Efros. Instructpix2pix: Learning to follow image editing instructions. In Proceedings the IEEF/CVE Conference on Computer Vision and Pattern Recognition, pp. 98392-18402, 2023.
Razvan Caramalau, Binod Bhattarai, and Tae-Kyun Kim. Sequential graph convolutional network for active learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 9583-9592, 2021.
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining, pp. 785-794, 2016.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 113-123, 2019.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings the IEEF/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A harge-scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. leee, 2009.
Lisa Dunlap, Clara Mohri, Devin Guillory, Han Zhang, Trevor Darrell, Joseph E Gonzalez, Aditi Raghunathan, and Anna Rohrbach. Using language extend to unscen domains. In The Eleventh International Conference on Learning Representations, 2022.
Lisa Dunlap, Alyssa Umino, Han Zhang, Jiezhi Yang, Joseph E Gonzalez, and Trevor Darrell. Diversify your vision datasets with automatic diffusion-based augmentation. arXiv preprint arXiv:2305.1622 16289, 2023.
Jiahui Gao, Renjie Pi, LIN Yong, Hang Xu, Jiacheng Ye, Zhiyong Wu, Weizhong Zhang, Xiaodan Liang, Zhenguo Li, and Lingpeng Kong. Self-guided noise-free data generation for efficient zero- shot learning. In The Eleventh International Conference on Learning Representations, 2022.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. In International Conference Learning Representations, 2015.
10

---

Published as a conference paper at ICLR 2024
Sachin Goyal, Ananya Kumar, Sankalp Garg, Zico Kolter, and Aditi Raghunathan. Finetune like you pretrain: Improved finetuning of zero-shot vision models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 19338-199999 19347, 2023.
Beliz Gunel, Jingfei Du, Alexis Conneau, and Veselin Stoyanov. Supervised contrastive learning for pre-trained language model fine-tuning. In International Conference Learning Represen- tations, 2020.
Bo Han, Quanming Yao, Xingrui Yu, Gang Niu, Miao Xu, Weihua Hu, Ivor Tsang, and Masashi Sugiyama. Co-teaching: Robust training of deep neural networks with extremely noisy labels. Advances in neural information processing systems, 31, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Ruifei He, Shuyang Sun, Xin Yu, Chuhui Xue, Wenqing Zhang, Philip Torr, Song Bai, and XIAO- JUAN QI. Is synthetic data from generative models ready for image recognition? In The Eleventh International Conference on Learning Representations, 2022.
Chih-Hui Ho and Nuno Nvasconcelos. Contrastive learning with adversarial examples. Advances in Neural Information Processing Systems, 33:17081-17093, 2020.
Daniel Ho, Eric Liang, Xi Chen, Ion Stoica, and Pieter Abbeel. Population based augmentation: Efficient learning of augmentation policy schedules. In International Conference on Machine Learning, pp. 2731-2744 PMLR, 2019.
Sheng-War Wei Huang, Che-Tsung Lin, Shu-Ping Chen, Yen-Y: Yi Wu, Po-Hao Hsu, and Shang-Hoo Hong Lai. Auggan: Cross domain adaptation with gan-based data augmentation. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 718-731, 2018.
Saachi Jain, Hannah Lawrence, Ankur Moitra, and Aleksander Madry. Distilling model failures as directions in latent space. In The Eleventh International Conference on Learning Representations, 2022.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in neural information processing systems, 33:1 116116661 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6980, 2014.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of tetthhewldd distribution shifts. In International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021.
Nikola Konstantinov and Christoph Lampert. Robust learning from untrusted sources. In Interna- tional conference on machine learning, pp. 3488-3490 PMLR, 2019.
Jan Kremer, Fei Sha, and Christian lgel. Robust active label correction. In International conference on artificial intelligence and statistics, pp. 308-316. PMLR, 2018.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Communications of the ACM, 60(6):84-90 90.84490.2000 2017.
Ananya Kumar, Aditi Raghunathan, Robbie Matthew Jones, Tengyu Ma, and Percy Liang. Fine- tuming can distort pretrained features and underperform outtof-distribution. In International Con- ference Learning Representations, 2021.
Hunter Lang, Aravindan Vijayaraghavan, and David Sontag. Training subset selection for weak supervision. In Advances in Neural Information Processing Systems, 2022.

---

Published as a conference paper at ICLR 2024
Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. Fast autoaugment. Advances in Neural Information Processing Systems, 32, 2019.
Chi-Heng Lin, Chiraag Kaushik, Eva L Dyer, and Vidya Muthukumar. The good, the bad and the ugly sides of data augmentation: An implicit spectral regularization perspective. arXiv preprint mrXiv:2210005000 2022.
Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and Qiang Liu, Bome! bilevel optimization made easy: A simple first-order approach. Advances in Neural Information Processing Systems, 35: 17248-17262, 2022.
Yang Liu, Shen Yan, Laura Leal-Taixé, James Hays, and Deva Ramanan. Soft augmentation for im- age classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16241-16250 2023.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit diffcrentiation. In International conference artificial intelligence and statisties, pp- 1540-1552. PMLR, 2020.
Xingjun Ma, Yisen Wang, Michael E Houle, Shuo Zhou, Sarah Erfani, Shutao Xia, Sudanthi Wijew- ickrema, and James Bailey. Dimensionality-drive learning with noisy labels. In International Conference Machine Learning, pp. 3355-3364. PMLR, 2018
Sayprascuth Mounsaveng, Issam Laradji, Ismail Ben Ayed, David Vazquez, and Marco Pedersoli. Learning data augmentation with online bilevel optimization for image classification. In Proceed- ings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 1691-1700, 2021.
Saypraseuth Mounsaveng, Issam Laradji, David Vázquez, Marco Perdersoli, and Ismail Ben Ayed. Automatic data augmentation learning using bilevel optimization for histopathological images. arXiv preprint arXiv:2307.1 1808, 2023.
Augustus Odena, Christopher Olah, and Jonathon Shlens. Conditional image synthesis with iliary classifier gans. In International conference on machine learning, pp. 2642-2651. PMLR, 2017.
Soma Onishi and Shoya Meguro. Rethinking data augmentation for tabular data in deep learning. arXiv preprint 2000::2335.00008 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8748-8763. PMLR, 2021.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Swami Sankaranarayanan, Yogesh Balaji, Carlos D Castillo, and Rama Chellappa. Generate to adapt: Aligning domains using generative adversarial networks. In Proceedings of the IEEE conference on computer vision and paltern recognition, pp. 8503-8512, 2018.
Jordan Shipard, Amold Wiliem, Kien Nguyen Thanh, Wei Xiang, and Clinton Fookes. Diversity is definitely needed: Improving model-agnostic zero-shot classification via stable diffusion. In Proceedings of the IEEE/VV Conference on Computer Vision and Pattern Recognition, pp. p 769- 778. 2023.
Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmentation for deep learning. Journal of big data, 6(1):1-48, 2019.
Patrice Y Simard, David Steinkraus, John C Platt, et al. Best practices for convolutional neural networks applied to visual document analysis. In Iedar, volume 3. Edinburgh, 2003.
12

---

Published as a conference paper at ICLR 2024
Gowthami Somepalli, Avi Schwarzschild, Micah Goldblum, C Bayan Bruss, and Tom Goldstein. Saint: Improved neural networks for tabular data via row attention and contrastive pre- pre-training. In NeurlPS 2022 First Table Representation Workshop. 2022.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, lya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15()):1999-198,, 2014.
Sunil Thulasidasan, Tanmoy Bhattacharya, Jeff Bilmes, Gopinath Chennupati, and Jamal Mohd- Yusof. Combating label noise in deep learning using abstention. In International Conference Machine Learning, pp. 6243. PMLR, 2019.
Brandon Trabucco, Kyle Doherty, Max Gurinas, and Ruslan Salakhutdinov. Effective data augmen- tation with diffusion models. In ICLR 2023 Workshop Mathematical and Empirical Under- standing of Foundation Models, 2023.
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-ucsd birds-200-2011 dataset. 2011.
Zhenyi Wang, Li Shen, Donglin Zhan, Qiuling Suo, Yanjun Zhu, Tiehang Duan, and Mingchen Gao. Metamix: Towards corruption-robus! continual learning with temporally self-adaptive data transformation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 24521-244 -22531, 2023.
Mitchell Wortsman, Gabriel IIharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rehecca Roelofs, Raphael Gontijo Lopes, Hannanch Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7959-7977 2022.
Sen Wu, Hongyang Zhang, Gregory Valiant, and Christopher Ré. On the generalization effects of linear transformations in data augmentation. In International Conference on Machine Learning, pp. 10410-10420. PMLR, 2020.
Jiancheng Yang, Rui Shi, and Binghing Ni. Medmnist classification decathlon: A lightweight au- toml benchmark for medical image analysis. In 2021 IEEE 18th International Symposium on Biomedical Imaging (ISBI), pp. 191-195, IEEE, 2021.
Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni. Medmnist v2-a large-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientific Data, 10(1):41, 2023.
Kaiwen Yang, Yanchao Sun, Jiahao Su, Fengxiang He, Xinmei Tian, Furong Huang, Tianyi Zhou, and Dacheng Tao. Adversarial auto-augment with label preservation: A representation learn- ing principle guided approach. Advances in Neural Information Processing Systems, 35:22035- 22048, 2022a.
Kaiwen Yang, Tianyi Zhou, Xinmei Tian, and Dacheng Tao. Seentity-disentanghee adversarial aug- mentation for self-supervised learning. In International Conference on Machine Learning, pp- 25364-25381. PMLR, 2022b.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy train strong classifiers with localizable features. In Proceed- ings of the IEEE/CVF international conference on computer vision, pp. 6023-6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Yifan Zhang, Bryan Hooi, Dapeng Hu, Jian Liang, and Jiashi Feng. Unleashing the power of con- trastive self- supervupe visual models via contrast-regularized fine-tuning. Advances in Neural Information Processing Systems, 34:29848 298600 2021.
Bo Zhao, Konda Reddy Mopuri, and Hakan Bilen. Dataset condensation with gradient matching. In International Conference Learning Representations, 2020.
13


Reviews:
[{"review_id": "NjAubkI7qx", "reviewer": "Reviewer_EBBs", "summary": "This paper presents a principled method for data augmentation. To this end, the paper presents a bilevel optimization framework for weighing and soft-labelling the augmented data in order to compensate for the adverse generalization effects of weak, strong and sometimes meaningless augmented examples. Although the impact of data augmentation for generalization, in particular deep learning frameworks, has been substantial, there is still a lack of principled ways of doing data augmentation. This paper has identified this gap and convincingly addressed the problem.", "questions": "I like the paper. Please see a few comments above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The paper is well-written and easy to understand. \n\nThe diagrams and the equations are easy to follow.\n\nThe experiments are performed on diverse datasets with various tasks, including medical imaging and tabular data.\n\nThe results are highly encouraging.", "weaknesses": "A few important previous works on sampling and purifying GAN synthetic data are relevant to this paper.  It is important to acknowledge and discuss their contributions in the paper. \n\nCaramalau, Razvan, Binod Bhattarai, and Tae-Kyun Kim. \"Sequential graph convolutional network for active learning.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.\nBhattarai, Binod, et al. \"Sampling strategies for gan synthetic data.\" ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rQpvcuWHjB", "reviewer": "Reviewer_GM6j", "summary": "The paper argues that data augmentations can suffer with two main issues - 1. The augmented samples can become out of distribution to the training distribution and 2) the augmented samples can belong to a different class than the original sample. To tackle the first issue, the authors propose to add sample weights (w_i) to the augmented samples. Samples which are farther from the training distribution can be assigned a smaller weight. To tackle the second issue, the authors propose to make the one-hot label as soft-label to capture the uncertainties. \nTo learn the sample weights and the soft-label the authors pose a bi-level optimization problem where in the inner loop, the model parameters are optimized over the training and augmented samples and in the outer loop the optimal augmentation parameters are optimized for. \n\nThe authors conduct experiments across three settings - 1. medical datasets, 2. tabular datasets and 3. for contrastive learning approaches. Across all the experiments the authors show improved performance on top of standard augmentations such as RandAug, Mixup and CutMix.", "questions": "I have already mentioned it in the weakness section", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The motivation in the paper about identifying the two issues with standard augmentation and then solving it by learning sample weights and soft-labels is really clear.", "weaknesses": "1. The main issue is a lack of proper baselines. Papers such as [1] have already explored using soft labels for augmentations where the softness is derived on the basis of augmentation strength. This paper's novelty thus gets limited. There is no comparison with [1] in any of the experiments. The authors should do a proper comparison with [1] and justify how their approach is better than it. \n\n2. To solidify the experimental results the authors should also experiment with stronger architectures and datasets such as ResNet-101 over ImageNet as done in [1].\n\nI am willing to update my ratings if my concerns are addressed. \n\nReferences - \n\n1. Soft Augmentation for Image Classification. Liu et al. https://arxiv.org/pdf/2211.04625.pdf", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0EtBJmJ3ed", "reviewer": "Reviewer_PN9Y", "summary": "This article contributes a workflow named as SAFLEX as data augmentation. Here is a summry:\n\n(1) Authors unveil a novel parametrization for learnable augmentation complemented by an adept bilevel\nalgorithm primed for online optimization.\n(2) Author's SAFLEX method is distinguished by its universal compatibility, allowing it to be effortlessly\nincorporated into a plethora of supervised learning processes and to collaborate seamlessly with an\nextensive array of upstream augmentation procedures.\n(3) The potency of authors' approach is corroborated by empirical tests on a diverse spectrum of datasets\nand tasks, all underscoring SAFLEX’s efficiency and versatility, boosting performance by1.2% on\naverage over all experiments.", "questions": "Is there any empirical experiments that SAFLEX can contribute to some other applicable downstream task like model training?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "They have considered experiments of different data types and model training as downstream tasks, which demonstrate their workflow as a robust one.", "weaknesses": "From a model perspective, this is a good one as topic of adaptive learning, though a  little bit off the topic of this conference. \nFrom data augmentation perspective, it is better to demo some more experiments in downstream task involves with high dimensional data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "obs1F2TGgJ", "reviewer": "Reviewer_NKYa", "summary": "This paper addresses the training of high-performance neural networks on small amounts of training data by optimizing data augmentation. Existing methods in this field optimize the transformation itself in the feature space, which limits the available data augmentation transformations and has high computational complexity. This paper differs from these existing approaches by optimizing the importance weights of the input features and the soft labels to be assigned to the augmented data, thereby meta-learning the data augmentation available in the existing data augmentation pipeline. Focusing on features and labels in data augmentation has not received much attention, and the idea is novel. Furthermore, the paper proposes a method to approximate the bi-level optimization of meta-learning with validation data to a single-level optimization by borrowing the idea of gradient matching and developing an efficient algorithm. Experiments verify the effectiveness of the proposed method on various datasets, tasks, and combinations of data augmentation methods. On the other hand, the paper does not provide an evaluation of the first-order approximation or the computation cost, and there is room for improvement in this aspect.", "questions": "- Do you think the proposed method can be applied to consistency-based semi-supervised learning with data augmentation, e.g., FixMatch [a]? The study of estimating the importance of samples and labels is well studied in the field of semi-supervised learning rather than in the field of data augmentation (e.g., FreeMatch [b]). If it can be shown that the scheme of the proposed method can be implemented in semi-supervised learning, the impact of this paper on the community will be even greater.\n\n[a] Sohn, Kihyuk, et al. \"Fixmatch: Simplifying semi-supervised learning with consistency and confidence.\" NeurIPS 2020.\n\n[b] Wang, Yidong, et al. \"Freematch: Self-adaptive thresholding for semi-supervised learning.\" ICLR 2023.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "+ The paper proposes a novel data augmentation optimization strategy that optimizes the importance and labels of input features. The idea is very interesting and original.\n+ The paper proposes a first-order approximation method to efficiently compute meta-learning that requires bi-level optimization.\n+ The paper applies and evaluates the proposed method not only on image datasets but also on table datasets. This evaluation is important in supporting the paper's claim that the method can be applied to any data augmentation pipeline.\n+ The paper provides experimental results on the recently widely used CLIP pre-trained models, effectively demonstrating the impact of the proposed method.", "weaknesses": "- Even though the paper proposes a first-order approximation method, it does not provide an evaluation of this method. In other words, the paper should provide a performance comparison with the usual bi-level optimization and a computation cost comparison with other data augmentation strategies such as Fast AutoAugment.\n- The writing of the paper is not necessarily of high quality. For example, Theorem 1 is very difficult to read because it contains multiple claims that make up the entire solution. Theorems and corollaries should be split for each claim, or if the propositions are ambiguous, they should be replaced with detailed explanations for each component, rather than in the form of a theorem. In fact, the proofs provided by the Appendix are almost obvious and make little theoretical contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)