Learning Tree-Structured Composition of Data Augmentation
Authors: Dongyue Li
Kailai Chen
Predrag Radivojac
Hongyang R. Zhang
Venue: Accepted by TMLR ()

Abstract: Data augmentation is widely used in scenarios where one needs to train a neural network given little labeled data. A common practice of augmentation training is applying a composition of multiple transformations sequentially to the data. Existing augmentation methods such as RandAugment rely on domain expertise to select a list of transformations, while other methods such as AutoAugment formulate an optimization problem over a search space of size $k^d$, which is the number of sequences of length $d$, given a list of $k$ transformation functions.

In this paper, we focus on designing efficient algorithms whose running time complexity is much faster than the worst-case complexity of $O(k^d)$, provably. We propose a new algorithm to search for a binary tree-structured composition of $k$ transformations, where each tree node corresponds to one transformation. The binary tree generalizes sequential augmentations, such as the one constructed by SimCLR. Using a top-down, recursive search procedure, our algorithm achieves a runtime complexity of $O(2^d k)$, which is much faster than $O(k^d)$ as $k$ increases above $2$. We apply the algorithm to tackle data distributions with heterogeneous subpopulations, by searching for one tree in each subpopulation, and then learn a weighted combination, leading to a forest of the trees.

We validate the proposed algorithms on numerous graph and image data sets, including a multi-label graph classification data set we collected. The data set exhibits significant variations in the sizes of graphs and their average degrees, making it ideal for studying data augmentation. We show that our approach can reduce the computation cost (measured by GPU hours) by 43% over existing augmentation search methods while improving performance by 4.3%. Extensive experiments on contrastive learning also validate the benefit of our approach. The tree structures can be used to interpret the relative importance of each transformation, such as identifying the important transformations on small vs. large graphs.

Published in Transactions on Machine Learning Research (06/2024)
Learning Tree-Structured Composition of Data Augmentation
Dongyuc Li
li. A.donggu.Shoorttensste...ttt
Northeastern University, Boston
Kailai Chen
chon.kalaathhortteenseemaaa edu
Northeastern University, Boston
Predrag Radivojac Northeastern University, Boston
predbug@hhotheesserr..d
Hongyang R. Zhang
ho. lo.shang@hhotthesster..cm
Northeastern University, Boston
Reviewed on OpenReview: http:///pprrniiieeeeee//o/oooona
Abstract
Data augmentation is widely used in scenarios where one needs to train neural network given little labeled data. A common practice of augmentation training is applying composition of multiple transformations sequentially to the data. Existing augmentation methods such as Rand Runmeng rely on domain expertise to select a list transformations ions, while other methods such as AutoAugment formulate an optimization problem over a search space size kd which is the number of sequences of length d, given a list of transformation functions.
In this paper, we focus on designing efficient algorithms whose running time complexity is much faster than the worst-case complexity of O(k*), provably. We propose a new algorithm to scarch for a binary tree-structured composition of le transformations, where cach tree node corresponds to one transformation. The binary tree generalizes sequential augmentations, such as the constructed by SimCLR. Using top-down, recursive search procedure, our algorithm achieves runtime complexity of O(24k), which is much faster than O(k") as k increases above 2. We apply the algorithm to tackle data distributions with heterogeneous subpopulations, by searching for one tree in each subpopulation, and then learn a weighted combination, leading to a forest of the trees.
We validate the proposed algorithms on numerous graph and image data sets, including a multi-label graph classification data set we collected. The data set exhibits significant variations in the sizes of graphs and their average degrees, making ideal for studying data augmentation. We show that our approach can reduce the computation cost (measured by GPU hours) by 43% over existing augmentation scarch methods while improving performance by 4.3%. Extensive experiments approach. The tree structures can be used to interpret the relative importance of each transformation, such as identifying the important transformations on small vs. large graphs.
contrastive learning also validate the benefit of our
1 Introduction
Data augmentation is technique to expand the training data set of a machine lcarning algorithm by transforming data samples with pre-defined transformation functions. Data augmentation is widely used in settings where only a small fraction of the unlabeled data is annotated, such as supervieee learning (Xie et al., 2020; Chen et al., 2020b). A common practice of augmentation training is apply sequence of transformations (Ratner al., 2017: Cubuk et al., 2019; Lim al., 2019; Zhang al., 2020). For instance, Auto (Cubuk et al., 2019) first applies random rotation to the images, then equalizes the histogram

---

Published in Transactions on Machine Learning Research (06/2024)
X
P1
A1
Permuting Edges
Generating
Subgraphs
P2
P2
A1, ... Ax
A2
A3
Input
P3
Ps P4
1
P4
Dropping Nodes
Permuting
Dropping
Transformations
Edges
Nodes
A4
As
A6
A7
Finding a tree-structured composition of transformations
2 Learning a forest of trees for stratified subpopulations
Figure 1: We illustrate the overall procedure of our algorithms. The input consists of transformation functions, denoted as AI,...,Ar Ak- Given data set, our algorithm constructs probablistie binary tree- structured composition of these transformations, as shown the left. Given input data æ, A is applied to map to (r)), with probability pı; otherwise, no transformation is applied, and remains unchanged. Let z' denote the output. In the next step, we will apply A₂ to x' with probability P2, or Ag to x' with probability 1-p2, etc. The second algorithm will first partition the entire data set into a few groups, e.g,, by the sizes of graphs as illustrated above. Then, the first algorithm is applied to learn one tree for each partition. These trees are weighted jointly to form a "forest" as the final augmentation scheme. A byproduct is that we can now measure the importance of each transformation in the tree of each group. For example, we find that permuting edges by randomly adding deleting fraction of edges works best for small graphs. For large graphs, generating subgraph by simulating a random walk works better. Notice that in an augmentation tree, if a branch only has a single child node, it means if the transformation is not applied, then we will not change the input or use any augmentation, which is the same as applying A(x) = x.
of each cell. The sequence transformations to apply depends on the downstream application; For instance, transformations that apply to colored images may not be suitable for black-and-white images. Thus, it would be desirable to have procedure that, given list of transformations and a target data set of interest, finds the most beneficial composition of transformations for that data set.
Various techniques have been introduced in prior work lo overcome the computational challenge of finding transformation compositions, which often involve solving highly complex optimization problem with a worst- case complexity of O(::)) (Ratner et al., 2017; Cubuk et al., 2019). Concretely, if there are & transformations, and the goal is to find the best sequence of length d, the search space would have Ad possible choices. The abovementioneed optimization methods suffer from a worst-case complexity of O(2a) because the search space includes all the kd choices. This paper aims to revisit this issue in range of settings of strong practical interest by designing search algorithms whose runtime complexity is much faster than this worst-case.
As an example, we start by considering data set of proteins with the goal of predicting their biological activity (function), a task relevant to biological discovery (Radivojac, 2022) and precision medicine (Rost et al., 2016). Each protein given as a graph, with nodes corresponding to amino acid residues and edges indicating spatial proximity between the residues in a protein's 's 3D structure. Protein function prediction can be seen here as a graph-level multi-label classification problem (Clark & Radivojac, 2011), where the presence of specific function indicates binary label and proteins can have multiple functions. Since the graphs can vary considerably in size and degree distribution, finding the right augmentation for this type of data set is particularly challenging. On related note, see also Zhang et al. (2022); Gao et al. (2023); Nguyen et al. (2023) for studies on image data sets with stratified subpopulations.
Existing approaches for finding compositionality either rely on domain expertise to identify the most relevant transformations, or tackle a ayper-parameter optimization problem over the scarch space (Wu et al., 2020). For instance, reinforcement learning- based methods are one type of approach (Cubuk et al., 2019; Luo et al., 2023), which define reward function given augmentation, and then search over the space of compositions
2

---

Published in Transactions on Machine Learning Research (06/2024)
to optimize the reward function. These methods require exploring search space of size k for finding sequence of length d. Therefore, in the worst case, the runtime complexity can still be O(I4).
The contribution of this paper is designing a faster algorithm to find a binary tree composition of k transformations of depth d, whose running time complexity is instead only O(2*k), provably. Our empirical results show that this complexity reduction does not. come at the cost of downstream performance. The algorithm conducts top-down, recursive search to construct a binary tree, where each node of the tree corresponds to one transformation. This tree structure leads to natural notion of importance score for each transformation, analogous to the feature importance score used in tree-based methods. Additionally, the algorithm can be used to tackle the abovementioned scenario where the underlying data involves a mixture of stratified subpopulations. See Figure 1 for an overall illustration of our approach.
We conduct extensive experiments on numerous data sets to validate our proposed algorithms. First, we apply our algorithm to newly collected graph classification data set generated using AlphaFold2 protein structure prediction APIs (Jumper al., 2021), containing 20,504 human proteins and 1,198 lypes of protein functions listed in Gene Ontology (Ashburner et al., 2000). Compared to existing data sets in Borgwardt et al. (2005) and Hu et al. (2020a), our data set is based & newer source and contains broader set of protein functions. We find that our algorithm can outperform RandAugment which does not search for composition by 4.3%. Compared to recent. augmentation optimization methods such as GraphAug (Tuo et al., 2023), our algorithm reduces its runtime by 43% and achieves 1.9% better performance. Second, we evaluate our algorithm in contrastive learning settings. Compared to the augmentation method of Chen et al. (2020b), the performance is on par with natural images such as CIFAR-10. On medical image classification data set, our algorithm can find better augmentation scheme that outperforms SimCLR by 5.9%. Finally, we justify the design of our algorithm through detailed ablation analysis. The details can be found in the experiments section.
To summarize, we list the main results of our paper as follows:
We design an algorithm whose worst-case rumning time complexity is O(2*k), where & is the number of input transformations and d is the length of the composition. This is significant improvement compared to the worst-case complexity of O(kd)) (e.g., when k > 10 and d > 3).
We use the algorithm to tackle data distributions that involve mixture of heterogencous subpopula- lions, by first partitioning the data set into several groups, e.g., according to the graph sizes. The algorithm is tested on newly collected graph classification dalaset and a few other benchmark data sets, reducing runtime while showing improved performance compared to several search methods.
We construct new multi-label graph classification data set, which contains a wide spectrum of graph sizes and degrees, making it ideal for studying data augmentation in future work.
1.1 Related Work
Motivated by the computational considerations around AutoAugment, active sampling has been used to accelerate the search procedure in data augmentation training (Wu et al., 2020). Besides, bilevel optimization (Benton et al., 2020), which applies a weighted training procedure on the augmentations, is another effective approach. Yang et al. (2022) introduce an adversarial training objective to find hard positive examples. Hounie et al. (2023) formulate data augmentation as an invariance-constraanne learning problem and leverage Monte Carlo Markov Chain (MCMC) sampling to solve it. Tt worth noting that the tree structures we are proposing in this work can be adopted as the base compositionality in all of these optimization frameworks. Further exploring the utility of tree structures would be an interesting direction for future work.
Apart from supervised learning, data augmentation plays crucial role in contrastive learning (You al., 2020). The invariance introduced by data transformation methods can serve as regularizer in self and semi-supervised learning techniques (Xie el al., 2020). Related to graph structures, there is a line of work on model robustness to graph size shifts by comparing local graph structures (Yehudai et al., 2021), or by applying regularization on representation distances between different sizes (Buffelli al., 2022) (sce the discussion in Ju et al. (2023) for further references). In particular, Li et al. (2023) and Nippani al. (2023) design multitask learning algorithms to improve the robustness of learning under distributional shifts data
3

