PaperlayerView exact revision on arXiv

Source-linked AI summary

Weighted Ensemble Self-Supervised Learning

Yangjun Ruan, Saurabh Singh, Warren Morningstar, Alexander A. Alemi, Sergey Ioffe, Ian Fischer, Joshua V. Dillon

arXiv:2211.09981v3cs.LGcs.AIstat.ML

TL;DR

In SSL, it is less clear where and how to ensemble because the aim is to learn useful representations from unlabeled data. The paper ensembles non-representation components during training using data-dependent weighted losses, and reports improved evaluation performance, particularly in few-shot settings.

  • Problem

    Because SSL aims to learn useful representations from unlabeled data, it is less obvious than in supervised learning where and how to ensemble.

    Source: §1 · p. 1

  • Method

    The method ensembles non-representation components during training and uses data-dependent importance weights in weighted cross-entropy losses to train a single representation encoder.

    Source: §1 · p. 1, §3.2 · p. 3

  • Results

    The ensembles consistently improve all evaluation metrics across the reported SSL methods and architectures, with particularly substantial few-shot gains on ImageNet-1K.

    Source: §5.2 · p. 9, §6 · p. 9

  • Takeaways & Limitations

    Ensembling throw-away SSL components can improve representation learning and downstream evaluation without extra evaluation cost.

    Source: §6 · p. 9

    The explored setting leaves encoder ensembling untested; including the representation encoder may improve performance but would increase training and evaluation computation.

    Source: §6 · p. 9

Abstract

from arXiv · show

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 efficient 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 find 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.

1 INTRODUCTION

The paper adapts ensembling to SSL by training multiple non-representation components to improve a single representation encoder, with data-dependent weighted cross-entropy losses. It reports gains over improved prior-art baselines and finds that greater ensemble-head diversity correlates with better downstream performance.

1 INTRODUCTION

  • The method replicates non-representation components, such as projection heads, and uses them during training to improve one representation encoder without downstream evaluation overhead.The ensemble is removed for downstream evaluation.Source: §1 · p. 1
  • Data-dependent importance weights in weighted cross-entropy losses support different ensemble weighting schemes; greater head diversity correlates with improved downstream performance.Source: §1 · p. 1
  • Up to 5.5 p.p. gain for 1-shot learning is reported for ensembling over improved prior-art baselines.The contribution summary cites Table 2 for this gain.Source: §1 · p. 2
  • Up to 6.3 p.p. improvement over prior-art baselines is reported from extensive ablation studies.Source: §1 · p. 2

2 BACKGROUND

This section frames SSL through maximum likelihood estimation and describes clustering-based student-teacher methods. It also sets out the teacher assumption used in its SSL perspective and motivates unnormalized quantities for deriving weighted ensemble losses.

2 BACKGROUND

  • The section frames SSL methods from a maximum-likelihood perspective as a basis for describing clustering-based methods and deriving ensemble variants.Source: §2 · p. 2
  • Unnormalized cross-entropy and entropy define an unnormalized KL divergence that enables flexible distribution weighting in the derivation of weighted ensemble losses.Source: §2 · p. 2
  • The SSL viewpoint treats the oracle conditional target as optimal but assumes the empirical target distribution is inaccessible or unreliable; the learned teacher need not recover the oracle.Source: §2 · p. 2
  • Clustering-based SSL uses a student whose encoder produces downstream representations, while projection heads and codebook entries characterize the training loss and are discarded at evaluation.Source: §2 · p. 2
  • The student assigns inputs to codebook centroids through soft clustering, while the teacher uses a stopped-gradient model and commonly an exponential moving average with a lower temperature.Source: §2 · p. 2
  • Data augmentation and regularization are described as essential for desirable invariances and preventing degeneracy, though omitted from the discussion as not directly relevant to the method.Source: §2 · p. 2

3 METHOD

The method ensembles SSL components during training while keeping the representation encoder shared, and uses data-dependent weighted cross-entropy objectives to train the heads. It explores weighting schemes that vary how teacher and student predictions interact.

