# The geometry of invariant learning: an information-theoretic analysis of data augmentation and generalization

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 58 of its 92 references that are in the map.
- First posted on arXiv 16 Feb 2026.

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

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Attention-based Deep Multiple Instance Learning](https://mlsanity.com/paper/1802.04712/) (2018)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [Spherical CNNs](https://mlsanity.com/paper/1801.10130/) (2018)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Synthetic Data from Diffusion Models Improves ImageNet Classification](https://mlsanity.com/paper/2304.08466/) (2023)
- [Information-theoretic analysis of generalization capability of learning algorithms](https://mlsanity.com/paper/1705.07809/) (2017)
- [Curriculum Learning: A Survey](https://mlsanity.com/paper/2101.10382/) (2021)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [Prioritized Training on Points that are Learnable, Worth Learning, and Not Yet Learnt](https://mlsanity.com/paper/2206.07137/) (2022)
- [A General Theory of Equivariant CNNs on Homogeneous Spaces](https://mlsanity.com/paper/1811.02017/) (2018)
- [D4: Improving LLM Pretraining via Document De-Duplication and Diversification](https://mlsanity.com/paper/2308.12284/) (2023)
- [Reasoning About Generalization via Conditional Mutual Information](https://mlsanity.com/paper/2001.09122/) (2020)
- [Information-Theoretic Generalization Bounds for SGLD via Data-Dependent Estimates](https://mlsanity.com/paper/1911.02151/) (2019)
- [Tightening Mutual Information Based Bounds on Generalization Error](https://mlsanity.com/paper/1901.04609/) (2019)
- [Deep Scale-spaces: Equivariance Over Scale](https://mlsanity.com/paper/1905.11697/) (2019)
- [Scale-Equivariant Steerable Networks](https://mlsanity.com/paper/1910.11093/) (2019)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [Preserving Statistical Validity in Adaptive Data Analysis](https://mlsanity.com/paper/1411.2664/) (2014)
- [Image Data Augmentation for Deep Learning: A Survey](https://mlsanity.com/paper/2204.08610/) (2022)
- [Sharper bounds for uniformly stable algorithms](https://mlsanity.com/paper/1910.07833/) (2019)
- [Generalization Error Bounds for Noisy, Iterative Algorithms](https://mlsanity.com/paper/1801.04295/) (2018)

## Abstract (arXiv)

Data augmentation is one of the most widely used techniques to improve generalization in modern machine learning, often justified by its ability to promote invariance to label-irrelevant transformations. However, its theoretical role remains only partially understood. In this work, we propose an information-theoretic framework that systematically accounts for the effect of augmentation on generalization and invariance learning. Our approach builds upon mutual information-based bounds, which relate the generalization gap to the amount of information a learning algorithm retains about its training data. We extend this framework by modeling the augmented distribution as a composition of the original data distribution with a distribution over transformations, which naturally induces an orbit-averaged loss function. Under mild sub-Gaussian assumptions on the loss function and the augmentation process, we derive a new generalization bound that decompose the expected generalization gap into three interpretable terms: (1) a distributional divergence between the original and augmented data, (2) a stability term measuring the algorithm dependence on training data, and (3) a sensitivity term capturing the effect of augmentation variability. To connect our bounds to the geometry of the augmentation group, we introduce the notion of group diameter, defined as the maximal perturbation that augmentations can induce in the input space. The group diameter provides a unified control parameter that bounds all three terms and highlights an intrinsic trade-off: small diameters preserve data fidelity but offer limited regularization, while large diameters enhance stability at the cost of increased bias and sensitivity. We validate our theoretical bounds with numerical experiments, demonstrating that it reliably tracks and predicts the behavior of the true generalization gap.