Learning equivariant tensor functions with applications to sparse vector recovery

W1.1: The paper focuses heavily on mathematical formulas and definitions but lacks intuitive explanations to clarify the need for equivariant tensor functions and their practical motivations.

Suggestion: Consider adding a high-level overview section that explains why equivariance is essential in certain applications and provides examples that demonstrate the demand for equivariant tensor functions in real-world scenarios.

Tensor functions appear in physical applications such as: polarization maps in cosmic microwave backgrounds, stress tensors in material sciences, and potential vorticity (a pseudoscalar or 0(−) -tensor field) in the Earth’s atmosphere (see Figure 1 of [1]). The equivariance appears because all these applications coming from physics have a coordinate-free description, and therefore are equivariant with respect to coordinate transformations.

[1] Gregory, W., Hogg, D. W., Blum-Smith, B., Arias, M. T., Wong, K. W., & Villar, S. (2023). GeometricImageNet: Extending convolutional neural networks to vector and tensor images. arXiv preprint arXiv:2305.12585.

W1.2: Key theoretical challenges and the novel aspects of the proposed method are not thoroughly discussed, leaving readers uncertain about the primary obstacles and contributions within the framework.

Suggestion: Consider adding a dedicated “Contributions” section that explicitly outlines how this method differs from existing approaches, helping readers understand the unique contributions of the framework.

See the section entitled “Our contributions” at the end of page 2.

W1.3: The paper does not clearly explain the necessity of "learning equivariant tensor functions," lacking specific application scenarios that showcase the unique benefits and suitability of these functions. It remains unclear why certain types of tensors in applications require equivariance.

Suggestion: Include examples of practical applications where specific types of tensors need to be equivariant and clarify how these scenarios validate the advantages of using equivariant tensor functions.

The application of sparse vector recovery is one example where one needs to learn an O(d)-equivariant tensor function. Other examples are the ones described in the answer of your question W1.1.

W2.1: The paper centers its applications on sparse vector recovery, a problem that is primarily theoretical and rooted in computer science, with limited overlap with current mainstream machine learning applications, such as computer vision or natural language processing.

Suggestion: Discuss the potential connections between sparse vector recovery and mainstream machine learning applications, such as compressed sensing, feature selection, or dimensionality reduction in deep learning, to enhance the method’s appeal. Assignee:

The sparse vector recovery problem has applications to dictionary learning, representation learning, and computer vision tasks. Section VI of [1] gives a thorough discussion of many applications of this problem.

[1] Qu, Q., Zhu, Z., Li, X., Tsakiris, M. C., Wright, J., & Vidal, R. (2020). Finding the sparsest vectors in a subspace: Theory, algorithms, and applications. arXiv preprint arXiv:2001.06970.

W2.2: The paper does not demonstrate if or how sparse vector recovery can extend to practical uses in deep learning or other prominent areas within machine learning, lacking case studies or examples of its effectiveness in real-world tasks.

Suggestion: Consider adding relevant cases or discussions on how the method could be integrated into neural network architectures or applied to tasks such as network pruning, compression, or efficient inference, to illustrate its practicality.

This is out of the scope of this paper.

W2.3: The real-world significance of the proposed method is unclear, potentially making it seem disconnected from practical contributions to machine learning, which could limit its appeal to a broader audience.

Suggestion: To increase its appeal, the authors could discuss the broader implications of their work, such as how advances in equivariant functions might contribute to more efficient or interpretable machine learning models over the long term.

Equivariant machine learning is a well-established field, with plenty of applications, including AI4physics, AI4science (drug design and protein folding as the most important applications), and graph neural networks. The specific contribution here is theoretical.

Limited Scope of Experiments: One of the main weaknesses of the paper is the limited scope of the experimental evaluation. The experiments are primarily focused on the sparse vector estima- tion problem, which, while important, represents a rather narrow application domain. The method should have broader applications, such as in physics and general relativity, but these areas are not sufficiently discussed or explored.

We agree with the reviewer that there are many exciting applications for our model in physics and general relativity, and we look forward to pursuing them in future work. Our main goal in this work was to present a strong theoretical framework along with experiments to show the utility of our methods.

Clarity: The paper is difficult to follow due to disorganized notation and unclear variable definitions. The use of coordinates is somewhat messy, which makes the math- ematical developments hard to track. Variables are often introduced without proper explanation or context, causing confusion. For example:

• In Definition 4, the indices need more explanation.

• In Theorem 1, how each tensor alk contracts with certain indices (dimen- sions) of $c_{l_1, \ldots, l_r}$ should be explained more clearly.

