Embedding Dimension of Contrastive Learning and $k$-Nearest Neighbors

We thank the reviewer for your thoughtful, detailed, and precise comments. Please see our reply below:

W1 Missing conclusion / future work section

Thanks, please see the global response above for a proposed conclusion / future work section.

W2 Exact $d=\\sqrt{m}$ in experiments

Thank you! We first would like to clarify that the behavior changes around $\\sqrt{m}$ per our theory, not exactly at $\\sqrt{m}$. We changed the text to avoid confusion.

We expanded experiments which accurately capture the regime around the $d=\\sqrt{m}$ threshold. In particular, we consider $m = c \\cdot d^2$ for $c=0.5, 0.75, 1, 1.5, 2, 2.5, 3$. The corresponding mean accuracies are $99\\%, 99\\%, 98\\%, 97\\%, 96\\%, 95\\%, 95\\%$.

W3 Discussion of bounds for the approximate setting

Thanks, we’ve added a discussion of the approximate setting, in which we give a full characterization of your suggested setting. It exhibits a surprising behavior: for $\\alpha \\ge 1/2$, one dimension always suffices, and for constant $\\alpha < 1/2$ it becomes as hard as the “exact” setting. We will be happy to outline it during the discussion period if you are interested.

W4 Description of contrastive learning

Thank you, we clarified that in the paper we consider one of the most common forms of contrastive learning.

W5: Consistency between using big-O notation and not

Thank you, we replaced $O(n^2)$ with $n^2$ in the footnote. Generally in this paper, we are interested in the asymptotic behavior.

W6: Resolution of Figure 3

Thank you, we added figures with higher resolutions.

W7: Experimental validation for the k-NN setting

We first would like to note that our construction for the upper bound for k-NN is theoretical - even for $k=1$ and mild $n$, the dimension significantly exceeds a trivial bound of $n$. In the general response, we show the experiments for k-NN using the setup similar to that from the paper, with the objective of reconstructing the original k-NN.

W8: Adding a line of intuition for linear independence

Thank you, we added the intuition that random perturbation is used to preserve independence.

Limitations

Thank you, we discuss the limitations in the conclusion (please refer to our general response).

Constants in O-notation

For $\\ell_2$ construction, the constants in O-notation are within a factor of 4. For the k-NN construction, the hidden constants are very large, which makes it mainly of theoretical interest. We specified in the conclusion that the k-NN construction can likely be improved.

Quadratic number of constraints

It’s indeed the case that the dimension might be as high as $\\Theta(n)$. But note that this is unavoidable in the worst case: using contrastive samples, one might recover full comparison information, and [CI24] show that, to preserve this information, $n/2$ dimensions are required in the worst case.

Technical Questions

Q1. Chernoff bound in 236 and 240

In 240, we have $L$ independent indicator random variables $A^j(x,y)$, each occurring w.p. $\\Omega(1/r)$. Denote $X = \\sum_j A^j(x,y)$. The expectation $E[X]$ is at least $E[X] = \\sum_j E[A^j(x,y)] = \\Omega(\\log{n})$. (since $L = \\Theta(r\\log{n})$). By Chernoff, $Pr[|X - E[X]| \\geq E[X]/2] < 2e^{-E[X]/12}$. In particular, $Pr[X = 0] < 2e^{-E[X]/12}$. Taking $L$ to be with sufficiently large constant, we get $Pr[X = 0] < 1/poly(n)$.

In 236, Fix $x$. For $N^-(x) = \\{y_1,...,y_{O(r)}\\}$ we define $O(r)$ independent indicator variables $X_1,...,X_{O(r)}$, where $X_i = 1$ if $y_i$ and $x$ have the same color, and $X_i = 0$ otherwise. Denote $X = \\sum_i X_i$. Notice that $E[X] = O(1)$. By Chernoff, $Pr[X > c\\log{n}] \\leq Pr[X > (1+(c/E[X])\\log{n})E[X]] < e^{-(c\\log{n}/2)} < 1/poly(n)$ (here $c$ is the constant of the load). Therefore, it holds for all $x$ w.h.p. by union bound.

Q2. Embedding dimension in Lemma 18

In the NCC embedding, we have $L = \\Theta(r \\log{n})$ blocks, where each block is a standard unit vector of dimension $\\Theta(r\\log{n})$. Intuitively, the reason each unit vector is taken from this dimension is that we need $\\Theta(r\\log{n})$ distinct standard unit vectors.

Q3. Do you require $|U(C_{j(y)}(x))| > J$?

The only thing we require in those lines is that $|U(C_{j(y)}(x))| > 1$, because we only need to pick some vector distinct from that of $y$ (and since $j(y) \\in J$, we are guaranteed all other backwards neighbors choose a different vector for this block).

