I MPROVING T RANSFORMATION I NVARIANCE IN C ONTRASTIVE R EPRESENTATION L EARNING
A BSTRACT
We propose methods to strengthen the invariance properties of representations obtained by contrastive learning.
While existing approaches implicitly induce a degree of invariance as representations are learned, we look to more directly enforce invariance in the encoding process.
To this end, we ﬁrst introduce a training objective for contrastive learning that uses a novel regularizer to control how the representation changes under transformation.
We show that representations trained with this objective perform better on downstream tasks and are more robust to the introduction of nuisance transformations at test time.
Second, we propose a change to how test time representations are generated by introducing a feature averaging approach that combines encodings from multiple transformations of the original input, ﬁnding that this leads to across the board performance gains.
Finally, we introduce the novel Spirograph dataset to explore our ideas in the context of a differentiable generative process with multiple downstream tasks, showing that our techniques for learning invariance are highly beneﬁcial.
I NTRODUCTION
Learning meaningful representations of data is a central endeavour in artiﬁcial intelligence.
Such representations should retain important information about the original input whilst using fewer bits to store it (van der
Maaten et al., 2009; Gregor et al., 2016).
Semantically meaningful representations may discard a great deal of information about the input whilst capturing what is relevant.
Knowing what to discard, as well as what to keep is key to obtaining powerful representations.
By deﬁning transformations that are believed a priori to distort the original without altering semantic features of interest, we can learn representations that are (approximately) invariant to these transformations (Hadsell et al., 2006).
Such representations may be more efﬁcient and more generalizable than lossless encodings.
Whilst less effective for reconstruction, these representations are useful in many downstream tasks that relate only to the semantic features of the input.
Representation invariance is also a critically important task in of itself: it can lead to improved robustness and remove noise (Du et al., 2020), afford fairness in downstream predictions (Jaiswal et al., 2020), and enhance interpretability (Xu et al., 2018).
Contrastive learning is a recent and highly successful self-supervized approach to representation learning that has achieved state-of-the-art performance in tasks that rely on semantic features, rather than exact reconstruction (van den Oord et al., 2018;
Hjelm et al., 2018; Bachman et al., 2019;
He et al., 2019).
These methods learn to match two different transformations of the same object in representation space, distinguishing them from contrasts that are representations of other objects.
The objective functions used for contrastive learning encourage representations to remain similar under transformation, whilst simultaneously requiring different inputs to be well spread out in representation space (Wang & Isola, 2020).
As such, the choice of transformations is key to their success (Chen et al., 2020a).
Typical choices include random cropping and colour distortion.
However, representations are compared using a similarity function that can be maximized even for representations that are far apart, meaning that the invariance learned is relatively weak.
Unfor- tunately, directly changing the similarity measure hampers the algorithm (Wu et al., 2018; Chen et al., 2020a).
We therefore investigate methods to improve contrastive representations by explicitly encouraging stronger invariance to the set of transformations, without changing the core selfsupervized objective; we look to extract more information about how representations are changing with respect to transformation, and use this to direct the encoder towards greater invariance.
To this end, we ﬁrst develop a gradient regularization term that, when included in the training loss, forces the encoder to learn a representation function that varies slowly with continuous transformations.
This can be seen as constraining the encoder to be approximately transformation invariant.
We demonstrate empirically that while the parameters of the transformation can be recovered from standard contrastive learning representations using just linear regression, this is no longer the case when our regularization is used.
Moreover, our representations perform better on downstream tasks and are robust to the introduction of nuisance transformations at test time.
Test representations are conventionally produced using untransformed inputs (Hjelm et al., 2018;
Kolesnikov et al., 2019), but this fails to combine information from different transformations and views of the object, or to emulate settings in which transformation noise cannot simply be removed at test time.
Our second key proposal is to instead create test time representations by feature averaging over multiple, differently transformed, inputs to address these concerns and to more directly impose invariance.
We show theoretically that this leads to improved performance under linear evaluation protocols, further conﬁrming this result empirically.
We evaluate our approaches ﬁrst on CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009), using transformations appropriate to natural images and evaluating on a downstream classiﬁcation task.
To validate that our ideas transfer to other settings, and to use our gradient regularizer within a fully differentiable generative process, we further introduce a new synthetic dataset called Spirograph.
This provides a greater variety of downstream regression tasks, and allows us to explore the interplay between nuisance transformations and generative factors of interest.
We conﬁrm that using our regularizer during training and our feature averaging at test time both improve performance in terms of transformation invariance, downstream tasks, and robustness to train–test distributional shift.
In summary, the contributions of this paper are as follows:
• We derive a novel contrastive learning objective that leads to more invariant representations.
• We propose test time feature averaging to enforce further invariance.
• We introduce the Spirograph dataset.
• We show empirically that our approaches lead to more invariant representations and achieve state-of-the-art performance for existing downstream task benchmarks.
P ROBABILISTIC FORMULATION OF CONTRASTIVE LEARNING
The goal of unsupervized representation learning is to encode high-dimensional data, such as images, retaining information that may be pertinent to downstream tasks and discarding information that is not.
To formalize this, we consider a data distribution p ( x ) on X and an encoder f θ : X → Z which is a parametrized function mapping from data space to representation space.
Contrastive learning is a self-supervized approach to representation learning that learns to make representations of differently transformed versions of the same input more similar than representations of other inputs.
Of central importance is the set of transformations, also called augmentations (Chen et al., 2020a) or views (Tian et al., 2019), used to distort the data input x .
In the common application of computer vision, it is typical to include resized cropping; brightness, contrast, saturation and hue distortion; greyscale conversion; and horizontal ﬂipping.
We will later introduce the Spirograph dataset which uses quite different transformations.
In general, transformations are assumed to change the input only cosmetically, so all semantic features such as the class label are preserved; the set of transformations indicates changes which can be safely ignored by the encoder.
Formally, we consider a transformation set T ⊆ { t : X → X} and a probability distribution p ( t ) on this set.
A representation z of x is obtained by applying a random transformation t to x and then encoding the result using f θ .
Therefore, we do not have one representation of x , but an implicit distribution p ( z | x ) .
A sample of p ( z | x ) is obtained by sampling t ∼ p ( t ) and setting z = f θ ( t ( x )) .
If the encoder is to discard irrelevant information, we would expect different encodings of x formed with different transformations t to be close in representation space.
Altering the transformation should not lead to big changes in the representations of the same input.
In other words, the distribution p ( z | x ) should place most probability mass in a small region.
However, this does not provide a sufﬁcient training signal for the encoder f
θ as it fails to penalize trivial solutions in which all x are mapped to the same z .
To preserve meaningful information about the input x whilst discarding purely cosmetic features, we should require p ( z | x ) to be focused around a single z whilst simultaneously requiring the representations of different inputs not to be close.
That is, the marginal p ( z ) =
E p ( x ) [ p ( z | x )] should distribute probability mass over representation space.
This intuition is directly reﬂected in contrastive learning.
Most state-of-the-art contrastive learning methods utilize the InfoNCE objective (van den Oord et al., 2018), or close variants of it (Chen et al., 2020a).
InfoNCE uses a batch x 1 , ..., x K of inputs, from which we form pairs of representations ( z 1 , z (cid:48) 1 ) , ..., ( z K , z (cid:48) K ) by applying two random transformations to each input followed by the encoder f θ .
In probabilistic language
[Equation]
conditionally
[Equation]
such that z i , z (cid:48)
i =
f θ ( t ( x )) , f θ ( t (cid:48)
( x )) for i.i.d. transformations t , t (cid:48) ∼ p ( t ) .
Given a learnable similarity score s φ : Z × Z → R , contrastive learning methods minimize the following loss
[Equation]
Written in this way, we see that the loss will be minimized when s φ ( z i , z (cid:48) i ) is large, but s φ ( z i , z (cid:48)
j ) is small for i (cid:54) =
j .
In other words, InfoNCE makes the two samples z i , z (cid:48)
i of p ( z | x = x i ) similar, whilst making
samples z i , z (cid:48)
j of p ( z ) dissimilar.
This can also be understood through the lens of mutual information, for more details see Appendix A.
In practice, the similarity measure used generally takes the form (Chen et al., 2020a)
[Equation]
where g φ is a small neural network and τ is a temperature hyperparameter.
If the encoder f θ is perfectly invariant to the transformations, then z i
= z (cid:48)
i and s φ ( z i , z (cid:48) i ) will be maximal.
However, there are many ways to maximize the InfoNCE objective without encouraging strong invariance in the encoder.
In this paper, we show how we can learn stronger invariances, above and beyond what is learned through the above approach, and that this beneﬁts downstream task performance.
I NVARIANCE BY GRADIENT REGULARIZATION
Contrastive learning with InfoNCE can gently encourage invariance by maximizing s φ ( z , z (cid:48) ) , but does not provide a strong signal to ensure this invariance.
Our ﬁrst core contribution is to show how we can use gradient methods to directly regulate how the representation changes with the transformation and thus ensure the desired invariance.
The key underlying idea is to differentiate the representation with respect to the transformation , and then encourage this gradient to be small so that the representation changes slowly as the transformation is varied.
To formalize this, we begin by looking more closely at the transformations T which are used to deﬁne the distribution p ( z | x ) .
Many transformations, such as brightness adjustment, are controlled by a transformation parameter .
We can include these parameters in our set-up by writing the transformation t as a map from both input space X and transformation parameter space U , i.e. t : X × U → X .
In this formulation, we sample a random transformation parameter from u ∼ p ( u ) which is a distribution on U .
A sample from p ( z | x ) is then obtained by taking z = f θ ( t ( x , u )) , with t now regarded as a ﬁxed function.
The advantage of this change of perspective is that it opens up additional ways to learn stronger invariance of the encoder.
In particular, it may make sense to consider the gradient ∇ u z , which describes the rate of change of z with respect to the transformation.
This only makes sense for some transformation parameters—we can differentiate with respect to the brightness scaling but not with respect to a horizontal ﬂip.
To separate out differentiable and non-differentiable parameters we write u
= α , β where α are the parameters for which it makes sense to consider the derivative ∇ α z .
Intuitively, this gradient should be small to ensure that representations change only slowly as the transformation parameter α is varied.
For clarity of exposition, and for implementation practicalities, it is important to consider gradients of a scalar function, so we introduce an arbitrary direction vector e ∈ Z and deﬁne
[Equation]
so that F : A × B × X × Z
→
R calculates the scalar projection of the normalized representation z / (cid:107) z (cid:107) 2 in the e direction.
To encourage an encoder that is invariant to changes in α , we would like to minimize the expected conditional variance of F [-e as α is changed, taking the expectation over different directions e . 
 That is, we wish-]with respect to α :
