This is the PyTorch implementation for our NeurIPS 2023 paper.

Junkang Wu, Jiawei Chen, Jiancan Wu, Wentao Shi, Xiang Wang & Xiangnan He. 2023. Understanding Contrastive Learning via Distributionally Robust Optimization. arxiv link

The implementation only requires a small modification to the InfoNCE code.

# pos     : exp of inner products for positive examples
# neg     : exp of inner products for negative examples
# N       : number of negative examples
# t       : temperature scaling
# mu      : center position
# sigma   : height scale

#InfoNCE
standard_loss = -log(pos.sum() / (pos.sum() + neg.sum()))

#ADNCE
weight=1/(sigma * sqrt(2*pi)) * exp( -0.5 * ((neg-mu)/sigma)**2 )
weight=weight/weight.mean()
Adjusted_loss = -log(pos.sum() / (pos.sum() + (neg * weight.detach() ).sum()))

The code can be found in image

The code can be found in sentence

The code can be found in graph

If you find this repo useful for your research, please consider citing the paper

@inproceedings{wu2023ADNCE,
  author = {Junkang Wu and Jiawei Chen and Jiancan Wu and Wentao Shi and Xiang Wang and Xiangnan He},
  title = {Understanding Contrastive Learning via Distributionally Robust Optimization},
  booktitle = {NeurIPS},
  year = {2023}
}

For any clarification, comments, or suggestions please create an issue or contact me ([email protected]).