3.1 WHERE TO ENSEMBLE?

  • Teacher and student heads are independently initialized, while teachers have zero gradients and all students share one temperature and all teachers another.Heterogeneous architectures and different teacher parameterizations are left for future work.Source: §3.1 · p. 3
  • The method ensembles projection heads and/or codebooks, not encoders, so the shared representation encoder is learned without extra downstream evaluation cost.The encoder ensemble is left for future work.Source: §3.1 · p. 3, Figure 2 · p. 3

3.2 HOW TO ENSEMBLE?

  • The weighted ensemble objective combines teacher–student cross-entropies using event-specific, data-dependent weights, with weight gradients stopped and each data point normalized to contribute equally.The objective is an importance-weighted mixture of unnormalized KL divergences and allows interactions across teacher/student pairs.Source: §3.2 · p. 3

3.3 HOW TO WEIGHT?

  • UNIF averages teacher/student pairs independently, PROB averages predictive distributions to couple heads, and ENT gives larger weights to lower-entropy teacher predictions.PROB favors confident student predictions; ENT emphasizes opinionated teacher heads for different inputs.Source: §3.3 · p. 4
  • The weighting family also includes equal-pair, disagreement-based, and low-variance-teacher schemes, while aligned variants of all schemes are possible.Early experiments with Eqs. (11) and (12) were inferior and largely omitted; many other schemes were left unexplored because of resource constraints.Source: §3.3 · pp. 4–5

4 RELATED WORK

Prior work spans SSL representation-learning approaches and ensemble techniques for supervised, semi-supervised, and SSL settings. This method differs from post-training ensembles by applying shared-encoder ensembling during SSL training.

Self-supervised learning

  • SSL methods include contrastive, clustering, bootstrapping, self-distillation, asymmetric-architecture, redundancy-reduction, and masked-image-modeling approaches.The authors note their method may apply to approaches with a projection head, including most discriminative approaches.Source: §4 · p. 5

Ensemble methods

  • Prior ensemble research studies performance and uncertainty, efficient partial-parameter-sharing ensembles, and post-training ensembles adapted from SSL models.The cited adapted-model approaches find benefits from higher divergence or weight averaging, while this method targets representation learning during SSL training and is complementary to post-training methods.Source: §4 · p. 5

5 EXPERIMENTS

Experiments compare DINO and MSN with original and retuned baselines, then assess ensemble performance on ImageNet-1K, transfer learning, and training overhead. The empirical study tests ensemble placement, weighting, head count, temperature, and projection-head capacity.

5.1 EMPIRICAL STUDY OF (hψ, µ)-ENSEMBLES

  • Ensembling both the projection head and codebook gives the largest gains for ENT and PROB; ensembling only the projection head also works well with ENT.Source: §5.1 · p. 6
  • ENT gains 3.4 p.p. in 1-shot and 0.9 p.p. in full-data evaluation; PROB adds 1 p.p. in full-data, while UNIF adds none and ENT-ST underperforms baseline.The comparisons are for ensembles of both the projection head and codebook.Source: §5.1 · p. 6
  • ENT learns the most diverse codebooks and UNIF the least diverse; ensemble diversity positively correlates with empirical performance.Similarity is measured between aligned codes across heads.Source: §5.1 · p. 7
  • With ENT, two heads already improve over baseline and gains continue until 16; ensembles also outperform non-ensembled heads when controlling for parameters.A too-powerful non-ensembled projection head significantly harms label efficiency.Source: §5.1 · pp. 7–8, Figure 3 · p. 7
  • Entropy-weighting temperature has a larger impact on few-shot results; values that are too high reduce head diversity, while values too low can over-specialize heads.Source: §5.1 · p. 7

5.2 IMPROVING SOTA RESULTS WITH ENSEMBLEING

  • Ensemble heads consistently improve ImageNet-1K downstream metrics across DINO*, MSN*, and ViT architectures, setting a new state of the art for few-shot evaluation.Table 2 covers full-data and few-shot evaluation.Source: Table 2 · p. 8, §5.2 · p. 9

