I MPROVING T RANSFORMATION I NVARIANCE IN C ONTRASTIVE R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
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 directly 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 endeavor 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 formed as 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.
Unfortunately, directly changing the similarity measure hampers the algorithm (Chen et al., 2020a; Wu et al., 2018).
We therefore investigate methods to improve contrastive representations by explicitly encouraging stronger invariance to the set of transformations, without changing the core self–supervized 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.
Empirically, we ﬁrst evaluate our approaches on CIFAR (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 the Spirograph dataset.
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 both using our regularizer during training and our feature averaging at test time 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 natural 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 computer vision tasks, it is typical to include resized cropping; brightness, contrast, saturation and hue distortion; greyscale conversion, and horizontal ﬂip.
In Section 6, we introduce the Spirograph dataset which uses quite different transformations.
In general, the transformations are assumed to act in a way that changes the input only cosmetically, so all semantic features such as the class label are preserved.
The transformations indicate 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 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 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 inputs are mapped to the same point.
To preserve meaningful information about the input x whilst discarding purely cosmetic features, we should require p ( z | x ) to be focused on a single point 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 independent given x i ,
[Equation]
so z i , z (cid:48)
i =
f θ ( t ( x )) , f θ ( t (cid:48)
( x )) where t , t (cid:48) ∼ p ( t ) are independent.
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 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, because the function g φ is not an injection.
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 encourages invariance by maximizing s φ ( z , z (cid:48) ) , but this does not provide a strong signal to learn invariance.
Our idea is that we can use gradient methods to regulate how the representation changes with the transformation.
Our key proposal 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 idea, 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 )) , and t is 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 makes 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 and deﬁne
[Equation]
which is the projection of the normalized representation z / (cid:107) z (cid:107) in the e direction.
To encourage an encoder that is invariant to changes in α , we would like to minimize, for a given x , the expected variance of F e as α is changed over different directions e .
That is, we wish to ensure that the conditional variance
[Equation]
is minimized.
A naive estimator of this quantity would be a direct nested Monte Carlo estimator (Rainforth et al., 2018) of sample variances
[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)
[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 gradient ∇ α z , forcing the representation to change slowly as the transformation is varied and thus inducing approximate invariance to the transformations [Equation]
In practice equation 11 is estimated using a batch x 1 , ..., x K as [Equation]
[Equation]
where x i , α i , β i , e i ,
∼ p
( x ) p ( α ) p ( β ) p ( e ) and α (cid:48) ij ∼ p
( α ) .
Note that here 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 e .
Our ﬁnal loss function is
[Equation]
[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 , α , β ) 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 introduce 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 the Spirograph dataset we introduce in Section 6, there may not be only one unique ‘untransformed’ version of x .
To this end, we introduce the idea of 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 , ..., β K ∼ 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 penalization, 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. 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.
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 for M ≤ M [Equation] we have
[Equation]
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 attempts to further understand and reﬁne the role of transformations (Tian et al., 2020).
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).
In our notation, previous work focuses on constraining x
(cid:55)→ z to have small Lipschitz constant which is beneﬁcial for adversarial robustness, in our work we focus α (cid:55)→ 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 4 generative factors and 6 nuisance transformations.
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 on this dataset, see
Appendix E, while code to generate it is provided in the supplement.
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.
6. G RADIENT REGULARIZATION LEADS TO STRONGLY INVARIANT REPRESENTATIONS
[Table]
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 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 Figure 2 we see that the gradient penalty strikingly reduces the conditional variance on CIFAR-10 compared to standard contrastive learning.
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 1, 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 classiﬁer for a ﬁnite number of steps, whilst the latter is a theoretical optimum value.)
Similar results for our other datasets are in Appendix F.
[Figure]
[Figure]
6. G RADIENT REGULARIZED REPRESENTATIONS PERFORM BETTER ON DOWNSTREAM TASKS AND ARE MORE ROBUST TO TEST TIME DATASET SHIFT
We now show that these more invariant representations perform better on downstream tasks.
For
CIFAR-10 and CIFAR-100, 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 Figure 3(a) and (b), we see that the test accuracy improves across the board with gradient regularization with about a 1% improvement on CIFAR-100 at all levels of semi-supervision.
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 1% better than using contrastive learning 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.
[Figure]
[Table]
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 [Table] 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)
[Table]
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.
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 supervised training with the same architecture on this dataset (95.1%, Chen et al.
(2020a)).
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 λ . 6. C OMPARISON OF FEATURE AVERAGING WITH ENSEMBLING We also provide an additional comparison between feature averaging and augmentation ensembling.
Here, predictions from multiple differently transformed versions of the same input are combined by
[Figure]
[Figure]
averaging prediction probabilities p ( y | x )
= 1 M (cid:80) Mm =1 p ( y | f θ ( t ( x , α m , β m ))) where p ( y | z ) is a linear classiﬁer and α m , β m are as in Section 4.
In Figure 7, we see that ensembling performs almost as well as feature averaging in terms of accuracy, but worse in terms of test loss.
6. G ENERALIZATION TO OTHER BASE METHODS
As previously mentioned, our approaches can be applied to a wide variety of contrastive learning algorithms.
To demonstrate this, we apply our ideas to MoCo v2 (Chen et al., 2020c).
We present results from a direct implementation, and when using gradient regularization and feature averaging.
Table 3 shows that, whilst MoCo v2 itself does not perform as well as SimCLR for this dataset, 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.
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.
We believe that the idea of constraining the gradient of representations with respect to transformation parameters is a fresh perspective on approximate invariance.
Less restricted and easier to implement than many other notions of invariance, we believe it has potential for other applications such as supervized learning. [Table] [Table]
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]
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
[Equation]
( α ) p ( α (cid:48) )
(cid:2) ( F b ( α , β , x ) − E
α
[ F b ( α , β , x )]
+ E α
[ F b ( α , β , x )] − F b ( α (cid:48) , β , x )) 2 (cid:3) = 12 E p ( α ) p ( α (cid:48) ) (cid:2) ( F b ( α , β , x ) − E α [ F b ( α , β , x )])
+ ( E α [ F b ( α , β , x )] − F b ( α (cid:48) , β , x )) 2 (cid:3) + E p ( α ) p ( α (cid:48) )
[( F b ( α , β , x ) − E α [ F b ( α , β , x )])
( E α
[ F b ( α , β , x )] − F b ( α (cid:48) , β , x ))]
= 12 E p ( α ) p ( α (cid:48) ) (cid:2) ( F b ( α , β , x ) − E α
[ F b ( α , β , x )])
+ ( F b ( α (cid:48) , β , x )] − E α
[ F b ( α , β , x )]) 2 (cid:3) =
Var p ( α ) [ F b ( α , β , x )] .
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
[Equation]
and for M ≤ M (cid:48) we have
[Equation]
Proof.
We have the softmax loss
[Equation]
or the square error loss
[Equation] [Equation]
We ﬁrst show that both loss functions considered are convex in the argument z .
To show this, 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, we consider drawing M (cid:48) ≥ 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 37 is by Jensen’s Inequality.
This completes the proof
[Equation]
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.
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 f b ) =
[Equation]
r b
[Equation]
=
[Equation]
[Equation] [Equation]
0 . 5) 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.
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 8(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]
[Equation]
For a grid of pixel, the intensity values are normalized so that the maximum intensity is equal to 1.
Figure 8(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 8(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 .
The the code used to generated the dataset is included in the supplement.
F E XPERIMENT DETAILS
Our experiments were implemented in PyTorch (Paszke et al., 2019) and ran on 8 Nvidia GeForce
GTX 1080 Ti GPUs.
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]
43) 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]
(45) 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
[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]
where the matrices operate on the three colour channels of x and in parallel over all spatial dimensions.
Each operation is followed by pointwise clipping of pixel values to the range [0 , 1] .
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 the table 4.
We call the resultant pairs ( x 1 , x (cid:48) 1 ) , ..., ( x K , x (cid:48) K ) .
Transformation
Distribution
Adjust brightness
U (0 . 6 , 1 . 4)
Adjust contrast
U (0 . 6 , 1 . 4)
Adjust saturation
U (0 . 6 , 1 . 4)
Adjust hue
U ( − 0 .
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.
[Table]
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]
F.2. 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 the table 7 on the train set and evaluate performance on the test set.
[Table]
F. A DDITIONAL EXPERIMENTAL RESULTS 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.
Figure 9 and Table 8 are the equivalents of Figureand Table 1 for 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 CIFAR-10 in 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.
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, 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.
[Table]
[Table]
We now turn to our experiments used to investigate robustness—we investigate scenarios when we change the distribution of transformation parameters α at test time, 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 the 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 10.
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.
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 con-
[Figure]
[Table]
[Figure] [Figure]
sider 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 12 is a visualization of the effect of varying h from 0 .
5 to 2 .
while other parameters are kept constant.
The ﬁgure 13 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) .