Optimization Can Learn Johnson Lindenstrauss Embeddings

Thank you for your detailed and thoughtful review of our paper. We appreciate your recognition of the strengths and innovation of our work. We understand your concerns regarding the practical efficiency of our method and the explicit determination of its complexity. While our paper primarily focuses on the theoretical foundations, we acknowledge the importance of practical applicability. This is why we included an experiment to showcase the superior performance of our method, even with access to only first-order information, demonstrating its computational efficiency.

We understand your questions but first, we want to clarify what we mean by "our approach". Our approach includes lifting the domain to the distributional space, then optimizing the distributions’ parameters and finally obtaining a deterministic solution using variance reduction. We view this as the main contribution of this paper as we provide a different optimization paradigm and show that optimizing indirectly using our framework can be provably better than direct optimization. We expand on this further below.

Q: [...] Is there any intuitive reason why it may be possible to formulate a relaxed differentiable objective function when the embedding guarantee must hold over an uncountable number of points?

A: There are many applications of our approach in theoretical results. Often times specialized randomized constructions are used for several tasks (distance preserving in $L_2$ or other norms, $L_0$ sampling, linear-sketching or other tasks as you mentioned). Our hope is to be able to recover these results directly via optimization methods and potentially derandomizing them via variance reduction. While in the case of JL the initial distribution (Gaussian) was very simple, in other applications much more clever constructions are required. Can optimization recover these embeddings in a principled way? While the specific formulation may vary, the underlying principle of optimizing distribution parameters to achieve low variance and deterministic solutions could be adapted to other contexts.

Q: How does learning a JL embedding relate to learning embeddings for application areas discussed in the introduction? In particular, how do you see the results of this paper affecting that line of work? [...].

A: This is an interesting question. As a first point, we are optimistic that since our contributions include effectively, a new perspective and analysis, it opens the door to studying embeddings of points that have some specific structure. Secondly, in broader terms, deep learning is a general framework applicable to many domains. Frequently in deep learning, we utilize encoder architectures to generate embeddings from data, for which we want certain properties to hold. In this sense, learning a JL embedding can be viewed as a special case of an encoder, where our focus is on preserving the $L_2$ distance. Our work contributes to the deep learning field by informing the design of algorithms of this general principle. Specifically, we show that optimizing directly can lead to “bad” solutions and to overcome this one can choose to optimize in the “richer” space of distributions.

In addition, we believe our results provide a crucial link between deep learning practitioners and theorists who offer theoretical guarantees through classical/randomized algorithms. Embeddings are central to the success of neural networks, driving much of their empirical achievements. However, the lack of provable guarantees for these embeddings remains a significant challenge. By introducing provable guarantees for embeddings via optimization, alongside an algorithmic approach that is still essentially as tractable as a direct optimization approach, our work bridges these two perspectives. It allows us to harness the power of optimization while ensuring that the embeddings maintain desired properties.

This combination not only enhances the reliability of embeddings used in deep learning but also provides insights into the underlying dynamics contributing to the empirical success of deep learning models. This approach can potentially illuminate why certain embeddings work well in practice and how they can be systematically improved with theoretical backing.

Thank you for the review and valuable comments. We appreciate your recognition of the strengths of our work and have addressed your concerns below.

You raise a valid point about the computational expense of second-order methods. We should clarify here that our primary result is that second-order stationary points of the objective function in Equation 4 satisfy the Johnson Lindenstrauss (JL) guarantee. We opted to use a deterministic second-order method in order to achieve an end-to-end derandomized construction of a JL matrix.

Having established these theoretical guarantees, in our empirical evaluations, we do not insist on deterministic computation as our focus is now different, namely to demonstrate that for practical instances that have non-worst-case structure our method significantly outperforms the randomized JL construction. Instead, for simplicity, we use Monte Carlo sampling and randomized first-order methods which are known to converge to second-order stationary points with high probability [Gao et al. 2017]. In practice, second-order methods can be prohibitively expensive but are not necessary for convergence to second-order stationary points if randomness is allowed.

Q: [...] Can authors explain why the relaxation of f* to f using the probability bound in Eq (2) and invoking union bound to obtain (3) does not reduce the max to a sum.

A: The primary objective of our paper is to use optimization-based methods to find matrices that satisfy the JL guarantee, ensuring that the maximum distortion over all points does not exceed a specified threshold $\epsilon$. As we mention, this is a different (and more challenging) objective than what spectral methods like PCA control. Our goal is to show that this can be accomplished deterministically, directly using optimization. We show, however, that if one tries to do this naively, e.g., by directly working in the space of projection matrices, then significant challenges arise (to put it simply, such an approach will not work) because the landscape of the optimization in this space is non-convex, with bad local minima.

Therefore, another approach is needed. We do not change the objective -- as you correctly point out, working with the maximum is important in getting the guarantees we wish to give. Instead, we change the optimization approach, and the space in which we optimize. To do this, we work in the space of solution samplers rather than in the space of projection matrices. So at any current step in the optimization, the current "solution" of the optimizer is not a deterministic projection matrix, but rather a distribution from which one samples such a matrix. For our framework, therefore, we define a new objective $f^*$ (Equation 2) which represents the probability of generating a matrix $A$ with maximum distortion exceeding $\epsilon$. This is the probability that a given solution sampler generates a projection matrix that does not match the JL guarantee. We clarify that we do not use the union bound to reduce the maximum distortion to the sum of all distortions. Instead, we reduce $f^*$ to $f$ (Equation 3), i.e. to the sum of probabilities of each point having distortion larger than $\epsilon$ using a randomly generated matrix $A$. To make it more clear, we note that when no randomness is involved, i.e. $\sigma^2=0$, we have that $f^*(M,0) = 0$ means that no projected data point has distortion larger than $\epsilon$ and $f(M,0) = 0$ measures how many projected data points have distortion larger than $\epsilon$. Additionally, it is important to see that, $f^*(M,0) = 0$ if and only if $f(M,0) = 0$.