5.3 MORE EVALUATIONS FOR (hψ, µ)-ENSEMBLES

  • On nine natural downstream datasets using ViT-S/16, ensembles consistently improve transfer performance for MSN* and give comparable results for DINO*.Source: §5.3 · p. 9, Table 3 · p. 9
  • Training overhead is relatively low because ensembled components account for a small share of computation, and removing them adds no evaluation overhead.The overhead benchmark uses DINO* ViT-B/16 on 128 TPUv3 cores.Source: §5.3 · p. 9, Table 4 · p. 9

6 CONCLUSION & DISCUSSION

The paper ensembles projection heads during training to improve representation learning, with no additional downstream evaluation cost. It identifies broader ensemble strategies and reduced computation as directions for future work.

6 CONCLUSION & DISCUSSION

  • Ensembling multiple projection heads with diversity-inducing losses improves SSL evaluation, particularly few-shot performance, without adding downstream evaluation cost.The heads help train a single representation encoder and are not used at evaluation.Source: §6 · p. 9
  • The approach may extend to BYOL or SimSiam projection/prediction heads and MAE decoders, although decoder ensembling adds training cost.The weighted losses also apply when the original loss can be reformulated as maximum likelihood, including MSE under multivariate normal distributions.Source: §6 · p. 9
  • Ensembling the representation encoder may improve performance but would increase computation during both training and evaluation.The authors also identify heterogeneous ensemble architectures and alternative fijy choices as future directions.Source: §6 · p. 9

A PSEUDOCODE

  • The pseudocode forms student and teacher log probabilities for each head, stacks them, computes an ensemble loss, then updates the student by gradients and teacher by momentum.Its entropy strategy weights per-head losses using teacher entropy.Source: A PSEUDOCODE · p. 15

B EXPERIMENTAL DETAILS

The appendix describes experiments and baseline reproduction, then reports improvements to DINO and MSN baselines. It attributes DINO’s few-shot gains to projection-head and collapse-avoidance changes.

B EXPERIMENTAL DETAILS

  • The appendix points to B.2 for SSL training and evaluation implementation details and hyperparameters.B.2.1 contains implementation details and B.2.2 lists hyperparameters.Source: B EXPERIMENTAL DETAILS · p. 16, §B.2 · p. 17

B.1 REPRODUCING & IMPROVING BASELINES

  • DINO* gains 2∼6 p.p. for 1-shot evaluation while maintaining full-data evaluation performance.Table 5 reports DINO/DINO* ViT-S/16 results after 800 epochs.Source: §B.1 · p. 16, Table 5 · p. 16

B.1.1 DINO

  • The reproduced DINO results were comparable with published numbers, validating the authors’ training and evaluation pipelines.Table 5 covers ViT-S trained for 800 epochs.Source: §B.1.1 · p. 16, Table 5 · p. 16
  • DINO* uses SK normalization with a smaller teacher temperature and codebook, while a 1024-hidden-dimension projection head improves few-shot performance without reducing full-data performance.The smaller codebook is 1024 or 4096 rather than 65536; the cited study reports similar performance with up to 1M rather than 16M parameters.Source: §B.1.1 · pp. 16–17, §B.1 · p. 16

B.1.2 MSN

  • MSN* surpasses public MSN results after adding ME-MAX and masking to DINO*; omitting SK normalization and tuning ME-MAX strength bridge the reproduction gap.The authors report that an exact reproduction following public MSN code fell well below public results.Source: §B.1.2 · p. 17, Table 6 · p. 17

B.2.1 IMPLEMENTATION DETAILS

The implementation uses a shared DINO-style training setup for DINO and MSN, with specified modifications for each baseline and the ensemble.

B.2.1 IMPLEMENTATION DETAILS

  • The common setup uses AdamW, batch size 1024, cosine learning-rate and weight-decay schedules, and two global plus ten student-local augmented views.MSN adds ME-MAX regularization and masking to the student global view.Source: §B.2.1 · p. 17
  • DINO* and MSN* share a 3-layer, 1024-hidden-dimension projection head, 256-dimensional normalized embeddings and codes, and teacher temperature τ = 0.025.DINO* uses SK normalization; MSN* uses ME-MAX without SK, with regularization strength selected from {3, 4, 5}.Source: §B.2.1 · p. 18

