Self-Supervised Learning with Kernel Dependence Maximization

This is the code for SSL-HSIC, a self-supervised learning loss proposed
in the paper Self-Supervised Learning with Kernel Dependence Maximization

Using this implementation should achieve a top-1 accuracy on Imagenet around
74.8% using 128 Cloud TPU v2/3.


To set up a Python3 virtual environment with the required dependencies, run:

python3 -m venv ssl_hsic_env
source ssl_hsic_env/bin/activate
pip install --upgrade pip
pip install -r ssl_hsic/requirements.txt



For pre-training on ImageNet with SSL-HSIC loss:

mkdir /tmp/ssl_hsic
python3 -m ssl_hsic.experiment \
--config=ssl_hsic/ \

This is going to pre-train for 1000 epochs. Change config to
for testing purpose. See
jaxline documentation for more
information on jaxline_mode.

If save_dir is provided in, the last checkpoint is saved and can
be used for evaluation.

Linear Evaluation

For linear evaluation with the saved checkpoint:

mkdir /tmp/ssl_hsic
python3 -m ssl_hsic.eval_experiment \
--config=ssl_hsic/ \

This is going to train a linear layer for 90 epochs. Change config to for testing.

Citing this work

If you use this code in your work, please consider referencing our work:

  title={Self-Supervised Learning with Kernel Dependence Maximization},
  author={Yazhe Li and Roman Pogodin and Danica J. Sutherland and Arthur Gretton},
  booktitle={Thirty-Fifth Conference on Neural Information Processing Systems},


This is not an official Google product.


View Github