Weighted Ensemble Self-Supervised Learning

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

ICLR 2023

Overview

Ensembling has proven a simple yet effective method for…

Self-supervised learning (SSL) has emerged as the dominating ML paradigm

2

• Improving model accuracy (left)
• Capturing predictive uncertainty (right)

Ovadia et al. Can You Trust Your Model’s Uncertainty? Evaluating Predictive Uncertainty Under Dataset Shift. NeurIPS 2019

• … yet mostly in supervised learning!

• Ensembling for SSL is underexplored
• It is not obvious where and how to ensemble

Overview

We develop an efficient ensemble method tailored for SSL

3

• Significantly improve downstream performance, particularly for few-shot learning
• No inference cost during downstream evaluation
• Potentially applicable to a wide range of SSL methods

Method - Motivation

Many SSL methods utilize a projection head for learning better representations

4

• Examples: SimCLR, BYOL, DINO, …

Chen at al. A Simple Framework for Contrastive Learning of Visual Representations. ICML 2020

Grill et al.Bootstrap your own latent: A new approach to self-supervised Learning. NeurIPS 2020

Caron et al. Emerging Properties in Self-Supervised Vision Transformers. ICCV 2021

• Crucial for downstream performance
• Thrown away during downstream evaluation ⇒ No inference cost!

Method - Where to Ensemble?

Only ensemble the projection heads (and optionally other non-encoder parts)

5

• Ensemble heads facilitate the learning of a single encoder
• No additional inference overhead

Method - How to Ensemble?

A data-dependant weighted objective for learning the ensemble

6

• Weight each teacher/student pair to enable preferential treatment
• Different weightings are used based on the inputs

Method - How to Weight?

We explore numerous weighting schemes

7

• Uniform weighting
• Treat all teacher/student pairs equally
• Probability weighting
• Favor student head with maximum prediction confidence
• Entropy weighting
• Favor teacher head with minimum prediction entropy
• …

Experiment - Setup

We demonstrate the effectiveness of our methods with an extensive study

8

• Two state-of-the-art SSL methods
• DINO (Caron et al., 2021)
• MSN (Assran et al., 2022)
• Multiple metrics on ImageNet-1K
• Decodability: Linear or k-NN evaluation
• Label efficiency: Few-shot evaluation (1/2/5-shot, 1% data)
• Improved baselines with careful tuning for fair comparison
• A smaller projection head significantly improves few-shot performance!

Experiment - Controlled Study

Better weighting schemes encourage more diverse ensembles

Comparison of different weighting schemes with DINO ViT-S/16

Visualization of ensemble diversity

9

Experiment - Improving the SOTA

Our method consistently improves the SOTA

DINO ViT-S/16

DINO ViT-B/8

MSN ViT-S/16

Our improvements include both baseline improvements (dark) and ensembling (light)

10

• Across all setups and metrics
• Substantial gains in few-shot evaluation