W EIGHTED E NSEMBLE S ELF -S UPERVISED L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Ensembling has proven to be a powerful technique for boosting model performance, uncertainty estimation, and robustness in supervised learning.
Advances in self-supervised learning (SSL) enable leveraging large unlabeled corpora for state-of-the-art few-shot and supervised learning performance.
In this paper, we explore how ensemble methods can improve recent SSL techniques by developing a framework that permits data-dependent weighted cross-entropy losses.
We refrain from ensembling the representation backbone; this choice yields an efﬁcient ensemble method that incurs a small training cost and requires no architectural changes or computational overhead to downstream evaluation.
The effectiveness of our method is demonstrated with two state-of-the-art SSL methods, DINO (Caron et al., 2021) and MSN (Assran et al., 2022).
Our method outperforms both in multiple evaluation metrics on ImageNet-1K, particularly in the few-shot setting.
We explore several weighting schemes and ﬁnd that those which increase the diversity of ensemble heads lead to better downstream evaluation results.
Thorough experiments yield improved prior art baselines which our method still surpasses; e.g., our overall improvement with MSN ViT-B/16 is 3.9 p.p. for 1-shot learning.
I NTRODUCTION
[Figure]
The promise of self-supervised learning (SSL) is to extract information from unlabeled data and leverage this information in downstream tasks (
He et al., 2020; Caron et al., 2021); e.g., semi-supervised learning (Chen et al., 2020a;b), robust learning (Radford et al., 2021; Ruan et al., 2022; Lee et al., 2021), few-shot learning (Assran et al., 2022), and supervised learning (Tomasev et al., 2022).
These successes have encouraged increasingly advanced SSL techniques (e.g., Grill et al., 2020; Zbontar et al., 2021;
He et al., 2022).
Perhaps surprisingly however, a simple and otherwise common idea has received limited consideration: ensembling.
Ensembling combines predictions from multiple trained models and has proven effective at improving model accuracy (Hansen & Salamon, 1990; Perrone & Cooper, 1992) and capturing predictive uncertainty in supervised learning (Lakshminarayanan et al., 2017; Ovadia et al., 2019).
Ensembling in the SSL regime is nuanced, however; since the goal is to learn useful representations from unlabeled data, it is less obvious where and how to ensemble.
We explore these questions in this work.
We develop an efﬁcient ensemble method tailored for SSL that replicates the non-representation parts (e.g., projection heads) of the SSL model.
In contrast with traditional “post-training” ensembling, our ensembles are only used during training to facilitate the learning of a single representation encoder, which yields no extra cost in downstream evaluation.
We further present a family of weighted crossentropy losses to effectively train the ensembles.
The key component of our losses is the introduction of data-dependant importance weights for ensemble members.
We empirically compare different choices from our framework and ﬁnd that the choice of weighting schemes critically impacts ensemble diversity, and that greater ensemble diversity correlates with improved downstream performance.
Our method is potentially applicable to many SSL methods; we focus on DINO (Caron et al., 2021) and MSN (Assran et al., 2022) to demonstrate its effectiveness.
Fig.
1 shows DINO improvements from using our ensembling and weighted cross-entropy loss.
In summary, our core contributions are to:
• Develop a downstream-efﬁcient ensemble method suitable for many SSL techniques (Sec. 3.1).
• Characterize an ensemble loss family of weighted cross-entropy objectives (Sec. 3.2).
• Conduct extensive ablation studies that improve the prior art baselines by up to 6.3 p.p.
(Sec. 5.1).
• Further improve those baselines with ensembling (e.g., up to 5.5 p.p. gain for 1-shot)
(Table 2).
B ACKGROUND
In this section, we frame SSL methods from the perspective of maximum likelihood estimation (MLE) and use this as the notational basis to describe the state-of-the-art clustering-based SSL methods as well as derive their ensembled variants in Sec. 3.
From Maximum Likelihood to SSL Denote unnormalized KL divergence (Dikmen et al., 2014) between non-negative integrable functions p, q by K [ p ( X ) , q ( X )]
=
H × [ p ( X ) , q ( X )] − H [ p ( X )] , where H ×
[ p ( X ) , q ( X )]
= − (cid:82) X p ( x ) log q ( x )d x + (cid:82) X
q ( x )d x − 1 is the unnormalized cross-entropy (with 0 log 0 = 0 ) and H [ p ( X )]
=
H × [ p ( X ) , p ( X )] .
These quantities simplify to their usual deﬁnitions when p, q are normalized, but critically they enable ﬂexible weighting of distributions for the derivation of our weighted ensemble losses in Sec. 3.2.
Let ν ( X, Y ) =
ν ( X ) ν
( Y | X ) be nature’s distribution of input/target pairs over the space X × Y and s ( Y | θ, X ) be a predictive model of target given the input parameterized by θ ∈ T .
Supervised maximum likelihood seeks the minimum expected conditional population risk with respect to θ ,
E ν ( X ) K [ ν ( Y | X ) [Equation]
[Equation]
Henceforth omit E ν ( X ) H [ ν ( Y | X )] since it is constant in θ.
Since ν ( X, Y ) is unknown, a ﬁnite sample approximation is often employed.
Denote a size- n i.i.d.
training set by D n = { x i } i ∈ [ n ] ∼ ν ⊗ n and empirical distribution by ˆ ν ( X, Y )
=
1 n (cid:80) x ∈D n ,y ∼ ν
( Y | x ) δ ( X − x ) δ ( Y − y )
where δ : R → { 0 , 1 } is 1 when x = 0 and 0 otherwise.
The sample risk is thus − 1 n (cid:80) x ∈D n H ×
[ˆ ν ( Y | x ) , s ( Y | θ, x )] .
In SSL, we interpret ν
( Y | x ) as being the oracle teacher under a presumption of how the representations will be evaluated on a downstream task.
This assumption is similar to that made in Arora et al.
; Nozawa et al.
We also assume ˆ ν
( Y | X ) is inaccessible and/or unreliable.
Under this view, some SSL techniques substitute ˆ ν ( Y | x ) for a weakly learned target or “teacher”, t ( Y | x ) .
We don’t generally expect t ( Y | x ) to recover ν
( Y | x ) ; we only assume that an optimal teacher exists and it is ν
( Y | x ) .
With the teacher providing the targets, the loss becomes − 1 n (cid:80) x ∈D n H × [ t ( Y | x ) , s ( Y | θ, x )] .
Teacher and student in clustering SSL methods Clustering SSL methods such as SWaV (Caron et al., 2020), DINO (Caron et al., 2021), and MSN (Assran et al., 2022) employ a student model characterized by proximity between learned codebook entries and a data-dependent code,
[Equation]
where the encoder r ω : X → Z produces the representations used for downstream tasks, and the projection head h ψ : Z → R d and codebook entries { µ y } y
∈Y ∈ R
d characterize the SSL loss.
Eq.
(2) can be viewed as “soft clustering”, where the input is assigned to those centroids that are closer to the projection head’s output.
The projection head and codebook are used during training but thrown away for evaluation, which is empirically found vital for downstream tasks (Chen et al., 2020a;b).
Hyperparameters τ ∈ R > 0 , c ∈ Z > 0 represent temperature and codebook size.
The teacher is deﬁned as t ( Y | x )
=
s ( Y | stopgrad( g ( θ )) , x ) where g : T → T .
Commonly g ( θ ) is an exponential moving average of gradient descent iterates and the teacher uses a lower temperature than the student.
To capture desirable invariances and prevent degeneracy, data augmentation and regularization (e.g., Sinkhorn-Knopp normalization (Caron et al., 2020), mean entropy maximization (Assran et al., 2022)) are essential.
As these are not directly relevant to our method, we omit them from further consideration.
M ETHOD
Ensembling is a technique that combines models to boost performance, and has been especially successful in supervised learning.
We are interested in ensembling methods that carry over this success to SSL approaches.
However, SSL has key differences, such as throw-away “projection heads”, from supervised learning that result in a multitude of possibilities for how to ensemble.
With this in mind, we propose ﬁrst where to ensemble, and then how to ensemble.
Those proposals result in an efﬁcient “peri-training” ensembling technique speciﬁcally tailored for SSL and a family of weighted ensemble objectives; we subsequently suggest different ways to select the weights.
3. W HERE TO ENSEMBLE ?
[Figure]
Two augmented inputs are encoded by the teacher/student into representations, and then processed by an ensemble of heads.
The loss for each head is weighted and summed into the ﬁnal loss.
Strike-through edges indicate stop-gradients.
See Appx. A for pseudocode. Denote the teacher/student ensembles by { t i ( Y | x ) } i ∈ [ m ] and { s ( Y | θ j , x ) } j ∈ [ m ] and deﬁne each as in Sec. 2; parameters θ = { θ j } j ∈ [ m ] ∈ T m are independently initialized, all students use one temperature and all teachers another.
We asymmetrically denote t
i
( Y | x ) and s ( Y | θ j , x ) to emphasize that teachers’ gradients are zero and that the students are distinct solely by way of θ i (cid:54) =
θ j .
Studying heterogeneous architectures and/or different teacher parameterizations is left for future work.
Recall that θ j parameterizes the encoder, projection head, and codebook parameters: θ j = ( ω j , ψ j , { µ jy } y ∈Y ) .
We further restrict T m such that ω i =
ω
j , i.e., we limit our consideration to ensembles of projection heads h ψ j and/or codebooks µ j but not encoders r ω j .
This choice makes our ensemble method inherently different from traditional supervised ensembling or encoder r ω ensembling: the ensembled parts are not used for evaluation but for improving the learning of non-ensembled representation encoder during training, thus it requires no change of downstream evaluation or computational cost.
Ensembling of r ω is left for future work.
3. H OW TO ENSEMBLE ?
We would like to extend the loss to support an ensemble of teacher/student pairs while respecting the MLE intuition of the loss as in Sec. 2.
Additionally, we want to facilitate data-dependent importance weights, thus enabling preferential treatment of some teacher/student pairs.
We therefore propose a weighted average (unnormalized) cross-entropy loss,
[Equation]
where w ijy = softmax (cid:16)(cid:110) 1 γ f ijy (stopgrad [Equation] x [Equation]
[Equation]
j ∈ [ m ] (cid:111)(cid:17)
[Equation]
The notation w ijY (cid:12) t i ( Y | x ) denotes a Hadamard product; i.e., the product of event-speciﬁc weights and probabilities for each y ∈ Y .
The hyperparameter γ is the temperature.
The function f ijy is deﬁned for brevity and discussed in the following section.
This objective admits generality and ﬂexibility for introducing various weighting schemes, as it supports potential interactions between all teacher/student pairs and allows the weights to be both model- and data-dependent.
Up to a constant independent of θ , it is an importance weighted average of (unnormalized) KL divergences between each teacher and each student; i.e., a mixture of MLE-like objectives.
We stop the gradient of w ijy to θ in order to keep the overall gradient a weighted average of students’ log-likelihood gradients, similar to Eq.
We also normalize the weights such that each data point equally contributes to the loss.
3. H OW TO WEIGHT ?
In this section, we present several instantiations of our losses with different weighting schemes.
We empirically show in Sec.
5 that the particular choice of weighting scheme is critical for the representation performance and the induced diversity of ( h ψ , µ ) -ensembles.
For simplicity we assume γ = 1 in this section.
We indicate with ⇐⇒ that a loss has the same arg min as Eq.
For additional analysis and discussion, see Appx.
D.
Uniform weighting (U NIF )
The simplest strategy is to treat different teacher/student pairs independently and average each with uniform weighting; i.e.,
[Equation]
This strategy introduces uniform weights w i = 1 m over ensemble elements.
The role of log δ ( i − j ) (here and elsewhere) is to sub-select corresponding teacher/student pairs rather than all m 2 pairs.
Probability weighting (P ROB )
An alternative to using the average cross-entropy loss (U NIF ) is to compute the cross-entropy loss of the average predictions whose gradient is weighted by w ijy (see
Appx.
D.1).
At γ = 1 , those gradient weights simplify into an average over the student probabilities:
[Equation]
Averaging the predictive distributions introduces correspondence between codes from different heads; thus different heads are no longer independent but instead cooperate to match the student to the teachers.
The loss favors student heads with more conﬁdent predictions (i.e., larger s ( y | θ j , x ) ).
Further motivation for averaging predictions comes from multi-sample losses studied in Morningstar et al.
Note that the joint convexity of (unnormalized) KL divergence implies that this loss is upper bounded by the U NIF loss up to some constant in θ (see Appx. D).
Although the P ROB strategy favors conﬁdent student predictions, the weights change as a function of y ∈ Y .
This may be in conﬂict with our intuition that SSL is like maximum likelihood (Sec. 2), since under that view, the teacher is responsible for weighting outcomes.
Entropy weighting (E NT )
Another way to favor heads with more conﬁdent predictions is to directly weight by their predictive entropies; i.e.,
[Equation]
where the weight w i = softmax i ( {− 1 γ H [ t i (cid:48)
( Y | x )] : i (cid:48) ∈ [ m ] } ) is inversely correlated with the entropy of teacher predictions.
In other words, the head whose teacher has a lower entropy (i.e., higher conﬁdence about its prediction) is given a larger importance weight for learning the representation.
Like P ROB , this strategy encourages “data specialists” by emphasizing strongly opinionated teacher heads for different inputs.
Like U NIF , different heads are treated more independent (than P ROB ), since interaction between different heads is introduced only through the weight computation.
By preferring low-entropy teachers we also favor low variance teachers; this aligns with the intuition that using a lower-variance teacher beneﬁts representation quality (Wang et al., 2022).
Countless other weighting schemes It is impossible to fully explore the space of weightings; the following might also be interesting to study in detail but were omitted due to resource constraints.
[Equation]
[Equation]
Note that “aligned” versions of all schemes are possible by using f ijy + log
δ
( i − j ) .
We did early experiments exploring Eqs. (11) and (12), but the results were inferior and are largely omitted below.
R ELATED W ORK
Self-supervised learning Recent work on self-supervised learning (SSL) focuses on discriminative or generative approaches.
Most discriminative approaches seek to learn augmentation-invariant representations by enforcing the similarity between augmented pairs of the same image while utilizing different techniques to avoid collapse.
Contrastive methods (Chen et al., 2020a;
He et al., 2020; Wu et al., 2018; Hjelm et al., 2018; Bachman et al., 2019; Tian et al., 2020) use a large number of negative samples with a noise-contrastive objective (Gutmann & Hyvärinen, 2010; Oord et al., 2018).
A large body of followup work eliminates the necessity of explicit negative samples with various techniques, including clustering assignment constraints (Caron et al., 2018; 2020; 2021;
Asano et al., 2019), bootstrapping (Grill et al., 2020) or self-distillation (Caron et al., 2021) inspired by mean teacher (Tarvainen & Valpola, 2017), asymmetric architecture design (Grill et al., 2020;
Chen & He, 2021), or redundancy reduction (Zbontar et al., 2021; Bardes et al., 2021).
Recent generative approaches that use masked image modeling as the pretraining task (Dosovitskiy et al., 2020; Bao et al., 2021;
He et al., 2022; Zhou et al., 2022; Xie et al., 2022) have achieved competitive ﬁnetuning performance.
Our method may be applicable to all of the above methods that have some sort of “projection head”, such as most of the discriminative approaches.
Ensemble methods Ensembling has been extensively studied for improving model performance (Hansen & Salamon, 1990; Perrone & Cooper, 1992; Dietterich, 2000) and uncertainty estimation (Lakshminarayanan et al., 2017; Ovadia et al., 2019) in supervised learning and semi-supervised learning (Laine & Aila, 2016).
A major research direction is to train efﬁcient ensembles with partial parameter sharing (Lee et al., 2015; Wen et al., 2020; Dusenberry et al., 2020; Havasi et al., 2020) or intermediate checkpointing (Huang et al., 2017; Garipov et al., 2018).
Our method also shares the encoder parameters across ensembles, which is closely related to multi-headed networks (Lee et al., 2015; Tran et al., 2020).
Ensemble methods for SSL are less explored.
Some recent work studies ensembles of supervised models adapted from pretrained SSL models.
Gontijo-Lopes et al.
(2022) conduct an empirical study of ensembles adapted from different SSL models and ﬁnd that higher divergence in SSL methods leads to less correlated errors and better performance.
Wortsman et al.
(2022) ensemble multiple ﬁnetuned models adapted from the same SSL model by averaging their weights, which boosts the performance without any inference cost.
Our method differs from them in that it (1) applies to the SSL training stage to directly improve representation quality, rather than aggregates multiple models in the post-training/ﬁnetuning stage; (2) introduces little training cost and no evaluation cost; and (3) is complementary to these post-training/ﬁnetuning ensembling methods.
E XPERIMENTS
We carefully study the impact of ( h ψ , µ ) -ensembles and our selected weighted ensemble losses (U NIF , P ROB , and E NT ) on smaller DINO models in Sec. 5.1.
Using what we learned in those experiments, in Sec.
5.2 we present new state-of-the-art results on ImageNet-1K on various metrics for multiple model sizes by ensembling both DINO- and MSN-based models.
Finally, we explore ensemble evaluations in the transfer learning setting in Sec.
Additional experimental details and results are in Appx.
B and Appx.
C, respectively.
Experimental setup We assessed the effectiveness of our method with two SSL methods: DINO (Caron et al., 2021) and MSN (Assran et al., 2022).
In order to ensure that we are comparing against strong baselines, we consider three different classes of baselines: (1) evaluation numbers reported in the original works (Caron et al.
(2021), Assran et al.
(2022), and Zhou et al.
for an additional baseline iBOT)
; (2) evaluation of our implementation using the hyperparameters reported in the original works (DINO only, for space reasons) to validate our implementation; and (3) evaluation of our implementation using the best hyperparameters that we found by tuning the baselines (DINO and
MSN) for fair comparisons.
In almost all models and evaluations, our retuned baselines give nontrivial performance improvements on top of previously reported numbers.
These type (3) baselines [Table] we label DINO ∗ and MSN ∗ , and we use them as the base models for our experiments with ( h ψ , µ ) - ensembles and weighted ensemble losses.
Appx.
B.2.1 describes the details for getting such strong performance for DINO ∗ and MSN ∗ .
In particular, we ﬁnd that the projection head has a crucial impact on label efﬁciency of representations and using a smaller head (3-layer MLP with hidden size 1024) signiﬁcantly improves few-shot evaluation performance (see Appx. C.2).
Evaluation metrics We compared models trained with and without our ( h ψ , µ ) -ensembles by measuring various evaluation metrics on ImageNet-1K (Deng et al., 2009).
The evaluation metrics reﬂect the decodability and the label efﬁciency of learned representations.
We measured the decodability with respect to both the linear classiﬁer following the common linear evaluation protocol and the k -NN classiﬁer following Caron et al.
We measured the label efﬁciency by evaluating the linear evaluation performance in few-shot settings, including 1% ( ∼ 13-shots) labeled data evaluation (Chen et al., 2020a) and 1-/2-/5-shot evaluations (Assran et al., 2022).
All evaluations used frozen representations of the teacher encoder – we did not ﬁne tune the models.
See Appx.
B.3 for details.
5. E MPIRICAL STUDY OF ( h ψ , µ ) - ENSEMBLES
Table 1 compares different strategies for where and how to ensemble .
Fig.
4 compares the impact of the weighted ensemble loss on ( h ψ , µ ) -ensemble diversity .
Fig.
3 shows the effect of increasing the number of ensembles , adjusting the temperature γ , and increasing baseline projection head parameters .
In these experiments, we used DINO ∗ with ViT-S/16 trained for 300 epochs as the base model.
We compared different ensemble methods applied to the base model with m = 16 heads which we found to work the best.
For the E NT strategy in Table 1, the entropy weighting temperature γ is set to 0 .
05 × log( c ) by default which is selected from { 0 . 0125 , 0 . 025 , 0 . 05 , 0 . 1 , 0 . 2 } × log( c ) , where the scale log( c ) gives the maximum entropy of the codebook size c .
For P ROB , we keep γ = 1 .
Where to ensemble We study the where question by ensembling either the projection head h ψ , the codebook µ , or both with the E NT and the P ROB ensemble strategies, as shown in Table 1.
We ﬁnd that ensembling both h ψ and µ provides the largest gains for both losses, probably due to the increased ﬂexibility for learning a diverse ensemble.
Interestingly, only ensembling h ψ also works well for the E NT strategy.
How to ensemble We study the how question by considering four different loss variants: U NIF ,
P ROB , E NT , and the variant of E NT with student entropy weighting.
We ﬁnd that when we ensemble both the projection head h ψ and the codebook µ , the E NT ensemble strategy leads to the most signiﬁcant gains (e.g., 3.4 p.p. gains for 1-shot and 0.9 p.p. gains for full-data).
The P ROB strategy also consistently improves the performance with a slightly larger gain (1 p.p.) in full-data evaluation.
In contrast, we see no gains for the U NIF strategy over the baseline.
We also study a variant of E NT that uses the student entropy (i.e., Eq.
(12) with the log δ
( i − j ) term) for the importance weights (denoted as E NT -S T ).
E NT -S T performs much worse than E NT and is even worse than the baseline.
ȰʌȰ ȰʌȲ Ȱʌȴ Ȱʌȶ Ȱʌȸ ȱʌȰ ȰʌȰ ȰʌȲ Ȱʌȴ Ȱʌȶ Ȱʌȸ ȱʌȰ  ÎÎ ľ Į ´ Î Ř ȷȴ ȷȵ ȷȴʌȴ ȷȴʌȷ ȷȴʌȹ ȷȵʌȱ ȷȵʌȳ ȷȵʌȲ ȷȵʌȳ ȶȳ ȶȴ ȶȵ ȶȳʌȴ ȶȴʌȴ ȶȴʌȷ ȶȵʌȳ ȶȵʌȵ ȶȵʌȳ ȶȵʌȴ ȱ Ȳ ȴ ȸ
ȱȶ ȳȲ ȶȴ ZľċʌĒë(ČĲØċÍĆØĲ ȴȱ ȴȲ ȴȳ ȴȴ ȴȰʌȶ ȴȲʌȱ ȴȳʌȱ ȴȳʌȵ ȴȴʌȰ ȴȳʌȷ ȴȳʌȹ :ľĆĆÔ´Ĺ´
[Figure] [Figure] We conjecture that this is because the student predictions typically have a larger variance than teacher predictions (Wang et al., 2022) especially when multi-crop augmentation (Caron et al., 2020; 2021) is applied to the student.
Similar experiments on Eq.
(11) and/or γ = 0 variants of P ROB also resulted in inferior performance (see Table 12).
Unif
Prob
Ent 0.
Analysis of ( h ψ , µ ) -ensemble diversity The previous experiments showed that the choice of ensemble weighting strategy has a large impact on performance.
We hypothesize that this choice substantially impacts the diversity of the codebook ensembles.
Since the codes in different heads may not be aligned, we align them by the similarity of their code assignment probabilities across different input images, which measures how the codes are effectively used to ‘cluster’ the data.
See Appx.
C. for detailed explanations and results.
In Fig.
4, we visualize the decay patterns of the similarity score between aligned codes (1.0 means the most similar) in a random pair of heads for each weighting strategy.
E NT decays the fastest and U NIF decays the slowest, indicating that E NT learns the most diverse codebooks while U NIF is least diverse.
This shows a positive correlation between the diversity of ( h ψ , µ ) -ensembles and the empirical performance of the ensemble strategies from Table 1.
Finally, for U NIF , we ﬁnd that different heads tend to learn the same semantic mappings even when randomly initialized; i.e., the code assignments in different heads become homogeneous up to permutation.
See Fig.
8 for a visualization.
Number of ( h ψ , µ ) -ensembles
We study the effect of increasing the number of ( h ψ , µ ) -ensembles m for E NT in Fig. 3a.
Having more ( h ψ , µ ) -ensembles boosts the performance until m = 16 .
Interestingly, using as few as m = 2 heads already signiﬁcantly improves over the baseline.
Effect of E NT temperature γ
Fig.
3b studies the effect of entropy weighting temperature γ for different evaluation metrics.
We observe that γ has a relatively larger impact on few-shot evaluation performance.
γ should be neither too high nor too low: a high temperature leads to under-specialization (i.e. less diversity) of heads similar to U NIF ( γ → ∞ ) and a low temperature may otherwise lead to over-specialization (i.e., only a single head is used for each input).
Comparison of different projection heads Our method linearly increases projection head parameters, thus a natural question is: Is the gain of ( h ψ , µ ) -ensembles due to the increased power (or number of parameters) in projection heads?
We answer this question with an empirical study of non-ensembled projection heads.
Speciﬁcally, we studied non-ensembled h ψ with (depth, width) searched over { 2 , 3 , 4 } × { 512 , 1024 , 2048 , 4096 } and measured the linear evaluation performance with different amounts of labeled data.
In Fig.
3c, we plot the 1%-data evaluation result with respect to the number of parameters of the projection head both for ensembled and non-ensembled baselines.
See Appx.
C.2 for detailed analysis and extra results for other metrics.
Our key ﬁndings are:
[Table]
A too powerful non-ensembled h ψ signiﬁcantly hurts the label efﬁciency of learned representations.
This result is similar to Chen et al.
(2020b), which found that probing from intermediate layers of projection heads (which can be viewed as using a shallower head) could improve semi-supervised learning (1%-/10% labeled data) results.
• The default head (3/2048, denoted as ‘Default’) used in recent SSL methods (SimCLRv2, DINO, MSN, etc.) does not perform as well in few-shot evaluations, probably because it is selected by looking at full-data evaluation metrics.
In contrast, our baseline (3/1024, denoted as ‘Our baseline’) signiﬁcantly improves few-shot evaluation performance.
• Our ( h ψ , µ ) -ensembles outperform all non-ensembled baselines and lead to consistent improvements in all evaluation metrics, despite the increase of parameters.
5. I MPROVING SOTA RESULTS WITH ENSEMBLEING
Next we apply ( h ψ , µ ) -ensembles to DINO ∗ and MSN ∗ and compare with the state-of-the-art results.
We experimented with model architectures ViT-S/16, ViT-B/16, ViT-B/8 trained for 800, 400,epochs respectively following Caron et al.
We include both the published results and our retuned versions to ensure strong baselines.
For clarity, we denote our method as “{baseline}- {ensemble strategy} (# of heads)”, e.g., DINO ∗ -E NT (4).
We tuned both baselines and our methods for all architectures.
We report the best hyperparameters for all models in Appx.
B.2.2.
Table 2 compares the results of ( h ψ , µ ) -ensembles and baselines.
We ﬁnd that ( h ψ , µ ) -ensembles with E NT consistently improve all evaluation metrics (full-data, few-shot) across both SSL methods (DINO ∗ , MSN ∗ ) and all architectures (ViT-S/16, ViT-B/16, ViT-B/8) over their non-ensembled counterparts.
The gains in few-shot evaluation are particularly substantial, providing a new state-ofthe-art for ImageNet-1K evaluation from ImageNet pretraining.
5. M ORE EVALUATIONS FOR ( h ψ , µ ) - ENSEMBLES
[Table]
Transfer learning In Table [Table] we compare the transfer learning performance of ( h ψ , µ ) -ensembles and non-ensembled baselines.
We used ViT-S-16 models trained on ImageNet-1K for 800 epochs and evaluated on 9 natural downstream datasets from Chen et al.
(2020a) with linear evaluation (details in
Appx.
B.3).
( h ψ , µ ) -ensembles lead to consistent improvements in transfer performance for MSN ∗ and comparable results for DINO ∗ .
[Table]
Training overhead In Table 4, we benchmark the computational overhead of ( h ψ , µ ) -ensembles at training time.
We used a medium sized model, DINO ∗ with ViT-B/16, trained with the same setting used in all of our experiments.
We benchmarked the wall-clock time and peak memory on 128 TPUv cores.
( h ψ , µ ) -ensembling is relatively cheap in terms of training cost because the ensembled parts typically account for a small portion of total computation, especially when the backbone encoder is more computationally expensive (e.g., ViT-B/8).
Again, we emphasize that there is no evaluation overhead when ( h ψ , µ ) -ensembles are removed.
C ONCLUSION & D ISCUSSION
We introduced an efﬁcient ensemble method for SSL where multiple projection heads are ensembled to effectively improve representation learning.
We showed that with carefully designed ensemble losses that induce diversity over ensemble heads, our method signiﬁcantly improves recent state-ofthe-art SSL methods in various evaluation metrics, particularly for few-shot evaluation.
Although ensembling is a well-known technique for improving evaluation performance of a single model, we demonstrated that, for models with throw-away parts such as the projection heads in SSL, ensembling these parts can improve the learning of the non-ensembled representation encoder and also achieve signiﬁcant gains in downstream evaluation without introducing extra evaluation cost.
Our ensemble method is applicable to many SSL methods beyond the two we explored. For example, one may consider generalization to BYOL (Grill et al., 2020) or SimSiam (Chen & He, 2021) that ensembles projection and/or prediction heads, or MAE (He et al., 2022) that ensembles the decoders (which introduces more training cost though). Our weighted ensemble losses can also be applied as long as the original loss can be reformulated as MLE for some t , s , and Y , e.g., the MSE loss in these methods is MLE under multivariate normal distributions. We hope our results and insights will motivate more future work for extending our method or exploring more ensemble techniques for SSL.
In future work, we also hope to remove three limitations of our setting.
First, considering ensembling strategies that include the representation encoder, r ω , may lead to further improvements in the performance of weighted ensemble SSL, at the cost of increased computation requirements during both training and evaluation.
Second, considering heterogenous architectures in the ensemble may further improve the learned representations (e.g., mixing Transformers with ResNets), whether the heterogeneity is in r ω , h ψ , or both.
Third, considering other possibilities for f ijy may also reveal performance gains and improve our understanding of the critical aspects that lead to good SSL representations, similar to what we learned about the importance of ensemble diversity.
R EPRODUCIBITLITY S TATEMENT
We include detailed derivations for all our proposed losses in Appx.
D. We report experimental details in Appx.
B, including the implementation details for reproducing the baselines (Appx. B.1), training and evaluating our methods (Appx.
B
A P SEUDOCODE
[Table]
[Table]
B E XPERIMENTAL D ETAILS
In this section, we provide details for our experiments.
In Appx.
B.1, we describe how we reproduced and improved the baseline DINO/MSN models.
We give the implementation details for SSL training and evaluation in Appx.
B.2 and Appx.
B.3 respectively.
All the hyper-parameters used in our experiments are in Appx.
B.2.2.
B. R EPRODUCING & I MPROVING B ASELINES
We carefully reproduced and further improved baseline methods (denoted as DINO ∗ and MSN ∗ respectively) with an extensive study and hyperparameter search (see Appx.
B.1).
In particular, we systematically study the projection head design (which we found is crucial for few-shot evaluation performance (Appx. C.2)) and different techniques for avoiding collapse used in both methods (Appx. C.1).
DINO ∗ performs signiﬁcantly better than DINO on few-shot evaluation (e.g., 2 ∼percentage point (p.p.)
gains for 1 shot) and maintains the full-data evaluation performance.
The main adjustments of DINO ∗ are: (i) A 3-layer projection head with a hidden dimension of(instead of 2048);
(ii) Sinkhorn–Knopp (SK) normalization (instead of centering) is applied to teacher predictions, combined with a smaller teacher temperature τ = 0 .
025 and codebook size c = 1024 or 4096.
MSN ∗ uses the same projection head as DINO ∗ and applies ME-MAX regularization without
SK normalization (which is applied in MSN by default).
Further details for DINO and MSN can be found below.
B.1. DINO
[Table]
Reproducing DINO
We carefully reproduced DINO with JAX following the ofﬁcial DINO implementation 1 .
In Table 5, we report the evaluation results of DINO using ViT-S trained withepochs following the exact training conﬁguration for ViT-S/16 in the ofﬁcial DINO code.
The ofﬁcial results of full-data evaluation and 1%-data evaluation are from Caron et al.
(2021), the other few-shot evaluation results are evaluated by Assran et al.
and also validated by us.
Note that for consistency of full-data linear evaluation, we report the results with both the [CLS] token representations of the last layer and the concatenation of the [CLS] token representations from the last 4 layers following Caron et al.
For 1-/2-/5-shots evaluation results, we report the mean accuracy and standard deviation across 3 random splits of the data following Assran et al.
As shown in Table 5, our reproduced results are all comparable with the published numbers which validates the implementation of our training and evaluation pipelines.
Improving DINO We improved the DINO baseline with a systematic empirical study of some important components.
We ﬁrst empirically compared different techniques for avoiding collapse (see Appx.
C.1)
and ﬁnd that Sinkhorn-Knopp (SK) normalization is a more effective and also simpler technique for encouraging codebook usage than the centering operation used in DINO.
We thus applied SK normalization, which enabled us to use a smaller teacher temperature τ = 0 .(instead of τ = 0 . 07 ) and a much smaller codebook size c = 1024 or 4096 (instead of 65536).
These modiﬁcations lead to similar performance as DINO with a much smaller codebook (up to 1M parameters, compared to 16M parameters for DINO).
Next we empirically studied the effect of projection heads for different evaluation metrics (see Appx. C.2), and found that the design of projection heads is crucial for few-shot evaluation metrics and an too power powerful projection head (e.g., the 3-layer MLP with a hidden dimension of 2048 used in DINO/MSN/etc.) could signiﬁcantly hurt the few-shot performance.
With an empirically study of projection head architectures, we found that a simply reducing the hidden dimension to 1024 could signiﬁcantly improves the few-shot evaluation performance while maintaining full-data evaluation performance.
The improved results of
DINO ∗ are shown in Table 5.
B.1. MSN
[Table]
We carefully implemented MSN by adding its main components,
i
[Table]
e., ME-MAX regularization and masking, to the DINO implementation (denoted as MSN ∗ ), which surpassed public results as shown in Table 6.
Note that the implementation of MSN ∗ does not exactly match the public implementation in the public MSN code 2 , where the main differences are:
• MSN applies ME-MAX with Sinkhorn-Knopp normalization by default (as in the released training conﬁguration), which we empirically ﬁnd does not work very well (see Table 9).
MSN ∗ does not apply SK normalization and tunes the regularization strength for ME-MAX.
• Some differences in implementation details, e.g., schedules for learning rate/weight decay, batch normalization in projection heads, speciﬁc data augmentations, etc.
MSN ∗ uses the exact same setup as DINO ∗ which follows original DINO implementation.
We initially tried to exactly reproduce the original MSN following the public MSN code, but the results are much below the public ones, as shown in Table 6.
Incorporating the two differences above bridges the gap and makes MSN ∗ surpass the public results.
B. P RETRAINING DETAILS
In this subsection, we provide the general implementation details in Appx.
B.2.1 and speciﬁc hyperparameters in Appx.
B.2.2 in Appx.
B.2.2 for reproducibility.
B.2. I MPLEMENTATION D ETAILS
Common setup We experimented with DINO (Caron et al., 2021) and MSN (Assran et al., 2022) models on ImageNet ILSVRC-2012 dataset (Deng et al., 2009).
We mainly followed the training setup in Caron et al.
In particular, all models were trained with AdamW optimizer (Loshchilov & Hutter, 2018) and a batch size of 1024.
The learning rate was linearly warmuped to 0.
( = 0.001 × batch size/512)
and followed a cosine decay schedule.
The weight decay followed a cosine schedule from 0.04 to 0.4.
The momentum rate for the teacher was increased from 0.996 to 1 with a cosine schedule following BYOL (Grill et al., 2020).
A stochastic depth (Huang et al., 2016) of 0. was applied without dropout (Srivastava et al., 2014).
The student temperature τ is set to 0.1.
As with DINO, we used the data augmentations of BYOL and multi-crop augmentation of SWAV (Caron et al., 2020).
In particular, 2 global views with a 224 × 224 resolution and crop area range [0.25, 1.0] were generated for the teacher and student, and another 10 local views with 96 × 96 resolution and crop area range [0.08, 0.25] were used as extra augmented inputs for the student.
For MSN, we used the exact same setup and incorporated its major component: 1) mean entropy maximization (ME-MAX) regularization; 2) masking as an extra augmentation applied to the student global view.
Main modiﬁcations We retuned the baselines (DINO ∗ and MSN ∗ ) as detailed in Appx.
B.1, and the main adjustments are as followed.
We used a 3-layer projection head with a hidden dimension of 1024.
The output embedding (i.e., ( h ψ ◦ r ω )( x ) ) and the codes (i.e., µ ) both have a dimension of 256 and are L 2 normalized.
For DINO ∗ , Sinkhorn–Knopp (SK) normalization was applied to teacher predictions.
For MSN ∗ , ME-MAX was used without SK normalization and the regularization strength was tuned over {3, 4, 5}.
For all models, we used teacher temperature τ = 0 .
025 which was linearly decayed from 0.05 for the ﬁrst 30 epochs.
The codebook size c was selected over {1024, 4096} for all models, and typically c = 4096 was selected for baseline methods and c = 1024 was selected for ours.
For our ( h ψ , µ ) -ensembles with E NT , entropy weighting temperature γ is linearly decayed from 0.5 to the speciﬁed value.
B.2. H YPER - PARAMETERS
We report the hyperparameters for training our models for reproducibility:
[Table]
[Figure]
B. E VALUATION P ROTOCALS
Few-shot linear evaluation We followed the few-shot evaluation protocal in Assran et al.
Speciﬁcally, we used the 1-/2-/5-shot ImageNet dataset splits 3 in Assran et al.
(2022) and 1% ( ∼ 13-shot)
ImageNet dataset splits 4 .
For given labelled images, we took a single central crop of size 224 × 224 without additional data augmentations
[Figure]
and extracted the output [CLS] token representations from the frozen pretrained model.
Then we trained a linear classiﬁer with multi-class logistic regression on top of the extracted representations.
We used the scikit-learn package (Pedregosa et al., 2011) for the logistric regression classiﬁer.
For all few-shot evaluations, we searched the L 2 regularization strength over {1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1, 1, 3, 10}.
Full-data linear evaluation We followed the linear evaluation protocal in (Caron et al., 2021).
Speciﬁcally, we trained a linear classiﬁer on top of the representations extracted from the frozen pretrained model.
The linear classiﬁer is optimized by SGD with Nesterov momentum (Nesterov, 1983; Sutskever et al., 2013) of 0.9 and a batch size of 4096 for 100 epochs on the whole ImageNet dataset, following a cosine learning rate decay schedule.
We did not apply any weight decay.
[Table]
During training, we only applied basic data augmentations including random resized crops of size 224 × 224 and horizontal ﬂips.
During testing, we took a single central crop of the same size.
For
ViT-S/16, Caron et al.
(2021) found that concatenating the [CLS] token representations from the last l (speciﬁcally, l = 4 ) layers (c.f. Appendix F.2 in Caron et al.
(2021)) improved the results by about 1 p.p.
We followed the same procedure, but reported linear evaluation results with both l =and l = 4 in Table 2 for consistency.
In our empirical study with ViT-S/16, we used the result with l = 1 .
For larger models (e.g., ViT-B/16), we followed Caron et al.
; Zhou et al.
(2022) to use the concatenation of the [CLS] token representation and the average-pooled patch tokens from the last l = 1 layer for linear evaluation.
For all linear evaluations, we searched the base learning rate over {4.8e-3, 1.6e-2, 4.8e-2, 1.6e-1, 4.8e-1, 1.6}.
Full-data k -NN evaluation We followed the k -NN evaluation protocal in Caron et al.
Wu et al.
Speciﬁcally, for each image in the given dataset, we took a single central crop of size 224 × 224 without additional data augmentations, and extracted the output [CLS] token representations from the frozen pretrained model.
The extracted representations are used for a weighted k -Nearest-Neighbor classiﬁer.
In particular, denote the stored training representations and labels as D = {
( z i , y i ) } Ni =1 .
For a test image with extracted representation z , denote the set of its top k -NN training samples as D k [ z ] ⊆ D and |D k
[ z ] | = k .
The k -NN set D k [ z ] is used to make the prediction for the test image with a weighted vote, i.e., ˆ y = arg max y (cid:16)(cid:80) ( z j ,y j ) ∈D k
[ z ] α j 1 y = y j (cid:17) , where 1 y =
y j is the one-hot vector corresponding to label y j and α j is the weight induced by the cosine similarity between z and z j , i.e., α j = exp (cid:16) 1 τ (cid:48)
z T z j (cid:107) z (cid:107)(cid:107)
z j (cid:107) (cid:17) .
We set τ (cid:48)
= 0 .
07 without tuning as in Caron et al.
; Wu et al.
For all k -NN evaluations, we searched k over {5, 10, 20, 50, 100} and found that k
= 10 or k = 20 was consistently the best.
Transfer evaluation via linear probing We mainly followed the transfer evaluation protocal in (Grill et al., 2020;
Chen et al., 2020a).
In particular, we used 9 of their 13 datasets that are available in tensorflow-datasets (tfd), namely Food-101 (Bossard et al., 2014), CIFAR10 (Krizhevsky et al., 2009), CIFAR100 (Krizhevsky et al., 2009), SUN397 scene dataset (Xiao et al., 2010), Stanford
Cars (Krause et al., 2013), Describable Textures Dataset (Cimpoi et al., 2014, DTD), Oxford-IIIT Pets (Parkhi et al., 2012), Caltech-101 (Fei-Fei et al., 2004), Oxford 102 Flowers (Nilsback & Zisserman,
[Table] [Table]
Following their evaluation metrics
[Table]
we reported mean per-class accuracy for Oxford-IIIT Pets,
Caltech-
[Table]
and Oxford 102 Flowers datasets and reported top-1 accuracy for other datasets.
We transferred the models pretrained on ImageNet (Deng et al., 2009) to these datasets by training a linear classifer on top of frozen representations.
In particular, we resized given images to 256 ×and took a single central crop of size 224 × 224 without additional data augmentations.
We extracted the output [CLS] token representations from the frozen pretrained model.
Then we trained a linear classiﬁer with multi-class logistic regression on top of the extracted representations.
We used the scikit-learn package (Pedregosa et al., 2011) for the logistric regression classiﬁer.
For all transfer evaluations, we searched the L 2 regularization strength over {1e-6, 1e-5, 1e-4, 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, 3e-1, 1, 3, 1e1, 3e1, 1e2, 1e3, 1e4, 1e5}.
C A DDITIONAL R ESULTS C. E MPIRICAL STUDY OF TECHNIQUES FOR AVOIDING COLLAPSE
Most self-supervised learning methods utilize some techniques to avoid collapse of representations with, e.g., contrastive loss (Chen et al., 2020a;
He et al., 2020), batch normalization (Grill et al., 2020), asymmetric architecture design with a predictor (Grill et al., 2020; Chen & He, 2021), etc.
In
DINO and MSN, a learnable codebook is used for the learning objective and different techniques are applied to encourage the effective codebook usage.
There are two potential cases of collapse (as discussed in Caron et al.
• Dominating codes .
This is the case of “winner-take-all”: only a small portion of codes are being predicted while others are inactive.
Typical solutions for avoiding this include applying Sinkhorn–Knopp normalization (Cuturi, 2013) as in SWaV (Caron et al., 2020), centering teacher logits as in DINO (Caron et al., 2021), and applying mean-entropy maximization regularization (ME-MAX) as in MSN (Assran et al., 2022).
Note that in MSN, ME-MAX is combined with Sinkhorn–Knopp normalization by default.
• Uniform codes .
This is the case where all codes are treated equally and the predictions reduce to be uniform over codes.
A simple and effective solution is to applying sharpening, i.e., using a lower temperature for computing the teacher prediction.
We systematically study different techniques in a uniﬁed setup.
In particular, we used DINO with the
ViT-S backbone, a 3-layer MLP projection head with hidden dimension 2048, and a codebook of size 4096 and dimension 256.
We applied different techniques to DINO and searched the teacher temperature in { 0 . 0125 , 0 . 025 , 0 . 05 } for each.
For ME-MAX, we searched regularization weight in
[Figure]
Ȱʌȶ Ȱʌȷ Ȱʌȸ Ȱʌȹ ȱʌȰ Z Ē Į ċ
´ Ć ô Š Ø Ô
Y
Ø Ĺ Į ô Î :ľĆĆʲÔ´Ĺ´ ȱ̑Ô´Ĺ´ ȱʲĲñĒĹ $Øë´ľĆĹ `ľĮÍ´ĲØ
[Figure]
For ME-MAX combined with Sinkhorn, we followed Assran et al.
and used default regularization weight of 1.0.
The results are in Table 10.
We observed that:
• DINO’s centering operation is not as strong as other techniques, and it favours a larger teacher temperature (e.g., 0.05).
It does not work well when the codebook size (4096) is not as large as the one used in the original DINO model (65536).
Switching to use Sinkhorn–Knopp normalization leads to much more improved performance, and matches the performance of original DINO (Table 5) with a much smaller codebook.
• MSN’s ME-MAX regularization is very effective, and leads to signiﬁcant improvements over others.
We also found it is sensitive to the regularization weight and teacher temperature (c.f. Table 10).
However, we observed that combining ME-MAX with Sinkhorn does not work well without tuning the regularization weight (which is recommended by Assran et al.
C. E MPIRICAL S TUDY OF P ROJECTION H EADS
In this subsection, we systematically study the effect of projection heads for different evaluation metrics.
In particular, we used DINO ∗ ViT-S/16 as the base model and used different projection heads with (depth, width) searched over { 2 , 3 , 4 } × { 512 , 1024 , 2048 , 4096 } .
All models are trained with 300 epochs using exact the same set of hyper-parameters.
We measured the linear evaluation performance with different amount of labeled data (i.e., full-data, 1%-data, 1-shot).
In Fig.
5a, we plot different evaluation metrics (normalized respectively by the best of each) versus the number of projection head parameters.
In Figs.
5b to 5d, we plot each unnormalized evaluation metric respectively for different heads as well as our ( h ψ , µ ) -ensembles.
Our key ﬁndings are:
[Figure]
• The projection head has a relatively larger impact on few-shot evaluation metrics, as reﬂected by the relative magnitudes of different metrics in Fig.
5a.
An too powerful non-ensemble projection head signiﬁcantly hurts the label efﬁciency of learned representations, reﬂected by a much larger drop in few-shot evaluation performance (up to 18 p.p. for 1-shot, 9 p.p. for 1%-data).
This result is also partially observed in Chen et al.
(2020b), where they found that probing from intermediate layers of projection heads (which can be viewed as using a shallower head) could improve the semi-supervised learning (1%-/10%) results.
• The optimal projection head for different metrics can differ a lot.
A weaker head improves label efﬁciency (few-shot performance), while a stronger (but not too strong) head improves linear decodability.
As a result, the default projection head (3/2048) that is widely used in SimCLR v2 (Chen et al., 2020b), DINO (Caron et al., 2021), iBOT (Zhou et al., 2022), MSN (Assran et al., 2022), etc., does not perform well in few-shot evaluations (as shown by the green cross denoted as ‘Default’), probably because it is selected by full-data evaluation metrics.
• There exist some projection heads that performs decently well on all evaluation metrics, e.g., the baseline model (3/1024) used in our experiments (pink star denoted as ‘Our base’).
• Compared to naively tuning projection head architectures, our ( h ψ , µ ) -ensembles (orange curves in Figs.
5b to 5d) consistently improve all metrics with different amount of labeled data, despite it also increases the number of parameters in projection heads.
Our ( h ψ , µ ) - ensembles outperform all non-ensembles, which also include the counterparts of probing from intermediate layers from the a deeper head (i
e., shallower heads).
C. E MPIRICAL STUDY OF ( h ψ , µ ) - ENSEMBLES
Are the gains of E NT purely from sharper teacher predictions?
Our E NT strategy assigns higher weights to the heads that predict with lower entropies, thus effectively uses sharper teacher predictions as the targets.
One may be curious about how this effect accounts for the gains of the E NT strategy.
We empirically answer this question by studying the non-ensemble baseline that uses a sharper teacher predictions in a data-independent manner (in contrast to E NT , which uses data-dependent entropy weights).
Speciﬁcally, we compare the non-ensemble DINO ∗ that use different teacher temperature τ ∈ { 0 . 005 , 0 . 01 , 0 .
025 , 0 . 05 } and also our DINO ∗ -E NT (16) with τ = 0 .
025 , as shown in Fig.
We ﬁnd that the teacher temperature has a big impact on evaluation results especially for few-shot evaluation.
Compared to our default baseline that uses τ = 0 .
025 , a lower temperature (e.g., τ = 0 .
01 ) can indeed improve the 1-shot performance (at the cost of worse full-data performance).
However, an too low temperature ( τ = 0 . 005 ) will hurt the performance.
Our DINO ∗ -
[Equation]
16) consistently outperform all the baselines, which implies the importance of selecting sharper teacher predictions in a data-dependent manner.
[Table]
[Table]
Comparison of different ensemble strategies and variants We present the full table of Tablethat includes all the metrics in Table 11.
The same observation holds for all metrics
[Table]
For all previous studies, we considered a speciﬁc instantiation of P ROB strategy, i.e., weight by student predicted probabilities f ijy
= log s ( y | θ j , x ) and γ = 1 , which has a nice interpretation of model average (see Sec. 3.3).
We also studied different variants of the P ROB strategy (see Appx.
D.1),
• P ROB -T E : weight by teacher f ijy = log
t i ( y | x ) and γ = 1 ; • P ROB -M AX : weight by student f ijy = log s j ( y | x ) and γ → 0 ; • P ROB -M AX -T E : weight by teacher f ijy = log t i ( y | x ) and γ →
Table 12 compares the downstream performance for all the variants.
We ﬁnd that the our P ROB (used in our empirical studies) performs better than other variants.
Interestingly, weighting by the teacher (P ROB -T E ) performs worse than P ROB .
We conjecture that this is because the important weights turn out to give a weighted average of teacher predictions as the surrogate target that is shared across all students (like P ROB ) but does not give effective preferential treatment across students which are directly optimized (unlike P ROB -T E ).
Furthermore, P ROB -M AX which sharpens the importance weights leads to training divergence.
This is probably because the student predictions have higher variance based on which sharp weights lead to unstable training.
In contrast, P ROB -M AX -T E which uses the (lower-variance) teacher gives reasonable results and comparable to P ROB -T E .
Number of ensembles for MSN ∗
In Fig.
7a, we study the effect of increasing the number of ( h ψ , µ ) -ensembles for MSN ∗ -E NT with ViT-S/16 trained for 800 epochs.
The scaling trend is similar to DINO ∗ -E
NT (Fig. 3a) and the gains start to diminish when the number of heads increases above 8.
Effect of E NT temperature γ for MSN ∗
Fig.
7b studies the effect of entropy weighting temperature γ for MSN ∗ -E NT .
We observed that MSN ∗ is more robust to small temperatures, and the
ȰʌȰ ȰʌȲ Ȱʌȴ Ȱʌȶ Ȱʌȸ ȱʌȰ ȰʌȰ ȰʌȲ Ȱʌȴ Ȱʌȶ Ȱʌȸ ȱʌȰ  ÎÎ ľ Į ´ Î Ř ȷȶ ȷȷ ȷȵʌȶ ȷȶʌȰ ȷȶʌȳ ȷȶʌȴ ȷȶʌȶ ȶȷ ȶȸ ȶȷʌȱ ȶȷʌȹ ȶȸʌȳ ȶȸʌȷ ȶȸʌȶ ȱ Ȳ ȴ ȸ
ȱȶ ZľċʌĒë(ČĲØċÍĆØĲ ȴȷ ȴȸ ȴȹ ȵȰ ȴȷʌȴ ȴȸʌȸ ȴȹʌȳ
ȵȰʌȱ ȴȹʌȹ :ľĆĆÔ´Ĺ´ ȱ̑Ô´Ĺ´ ȱʲĲñĒĹ (
[Figure]
) Scaling of ensembles 0.
A cc u r a c y 76 76.
0.07 0.1 Entropy Weight Temperature 49 48.
48.5 48.3 Full data 5-shot 1-shot Baseline
[Figure]
best γ = 0 .
01 is smaller than that of DINO ∗ ( γ = 0 . 05 ).
When the temperature is too high, the performance drops as a result of under-specialization (i.e., less diversity) as with DINO ∗ .
C. A NALYZING ( h ψ , µ ) - ENSEMBLE DIVERSITY
Visualizing ( h ψ , µ ) -ensemble similarity
We analyze the diversity between different heads by visualizing the similarity matrix between their codes.
Directly measuring the similarity between codes in two heads could not work, because 1) they may live in different subspaces because of the ensembled projection heads; 2) they may not align in the natural order but in a permuted order.
Therefore, we seek to align codes between different heads by how they are effectively used to ‘cluster’ the data.
In particular, we use a set of randomly sampled inputs (cid:8) x
i (cid:9) i ∈ [ b ] of size b = 51200 to obtain an empirical code assignment matrix A j ∈ R b × c for each ( h ψ , µ ) -ensemble j ∈ [ m ] , where the i -th row of A j corresponds to the teacher predictions t j ( Y | x i ) .
For the k -th code in the head j , we extract the k -th column from A j (i.e., its empirical assignment) as its embedding.
For two codes, we measure their similarity by the cosine similarity between their embeddings.
For a pair of heads j and j (cid:48) , we align their codes using the Hungarian algorithm (Kuhn, 1955) to maximize the sum of cosine similarity.
After that, we plot the similarity matrix which is aligned and reordered by the similarity value on the diagonal (in an descending order).
Note that it is not necessary to do the alignment procedure for the P ROB strategy since it is naturally aligned because of the direct distribution averaging over ( h ψ , µ ) -ensembles, but we did for fair comparison with other strategies.
We applied the same procedure for different ensemble weighting strategies using DINO ∗ with( h ψ , µ ) -ensembles.
We randomly picked a pair of heads and visualize the similarity matrix before (top row) and after (bottom row) the alignment-reordering setup in Fig.
We found that before the alignment procedure, the similarity matrix of the P ROB strategy already mostly aligns because it explicitly introduces code correspondence between different heads.
Furthermore, by analyzing the similarity decay pattern on the diagonal, it is clear that E NT learns the most diverse ( h ψ , µ ) -ensembles while U NIF learns the least ones, which may explain the difference of their empirical performance.
For completeness, we also include the visualization of aligned similarity matrices for all pairs of heads in Figs.
9 to 11, the observations are the same.
[Figure] [Figure]
[Figure] [Figure]
[Figure]
D A NALYSIS D. D ERIVATIONS
In this subsection, we provide derivations for some non-trivial losses that we explore within our framework.
Recall that our weighted cross-entropy loss is of the form,
[Equation]
where w ijy = softmax [Equation] Fuethermore, observe that,
[Equation]
This indicates that the proposed weighted ensemble SSL loss is simply a reweighted log-likelihood loss.
We use this fact in our derivation of probability weighting (P ROB ) loss.
Uniform weighting (U NIF )
Our U NIF strategy in Eq.
(6) uses f ijy =
log δ
( i − j ) which gives w ijy = 1 m δ ( i − j ) (for any choice of γ ), thus the loss,
[Equation]
This loss assigns equal weights to m pairs of pairwised student/teacher.
An straightforward generalization is to assign equal weights to all possible pairs ( m 2 ) of student/teacher with f ijy = 0 and w ijy = 1 m 2 , which gives the U NIF -A LL loss,
[Equation]
Probability weighting (P ROB )
Recall our P ROB loss in Eq.
(7) has the form,
[Equation]
We derive its equivalence with our general loss with f ijy = log s
( y | θ j , x ) and γ = 1 in terms of the gradients,
[Equation]
[Equation]
) where w ijy =
s ( y |
[Equation]
j ,x ) [Equation] [Equation] ,x ) (or equivalently, f ijy = log s ( y | θ j , x ) and γ = 1 ).
The last equality is because w ijy is stopped gradient with respect to θ .
This is the same analysis as done in Burda et al.
The above formation establishes the equivalence of gradients between two losses, which implies the same behavior (e.g., optimum) using gradient-based optimization, as the common practice of deep learning.
We also generalize this loss to some variants which we explore in Table 12.
A “dual” variant is to use teacher predictions f ijy = log t
i ( y | x ) instead of student ones; this implies w ijy =
[Equation]
) and the P ROB -T E loss,
[Equation]
Note that this simply reduces to use a weighted teacher predictions
[Equation]
x ) as the surrogate target that is shared across all students.
Another generalization is to use “hard” weighting, i.e., γ → 0 , which gives the P ROB -M AX loss that only assigns weight to the most conﬁdent student,
[Equation]
where w ijy
[Equation]
δ
[Equation]
This loss reduces to a generalization of multiple choice learning (Guzman-Rivera et al., 2012) used in multi-headed networks (Lee et al., 2015) in our ensemble SSL setup.
Similarly we can also derive the dual variant of it that uses the teacher predictions, which is omitted here for brevity.
Entropy weighting (E NT )
The derivation of E NT loss in Eq.
(9) is similar to the U NIF loss but applies an entropy weights.
Recall that we use f ijy = − H
[ t i ( Y | x )]
+ log δ ( i − j ) , which gives w ijy = [Equation] )] : i [Equation]
∈ [ m ] } ) and,
[Equation]
One can also generalizes it to its dual variant which uses the student entopies, i.e., f ijy = − H
[ s ( Y | θ j , x )]
+ log δ ( i − j ) , which gives the E NT -S T loss,
[Equation] NT
[Equation]
D. R ELATING SOME LOSSES
Here, we relate some losses derived above.
Speciﬁcally, we relate the uniform weighting (U NIF , U NIF -A LL ) and probability weighting (P ROB ) in Appx. D.2.1, and relate entropy weighting (E NT ) and variance weighting in Appx. D.2.2.
We ﬁrst establish the relation between U NIF and P ROB using the joint convexity of unnormalized KL divergence and the fact that our weighted cross-entropy loss is a weighted unnormalized KL divergence up to some constant in θ . In particular, the joint convexity of unnormalized KL divergence can be shown by combining the facts that Csiszàr f -divergences are jointly convex (Proposition 1 in
Dragomir (2013)) and unnormalized KL divergence corresponds to the convex generator, f ( u ) = u log u − u + 1 , as required by the proposition.
First, our weighted cross-entropy loss is unnormalized KL divergence up to some constant in θ :
[Equation]
Therefore, the joint convexity of (unnormalized) KL divergence directly implies an ordering of the loss up to some constant in θ , i.e.,
[Equation]
Furthermore, we can also relate P ROB and U NIF -A LL using the fact that the (unnormalized) crossentropy H × [ p ( X ) , q ( X )] is linear in the ﬁrst argument p but convex in the second argument q , which implies, [Equation]
D.2. E NTROPY & V ARIANCE W EIGHTING
Suppose p ( X ) is a discrete distribution (normalized) on X
=
[ c ] .
It can be shown that,
[Equation]
where V ar p
[ X ] = (cid:80) x ∈ [ c ] p ( x )( x − µ ) 2 and µ = E p [ X ] = (cid:80) x ∈ [ c ] p ( x ) x (Theorem 9.7.1, Cover &
Thomas (1999)).
Note, a tighter bound (Mow, 1998) also exists but it places stronger restrictions on p .
This relationship suggests that choosing weights proportional to exp( − H [ t i ( Y | x )])
(as in E NT ) is potentially related to choosing weights proportional to weighting by variance ( V ar [Equation] i
[Equation] ]+ [Equation] where ( (cid:15) = 112 ) .