Generalized Contrastive Alignment (GCA) - NeurIPS 2024

Generalized Contrastive Alignment (GCA) provides a robust framework for self-supervised learning tasks, supporting various datasets and augmentation methods.

This is the repo for the paper:

Your Contrastive Learning Problem is Secretly a Distribution Alignment Problem

---

Environment Setup

To set up the required environment, follow these steps:

# Create and activate the environment
conda create -n GCA python=3.11.9 -y
conda activate GCA
# Install dependencies
pip install hydra-core numpy==1.26.4 matplotlib seaborn scikit-image scikit-learn \
pytorch-lightning==1.9.5 torch==2.2.1 torchaudio==2.2.1 \
torchmetrics==1.4.2 torchvision==0.17.1

---

References

• SimCLR CIFAR10 Implementation by Damrich et al.: GitHub Link
• SimCLR by Ting Chen et al.: GitHub Link
• IOT in Liangliang Shi, et al. "Understanding and generalizing contrastive learning from the inverse optimal transport perspective." ICML, 2023.

---

Usage Instructions

Supported Tasks

The framework supports the following tasks:

• simclr
• hs_ince
• gca_ince
• rince
• gca_rince
• gca_uot

Supported Datasets

The following datasets are supported:

• SVHN
• imagenet100
• cifar100
• cifar10

Strong Data Augmentation

You can configure data augmentation using the strong_DA option:

• None (standard augmentation)
• large_erase
• brightness
• strong_crop

---

Pretraining a Model

To pretrain a model using self-supervised learning, run the following script:

python ssl_pretrain.py \
--config-name "simclr_cifar10.yaml" \
--config-path "./config/" \
task=gca_uot \
dataset=CIFAR10 \
dataset_dir="./datasets" \
batch_size=512 \
seed=32 \
backbone=resnet18 \
projection_dim=128 \
strong_DA=None \
gpus=1 \
workers=16 \
optimizer='SGD' \
learning_rate=0.03 \
momentum=0.9 \
weight_decay=1e-6 \
lam=0.01 \
q=0.6 \
max_epochs=500 \
r1=1 \
r2=0.2

---

Linear Evaluation

To evaluate the pretrained model with a linear classifier, use the following script:

python linear_evaluation.py \
--config-name="simclr_cifar10.yaml" \
--config-path="./config/" \
task=gca_uot \
dataset=cifar10 \
batch_size=512 \
seed=64 \
backbone=resnet18 \
projection_dim=128 \
strong_DA=None \
lam=0.01 \
q=0.6 \
load_epoch=500

---

Configuration Options

• Task: Specify the self-supervised learning task (e.g., gca_uot).
• Dataset: Choose from supported datasets (e.g., cifar10).
• Data Augmentation: Use strong_DA to set augmentation type.
• Training Parameters:
  • batch_size: Batch size for training.
  • backbone: Backbone architecture (e.g., resnet18).
  • projection_dim: Dimension of the projection head.
  • lam and q: Regularization and scaling parameters.
  • max_epochs: Maximum number of epochs for training.

---

Notes

• Ensure that the dataset_dir contains the datasets in the correct structure.
• Customize parameters in the scripts to fit your experimental needs. As an example for SVHN,
• We modify the linear_evaluation to make it able to run on new pytorch version 2.9.0.

python ssl_pretrain.py \
--config-name "simclr_svhn.yaml" \
--config-path "./config/" \
task=gca_uot \
dataset=SVHN \
dataset_dir="./datasets" \
batch_size=512 \
seed=48 \
backbone=resnet18 \
projection_dim=128 \
strong_DA=None \
gpus=1 \
workers=16 \
optimizer='Adam' \
learning_rate=0.03 \
momentum=0.9 \
weight_decay=1e-6 \
lam=0.01 \
q=0.6 \
max_epochs=500 \
r1=1 \
r2=0.01

Citation

If you find this repository helpful for your research, please cite our paper:

@article{chen2025your,
title={Your contrastive learning problem is secretly a distribution alignment problem},
author={Chen, Zihao and Lin, Chi-Heng and Liu, Ran and Xiao, Jingyun and Dyer, Eva},
journal={Advances in Neural Information Processing Systems},
volume={37},
pages={91597--91617},
year={2025}
}