Optimal VC Dimension of Contrastive Learning with Margin

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 $\mathit{\text{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, $\mathrm{\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 $\mathit{\text{contrastive}}\mathit{\text{learning}}\mathit{\text{with}}\mathit{\text{a}}\mathit{\text{margin}}$ can be improved. For a margin parameter $\alpha >0$, a triplet $(i,j^{+},k^{-}{)}_{\alpha }$ is satisfied by the embedding $\phi :[n]\to {\mathbb{R}}^d$, if $\mathrm{\parallel }\phi (i)-\phi (k){\mathrm{\parallel }}_2>(1+\alpha )\cdot \mathrm{\parallel }\phi (i)-\phi (j){\mathrm{\parallel }}_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\mathrm{/}{\alpha }^2)$, improving on the previous bound of $O(n\log (n)\mathrm{/}{\alpha }^2)$. We also establish that the bounds are optimal up to constant factors, by providing a matching lower bound of $\mathrm{\Omega }(\frac{n}{{\alpha }^2})$ (the previously known lower bound was $\mathrm{\Omega }(\frac{n}{\alpha })$), for $\alpha \ge \max (n^{-1\mathrm{/}2},d^{-1\mathrm{/}2})$.