Non-negative Contrastive Learning

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.14. Update codes for calculating class consistency, sparsity, and image retrieval precision in main_eval.py.
• 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. 🎉

TLDR

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.

Installation

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

Pretraining

Pretrain with the default configuration files using the following command.

CIFAR-100 / CIFAR-10

# 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.

ImageNet-100

# 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.

Linear Evaluation

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.

Full finetuning

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

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

Feature Selection

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=512

where the argument selected_dims configures the dimensions of selected features.

Consistency, Sparsity, and Image Retrieval Precision

For offline evaluating the class consistency, the sparsity and the mAP for image retrieval, run the following command:

python3 main_eval.py \
    --config-path scripts/eval/{dataset} \
    --config-name simclr.yaml \
    resume_from_checkpoint=path/to/pretrained/feature/extractor

Pretrained Checkpoints

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

	CIFAR-10	CIFAR-100	ImageNet-100
checkpoints	download	download	download

Citing this work

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},
}

Acknowledgement

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