B.2.2 HYPER-PARAMETERS

  • The appendix reports DINO* training hyperparameters for reproducibility.Table 7 is identified as the DINO* training hyperparameter table.Source: §B.2.2 · p. 18, Table 7 · p. 18

B.3 EVALUATION PROTOCALS

Evaluation uses frozen pretrained representations with linear classifiers or weighted k-nearest-neighbor voting across few-shot, full-data, and transfer settings. The appendix also reports findings on collapse-avoidance techniques and ME-MAX sensitivity.

B.3 EVALUATION PROTOCALS

  • Few-shot evaluation trains logistic regression on frozen [CLS] representations from the 1-, 2-, and 5-shot and 1% ImageNet splits.The protocol uses a central crop without additional augmentation and searches L2 regularization strength.Source: §B.3 · p. 18
  • Full-data linear evaluation trains a linear classifier on frozen representations, using SGD with Nesterov momentum for 100 epochs on ImageNet.For larger models such as ViT-B/16, the feature uses the [CLS] representation concatenated with average-pooled patch tokens from the last layer.Source: §B.3 · pp. 18–19
  • Full-data k-NN evaluation predicts by weighted voting over top-k training representations, with k searched over {5, 10, 20, 50, 100}.The authors found k = 10 or k = 20 consistently best.Source: §B.3 · p. 19
  • The collapse-avoidance study reports improved DINO performance with SK instead of centering, while Table 10 identifies ME-MAX as hyperparameter-sensitive.The table caption also reports that ME-MAX significantly improves few-shot evaluation.Source: Table 9 · p. 20, Table 10 · p. 20
  • Transfer evaluation uses linear probes on frozen ImageNet-pretrained representations; mean per-class accuracy is reported for Pets, Caltech-101, and Flowers, and top-1 accuracy for other datasets.The evaluation protocols draw on nine datasets available in TensorFlow Datasets.Source: §B.3 · pp. 19–20

C.1 EMPIRICAL STUDY OF TECHNIQUES FOR AVOIDING COLLAPSE

The study compares techniques for avoiding codebook collapse in DINO and MSN. Sinkhorn–Knopp and MSN’s ME-MAX regularization perform strongly, while the outcomes depend on the technique and its settings.

C.1 EMPIRICAL STUDY OF TECHNIQUES FOR AVOIDING COLLAPSE

  • The study distinguishes dominating-code collapse, where few codes are predicted, from uniform-code collapse, where predictions become uniform.The cited approaches include Sinkhorn–Knopp for dominating codes and lower teacher temperature for uniform codes.Source: §C.1 · p. 20
  • The DINO comparison used a ViT-S backbone, a 3-layer projection head, a 4096-entry codebook, and searched teacher temperatures from 0.0125 to 0.05.ME-MAX regularization weights were also searched; the ME-MAX-plus-Sinkhorn setting used the default weight of 1.0.Source: §C.1 · pp. 20–21
  • Sinkhorn–Knopp substantially improves DINO with a 4096-code codebook, matching original DINO, while MSN’s ME-MAX is highly effective but sensitive to regularization weight and teacher temperature.Combining ME-MAX with Sinkhorn does not work well without tuning the regularization weight.Source: §C.1 · p. 21

C.2 EMPIRICAL STUDY OF PROJECTION HEADS

Projection-head architecture affects evaluation metrics differently, particularly few-shot performance. The study finds that the proposed ensembles improve evaluation across labeled-data amounts, while very powerful non-ensemble heads can hurt label efficiency.

C.2 EMPIRICAL STUDY OF PROJECTION HEADS

  • For DINO* ViT-S/16 trained for 300 epochs, a lower teacher temperature can improve few-shot performance, yet it still underperforms the DINO*-ENT ensemble with 16 heads.The cited comparison uses τ = 0.025 for DINO*-ENT.Source: Figure 6 · p. 22
  • An overly powerful non-ensemble head reduces performance by up to 18 p.p. for 1-shot and 9 p.p. for 1%-data evaluation.The study reports that projection-head choice has a relatively larger impact on few-shot metrics.Source: §C.2 · p. 22
  • A weaker head improves few-shot label efficiency, whereas a stronger but not overly strong head improves linear decodability, so the default 3/2048 head performs poorly in few-shot evaluation.The 3/1024 baseline used in the experiments performs decently across all evaluation metrics.Source: §C.2 · p. 22
  • The (hψ, µ)-ensembles outperform all tested non-ensemble heads across evaluation metrics and labeled-data amounts, including intermediate-layer probing.This holds despite the ensembles increasing projection-head parameters.Source: §C.2 · p. 22

