HyperInvariances: Amortizing Invariance Learning
Abstract
Providing invariances in a given learning task conveys a key inductive bias that can lead to sampleefficient learning and good generalisation, if correctly specified.
However, the ideal invariances for many problems of interest are often not known, which has led both to a body of engineering lore as well as attempts to provide frameworks for invariance learning.
However, invariance learning is expensive and data intensive for popular neural architectures.
We introduce the notion of amortizing invariance learning.
In an up-front learning phase, we learn a low-dimensional manifold of feature extractors spanning invariance to different transformations using a hyper-network.
Then, for any problem of interest, both model and invariance learning are rapid and efficient by fitting a low-dimensional invariance descriptor an output head.
Empirically, this framework can identify appropriate invariances in different downstream tasks and lead to comparable or better test performance than conventional approaches.
Our
HyperInvariance framework is also theoretically appealing as it enables generalisation-bounds that provide an interesting new operating point in the trade-off between model fit and complexity.
1. Introduction
Exploiting symmetries is crucial to the success of many intelligent systems.
Convolutional neural networks most famously provide automatic translation equivariance, which improves learning for a large set of problems where this is an appropriate assumption.
Symmetries can be seen as an example of inductive bias, which constrains the search space or hypothesis class, of a learner.
When well matched to the problem at hand, this leads to improvements in learn- ing speed and generalisation and robustness to spurious correlations (Geirhos et al., 2020).
The importance of symmetries and invariances has led to research on architectures to enforce specific inductive biases’, such as rotation and scale equivariance (Worrall et al., 2017;
Worrall & Welling, 2019) (cf: translation).
However, for many problems the ideal invariances are not known a-priori.
This has led to a growing body of work on providing the ability to learn invariances, for example, by MAP learning (Benton et al., 2020), marginal-likelihood learning (Immer et al., 2022) and meta-learning (Zhou et al., 2021) (finding a common invariance that works well for a given problem family).
Perhaps the most popular way to imbue general architectures with particular invariances is using data augmentation.
Optimising for the same output under a given transformation leads to invariance to that transformation (Benton et al., 2020).
However, while flexible and simple to implement, this is extremely costly and data intensive.
In a distinct body of work, research has identified a set of augmentations (and hence implicitly invariances) which are widely useful in representation learning for many popular tasks in computer vision (Cubuk et al., 2019).
In the explosively growing field of self-supervised learning, this is exemplified by the suite of augmentations used by popular learners such as SimCLR (Chen et al., 2020), which have been analysed in terms of the resulting invariances (Ericsson et al., 2021).
However, while these augmentation suites are effective for many problems, they are not ideal for all problems (Ericsson et al., 2021; Xiao et al., 2021).
For example, while object recognition may prefer rotation invariance, dense prediction may prefer rotation equivariance.
In this paper we explore the idea of amortising invariance learning.
In an up-front step, we aim to train a feature extractor that compactly encodes a range of possible invariances.
More specifically, we parameterize the feature extractor in terms of a hypernetwork (Ha et al., 2017), which is conditioned on an invariance descriptor.
This HyperInvariance architecture thus defines a low-dimensional manifold of feature extractors defined by hypernetwork inputs, with points on the manifold corresponding to different invariance properties.
In a task-specific step, we take this frozen (hyper) feature extractor and learn a new prediction head as well as the hypernetwork inputs, which correspond to choice of invariance for this task.
In practice this framework means that new tasks can be solved – including both model and invariance learning – quickly, data-efficiently, and with few learnable parameters.
Because invariances are selected out of a low-dimensional set, they can be chosen very efficiently compared to existing approaches that train general neural architectures for invariances from scratch (Raghu et al., 2021;
Immer et al., 2022; Zhou et al., 2021).
From a learning-theoretic point of view, HyperInvariance is also appealing.
Generalisation bounds guarantee testing error in terms of empirical risk (training error) plus model complexity.
Standard frameworks provide different trade offs between these terms.
Flexible deep models provide a good train fit, but poor model complexity, while shallow models provide a weaker train fit, but limited model complexity.
HyperInvariance provides an interesting new operating point.
By squeezing a rich and relevant variety of feature extractors into a low-dimensional space of hypernetwork inputs, we can provide an improved training fit with limited additional model complexity.
2. Related Work Invariance Learning
Invariances have been created by mean embeddings (Lyle et al., 2020) and learned by MAP (Benton et al., 2020), marginal likelihood (Immer et al., 2022), or meta learning (Zhou et al., 2021; Cubuk et al., 2019; Raghu et al., 2021) – where gradients from the validation set are backpropagated to update the invariances or augmentation distributions.
These mostly aim to learn taskspecific symmetries, except (Zhou et al., 2021) who learn symmetries for a given family of tasks.
All these approaches are highly data and compute intensive due to the substantial effort required to train such symmetries into general purpose architectures.
Our HyperInvariance framework amortises the cost of invariance learning so that it is quick and cheap to learn task-specific invariances downstream.
Invariances in Self Supervision
Self-supervised methods (Jing & Tian, 2021; Ericsson et al., 2022) often rely on contrastive augmentations (Chen et al., 2020).
Their success has been interpreted as engendering invariances (Ericsson et al., 2021; Wang & Isola; Purushwalkam Shiva Prakash &
Gupta, 2020) through these augmentations, which in turn provide good inductive bias for downstream tasks.
While self-supervision sometimes aspires to providing a single general purpose feature suited for all tasks in the guise of foundation models (Bommasani et al., 2021), studies have shown that different augmentations (invariances) are suited for different downstream tasks, with no single feature being optimal for all tasks.
This leads to the tedious need to produce and combine an ensemble of features (Xiao et al., 2021;
Ericsson et al., 2021), or task-specific self-supervised pretraining (Raghu et al., 2021), which is extremely costly.
Our
HyperInvariance breaths new life into the notion of general purpose features by defining a parametric feature extractor that spans an easily accessible range of invariances.
3. Methodology Pre-train
We begin by considering a set of pre-training tasks T train with dissimilar invariance requirements.
Let the dataset corresponding to task t be D t
= { x ti , y ti } n t
i
=1 , where n t is the number of samples per-task.
Let f θ and Φ train = { ϕ t } t ∈T train be the shared feature extractor and task-specific decoder weights for training tasks respectively.
Predictions for a task t is given by ˆ y t = ⟨ ϕ t , f θ ( x t ) ⟩ .
During pre-training the invariance descriptor i t for each task is assumed observed.
This is a vector i ∈ [0 , 1] K over K possible transformations, where i k
=
and i k = 0 indicate invariance and sensitivity to the k th factor respectively.
We employ hypernetworks to generate weights for the given encoder given the desired invariance descriptor i t .
In this setup, we use a hypernetwork which is shared between all training tasks to generate weights for the encoder to satisfy task-specific invariance requirements.
The hypernetwork h , parametrized by W generates weights for the encoder given a invariance hyper-parameter i t for task t as θ = h W ( i t ) .
Thus, predictions for task t are
[Equation]
In the pre-training stage, the hypernetwork and the corresponding task-specific parameters are updated to optimize the loss of the training tasks.
[Equation]
Downstream
We next consider a set of downstream tasks
T test , that can be solved optimally with different invariance requirements which are unknown to us.
We denote the training data available for a downstream task t ′ ∈ T test as D t ′ = { x t ′ i , y t ′ i } n t ′ i =1 .
In the downstream task training stage, we employ the optimal hypernetwork learned from the pre-training stage to make predictions for the test-tasks.
For a downstream task t ′ , let ˆ y t ′ =
⟨ ϕ t , f h ⋆ ( i t ) ( x t ′ )
⟩ be the predicted output given a invariance hyper-parameter i
t .
Subsequently, we evaluate the optimal invariance hyperparameters and prediction heads Φ test by minimzing the task-specific loss of the training set,
[Equation]
Continuous vs Discrete Invariances Estimates
This setup provides the option of (i) using continuous invariance estimates i ∗ t and Φ ∗ t directly, or (ii) discretizing invariance estimates, ˜ i ∗ =
round ( i ∗ t ) , before re-learning Φ .
The latter may provide an advantage if i ∗ t can be stuck in a local minima, and also provides a learning theoretic advantage.
Theoretical Analysis
We provide some theoretical analysis into the potential of our approach to overfit when applied to novel downstream tasks not seen during pre-training in the theorem below.
Theorem 3.1.
For 1-Lipschitz loss function, L , ff for all ϕ
we have that ∥ ϕ ∥ ≤ B and ∥ f ϕ ( x ) ∥ ≤ X , the following holds with probability 1 − δ
[Equation]
where I = { 0 , 1 } d is the space of possible invariance hyperparameters.
The proof is in the appendix.
This theorem demonstrates that the overfitting behaviour of our approach when applied to novel tasks scales similarly to conventional linear models (i.e., 2 XB √ n ), but due to the more flexible feature extractor we have a potential to obtain a much better fit on the training data, and therefore a better guarantee on test error.
In contrast, formal analysis of the overfitting behaviour of deep neural networks cannot obtain as tight bounds, instead depending exponentially on the depth of the network due to the product of norms of the weights in each layer (Bartlett et al., 2017; Golowich et al., 2018; Long & Sedghi, 2020).
4. Experiments and Results
We conduct synthetic and real-world learning experiments.
Synthetic experiments: Firstly, we perform a set of simple experiments consisting of hand-designed tasks with distinct invariance requirements trained jointly in a multitask learning setup.
Pre-Train: We consider color and rotation invariance on a Colored-Rotated MNIST dataset, which is generated by randomly coloring one of the RGB channels of grayscale Rotated MNIST images.
In our experiments, Rotated MNIST is generated by rotating
MNIST images by randomly selecting an angle from the set {− 90 , − 60 , − 30 , 0 , 30 , 60 , 90 } .
We consider three tasks: (i) digit prediction, (ii) rotation prediction and (iii) color prediction.
In this setup, digit prediction is ideally color and rotation invariant, while rotation angle prediction is rotationsensitive and color-invariant and color prediction is rotationinvariant and color-sensitive.
Following the convention of assigning invariance descriptors, we assign the invariance hyper-parameters to be i =
[1 , 1] , i =
[1 , 0] and i =
[0 , 1] for digit, color, and rotation prediction respectively.
We use
Augerino (Benton et al., 2020) augmentations to teach the invariances corresponding to each hyper-parameter setting. [Table] [Table] [Table]
With these tasks, we learn a hypernet derived feature encoder and distinct readout heads (Eqs. 1 and 2).
Downstream:
As a downstream task, we consider digit prediction and rotation prediction on a new domain of Colored-Rotated KMNIST, generated in the same way as the pre-training dataset to examine if the HyperInvariance model can learn appropriate invariances with limited data.
Here the (hyper) feature extractor is frozen, and a new readout head and hyperparameters (Eq 3) are learned.
Competitor:
For a baseline we do multi-task learning (MTL) among source tasks with a conventional shared feature extractor, and learn a new head for the downstream task.
Contrastive Learning experiment: We next evaluate our
HyperInvariance framework with a real-world contrastive learning experiment, taking SimCLR (Chen et al., 2020) as a representative state of the art learner to build upon.
To define a set of invariances of interest, we borrow from Ericsson et al.
(2021) who extract two subsets of SimCLR augmentations denoted dorsal and ventral .
Ericsson et al.
(2021) showed that contrastive models trained to maximize similarity between images and their dorsal/ventral augmented counterparts learn representations invariant to corresponding transformations.
Pre-train: We instantiate our framework with SimCLR (hereafter referred to as HyperSimCLR) and a ResNet18 architecture and train on STL-10 (Coates et al., 2011).
During training, we assign invariance hyperparameters as i =
[1 , 1] , i =
[1 , 0] and i =
[0 , 1] for default, ventral and dorsal augmentations respectively.
For every i , the hypernetwork is trained to optimize the contrastive loss for an image and the counterpart with augmentation corresponding to i .
Downstream Given the learned frozen
HyperSimCLR-ResNet18 feature, we train linear readout and hyperparameters for new tasks.
For recognition we evaluate CIFAR10 classification, and for regression, we evaluate 300W facial interest point detection (Sagonas et al., 2016).
4.1. Results
With the experiment setup above, we aim to answer the following three research questions:
Does the HyperNetwork learn to generate invariant features?
To examine if the hypernetwork has learned features that are provide a specified invariance, we measure invariance by evaluating the cosine similarity between images and their augmented versions when providing different hypernetwork inputs.
Figure 1(left) reports the invariance to rotation and color in the synthetic multi-task experiment.
At each x-coordinate we pass [ i, 1 − i ] to the hypernetwork to interpolate between rotation and color invariance.
The trend shows that invariances broadly vary between minimum and maximum as a function of hyperparameter i .
Figure 1(middle) performs the corresponding evaluation for dorsal and ventral augmentations in the SimCLR-ResNet18 experiment.
Again we see that measured invariance is a near monotonic function of the hyperparameter.
These results show that the hypernet can synthesize features with a desired invariance for both shallow and deep CNNs, and that both supervised multi-task and self-supervised pre-training can be used.
Can the preferred invariance of a downstream task be identified?
Given our pre-trained hypernet (on MNIST and STL10 respectively) we study downstream tasks (KMNIST and CIFAR, 300W respectively) by learning a linear readout and invariance hyperparameters.
First, for the synthetic experiment, we solve downstream KMNIST with a variety of invariance parameters, interpolating along [ i, 1 ] (color invariance).
We report the training loss for digit prediction and color prediction.
The results in Figure 1(right) show that loss is a clean monotonic function of the corresponding invariance: Digit prediction is best with maximum color invariance; color prediction is obviously best with minimum color invariance.
Both curves are quite smooth, so we expect optimising loss wrt the invariance parameter (x-axis) will discover a good invariance parameter.
Next we check if we can optimise downstream tasks wrt preferred invariance parameter
(Eq. 3).
The results for synthetic digit prediction and color prediction are shown in Tablefor N training samples per class.
Inspecting Table 1 for digit prediction, we see that invariances are tending towards [1 , 1] (and always round to [1 , 1] ) correctly identifying that digit prediction should be both color and rotation invariant.
For rotation prediction in Table 1, invariances are tending towards [0 , 1] (and always round to [0 , 1] ) confirming that color prediction should be only rotation invariant.
Quantitative results: For the synthetic experiment results in Table 1 shows that (i)
HyperInvariance outperforms the multi-task baseline in accuracy, especially at low samples per class – demonstrating the benefit of the correctly discovered invariance as an inductive bias.
(ii)
Discretizing the invariances usually performs only slightly worse ( A I ∗ vs
A i ∗ ), while providing theoretical benefits.
For the HyperSimCLR experiment in Table 2, we can see that for CIFAR-10: (i) Our HyperSimCLR selects both dorsal and ventral invariance (ii) it performs similarly to regular SimCLR baseline and both outperform the dorsaland ventral-alone baselines.
Meanwhile for 300W: (i) HyperSimCLR clearly outperforms regular SimCLR thanks to the ability to tune invariances to a moderate amount with a small preference for Dorsal invariances, and even slightly outperforms DorsalSimCLR which was previously best at this task.
Overall, this solves the problem in (Ericsson et al., 2021) where specific choice of dorsal, ventral, or regular
SimCLR had to be made on a per-problem basis.
A single
HyperSimCLR feature can support different tasks through amortised invariance learning.
4.2. Discussion
We have shown that multiple invariances can be pre-learned by a single (hyper) feature extractor.
This enables down- stream tasks to easily select a useful invariance.
This provides an interesting new avenue of study for general purpose features suited for diverse downstream tasks.
Acknowledgements
R.C. was supported by Samsung AI Research, Cambridge.
A. Additional Discussion On Theory
We can contrast the result in Theorem 3.1 with the corresponding result for standard approaches.
The corresponding result for the popular protocol of linear readout from a fixed feature would reduce the third term in the right-hand side of the bound by eliminating the dependence on | I | .
However, it would also worsen the empirical risk (first term, RHS).
This is illustrated in Table 3 where we also show the train performance for HyperInvariance and the MTL baseline – HyperInvariance provides a clearly better training fit.
Meanwhile, the corresponding result for the other popular protocol of fine-tuning the whole feature extractor would improve the first empirical risk term in the bound, but introduce an exploding complexity term in the bound.
In particular, the second term of the bound would be proportional to a product of norms of the weight matrices in each layer, thus scaling exponentially with the depth of the network.
See, e.g., Golowich et al.
(2018) for a demonstration of this in the framework of Rademacher complexity-based analysis.
This is why we describe HyperInvariance as providing an interesting new theoretical operating point between two popular protocols.
Connection to Empirical Results
With regard to the discretization operation, our current experiments report binary discretization of the continuously estimated invariance parameter, which works well in the synthetic experiments and
CIFAR-10.
We remark that our framework and Theorem are compatible with any quantization strength, for example ternary or higher.
This may provide a better tradeoff between empirical performance + theoretical guarantees for benchmarks like 300W, where the strong results for our continuous model ( A i ∗ ) in Tab 2 suggest that a moderate amount of invariance is preferred.
A.1. Proof of Theorem 3.
Proof.
There is a one-to-one mapping between elements of the finite set of invariance hyperparameters and potential feature extractors, implying that there is also a finite number, | I | , of potential feature extractors for a novel task.
For a linear model coupled with a pre-selected feature extractor, one can apply the standard generalisation bound for linear models based on (empirical)
Rademacher complexity (Bartlett & Mendelson, 2002) to obtain, with probability 1 − δ ,
[Equation]
Our result follows from using the union bound to optimise over the choice of feature extractor.
B. Additional Details B.1. Synthetic experiments
Hypernetwork:
For the multi-task learning experiments, we fix f θ to be a single convolution layer with 16 filters of kernel size 5 × 5 .
This convolution is performed with a stride of 2 and is followed by batch normalization and ReLU before passing the features to the task-specific heads.
The hypernetwork is designed as a two-layer network generates weights for the convolution layer given binary invariance descriptors .
We design the hypernetwork such that w 1 ∈ R k × d h , b 1 ∈ R d h , w 2 ∈ R d h × d out and b 2 ∈ R d out .
Here, we choose d h = 40 and subsequently d out = 1200 which is equal to the total number of weights of the convolution layer.
For colored and rotated MNIST images, the output of the hypernetwork is further resized as 3 × 16 × 5 × 5 .
Training Details: For digit prediction, color predictions, and rotation prediction the output cardinality of the task-specific prediction heads is 10, 3, and 7 respectively.
During the pre-training stage, the hypernetwork and the task-specific weights are trained for 200 epochs with the Adam optimizer with a learning rate of 5 e − 4 along with a cosine annealing learning rate scheduler.
During the downstream task learning stage, the invariance hyper-parameters and task-specific decoders are trained with the Adam optimizer with the same learning rate as above.
[Table]
N i ⋆
A train i ⋆
A i ⋆
A trainMTL
A MTL i ⋆
A train i ⋆
A i ⋆
A trainMTL
A MTL [61, 65] 98.
[Table] [Table]
[Table] [Table] [Table]
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
B.2. Hyper-SimCLR
Hypernetwork: For contrastive learning experiments, the hypernetwork is trained to generate the parameters of the convolution layers of the Resnet18 model.
Unlike the synthetic experiments, the hypernetwork generates kernels for multiple convolution layers of the Resnet18 model.
For SimCLR experiments with two types of invariances ( k = 2 ), the hypernetwork is designed as a two layer-network with and d h = 64 such that w 1 ∈ R 2 × 64 and b 1 ∈ R 64 .
However, to generate a convolution layer with parameters denoted by { θ l } 18 l =1 of the Resnet18 architecture, we learn a separate set of { w l 2 , b l 2 } 18 l =1 such that θ l = w l 2 T ( σ ( w T 1 i ) + b 1 ) + b l 2 .
We direct the readers to (Ha et al., 2017) for more details on the hypernetwork architecture that generates all convolution kernels for Resnet18.
Augmentations: In this section of the appendix, we provide details about the augmentation policies used to train HyperSimCLR (Ericsson et al., 2021).
Ventral augmentation is a set of spatial transformations including random resized crop and random horizonal flip.
Dorsal augmentation is a set of appearance changing augmentations consisting random grayscaling, random color jitter and gaussian blurring.
Finally, the default augmentations is a combination of both ventral and dorsal augmentations.
Datasets: We pretrain the Hyper-SimCLR model on the unlabeled split of the STL-10 dataset consisting of 100000 images with each image is of size 96 × 96 .
We resize these images to 224 × 224 to match the image size of the downstream tasks.
To evaluate the Hyper-SimCLR model on downstream tasks, we perform experiments on 300W landmark detection and
CIFAR-10 image classification.
On 300W use only the indoor sets where we use 40% of the images to form a test set.
For
CIFAR-10, 80 % of the labeled data available is used for learning invariances.
Training details: During the pre-trainig stage, the hypernetwork is trained for 200 epochs with the Adam optimizer with a learning rate of 3 e − 4 , and weight decay coefficient of 1 e − 4 along with a cosine annealing learning rate scheduler.
During the fine-tuning stage, the invariance hyper-parameters and corresponding task-specific prediction heads are trained forepochs with the Adam optimizer using a learning rate of 3 e − 4 with a multi-step learning rate scheduler that decays the learning rate of each parameter group by γ = 0 .
1 after every 10 epochs.
In this stage, we apply weight decay