Improving Self-Supervised Learning by Characterizing Idealized Representations
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Despite the empirical successes of self-supervised learning (SSL) methods, itis unclear what characteristics of their representations lead to high downstreamaccuracies.
In this work, we characterize properties that SSL representations shouldideally satisfy.
Speciﬁcally, we prove necessary and sufﬁcient conditions suchthat for any task invariant to given data augmentations, desired probes (e.g., linearor MLP) trained on that representation attain perfect accuracy.
These requirementslead to a unifying conceptual framework for improving existing SSL methods andderiving new ones.
For contrastive learning, our framework prescribes simple butsigniﬁcant improvements to previous methods such as using asymmetric projectionheads.
For non-contrastive learning, we use our framework to derive a simple andnovel objective.
Our resulting SSL algorithms outperform baselines on standardbenchmarks, including SwAV+multicrops on linear probing of ImageNet.
Introduction
We study self-supervised learning (SSL), where the goal is to learn representations from minimalsupervision, such that simple probes trained on these representations achieve high downstreamaccuracy.
Recently, there has been many different SSL methods achieving impressive empirical results(e.g.
SimCLR [1], SwAV [2]) using label-preserving augmentations (e g. cropping or color jittering)as supervision.
We dub this setting invariant SSL (ISSL).
Despite these empirical successes, it remainsunclear how these various SSL methods relate to one another, how to improve them and how to derivenew ones.
Our goal is to provide a simple conceptual framework to think about those questions.
To derive such a framework, we ask ourselves: what are the ideal requirements that ISSL representations should aim to satisfy?
We prove necessary and sufﬁcient requirements to ensure that probes froma speciﬁed family, e.g. linear or multi-layer perceptron (MLP), perfectly classify any task that is invariant to desired data augmentations.
This complements theoretical work in ISSL [3–8], which analyzespeciﬁc ISSL algorithms.
Our work instead focuses on properties of representations that should serveas a goal for any ISSL algorithm.
These ideal properties are: (i) desired predictors should be able todistinguish positive and negative examples from the representation; (ii) the dimensionality of the representation should be sufﬁciently large; (iii) augmented inputs should map to the same representations.
Previous ISSL methods can be seen as approximations of our ideal requirements.
Using our requirements, we derive simple improvements to those approximations as well as new ISSL objectives.
Ourtheory thus results in a unifying conceptual ISSL framework, with practical prescriptions including:
• improvements to existing methods, such as increasing the dimensionality of representations and
using asymmetric projections heads, which lead to 5% point gains on TinyImageNet;
• a novel non-contrastive ISSL objective that outperforms all baselines, includingSwAV+multicrops, by at least 1% point on linear classiﬁcation of ImageNet;• extensions of ISSL algorithms to learn representations that are better suited for non-linear probes.
Problem statement
[Figure]
2.1 Invariant
[Figure]
learns an interest t (
Y )
Supervision for ISSL comes from unlabeled data p X and label-preserving augmentations.
Augmentations are ways of sampling “positive” x, x + examples that are equivalent for downstream classiﬁcationtasks.
We formalize this using an equivalence relation x ∼ x + that partitions the inputs X intoequivalent classes [ x ] ∈ X / ∼ , and we consider the following downstream tasks T whose labelingsare deterministic and constant within these classes (we allow stochastic labeling in appendices).
Deﬁnition 1.
The ∼ - invariant tasks T , is the set of all input-label distributions p t
( X, Y ) such thatthe labeling p t
( Y | X ) is deterministic and invariant to ∼ , i.e.,for all p t ∈ T , x, x + [Equation]
[Equation]
As an illustration, consider the 3 classes (triangle, square, and circle) shown in Fig.
Then T consistsof all tasks that are predictable from those shapes, e.g., recognizing shapes with vertices (blue/orangein Fig. 1) or recognizing the shape (yellow/red/purple in Fig. 1).
Importantly, equivalence classes (hereshapes) are different from—and essentially reﬁnements of—downstream classes
Y (here colors). Notethat equivalence relations can model arbitrary transformations including cropping and adding Gaussian noise, which contrasts with typical restrictions to augmentations deﬁned by group actions[9–11].
2.2 Idealized representations for ISSL
In this section, we deﬁne the optimal encoders, those that induce idealized representations that ISSLshould be striving for.
Although such encoders exist, they will likely not be learned in practice.
Thoseidealized representations will nevertheless allow us to derive practical algorithms in Sec. 4.
The goal of ISSL is to learn representations from which (typically simple) probes f ∈ F classify welldownstream tasks p t ∈ T , i.e., they achieve low 0-1 risk R t ( φ, f ) :=
E p t ( X,Y ) [ 1 [ Y (cid:54) =
f
( φ ( X )))]] .
Ideally, for any task of interest p t ∈ T there will be a probe that can classify it perfectly.
Deﬁnition 2.
An encoder φ is population optimal for T , F , denoted as φ ∈ Φ pop , iff predictorsrealize the Bayes error on any invariant task, i.e., for all p t ∈ T we have inf f ∈F R t ( φ, f ) = 0 .
Population optimality ensures the existence of perfect downstream probes, which would be learnedwith inﬁnite downstream data.
In practice, however, predictors will be trained from ﬁnite datasets
D t of possibly small size n with empirical risk minimization.
When n is small a ﬁtted probe (ERM)ˆ f ∈ (cid:98) F ( D t , φ ) :=
arg min f ∈F |D t | − 1 (cid:80) x,y ∈D t 1 [ y (cid:54) =
f ( φ ( x )))] could be a terrible populationpredictor even when the underlying encoder is population optimal.
Ideally, representations wouldthus also guarantee that any ERM performs as well as possible for any desired task and dataset sizen .
This suggests minimizing the following worst-case expected risk over tasks and ERMs.
[Equation]
77 Deﬁnition 3.
An encoder φ ∗ is sample optimal for T , F iff if it is population optimal and minimizesthe worst-case expected risk of ERMs for arbitrary sample sizes, i.e.,for all n ≥ 1 :
[Equation]
[Figure]
Theoretical framework for linear probes F3.1 Characterizing optimal encoders for linear ISSL
In this section, we characterize all sample-optimal encoders (Def. 3), with simple properties that giveinsights into the form of the idealized representation.
The key for our theory is that any ∼ -invariantfunction g can be written as a composition between some function c g and a maximal invariant M[12], i.e., g = c g ◦ M where M : X → { 1 , . . .
, |X / ∼|} indexes the equivalence class:for any x, x + ∈ X :
[Equation]
⇐⇒ M ( x [Equation] =
M ( x + ) .
To build intuition for the ﬁnal characterization, let us ﬁrst discuss population optimal encoders (Def. 2)for unconstrained probes, then linear probes, and ﬁnally sample optimality (Def. 3).
Unconstrained probes.
By Def. 1, labels of downstream tasks p t ∈ T are ∼ -invariant.
Labels canthus be written as some function c t of M , i.e., arg max y p t
( y | X )
=
c t ( M ( X )) .
This shows that
M ( X ) contains all and only information about desired tasks T .
If probes are unconstrained, anencoder is population optimal iff M ( X ) is predictable/classiﬁable, i.e., there exist an h M such that
M ( X ) =
h M ( φ ( X )) .
Indeed, this ensures that the probe deﬁned by c t ◦ h M can classify the task p t .
Linear probes F .
The problem with constrained probes is that they might not be able to use theinformation about desired tasks.
In particular, the previous probe might not be linear c t ◦ h M (cid:54)∈ F .
As an illustration, consider 4 equivalence classes: cats, dogs, oranges, and lemons.
Fig.
2b shows arepresentation from which M ( X ) is predictable but invariant tasks are not linearly classiﬁable.
Infact, even when M ( X ) is linearly predictable, i.e., h M ∈ F as in Fig.
2c, downstream labels mightnot be, i.e., c t ◦ h M (cid:54)∈ F as shown in Fig.
2d.
By standard VC dimension arguments [13, 14], suchbinary task is linearly predictable if the representation’s dimensionality d is one less than the numberof equivalence classes
|X / ∼| − 1 , as in Fig.
2e.
Building on this intuition we prove that populationoptimal encoders are essentially 1 those that induce d ≥ |X / ∼| − 1 dimensional representations fromwhich M ( X ) is linearly predictable.train test
Sample optimality.
Although population optimality ensures the existence of a perfect linear probe,
ERM probes trained on ﬁnite samples might stillbe bad due to generalization issues (Fig. 3 left).
Intuitively, one can remove such bad ERMs bymapping equivalent examples to the same representation, i.e., by using invariant encoders.
Indeed,this ensures that ERMs that correctly predict oneexample in an equivalence class also correctlypredict all the other ( Fig. 3 right).
We prove thatsuch invariance of population optimal encoders is necessary and sufﬁcient for sample optimality.
Putting all together gives the following necessary and sufﬁcient properties of sample-optimal encoders.
Theorem 1.
An encoder φ ∗ is sample optimal for ∼ -invariant tasks T and linear F if and only if
• F -predictability of M : there exists a max.
invariant M and an f ∈ F s.t.
M ( X ) =
f ( φ ∗ ( X )) ;
• Invariance : φ ∗ is ∼ -invariant,
i
e., for any x, x + ∈ X we have x ∼ x + = ⇒ φ ∗ ( x ) =
φ ∗ ( x + ) ;
• Dimensionality : the effective dimensionality of the representations is at least one less than the
number of equivalence classes, i.e., dim(span( { φ ∗ ( x ) | x ∈ X} )) ≥ |X / ∼| − 1 .
3.2 The impact of augmentations on downstream performance
Let us compute the worst-case excess risk W n ( φ, F , T ) of sample optimal encoders and show itsdependence on the invariance structure.
The key is that since sample optimal encoders are invariant(Theorem 1), ERMs only need to be trained on a single example per equivalence class to perfectly classify all other examples from that class.
The risk then depends on the proportion of equivalence classesseen when training the probe, which can be computed in closed form as a function of the number ofequivalence classes |X / ∼| and downstream samples n .
For other similar results refer to Appx.
B.3.
Proposition 2.
The worst expected risk of a sample optimal φ ∗ for ∼ -invariant tasks T and any F is
[Equation]
Prop. 2 shows that fewer equivalence classes, i.e., coarser ∼ , leads to better downstream sampleefﬁciency.
Of course, the convergence rate will be slower for practical encoders than for sampleoptimal ones.
The result nevertheless suggests that good augmentations should be as strong aspossible (induce coarser ∼ ), while being label-preserving.
An example of coarse augmentationsare those in CLIP, [15] which map many images to similar sentences.
Practical ISSL objectives for linear probes F
The main remaining question is how to practically enforce requirements from Theorem 1.
In thissection, we derive simple objectives that learn optimal encoders in ideal settings (inﬁnite data, perfectoptimizers, universal approximators).
In the process, we shed light on why previously proposed ISSLmethods work, and how to improve them in practice.
same-class examples and maximize the angle across classes.
Our key insight is that we can learn sample-optimal encoders by jointlytraining the encoder and a logistic regression predicting M ( X ) from therepresentations.
Indeed, as shown in the neural collapse literature [16–21],the loss of logistic regression is minimized when representations form asimplex equiangular tight frame (sETF) as illustrated in Fig.
Namely,when: (i) M ( X ) is linearly predictable; (ii) the representation of same-classexamples collapse; (iii) the angle between representations of different classesis maximal.
Intuitively, each of those properties implies a requirement from
Theorem 1, namely: (i) corresponds to F -predictability and comes from thelinearity of logistic regression; (ii) corresponds to invariance and comes fromthe fact that log loss decreases variance due to Jensen’s inequality; (iii) impliesmaximal effective dimensionality and comes from the symmetry of log loss.
The previous discussion suggests optimizing the following ∼ -ISSL log loss , which simplycorresponds to multinomial logistic regression of M ( X )
[Equation]
, (6) where w maps classes to weights, e.g., by indexing a weight matrix.
Following previous work onneural collapse, we assume for this section that classes are equiprobable p X
([ x ]) = 1 / |X / ∼| , and thatclassiﬁer’s weights and representations are unit-normalized, i.e., w ∈ W 1 := { 1 , . .
. , |X / ∼|}
→ Sand φ ∈ Φ 1 := X → S where S denotes the ( d − 1) -sphere (for relaxations see e.g. [22, 23]).
We thenhave that any high dimension encoder trained to minimize the ISSL log loss will be sample-optimal.
Proposition 3.
Let p
X be a distribution with support X and equiprobable
equivalence classesp X ([ x ]) = 1 / |X / ∼| , ∀ x ∈ X .
If d ≥ |X / ∼| − 1 then any unit-normalized encoder that minimizes the∼ -ISSL log loss φ ∗ ∈ arg min φ ∈ Φ 1 L I ( φ ; p X ) is sample-optimal for ∼ -invariant tasks and linear F .
Prop. 3 shows that the ISSL log loss is a perfect pretext task in ideal settings.
This suggests optimizingthe ISSL log loss in practice and provides a formal relation between self-supervised and classicalsupervised learning.
The challenge is that we typically neither have access to the maximal invariant
M ( X ) nor the number of equivalence classes |X / ∼| required to compute the denominator of Eq.
Instead, knowledge about the equivalence structure comes from data augmentations A( ˜ X | X ) fromwhich we can sample examples that are equivalent to the input ˜ X ∼ X .
Having established our framework, we can re-interpret previous ISSL methods as practical approximations of the ISSL log loss using data augmentations (e.g., SimCLR, SwAV, DINO [24], [2],
SimSiam [25], MOCO [26]).
These approximations are nevertheless suboptimal, in that none of themlearn sample-
(nor population-) optimal encoders in idealized settings.
By directly deriving the ISSLmethods from Eq.
(6), we will prescribe improvements to previous ISSL objectives that ensure thatsample-optimal encoders are learned in idealized settings.
We broadly categorize previous methodsinto two families depending on whether they select the desired number of equivalence classes |X / ∼|or whether it is implicitly inferred from the augmentations.
We call these approaches distillation andcontrastive ISSL respectively.
For derivations and theoretical claims see Appx.
C.
4.1 Contrastive ISSL (CISSL)
augment represent project (a) SimCLR augment represent project (b) CISSL
This ensure linear predictability of downstream tasks, which will be computed by W t (cid:62) z .
Inspired by previous contrastive objectives[1, 27–29], we show how to optimize the ISSLlog loss using data augmentations to implicitlyreplace the maximal invariant M ( X ) and thenumber of classes |X / ∼| .
The resulting objective, dubbed CISSL , corresponds to SimCLRusing only a projection head g on one branch.
See Fig.
This asymmetry is necessary tolearn population optimal encoders for linear F .
CISSL bypasses the need for the maximal invariant by noticing that it only appears in the ISSLlog loss through the mapping to class weightsw ( M ( X )) .
CISSL thus aims to directly learna function g mapping equivalent (augmented)examples ˜ X to w ( M ( ˜ X )) .
Such mappingexists since M is invariant, e.g., g := w ◦ M .
When augmentations satisfy the Markov Chain ˜
X − M ( X ) − X , we show that we can replace theclass weights w ( M ( X )) in the ISSL loss by g ( ˜ X ) , where g is optimized over.
Intuitively, this isbecause predicting ˜ X contains only information about M ( X ) due to the data processing inequality.
The only remaining challenge is computing the denominator of Eq.
(6) without summing over theunknown classes X / ∼ .
To do so we use ranking conditional noise contrastive estimation (NCE;[30–32]).
NCE replaces classiﬁcation of equivalence classes by classiﬁcation of positives in a batch˜X :=
{ ˜ X + , ˜ X − 1 , . . .
˜ X − k } where the positive ˜ X
+ is sampled from the conditional A( ˜ X | X ) , whilethe k negatives ˜
X − i come from the marginal A( ˜ X ) =
E p X
[A( ˜ X | X )] .
Using Monte Carlo (MC)estimates with an unlabeled dataset D i.i.d.
∼ p
X we get our ﬁnal empirical CISSL objective
[Equation]
By Prop. 3 and NCE’s consistency [32], encoders trained with CISSL φ ∗ ∈ arg min φ ∈ Φ 1 ˆ L C ( φ ; D )are sample optimal for linear F in ideal settings ( |D| → ∞ and unconstrained g ) and when ˜ X −
M ( X ) − X .
While consistency holds for k ≥ 1 , more negatives k improves statistical efﬁciency [32].
Typically ˜ X, X take value in the same space (e.g. images).
If so, we can tie parameters by encoding˜ X , i.e., we can replace g by g (cid:48) ◦ φ where g (cid:48) :
R d → S is called a projection head.
The logits inside of
Eq.
(7) are then computed by g (cid:48)
( φ ( ˜ X )) (cid:62) φ ( x ) as shown in Fig.
5b.
This is very similar to SimCLR’sobjective g (cid:48)
( φ ( ˜ X )))
(cid:62) g (cid:48)
( φ ( x )) shown in Fig.
5a.
The difference is that, by projecting the currentrepresentation, SimCLR learns encoders that are not even population optimality for linear F .
Incontrast, CISSL learns sample optimal encoders in idealized settings.
Intuitively, this is because
CISSL trains the representations in the same way as they will be used in downstream tasks p t .
Indeed,representations φ ( X ) will be dotted with the downstream tasks’ weights W t (cid:62) φ ( x ) to compute logits.
In ISSL, representations φ ( x ) —rather than their projections g (cid:48)
( φ ( x )) —should thus be dotted with
ISSL weights w ( M ( X ))
=
g (cid:48) ( φ ( ˜ X ))) to compute logits.
CISSL thus derives, from ﬁrst principles,the necessary but asymmetric use of projection heads.
4.2 Distillation ISSL (DISSL)
[Figure]
The student (bottom branch) distills the teacher by predicting ˆ M .
In practice, CISSL requires contrasting many negatives.
Toavoid contrastive estimation, we can directly approximate
M ( X ) —instead of the weight w ( M ( X )) —so that thedenominator of Eq.
(6) can be computed exactly.
Theresulting method, dubbed DISSL (Fig. 6), is a simpler andtheoretically-grounded version of non-contrastive losses (e.g.
SwAV, DINO, SimSiam).
The challenge and main differencebetween each of those methods is how they estimate M ( X ) .
As M ( X ) is discrete, we use a conditional categoricaldistribution q ( ˆ M | X ) to estimate M ( X ) with a randomvariable ˆ M .
Replacing terms in the ISSL log loss then givesthe following objective.
Differences with Eq.
(6) are in red.
[Equation]
To highlight similarities with previous methods [2, 24, 33] we refer to Eq.
(8) as distilling theteacher q
( ˆ M | X ) into the student s φ,w ( ˆ M | X ) .
Importantly, Eq.
(8) is exactly the ISSL log loss ifthe teacher outputs a maximal invariant random variable, i.e., ˆ M = M ( X ) .
By Prop. 3, distillationthus learns sample-optimal encoders when the teacher takes at least C ≥ |X / ∼| values and satisﬁesthe following requirements, which correspond to the deﬁnition of maximal invariance (Eq.
Deterministic the teacher is a deterministic distribution max m ∈{ 1 ,...,C } q ( m | X ) = 1 ;
Invariant the teacher maps positives together x ∼ x + = ⇒ q
( ˆ M | x ) =
q ( ˆ M | x ) ;
Maximal the teacher maps negatives separately x (cid:54)∼ x − = ⇒
q
( ˆ M | x ) (cid:54) =
q ( ˆ M | x − ) .
There are many ways of enforcing these desired requirements for the teacher.
In the following, weuse information-theoretic quantities (entropies and KL divergences) as those are interpretable andcan be computed in closed form for categorical distributions such as our teacher and student.
Determinism and Invariance.
Both properties hold if and only if for any equivalent example x ∼ x +the cross-entropy of the teacher’s outputs is minimized E q ( ˆ M | x ) [ − log q
( ˆ M | x + )]
= 0 .
Indeed,zero cross-entropy simultaneously minimizes the conditional entropy H
[ ˆ M | x ] = 0 (equivalent todeterminism) and the KL divergence D KL [ q ( ˆ M | x ) (cid:107) q
( ˆ M | x + )]
=
(equivalent to invariance).
Maximality.
A natural way of enforcing maximality is to train the teacher to predict differentlynegatives x (cid:54)∼ x − .
This requires accessing batches of negatives, which we wanted to avoid with
DISSL.
Instead, assume that we know (or have a prior on) the true distribution of equivalence classesp ( M ( X )) .
Then a deterministic and invariant teacher is maximal if and only if the KL divergencebetween its marginal q ( ˆ M ) =
E p X [ q ( ˆ M | X )] and the true one is minimized D KL
[ q ( ˆ M ) (cid:107) p ( M ( X ))]
=0 .
We can thus avoid contrasting negatives using (a prior on) the distribution of equivalence classes.
Using a Lagrangian relaxation (with multipliers λ, β ) of the teacher’s constraints with an MCestimate of the distillation loss (Eq.
we get the following empirical DISSL loss ˆ L D ( φ ; D ) :=
[Equation]
Algorithm
Require: head g , weights W , enc.
φ , batched inputs X , aug.
A , hyp.
β , λ .
1: ˜ X, ˜ X + ← sample
(A( ˜ X | X ) , 2) 2: q
( ˆ M | ˜ X )
= softmax( g ( φ ( ˜ X ))
q ( ˆ M | ˜ X + ) =
softmax( g ( φ ( ˜ X + )) 4: s ( ˆ M | ˜ X + ) =
softmax( W T φ ( ˜ X + )) 5: q ( M ) = batch _ avg( q ( M | ˜ X )) 6: mxml =
(cid:80) m q ( m ) log q ( m ) 7: det _ inv =
(cid:80) m q ( m | ˜ X ) log q ( m | ˜ X + ) 8: dstl = (cid:80) m
q ( m | ˜ X ) log s ( m | ˜ X + ) 9: return λ ∗ mxml − β ∗ det _ inv − dstl
In ideal settings ( |D| → ∞ , unconstrained q , knownp ( M ( X )) ) and for λ, β → ∞ we show that encoderstrained with DISSL are sample optimal for linear F andtasks that are invariant to the equivalences deﬁned by theconnected components of the augmentation graph [8].
Ifthe marginal is unknown, we follow the MaxEnt [34] anduse a uniform prior p ( M ( X ))
=
Unif(1 , C ) .
The KL in
Eq.
(9) then corresponds to maximizing entropy H [ ˆ M ] .
As with CISSL, when ˜ X, X are in the same space, wetie parameters by encoding ˜ X before the teacher, i.e.,q ( ˆ M | X ) =
softmax( g (cid:48) ( φ ( X ))) for a head g (cid:48) :
R d → 262 R C .
Algorithm 1 and Fig.
6 illustrate DISSL with a uniform prior, and a marginal q ( ˆ M ) estimated from batches.
DISSL is one of the many distillation methods that can be derived from our teacher’s requirements.
These requirements generally lead to a framework for deriving, comparing, and analyzing distillationobjectives.
In Appx.
E we provide a taxonomy of 12 previous SSL methods from this perspective—none of which recover sample optimal encoders.
Typically, previous methods favor: (i) determinismby “sharpening” the teacher q
( ˆ M | X ) with a temperature parameter; (ii) invariance by makingthe teacher’s and student’s similar for equivalent inputs q ( ˆ M | x + ) ≈ s φ,w ( ˆ M | x ) ; (iii) maximalitythrough optimization tricks (e.g. stop gradients or momentum encoder) to avoid a constant teacher, i.e.,referred to as “collapsing” [25, 26, 35, 36].
Note that avoiding constant collapse q ( ˆ M | x (cid:48) )
=
q ( ˆ M | x )for any x, x (cid:48) ∈ X is insufﬁcient, e.g., mapping half the inputs to a constant is also undesirable.
Maximality formalizes the desired requirement: non-equivalent examples should not be mapped together.
ISSL for non-linear predictors F +
Linear probes F are standard for evaluating representations [28, 37, 38].
However, if the goal is tomaximize performance it is natural to consider non-linear predictors F + .
The question then becomeshow should we learn optimal representations for any F + ?
In Appx.
B.5 we extend our frameworkto certain non-linear probes, such as neural networks, that separate into a hidden component h andlinear layer,
i.e., f ( · ) =
W T h ( · ) .
We informally summarize the theory and implications below.
Theory.
There are two main differences between our characterization of optimal encoders for linearand non-linear probes.
First, encoders should ensure predictability of M ( X ) for the desired F + .
Second, the dimensionality requirement decreases with the complexity of F + , e.g., for universal F + itis d ≥ 1 .
More generally, the necessary dimension and the sufﬁcient dimension for optimality do notcoincide and respectively depend on the predictors’ VC dimension [13] and memory capacity [39].
Implications.
Both theoretical differences lead to direct implications for non-linear ISSL.
First, wecan decrease the dimensionality when performing ISSL for more complex non-linear probes.
Second,we should use a projection head that ensures the predictability of the maximal invariant with using thedesired probes F + .
Similarly to the linear case, we should match how the representation is used inthe ISSL log loss and downstream tasks.
We should thus apply the non-linear predictor h in a similarasymmetric way.
For CISSL we should thus compute the right hand side of Fig.
5 as g ( φ ( ˜ X ))
h ( φ ( x )) .
For DISSL, we should change the line 4 of Algorithm 1 to s ( ˆ M | ˜ X + ) =
softmax( W T h ( φ ( ˜ X + ))) .
Summary of insights and relation to previous work
Let summarize our framework’s main insights and their relation to previous work.
Details in Appx.
E.
DISSL.
DISSL is related to DINO and SwAV, as each method distill a categorical teacher but differ inhow they enforce its maximality.
DINO does it implicitly through optimization tricks, e.g., the teacheris an exponential moving average of the student and uses stop-gradients.
SwAV enforces maximalityexplicitly by equiprobably clustering a queue of examples using Sinkhorn’s algorithm [40] and stopgradients.
DISSL is also explicit but uses a simpler and more efﬁcient max entropy regularization,which does not require queues or stop-gradients.
Furthermore, both SwAV’s and DINO’s student usenon-linear projection head, which do not ensure good linear probing (see Sec. 4.1).
Dimensionality.
Theorem 1 suggests increasing the representation’s dimensionality to ensure probescan classify all invariant tasks.
To our knowledge increasing the dimensionality of SSL representationshas never been investigated (note: this is different from the projection’s dimensionality in [1, 41]).
The closest empirical work concerns the beneﬁts of avoiding dimensionality collapse [42, 43].
In thetheoretical literature, the importance of large dimension has been indirectly suggested by the choiceof assumptions (e.g. [8, 44]) but the exact necessary and sufﬁcient dimension is novel.
Augmentations.
Prior work have suggested that good augmentations remove nuisance informationbut preserve labels [12, 45–49].
Prop. 2 gives a new perspective on the beneﬁts of such augmentations,by proving the exact relation between optimal sample efﬁciency and the number of equivalences.
Projection heads.
Secs.
4 and 5 theoretically show how to chose projection heads, namely, oneshould be as large as possible while the other should have the architecture of downstream probes.
Toour knowledge, we are the ﬁrst to relate the probing family and the projection head.
Furthermore, ourtheory is the ﬁrst to suggest why projection heads improve performances (e.g. [1, 25, 50]) in general
SSL (note: this is unrelated to adding heads to avoid collapsing in non-contrastive learning [36, 51]).
Theoretical SSL.
Theorem 1 characterizes optimal encoders for ISSL.
This contrasts with standardtheoretical work in SSL that typically analyze speciﬁc SSL algorithms [3, 5, 6, 8, 52–54] and / orfocus on upper-bounding downstream performance [4].
The advantage of our theory is that it providesa simple unifying framework from which we can derive a variety of SSL methods and actionableinsights.
The downside is that by distancing itself from speciﬁc algorithms, the framework does notprovide guarantees outside of ideal settings.
We thus see our theory as complementary to prior work.
Experiments
The experiments in the main paper focus on evaluating our frameworks’ prescriptive implications.
Formore results see Appx.
H.
For experimental details see Appx.
G. Code is at anonymous .
In summary, our experimental results show that: (i) CISSL/DISSL outperform their respectivebaselines on standard benchmarks as suggested by Sec. 4; (ii) increasing dimensionality improvesdownstream probing as suggested by Theorem 1; (iii) coarser augmentations improve sample efﬁciency as in Prop. 2; (iv) smaller ISSL log loss improves downstream performance as suggested by
Prop. 3; (v) projection heads should be related to probing family as discussed in Secs.
4 and 5.
For the ﬁrst experiments, we use TinyImageNet [55], 300 pretraining epochs, and ResNet18s.
Notethat for ResNet18s the dimensionality of the representation ( res5avg ) is d = 512 .
For contrastivebaselines we use SimCLR.
For distilling baselines we use the best over DINO, SwAV, SimSiam.
[Table]
CISSL/DISSL improve linear probing
Table 1 shows that each of our prescriptions significantly improve distillation and contrastive ISSL.“Ours” vs “Baselines” shows that removing one thenon-linear projection heads (as in Fig. 5) gives a1% gain in contrastive learning, while our new distillation objective achieves 3% gains compared tothe best distillation baseline.
“dim. ↑ ” shows thatincreasing the dimensionality of the representationfrom 512 to 2048 further improves linear probingby 1 – 2% .
Finally, “coarse aug.” shows that coarseraugmentations (we double the strength of standard TinyImageNet augmentations [35] and add blurring), we further improve performance by 3 – 4% .
All together our framework’s insights signiﬁcantlyimproves linear probing accuracy: 5% gains for contrastive learning and 8% for distillation.
Coarse augmentations improve sample efﬁciency.
Prop. 2 suggests that coarser label-preservingaugmentations improve downstream sample efﬁciency.
We test that by introducing a hyperparameterα that controls the strength of standard augmentations.
Speciﬁcally, α increases the probability of grayscaling ( p= 0 .
1 α ), each color jittering parameter ( brightness= 0 .
4 α , contrast= 0 .
1 α ,saturation= 0 .
2 α , hue= 0 .
1 α ), and the cropping interval ( scale= [0 . 1 /α, 1] ).
Fig.
7a shows fourdifferent DISSL encoders α ∈ { 0 . 25 , 0 . 5 , 1 , 2 } evaluated at varying downstream samples sizes.
Assuggested by the theory, stronger label-preserving augmentations improve sample efﬁciency.
[Figure]
ISSL log loss correlates with performance.
Previous work [44] suggested that contrastive lossesmay not predict downstream performance.
On the contrary, Prop. 3 shows that minimal ISSL log lossimplies the optimality of encoders if d is large.
This suggests that, for ﬁxed augmentations, ISSLlog loss is highly related to performance.
To test this relation, we trained 80 CISSL models withvarious hyperparameters (e.g. optimizer, projection head, learning rate) while ﬁxing augmentationsand negatives k .
Fig.
7b shows that the test ISSL loss and downstream accuracy are highly correlated(gray points).
In Appx.
H.3 we show similar correlation with the training ISSL loss.
Longer training decreases ISSL log loss.
An efﬁcient way of decreasing ISSL log loss is longertraining (blue points in Fig.
7b only vary the number ISSL epochs).
Prop. 3 thus provides anexplanation of the well-known performance gains of longer ISSL training [1, 56].
Namely, longertraining decreases ISSL log loss, which results in representations that are closer to optimality.
Increasing dimensionality improves performance.
Theorem 1 shows that linear ISSL requires adimensionality of d = |X / ∼| − 1 to ensure predictability of all invariant tasks.
In practice, thenumber of equivalence classes is most likely very large.
Although such dimensionality is impractical,it does suggest increasing the dimensionality beyond the d = 512 of ResNet18s.
To test this, we train
CISSL with a larger dimensionality by increasing the output channels before the average poolinglayer.
To keep the number of parameters comparable, we use bottlenecks or low-rank linear layersbefore and after the representations layer as detailed in Appx.
G. Fig.
7c shows that increasing thedimensionality has a signiﬁcant effect on downstream performance for linear F (blue).
In Appx.
H.3we show that, contrary to popular beliefs, the gains when going from ResNet18s to ResNet50s are dueto an increase in dimensionality (512 to 2048) rather than the number of parameters (11M to 23M).
Previous results concern linear F which are standard for ISSL evaluation.
In practice, one wouldlikely use more complex probes F + to improve results.
In Sec.
5 we discussed how to perform ISSLfor F + .
In the following, we evaluate ISSL for MLPs F + with two hidden layers of 2048 units.
[Table]
Performing ISSL for MLPs F + .
Table 2 showsthat MLP probes (“ F EVAL : MLP”) outperforms linear ones (“ F EVAL : lin.”
) even in few-shot regimes(10 shot).
Furthermore, when predicting with MLPprobes it is desirable to train the ISSL encoders for
MLPs (“ F ISSL : MLP”) instead of linear probes (row4 vs 3).
As discussed in Sec. 5, the only differencebetween DISSL for different probes is that the student’s representation is projected using a head withthe same architecture as downstream probes.
Dimensionality affects less MLPs F + .
In Sec.
5 we saw that the optimal dimensionality d decreasesfor more complex predictors F + .
Fig.
7c indeed shows that larger dimensions result in smaller gainswhen the training and evaluation probe are MLPs F + (orange) compared to linear ones F (blue).
Given encouraging TinyImageNet results, we investigated whether our objectives can be used as replacement for standard objectives on ImageNet [57].
We speciﬁcally replaced SimCLR by our CISSLand SwAV by our DISSL in VISSL ’s codebase [58] without any tuning.
Due to computational constrains we mostly consider 100 epochs of pretraining and do not use augmentations to train the probes.
[Table]
Our models outperform baselines on ImageNet.
Table 3 shows that CISSL’s small modiﬁcationsof SimCLR lead to important gains ( 1 . 8% ).
Both CISSL and our novel DISSL objective signiﬁcantlyoutperform all the considered baselines including SwAV by a 3% when using standard augmentations.
Given encouraging results on standard augmentations, we investigated the performance of our modelcompared to near a SOTA model: SwAV trained with their special multi-crop augmentations.
[Table]
DISSL is competitive with SOTA models at scale.
Table 4 shows that DISSL outperforms SwAV on 100 epochsand performs at least as well for 400 epochs.
Combinedwith Table
this suggests that DISSL is robust and workswell in different settings.
This is encouraging given thelack of tuning and the simplicity of DISSL’s out-of-the-boxobjective.
In contrast, SwAV requires stopping-gradients,storing a queue, running Sinkhorn’s algorithm, andfreezing certain parameters during initial training steps.
Summary and Outlook
We presented a simple conceptual framework to understand, compare and improve ISSL methods.
Inspired by recent work on supervised [59] and robust [60] learning, we derived such framework bystudying algorithmic-agnostic goals for ISSL.
In particular, Theorem 1 provides the minimal andsufﬁcient requirements for optimal encoders, and Prop. 2 gives the resulting probing risk as a functionof pretraining augmentations.
On the algorithmic side, Sec. 4 uses our framework to derive simplerand theoretically-motivated variants of prior SSL objectives.
Altogether our framework providesmany actionable insights for ISSL.
In particular, how to chose: SSL algorithm, the dimensionality ofthe representations, the architecture of projection heads, and the augmentations.
On the empirical side,
Sec.
7 shows that each of our prescription leads to signiﬁcant gains for contrastive and non-contrastive
SSL on TinyImageNet.
As a result, our ﬁnal SSL algorithms outperforms strong baselines such as
SwAV+multicrop trained on linear probing of ImageNet.
There are many limitations that should be addressed for a more realistic prescriptive framework for
ISSL.
See Appx.
F for more details.
First, our theory is binary, e.g., augmentations are either labelpreserving or not, and encoders are optimal or not.
Although our framework likely captures the rightintuition, nuance would be more realistic.
Second, we consider unconstrained encoders φ , dimensionality, and inﬁnite unlabeled data.
Such assumptions follow recent empirical trends in increasing dataand models [61, 62], but it means that our framework currently cannot improve the computational ordata efﬁciency of the pretraining phase.
Finally, we consider all invariant tasks T even though only asubset of those are meaningful.
Optimal representations for this subset will be different, and deﬁningthis subset likely requires considering inductive biases of encoders [44].
Despite those limitations,