Tight PAC-Bayesian Risk Certificates for Contrastive Learning

TL;DR

This paper develops non-vacuous PAC-Bayesian risk certificates for contrastive learning, specifically addressing the dependence issues in SimCLR and providing tighter bounds verified on CIFAR-10.

Contribution

It introduces novel PAC-Bayesian bounds tailored for contrastive learning frameworks like SimCLR, accounting for data dependence and augmentation effects.

Findings

Tighter risk certificates for contrastive loss and downstream tasks.

Effective handling of dependence induced by positive pair reuse in SimCLR.

Experimental validation on CIFAR-10 demonstrating improved bounds.

Abstract

Contrastive representation learning is a modern paradigm for learning representations of unlabeled data via augmentations -- precisely, contrastive models learn to embed semantically similar pairs of samples (positive pairs) closer than independently drawn samples (negative samples). In spite of its empirical success and widespread use in foundation models, statistical theory for contrastive learning remains less explored. Recent works have developed generalization error bounds for contrastive losses, but the resulting risk certificates are either vacuous (certificates based on Rademacher complexity or $f$-divergence) or require strong assumptions about samples that are unreasonable in practice. The present paper develops non-vacuous PAC-Bayesian risk certificates for contrastive representation learning, considering the practical considerations of the popular SimCLR framework. Notably, we take into account that SimCLR reuses positive pairs of augmented data as negative samples for other data, thereby inducing strong dependence and making classical PAC or PAC-Bayesian bounds inapplicable. We further refine existing bounds on the downstream classification loss by incorporating SimCLR-specific factors, including data augmentation and temperature scaling, and derive risk certificates for the contrastive zero-one risk. The resulting bounds for contrastive loss and downstream prediction are much tighter than those of previous risk certificates, as demonstrated by experiments on CIFAR-10.