---

Published in Transactions on Machine Learning Research (06/2024)
imbalance. Our work contributes to this line of literature by developing simple, efficient data augmentation search procedure. Besides, there has been developments on novel use cases of higher-order graph structures in graph neural networks (Chien et al., 2022), and spectral clustering within a Superimposed Stochastic Block Model (Paul et al., 2023). Specific to graph contrastive learning, little is known regarding the use of higher-order graph structures. Further fleshing this out may be a promising avenue for future work.
Even though data augmentation has generally been quite useful in practice, the theoretical understanding of data angmentation training is relatively scarce. Part of the challenge is that the data augmentation training paradigm violates the independent sampling assumption typically required by supervised learning. For instance, suppose one would like to understand how sequence of transformations like rotation, cropping, etc., to an image, affects the rownstream performance. This would require modeling such transformations within the data augmentation paradigm. There have been few developments in this direction. Dao et al. (2019) model the behavior of stochastic augmentation in a Markov chain, and reason about its behavior more preciscly using kernel methods. Chen et al. (2020a) develop a group-thcoretin framework modeling data augmentation. Their framework applies to the family of label-preserving transformations. Wu et al. (2020) categorize the generalization effects of linear transformations using bias-variance decomposition and study their effects on the ridge estimator in an over- parraaaaatter linear regression setting. Shao et al. (2022)
develop a PAC-learning framework for transformation -variaat hypothesis spaces.
Specific to the approach taken by this paper, we note that greedy recursive partitioning is a known method for constructing decision trees. There are studies, albeit in a very different setting, on the approximation ratio of the greedy search for constructing trees (e.g., Adler & Heeringa (2008); Gupta et al. (2017)). It an interesting question to revisit. these results in the setting of designing tree data augmentation.
Organization. We start with our problem setup in Section 2. Then, we describe the proposed algorithms in Section 3. We will also give some examples of the algorithms to illustrate their design. Next, we present the empirical evaluations in Section 4, with a detailed comparison to existing approaches. The Appendix includes theoretical derivations and analysis omitted from the main text.
2 Learning Composition of Data Augmentation
This section describes the problem setup for learning composition of transformations. Consider prediction problem where the goal is to map an input feature vector E X to label E y. We have access to a data set P, which includes a list of examples (x, y) drawn independently from an unknown distribution P over X x V. Given & transformation functions, denoted as A1...., AA, each transforms an input feature vector from X' - X. The problem is to find a composition Q of a subset of the functions, such that a model Se is learned to minimize the loss, denoted as f, of the augmented examples. For example, the case that is the cross-entropy loss over C classes, it maps from RC y + R. We will design Q with probabilistic distribution over sequences of compositions. Concretely, let be sequence sampled from Q. The learning objective can be written as:
L(Se:Q) E E [((((T())),)) (e.))PP~6
When the underlying distribution involves several groups of stratified subpopulations, we can write the above objective as a mixture of losses over the groups. Let m denote the mumber of groups and let g = {1..... m} denote the set of group labels. Let Pg denote the data distribution of P restricted to group g. For our context, it is sufficient to think that the group labels are available during training and testing time. For this setting, the learning objective becomes
m
() 9g Lc(oo: where Lz(fo:Q) = E E [((fo((()) (x.y)~,, TQ
(1)
9=1
and 4g E 10 indicates the proportion of examples coming from group g. Besides this empirical risk minimization sctup, an alternative min-max optimization is studiod in robust optimization (Zhang al., 2022; Nguyen al., 2023). We note that the ideas to be presented next can also be used in this min-max formulation.
4

---

Published in Transactions on Machine Learning Research (06/2024)
One type of compositionality is using a probabilistic sequence. We illustrate such an example in Figure 2, used in SimCLR (Chen et al., 2020b). This example applies "random cropping" with probability p. With probability 1-2 p, the input not changed, and is pushed to "color distortion,"" " and finally "Gaussian blurring." The probability value of each transformation indicates the likelihood of applying it that step. More generally, consider a sequence of length d, denoted as (A1, P1). (A2, .2)..... (As, (A), where each A, is associated with Pi E [0, 1]. The transformations are applied sequentially from the first node to the last node, each in a probabilistic way. For example, if A, is not applied, then the input to A; remains unchanged and is passed to the next node, similar to a skip connection, and so on.
p +
Random
Cropping
Color
Distortion
p
Gaussian
Blur
Figure 2: Illustrating a sequential aug- mentation scheme (Chen et al., 2020b).
In this paper, we will consider the family binary treo-structured compositions, which naturally generalizes the sequence illustrated in Figure? 2. We define the composit as binary tree. With depth d, there will be 2d nodes in total. Each node is associated with transformation A and the probability value Pi indexed from = 1,...,2" 2d 1. Given such a tree, we apply the transformations from the root node to one of the leaf nodes. For instance, after applying A,, the augmentation proceeds to either the left or right node. A special case is when the node itself the identity mapping. In such cases, the transformation will terminate at that point.
We now elaborate on the computational cost of existing methods. AutoAugment (Cubuk al., 2019) uses reinforcement learning the search method to find composition that optimizes the validation performance. This method takes many GPU hours to search for a composition with depth d out of k 16 augmentations on CIFAR-10. Fast AutoAngment (Lim et al., 2019) uses Bayesian optimization as the search method. Both of these two methods use worst-case search space of (la*). Instead of searching in a discrete space, Chatzipantazis et al. (2023) parameterizes the composition space with O continuons variables, and then uses SGD to optimize the model weights and these variables jointly.
3 Our Proposed Algorithms
We now present our algorithm for learning tree compositions. Our algorithm conducts a greedy search from the root to one of the leaves. We will provide illustrative examples to validate this algorithm. Then, we extend the algorithm to a setting involving multiple subpopulations. We will design a weighted, probabilistic combination of the trees in forest.
3.1 Learning Tree-Structured Composition
We consider lop-down binary search procedure inspired by the recursive binary splilling of building a decision tree. The procedure searches from the root of tree and iteratively finds one augmentation in one of the empty leaf nodes. A node would indicate a new branch of two nodes to build the tree further.
For a particular node the tree, we search for the choice of A, and p; that leads to the best improvement the validation performance, such as cross-cntropy loss. Each search itcrates over all possible transformations AI,..., A& and a discrete list of probabilities. We include A(x) x, namely identity mapping, to one of the l transformations. Naively optimizing the validation performance requires training O(k) models for each choice and choosing one via cross-validation This can be slow when & is large.
We enhance the efficiency of each search step by employing a density matching technique, without repeatedly training multiple models. This method has been used in prior work (Lim et al., 2019). Specifically, we train one model with the current tree, denoted as før Then, for each choice of A and p, we incorporate them position of the current composition and evaluate the performance of fa- on augmented examples generated by applying the new composition to the validation set. Denote tree composition T. We measure its
5

---

Published in Transactions on Machine Learning Research (06/2024)
performance validation set of size n as:
a Lval (Ta Σ TWr
E
(T(a))
(2)
Thus, each search only trains one model, with O(k) evaluations of adding each transformation to position of T. Empirically, we find that density matching identifies the same tree as fully-trained models on a protein graph data set and another image classification dataset. We evaluate the relative residual of squares error between (T') and that of training model for each transformation. We find that the gap is S 0.7%.
This concludes the search procedure for single tree node. We will repeat the same procedure for all the other nodes of the tree:
When child node is added with transformation A, and probability Ps, the other child node of the same parent node should find a transformation from one of the k input ones, and apply the chosen transformation with probability 1 - p;. In other words, the probability value - pt will be fixed.
The process continucs until the tree reaches a pro-specified depth d, or until L longer improves after adding any transformation.
We summarize the complete procedure in Algorithm 1.
Algorithm A top-down recursive search procedure for finding binary tree-structured data augmentation scheme
Input: & transformation functions A1, A2, .2...,,..44 As (including identity mapping), training and validation Require: Maximum tree depth d; A list H of probability values
Output: A probabilistic, binary treo-structured composition T of A2, ................................ As}
1: Initialize T = f}, V = ={1}, L 0c
2: while depth(?') v d and V is not empty do Randomly choose i V, let A,, pr, (.) Bulld-one-node(i, T), then remove i from V Add A,, P; to index of T /* If A, is the identity mapping, it means improvement is found */
5: if (1) < Lval then
Update Lval (00 and V - V2,,
end if
8: end while return T
10: procedure Build-one-nodeli. T)
11: Train a model with T' denoted as 0"
12: if i's has a sibling node that is in T then
13: Let p -1- Pkr where the index of the sibling node of i
14: for j E [I.....k] do
15: Add p) to at position i of T and evaluate Lval (T')
16:
end for
17:
else
18:
for j E fl...., ,}} and E H do
19: Add (Aj,P) to position i of and evaluate Leva (T)
20: end for
21: end if
22: return (Aj.)) achieving the lowest Lvat
23:
end procedure
Importance scores.
The tree structures allow us interpret the relative importance of each transformation, analogous to feature importance scores used in tree-based methods. Concretely, measure the usefulness

---

Published in Transactions on Machine Learning Research (06/2024)
of each transformation by the amount by which it reduces Lyal. added across all of its appearances in the tree. Thus, a higher score indicates that a transformation contributes more to reducing the loss. We will refer to this as the importance score of transformation and measure it in the experiments later on.
3.1.1 Running Time Analysis
We examine the rumning time of building single tree. At depth d, the number of search steps is at most 24 - 1, the maximum size of the tree. Each step involves training one model and iterating over the list of k transformations and probability values. Thus, the overall procedure takes ((24k) time to complete the search. By contrast, the worst-case complexity of existing methods scales as O(k") which is the space of all possible sequences of length k. We present the comparison Table
Table 1: Summary of comparisons between our approach and previous optimization algorithms to search for compositions of data augmentations. We use d to denote a composition's 's length and & denote the number of transformations to the input..
Method
Running Time
Compositionality
AutoAugment (Cubuk et al., 2019)
Reinforcement Learning
O(k")
Sequences
Fast AutoAugment (Lim al., 2019)
Bayesian Optimization
O(k")
Sequences
SCALE (Chatzipantavis et al., 2023)
Stochastic Gradient Descent
((k4)
Sequences
Algorithm (Ours)
Binary Search
0(24k
Forest of Trees
3.1.2 Illustrative Examples
Next, we give an example of running our algorithm before getting further into the technical details. We will illustrate its runtime as well as previous methods's 's runtime. We train WideResNNet-28-10 on CTFAR-10 and CIFAR-100 in supervised learning, following the setting of AutoAugment. We use 16 transformations, cach with five values of perturbation scales (Cubuk et al., 2019), including ShcarX/Y, TranslatcX/Y, Rotate, AutoContrast, Invert, Equalize, Solarize, Posterize, Contrast, Color, Brightness, Sharpness, Cutout, and 1.0. .... We use reduced set of randomly sampled 4, 000 examples as the training set during the search. After the search, we train a model on the full data set and report the test performance. For each algorithm, we report the runtime using an Nvidia RTX 6000 GPU.
Sample Pairing. As for H, we set it as 0, 0,0.1,0.2..... 0.2,...,
The results are shown in Table Our algorithm uses 4.7 GPU hours to complete, reducing the runtime by 51% compared to the existing methods. Meanwhile, the test performance remains close the methods.
Table 2:
We compare our algorithm with several existing approaches, trained from randomly initialized Wede 10 on CIFAR-10 and CIFAR-100. We report the number of GPU hours of the search single GPU, We also report the error rate the test set, averaged over five random seeds.
CIFAR-10
CIFAR-100
GPU hours
Error rates (%)
GPU hours
Error rates (%)
AutoAugment
25000
2.6+1.11
25000
17.1+0.33
Fast AutoAugment
19.2
2.7+0.1
19.6
17.3+0.2
SCALE Algorithm (Ours)
9.6
3.3+0.1
9.6
17.3+0.1
4.7
3.110.1
4.7
17.310..1
Next, we show example of using augmentation composition for training ResNet-50 on CIFAR-10 in contrastive learning (Chen et al., 2020b). We use seven transformations, including Random Cropping, Cutout, Rotation, Color Distortion, Sobel filtering, Gaussian noise, Gaussian blur, and the same H.
Figure 3 shows the tree found by our algorithm. Interestingly, this tree shares three nodes, as in SimCLR, which carefully designed by domain experts. The tree begins by applying random cropping and color distortion, similar to SimCLR augmentation. Then, the tree uses Gaussian blur one branch but adds rotation another branch. As shown in Table 2, this is comparable to the tree illustrated in Figure 2.
7

---

Published in Transactions on Machine Learning Research (06/2024)
p == 1.0
Random
Cropping
p = 0.6
p 0.4
Color
Identity
Distortion
Mapping
p p= 0.4
p 0.6
Gaussian Blur
Rotation
Figure 3: Illustrating the binary tree returned by Algorithm 1, conducted on CIFAR-10. On the right branch, no further transformation is applied after the identity mapping.
Lastly, we also find that the top-down search performs comparably to the exhaustive search of all possible trees of depth two. On several image and graph classification data sets, the greedy search is only 0.4% below the test result of the exhaustive search on average.
3.2 Learning Forest of Trees
With the efficient search procodure, we proceed to study the case of learning from multiple subpopulations. We start by presenting a motivating example, showing that learning a single tree does not suffice. We conduct experiments on a protein graph classification data set. The problem involves multiple binary-labeled classifications, aiming to classify proteins into their protein functions, each treated as a zero-one label. Since labeling the protein functions requires domain expertise, we could only obtain about - labels of all the unlabeled data. We will design data augmentation schemes to generate synthetic labels based on the labeled examples. The dataset exhibits significant variations in graph sizes and degrees. The sizes of the graphs range from 16 to 34,350 nodes, and the average degrees range from 1.8 to 9.4. In the experiment, data set into 16 groups, which corresponds to 16 intervals of sizes and average degrees. We will use graph neural networks as the base learner (Hu et al., 2020b). We consider four augmentation methods, including DropNodes, Permutelidges, Subgraph, and MaskingNodes, as in You et al. (2020). We use the same list H.
will divide the
We first find a tree by minimizing the average loss. We also apply this to each group and compare it with a model trained without augmentation. For graphs with less than 200 nodes or an average degree of less than 0.1, just minimizing the average loss can lead to worse result than stand-alone training.
Next, we search for one tree in each group, across all groups. For each group, we run Algorithm 1 to construct a tree. Notice that the data we use for executing this step is separate across groups. Hence, there is no synchronization performed on different groups. In other words, the models trained in cach group will not be reused training other groups (or in the weighted training slep later). This leads to some heterogeneily in the tree structure we obtain for different groups. To give an example, in Figure 4, we plot two trees, corresponding to group of the smallest sizes and another group of the largest sizes. The latter uses larger scales and one more step. We find that "Permuting Edges" and "Dropping Nodes" are used on the small-sized graphs, which remove randomly sampled edges and nodes, respectively. Meanwhile, all the transformations are used on the large graphs, including "Generating Subgraph," which extracts a subgraph by random walks, and "Masking Nodes," " which randomly masks a fraction of node attributes to zero We also report the importance score of each transformation. Figure 4a, "Permuting Edges" " has the highest. importance score. Whereas in Figure 4b, "Generating subgraph" has the highest importance score.
We give another example from an image classification data set. This dataset contains three groups of images with different colors and backgrounds. We illustrate the trees found for one group of colored images and another group of black-and-whit images. We note distinct differences between the two trecs. Transformations that alter the color distributions, such as Auto Costrast, Solarize, and Color Enhancing, are applied to the colored images. In contrast, transformations applied to Black-and-white images modify the grayscale or shape, including Equalize, TranslatcY, and Posterizo. Additionally, for the colored images, the Auto Contrast. transformation has the highest importance score, while for the black-and-white images, Equalize has the highest importance score. These observations validate that augmentation schemes vary across groups.
8

---

Published in Transactions on Machine Learning Research (06/2024)
= 01.6
0.4
p
Generating Subgruph
Identity Mupping
Score: 0.26
Score:
Score: Permutting Edges
Identity Mapping
p == 0.6
0.4
Permutting Edges
Dropping Nodes
Score: 0.04
Soore: 0.03
0.6
0.4
p 0.6
0.4 p= 0.2
0.8
Dropping Nodes
Identity Mapping
Score:
0.02
Score: 0
Dropping Nodes
Identity Mupping
Permuting Edges
Masking
Score: 0.01
Nodes
(a) Augmentation tree found on small graphı with the number nodes less than 200.
(b) Augmentation tree found larger graphs with the number of nodes larger than 600.
Figure The augmentation trees found from different groups can vary dramatically. classification data set, the augmentation the tree found on large graphs (right). We also report each augmentation's importance score, computed from the validation set. To clarily, after the identity map, i.e., A(x) x, no further transformation will be applied.
protein graph on small graphs (left) involves fewer augmentation steps than
p p= 0.8
= 0.2
= 0.6
p 0.4
Auto Contrast
Identity Mapping
Identity Mapping
Score: 0.21
Score:
Score: 0.64
Equalize
Score:
p 0.4
p 0.6
p =
= 0.2
Color Enhancing
Score: 0.11
Solarize
Score: 0.09
Score: 0.25
Translate
Posterize
Score: 0.09
(a) Augmentation tree found on colored images.
(b) Augmentation tree found on black-white images.
Figure 5: Illtsratting the trees found between colored images and black-white images on an image classífication data set. The tree on the left involves different transformations compared to the right.
3.2.1 Learning the Weight of Each Tree
With one tree learned for each group, we next unily the trees through weighted training of a new model fo and the weight of each group. For each group g, denote the augmentation tree as Qg. for = m. Consider a bilevel optimization problem:
min L(fe such that 0" e arg min Σ w (foi w.0w
(3)
where the loss of group i is associated with weight E R(0, 1), and Em We = 1. We write objective (3)
as function of f":
min ((fo)) = Σ 9o wE2m
Ly(so=: Q,).
(4)
9=1
We derive the gradient of objective (4) with respect to w. This is achieved by applying the chain rule and computing the gradient of 4" for any i 1,..., m.
ai((f=)) = Σ 1g Voigg(oon:
d; ==
H Ho where H = wy Vi (fo=; L
(5)
Ow,
We defer the complete derivation to Appendix A. d, can be viewed as similarity value of group i compared to the averaged group, normalized by Hessian inverse.

---

Published in Transactions on Machine Learning Research (06/2024)
We will use alternating minimization to update both 0 and w. Let the model parameters and weights at iteration be g and w(9)
With fixed w(,, update the model parameters from ( to g((+1) by running a SGD steps:
min g=1
(1) Lg(fo: Qg).
9
With fixed g(++1), obtain the gradient af from equation (5). Then, update un as:
(t)
(- -nnd("))
w exp
(1-1)
(6)
Σド=1 m
wi exp (-ndf"")
Taken together, we summarize the complete procodure for learning unified augmentation scheme in Algorithm
2. To recap, the final augmentation scheme is a probabilistic mixture of trees, one for each group.
Algorithm Learning a forest of trees from a mixture of subpopulations
Input: & transformations A2, As; Training/Validation Vaalidaiin splits of m groups
Require: Number of iterations S, SGD steps a, learning rate
Output: m trees with a probability value for each tree: (Q1, w1), (Q2, wy), (Qm: Wn); Model weight 0*
1: for g= m do
2: Compute an augmentation scheme Qg for group g using Algorithm
end for
Initialize model parameters (10); Set weight variables w(t) the uniform proportions [1/m, ... 1/m]
for t = o...., do
6: Update g(t) with a SGD steps to get ((++1)
7: Update ((+1)) from (()) according to equations (5) and (6)
8: end for
9: return (Q1,w) (5)), (Q2, w/ ) (Qm, w g(S)
In Appendix 2, we provide generalization bound to justify the consistency of this procedure. We demonstrate that the generalization error the model trained by Amgorithm 2 is bounded by a bias increase term from transferring from across different groups, and a variance reduction term from data augmentation. The proof technique involves carefully analyzing the bilevel optimization algorithm using covering numbers, which builds on the transfer exponent framework of Chen et al. (2022) and the work of Hanneke Kpotufe (2019).
3.2.2 Running Time Analysis
Next, we examine the runtime of the weighted training step. Compared to SGD, the algorithm includes updating the group weights using equations (5) and (6). This step involves computing Hessian inverse and the product between Hessian inverse with the gradient of each group. To avoid explicitly computing Hessian inverse, we the following method:
First, estimate the inverse Hessian-graddent vector s = H H- We apply conjugate gradient with stochastic estimation. We samples n data points (1),,,))) recursively computing
l, where v = En 9 VoLy
- (f
and takes A as the final estimate of H 11.. For data points and 0 E R2, this procedure takes (mpp) time.
Secondly, compute d, -
- Vota(So== Qt). Since there groups, this step takes 0(mp) time.
10

---

Published in Transactions on Machine Learning Research (06/2024)
Taken together, updating the group weights takes O(np + mp) time. In our implementation, conduct the above estimation using one batch of b data points from each group, leading to running time of O(bmp + mp) = Obmp)).
Empirically, we find that the running time is comparable to SGD. To illustrate, we train a three-layer graph neural network for protein graph classification. Rumning the above update takes 1.8 GPU hours, which is 1.1x of SGD, which takes 1.6 GPU hours. On another image classification data set for training ResNet-50, the runtime is 0.8 GPU hours, 1.3x of SGD, which takes 0.6 hours.
Lastly, we illustrate the convergence behavior in Figure 6. The loss curves on both data sets behave similarly to SGD, with weighted training achieving a lower test loss.
1.5
Tist loss SGD
Test luss: NGD
Alg.
Tes lass: A.g.
1.0
Training lest SGD
Training loss SGD
Loss
Loss
Triniing loer: Alg.
Training Alg.
0.5
19
0.0
0
25
50
75
100
10
20
Number of epochs
Number of epochs
(a) Training a three-layer GNN for graph classification
(b) Training a ResNet-50 for image classification
Figure 6: Showing the loss curves of both weighted training and SGD. The left figure shows the results of training a three-layer neural network on protein graph classification. The right figures show the result of training ResNet-50 on image classification.
4 Experiments
In this section, we evaluate the performance of Algorithm for numerous data sets in terms of both the runtime and the performance of the augmentation. We will consider both supervised and self-supervised lcarning settings. In particular, we will compare method with existing optimization methods for augmentation training. We will also compare with schemes designed based on domain knowledge, such as RandAugment (Cubuk et al., 2020) and SimCLR (Chen et al., 2020b). The results are broadly relevant to learning with little labeled data, with a particular focus on graphs and images. We summarize a few key findings as follows:
In supervised learning settings with both graph data sets (such protein graphs) and image data sets (such as medical images), our approach can outperform augmentation search methods the graph classification data set by 1.9% with 43% less runtime and on the image classification data set by 3.0% with 38% less runtime. Compared to RandAugment which does not search for a composition, we note 4.3% and 5.7% better performance, respectively.
For contrastive learning settings, we can now get up to 7.4% improvement across eight graph and image data sets. The runtime is also roduced by 32%. Compared to SimCLR (sce Figure 3), the augmentation scheme can deliver comparable performance on CIFAR-10. On a medical image classification task, we now find new scheme that outperforms SimCLR by 5.9%.
We verify that the improvement is statistically significant, by conducting the Wilcoxon signed-rank test on the performance of our algorithm compared to the basclines across ten datasets. When comparing our hm with the recent augmentation search method and with Rand Augment, the test yields p-values of 0.0025 and 0.0019, respectively. This indicates that the probability of our algorithu showing no improvement is less than 1%.
Lastly, we provide ablation studies to justify the design, including finding the tree for each group and the weighted training, for achieving the final performance.
11

---

Published in Transactions on Machine Learning Research (06/2024)
Our code for reproducing the experiments is available at augmentation, which also includes instructions for loading the new dataset.
4.1 Experimental Setup
We apply our algorithm to two settings. The first setting is supervised learning, which uses data augmentation to generate labeled examples. The second setting is contrastive learning, which uses data augmentation to generate contrastive examples.
Data sets. For supervised learning, we construct a graph classification data set for protein function prediction collected from the AlphaFold Protein Structure Database. This data set includes 20,504 human protcins structured undirected graphs, where nodes correspond to amino acid residues labeled with one of 20 amino acid types. The edges the spatial distances between two amino acids thresholded 6λ between two Co atoms. The problem is multi-label classification containing 1,198 protein functions, each function viewed as a binary label. We split the data set into 16 groups using intervals of graph sizes and average degrees, The idea is to group graphs of similar sizes together. To determine the opt number of groups, we experiment with dividing the graphs into 4,9, 9, 16, and 25 groups, corresponding to splitting graph sizes and average degroes into 3,3,4 4, and 5 intervals, respectively. Increasing m beyond 16 did not lead to any benefit. Therefore, we restrict m to less than 16.
The main difference between our data set and previous data sets (Borgwardt et al., 2005; Hu et al., 2020a) is the coverage of protein functions. The data set from Borgwardt et al. (2005) predicts whether a protein is an enzyme as binary classification task. The data set from Hu et al. (2020a) classifies a protein into 37 taxonomic groups as different species as a multiclass classífication problem. Our data set involves 1,198 protein functions using Gene Ontology (GO) annotation. Another difference is constructing the graphs. Borgwardt al. (2005) and our data set abstract the structure of the protein itself as a graph. Hu al. (2020a) constructs the graph based on relationship between proteins in protein association networks. The third difference is that the graphs in our data set exhibit more significant variations in their sizes, ranging from 16 to 34, 350. This comparison is summarized in Table 3.
Table 3: Comparison of our protein graph classification data set to two existing data sets. Our data set is constructed from newer data source and covers broader range of (protein) functions than previous ones.
PROTEINS
OGBG-PPA
Our Data Set
# Graphs
1113
158, 100
20, 501
# Functions
2
37
1, 198
# Nodes
4 620
~ 300
16 34,
Avg. Degree
3.4 10.1
2.0 240.9
1.8 9.4
Category
Binary classification: whether or not the protein is an enzyme
Multi-class classification for 37 species, e.g., mammals, bacte- rial families, archaeans
Multi-label classification for
1,198 GO terms that describe biological functions of proteins
Graph Type
Attributed and undirected graphs, where nodes represent secondary structure elements (SSE), and edges represent spa- tial closeness
Undirected graphs of protcin associations, where nodes rep- resent proteins and edges in- dicate biologically meaningful associations between proteins.
Undirected graphs, where nodes represent the amino acids and edges represent their spatial distance thresholded at 6Å between two Ca atoms
Data Source
Protein data bank
Protein association network
AlphaFold APIs
Next, we consider an image classification task using the iWildCam data set from the WILDS benchmark (Beery et al., 2021), where images are collected from different camera traps with varied illumination, camera angles, and backgrounds. The problem is multi-class classification: given a camera trap image, predict the labels as one of 182 animal species. We evaluate the macro-Fi score on this data set, which is the averaged Fi score over every class. We consider three groups from three cameras with the most images.
hhtpps//alphaaidddeeaaaa..
12

---

Published in Transactions on Machine Learning Research (06/2024)
For contrastive learning, we consider image classification, including CIFAR-10 and a medical image data set containing eye fundus images for diabetic retinopathy classifications. This is a multi-class classification problem: given an eye fundus image, predict the severity of diabetic retinopathy into five levels. We consider three groups based on three hospitals, which differ in the patient population and imaging devices. The sources are available online: Messidor, 2 APTOS, 3 and Jinchi
We also consider six graph classification data sets from TUdatasets (Morris et al., 2020), including NCI1, Proteins, DD, COLLAB, REDDIT, and IMDB, for testing graph contrastive learning. The first three data sets contain small molecular graphs where nodes represent atoms and edges represent chemical bonds. The goal is to predict the biological properties of molecular graphs. The latter three data sets contain social networks abstracted from scientific collaborations, online blogs, and actor collaborations. The goal is to predict the topic of the network. We split each data set into four groups by their sizes and average degrees.
Baselines. We compare our approach with pre-specified data augmentation methods, including ones designed for specific group shifts, and optimization methods that search for compositions. For supervised learning, we consider RandAugment (Cubuk al., 2020), Learning Invariant Predictors with Selective Angmentation (LISA) (Yao et al., 2022), and targeted augmentations (Gao et al., 2023). We also consider optimization methods including GraphAug (Luo et al., 2023) on graph data sets and Stochastic Compositional Augmentation Learning (SCALE) (Chatzipantazis et al., 2023) on image data scts.
For contrastive learning, we consider SimCI. (Chen et al., 2020b) and LISA (Yao et al., 2022) on image data sets, InfoGraph (Sun et al., 2020), RandAugment (Cubuk et al., 2020), and GraphCL (You et al., 2020) on graph data sets. We also consider Joint Augmentation Optimization (JOAO) (You et al., 2021) on graph data sets, and SCALE (Chatzipantazis et al., 2023) on image data sets.
Implementations. We use four graph transformation functions from You et al. (2020), including DropN- odes, PermuteEdges, Subgraph, and MaskingNNdess. Each method modifies fraction of nodes, edges, node fcatures. Specifically, Drop.ooees deletes a randomly sampled set of nodes and their connections. PermuteEdges adds and deletes a randomly sampled set of edges. Subgraph extracts a subgraph generated by random walks from randomly sampled nodes in the graph. NodeMasking masks a randomly sampled fraction of node attributes to zero. We use 0.2, 0.3, 0.4, 0.5 as their perturbation magnitudes. In total, h = 20.
For images, we use sixteen image transformations including ShearX/Y, TranslateX/Y, Rotate, AutoContrast, Invert, Equalize, Solarize, Posterize, Contrast, Color, Brightness, Sharpness, Cutout, and Sample Pairing. We consider five values of perturbation magnitude uniformly spaced between intervals. In total, k = 80.
We use a thyee-layer graph neural network on graph data sets. We use pretrained ResNet-50 on image data sets. In terms of hyperparameters, we search the maximum depth up to and H between [0, 1]. For weighted training, we adjust the learning rate 7) between 0.01, 0.1, 1.0 and the SGD steps a between 25, 50, 100.
4.2 Experimental Results
Supervised learning. We first discuss the results in supervised learning settings, as reported in Table 4. We compare our algorithm with two types of baselines, including pre-specified augmentation methods and optimization methods that search for the composition.
On the protein graph classification data set, our algorithm oulperforms RandAugment, which randomly samples a composition of augmentations for each batch of data, and LISA, which applies the mixup technique between input examples, by 4.3% and 2.3%, respectively. This shows the benefit of searching augmentations over pre-specified schemes. In addition, we also see an improvement of 1.9% over GraphAug.
htttts//wwwwaaciis.ce
httppss////kgll.ooo///oo/// Shtnep///aaggl..oooommmmtti//ppp /compettiionns/ppp2222 4 4https:/ /fggsharehhhahommmmmmm
13

---

Published in Transactions on Machine Learning Research (06/2024)
Table 4: We compare our algorithm with several existing data augmentation schemes on a protein graph classification data set (left) and a wildlife image classification data set (right). In particular, the left-hand side shows the average test AUROC scores for protein function prediction. The right shows the test macro Fi score the image classification data set. We report the averaged rosults over five random seeds.
Graph Classification
Image Classification
Training Set Size
12,302
6,568
Validation Set Size
4,100
426
Testing Set Size
1,102
789
# Classes
1,198
182
Metric
AUROC
Macro Fi
Empirical Risk Min.
71.3 + 0.3
52.4 1.1
Randengment
71.8 + 0.8
58.9 = 0.4
LISA
73.2 + 0.3
59.7 0.3
Targeted Augmentation
56.3 = 0.4
GraphAug
73.5 + 0.3
SCALE
60.4 + 0.5
Algorithm (Ours)
74.9 +: 0.4
62.3 + 0.6
On the image classification data set, our algorithm outperforms pre-specified data augmentation methods, including RandAngmen,, targeted augmentation, and LISA, by 6.8% on average. The improvement over SCALE is about 3.0%.
Next, we report the runtime. Recall that we first find augmentation composition in each group and then learn weighted combination of them. This algorithm takes 8.2 hours on the graph data set and 2.6 hours the image data set. This is 44% and 38% less than the optimization methods on the graph and image data sets, which use 14.4 and 4.2 hours, respectively. Moreover, as discussed in Section 3.2, the weighted training procedure of our algorithm takes a comparable runtime as SGD.
Contrastive learning. Next, we report the results of contrastive learning. Again, we consider pre-specified augmentation methods, such as SimCLR, and existing optimization methods. We evaluate trained models by linear evaluation using an SVM classifier on the contrastive features.
We first report the results of contrastive learning on image data sets in Table On CIFAR-10, which contains natural images, our algorithm performs on par with SimCLR designed by domain experts. On medical image data set,, we observe that the SimCLR augmentation scheme does not work on medical images and even decreases the performance compared to just using the pretrained network's features. By finding compositions of augmentations, our algorithm improves over SimCLR by 5.9%. In both data sets, our algorithm outperforms HandAugment by 2.4% on average. Moreover, our algorithm improves an optimization method, SCALE, by 1.1% on both image data sets on average. Our algorithm also takes 32% less runtime than the optimization method.
For graph contrastive learning, report the 10-fold cross-validation results in Table 6. Compared to pre-specified schemes, algorithm outperforms Randengnent by up to 20% and GraphCL by up to 17% across six data sets. Compared to JOAO, our algorithm improves by 7.1% %. Our algorithm benefits from splitting graphs into groups with similar sizes and degrees before applying graph contrastive learning.
Ablation studies. We verify that both parts of the algorithm contribute to the results, including finding one tree per group and weighted training. First, we remove the tree for cach group and instead find single tree for all groups. This resulted in a performance decrease of 0.6% and 1.0% on the protein graph and image classification data set, respectively. Second, we remove weighted training and replace with uniform weighting. This leads to performance decrease of 1.0% and 1.6%, respectively. From analyzing model features, we find that our approach results in more similar features between groups than single tree. For details, see A Appendix C.
14

---

Published in Transactions on Machine Learning Research (06/2024)
Table 5: We compare our algorithm with several augmentation methods for contrastive learning on images. We report the accuracy on CIFAR-10 and a medical image classification data set (Messidor). We conduct contrastive learning using ResNet-50 pretrained on ImageNet and evaluate the test accuracy using linear evaluations of last-layer features. The results averaged over five random sceds.
Natural Image Classification
Medical Image Classification
Training Set Size
45,000
10,659
Validation Set Size
5,000
2,670
Test Size
5,000
3,072
# Classes
10
Pretrained Features
88.9% 0.0
70.1% 1% 0.0
SimCLR
92.8% 1: 0.5
69.1% 1.1
Rand Angment
89.8% 0.7
72.4% 0.5
LISA
72.1% + 0.5
SCALE
91.3% 0.3
72.4% 0.2
Algorithm (Ours)
93.0% + 0.4
73.2% + 0.3
Table We compare our algorithm with several existing augmentation methods graph contrastive learning. We report the test classification accuracy on several graph prediction data sets. We first conduct contrastive learning using three-layer graph neural networks and evaluate the test accuracy using linear evaluations of last-layer features. The results are based on ten-fold rrossvaaldation.
NC11
PROTEINS
DD
COLLAB
REDDIT-B
IMDB-B
Graphs
4,110
1,113
1,178
5,000
2,000
1,000
Classes
2
2
2
3
2
2
InfoGraph
76.2%
74.4%
72.8%
70.6%
82.5%
73.0%
RandAugment
62.0%
72.2%
75.7%
58.1%
76.3%
55.2%
GraphCL
77.8%
74.3%
78.6%
71.3%
89.4%
71.1%
JOAO
78.0%
74.5%
77.4%
69.5%
86.4%
70.8%
Algorithm 2 (Ours)
79.3%
77.7%
78.9%
78.4%
89.4%
74.4%
In Algorithm 1, we require the maximum depth of the tree d and the probabilities H. In Algorithm 2, we require the number of SGD steps a, and learning rate 17 for updating combination weights. In each ablation study, vary one hyper-parameter while keeping the others constant. The fixed hyperparameters are follows: tree depth of 4, number of probabilities of 10, SGD steps of 50, and learning rate of 0.1. Table 7 presents the findings. Setting d = and H = 10 yields the best results. Additionally, using a = 50 and = 0.1 yields the best results. We also note that these are effective other settings. Therefore, we use these values as the default in all of our experiments.
4.3 Extension and Discussion
Extension. To further illustrate the application of our approach, we consider semi-supervised lcarning and tree augmentation for consistency regularization (Xie et. al., 2020).
We evaluate approach on CIFAR-10 and SVHN. For CIFAR-10, use 1,000 labeled and 44,000 unlabeled examples. For SVHN, we use 1,000 labeled and 64,932 unlabeled examples. we also consider Chest-X-Ray which contains frontal-view X-ray images labeled with 14 diseases, each viewed as a zero-one label (Wang et al., 2017; Rajpurkar et al., 2017). We use 1,000 labeled and 9,000 unlabeled examples. We train a randomly initialized Wide-ResNot-28-1U1 on all datasets using SGD with learning rate of 0.03 and 100,000 gradient update steps, following Xie et al. (2020). To determine the composition of data augmentation, we use the same 16 types of transformations in other image data sets.
15

---

Published in Transactions on Machine Learning Research (06/2024)
Table 7: Ablation study of varying maximum tree depth d, number of probability values |H|, SGD steps a, and learning rate 7 in Algorithm 2. We report the average AUROC for the protein graph classification data set and the Macro-Fi for the wildlife image classification data set (on validation sets). The results are averaged over five random sceds.
Graph Classification (AUROC)
Image Classification (macro Fi)
d=2
d=3
d = 4
2
d
72.4% 0.1
74.2% + 0.2
74.5% 0.4
61.6% + 0.6
65.0% + 0.3
65.3% +: 0.6
||| 5
[ = 10
[7] 20
[ =
| = 10
[ = 20
72.9% + 0.2
74.5% = 0.4
74.5% + 0.6
64.7% + 0.3
65.3% = 0.6
65.3% + 0.4
25
a 50
100
25
50
100
74.0% + 0.2
74.5% 0.4
73.4% 1 0.8
64.8% 1 0.3
65.3% + 0.6
64.2% + 0.4
17 == 1.0
1) 0.1
1 = 0.01
17 = 1.0
?
0.1
1) = 0.01
71.0%
74.5% 0.4
72.4% + 0.8
60.5%
0.7
65.3%
0.6
64.3% 0.5
Table reports the test performance on three data sets. We verify that our algorithm can outperform UDA, which uses RandAugment, by 1.2% averaged over the data sets.
Table 8: We apply our approach for semi-supervised learning. We compare to Unsupervised Data Augment a ion (UDA) (Xie et al., 2020), which uses RandAugmen as the base augmentation scheme. We report the results from training a randomly initialized WidcResNNet-28-10 averaged over five random soeds.
CIFAR-10
SVHN
Chest-X-Ray
Labeled Training Set Size
1,000
1,000
1,000
Unlabeled Training Set Size
44,000
64.932
9,000
Validation Set Size
5,000
7,325
7,846
Test Set Size
10,000
26,032
22,433
Metrics
Error Rates (L)
Error Rates (L)
AUROC (t)
UDA
4.9% 0.1
2.7% + 0.1
73.3% 0.5
Algorithm 1 (Ours)
4.5% + 0.1
2.2% + 0.1
75.2% + 0.6
Discussion. One justification for our empirical findings is that the trees generalize the sequences, which have been the focus of prior work (Ratner et al., 2017; Wu et al., 2020; Xie et al., 2020; Chen et al., 2020b). Since the search now happens in this generalized set, the results are expected to be on par, if not better. It is an interesting research question to understand why the greedy construction of trees works well empirically. See a recent work (Deng & Hsu, 2024) that develops a new framework for formally reasoning about learning hierarchical groups in multi-group learning.
5 Conclusion
This paper designs an algorithm for learning binary tree composition of data augmentation. The algorithm uses top-down recursive search method to find tree with reduced running time complexity. Experiments validate that the algorithm reduces the runtime compared to existing search methods without decreasing downstream performance. Next, the algorithm is extended to the case of learning under group shifts. The algorithm first finds one tree for each group to account for helerogeneous features and then reweights each tree into a forest of trees. Extensive experiments show the empirical benefits of this approach over existing augmentation methods. The algorithm can be readily extended to tackle different loss metrics such as the worst-group loss. It may also be worth revisiting out-of-domain generalization where heterogencous domains are present through the design of automatic data augmentation.
16

---

Published in Transactions on Machine Learning Research (06/2024)
Acknowledgement
We would like to thank Michelle Velyunskiy for her work on collecting the protein graphs data set during the initial stage of this project. Thanks to Jinhong Yu for setting up some of the computational environments in the experiments. We are very grateful to the anonymous reviewers and the action editor for various insightful comments that have resulted in a significant improvement of this work.
References
Micah Adler and Brent Heeringa. Approximating optimal binary decision trees. In International Workshop on Approzimation Algorithms for Combinatorial Optimization, pp. 1-9. Springer, 2008. 4
Michael Ashburner, Catherine A. Ball, Judith A. Blake, David Botstein, Heather Butler, J. Michael Cherry, Allan P. Davis, Kara Dolinski, Selina S. Dwight, Janan T. Eppig, Midori A. Harris, David P. Hill, Laurie Tarver, Andrew Kasarskis, Suzanna Lewis, John C. Matese, Joel E. Richardson, Martin Ringwald, Gerald M. Rubin, and Gavin Sherlock. Gene ontology: tool for the unification of biology. The Gene Ontology Consortium. Nat Genet, 25(1):25 29, 2000. 3
Sara. Beery, Arushi Agarwal, Elijah Cole, and Vighnesh Birodkar. The iwildcam 2021 competition dataset. CVPR, 2021. 12
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew G Wilson. Learning invariances in neural networks from training data. NeurIPS, 2020. 3
Karsten M Borgwardt, Cheng Soon Ong, Stefan Schönaner, SVN Vishwanathan, Alex J Smola, and Hans-Petor Kriegel. Protein function prediction via graph kernels. Bioinformatics, 2005. 3, 12
Davide Buffelli, Pictro Liò, and Fabio Vandin. Sizcshiftreg: a regularization method for improving size-
generalization in graph neural networks. NewrIPS, 2022. 3
Evangelos Chatzipantazis, Stefanos Pertigkiozoglou, Kostas Daniliddis, and Edgar Dobriban. Learning augmentation distributions using transformed risk minimization. TMLR, 2023. 5, 7, 13
Shuxiao Chen, Edgar Dobriban, and Jane H Tee. A group-theoretic framework for data augmentation. The Journal of Machine Learning Research, 21 (1(1):9885 9955, 2020a. 4
Shuxiao Chen, Koby Crammer, Hangfeng He, Dan Roth, and Weijie J Su. Weighted training for cross-tast task lcarning. ICLR, 2022. 10, 21
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020b. 1, 3, 5, 7, 11, 13, 16
Eli Chien, Chao Pan, Jianhao Peng, and Olgica Milenkovic. You are allset: A multisct function framework for hypergraph neural networks. TCLR, 2022. 4
Wyalt T Clark and Predrag Radivojac. Analysis of protein function and its prediction from amino acid sequence. Proteins: Structure, Function, and Bioinformatics, 2011. 2
Ekin D Cubuk, Barret. Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation policies from data. CVPR, 2019. 1, 2, 5, 7
Ekin Cubuk, Barret Zoph, Jonathon Shlens, and Quoc Le. Randaugment: Practical data augmentation with no separate scarch. CVPRW, 2020. 11, 13
Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data angmentation. In International conference on machine learning, pp. 1528-1537. PMLR, 2019.
Samuel Deng and Daniel Hsu. Multi-group learning for hierarchical groups ICML, 2024. 16
17

---

Published in Transactions on Machine Learning Research (06/2024)
Irena Gao, Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Out-of-domain robustness via. targeted augmentations. ICML, 2023. 2, 13
Anupam Gupta, Viswanath Nagarajan, and R Ravi. Approximation algorithms for optimal decision trees and adaptive tsp problems. Mathematics of Operations Research,
2017. 4
Steve Hanneke and Samory Kpotufe. On the value of target data in transfer learning. NearIPS, 2019. 10, 21
Ignacio Hounie, Luiz FO Chamon, and Alejandro Ribeiro. Automatic data augmentation via invariance- constrained lcarning. In ICML, 2023. 3
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. In NewrIPS, 2020a. 3, 12
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovec. Strategies for pre-training graph neural networks. ICLR, 2020b. 8
Haotian Ju, Dongyue Li, Aneesh Sharma, and Hongyang R Zhang. Generalization in graph neural networks: Improved pac-bayesian bounds on graph diffusion. In AISTATS, 2023.
John Jumper, Richard. Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunkkoo,, Russ Bates, Augustin Zidek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, (06(7773):583-333 2021. 3
Dongyue Li, Haotian Ju, Aneesh Sharma, and Hongyang R Zhang. Boosting multitask learning on graphs through higher-order task affinities. In KDD, 2023. 3
Sungbin Lim, Ildoo Kim, Tacsup Kim, Chiheon Kim, and Sungwoong Kim. Fast autoaugment. In NeurIPS, 2019. 1, 5, 7
Youzhi Luo, Michael McThrow, Wing Yee Au, Tao Komikado, Kanji Uchino, Koji Maruhashi, and Shuiwang Ji. Automated data augmentations for graph classification. ICLR, 2023. 2, 3, 13
Christopher Morris, Nils M Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. arXiv preprint orXiv:2007.080 086663 2020. 13
Thien Hang Nguyen, Hongyang R Zhang, and Huy Nguyen. Improved group robustness via classifier retraining on independent splits. Transactions on Machine Learning Research, 2023. 2, 4
Abhinav Nippani, Dongyue Li, Haotian Ju, Haris Koutsopoulos, and Hongyang Zhang. Graph neural networks for road safety modeling: Datasets and evaluations for accident analysis. NewrIPS, 2023.
Subhadeep Paul, Olgica Milenkovic, and Yuguo Chen. Highor-order spectral clustering under superimposod stochastic block models. Journal of Machine Learning Research, 24(320):1 58, 2023. 4
Predrag Radivojac. Advancing remote homology detection: a step toward understanding and accurately predicting protein function. Cell Systems, 2022. 2
Pranav Rajpurkar, Jeremy Irvin, Kaylie Zhu, Brandon Yang, Hershel Mehta, Tony Duan, Daisy Ding, Aarti Bagul, Curtis Langlotz, Katie Shpanskaya, et al. Chexnet: Radiologist-leva pneumonia detection on chest x-rays with deep learning. arXiv preprint axXiv:1711.055 15225, 2017. 15
Alexander J Ratner, Henry Ehrenberg, Zeshan Hussain, Jared Dunnmon, and Christopher Ré. Learning compose domain-specific transformations for data augmentation. NewrIPS, 2017. 1, 2, 16
18

---

Published in Transactions on Machine Learning Research (06/2024)
Burkhard Rost, Predrag Radivojac, and Yana Bromberg. Protein function in precision medicine: deep
understanding with machine learning. FEBS Letlers, 2016. 2
Han Shao, Omar Montasser, and Avrim Blun. A theory of pac learnability under transformation invariances. Advances in Newral Information Processing Systems, 55:13989-14001, 2022.
Fan-Yun Sun, Jordan Hoffmann, Vikas Verma, and Jian Tang. Infograph: Unsupervised and semi-supervised
graph-level representation learning via mutual information maximization. JCLR, 2020. 13
Xiaosong Wang, Yifan Peng, Le Lu, Zhiyong Lu, Mohammadhadi Bagheri, and Ronald M Summers. Chestx-ray8: Hospital-scale chest x-ray database and benchmarks on weakly-supervised classification and localization of common thorax diseases. In CVPR, 2017. 15
Sen Wu, Hongyang R Zhang, Gregory Valiant, and Christopher Ré. On the generalization effects of linear transformations in data augmentation. In ICML, 2020. 2, 3, 4, 16
Qizhe Xie, Zihang Dai, Eduard Hovy, Thang Luong, and Quoe Le. Unsupervised data augmentation for consistency training. NeurIPS, 2020. 3, 15, 16
Kaiwen Yang, Yanchao Sun, Jiahao Su, Fengxiang He, Xinmei Tian, Furong Huang, Tianyi Zhou, and Dacheng Tao. Adversarial auto-augment with label preservation: representation learning principle guided approach. NewrIPS, 2022.
Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving out-of-distribution robustness via selective augmentation. In ICML, 2022. 13
Gilad Ychudai, Ethan Fetaya, Eli Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks. In ICML, 2021. 3
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations, NewrIPS, 33, 2020. 3, 8, 13
Yuning You, Tianlong Chen, Yang Shen, and Zhangyang Wang. Graph contrastive learning automated. In ICML, 2021. 13
Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Ré. Correct-n-contrast: A contrastive approach for improving robustness to spurious correlations. ICML, 2022. 2, 4
Xinyu Zhang, Qiang Wang, Jian Zhang, and Zhao Zhong. Adversarial autoaugment. In ICLR, 2020.
19

---

Published in Transactions on Machine Learning Research (06/2024)
A Derivation of Equation (5)
By regarding 8t as function of w and applying the chain rule, we have the following:
oi(so") (vol(
80
(7)
fw
Ow
By the definition of 8", we must have that 4" is minimizer of E Wy (Jo Qg)- Thus, the gradient with respect to 0 must be zero, given w. Next, we write the gradient as:
F(w, 0) Vo Σ "ig Lg(0:
Σ "g Volyy(o:Q.))
(8)
For a stationary point (m) of equation (4), we have F(w, 0") = 0. By the implicit function theorem," 5 if F(w, 8) is differentiable and (F(w.)) is invertible, then for (w, near some (w, 0) satisfying F(ù, 0) 0, there exists a function mapping w - 9" which is continuous and continuously differentiable and the derivative of this mapping over w is given by:
80
@F 0)
)
aF aF(w, 0) Ow
(9)
Ow
00
10m)
g00
Replacing equation (9) into equation (7), we get
af.(Jaw)
Σ lg Votylloo:
Σ Uy Votyg(ow
Voli(Jou; Qi),
(10)
Ow;
which concludes the proof of equation (5).
B Consistency of the Weighted Training Procedure
This section derives a generalization bound for the bilevel optimization algorithm. The result shows that when the number of samples goes to infinity, the gap between training and test errors will shrink to zero, )m plus discrepancy distance term between the source and target tasks. Let {Pg) denote the groups of augmented data distributions. Concretely, for group 9, we have ng i.i.d samples. Let Pg = (Tg,i, 22 denote the augmented set, by applying Q, to the input samples from Pg. We denote Po as a target data distribution. Given bounded loss function f, let Lyff) and (fe) denote the expected and empirical loss on group respectively. We denote the minimizer of the weighted empirical loss as
i arg min PC0
Σ w. (fo).
(11)
We denote the optimal representation as:
0 arg min Lo())).
0gee
We denote the w-weigtted optimal representation as:
TTh E arg min
Wy (.)
BE0
Since 0= may not unique, we introduce the function space 0= E θ to include all w-weighted ~".
http::/www.ikkkeeiiaaaa////////s function theoren
20

---

Published in Transactions on Machine Learning Research (06/2024)
Definition B.1 ((p, (.)-traasfer. A representation E e is (p. cr)-traanferable from w-weighted souree subsels to the larget subsel, if there erists p > 0, Cp O such that for any 4 E O we have
Lo(0) Lo(O - Cp
Σ Wig (Lg(0) Lg 1)(0")))
Next, we state two technical assumptions. The first assumption describes the loss function's Lipschitz continuity.
Assumption B.1. Let the loss function é X x y + [0 be C-Lipschitz in a E H, meaning for all 22,32 E X and E y, the following holds
[((x1,8) - ((22,))) E C. C. x2|
The second assumption describes the covering size of the functional class.
Assumption B.2. There erist constants Ce and ve greater than 0 such that for any probability measure Qx on X' _ Rd, we have
1(xx01)(0x)))) < ((e) for any > O,
( Ce
de
where ((; Φ; L2 (Qxx) is the minimum number of L (2xx) balls with radius required cover the entire space. Here, the 1.2 (Qxx distance between two vector-valued functions 0 and ub is defoned
12((x) - ( (())))
Given a weight vector w, let us define
77
dist
Pg. Po)
sup Lo(O") Lo())).
j= c
which represents the distance between the source w-weighted and target groups. Now, we are ready to state the result.
Theorem B.3. Let θ, @t 0% be defined the above equations with a fixed w. Suppose Assumptions B.1 and B.2 both hold. Let d E (0, 1) be a fixed real number. Then, for any representation ① 6"r such that Ô is (p, tranfsfasfere,l,, with probability at least 1 - 5, have:
m Σ
(CiV log(CoC) Co log(o-1) 1oagoo
(0) [o(()) S dist
W Pug.o Po + Co
VN.
where No, = (Em un Ci and C2 are two conslants that do not grow with the size of the inpud.
1=
At a high level, the first distance lerm represents the bias increase by transferring across groups, and the second variance term is the empirical error of learning an imperfect representation 0", This term reduces under pooling: Consider a scenario where n, is the same across groups, then through No,, the second reduces by factor of m- The proof technique is based on carefully examining the weighted training procedure using covering numbers (Chen et al., 2022; Hanneke & Kpotufe, 2019).
21

---

Published in Transactions on Machine Learning Research (06/2024)
Proof. Let @ be the domain for which 0 lies in. Consider bounded loss function L. To establish the result, we begin by decomposing the difference between (o( and Lo Lo(D)). We proceed with the following calculations:
Lo( Lo( = Lo(O) - Lo(d + Lo( - Lo(())
adding and subtracting Lu (8
m
S C
Σ 9=1
Wg (Lg(0) ll((0M))
sup 6th Ee"
(Lo(0 Lo(()))
(12)
upper bounded by the supremum
C
Σ Wg
W (Lg(0) -b000)))(000))
dist Ewg CigPan Po
Now we focus on the first term from above, which is Σg= 4g (Lg(0) Lg((())) have
(Tg - l.g(0"))
Σ
(L) (,)(0) + [0) - f.g (000)) ) + (.g(0rr) - L.g (0
wg
g=1
Wg (Lg(() - (6) - Lg - L Lg(
(13)
g=1
m
Wi (L) (0) (0 Lo Lo(0aM)
sup 8E9
1 (L) - ((fe(x Vaii) )= + E( (X (o.s)))))))) - Lo (0((0)))
sup AE
Wy
12g
Denote the last equation above as ((Zzg,)) }), where Zy = {x Ugs}. Step (13) by noting that L.y(0) (0w( is most zero according to equation (11).
Next, will apply McDiarmid's inequality. Let us fix two indices S g S and 1 ig ng. Let define (Zgra} } by replacing ig' with another Egiiig = (Rgrig Ugr E X x V. Given that {Zg.i} and (āgr.t differ by only one element, we have
ng
[G((zz))) G((Zg()))
Wi (L) ((J((xg.)::9.)) Vg.)) E(Sg (ig,i))99.)) (04))
9700
+
Σ
Wy'
(Lg (6) e(folxg .). + E( (Xg (). - Ly (Bw))
ng
+
Wg'
(T)) - E(Fo(xg' Ag bgy) ( + l (x (xg') Us' Log (fro)
Ng
Wg'
Ly (0) - e(Jo(xx' ( Ma' ( ( - E (x 3:M
1'''
Awg' Nggg
+ (g) - (xfo i Fg'
5. ) (0g) e( (Xg' We
The last step above is based on the fact that the loss function is bounded between 0 and 1. Thus, we derive the following result:
2ea
Pr (G(((g))) B[G((2g+))) 6) Sexp
16 >0
(14)
22 g=1
22

---

Published in Transactions on Machine Learning Research (06/2024)
Recall our previous definition that No (>n=1
By sctting ö as
2
= exp exp
> 1=
10 2
get that 2
(1/6) N
Thus, Equation (14) can be equivalently stated as:
log(1/6)
2/2
(15)
Nw
which holds with a probability of at least 1 5 for every E (0, 1).
Next, for the function space θ let
Wa (-E(Jo(zg.)) Vy.)) VOE
Ma N
Tg.
Tlg
where rgg are independent Rademacher random variables. For any 0,, 02 O, define d as
ng
((11,02) N
(e(fo, (Xga), Ug.1) (xg.)).
Now, we justify why this represents random process with sub-Gaussian increments. In the definition of Mo; the Rademacher random variable Tg introduces randomness from the sign of each term. Ma can thus be understood as an empirical average over random sign flips. Additionally, d (01, 02) computes the squared difference in losses under 0 and 02. The squared term ensures that this quantity is non-negative and gives a measure of " "distance" between the two functions in terms of their losses. Combining both observations, we have that the difference Ma. - Ma, a random process with increments characterized by the metric d². Morcover,
E exp (A(Moo
Mo.))] < exp 2 ( (0)) A A 0,0 0,02 e O.
This result shows that the tail probabilities of the process Mo. Me₃ decay at least as fast as those of a Gaussian process, justifying the sub-Gaussian property.
Next, Dudley's entropy integral incquality conditioned the randomness of Zg to obtain
E sup MM MM s 8/2 (' Neg
(16)
Loee
Recall the definition of the covering number from Assumption B.2. Given that
ng
V N 2 No x
We ( (Xg.), Vo.s) - E(So (x),(x))), Waril)
sup BEB
To.
ng
Eq. (16) can be transformed to
16 E [[(((gg.))
(' VlogNNcs O didc.
(17)
VN
Next, we consider the covering number of e. We define
2 ng
8x
where 8x,, represents a point mass Xg.
Tg
ng
s=1
23

---

Published in Transactions on Machine Learning Research (06/2024)
Denote Ne (ee; (Q))))). Let {8(),.. (80.....00N.)) S O be an e-covering of @ with respect to (2))). This implies that for any E O there exists E {1,... Ne) such that
w [x( Lecs (xg,))l s e
M Jesn (2)00) =
m
In conclusion, from the previous steps, we have:
m Bg
((0,050) No
(x,((x)):) - ((for) (g,)))
w
N
|l (x Socs (x1
= cf - far $
S Ce..
(18)
From the above, we have that
(
Ce e
(e uo (Iog(CoC)+ +lg ())
N(LI;; S..
Applying the above to Eq. (17), we get
16v2
VN L.' /ue (19[(C60++ + +log ( de
B[G((Zgat)
16 ve log(CoC)
1+ S'' V log
1
de
(19)
VNw
Further applying the above to Eq. (15), we get
G((zg,))) 16/2, Uh Cog(CoC)
1+ f/' | log (2) de + +2v2
log(1/5)
VNo
N
Finally, from Eq. (12), we can conclude that:
La(0)- - Lo(0)) s c. 16v2/ un (CoC)
1+ 1.' Vlog (2) de 2
log(1/8) Nw.
+ dist
> 9=1
Wa PupPor Po
VN
The proof of Theorem B.3 is thus finished.
7
C Illustration of Feature Similarity
In this section, we explore how our algorithm affects the learned features across groups. We use the protein graph dataset an example. For every pair of groups and j, we compute feature similarity score s(i. between the last-layer features' covariance matrices. For group denote X, R the feature vectors of n4 samples with dimension d. Denote the covariance matrix as X X, We use the rank-r; approximation to the covariance matrix Uk., Di,,, Urry U.T. where r; is chosen to contain 99% of the singular values. Then, we measure the similarity as
s(i.)) = [[UU... D..33 Ug 11/2 l1
(20)
D 1/2 U j.ry D D1/211
In particular, higher values of si. indicate greater similarity between i and j.
24

---

Published in Transactions on Machine Learning Research (06/2024)
We verify that after splitting the graphs based sizes and average degrees, the feature similarity score indeed varies. In particular, for every pair of two groups, train model on their combined data set and measure the feature similarity score between them. At the same time, we compute the differences in average graph sizes and average degrecs between the two groups. Figure 7a shows that the feature similarity score drops as the disparity between graph sizes grows. Figure 7b shows qualitatively similar results when we measure the disparity by average degrees.
Our method can improve the feature similarity between groups; this is shown in Figure 7c. Here measure the averaged feature similarity score, over every pair of groups.
Feature Similarity
0.9
Reature Similority
0.8
Feature similarity
0.8
0.7
0.7
0.6
0.6
0.5
200
400
600
2
6
ERM
GraphAus
LISA
Alg.
Difference in number of nodes
Difference in average degrees
(a) Varying graph sizes
(b) Varying average degrees
(c) Averaged feature similarity score
Figure 7: We observe that groups with larger differences in graph sive average degree exhibit more distinet features in trained model. Using augmentation can result in higher similarities between groups than not using it. In particular, our algorithm improves the similarity score of equation (20) between different groups. Here, ERM refers to empirical risk minimization.
25


Reviews:
[{"review_id": "scym6EZuGv", "reviewer": "Reviewer_aMpo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n1.\tThe proposed algorithm is clearly explained and the idea of an efficient search over tree-structured compositions is sensible and effective. \n\n2.\tThe experiments are extensive and thorough and demonstrate convincingly that the proposed method works well. Ablation studies are provided to demonstrate that the key components of the proposed procedure are all important in contributing to the performance improvement.\n\nWeaknesses\n\n1.\tThe proposed method is a greedy search procedure. In addition, the validation performance during the search is approximated, so that even the greedy search is over an approximate objective. There are no theoretical results characterizing the quality of the approximation or the ability of the overall algorithm to approximate the optimal solution. Theoretical results quantifying the generalization capability of the proposed bilevel optimization procedure are included, but these results, while commendable, do not address the greedy search and approximation issues. \n\n2.\tThere is minimal in-depth experimental analysis to establish why the method works well.  The discussion in the main body of the paper does not go much beyond “Our method is better…” There is some investigation of the similarity of the selected features for different groups in Appendix B. Diagnostic experimental results would be a welcome inclusion in the paper.\n\n3.\tThere is no reporting of variability; there are no statistical tests.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents efficient algorithms for searching over a space of tree-structured compositions of transformations for the purpose of data augmentation. The algorithm has dramatically reduced worst-case runtime complexity, compared to existing approaches, for settings where there is a relatively large number of transformations to choose from. The paper includes an extension to address data distributions with heterogeneous subpopulations. The experiments, conducted on several graph and image datasets, provide support for the claimed advantages of the method, namely demonstrating that the proposed method can reduce computation substantially while achieving a small performance improvement. The paper discusses how the adoption of a tree structure can provide insights into which transformations are most effective for different types of data.  As a final contribution, the paper reports on a newly collected graph dataset.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.\tThe description of the proposed procedure (single tree) is brief (<1page). The presentation should be expanded considerably, with better justification and more in-depth discussion of the design decisions and their validity. \n1a. What is the justification for training only one model at each step? What is the nature of the resultant approximation and what are the implications and ramifications? Why is it considered that this is a reasonable approximation? Please add more discussion on this point. Currently there is just a passing reference to another paper.\n1b. Is there an understanding of how well or poorly the greedy search performs? What aspects of the problem influence this?\n \n2.\tAppendix C provides a generalization bound, but there is no discussion of this result in the main paper. Even the appendix does not clearly present the implications of the result with respect to the overall search procedure. Why are these results downplayed to this extent? They seem to embody a substantial effort (4 pages of theoretical work) and provide interesting insights. One sentence in passing in the main paper seems odd. Please add a paragraph in the main text that describes the nature of the result and the context. \n\n3.\tThere are no citations in Appendix C. Please provide some discussion about related work – does this proof build on work in other papers? Is it related to results in the transfer learning literature? Is there use of proof techniques from other papers? The complete absence of citations is surprising. \n \n4.\tThe approach for learning a forest of trees relies on identifying suitable groups and selecting the number of groups. Could the authors add text to explain how this is achieved? In the provided example the choice is “16 intervals of sizes and average degrees” but there is no explanation of how the number 16 is chosen or how the sizes and partitions are selected.\n\n5.\tPlease add quantification of the variability of the proposed technique. This would be particularly useful in this setting to get a better understanding of the robustness of the greedy search (whether it is prone to find local minima). Please conduct appropriate statistical significance tests (e.g., Wilcoxon pair) to demonstrate that the reported differences are meaningful. \n  \n6.\tMinor: I do not understand why Section 3.1.2 is included in the methodology (Our Proposed Algorithms). These aren’t examples that provide any further insight into how the algorithm works. They are performance experiments and belong in the experiments and results section. Please move them accordingly or justify what value they bring to the methodology section.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "FPRi1UAiDD", "reviewer": "Reviewer_BzR9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:** The paper is well motivated and presented. The proposed algorithms demonstrate superior performance as compared to existing methods. The appendix provides further mathematical details to strengthen the arguments.\n\n**Weaknesses:** \n* The discussion of the augmentation algorithms is somewhat confusing, especially in the treatment of \"no transformation\" and \"identity mapping\". See Requested Changes.\n* I didn't find a link to the implementation of the algo. \n  * This significantly reduces opportunities for further re-use and development of the algorithm.\n  * Diminishes trust in the results as they cannot be re-produced\n* I didn't find a link to access the newly created dataset.\n  * This significantly reduces opportunities for further re-use and development of the algorithm.\n  * Diminishes trust in the results as they cannot be re-produced", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes an algorithm for tree-based search for composition of data augmentations. It further proposes a weighted scheme for learning a model over sub-populations within the data, each with specific augmentation tree. The algorithm excels in runtime as compared to existing augmentation-search methods and improves (or is on par) with these in respect to supervised and self-supervised performance metrics. In the experimental section, the paper also introduces a new dataset for classification of protein graphs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "*1. Please clarify the augmentation procedure to address these concerns:*\n* I am confused by the construction of the trees. If I understand these re binary trees (in the sense of every node having two children). However, in Figure 1, the first graph in subfigure 2 seems to have only single child at every step. \n* This gets even more confusing when we consider \"No transformation\" or \"Identity mapping\". Are these included as transformations on its own in the list of possible transformations? Or how can these be introduced into the tree otherwise.\n* Another confusing thing is the binary split probability and identity of the alternative. It seems to me from Algo 1, that at each step a single transformation is picked, tested and placed in the tree as a child with appropriate probability based on testing its performance over a validation set. How do I decide on its alternative? That is if augmentation $A_i$ is chosen with probability $p_i$ what happens with remaining probability $1-p_i$. What is the alternative to $A_i$? No transformation and stop growing the tree? Or Identity transformation so skip and go to next level? This would seem natural (and seems to be suggested in Figure 2) but in Figure 3 and 4 it seems the alternative is none of these. \n\n*2. Please clarify the model learning description to clarify the below:*\n* I gathered that the algorithm operates in two independent steps. First the augmentation trees are constructed for every sub-population, then a common model is learned by the weighted scheme. The trees for each sub-population may be different. Are the individual models of the subpopulations also different at this step or are they somehow synchronised through the updates of step 9. in Algo 1? (I do understand that Algo 2 learns a common model.)\n* If they are not synchronized, I would be interested to know what is the effect of these models possibly diverging and then being forced to align through the second step. The augmentation trees are tuned for the model of Algo 1 and this may in the end be rather different than Algo 2. Please comment.\n* Minor: the fact that the algo learns a single model over the sub-populations is correctly stated in the equations but is somewhat difficult to understand at first sight. Please state it clearly somewhere early in the text to help the reader to follow.\n\n*3. Minor changes:*\n* Section 2, first para: I believe the loss function should read $l: \\mathcal{Y} \\times \\mathcal{Y} \\to \\mathbb{R}$. That is mapping from Y spaces to R, not from X and Y to R.\n* Section 3.2, second para, second sentence: \"Since labeling the protein ... \" - please elaborate. I do not understand what does this mean.\n* Figure 5: please add titles to the graphs (to recognize quickly what they contain without having to read the caption)", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CasBhofZF7", "reviewer": "Reviewer_uCps", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n1. The introduction of a top-down recursive algorithm for finding tree-structured compositions of data augmentations is a significant strength. It offers a novel approach to handling the combinatorial explosion problem in data augmentation space, reducing the computational complexity from \\(O(k^d)\\) to \\(O(2^{dk})\\), where \\(k\\) is the number of transformations and \\(d\\) is the depth of the composition.\n\n2. The extension of the algorithm to manage heterogeneous subpopulations by learning a forest of trees tailored to each subpopulation is particularly commendable. This approach not only enhances performance but also provides insights into the importance of specific augmentations for different data characteristics.\n\n3. The comprehensive testing across several datasets, including a novel multi-label graph classification dataset derived from AlphaFold2, showcases the algorithm's effectiveness and efficiency. The demonstrated reduction in computation time and improvement in performance metrics are convincing arguments for the algorithm's utility.\n\n### Weaknesses\n\n1.  While the algorithm shows impressive results on the datasets tested, the paper could benefit from a broader evaluation across a more diverse set of domains. This would help ascertain the algorithm's generalizability and effectiveness in various practical applications beyond the ones explored.\n\n2. Although the tree structures allow for some level of interpretation regarding the importance of different augmentations, the paper could further delve into the explainability aspect of the proposed method. Providing deeper insights into why certain transformations are more beneficial for specific subpopulations or tasks would enhance the paper's contribution.\n\n3. While the paper introduces an innovative algorithm with demonstrated efficiency and effectiveness, an exploration of the sensitivity of the algorithm's performance to its hyperparameters (e.g., depth of the tree, number of transformations) could provide valuable insights for practical implementation. A detailed analysis could help users better understand the trade-offs involved and how to tune the algorithm for specific applications.\n\n4. Providing more detailed implementation guidelines, code availability, and potentially a user-friendly interface or library could significantly enhance the paper's impact. This would allow a broader range of researchers and practitioners to apply the proposed method to their own problems without needing to navigate the complexities of algorithm implementation from scratch.\n\nIn conclusion, while the paper presents a significant advancement in the field of data augmentation with its innovative algorithm, addressing the above weaknesses could further solidify its contributions and facilitate wider adoption and adaptation of the proposed methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a top-down recursive algorithm for efficiently searching within the space of tree-structured composition of transformations, where each node in the tree corresponds to one transformation. This algorithm significantly reduces the computational complexity from O(k^d) to O(2^{dk}), where k is the number of input transformations and d is the depth of the composition. This improvement is particularly notable as k increases beyond 2. The algorithm also address problems with data distributions with in heterogeneous subpopulations. It does so by partitioning the dataset into groups, finding one tree for each subpopulation, and then learning a weighted combination of these trees. This method not only improves computational efficiency but also enhances the performance by adapting to the variations within the data.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Some suggestions - \n* To ensure the generalizability of the proposed algorithm, it is essential to evaluate its performance across a wider array of datasets, particularly in domains significantly different from the ones tested.\n* Providing comprehensive implementation details will help in better understanding.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The problem of data augmentation is of significant importance across the spectrum of machine learning. The authors propose what I believe to be an intuitively clear and interesting approach for selecting sample transformations for creating a new training point instance. Although the approach of a top-down tree-based search is fairly simple, the added mixture of tree transforms idea for different data splits makes it rather interesting. The whole approach reduces the augmentation process complexity significantly. 

There are extensive experimental results that show the utility of the method, but there are no theoretical results that can provably establish the improvement of the method over other approaches. I find the lack of analyses reasonable due to the difficulty of the problem.

The three reviews were mixed. Two reviewers were positive while one reviewer requested more information about the method, some analytical verifications and more data set evaluations. I agree with the latter that when there is no theoretical analysis, significant experimental evidence should be included. I personally find the results on Alphafold protein data very interesting, but to reach out to a broader audience, the authors should include 1-2 more examples from different applications domains. 

Also, the revision of the paper seemed a bit rushed and I noticed inconsistencies in the new text wrt the old text. For example, the authors used \to in the original submission to indicate mappings, and then reverted to the word "to" in some of the red text. Another thing that I see as a minor weakness is that certain claims are made in a "nonscientific manner" (like "solving highly complex optimization problem"  - please mention complexity right there and also add a reference).

So, my only requirements are to add at least one more data set example and polish up the writing in the newly added text.
Decision: Accept (poster)