Ultimately, we show that our process converges to a distribution with no variance, hence, a specific projection matrix that satisfies the JL guarantee.

Thanks for the response. We answer your follow-up questions below.

Q: is there a reason why a first-order optimizer would not work and we actually need a second-order (more expensive) optimizer? First order optimizers are not necessarily stochastic / randomized.

A simple answer to this is that second-order optimization is necessary because, without it, using first-order optimization without any randomization results in the optimization getting stuck at the initial point. Generally, to avoid getting stuck throughout the optimization process, you either need second-order information or first-order information combined with randomization. Stochasticity ensures that progress can always be made.

Q: If we remove the $\sigma^2$ term from Equation (4), what estimator do we get? is it not SVD?

If we remove the $\sigma^2$ term from Equation (4), we would not obtain SVD. Equation (4) provides a bound on the probability of generating a matrix that achieves worst-case distortion larger than a specified threshold $\epsilon$. In contrast, PCA, which relies on SVD, finds an embedding that minimizes average distortion of distances between the original and the embedded vectors.

By fixing (or as you mentioned, removing) $\sigma^2$, we would essentially be looking to optimize the mean matrix $M$ for that specific $\sigma^2$. However, completely removing $\sigma^2$ from the optimization process would create an undesirable situation. While we might obtain a “better” mean matrix $M$, we would lose the intended derandomization at the end of the optimization, effectively resulting in $\sigma^2 = 0$.

Q: [...] or is it theoretically proven and related to the objective function that we are optimizing?

A: It is related to the objective function and theoretically proven, so changing the learning rate won’t resolve it. More generally, to avoid getting stuck at any point during the optimization process, we require either a second-order method or a first-order method with some randomness. We expand further on this in the next question.

Q: Can you share a reference or a more formal statement on this?

A: Our main result is that the second-order stationary points of Equation (4) satisfy the JL guarantee (Theorem 2). To achieve this, we employ a deterministic second-order method, as detailed in our paper, and thus we achieve a truly derandomized construction of a JL matrix.

Additionally, it has been shown that first-order methods with randomness also converge to second-order stationary points with high probability, see for example the work of Jin et al. (2017) on saddle points. Thus, whether using a deterministic second-order method or a first-order method with some degree of randomness, we can reach a second-order stationary point, which is precisely what is required to ensure the JL guarantee.

We hope this clarifies your questions.

Thank you for your constructive comments and for appreciating our analysis. We agree with you that there are not many results analyzing the landscape of the learned sketching matrix. We address your questions below.

Q: [...] it gives a deterministic construction of the JL lemma, or it gives a better optimization way and it works well empirically? [...].

A: The answer is both. Our main contribution is that one can use direct optimization to efficiently find "good" embeddings. We propose a novel optimization objective and analyze its performance from two distinct angles:

• First, we take a worst-case perspective where we theoretically show that optimization can obtain matrices that are always at least as good as the guarantees of the Johnson Lindenstrauss (JL) lemma. This establishes that our method is worst-case optimal, given that the JL lemma guarantees cannot be improved.
• Secondly, we demonstrate that for practical instances that have non-worst-case structure our method significantly outperforms the randomized JL construction. This is intuitively right as the randomized construction is data agnostic.

Along the way, we establish that naive methods do not have these guarantees as the space of projection matrices is non-convex and has bad local minima. Instead, we need to work in a slightly larger space of samplers, and there, our optimization method will find a deterministic solution.

In summary, our theoretical guarantees state that we can achieve the JL guarantee in the worst case, but we showcase that in practice we can go beyond that, if the data allows it.

Q: The equation (4) is about probability, and B.1 says they use Monte Carlo sampling, however, would it means that the proposed method still contains some randomness part?

A: As mentioned above our main contribution is the analysis of this novel objective. While this objective (Equation 4) involves probabilities, they have a closed form expression and can be calculated exactly without any sampling or randomness. Our main result is that second-order stationary points of this objective function satisfy the JL guarantee. Using a deterministic second-order method, we achieve a truly derandomized construction of a JL matrix.

Having established these theoretical guarantees, in our empirical evaluations, we do not insist on deterministic computation as our focus is now different (as explained in the previous question). Instead, for simplicity, we use Monte Carlo sampling and randomized first-order methods which are known to converge to second-order stationary points with high probability [Gao et al. 2017]. In practice, second-order methods can be prohibitively expensive but are not necessary for convergence to second-order stationary points if randomness is allowed.

Q: In the experiment section, the paper compares the proposed method with the JL lemma. It will make this stronger if the comparison with equation (1) is also given.

A: Thank you for the suggestion. However, as we show in Theorem 1, optimizing Equation 1 directly in the matrix space can lead to bad solutions with high distortion, whereas our method is guaranteed to succeed. Since we showed that taking the approach of Equation 1 can (provably) fail, we decided that it would not provide a fair comparison and would distract from the comparison with the randomized JL construction.

Thank you for your thoughtful feedback and for recognizing the innovation in our work.

Q: Could you explain the statement in lines 215-216? Are the values of 1/(3n) and 1/3 derived based on the chosen value of $\epsilon$?

A: Yes, that's exactly correct: you can choose $\epsilon$ appropriately to get a $1/(3n)$ probability for any data point to violate the threshold constraint. Thus, if you sum these $n$ probabilities you get $1/3$.

Thank you for your keen observations on the paper's typo and notation inconsistencies, we fixed it and made it consistent throughout.