This repostiory contains pretrained weights from and the original implementation of Improving Self-Supervised Learning by Characterizing Idealized Representations, which derives a simple uniying framework for invariant self-supervised learning (ISSL). Our framework provides actionable insights into ISSL that lead to important empirical gains such as how to:
- Simplify non-contrastive ISSL using our DISSL objective (no momentum encoders / no stop-gradients / ... )
- Choose the dimensionality of representations
- Choose the architecture of projection probes
- Choose the augmentations
The following provides the code load our ImageNet pretrained models, to reproduce our key results, and minimal notebook implementations of our DISSL 
dissl_new.mov
Our DISSL objective is a very simple non-contrastive objective that outperforms previous baselines.
We release our pretrained weights on torch hub. To load any of our model use:
import torch
model = torch.hub.load('YannDubs/Invariant-Self-Supervised-Learning:main',
'dissl_resnet50_d8192_e800_m8')Here are all available models with their respective linear probing performance on ImageNet. They are all ResNet50 trained with a batch size of 2560 and 16fp on 8 A100.
| Epochs | Dimensionality | Multi-crop | ImageNet top-1 acc. | ImageNet top-5 acc. | TorchHub name | Weights |
|---|---|---|---|---|---|---|
| 100 | 2048 | 2x224 | 66.9 | 87.5 | dissl_resnet50_dNone_e100_m2 | model |
| 100 | 8192 | 2x224 | 68.9 | 88.5 | dissl_resnet50_d8192_e100_m2 | model |
| 100 | 8192 | 2x160 + 4x96 | 70.7 | 88.5 | dissl_resnet50_d8192_e100_m6 | model |
| 400 | 2048 | 2x224 | 71.1 | 90.2 | dissl_resnet50_dNone_e400_m2 | model |
| 400 | 2048 | 2x160 + 4x96 | 73.0 | 91.3 | dissl_resnet50_dNone_e400_m6 | model |
| 400 | 8192 | 2x160 + 4x96 | 74.0 | 91.9 | dissl_resnet50_d8192_e400_m6 | model |
| 800 | 8192 | 2x224 + 6x96 | 73.9 | 91.9 | dissl_resnet50_d8192_e800_m8 | model |
For an example of how to use the pretrained models see:
We also provide a minimal DISSL implementation:
To reproduce our key TinyImageNet results you need to install ISSL (see below) and run the desired script in bin/tinyimagenet/*.sh.
To run the script without slurm use bin/tinyimagenet/*.sh -s none.
If you want to use slurm then you need to define the desired configs in config/server for an example see nlprun or vector which can be called using bin/tinyimagenet/*.sh -s <server_name>.
Installation
- Clone repository
- Install PyTorch >= 1.9
pip install -r requirements.txt.
If you have issues try installing out exact dependencies using conda env update --file environment.yaml.
For our ImageNet models we used VISSL. The exact commands can be seen on this (still uncleaned/undocumented) VISSL fork and we aim to incorporate DISSL in the main VISSL soon.
The right column in Table 1 of our paper shows empirically that DISSL outperforms DINO on TinyImageNet.
To reproduce a similar table (single seed) run bin/tinyimagenet/table1_distillation.sh -s none (no servers/slurm).
Once the script is finished you can collect and print the results by running bin/tinyimagenet/table1_distillation.sh -v "" -a is_force_gpu=False.
You should get the following results printed:
| Model | TinyImageNet Linear probing acc. |
|---|---|
| DINO | 43.3% |
| DISSL | 45.1% |
| + dim. | 48.0% |
| + epochs | 49.0% |
| + aug. | 50.7% |
Training curves: 
In our paper we characterize exactly the minimal and sufficient dimensionality depending on the probing architecture.
For linear probes it's much larger than standard dimensionalities, which suggests that one would gain important gains by increasing dimensionality.
Figure 7c of our paper shows empirically that this is indeed the case.
To reproduce a similar figure (single seed) run bin/tinyimagenet/fig7c_dimensions.sh -s none.
Once the script is finished you can collect and print the results by running bin/tinyimagenet/fig7c_dimensions.sh -v "" -a is_force_gpu=False.
The following figure will then be saved in results/exp_fig7c_dimensions/fig7c.pdf.
Training curves:
projection_heads_new.mov
In our paper, we prove that one of the two projection heads needs to have the same architecture as the dowsntream probe. This is to ensure that the SSL representations are pretrained the same way as they will be used in downstream tasks.
This is the difference between our CISSL and SimCLR.
The left column in Table 1 of our paper shows empirically that this improves performance.
To reproduce a similar table (single seed) run bin/tinyimagenet/table1_contrastive.sh -s none (no servers/slurm).
Once the script is finished you can collect and print the results by running bin/tinyimagenet/table1_contrastive.sh -v "" -a is_force_gpu=False.
You should get the following results printed:
| Model | TinyImageNet Linear probing acc. |
|---|---|
| SimCLR | 45.2% |
| CISSL | 45.8% |
| + dim. | 47.6% |
| + epochs | 48.7% |
| + aug. | 51.2% |
Training curves:
We also provide a minimal CISSL implementation:
In our paper we characterize exactly optimal sample efficiency as a function of how coarse the equivalence class induced by augmentations are.
In particular, our theory suggests that stronger label-preserving augmentations improve performance.
Figure 7a of our paper shows empirically that this is indeed the case.
To reproduce a similar figure (single seed) run bin/tinyimagenet/fig7a_augmentations.sh.
The following figure will then be saved as results/exp_fig7a_augmentations/fig7a.pdf.
Training curves: