Source-linked AI summary
Learning Invariances in Neural Networks
Gregory Benton, Marc Finzi, Pavel Izmailov, Andrew Gordon Wilson
TL;DR
The appropriate symmetries for a task, and how strongly to enforce them, are often unknown. Augerino learns a distribution over augmentations from training data alone; experiments report recovery of ground-truth invariances and improved performance across varied tasks.
Problem
The suitable constraints for a task, and their extent, are often unknown, while standard learning objectives select for flexibility rather than constraints.
Source: §1 · p. 1Method
Augerino learns a distribution over augmentations and trains on augmented data to discover invariances and equivariances from training data alone.
Source: §1 · p. 1Results
Augerino recovers ground-truth invariances, including soft invariances, and outperforms specialized baselines and competing augmentation schemes across several tasks.
Source: §10 · p. 9Takeaways & Limitations
The learned augmentation distributions provide an interpretable representation of the dataset and are associated with improved generalization.
Source: §5 · p. 7, §10 · p. 9For equivariant predictions, the stated condition is a uniform distribution over transformations.
Source: §3.1 · p. 4
Abstract
from arXiv · showhide
Invariances to translations have imbued convolutional neural networks with powerful generalization properties. However, we often do not know a priori what invariances are present in the data, or to what extent a model should be invariant to a given symmetry group. We show how to \emph{learn} invariances and equivariances by parameterizing a distribution over augmentations and optimizing the training loss simultaneously with respect to the network parameters and augmentation parameters. With this simple procedure we can recover the correct set and extent of invariances on image classification, regression, segmentation, and molecular property prediction from a large space of augmentations, on training data alone.
1 Introduction
The paper motivates data-driven symmetry discovery, introduces Augerino, and extends augmentation averaging to equivariant targets and affine transformations.
1 Introduction
- Existing architectures often hard-code symmetries, while the appropriate constraints and their degree are unknown; Augerino learns them from training data alone.Source: §1 · p. 1
2 Related Work
- Prior work hard-codes transformations, learns them through Gaussian-process marginal likelihood, or searches augmentation policies; Augerino instead optimizes transformation distributions directly with training loss.Source: §2 · p. 2
3 Augerino: Learning Invariances through Augmentation
- Augerino jointly learns network weights and augmentation-distribution parameters by averaging predictions over sampled transforms and optimizing training loss with regularization favoring broader transformations.Source: §3 · pp. 2–4
3.1 Extension to Equivariant Predictions
- For transformed targets, Augerino averages inversely transformed outputs; equivariance holds for linear image actions when the distribution is right invariant, including uniform distributions.Source: §3.1 · p. 4
3.2 Parameterizing Affine Transformations
- Affine transformations are parameterized through bounded uniform draws over Lie-algebra generators, and L2 regularization recovers the correct invariance level across a wide range of λ.The parameterization includes translations, rotations, scalings, and shears.Source: §3.2 · pp. 4–5
4 Shades of Invariance
Augerino learns the extent of rotational invariance from training data, recovering full, partial, or no invariance across the evaluated tasks. Its learned rotations reach 98.9% test accuracy on rotMNIST, match the partial rotational invariance in Mario and Iggy, and avoid rotation invariance when predicting Olivetti face angles.
4 Shades of Invariance
- Across the experiments, Augerino learns full, soft, and no rotational invariance, with soft invariance identified as especially representative of real-world problems and difficult to encode a priori.The examples are rotMNIST, Mario and Iggy classification, and rotation prediction on Olivetti faces.Source: §4 · p. 5, §4.1 · p. 5, §4.2 · p. 6
4.1 Full Rotational Invariance: rotMNIST
- 98.9% test accuracy on rotMNIST accompanies a learned distribution over rotations, with regularization level having little effect on performance.The experiment targets invariance to rotations because rotMNIST images are randomly rotated.Source: §4.1 · p. 5
4.2 Soft Invariance: Mario & Iggy
- Augerino learns the Mario and Iggy labels’ partial invariance to rotations up to π/4 radians.The labels depend on both character identity and whether the rotation lies in the upper or lower half-plane.Source: §4.2 · p. 6, Figure 3 · p. 6
4.3 Avoiding Invariance: Olivetti Faces
- On Olivetti faces, Augerino narrows rotational augmentation to a support 0.003 radians wide, while the labels encode image rotation.The model is trained on 9,000 images generated from 30 people; an E(2)-steerable network fails because rotation invariance prevents predicting rotation.Source: §4.3 · pp. 6–7, Figure 4 · p. 6
5 Why Augerino Works
Regularization encourages Augerino to expand its augmentation distribution through regions where training loss stays flat, stopping when harmful transformations raise the loss. In the Mario and Iggy experiment, distributions initialized too wide or too narrow converge toward the appropriate rotation width.
5 Why Augerino Works
- Regularization pushes augmentation width through flat-loss regions but does not exceed the data’s invariance range unless its strength is extreme.The loss rises sharply beyond the range of transformations that preserve the data’s invariance.Source: §5 · p. 7
- On Mario and Iggy, regularized loss is optimal near the true invariance width of π/2 and rises quickly beyond it.Without regularization, loss is almost flat up to π/2; regularization encourages larger rotation ranges.Source: §5 · p. 7
- Augerino solutions are substantially flatter than those from standard training and provide better generalization.The paper points to Appendix G, Figure 9 for the flatness comparison.Source: §5 · p. 7
- Augerino converges to the appropriate rotation width from initial distributions that are too wide, too narrow, or approximately correct.The figure shows the distribution over rotation augmentations across training iterations.Source: Figure 6 · p. 8
6 Image Recognition
The CIFAR-10 experiments compare Augerino with no augmentation, fixed common augmentations, and data-based augmentation schemes. Augerino is competitive with advanced data-based schemes, while regularization can help but is not required for accurate models.
6 Image Recognition
- The CIFAR-10 comparison reports test accuracy for models trained with different augmentation schemes.Augerino with ncopies = 4 provides a performance boost with minimally increased training time.Source: Table 1 · p. 8, §6 · p. 8
- Augerino is competitive with advanced models that seek data-based augmentation schemes on CIFAR-10.The comparison includes no augmentation, fixed common augmentations, and Fast AutoAugment; Augerino can use one augmentation per training input and more at test time.Source: §6 · p. 8, Table 1 · p. 8
- Including regularization improves performance, but accurate models can also be trained without it.The authors report that excessively large regularization can make training unstable, with greater variance in achieved accuracy.Source: §6 · p. 8
7 Molecular Property Prediction
On QM9's HOMO task, Augerino learns affine transformations of atomic coordinates and improves on using no augmentation. It is competitive with domain-informed random rotation and translation augmentation, for both non-equivariant and translationally equivariant LieConv models.
7 Molecular Property Prediction
- Augerino learns affine transformations of the atoms' 3D coordinates, using a uniform distribution over transformation generators.The experiments use the LieConv-Trivial model and LieConv-T(3), which has 3D translational equivariance.Source: §7 · pp. 8–9
- On QM9's HOMO task, Augerino performs much better than no augmentation and is competitive with hand-chosen SE(3) random rotation and translation augmentation.The task predicts the highest occupied molecular orbital energy from atomic coordinates and charges; test performance is reported as MAE in meV.Source: §7 · pp. 8–9, Table 2 · p. 9
- Augerino is useful with both LieConv-Trivial and LieConv-T(3), suggesting it can complement architectural equivariance.Source: §7 · p. 9
8 Semantic Segmentation
The rotCamVid experiment tests Augerino in semantic segmentation, where targets are aligned with inputs and the network should be equivariant to transformations in the data. Augerino recovers rotational augmentation while matching the baseline's performance.
8 Semantic Segmentation
- rotCamVid rotates each training and test image by a random angle, with one fixed angle per image and no differently rotated copies of the same image.Source: §8 · p. 9
- The semantic-segmentation experiment uses FC-Densenet and trains Augerino with a Gaussian distribution over random rotations and translations.Source: §8 · p. 9
- Augerino recovers rotational augmentation on rotCamVid while matching the baseline's performance.Source: §8 · p. 9
9 Color-Space Augmentations
Augerino learns brightness and contrast adjustments on STL-10, attaining slightly higher accuracy than the baseline. The conclusion reports recovery of ground-truth and soft invariances across tasks, while the broader-impact discussion anticipates primary use in machine learning.
9 Color-Space Augmentations
- On STL-10, Augerino reaches 89.7 ± 0.3% accuracy versus 89.0 ± 0.35% for the baseline while learning invariance to noticeable brightness and contrast changes.Both models use standard spatial augmentation; the reported mean and standard deviation are across three independent runs.Source: §9 · p. 9
10 Conclusion
- The conclusion reports recovery of ground-truth invariances, including soft invariances, and increased performance across molecular prediction, segmentation, and classification tasks.The authors describe the learned augmentation distributions as interpretable and accurate.Source: §10 · p. 9
Broader Impacts
- The authors characterize the work as methodological and anticipate that Augerino will primarily be used within the machine-learning community.Source: Broader Impacts · p. 10
Appendix
The appendix explains how averaging over a transformation distribution forms an approximately invariant model, discusses partial-support limits, and gives the affine Lie group generators and their transformations.
Appendix
- Augerino’s appendix covers invariant-model formation, affine Lie group generators, and supplementary experimental and training details.It also includes color-space, molecular-property, and solution-width analyses.Source: Appendix · p. 13
A Forming The Invariant Model
- Averaging model outputs over transformations in the distribution’s support forms an approximately invariant model.With partial support, invariance is not guaranteed by the averaging equation alone; training can nevertheless yield approximate insensitivity to the sampled augmentations and generalize it to unseen test data.Source: A Forming The Invariant Model · p. 13
- Augerino learns the maximal rotation range that does not hurt training performance, such as the observed [−π/4, π/4] range for rotated digit 6s.Rotations beyond that range may move some examples below the upper half plane, where they may be more correctly labelled as 9.Source: A Forming The Invariant Model · p. 13
B Lie Group Generators
- Exponentiating the six 2D affine Lie group generators yields transformations for x- and y-translations, rotation, x- and y-scaling, and shearing.Source: B Lie Group Generators · p. 14
C Semantic Segmentation: Details
The segmentation experiment applies Augerino to rotCamVid using a rotation-and-translation augmentation distribution and compares its pixel-wise accuracy with a standard-augmentation baseline.
C Semantic Segmentation: Details
- The experiment applies Augerino to semantic segmentation on rotCamVid, with the learned augmentations shown in Figure 7.Source: C Semantic Segmentation: Details · p. 14
- The rotCamVid dataset contains one randomly rotated copy of each CamVid image, rather than rotational augmentation applied during training.Padding pixels created by rotation are ignored when computing training loss and test accuracy.Source: C Semantic Segmentation: Details · p. 14
- For segmentation, Augerino uses a Gaussian-parameterized distribution covering rotations and translations, with trainable μ and Σ.The sampled affine transformation matrix scales translation by 2/(w + h), where w and h are image width and height.Source: C Semantic Segmentation: Details · p. 14
- 69.8% pixel-wise segmentation accuracy for Augerino versus 68.7% for the standard-augmentation baseline.Source: C Semantic Segmentation: Details · p. 14
D Training Details
The appendix gives training settings and describes a color-space experiment that learns brightness and contrast augmentation ranges on STL-10.
D Training Details
- Networks in Sections 4 and 6 are trained for 200 epochs with initial learning rate 0.01, cosine scheduling, and batch size 128.Classification uses cross entropy; regression uses mean squared error except QM9, which uses mean absolute error.Source: D Training Details · p. 14
- Training uses one sampled augmentation per input, with multiple augmentations at test time to increase accuracy at minimally increased training time.Source: D Training Details · pp. 14–15
E Color-Space Augmentations: Details
- The STL-10 color-space experiment learns uniform-distribution ranges for brightness and contrast adjustments applied sequentially and independently.Figure 8 shows the learned distribution; Augerino learns broad color and contrast adjustments while matching baseline performance.Source: E Color-Space Augmentations: Details · p. 15, Figure 8 · p. 15
F QM9 Experiment
The appendix describes the QM9 affine-transformation setup and reports that Augerino finds flatter CIFAR-10 solutions associated with better generalization.
F QM9 Experiment
- The QM9 experiment uses nine 3D affine generators: three translations, three rotations, two squeezings, and one scaling.The bounds of the uniform distribution are parameterized, with regularization strength 10^-3.Source: F QM9 Experiment · p. 15
G Width of Augerino Solutions
- Augerino enables flatter CIFAR-10 loss-surface solutions, corresponding to better data compression and generalization.The authors trained 10 models with and without Augerino and used effective dimensionality and sensitivity to parameter perturbations to assess solution flatness; lower effective dimensionality implies flatter optima and better generalization [30].Source: G Width of Augerino Solutions · p. 15, Figure 9 · p. 16
Related papers
This paper cites
- Auto-Encoding Variational Bayes2013
- Spatial Transformer Networks2015
- DARTS: Differentiable Architecture Search2018
- Temporal Ensembling for Semi-Supervised Learning2016
- General $E(2)$-Equivariant Steerable CNNs2019
- The One Hundred Layers Tiramisu: Fully Convolutional DenseNets for Semantic Segmentation2016