• In Example 1, there should be more explanation of why only the generic elements of the G4 group are used while others are contracted.

• In Lemma 1, the meaning and constraints of ασ need clarification. Moreover, the progression from theoretical concepts to experimental appli- cations lacks smoothness. The presentation could be improved by organizing notation more systematically and clearly defining all variables.

Thank you for highlighting these confusing areas, we have added details to clarify the meaning. In particular, we have included a brief example following Definition 4, and we have added an explanation of Theorem 1 to help give an intuitive understanding. Following this example, we have referenced Appendix D justifying why the generic elements of the G4 group suffice. In Lemma 1, the values $\alpha_\sigma$ and $\beta_\sigma$ can take any real value which we clarified in the text. To strengthen the connection between the theory and the experiments, we added the full derivation of our SparseVectorHunter model (Equation (28)) from Corollary 1 to the Appendix.

The statement of Theorem 1 introduces a complex expression involving linear combinations over isotropic tensors. Could the authors provide more intuitive explanations or mathematical nature of this theorem?

Think of each term Theorem 1 as combining r of the input tensors with the tensor product, then mapping them to the appropriate output tensor order and parity with a linear map. However, since a linear map from a $k(p)$-tensor to a $k'(p')$-tensor can always be written as a tensor product with a $(k+k’)(pp’)$-tensor followed by a $k$-contraction, that is what we do where $c_{l_1,\ldots,l_r}$ is that $(k+k’)(pp’)$-tensor. Additionally, since the function must be $O(d)$-equivariant, the new tensor that we introduce $c_{l_1,\ldots,l_r}$ must be $O(d)$-isotropic. The sums in the expression then allow us to add up all combinations of $r$ input tensors for $r=0$ to $r=R$. The theorem merely says that this polynomial is enough to construct all tensor equivariant polynomials. We will add an explanation like this to the paper.

The authors state that parameterizing all permutation-invariant polyno- mial functions may be as challenging as solving the graph isomorphism problem. Could the authors provide an estimate of the computational complexity or approximate methods for tackling this problem? can we push this into high-dimensional settings?

Just to clarify, in Remark 3 we meant that finding all polynomial functions in $n \times n$ matrices that are invariant with respect to simultaneous permutations of rows and columns will distinguish graphs up to isomorphism. If one considers $d$-dimensional point clouds of n points $P \in \mathbb R^{d\times n}$ and asks for polynomials that are simultaneously equivariant with respect to orthogonal transformations O(d) and permutations of the n points, one obtains polynomials on the Gram matrix $P^\top P$ that are invariant with respect to simultaneous permutations of rows and columns. We don’t address this problem in this paper, we say that a possible interesting extension would be to address this problem using the formulation of Corollary 1. A recent paper [1] provides a characterization of such polynomial functions (and more!) using different techniques. It also addresses the high computational complexity of the problem and restricts the machine learning models to low-degree polynomials due to this issue.

[1] Puny, O., Lim, D., Kiani, B., Maron, H., & Lipman, Y. (2023, July). Equivariant polynomials for graph neural networks. In International Conference on Machine Learning (pp. 28191-28222). PMLR.

If I understand the proposed approach to the sparse vector recovery problem, for the approach to work one actually need to observe many data samples to be able to train neural networks that define coefficients for the unknown polynomial $h: L \to v_1$. This may be a strong assumption in many applications and is different from SOS-based methods that can recover $v_1$ from observing only one subspace $L$ and does not require to go through the learning process.

While SOS-based methods only require a single subspace, they make strong distributional assumptions on the input data. That's the reason why it doesn't perform well on the MNIST example, since the MNIST data doesn't satisfy the iid spherical gaussian assumption. By contrast, the machine learning approach allows us to learn a data-driven method that can adapt to a wider variety of settings and perform better. Thus our model is more robust to covariance model misspecification.

The proof of the main theorem follows standard steps for this sort of problems: problem is easily reduced to homogenous degree $r$ monomials and for homogeneous degree monomials the characterization is obtained by using a somewhat standard group averaging technique. In this sense, the proof of the main result is somewhat "standard" for the literature.

We note that our proofs follow the usual pathways in invariant theory. However, we prove the results because we didn’t find them in the given form anywhere in the literature. Moreover, the results for real-valued invariants are particularly scarce in the literature that focuses mainly on complex-valued invariants.

Can you provide a bit more insight into what are the requirements for the training set for your method to be able to train the neural networks defining coefficients in the parametrization of $h$ in Eq 28? Do you need to assume that all training samples have the same planted sparse vector ? Or is it sufficient to assume that all such vectors are coming from some distribution with nice properties?