Q4,Q6,Q7. In line 270, we seem to need $J=r$ rather than only $J=O(r)$. In Cor 19: $\\sum_{i=1}^{\\log{m'}} 2^i = 2(m'-1)$ and not $m'$. In Line 294, $\\gamma$ should be $\\gamma = (1-p)p$.

Thank you for pointing out these issues. These are typos arised due to refactoring the proofs. We fixed them, which did not affect correctness or the final result.

• Q4: In Lem 18, the exact distance should be $L - 1$ or $L$ (not $r-1$ or $r$ as written) - each $i \\in [L] \\setminus \\{j(y)\\}$ contributes $1$ to the distance (see lines 257-262 and 263-265).
• Q6: we fixed the sum.
• Q7: $\\gamma$ should be indeed $\\gamma = p(1-p)$. We also changed $\\gamma$ to $\\gamma/k$ in a few expressions due to this typo. This doesn't affect dependence on $k$ in the final bound.

Q4, Q6, Q7 are now fixed.

Q5. Basis of the log in Cor 19

By $\\log m'$ we denote the binary logarithm of $m'$ (i.e. $\\log_2$). We added the base of the logarithm in this place in all equations.

Q8. Repeating what $w(x,y)$ means in line 317

We added a clarification that “$w(x,y)$ is the ranking of edge $\\{x, y\\}$ if the edges are sorted by the decreasing order of distances”

Q9. Should inequalities in line 59 and fact 51 be strict?

Thank you! The inequalities between distances in the paper are strict. $\\le$ is a typo, which we fixed.

Please let us know whether we addressed your concerns, and thank you again for your comprehensive review!

Thank you very much for your detailed response, which clarified many of my issues. I have a few remaining follow-up questions / comments.

W3: This dependence on $\alpha \leq 0.5$ sounds intriguing! I'd be interested in a reference and suggest to add this for context to the background section of the paper.

W7: Can you give a rough guidance on the (n,k) regime in which your kNN bound would be an improvement over the trivial bound of $n$? I.e. how large is the constant in the big-O notation? I appreciate the additional experiments on the $k$-NN graph. Perhaps a table showing the required dimension like you provide in the general comment for $\ell_2$ embeddings would be useful.

Q1 Chernoff bounds I got the argument for line 240. Unfortunately, I am not very familiar with Chernoff bounds and for line 236, I only found the statement $P(X \geq (1+\delta)E(x))\leq e^{-\delta^2E(X)/(2+\delta)}$ online which looks similar to what you use, but is not quite the same. Could you please state the version of the Chernoff bound that you are using and a reference?

General comment on $\ell_2$ embeddings: I am surprised that the required dimension decreases with $n$ for certain $m$-values. Could you please give some intuition why this happens?

Thank you again for your comments!

W3. Could you please comment on why you can choose $m$ (or $c$) to make (*) true while maintaining $m \\le n^2$, so that the triplets are not contradictory? I also do not get the last bit on choosing $n = \\sqrt{m}$.

Let us present the construction in a slightly different way, which hopefully answers both of your questions.

Construction. Consider a set of $n$ points. On this set of $n$ points, according to Lemma 41, there exists a set of $m = {n - 1 \\choose 2} \\approx n^2 / 2$ triplets $C$, such that for every labeling of $C$ there exists a metric consistent with this labeling.

We will show that, in general, $d = \\Omega(n)$ is required so that for any labeling of $C$ there exists an embedding into a $d$-dimensional $\\ell_2$-space. Since for this construction $n = \\Theta(\\sqrt{m})$, this implies that $d = \\Omega(\\sqrt{m})$ is required in general to embed a set of $m$ non-contradictory constraints.

Proof. For the sake of contradiction, assume that there is a sufficiently large $n$ such that every labeling of the above set of triplets can be embedded using dimension $d < c n$ for some constant $c < 1$ (to be specified later). As mentioned in the previous reply, the amount of possible labelings consistent with an embedding of dimension $d$ is at most $(8 em / (nd))^{nd} \\approx (4e n / d)^{nd}$, where $\\approx$ ignores low-order terms in the exponent.

As we showed in the previous reply, not all labelings are embeddable as long as the above expression is less than $2^{\\alpha m} \\cdot {m \\choose \\alpha m} \\approx 2^{m(\\alpha - H(\\alpha))} \\approx 2^{n^2 (\\alpha - H(\\alpha)) / 2}$

When $d < cn$, we have $(4e n / d)^{nd} < (4e / c)^{c n^2} = 2^{n^2 c \\log_2 (4e/c)}$

for a sufficiently large $n$.

Then, if $c \\log_2 (4e/c) < (\\alpha - H(\\alpha)) / 2$ (such $c$ exists), we have that $(4e n / d)^{nd} < 2^{\\alpha m} \\cdot {m \\choose \\alpha m}$ for a sufficiently large $n$, implying that not all labelings of $C$ are embeddable. This proves that $d = \\Omega(n)$ is required.

Since in this construction $n = \\Theta(\\sqrt{m})$, this means that $d = \\Omega(\\sqrt{m})$ is required in general.

W7. Thanks for stating the constant! I do not mind the result being largely theoretical, especially with the additional discussion on why the poly-logarithmic dependence is surprising. Nevertheless, I encourage you to state the constant and the $n$-regime in which your bound becomes practically useful in the revision.

Thank you, we will specify the constants and add a short discussion about the constants and the practical regime at the conclusion of the paper.

Thank you for your comments.

... Whether the embedding is actually obtained ... To compute $F$, what computation do we really need?

Yes, all embedding constructions are explicit. The construction for $\\ell_2$ is in Section 2, and the construction for k-NN embeddings is in Section 3.

Here, there are no encoders.

We would appreciate it if you could please clarify what you mean by encoders here.

Can we practically set the valid $h$ beforehand

We are not exactly sure we understand this question. By $h$, we assume you mean the size parameter in the construction for $\\ell_2$. We don’t use this parameter in the k-NN construction, which your questions seems to be about.

Thanks a lot for your positive feedback, thoughtful review and the suggestions.

Complexity of Proofs

Thanks a lot for pointing this out, we’ve substantially updated the exposition by adding illustrations of the key concepts (you can find some examples in the global response to all reviewers), as well as illustrations for some of the most complicated parts in Sections 2 and 3. We will also simplify the notation in the technical sections and add the intuitive explanations. To give an example, in Section 2, the arboricity $r$ implies that there exists an order of vertices so that each vertex has $O(r)$ preceding neighbors. We use it when constructing both parts of the embedding:

• First, during construction of $F_1$, we want to control the inner products between the vertex and its preceding neighbors. This leads to at most $O(r)$ linear equations, and, by guaranteeing linear independence, it implies that $O(r)$ variables (that is, the embedding dimension of $O(r)$) is sufficient to guarantee the existence of the solution.
• Second, during construction of $F_2$, for each vertex we update a single coordinate to normalize the vector. We want to avoid changes in the inner products between neighbors, and hence for each vertex we select the coordinate different from that of any preceding neighbor. Since there are at most $O(r)$ preceding neighbors, we require only $O(r)$ coordinates.

Empirical Validation

As per your suggestion, we added a major update to the experimental section by focusing on the various settings (you can find it in the global response above). In particular, the response includes:

• $\ell_2$ construction based on Section 2;
• k-NN experiments;
• experiments on LLMs.

Please note that in Figures 3a) and 3b), we consider different values of $m$, which corresponds to a wide range of data scales, having as much as $10^6$ samples (in the global response, we consider up to $10^7$ samples).

