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})$.

Data Structures and Algorithms (cs.DS); Machine Learning (cs.LG); Machine Learning (stat.ML)

arXiv:2609.38834 [cs.DS]

(or arXiv:2609.38834v1 [cs.DS] for this version)

https://doi.org/10.48550/arXiv.2609.38834 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.