# Learning Tree-Structured Composition of Data Augmentation

arXiv: https://arxiv.org/abs/2408.14381 · HTML: https://mlsanity.com/paper/2408.14381/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 48 references that are in the map.
- First posted on arXiv 26 Aug 2024.

## Builds on (most cited first, 30 of 37)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [Unsupervised Data Augmentation for Consistency Training](https://mlsanity.com/paper/1904.12848/) (2019)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [TUDataset: A collection of benchmark datasets for learning with graphs](https://mlsanity.com/paper/2007.08663/) (2020)
- [CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning](https://mlsanity.com/paper/1711.05225/) (2017)
- [InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization](https://mlsanity.com/paper/1908.01000/) (2019)
- [Graph Contrastive Learning Automated](https://mlsanity.com/paper/2106.07594/) (2021)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Correct-N-Contrast: A Contrastive Approach for Improving Robustness to Spurious Correlations](https://mlsanity.com/paper/2203.01517/) (2022)
- [You are AllSet: A Multiset Function Framework for Hypergraph Neural Networks](https://mlsanity.com/paper/2106.13264/) (2021)
- [From Local Structures to Size Generalization in Graph Neural Networks](https://mlsanity.com/paper/2010.08853/) (2020)
- [Adversarial AutoAugment](https://mlsanity.com/paper/1912.11188/) (2019)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [Learning to Compose Domain-Specific Transformations for Data Augmentation](https://mlsanity.com/paper/1709.01643/) (2017)
- [The iWildCam 2021 Competition Dataset](https://mlsanity.com/paper/2105.03494/) (2021)
- [On the Value of Target Data in Transfer Learning](https://mlsanity.com/paper/2002.04747/) (2020)
- [On the Generalization Effects of Linear Transformations in Data Augmentation](https://mlsanity.com/paper/2005.00695/) (2020)
- [Generalization in Graph Neural Networks: Improved PAC-Bayesian Bounds on Graph Diffusion](https://mlsanity.com/paper/2302.04451/) (2023)
- [SizeShiftReg: a Regularization Method for Improving Size-Generalization in Graph Neural Networks](https://mlsanity.com/paper/2207.07888/) (2022)
- [Out-of-Domain Robustness via Targeted Augmentations](https://mlsanity.com/paper/2302.11861/) (2023)
- [Weighted Training for Cross-Task Learning](https://mlsanity.com/paper/2105.14095/) (2021)
- [A Theory of PAC Learnability under Transformation Invariances](https://mlsanity.com/paper/2202.07552/) (2022)
- [Automated Data Augmentations for Graph Classification](https://mlsanity.com/paper/2202.13248/) (2022)
- [Identification of Negative Transfers in Multitask Learning Using Surrogate Models](https://mlsanity.com/paper/2303.14582/) (2023)
- [Automatic Data Augmentation via Invariance-Constrained Learning](https://mlsanity.com/paper/2209.15031/) (2022)

## Built on by (most cited first, 1 of 1)

- [Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach](https://mlsanity.com/paper/2409.19458/) (2024)

## Abstract (arXiv)

Data augmentation is widely used for training 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 randomly sample from a list of pre-selected transformations, while methods such as AutoAugment apply advanced search to optimize over an augmentation set of size kd, which is the number of transformation sequences of length d, given a list of k transformations. In this paper, we design efficient algorithms whose running time complexity is much faster than the worst-case complexity of O(kd), 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 SimCLR augmentation scheme for contrastive learning. Using a top-down, recursive search procedure, our algorithm achieves a runtime complexity of O(2d k), which is much faster than O(kd) as k increases above 2. We apply our algorithm to tackle data distributions with heterogeneous subpopulations by searching for one tree in each subpopulation and then learning a weighted combination, resulting in a forest of trees. We validate our proposed algorithms on numerous graph and image datasets, including a multi-label graph classification dataset we collected. The dataset 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 by 43% over existing search methods while improving performance by 4.3%. 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.