Optimal VC Dimension of Contrastive Learning with Margin

Abstract

Contrastive learning is a successful paradigm for learning $d$-dimensional geometric representations from a collection of ``anchor--positive--negative''triplets $(i,j^{+},k^{-})$, indicating that ``item $i$ is closer to $j$ than to $k$.''Despite its success, understanding why contrastive learning leads to representations of high \textit{generalization} quality---beyond the often pessimistic predictions from PAC-learning---remains a central question. Recently, \citet*{alon2024optimal} proved that, for PAC-learning $d$-dimensional Euclidean representations of $n$-point datasets, $\Theta(\min(nd, n^2))$ triplets are necessary and sufficient, while they posed as an open question whether their VC dimension bounds for the more realistic setting of \textit{contrastive learning with a margin} can be improved. For a margin parameter $\alpha>0$, a triplet $(i,j^{+},k^{-})_{\alpha}$ is satisfied by the embedding $\phi:[n]\rightarrow \mathbb{R}^{d}$, if $\|\phi(i)-\phi(k)\|_2>(1+\alpha)\cdot\|\phi(i)-\phi(j)\|_2$. In this work, we resolve their question by proving that the VC dimension of contrastive learning under any margin $\alpha\in(0,1)$ is in fact $O(n/\alpha^2)$, improving on the previous bound of $O(n\log(n)/\alpha^2)$. We also establish that the bounds are optimal up to constant factors, by providing a matching lower bound of $\Omega(\frac{n}{\alpha^2})$ (the previously known lower bound was $\Omega(\frac{n}{\alpha})$), for $\alpha\geq \max(n^{-1/2},d^{-1/2})$.