# Learning Augmentation Distributions using Transformed Risk Minimization

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

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 70 of its 166 references that are in the map.
- First posted on arXiv 16 Nov 2021.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Wide Residual Networks](https://mlsanity.com/paper/1605.07146/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Conditional Generative Adversarial Nets](https://mlsanity.com/paper/1411.1784/) (2014)
- [Spatial Transformer Networks](https://mlsanity.com/paper/1506.02025/) (2015)
- [Matterport3D: Learning from RGB-D Data in Indoor Environments](https://mlsanity.com/paper/1709.06158/) (2017)
- [Random Erasing Data Augmentation](https://mlsanity.com/paper/1708.04896/) (2017)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Vision-and-Language Navigation: Interpreting visually-grounded navigation instructions in real environments](https://mlsanity.com/paper/1711.07280/) (2017)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [Beyond the Nav-Graph: Vision-and-Language Navigation in Continuous Environments](https://mlsanity.com/paper/2004.02857/) (2020)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning](https://mlsanity.com/paper/2002.06715/) (2020)
- [Invariant and Equivariant Graph Networks](https://mlsanity.com/paper/1812.09902/) (2018)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [Data Augmentation Generative Adversarial Networks](https://mlsanity.com/paper/1711.04340/) (2017)
- [Deep Anomaly Detection Using Geometric Transformations](https://mlsanity.com/paper/1805.10917/) (2018)
- [Gauge Equivariant Convolutional Networks and the Icosahedral CNN](https://mlsanity.com/paper/1902.04615/) (2019)

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

- [Robust Fine-Tuning of Deep Neural Networks with Hessian-based Generalization Guarantees](https://mlsanity.com/paper/2206.02659/) (2022)
- [Latent Space Symmetry Discovery](https://mlsanity.com/paper/2310.00105/) (2023)
- [iDECODe: In-distribution Equivariance for Conformal Out-of-distribution Detection](https://mlsanity.com/paper/2201.02331/) (2022)
- [A Theory of PAC Learnability under Transformation Invariances](https://mlsanity.com/paper/2202.07552/) (2022)
- [Automatic Data Augmentation via Invariance-Constrained Learning](https://mlsanity.com/paper/2209.15031/) (2022)
- [Relaxed Equivariance via Multitask Learning](https://mlsanity.com/paper/2410.17878/) (2024)
- [SymmPI: Predictive Inference for Data with Group Symmetries](https://mlsanity.com/paper/2312.16160/) (2023)
- [Discovering Symmetry Groups with Flow Matching](https://mlsanity.com/paper/2512.20043/) (2025)
- [CODiT: Conformal Out-of-Distribution Detection in Time-Series Data](https://mlsanity.com/paper/2207.11769/) (2022)
- [Learning to Transform for Generalizable Instance-wise Invariance](https://mlsanity.com/paper/2309.16672/) (2023)
- [Learning Tree-Structured Composition of Data Augmentation](https://mlsanity.com/paper/2408.14381/) (2024)
- [Optimizing Data Augmentation through Bayesian Model Selection](https://mlsanity.com/paper/2505.21813/) (2025)
- [Approximate Nullspace Augmented Finetuning for Robust Vision Transformers](https://mlsanity.com/paper/2403.10476/) (2024)

## Abstract (arXiv)

We propose a new Transformed Risk Minimization (TRM) framework as an extension of classical risk minimization. In TRM, we optimize not only over predictive models, but also over data transformations; specifically over distributions thereof. As a key application, we focus on learning augmentations; for instance appropriate rotations of images, to improve classification performance with a given class of predictors. Our TRM method (1) jointly learns transformations and models in a single training loop, (2) works with any training algorithm applicable to standard risk minimization, and (3) handles any transforms, such as discrete and continuous classes of augmentations. To avoid overfitting when implementing empirical transformed risk minimization, we propose a novel regularizer based on PAC-Bayes theory. For learning augmentations of images, we propose a new parametrization of the space of augmentations via a stochastic composition of blocks of geometric transforms. This leads to the new Stochastic Compositional Augmentation Learning (SCALE) algorithm. The performance of TRM with SCALE compares favorably to prior methods on CIFAR10/100. Additionally, we show empirically that SCALE can correctly learn certain symmetries in the data distribution (recovering rotations on rotated MNIST) and can also improve calibration of the learned model.