[Equation]
where we have exploited independence to write p ( x , β , e ) = p ( x ) p ( β ) p ( e ) . Deﬁning V requires a distribution for e to be speciﬁed. For this, we make components of e independent Rademacher random variables, justiﬁcation for which is included in Appendix B.
A naive estimator of V can be formed using a direct nested Monte Carlo estimator (Rainforth et al., 2018) of sample variances, which, including Bessel’s correction, is given by V
[Equation]
where x i , β i , e i ∼ p ( x ) p ( β ) p ( e ) and α ij ∼ p
( α ) .
However, this estimator requires LK forward passes through the encoder f
θ to evaluate.
As an alternative to this computationally prohibitive approach, we consider a ﬁrst-order approximation 2 to F
[Equation]
and the following alternative form for the conditional variance (see Appendix B for a derivation) Var
[Equation]
Combining these two ideas
[Equation]
we have
[Equation]
Here we have an approximation of the conditional variance V that uses gradient information.
Including this as a regularizer within contrastive learning will encourage the encoder to reduce the magnitude of the conditional variance V , forcing the representation to change slowly as the transformation is varied and thus inducing approximate invariance to the transformations.
An unbiased estimator of equation 11 using a batch x 1 , ..., x K is
[Equation]
where x i , α i , β i , e i ,
∼ p
( x ) p ( α ) p ( β ) p ( e ) , α (cid:48) ij ∼ p
( α ) .
We can cheaply use a large number of samples for α (cid:48) without having to take any additional forward passes through the encoder: we only require K evaluations of F .
Our ﬁnal loss function is
[Equation]
where λ is a hyperparameter controlling the regularization strength.
This loss does not require us to encode a larger number of differently transformed inputs.
Instead, it uses the gradient at ( x , α , β , e ) to control properties of the encoder in a neighbourhood of α .
This can effectively reduce the representation gradient along the directions corresponding to many different transformations.
This, in turn, creates an encoder that is approximately invariant to the transformations.
B ETTER TEST TIME REPRESENTATIONS WITH FEATURE AVERAGING
At test time, standard practice (Hjelm et al., 2018; Kolesnikov et al., 2019) dictates that test representations be produced by applying the encoder to untransformed inputs (possibly using a central crop).
It may be beneﬁcial, however, to aggregate information from differently transformed versions of inputs to enforce invariance more directly, particularly when our previously introduced gradient regularization can only be applied to a subset of the transformation parameters.
Furthermore, in real-world applications, it may not be possible to remove nuisance transformations at test time or, as in our Spirograph dataset, there may not be only one unique ‘untransformed’ version of x .
To this end, we propose combining representations from different transformations using feature averaging.
This approach, akin to ensembling, does not directly use one encoding from the network f θ as a representation for an input x .
Instead, we sample transformation parameters α 1 , ..., α M ∼ p ( α ) , β 1 , ..., β M ∼ p ( β ) independently, and average the encodings of these differently transformed versions of x to give a single feature averaged representation
[Equation]
Using z ( M ) aggregates information about x by averaging over a range of possible transformations, thereby directly encouraging invariance.
Indeed, the resulting representation has lower conditional variance than the single-sample alternative, since
[Equation]
Further, unlike gradient regularization, this approach takes account of all transformations, including those which we cannot differentiate with respect to (e.g. left–right ﬂip).
It therefore forms a natural test time counterpart to our training methodology to promote invariance.
We do not recommend using feature averaged representations during training.
During training, we need a training signal to recognize similar and dissimilar representations, and feature averaging will weaken this training signal.
Furthermore, the computational cost of additional encoder passes is modest when used once at test time, but more signiﬁcant when used at every training iteration.
As a test time tool though, feature averaging is powerful.
In Theorem 1 below, we show that for certain downstream tasks the feature averaged representation will always perform better than the single-sample transformed alternative.
The proof is presented in Appendix C.
Theorem 1.
Consider evaluation on a downstream task by ﬁtting a linear classiﬁcation model with softmax loss or a linear regression model with square error loss on with representations as features.
For a ﬁxed classiﬁer or regressor and M (cid:48) ≥ M we have
[Equation]
Empirically we ﬁnd that, using the same encoder and the same linear classiﬁcation model, feature averaging can outperform evaluation using untransformed inputs.
That is, even when it is possible to remove the transformations at test time, it is beneﬁcial to retain them and use feature averaging.
R ELATED WORK
Contrastive learning (van den Oord et al., 2018; H´enaff et al., 2019) has progressively reﬁned the role of transformations in learning representations, with Bachman et al.
applying repeated data augmentation and Tian et al.
using Lab colour decomposition to deﬁne powerful selfsupervized tasks.
The range of transformations has progressively increased (Chen et al., 2020a;b), whilst changing transformations can markedly improve performance (Chen et al., 2020c).
Recent work has attempted to further understand and reﬁne the role of transformations (Tian et al., 2020).
The idea of differentiating with respect to transformation parameters dates back to the tangent propagation algorithm
(Simard et al., 1998; Rifai et al., 2011).
Using the notation of this paper, tangent propagation penalizes the norm of the gradient of a neural network evaluated at α = 0 , encouraging local transformation invariance near the original input.
In our work, we target the conditional variance (Equation 6), leading to gradient evaluations across the α parameter space with random α ∼ p ( α ) and a regularizer that is not a gradient norm (Equation 12).
Our gradient regularization approach also connects to work on gradient regularization for Lipschitz constraints.
A small Lipschitz constant has been shown to lead to better generalization (Sokoli ´ c et al., 2017) and improved adversarial robustness (Cisse et al., 2017; Tsuzuku et al., 2018; Barrett et al., 2021).
Previous work focuses on constraining the mapping x
(cid:55)→ z to have a small Lipschitz constant which is beneﬁcial for adversarial robustness. In our work we focus on the inﬂuence of α on z , which gives rise to transformation robustness.
Appendix D provides a more comprehensive discussion of related work.
E XPERIMENTS 6. D ATASETS AND SET - UP
The methods proposed in this paper learn representations that discard some information, whilst retaining what is relevant.
To more deeply explore this idea, we construct a dataset from a generative process controlled by both generative factors of interest and nuisance transformations .
Representations should be able to recover the factors of interest, whilst being approximately invariant to transformation.
To aid direct evaluation of this, we introduce a new dataset, which we refer to as the Spirograph dataset. Its samples are created using four generative factors and six nuisance transformation parameters.
Figure 1 shows two sets of four samples with the generative factors ﬁxed in each set.
Every Spirograph sample is based on a hypotrochoid—one of a parametric family of curves that describe the path traced out by a point on one circle rolling around inside another.
This generative process is fully differentiable in the parameters, meaning that our gradient regularization can be applied to every transformation.
We deﬁne four downstream tasks for this dataset, each corresponding to the recovery of one of the four generative factors of interest using linear regression.
The ﬁnal dataset consists of 100k training and 20k test images of size 32 × 32 .
For full details of this dataset, see Appendix E.
As well as the Spirograph dataset, we apply our ideas to CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009).
We base our contrastive learning set-up on SimCLR (Chen et al., 2020a).
To use our gradient regularization, we adapt colour distortion (brightness, contrast, saturation and hue adjustment) as a fully differentiable transformation giving a four dimensional α ; we also included random cropping and ﬂipping but we did not apply a gradient regularization to these.
We used ResNet50 (He et al., 2016) encoders for CIFAR and ResNet18 for Spirograph, and regularization parameters λ = 0 .
1 for CIFAR and λ = 0 .
01 for Spirograph.
For comprehensive details of our set-up and additional plots, see Appendix F. For an open source implementation of our methods, see https://github.com/ae-foster/invclr .
6. G RADIENT REGULARIZATION LEADS TO STRONGLY INVARIANT REPRESENTATIONS
We ﬁrst show that our gradient penalty successfuly learns representations that are more invariant to transformation than standard contrastive learning.
First, we estimate the conditional variance [Figure] of the representation that was used as the starting point for motivating our approach, i.e. Equation 6, using the slower, but more exact, nested Monte Carlo estimator of Equation 7 to evaluate this.
In Figurewe see that the gradient penalty strikingly reduces the conditional variance on CIFAR-10 compared to standard contrastive learning. [Table]
As an additional measure of representation invariance we ﬁt a linear regression model that predicts α from z , for which higher loss indicates a greater degree of invariance.
We also compute a reference loss: the loss that would be obtained when predicting α using only a constant.
In Table we see that unlike standard contrastive learning, after training with gradient regularization the linear regression model cannot predict α from z any better than using a constant prediction.
The loss is actually higher than the reference value because the former is obtained by training a regressor for a ﬁnite number of steps, whilst the latter is a theoretical optimum value.
Similar results for other datasets are in Appendix F.
6. G RADIENT REGULARIZATION FOR DOWNSTREAM TASKS AND TEST TIME DATASET SHIFT
We now show that these more invariant representations perform better on downstream tasks.
For
CIFAR, we produce representations for each element of the training and test set (by applying the encoder f θ to untransformed inputs).
We then ﬁt a linear classiﬁer on the training set, using different fractions of the class labels.
This allows us to assess our representations at different levels of supervision.
We use the entire test set to evaluate each of these classiﬁers.
In Figures 3(a) and (b), we see that the test accuracy improves across the board with gradient regularization.
For Spirograph, we take a similar approach to evaluation: we create representations for the training and test sets and ﬁt linear regression models with representations as features for each of the four downstream tasks.
In Figure 3(c), we see the test loss on each task with the baseline scaled to 1.
Here we see huge improvements across all tasks, presumably due to the ability to apply gradient regularization to all transformations (unlike for CIFAR).
We further study the effect of transformation at test time, showing that gradient penalized representations can be more robust to shifts in the transformation distribution.
For CIFAR-10, we apply colour distortion transformations at test time with different levels of variance.
By focusing on colour distortion at test time, we isolate the transformations that the gradient regularization targetted.
In
Figure 4(a)
we see that when the test time distribution is shifted to have higher variance than the training regime, our gradient penalized representations perform better than using contrastive learn- [Figure] [Figure] ing alone.
For Spirograph, we investigate changing both the mean of the transformation distribution, moving the entire test distribution away from the training regime and increasing the variance of transformations to add noise.
Results are shown in Figure 4(b) and (c).
In 4(c) in particular, we see that gradient regularized representations are robust to a greater level of distortion at test time.
6. F EATURE AVERAGING FURTHER IMPROVES PERFORMANCE
We now assess the impact of feature averaging on test time performance.
For CIFAR, we apply feature averaging using all transformations, including random crops etc.
and compare with the standard protocol of using untransformed inputs to form the test representations.
Figures 5(a) and (b) show that feature averaging leads to signiﬁcant improvements.
This adds to the result of Theorem 1, which implies that test loss decreases as M is increased.
In Figure 5(c), we see that feature averaging has an equally beneﬁcial impact on Spirograph.
It is interesting to note that in both cases there is still signiﬁcant residual beneﬁt from gradient regularization, even with a large value of M .
6. O UR METHODS COMPARE FAVOURABLY WITH OTHER PUBLISHED BASELINES
Our primary aim was to show that both gradient regularization and feature averaging lead to improvements compared to baselines that are in other respects identical.
Our methods are applicable to almost any base contrastive learning approach
[Table]
and we would expect them to deliver improvements across this range of different base methods.
In Table
[Table]
we present published baselines on
[Table]
[Figure] CIFAR datasets, along with the results that we obtain using our gradient regularization and feature averaging with SimCLR as a base method. This is the default base method that we recommend, and that was used in our previous experiments. Interestingly, the best ResNet50 encoder from our experiments achieves an accuracy of 94.9% on CIFAR-10, which outperforms the next best published result from the contrastive learning literature by almost 1%, and 75.1% on CIFAR-100, an almost 5% improvement over a signiﬁcantly larger encoder architecture.
As such, we see our results actually provide performance that is state-of-the-art for contrastive learning on these benchmarks.
In fact, our performance increases almost entirely close the gap to the state-of-the-art performance for fully supervized training with the same architecture on CIFAR-10 (95.1%, Chen et al.
(2020a)).
To demonstrate that our ideas generalize to other contrastive learning base methods, we apply our ideas to MoCo v2 (Chen et al., 2020c).
Table 3 shows that, whilst MoCo v2 itself does not perform as well as SimCLR on CIFAR-100, the addition of gradient regularization and feature averaging still leads to signiﬁcant improvements in its performance.
Table 3 further illustrates that both gradient regularization and feature averaging contribute to the performance improvements offered by our approach and that our techniques generalize across diffrent encoder architectures.
6. H YPERPARAMETER SENSITIVITY
As a further ablation study, we investigated the sensitivity of our method to changes in the gradient regularization hyperparameter λ (as deﬁned in Equation 13).
In Figure 6(a) we see that, as expected, the conditional variance of representations decreases as λ is increased.
The downstream task performance Figure 6(b) similarly improves as we increase λ , reaching an optimum around λ = 10 − 3 , before beginning to increase due to over-regularization.
We see that a wide range of values of λ deliver good performance and the method is not overly sensitive to careful tuning of λ .
C ONCLUSION
Viewing contrastive representation learning through the lens of representation invariance to transformation, we derived a gradient regularizer that controls how quickly representations can change with transformation, and proposed feature averaging at test time to pull in information from multiple transformations.
These approaches proved beneﬁcial on downstream tasks.
Our work therefore provides evidence that invariance is central to the success of contrastive learning methods
[Table]
and that there is scope to further improve upon these methods using invariance as a guiding principle.
[Table]
A CKNOWLEDGMENTS
AF gratefully acknowledges funding from EPSRC grant no.
EP/N509711/1.
AF would also like to thank Benjamin Bloem-Reddy for helpful discussions about theoretical aspects of this work.
A M UTUAL INFORMATION
In Section 2, we saw that the InfoNCE objective equation 3 fulﬁlls the need to make p ( z | x ) tightly focused on a single point whilst simultaneously requiring p ( z ) to be well spread out over representation space.
In this appendix, we show that the same general principle of making p ( z | x ) tightly focused on a single point whilst simultaneously requiring p ( z ) to be well spread out over representation space connects to mutual information maximization.
To establish the connection to mutual information, we take the differential entropy as our measure of ‘spread’.
Recall the differential entropy of a random variable w is
[Equation]
[Equation] We then translate our intuition to make p ( z | x ) tightly focused on a single point whilst simultaneously requiring p ( z ) to be well spread out over representation space into requiring E p ( x ) [ H [ p ( z | x )]] to be minimized whilst H [ p ( z )] should be simultaneously maximized.
This suggests the following loss function
[Equation] [Equation]
which is the (negative) mutual information between x and z .
Note that in this formulation, it is the distribution p ( z | x ) as much as the InfoMax principle which determines how this loss will behave.
Finally, there is a clear connection between the InfoNCE loss and mutual information, speciﬁcally the InfoNCE loss is, in expectation and up to an additive constant, a lower bound on I ( x ; z ) van den
Oord et al.
; Poole et al.
B M ETHOD B. A N ALTERNATIVE VARIANCE FORMULA
We present a derivation of our alternative formula for the variance (dropping the conditioning from the notation for conciseness)
[Equation]
e )) 2 (cid:3) = 12 E p ( α ) p ( α (cid:48) ) (cid:2) ( F ( α , β , x , e ) − E α
[ F ( α , β , x , e )] + E α
[ F ( α , β , x , e )] − F ( α (cid:48) , β , x , e )) 2 (cid:3) = 12 E p ( α ) p ( α (cid:48) ) (cid:2) ( F ( α , β , x , e ) − E α[-[ F b ( α , β , x )]) 
-]
[ F ( α , β , x , e )]) 2 + ( E α [ F ( α , β , x , e )] − F ( α (cid:48) , β , x , e )) 2 (cid:3) + E p ( α ) p ( α (cid:48) )
[( F ( α , β , x , e ) − E α [ F ( α , β , x , e )])( E α [ F ( α , β , x , e )] − F ( α (cid:48) , β , x , e ))]
= 12 E p ( α ) p ( α (cid:48) ) (cid:2) ( F ( α , β , x , e ) − E α
[ F ( α , β , x , e )])
+ ( F ( α (cid:48) , β , x , e )]
− E α [ F ( α , β , x , e )]) 2 (cid:3) =
Var p ( α ) [ F ( α , β , x , e )] .
B. M OTIVATING THE R ADEMACHER DISTRIBUTION
We are interested in the conditional variance of z with respect to α , but as z is a vector valued random variable we properly need to consider the conditional covariance matrix Σ
=
Cov α ( z | x , β ) .
We henceforth consider x , β to be ﬁxed.
To reduce conditional variance in all directions, it makes sense to reduce the trace Tr Σ .
Due to computational limitations, we cannot directly estimate this trace at each iteration, instead we must estimate Var( e · z ) =
e (cid:62) Σ e .
However, by carefully selecting the distribution for e we can effectively target the trace of the covariance matrix by taking the expectation over e .
Speciﬁcally, suppose that the components of e are independent Rademacher random variables ( ± 1 with equal probability).
Then
[Equation]
C T HEORY
We present the proof of Theorem 1 which is restated for convenience.
Theorem 1.
Consider evaluation on a downstream task by ﬁtting a linear classiﬁcation model with softmax loss or a linear regression model with square error loss on with representations as features.
For a ﬁxed classiﬁer or regressor and M (cid:48) ≥ M we have
[Equation]
Proof.
We have the softmax loss
[Equation]
or the square error loss
[Equation]
We ﬁrst show that both loss functions considered are convex in the argument z .
To show this [Equation] we ﬁx 0 ≤ p = 1 − q ≤ 1 .
For softmax loss, we have
[Equation]
by H ¨ older’s
Inequality
[Equation]
and for square error loss we have
[Equation]
For the inequality in the Theorem
[Equation] we consider drawing M [Equation] ≥ M samples, and randomly choosing an M -subset.
Let S represent this subset and let z ( M ) S represent the feature averaged representation that uses the subset S .
We have
[Equation]
where the inequality at equation 34 is by Jensen’s Inequality.
This completes the proof.
We provide an informal discussion of other theoretical results that may relate to our work.
Lyle et al.
(2020) explored PAC-Bayesian approaches to analyzing the role of group invariance in generalization of supervized neural network models.
The central bound based on Catoni (2007) is given in Theorem 1 of Lyle et al.
and depends on the empirical risk
ˆ R (cid:96) ( Q, D n ) and the KL term
KL ( Q || P ) which represents the PAC-Bayesian KL divergence between distributions on hypothesis space.
Theorem 7 of Lyle et al.
(2020) shows KL ( Q ◦ || P ◦ ) ≤ KL ( Q || P ) , where Q ◦ and P ◦ are formed by symmetrization such as feature averaging over the group of transformations.
In our context, although the transformations do not form a group, we could still consider a symmetrization operation with feature averaging.
If the symmetrization does not affect the empirical risk, then
Theorem 9 of Lyle et al.
(2020) would apply to our setting and we would be able to obtain a tighter generalization bound for our suggested approach of feature averaging.
D R ELATED WORK D. T HE ROLE OF TRANSFORMATIONS IN CONTRASTIVE LEARNING
Recent work on contrastive learning, initiated by the development of Contrastive Predictive Coding (van den Oord et al., 2018; H´enaff et al., 2019), has progressively moved the transformations to a more central position in understanding and improving these approaches.
In Bachman et al.
(2019), multiple views of a context are extracted, on images this utilizes repeated data augmentation such as random resized crop, random colour jitter, and random conversion to grayscale, and the model is trained to maximize information between these views using an InfoNCE style objective.
Other approaches are possible, for instance Tian et al.
(2019) obtained multiple views of images using
Lab colour decomposition.
In SimCLR (Chen et al., 2020a;b), the approach of applying multiple data augmentations (including ﬂip and blur, as well as crops, colour jitter and random grayscale) and using an InfoNCE objective was simpliﬁed and streamlined, and the central role of the augmentations was emphasized.
By changing the set of transformation operations used, (Chen et al., 2020c) were able to improve their contrastive learning approach and achieve excellent performance on downstream detection and segmentation tasks.
Tian et al.
(2020) has studied what the best range of transformations for contrastive learning is.
The authors found that there is a ‘sweet spot’ in the strength of transformations applied in contrastive learning, with transformations that are too strong or weak being less favourable.
Winkens et al.
(2020) showed that contrastive methods can be successfully applied to out-of-distribution detection.
We note that for tasks such as out-of-distribution detection, transformation covariance may be a more relevant property than invariance.
D. G RADIENT REGULARIZATION TO ENFORCE L IPSCHITZ CONSTRAINTS
Constraining a neural network to be Lipschitz continuous bounds how quickly its output can change as the input changes.
In supervized learning, a small Lipschitz constant has been shown to lead to better generalization (Sokoli ´ c et al., 2017) and improved adversarial robustness (Cisse et al., 2017;
Tsuzuku et al., 2018).
One practical method for constraining the Lipschitz constant is gradient regularization (Drucker & Le Cun, 1992; Gulrajani et al., 2017).
Lipschitz constraints have also been applied in a self-supervized context: in Ozair et al.
(2019), the authors used a Wasserstein dependency measure in a contrastive learning setting by using gradient penalization to ensure that the function x , x (cid:48) (cid:55)→ s φ ( f θ ( x ) , f θ ( x (cid:48) )) is 1-Lipschitz.
Our work uses a gradient regularizer to control how quickly representations can change, but unlike existing work we focus on how representations change with α as x is ﬁxed, instead of how they change with x .
D. G ROUP INVARIANT NEURAL NETWORKS
A large body of recent work has focused on designing neural network architectures that are perfectly invariant, or equivariant, to a set of transformations T in the case when T forms a group.
Cohen &
Welling (2016) showed how convolutional neural networks can be generalized to have equivariance to arbitrary group transformations applied to their inputs.
This can apply, for instance, to rotation groups on the sphere (Cohen et al., 2018), rotation and translation groups on point clouds Thomas et al.
(2018), and permutation groups on sets (Zaheer et al., 2017).
Group transformations are a powerful theoretical framework.
Transformations that form a group cannot remove information from the input (because they must be invertible) and can be composed in any order.
This means that the more general transformations considered in our work cannot form a group—they cannot be composed (repeated decreasing of brightness to zero is not allowed) or inverted (crops are not
[Figure]
f g
[Equation]
=
[Equation]
[Equation] [Equation] [Equation] [Equation]
invertible).
We have therefore considered methods that offer aspects of invariance under much more general transformations.
E S PIROGRAPH DATASET
We propose a new dataset that allows the separation of generative factors of interest from nuisance transformation factors and that is formed from a fully differentiable generative process.
A standalone implementation of this dataset can be found at https://github.com/rattaoup/ spirograph .
Our dataset is inspired by the beautiful spirograph patterns some of us drew as children, which are mathematically hypotrochoids given by the following equations
[Equation]
Figure 7(a) shows an example.
To create an image dataset from such curves, we choose 40 equally spaced points t i with t 1 = 0 and t 40 = 2 π , giving a sequence of points ( x 1 , y 1 ) , ..., ( x 40 , y 40 ) on the chosen hypotrochoid.
For smoothing parameter σ , the pixel intensity at a point ( u, v ) is given by
[Equation]
For a grid of pixel, the intensity values are normalized so that the maximum intensity is equal to 1.
Figure 7(b) shows the pixel intensity with σ = 0 .
Finally, for a foreground colour with RGB values ( f r , f g , f b ) and background colour ( b r , b g , b b ) , the ﬁnal RGB values at a point ( u, v ) is
[Equation]
The ﬁnal coloured sample image is shown in Figure 7(c).
The Spirograph sample is fully speciﬁed by the parameters m, b, h, σ, f r , f g , f b , b r , b g , b b .
In our experiments, we treat m, b, σ, f r as parameters of interest.
We treat h and the remaining colour parameters as nuisance parameters.
That is, we take x = ( m, b, σ, f r ) and α = ( h, f g , f b , b r , b g , b b ) and the transformation t ( x , α ) is the full generative process described above.
There are no additional parameters β for this dataset.
Figure 1 shows two sets of four samples from the Spirograph dataset, in each set the generative factors of interest are ﬁxed and the nuisance parameters are varied.
In general for the Spirograph dataset, the distinction between generative factors of interest and nuisance parameters can be changed to attempt to learn different aspects of the data.
The transformation t is fully differentiable, meaning that we can apply gradient penalization to all the nuisance parameters of the generative process.
In our experiments, we took the following distributions to sample random values of the parameters: m ∼ U (2 , 5) , b ∼ U (0 . 1 , 1 . 1) , h ∼ U (0 . 5 , 2 . 5) , σ ∼
U (0 . 25 , 1) , f r , f g , f b ∼ U (0 . 4 , 1) , b r , b g , b b ∼ U (0 , 0 . 6) .
We synthesized 100,000 training images and 20,000 test images with dimension 32 × 32 .
F E XPERIMENT DETAILS
Our experiments were implemented in PyTorch (Paszke et al., 2019) and ran on 8 Nvidia GeForce
GTX 1080 Ti GPUs.
See https://github.com/ae-foster/invclr for an implementation of our approaches.
F. D IFFERENTIABLE COLOUR DISTORTION
We want to improve the representations learned from contrastive methods by explicitly encouraging stronger invariance to the set of transformations.
Our method is to restrict gradients of the representations with respect to certain transformations.
Ensuring that the transformations are practically differentiable within PyTorch (Paszke et al., 2019) required a thorough study of the transformations.
The subset of transformations we apply gradient regularization to includes colour distortions which are conventionally treated as a part of data preprocessing.
Rewriting this as a differentiable module within the computational graph allows us to practically compute the gradient regularizer of equation 11.
We will consider adjusting brightness, contrast, saturation, hue of an image.
In fact, most of these transformations are simply linear transformations of the original image.
First, the brightness adjustment is simply deﬁned as
[Equation]
(40) when α brt is a scale factor.
If we write x = r , g , b , for the three colour channels of x , then greyscale conversion on x is
[Equation]
Next, we can see adjusting saturation of x as a linear combination of x and x gs , the greyscale version of x
[Equation]
when α sat is a scale factor.
Adjusting contrast of x is also a linear combination of x and mean ( x gs ) , which the mean over all spatial dimensions of x gs .
With a scaling parameter α con we have
[Equation]
We utilize a linear approximation for hue adjustment.
We perform hue adjustment by converting to the YIQ colour space, and then applying rotation on the IQ components.
The transformation between RGB and YIQ colour space is given by the following linear transformation
[Equation]
Note that the Y component is exactly the greyscale version x gs deﬁned above.
We transform YIQ back to RGB by (cid:32)
[Equation]
In YIQ format, we can adjust hue of an image by θ = 2 πα hue by multiplying with a rotation matrix
[Equation] [Equation]
Therefore, our hue adjustment is given by
[Equation]
[Equation]
where the matrices operate on the three colour channels of x and in parallel over all spatial dimensions.
[Table]
U (0 . 6 , 1 . 4)
[Table]
U (0 . 6 , 1 . 4)
[Table]
U (0 . 6 , 1 . 4) [Table] U ( − 0 . [Table]
F. S ET - UP
Our set-up is quite similar to the setup in Chen et al.
(2020a) with two main differences: we treat colour distortions as a differentiable module while in Chen et al.
(2020a)
the transformation was performed in the preprocessing step, we add the gradient penalty term in addition to the original loss in Chen et al.
(2020a).
F.2. T RANSFORMATIONS
First, for a batch x 1 , ..., x K of inputs, we form a pair of ( x 1 , x (cid:48) 1 ) , ..., ( x K , x (cid:48) K ) by applying two random transformations: random resized crop and random horizontal ﬂip for each input.
We then apply our differentiable colour distortion function which is composed of random colour jitter with probability p = 0 .
8 and random greyscale with probability p = 0 .
(Colour jitter is the composition of adjusting brightness, adjusting contrast, adjusting saturation, adjusting hue in this order.)
We sample α that controls how strong the adjustment is for each image from the distribution that is shown in Table 4.
We call the resultant pairs ( x 1 , x (cid:48) 1 ) , ..., ( x K , x (cid:48) K ) .
F.2. C ONTRASTIVE LEARNING
Similar to Chen et al.
(2020a), we use the transformed ( x 1 , x (cid:48) 1 ) , ..., ( x K , x (cid:48) K ) as an input to an encoder to learn a pair of representations ( z 1 , z (cid:48) 1 ) , ..., ( z K , z (cid:48) K ) .
The ﬁnal loss function that we use for training is equation 13.
Table 5 shows all hyperparameters that were used for training.
The small neural network g φ is a MLP with the two layers consisting of a fully connected linear map, ReLU activation and batch normalization.
We use LARS optimizer (You et al., 2017) and apply cosine annealing to the learning rate.
F.2. G RADIENT REGULARIZATION
In this part, we explain our setup for calculating the gradient penalty as in equation 12.
We sample a random vector e with independent Rademacher components and independently for each sample in the batch.
We generate L samples of α for each element of the batch to compute the regularizer.
Finally, we clip the penalty from above to prevent instability at the onset of training.
In practice, this meant the gradient regularization was not enforced for about the ﬁrst epoch of training.
Tableshows hyperparameters that we used within gradient penalty calculation.
[Table]
[Table]
F
[Table]
E VALUATION
We use our representations as features in linear classiﬁcation and regression tasks.
We train these linear models with L-BFGS with hyperparameters as shown in Table 7 on the train set and evaluate performance on the test set
. F.2. M O C O V
To empirically demonstrate that our ideas transfer to alternative base contrastive learning methods, we applied both gradient regularization and feature averaging to the MoCo v2 (Chen et al., 2020c) base set-up.
We also explored two different ResNet (He et al., 2016) architectures.
We closely followed the MoCo v2 implementation at https://github.com/facebookresearch/moco .
As for SimCLR, we adapted the transformations to be a differentiable module.
We also made adaptations for CIFAR-100 in an identical way as in our previous experiments.
As in MoCo v2, we removed batch normalization in the projection head g φ ; we used SGD optimization with learning rate 0 .
06 for a batch size of 512, and used the MoCo parameters K = 2048 and m = 0 .
99 for
ResNet18 and K = 4096 , m = 0 .
99 for ResNet50.
We did not conduct extensive hyperparameter sweeps, but we did investigate larger values of K which did not lead to improved performance on
CIFAR-100.
(In particular, the original settings K = 65536 , m = 0 .
999 appeared to perform less well on this dataset.)
Other hyperparameters and settings were identical to Chen et al.
(2020c).
We did 3 independent runs with a ResNet18 and 2 runs with a ResNet50.
We conducted linear classiﬁcation evaluation with ﬁxed representations in exactly the same way as for our other experiments.
Feature averaging results used M = 40 .
F.2. C OMPUTATIONAL COST
We found that gradient regularization increased the total time to train encoders by a factor of at most 2 .
For feature averaging at test time with a ﬁxed dataset, the computation of features z ( M ) is an O ( M ) operation, whilst the training and testing of the linear classiﬁer is O (1) .
Training time remained by far the larger in all experiments by orders of magnitude.
F. A DDITIONAL EXPERIMENTAL RESULTS F.3. C OMPARISON WITH ENSEMBLING
Feature averaging is an approach that bears much similarity with ensembling.
To experimentally compare these two approaches, we applied both approaches to encoders trained on CIFAR-10.
To provide a suitable comparison with feature averaging using z ( M ) we ﬁrst a trained a linear classiﬁer p ( y
| z ) using an M -fold augmented dataset of representations with a standard cross-entropy loss using L-BFGS optimization using the same weight decay as for feature averaging.
For CIFAR-10, which has a training set of length 50000 , the feature averaging classiﬁer was trained usingaveraged representations, whereas the ensembling classiﬁer was trained with 50000 M examples using data augmentation.
At test time, we averaged prediction probabilities using M different rep-
[Figure]
[Figure]
resentations of each test image.
Speciﬁcally, if p ( y | z ) is the classiﬁer trained by the aforementioned procedure and α 1 , ..., α M ∼ p ( α ) , β 1 , ..., β M ∼ p ( β ) are independent transformation parameters, the probability of assigning label y to input x is given by
[Equation]
The results outlined in Figure 8 show that ensembling gives very similar performance to feature averaging in terms of accuracy, but is signiﬁcantly worse in terms of loss.
We can understand this result intuitively because ensembling includes probabilities from every transformed version of the input (including where the classiﬁer is uncertain or incorrect) whereas feature averaging combines transformations in representation space and uses only one forward pass of the classiﬁer.
More formally, the difference in test loss makes sense in light of Theorem 1.
Figure 9 shows additional results obtained using representations trained with standard SimCLR on CIFAR-10.
We see the same pattern—a similar test accuracy but worse test loss when using augmentation ensembling.
F.3. G RADIENT REGULARIZATION LEADS TO STRONGLY INVARIANT REPRESENTATIONS
We ﬁrst show that our gradient penalty successfully learns representations that have greater invariance to transformation than their counterparts generated by contrastive learning.
We consider two metrics: the conditional variance targetted directly by the gradient regularizer, and the loss when z is used to predict α with linear regression.
Table 8 and Figure 10 are the equivalents of Table 1 and Figurefor CIFAR-100, showing the conditional variance and the regression loss for predicting α respectively.
In Table 9 we present the same results for Spirograph.
We see very similar results to CIFARin both cases—the gradient penalty dramatically reduces conditional variances, and prediction of α by linear regression gives a loss that is better than a constant prediction only for standard contrastive representations.
[Table] [Table]
F.3. G RADIENT REGULARIZED REPRESENTATIONS PERFORM BETTER ON DOWNSTREAM TASKS AND ARE ROBUST TO TEST TIME TRANSFORMATION
For downstream performance on Spirograph, we evaluate the performance of encoders trained with gradient regularization and without gradient regularization on the task of predicting the generative parameters of interest.
In our set-up [Table] we use 100,000 train images and 20,000 test images and train the encoders on the training set for 50 epochs.
For evaluation, we train a linear regressor on the representations from encoders to predict the actual generative parameters.
Setting for the linear regressor is shown in the Table 7.
To accompany the main results in Figure 3(c), we include the exact values used in this ﬁgure in Table 10.
We now turn to our experiments used to investigate robustness—we investigate scenarios when we change the distribution of transformation parameters α at test time [Table] but use encoders that were trained with the original distribution.
We investigate on both CIFAR and Spirograph datasets.
For CIFAR, we chose to vary the distribution of parameters for colour distortions at test time.
We could write the distribution of parameter of brightness, saturation, contrast as U (1 − 0 .
8 S, 1 + 0 .
8 S ) and the distribution of hue as U ( − 0 .
2 S, 0 .
2 S ) where S is a parameter controlling the strength of the distortion.
In the original setup as in Table 4, we have S = 0 .
By varying the value of S used at test time, we can increase the variance of the nuisance transformations, including stronger transformations than those that were present when encoders were trained.
This is visualized in Figure 12.
Figure 11 is a companion plot for Figure 4(a) applied on CIFAR-100.
We see broadly similar trends—our representations outperform those from standard contrastive learning across a range of test time distributions. [Table] [Figure] [Figure]
[Figure]
For robustness on Spirograph, recall α =
( h, f g , f b , b r , b g , b b ) when h ∼ U (0 . 5 , 2 . 5) , f g , f b ∼
U (0 . 4 , 1) , b r , b g , b b ∼ U (0 , 0 . 6) .
We chose to vary the distribution of the background colour ( b r , b g , b b ) and the distribution of h which is a structure-related transformation parameter.
We consider two approaches, mean shifting where we shift the uniform distribution by S , for example
U ( a, b ) → U ( a + S, b + S ) and we consider changing variance where we increase the range of the uniform distribution by 2 S , for example, U ( a, b ) → U ( a − S, b + S ) .
We compare performance of the trained encoders at epoch 50 on predicting the generative parameters x
= ( m, b, σ, f r ) and we use the same setting for linear regressors as in Table 7.
Figure 13 is a visualization of the effect of varying h from 0 .
5 to 2 .
while other parameters are kept constant.
The ﬁgure 14 shows the effect of varying the background colour of an image by adding
S = 0 .
45 to each of the background RGB channels.
For varying the distribution of h , we consider shifting the mean of h ∼ U (0 . 5 , 2 . 5) by S = ± 0 .
1 , ± 0 .
3 , ± 0 .
5 and increasing the variance of h by S = 0 .
For the distribution of the background colour ( b r , b g , b b ) , we consider shifting the distribution of ( b r , b g , b b ) by S = 0 .
4 and increasing variance by the same amount.
We note that ( b r , b g , b b ) controls the background colour of an image, so we are varying the 3 distributions at the same time.
Since, the foreground colour f r , f g , f b ∼ U (0 . 4 , 1) , we are shifting the distribution of ( b r , b g , b b ) toward ( f r , f g , f b ) and will make background and foreground colour more similar.
For example, with S = 0 . 4 , when we apply a mean shift we change the distribution of ( b r , b g , b b ) to b r , b g , b b ∼ U (0 . 4 , 1) , and when we increase the variance the distribution becomes b r , b g , b b ∼ U (0 , 1) .