C.3 EMPIRICAL STUDY OF (hψ, µ)-ENSEMBLES

The empirical studies compare ensemble strategies, weighting variants, head counts, and head diversity. ENT and PROB improve over the non-ensemble baseline, whereas UNIF does not; ENT also learns the most diverse ensembles.

C.3 EMPIRICAL STUDY OF (hψ, µ)-ENSEMBLES

  • DINO*-ENT with 16 heads consistently outperforms non-ensemble baselines, so its gains are not explained solely by using sharper teacher predictions.Lowering the non-ensemble teacher temperature can improve 1-shot but worsens full-data performance; τ = 0.005 hurts performance.Source: §C.3 · p. 22
  • ENT and PROB significantly improve over the non-ensemble baseline, while UNIF provides no gains; ensembling the whole projection head works best.The results report means and standard deviations across three initialization seeds.Source: §C.3 · p. 23, Table 11 · p. 23
  • The tested PROB strategy performs better than its variants; PROB-TE performs worse, and PROB-MAX diverges during training.The authors conjecture that sharp weights based on higher-variance student predictions make PROB-MAX unstable.Source: §C.3 · p. 23, Table 12 · p. 23
  • For MSN*-ENT, gains from adding ensemble heads begin to diminish above 8 heads.This scaling trend is reported for ViT-S/16 trained for 800 epochs.Source: §C.3 · p. 23, Figure 7 · p. 24
  • MSN* is more robust to small entropy-weighting temperatures than DINO*, with best γ = 0.01 versus DINO*’s γ = 0.05.When γ is too high, performance drops alongside reduced ensemble diversity.Source: §C.3 · pp. 23–24

C.4 ANALYZING (hψ, µ)-ENSEMBLE DIVERSITY

  • Code similarity is measured after aligning code assignments across heads; ENT learns the most diverse ensembles and UNIF the least.The method uses empirical teacher assignment embeddings, cosine similarity, and Hungarian alignment.Source: §C.4 · p. 24, Figure 8 · p. 25

D ANALYSIS

The analysis derives variants of the weighted ensemble loss and relates several weighting schemes. It establishes loss relationships through convexity and connects entropy weighting with variance weighting under a discrete-distribution assumption.

D.1 DERIVATIONS

  • The weighted ensemble loss is a reweighted log-likelihood, and PROB weighting has gradients equivalent to its general-loss formulation when weights are stopped with respect to θ.The derivation says gradient-based optimization therefore has the same behavior, such as the same optimum.Source: §D.1 · pp. 27–28
  • PROB-TE uses teacher predictions to form a weighted teacher surrogate target shared across students, while PROB-MAX assigns weight only to the most confident student.The dual hard-weighting variant is related to multiple-choice learning and multi-headed networks [Lee et al., 2015].Source: §D.1 · p. 28
  • ENT weights heads using teacher entropy, while ENT-ST is the corresponding variant using student entropies.The derivation expresses these weights through a softmax over negative entropy scaled by γ.Source: §D.1 · p. 28

D.2.1 UNIFORM & PROBABILITY WEIGHTING

  • Joint convexity of unnormalized KL divergence orders UNIF and PROB losses up to a constant in θ; cross-entropy’s linearity and convexity also relate PROB to UNIF-ALL.The analysis treats weighted cross-entropy as unnormalized KL divergence up to a θ-constant.Source: §D.2.1 · p. 29

D.2.2 ENTROPY & VARIANCE WEIGHTING

  • For a discrete normalized distribution, the analysis relates entropy weighting by exp(−H) to inverse square-root variance weighting with ε = 1/12.The passage notes that a tighter bound exists but imposes stronger restrictions on the distribution.Source: §D.2.2 · p. 29

Related papers

Loading 2211.09981v3…