cs.LGFeb 16, 2026

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

Authors: Abdelali Bouyahia, Frédéric LeBlanc, Mario Marchand

Organizations: Département d’informatique et génie logiciel Université Laval Québec, QC, Canada · Institut intelligence et données Université Laval Québec, G1V 0A6, Canada

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: Information-theoretic generalization bound components under varying augmentation conditions. Top row: Evolution of the distribution shift (KL-divergence), orbit-averaged mutual information, augmentation mutual information, and total bound (scaled by RR ) as a function of augmentation variance t2t^{2} . Increasing augmentation strength amplifies distribution shift and total bound while reducing orbit-averaged mutual information. Bottom row: Dependence of the same four components on the number of augmentations per sample nn , for different training set sizes mm . The augmentation mutual information term decays rapidly with nn , leading to a tighter overall bound for larger nn and m.
Figure 2: Empirical evaluation of the generalization bound in Theorem 5.2 on MNIST and FashionMNIST with affine augmentations. Specifically, the transformation transforms.RandomAffine(degrees = strength × 10, translate = (strength, strength)) randomly rotates each image within ±(10 × strength) degrees and translates it by up to a fraction strength of the image width and height. (a, c) Estimated bound versus empirical generalization gap as a function of augmentation level. The bound increases consistently with transformation strength and closely follows the observed generalization behavior. (b, d) Contribution of individual terms; KL-divergence, orbit-averaged mutual information, and per-example augmentation mutual information; showing that the KL-divergence term dominates while mutual information components remain relatively stable across augmentation levels.
More in this topic

Data Augmentation

Data augmentation expands or transforms training examples to improve learning, data efficiency, and robustness to distribution shifts.

See all 249 papers

Explore similar work

CardsList
Jun 23, 2026

Data Augmentation: A Fourier Analysis Perspective

Data augmentation is a simple and model-agnostic approach for exploiting known invariances in learning problems. Given a group acting on the input space, one augments the training set with transformed copies of each sample. Because it exploits symmetries without modifying the underlying learning algorithm, data augmentation can be applied broadly across learning methods. However, this universality comes at a computational cost: when the group is large, full group-sized augmentation quickly becomes computationally infeasible. This raises a fundamental question: Can partial data augmentation achieve the same statistical benefits as full augmentation in terms of generalization and sample complexity? We develop a general framework for investigating this question using Fourier analysis and the representation theory of finite groups. We show that, for a broad class of classical learning problems, partial data augmentation based on a randomly sampled subset of group elements achieves the same minimax rates as full augmentation, up to an approximation error that vanishes as the subset size increases. Our results provide a theoretical explanation for why partial augmentation can retain the statistical benefits of full augmentation despite enforcing symmetry only approximately, and shed light on a recently raised question in learning with symmetries: whether statistically optimal learning under general group invariances can be achieved using computationally scalable methods. Moreover, we prove a complementary impossibility result: enforcing exact invariance via data augmentation requires averaging over the entire group, and cannot be achieved by any strict subset when the hypothesis space is sufficiently expressive. Together, these results provide a unified perspective on full and partial data augmentation, as well as exact and approximate symmetry enforcement.

Behrooz Tahmasebi, Melanie Weber, Stefanie Jegelka

Harvard John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA 02138, USA. · Technical University of Munich (CIT, MCML, MDSI) · MIT Computer Science and Artificial Intelligence Laboratory (CSAIL)

Data AugmentationStatistical Learning Theory
AskPDF
Oct 20, 2026

Understanding Data Augmentation under Label Noise

Data augmentation (DA) has become a fundamental component of modern machine learning systems and is widely regarded as an effective strategy for improving robustness and generalization. Despite its empirical success, the theoretical impact of DA under label noise remains poorly understood. In this work, we first theoretically explain when DA is beneficial and when it becomes detrimental, and provide sufficient conditions for distinguishing between these two regimes under both the i.i.d. learning setting and the more challenging domain generalization scenario. Then, we reveal a counterintuitive phenomenon: DA essentially acts as an uncertainty amplifier in the learning process, decreasing model stability and generalization performance. This finding is supported by the proposed learning and generalization bounds. Our analyses do not depend on specific DA techniques and require only mild assumptions, indicating that the results are applicable to a broad range of machine learning scenarios. Namely, we provide a unified theoretical foundation for understanding, analyzing, and designing DA methods in the presence of label noise.

Hong Zheng, Fei Teng

School of Computing and Artificial Intelligence Southwest Jiaotong University Chengdu, 611756, China · School of Computing and Artificial Intelligence and Engineering Research Center of Sustainable Urban Intelligent Transportation Southwest Jiaotong University and Ministry of Education Chengdu, 611756, China

Data AugmentationLearning with Noisy LabelsStatistical Learning Theory
AskPDF
May 11, 2026

Characterizing the Generalization Error of Random Feature Regression with Arbitrary Data-Augmentation

This paper aims at analyzing the regularization effect that data augmentation induces on supervised regression methods in the proportional regime, where the number of covariates grows proportionally to the number of samples. We provide a tight characterization of the test error, measured in mean squared error, in terms only of the population quantities of the true data, as well as first and second order statistics of the augmentation scheme. Our results are valid under misspecified feature maps, and for any network architecture where only the last readout layer is trained, and the rest of the network is either frozen or randomly initialized. We specify our results in the case of Gaussian data, and show that our asymptotic characterization is tight in this setting.

Lucas Morisset, Alain Durmus, Adrien Hardy

Qube Research & Technologies · École Polytechnique

Neural Network GeneralizationHigh-Dimensional AsymptoticsData AugmentationRandom Feature MethodsStatistical Learning Theory
AskPDF