Application for non-expert audiences

We clarified the technical sections to make them more accessible. We have added a conclusion section (see a comment above) to summarize our main contribution, which we believe should be relevant to non-experts. In particular, our results provide guidance for the choice of embedding dimension in contrastive learning and k-NN applications, which we believe can be understood without diving into the theoretical details.

Bound Tightness

The bounds are likely to be tight when the set of constraints is dense, i.e. when there exists a subset of $\\Omega(\\sqrt{m})$ elements so that the samples are focused on these elements.

For k-NNs, the lower bounds are likely to be tight in many cases, unless the dataset has some degeneracy. It is likely that the upper bounds for k-NNs are not tight, and the dependence on $k$ can be improved substantially.

Please let us know if our response addresses your concerns.

Thanks a lot for your careful review. We’ve updated the experimental section to provide a more precise validation of the theoretical bounds (please the global response posted to all the reviewers above).

Regarding the downstream applications of our embeddings, the answer to this question is two-fold:

• k-NN embeddings themselves represent the main downstream task in their context. Such embeddings are used for recommendation systems, ranking, etc. Hence, preserving Top-k results in this context is exactly aligned with the objective of the downstream task.
• The primary goal of our work is not to develop a scheme for computing embeddings better than the existing deep learning pipelines. It is our goal, however, to provide theoretical guidance for the choice of the embedding dimension in such architectures and do this in a theoretically rigorous fashion. Hence, we don’t expect the exact embedding constructions proposed in our theoretical bounds to perform better than the empirical embeddings which might be able to better capture the intrinsic structure of the data. It is only our goal to argue that the dimensionality used by these empirical constructions indeed suffices.

Regarding conclusion/future work, please refer to the global response.

Please let us know whether our response addresses your concerns.