Official PyTorch implementation of the ICLR 2024 paper Non-negative Contrastive Learning by Yifei Wang*, Qi Zhang*, Yaoyu Guo, and Yisen Wang.

Links: Wandb Logs | Checkpoints

Updates:

• 2024.04.01. Add logging of feature stats & wandb pretraining logs. With these stats, we can directly observe the differences between CL and NCL training.
• 2024.03.19. Code is released. 🎉

With non-negative constraints on contrastive features, NCL can significantly enhance feature interpretability, sparsity, and disentanglement over standard contrastive learning, while improving (at least maintaining) its performance on classical tasks.

For implementation, the difference between Contrastive Learning (CL) and Non-negative Contrastive learning (NCL) is minimal. Apart from argparser nuances, NCL only adds one line upon existing methods:

    z = torch.nn.functional.relu(z)

which applies a non-negative transformation (eg ReLU) to the output features z to enforce feature non-negativity. This is all we need to convert a standard CL method (e.g., SimCLR) to a non-negative version, and deliver all the magic.

The codebase is built upon a previous version of solo-learn (the version on Sep 27, 2022). To avoid unexpected errors, first create a Python3.8 environment, and then install the reposoity as below.

# clone the repository
git clone https://github.com/PKU-ML/non_neg
# create environment
conda create -n non_neg python=3.8
conda activate non_neg
# install dependences
cd non_neg
pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist --extra-index-url https://download.pytorch.org/whl/cu113

Pretrain with the default configuration files using the following command.

# SimCLR
python3 main_pretrain.py \
    --config-path scripts/pretrain/cifar \
    --config-name simclr.yaml
# NCL
python3 main_pretrain.py \
    --config-path scripts/pretrain/cifar \
    --config-name ncl.yaml

The default setting is for CIFAR-100. For CIFAR-10, override data.dataset=cifar10. Meanwhile, change the experiment name accordingly to avoid collision, e.g., name=simclr-resnet18-cifar10-ncl. Other experiments follow the same setting.

# SimCLR
python3 main_pretrain.py \
    --config-path scripts/pretrain/imagenet-100 \
    --config-name simclr.yaml
# NCL
python3 main_pretrain.py \
    --config-path scripts/pretrain/imagenet-100 \
    --config-name ncl.yaml

By default, we use non_neg=rep_relu for CIFAR-10 and CIFAR-100 and non_neg=relu for ImageNet-100.

After that, for linear evaluation, run the following command:

python3 main_linear.py \
    --config-path scripts/linear/{dataset} \
    --config-name simclr.yaml \
    pretrained_feature_extractor=path/to/pretrained/feature/extractor

Here dataset={cifar,imagenet100}. We use the argument pretrained_feature_extractor to configure the path of the pretrained checkpoints.

And for fine-tuning evaluation, run the following command:

python3 main_linear.py \
    --config-path scripts/finetuning/{dataset} \
    --config-name simclr.yaml

And for offline linear probing with selected dimensions, run the following command:

python3 main_linear.py \
    --config-path scripts/selected \
    --config-name simclr.yaml \
    selected_dims=256

where the argument selected_dims configures the dimensions of selected features.

The following table provides the pre-trained checkpoints for CL and NCL.

If you find the work useful, please cite the accompanying paper:

@inproceedings{
wang2024nonnegative,
title={Non-negative Contrastive Learning},
author={Yifei Wang and Qi Zhang and Yaoyu Guo and Yisen Wang},
booktitle={ICLR},
year={2024},
}

Our codes borrow the implementations of SimCLR in the solo-learn repository: https://github.com/vturrisi/solo-learn