Summary
This work shows that a deterministic optimization procedure can find a matrix $A$ that satisfies the Johnson Lindenstrauss guarantee. That is, a matrix $A$ maps a set of $n$ vectors to a lower dimensional space while preserving all pairwise distances up to some chosen multiplicative distortion. Typically, $A$ is constructed by sampling it from a random matrix distribution with i.i.d. entries. The authors prove that attempting to directly optimize the entries of $A$ through an optimization procedure by minimizing the maximum distortion is prone to being stuck at local minima. However, the authors show that by optimizing the mean of each entry and the entry-wise variance of the distribution $A$ is sampled from, one can maintain a fixed probability of $A$ being a JL-embedding while at the same time guaranteeing the entry-wise variance ‘sigma’ goes to zero. They then show that, when ‘sigma’ is sufficiently small, one may use the optimized expectation of $A$ as the embedding matrix while only slight increasing the maximum distortion, thereby deterministically finding the desired JL embedding matrix $A$. They show that $\rho$-SOSPs (second order stationary points) have sufficiently low variance when $\rho$ is small, and finally show that a method for finding $\rho$-SOSPs suffices to solve the designed optimization problem.
Strengths
Overall, the paper is clearly written and well-motivated. The intuition of the approach and analysis is easy to follow.
The key idea of optimizing the parameters of a random matrix distribution to preserve the JL-embedding property while reducing the entry-wise variance seems like an innovative approach. The authors point out the original space of matrices is contained in this larger probabilistic space, since a deterministic matrix $A$ is equivalent to having mean $A$ and zero variance. Hence, this can be seen as a probabilistic relaxation of the original matrix optimization problem. I have not seen this type of relaxation used in the field of matrix sketching or more generally randomized numerical linear algebra before, and I believe it may be useful for other problems in the area. I am not very familiar with diffusion models, so I cannot speak on the novelty of the approach regarding that area.
The empirical results are also strong in the sense that they show this procedure for constructing a JL embedding tends to achieve a much lower distortion factor than randomized constructions for a fixed dimension.
Weaknesses
The iterative method to find the matrix $A$ takes $\operatorname{poly}(n, k, d)$ steps, i.e., the complexity is proven to be polynomial but not explicitly determined. Since the paper is primarily theoretical with only limited experiments, it is unclear how efficient this method is in practice.
While the results seem very interesting theoretically, the paper could be strengthened by pointing out some practical applications where this improved deterministic JL embedding would be useful. In the applications I am familiar with, oblivious JL embeddings are needed due to the large number of points in the high-dimensional space (e.g., preserving k-means loss). The authors point to embeddings in deep learning as motivation. It is unclear to me as to how the authors expect progress in understanding deterministic JL embeddings to relate to these embeddings in deep learning. Additional clarification of this point would be helpful.
Questions
In the conclusion, you mention the potential for this approach in applications beyond the Johnson Lindenstrauss setting. In your approach for the JL setting, you upper bound the failure probability of the distortion guarantee via the union bound in eqn. (3). This formulation of the objective function seems difficult to translate to other sketching guarantees on $A$ (e.g., projection cost preservation, L2-subspace embedding, affine embedding). 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?
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? As mentioned above, I think it would be helpful to expand on the link between your result and the motivation of deep learning embeddings given in the intro.