For the synthetic data experiments, we pick a sampling scheme for the sparse vector (accept/reject, Bernoulli-Gaussian, corrected Bernoulli-Gaussian, or Bernoulli-Rademacher), a covariance matrix type (random, diagonal, or identity), then sample the covariance matrix if it is not the identity. Then for each data point, we sample the sparse vector in the manner prescribed and sample the noise vectors from a normal distribution with zero mean and the covariance matrix. So all the planted sparse vectors are different but sampled in the same way, and all the noise vectors are different but have the same covariance. We clarified some of these details in the text.

How many training samples do you anticipate needing to recover $n$-dimensional planted sparse vector?

While the sample complexity of this problem remains an open question, recent work [1] proves that imposing the symmetries of the problem reduces the dimension of the input space by the dimension of the group. For the sparse vector recovery problem, that would reduce the input dimension from $nd$ to $nd - d(d-1)/2$. For the synthetic data experiments we use 5,000 training points and for MNIST we use 900 training points.

[1] Behrooz Tahmasebi and Stefanie Jegelka. The exact sample complexity gain from invariances for kernel regression. Advances in Neural Information Processing Systems, 36, 2023.

Are there some applications about your proposed problem to recover a sparse vector?

Recovering a sparse vector is an important subtask in many applications in machine learning, such as computer vision, representation learning, and scientific imaging. See section VI of [1] for an overview.

[1] Qing Qu, Zhihui Zhu, Xiao Li, Manolis C. Tsakiris, John Wright, and Ren´e Vidal. Finding the sparsest vectors in a subspace: Theory, algorithms, and applications. arXiv preprint arXiv:2001.06970, 2020.

The representation of h is only related to matrices, but not a complex tensor. Your definitions and analysis are all based on tensors.

This is correct, the theory we develop in sections 3 and 4 is for general tensor functions of any parity and any order that are equivariant to the orthogonal group, the indefinite orthogonal group, and the symplectic group. We demonstrate the technique with one application where the function has 1-tensor (vector) inputs, a 2-tensor (matrix) output, and is equivariant to the orthogonal group. This is one particular case of our theoretical results, but there are many possible other cases that we plan on exploring in future work.

Can you give a details how to get (28) from Corollary 1?

Thank you for pointing out this gap in exposition. We have added a full derivation in the Appendix, and referenced it from the main text. In short, we write out Corollary 1 for output order 2 which allows us to break the sum over t and the sum over σ into separate cases and simplify the whole expression. We then swap from tensor notation to standard matrix and vector notation to make the equations clearer for readers who are primarily interested in the application.

Why you set h to be a equivariant function? What is the advantages?

We set $h$ to be an equivariant function due to the inherent symmetry of the problem, which offers an advantage over non-equivariant ML methods by reducing the sample complexity. Concretely for Problem 1, the input can be any orthonormal basis of span$(v_0, …, v_{d-1})$ for a given sparse vector $v$. In the language of equivariance/invariance, this means that the function that estimates the sparse vector $\hat{v}$ in equation (25) must be invariant to $O(d)$. A sufficient condition for this $O(d)$-invariance is that $h$ is $O(d)$-equivariant (as we show in Appendix G), and we conjecture that this is a necessary condition as well. Thus $h$ is an $O(d)$-equivariant function in the problem, so by learning an explicitly equivariant function, we hope to outperform non-equivariant learned methods. Our hopes are confirmed by our numerical experiments where our equivariant learned models do consistently and often significantly better than non-equivariant learned models.

In the numerical experiments about mnist, do you get the similar results if you choose $v_1,...,v_{d-1}$ to be some other vectors which is not so related to the sparse vector?

We explore this scenario with the Random Subspace (RND) noise scheme, and we do see similar results. We set the original MNIST digit to be $v_0$, and we sample $v_1,...,v_{d-1}$ from a normal distribution with zero mean and identity covariance. This experimental setup is similar to the synthetic data experiments, and the $v_1,...,v_{d-1}$ are indeed unrelated to the sparse vector. In Table 2 we see that our model SVH (and SVH-Diag) both modestly outperform the SOS methods.

How to get the orthonormal basis for the spanned subspace? Do you get the similar results if you choose different orthonormal basis?

The orthonormal basis is different every time because it is randomly generated. We construct the $n \times d$ matrix $B$ whose columns are the vectors $v_0,...,v_{d-1}$, then we right multiply by a random orthogonal matrix $O$ before taking the QR decomposition. The matrix $Q$ of the decomposition is an orthonormal basis of the spanned subspace as we prove in Appendix I.4. We have added a pointer to this Appendix to the main text.